# Improving LLM Final Representations with Inter-Layer Geometry

arXiv: https://arxiv.org/abs/2603.22665 · HTML: https://mlsanity.com/paper/2603.22665/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 45 references that are in the map.
- First posted on arXiv 24 Mar 2026.

## Builds on (most cited first, 30 of 31)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://mlsanity.com/paper/1908.10084/) (2019)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [BERT Rediscovers the Classical NLP Pipeline](https://mlsanity.com/paper/1905.05950/) (2019)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [GNNExplainer: Generating Explanations for Graph Neural Networks](https://mlsanity.com/paper/1903.03894/) (2019)
- [MTEB: Massive Text Embedding Benchmark](https://mlsanity.com/paper/2210.07316/) (2022)
- [On the Bottleneck of Graph Neural Networks and its Practical Implications](https://mlsanity.com/paper/2006.05205/) (2020)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Efficient Intent Detection with Dual Sentence Encoders](https://mlsanity.com/paper/2003.04807/) (2020)
- [Layer by Layer: Uncovering Hidden Representations in Language Models](https://mlsanity.com/paper/2502.02013/) (2025)
- [MTOP: A Comprehensive Multilingual Task-Oriented Semantic Parsing Benchmark](https://mlsanity.com/paper/2008.09335/) (2020)
- [Not All Layers of LLMs Are Necessary During Inference](https://mlsanity.com/paper/2403.02181/) (2024)
- [Expander Graph Propagation](https://mlsanity.com/paper/2210.02997/) (2022)
- [DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging](https://mlsanity.com/paper/2402.02622/) (2024)
- [Investigating Societal Biases in a Poetry Composition System](https://mlsanity.com/paper/2011.02686/) (2020)
- [BERTnesia: Investigating the capture and forgetting of knowledge in BERT](https://mlsanity.com/paper/2106.02902/) (2021)
- [MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections](https://mlsanity.com/paper/2502.12170/) (2025)
- [Graph Neural Networks Use Graphs When They Shouldn't](https://mlsanity.com/paper/2309.04332/) (2023)
- [Cayley Graph Propagation](https://mlsanity.com/paper/2410.03424/) (2024)
- [Depth-Wise Attention (DWAtt): A Layer Fusion Method for Data-Efficient Classification](https://mlsanity.com/paper/2209.15168/) (2022)

## Built on by (most cited first, 1 of 1)

- [Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management](https://mlsanity.com/paper/2605.06472/) (2026)

## 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.