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,Zn) , 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 ×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.
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
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 ×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
FC
OOM
90.4
99.2
93.2
80.2
GPU
Cayley
269
241
141
97
77
FC
OOM
8260
2978
1583
923
Table 4: Token-level aggregation with Gemma2-2B. Cayley-Encoder uses 10–30 × less GPU memory while achieving comparable or higher accuracy. OOM: out of memory.
Bank.
Emot.
MT-D
MT-I
Poem
Pythia
90.1 ± .3
74.9 ± .8
98.7 ± .1
94.6 ± .4
71.3 ± 1.8
Gemma
92.4 ± .2
73.1 ± 1.5
99.1 ± .1
96.2 ± .1
82.3 ± .5
Llama
92.9 ± .1
72.8 ± 1.3
98.9 ± .1
96.4 ± .2
79.2 ± 1.3
Table 5: Robustness to layer-to-node assignment (mean ± 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 ( ×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
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
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
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
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
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
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.
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.
With the growing number of LLMs available, selecting which one to use for a prompt has become a key question. A large part of the routing literature relies on semantically embedding prompts into a vector space and then applying a lightweight classifier to the embedded representations. These methods perform robustly in-distribution but they struggle to generalize to unseen tasks. By studying the geometry of semantic embeddings, we trace a clear failure mode: features that determine the neighbors of unseen prompts are an unreliable proxy for LLM selection. We therefore replace the semantic embedding with the following map: we run a partial forward pass of Qwen2.5-7B-Instruct on the prompt and use the hidden state of its last token at a deep layer as the representation. When applying kNN and other simple classifiers to this new representation, we obtain a clear result in generalization: on average across cost penalties, their performance improves. Moreover, kNN and another simple classifier that we propose outperform every evaluated semantic-embedding baseline. Our two simple methods also beat a recent SoTA hidden-states based router on RouterBench.
Large language models (LLMs) perform inference by following a fixed depth and order, non-recurrent execution of all layers. We reveal the wide existence of training-free, flexible, dynamic program-of-layers (PoLar), where pretrained layers can be packed as modules and then skipped or looped to form a customized program for each input. For most inputs, substantially shorter program executions can achieve the same or better accuracy, while incorrect predictions of the original LLM can be corrected by alternative programs with fewer layers. These observations indicate that inference admits multiple valid latent computations beyond the standard forward pass. To efficiently achieve PoLar in practice, we propose a lightweight PoLar prediction network, which learns to generate execution programs that dynamically skip or repeat pretrained layers for each input. Experiments on mathematical reasoning benchmarks demonstrate that PoLar consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers, and that these gains persist under out-of-distribution evaluation. Our results suggest that fixed-depth execution captures only a narrow subset of an LLM's latent reasoning capacity.
Ziyue Li, Yang Li, Tianyi Zhou
University of Maryland, College Park, MD, USA · 2MBZUAI, Abu Dhabi, UAE.