Improving LLM Final Representations with Inter-Layer Geometry
Not on the map. Why?
Abstract (arXiv)
The standard in LLM-based prediction is to use the final-layer representation as the input to a downstream predictor. However, intermediate layers may encode complementary task-relevant signals. Existing approaches therefore either search for the best layer for each task or apply expensive attention-based mechanisms to learn inter-layer aggregation. In this work, we first show that such complexity is unnecessary: a lightweight Graph Neural Network over a fully connected graph of LLM layers is more efficient and achieves significantly stronger predictive performance than existing approaches. We then introduce the Cayley-Encoder, which further improves both efficiency and predictive performance by replacing the fully connected graph with a Cayley graph over SL(2, Zn). These Cayley graphs provide a mathematically grounded topology that is sparse, regular by construction, and has low diameter. This enables effective communication across layers while constraining the aggregation structure and reducing the risk of GNN overfitting. In an evaluation of Cayley-Encoder across 13 tasks and 9 LLMs, Cayley-Encoder consistently outperforms baselines, achieving improvements of up to 40 percentage points in accuracy, while introducing at most 0.1% additional parameters relative to the LLM size. We further show that Cayley-Encoder is effective in few-shot regimes. Finally, we show that Cayley-Encoder outperforms LoRA fine-tuning while operating on the frozen LLM. We conclude with an explainability analysis showing that multiple layers contribute meaningfully to the final prediction, supporting our hypothesis.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 31 of its 45 references that are in the map. First posted on arXiv 24 Mar 2026.
Builds on
31 of 45 references.
References · 31
- 2017Attention Is All You Needcited by 59,932
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
- 2014Adam: A Method for Stochastic Optimizationcited by 36,960
- 2020Language Models are Few-Shot Learnerscited by 32,199
- 2016Semi-Supervised Classification with Graph Convolutional Networkscited by 11,282
- 2019Sentence-BERT: Sentence Embeddings using Siamese BERT-Networkscited by 9,049
- 2018How Powerful are Graph Neural Networks?cited by 4,355
- 2017Neural Message Passing for Quantum Chemistrycited by 3,460
- 2019Optuna: A Next-generation Hyperparameter Optimization Frameworkcited by 2,455
- 2018Deep contextualized word representationscited by 2,285
- 2023Pythia: A Suite for Analyzing Large Language Models Across Training and Scalingcited by 1,734
- 2017Deep Setscited by 1,394
19 more
- 2019BERT Rediscovers the Classical NLP Pipelinecited by 1,046
- 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 994
- 2019GNNExplainer: Generating Explanations for Graph Neural Networkscited by 735
- 2020On the Bottleneck of Graph Neural Networks and its Practical Implicationscited by 624
- 2022MTEB: Massive Text Embedding Benchmarkcited by 623
- 2019Reducing Transformer Depth on Demand with Structured Dropoutcited by 434
- 2020Efficient Intent Detection with Dual Sentence Encoderscited by 419
- 2025Layer by Layer: Uncovering Hidden Representations in Language Modelscited by 293
- 2020MTOP: A Comprehensive Multilingual Task-Oriented Semantic Parsing Benchmarkcited by 156
- 2024Not All Layers of LLMs Are Necessary During Inferencecited by 95
- 2022Expander Graph Propagationcited by 75
- 2024DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averagingcited by 48
- 2020Investigating Societal Biases in a Poetry Composition Systemcited by 47
- 2021BERTnesia: Investigating the capture and forgetting of knowledge in BERTcited by 40
- 2025MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connectionscited by 27
- 2023Graph Neural Networks Use Graphs When They Shouldn'tcited by 25
- 2024Cayley Graph Propagationcited by 21
- 2022Depth-Wise Attention (DWAtt): A Layer Fusion Method for Data-Efficient Classificationcited by 9
- 2026Towards Improved Sentence Representations using Token Graphscited by 1
Built on by
Citing papers per year: 2026: 1.
1 paper.