Improving LLM Final Representations with Inter-Layer Geometry

Tom Ulanovski; Eyal Blyachman; Maya Bechler-Speicher

cs.CLMar 24, 2026

Improving LLM Final Representations with Inter-Layer Geometry

Authors: Tom Ulanovski, Eyal Blyachman, Maya Bechler-Speicher

Organizations: Blavatnik School of Computer Science, Tel Aviv University · Tel Aviv University · Meta AI

Abstract

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.

Figures & tables

Figure 1: An overview of Cayley-Encoder. Layer representations are mapped into nodes in a Cayley Graph over the Special Linear group $SL(2,\mathbb{Z}_{n})$ , which is then fed into a GNN to learn the final inter-layer representation.

Base Model	Method	Banking77	Emotion	MTOPDomain	MTOPIntent	PoemSentiment
Pythia-410m	Last Layer	61.17	33.48	80.88	66.97	42.40
	Best Single Layer	66.67	35.02	83.78	71.18	45.67
	MLP Last Layer	83.84	33.99	96.97	83.75	75.00
	Weighted	58.50	28.26	79.79	61.68	42.60
	DWAtt	83.23	58.60	98.03	91.66	70.87
	Deepset	84.23	47.89	97.59	92.21	73.37

Table 1: Performance comparison across 3 LLMs on 5 classification tasks. Cayley-Encoder achieves the best score in 10 out of 15 configurations, with its advantage most pronounced on larger LLMs. In all but three cases, FC-Encoder and Cayley-Encoder achieve the top two scores. We highlight in bold the best result and in blue the second best per column.

Base Model	Method	STSBenchmark	STS12	STS13	STS14	STS15	STS16	BIOSSES	SICK-R
Pythia-410m	Last-Layer	39.12	46.96	47.00	41.45	49.32	50.37	67.30	52.55
	Best-Layer	53.53	50.62	59.27	51.61	65.59	58.02	74.80	58.26
	MLP Last-Layer	25.54	42.14	36.16	33.28	37.75	39.84	59.68	43.23
	Weighted	41.81	47.12	49.49	44.21	53.72	52.76	67.79	55.09
	DWAtt	49.40	52.51	44.57	47.46	52.51	44.80	45.29	52.43
	Deepset	39.48	52.11	42.68	43.99	48.32	40.48	44.72	44.84

Table 2: Performance comparison across 3 LLMs on 8 STS tasks (Spearman correlation

\times 100

). Cayley-Encoder achieves the highest score in 17 out of 24 configurations. We highlight in bold the best result and in blue the second best per column.

Panel of Figure 2 — Figure 2: Performance across the Pythia models ranging from 14M to 2.8B parameters. Cayley-Encoder consistently outperforms all baselines across all model sizes, with smaller LLMs achieving performance comparable

Panel of Figure 3 — Figure 3: Few-shot learning analysis using Pythia-410M. Both encoders outperform all baselines with as few as 8 samples per label on most tasks. Cayley-Encoder consistently achieves the highest accuracy on Banking77 and the MTOP tasks across all training set sizes.

Base Model	Method	Banking77	Emotion	MTOPDomain	MTOPIntent	PoemSentiment	STSBench.	STS12	STS13	STS14	STS15	STS16	BIOSSES	SICK-R
Pythia-410m	FC-Encoder	90.65	75.61	98.68	95.04	75.77	54.2	50.1	53.0	51.8	63.9	57.1	58.6	57.1
	Cayley-Encoder	89.43	73.83	98.77	94.72	70.87	55.8	56.0	60.2	56.7	67.0	58.6	56.6	55.3
	LoRA	82.96	91.57	96.64	88.25	70.96	54.1	44.6	51.7	52.3	64.8	45.8	64.3	52.0
Gemma2-2B	FC-Encoder	92.39	79.90	99.07	95.34	77.12	62.9	61.8	55.4	63.9	74.7	63.0	66.0	60.1
	Cayley-Encoder	92.58	69.60	99.16	96.43	83.27	61.6	60.4	64.0	64.6	73.9	62.7	64.1	60.3
	LoRA	56.77	70.89	89.65	80.83	61.92	58.7	49.1	57.8	60.4	69.9	51.8	58.8	54.2

Table 3: Cayley-Encoder vs. LoRA on classification tasks, measured by accuracy, and STS tasks, measured by Spearman correlation

\times 100

. Cayley-Encoder outperforms LoRA in 13 out of 15 classification and 21 out of 24 STS configurations while keeping the LLM frozen. We highlight in bold the best result per task.

		Bank.	Emot.	MT-D	MT-I	Poem
Acc.	Cayley	92.6	89.7	99.1	96.3	86.4
Acc.	FC	OOM	90.4	99.2	93.2	80.2
GPU	Cayley	269	241	141	97	77
GPU	FC	OOM	8260	2978	1583	923

Table 4: Token-level aggregation with Gemma2-2B. Cayley-Encoder uses 10–30

\times

less GPU memory while achieving comparable or higher accuracy. OOM: out of memory.

	Bank.	Emot.	MT-D	MT-I	Poem
Pythia	90.1 $\pm$ .3	74.9 $\pm$ .8	98.7 $\pm$ .1	94.6 $\pm$ .4	71.3 $\pm$ 1.8
Gemma	92.4 $\pm$ .2	73.1 $\pm$ 1.5	99.1 $\pm$ .1	96.2 $\pm$ .1	82.3 $\pm$ .5
Llama	92.9 $\pm$ .1	72.8 $\pm$ 1.3	98.9 $\pm$ .1	96.4 $\pm$ .2	79.2 $\pm$ 1.3

Table 5: Robustness to layer-to-node assignment (mean

\pm

std over 10 random permutations).

	Bank.	Emot.	MT-D	MT-I	Poem
Cayley	89.4	73.8	98.8	94.7	70.9
Cayley (PE)	90.3	75.5	98.7	94.6	69.2
FC	90.7	75.6	98.7	95.0	75.8
FC (PE)	86.3	72.1	98.4	94.1	68.4

Table 6: Effect of adding learnable layer-position encodings. Layer identity information degrades performance in most cases.

Figure 4: Per-layer contribution to the final prediction for Cayley-Encoder across classification tasks using Gemma2-2B. We estimated via GNNExplainer and averaged over all test instances. Error bars show standard deviation. All layers contribute roughly equally.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Pythia-410M	Gemma2-2B	Llama3-8B
Frozen Parameters	410M	2B	8B
Weighted	25	27	33
MLP	394K (0.10%)	721K (0.036%)	1.18M (0.015%)
Deepset	394K (0.10%)	721K (0.036%)	1.18M (0.015%)
FC-Encoder	395K (0.10%)	722K (0.036%)	1.18M (0.015%)
Cayley-Encoder	395K (0.10%)	722K (0.036%)	1.18M (0.015%)

Appendix

Table 7: Trainable parameters added by each method. All base model parameters remain frozen. Percentages show overhead relative to base model size.

Table 8: Hyperparameter search spaces. All methods use Optuna with MedianPruner (5 startup trials, 10 warmup steps) and 20-50 trials per study.

	Classification						STS
Method	Banking77	Emotion	MTOPDom.	MTOPInt.	PoemSent.	Avg	STS-B	STS12	STS13	STS14	STS15	STS16	BIOSSES	SICK-R	Avg
Cayley-Encoder	89.43	73.83	98.77	94.72	70.87	85.52	55.8	56.0	60.2	56.7	67.0	58.6	56.6	55.3	58.3
Cayley-Encoder (PE)	90.25	75.52	98.72	94.58	69.23	85.66	55.3	51.4	54.8	48.2	63.0	53.2	44.8	54.2	53.1
FC-Encoder	90.65	75.61	98.68	95.04	75.77	87.15	54.2	50.1	53.0	51.8	63.9	57.1	58.6	57.1	55.7
FC-Encoder (PE)	86.33	72.07	98.36	94.10	68.37	83.85	55.3	52.5	55.9	54.2	65.0	56.3	56.5	56.6	56.5

Appendix

Table 9: Positional encoding (PE) ablation on Pythia-410m. Classification: accuracy (%), STS: Spearman correlation (

\times 100

). Bold : best per column within each encoder type.

Method	Pythia-410M	Gemma2-2B	Llama3-8B
Frozen Parameters	410M	2B	8B
LoRA ( $r=3$ )	442K (0.11%)	1.20M (0.046%)	2.56M (0.032%)
LoRA ( $r=16$ )	2.36M (0.58%)	6.39M (0.25%)	13.6M (0.17%)
Cayley-Encoder	395K (0.10%)	722K (0.036%)	1.18M (0.015%)

Appendix

Table 10: Trainable parameters for LoRA (ranks 3 and 16) compared to Cayley-Encoder. Percentages show overhead relative to base model size.

Base Model	Graph Type	Banking77	Emotion	MTOPDomain	MTOPIntent	PoemSentiment
Pythia-410m	Random-Regular	90.38	75.15	98.80	94.73	70.77
Pythia-410m	Cayley-Encoder	89.43	73.83	98.77	94.72	70.87
Gemma2-2B	Random-Regular	92.16	76.50	99.12	96.21	82.31
Gemma2-2B	Cayley-Encoder	92.58	69.60	99.16	96.43	83.27
Llama3-8B	Random-Regular	92.23	72.99	98.90	96.29	79.81
Llama3-8B	Cayley-Encoder	92.85	73.43	99.03	96.46	79.04

Appendix

Table 11: Ablation isolating the contribution of the Cayley-Encoder-graph structure. Our Cayley-Encoder-Encoder uses a 4-regular topology; we ask whether its gains over FC-Encoder come from sparsity alone or from the specific Cayley-Encoder connectivity. Replacing the Cayley-Encoder graph with a random 4-regular graph in classification tasks shows comparable results. Cayley-Encoder-Encoder is better in 9 out of 15 cases. Bold : best per column.

Base Model	Method	STSBenchmark	STS12	STS13	STS14	STS15	STS16	BIOSSES	SICK-R
Pythia-410m	Random-Regular	53.19	51.69	56.81	52.82	62.19	54.27	47.81	55.66
Pythia-410m	Cayley-Graph	55.84	55.97	60.25	56.65	66.99	58.63	56.59	55.34
Gemma2-2B	Random-Regular	59.85	52.57	58.25	58.73	72.40	61.11	62.62	58.34
Gemma2-2B	Cayley-Graph	61.62	64.36	63.98	64.59	73.94	62.69	64.07	60.32
Llama3-8B	Random-Regular	54.46	44.55	48.92	46.02	65.47	55.55	57.68	53.51
Llama3-8B	Cayley-Graph	63.05	65.31	63.53	70.17	76.96	63.13	55.32	60.74

Appendix

Table 12: Ablation isolating the contribution of the Cayley-graph structure. Our Cayley-Encoder uses a 4-regular topology; we ask whether its gains over FC-Encoder come from sparsity alone or from the specific Cayley connectivity. Replacing the Cayley graph with a random 4-regular graph degrades performance on 22 of 24 task–model pairs across 3 LLMs and 8 STS benchmarks, showing that the structured connectivity – not the regularity itself – drives the improvement. Bold : best per column.

Task Category	Dataset	License
Classification	Banking77	CC BY 4.0
	Emotion	CC BY 4.0
	MTOP (Domain/Intent)	CC BY-SA 4.0
	Poem Sentiment	Apache License 2.0
STS	STSBenchmark	CC BY-SA 4.0
	STS12–16	CC BY-SA 4.0

Appendix

Table 13: Verified licenses for the models and datasets used in this work.

Explore similar work

Causal and Interpretable Structures in LLM Compositional Tasks
Sep 28, 2026Gurbir Arora, Toni J. B. Liu, Jiajun Bao +2LLM Interpretability Causal Interventions in Language Models
Cornell University, USA · Goodfire AI, USA
Embedding Geometry Shapes LLM Routing: Tracing Failure Modes and Improving Generalization
Oct 13, 2026Leonardo Torti, Emanuele Rimoldi, Tomer Galanti +2LLM Routing Representation Geometry in Language Models
Massachusetts Institute of Technology. · Texas A&M University.
Skip a Layer or Loop It? Learning Program-of-Layers in LLMs
Jun 4, 2026Ziyue Li, Yang Li, Tianyi ZhouLLM Inference LLM Inference Acceleration
University of Maryland, College Park, MD, USA · 2MBZUAI, Abu Dhabi, UAE.

Sep 28, 2026J/K move · Enter open · S save

Causal and Interpretable Structures in LLM Compositional Tasks

Gurbir Arora, Toni J. B. Liu, Jiajun Bao, Raphaël Sarfati+1

Cornell University, USA · Goodfire AI, USA

LLM Interpretability Causal Interventions in Language Models Representation Geometry in Language Models Compositional Representation Learning

Open paper AskPDF arXiv

Large language models are able to solve tasks whose answers depend on not only individual input tokens, but also on relations among them. How is such relational information represented and processed across transformer layers? We study activations from ensembles of prompts that require inferring relationships between three tokens corresponding to a cyclic concept (months, hours, weekdays, and musical notes) to correctly predict the next token. Across model families (Llama, Qwen, Gemma, and Mistral) and cyclic concepts, we find a consistent layerwise progression in how the joint dependence among the tokens is geometrically organized and causally used: intermediate layers use a joint representation based on the inferred relationship between two tokens, while later layers use a joint representation associated with all three tokens to correctly complete the task. We also find other relationships between tokens that are geometrically structured but remain causally inert in the next-token prediction. Crucially, when taken together, these geometric and causal investigations reveal the representation-level mechanism that progressively organizes and composes the relational information to form the answer. More surprisingly, restricting the models to such causally relevant joint representations improves next-token prediction accuracy.

Abstract

Explore similar work

Causal and Interpretable Structures in LLM Compositional Tasks

Embedding Geometry Shapes LLM Routing: Tracing Failure Modes and Improving Generalization

Skip a Layer or Loop It? Learning Program-of-Layers in LLMs