Improving LLM Final Representations with Inter-Layer Geometry¶

会议: NeurIPS 2026
arXiv: 2603.22665
代码: https://github.com/tomulanovski/ILSE
领域: 可解释性
关键词: 大语言模型表示、层间几何结构、图神经网络、Cayley 图、特征聚合

一句话总结¶

将冻结 LLM 的中间层特征聚合建模为图学习问题，基于特殊线性群 $SL(2, \mathbb{Z}_n)$ 构建稀疏正则、低直径的 Cayley 图并用轻量 GNN 聚合，仅增加至多 0.1% 参数量即全面超越最后一层抽取、注意力层聚合及 LoRA 微调。

研究背景与动机¶

在大语言模型（LLM）的下游表征迁移范式中，最普遍的做法是将预训练模型完全冻结作为特征提取器，直接提取最后一层（Last Layer）的隐藏状态输入轻量分类头或投影头。然而，Transformer 的表征特性分析早已表明，模型深度承载了层级化的互补信息：底层捕捉词汇与表层形态，中层沉淀句法骨架与事实知识，高层则更偏向语义抽象与特定预训练目标。在许多具体任务中，中间层的特征质量甚至稳定超越最后一层。

为了打破对单一最后一层的依赖，现存的跨层聚合方案主要走向两个极端：要么通过繁琐的网格搜索逐层独立验证以选取单一最佳层（Best Layer），这破坏了不同层信息的互补性；要么引入复杂的跨层注意力机制（如 Depth-Wise Attention, DWAtt），但这不仅参数量膨胀数倍，本质上仍以最后一层作为 Query 进行非对称锚定，将浅层退化为从属的 Key/Value，无法实现层与层之间的平等交互与深层协同。

针对这种非对称锚定与参数开销矛盾，本文另辟蹊径，将层间信息融合建模为图结构上的消息传递。核心 idea：将冻结 LLM 的全部网络层映射为图节点，利用特殊线性群 $SL(2, \mathbb{Z}_n)$ 构建兼具稀疏性、正则性与对数级通信直径的 Cayley 图，通过轻量级 GNN 实现所有网络层在几何对称结构下的平等交互与全局信息聚合。

方法详解¶

整体框架¶

Cayley-Encoder 的处理流程直接作用于预训练且冻结参数的 LLM 之上。给定输入序列，模型首先提取所有 $L$ 个隐藏层的 token 向量，并通过 token 维度的均值池化（Mean-Pooling）为每一层生成固定维度的层表征向量 $z_\ell$。接着，这 $L$ 个层向量被随机分配映射到预构建的 $SL(2, \mathbb{Z}_n)$ Cayley 图节点上；若图的阶数大于层数，则补充虚拟节点保证对称性。随后，在固定拓扑结构的图上执行 1~2 层轻量 MPNN（如 GIN/GCN）迭代传递跨层消息。最后，通过对非虚拟节点进行均值或求和池化得到全局图向量，输入下游预测头完成分类或相似度计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列与冻结 LLM<br/>提取各层隐藏状态"] --> B["层表征均值提取<br/>Token 级 Mean-Pooling"]
    B --> C["SL(2, Z_n) 拓扑映射<br/>构建 4-正则 Cayley 图与虚节点补充"]
    C --> D["正则几何消息传递<br/>轻量 GIN/GCN 跨层特征扩散"]
    D --> E["非虚拟节点池化<br/>输出下游图级任务表征"]

关键设计¶

1. 层表征提取与图节点初始化：解耦序列长度的固定尺寸层图 针对长文本或动态序列带来的维度爆炸痛点，本文首先将每一层输出的隐藏张量在序列维度上解耦。对于具有 $L$ 层、输入长度为 $T$、隐藏维度为 $d$ 的 LLM，第 $\ell$ 层的表示 $H_\ell \in \mathbb{R}^{T \times d}$ 被均值池化为单一的层向量 $z_\ell \in \mathbb{R}^d$： $$z_\ell = \frac{1}{T} \sum_{t=1}^T H_\ell[t, :]$$ 这一设计不仅与通用句子表征提取范式保持一致，更使得构建出的层间几何图节点规模严格由模型层数 $L$ 决定，而与样本 token 长度 $T$ 完全无关。每个层节点 $v_\ell$ 的初始特征即赋值为 $h_\ell^{(0)} = z_\ell$。

2. 基于特殊线性群 $SL(2, \mathbb{Z}_n)$ 的 Cayley 图构建：规避过拟合与信息瓶颈的理想拓扑 若直接采用全连接图（FC-Encoder），每对层间直接连边虽然提供了充分连通性，但在层数较多或扩展到 token 级图时边数膨胀为 $O(L^2)$，极易导致 GNN 对无信息拓扑结构的过拟合；而任意启发式稀疏图又容易引发信息压缩瓶颈（Over-squashing）。为此，本文引入特殊线性群 $SL(2, \mathbb{Z}_n)$（模 $n$ 意义下行列式为 1 的 $2 \times 2$ 整数矩阵群），并采用标准 Margulis 生成元集 $S = \{A, A^{-1}, B, B^{-1}\}$： $$A = \begin{pmatrix} 1 & 1 \\ 0 & 1 \end{pmatrix}, \quad B = \begin{pmatrix} 1 & 0 \\ 1 & 1 \end{pmatrix} \pmod n$$ 群元素总数即图节点数为 $|V_n| = n^3 \prod_{p|n, p \text{ prime}} (1 - \frac{1}{p^2})$。选取满足 $|V_n| \ge L$ 的最小整数 $n$，将 $L$ 个层表征随机指派给前 $L$ 个节点，剩余 $|V_n| - L$ 个节点置为零向量作为虚拟中继节点（Virtual Nodes）。该拓扑天然具备全局 4-正则（4-regular）与点传递性（Vertex-transitive），各节点拓扑完全对称平权；同时具备对数级直径 $O(\log |V|)$，确保任意两层之间可在极少步数内完成高效通信，从数学本质上兼顾了稀疏性与强扩展性（Expander 属性）。

3. 对称消息传递与去虚拟化聚合：消除层偏见的全局表征提炼 为打破以往注意力机制对特定层（如最后一层）的非对称锚定，消息传递阶段让所有节点在 Cayley 拓扑约束下平等交换信息： $$m_v^{(k)} = \sum_{u \in \mathcal{N}(v)} M^{(k)}\left(h_v^{(k-1)}, h_u^{(k-1)}\right), \quad h_v^{(k)} = U^{(k)}\left(h_v^{(k-1)}, m_v^{(k)}\right)$$ 其中 $M^{(k)}$ 和 $U^{(k)}$ 分别为可学习的邻居消息函数与状态更新函数（实验中采用表达能力完备的 GIN 或轻量 GCN，隐藏维度设为 256）。经过 $K \in \{1, 2\}$ 轮迭代后，为防止充当中继通信的虚拟节点引入噪声，最终的全局图表示仅对非虚拟节点（真实层节点）进行池化： $$h_G = \text{Pool}\left(\{h_\ell^{(K)} \mid \ell \in \{1, \dots, L\}\}\right)$$ 这一设计使各层特征得以在大局观下充分协同融合，完全避免了单一主导层引起的表征坍缩。

损失函数 / 训练策略¶

保持底层 LLM 全部权重严格冻结。对于文本分类任务，图表征 $h_G$ 输入单层线性分类器，采用标准交叉熵损失（Cross-Entropy Loss）进行端到端优化；对于语义文本相似度（STS）任务，句子对分别经过冻结 LLM 和 Cayley-Encoder 编码得到表示 $h_G^{(1)}, h_G^{(2)}$，计算余弦相似度并使用均方误差损失（MSE Loss）拟合真实标注分数。训练全流程采用 Adam 优化器，学习率与权重衰减在 $[10^{-4}, 10^{-3}]$ 范围内通过 Optuna 贝叶斯优化联合剪枝自适应确定。

实验关键数据¶

主实验¶

论文在 MTEB 基准的 5 个分类任务以及 8 个 STS 语义相似度任务（STSBenchmark 训练 + 7 个零样本迁移数据集）上，全面横跨 Pythia-410M、Gemma2-2B、Llama3-8B 等模型进行测试。下表汇聚了 3 种主流架构在 5 个分类任务上的准确率（Accuracy, %）表现：

基础模型	聚合方法	Banking77	Emotion	MTOP Domain	MTOP Intent	Poem Sentiment
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
	DWAtt	83.23	58.60	98.03	91.66	70.87
	FC-Encoder	90.65	75.61	98.68	95.04	75.77
	Cayley-Encoder	89.43	73.83	98.77	94.72	70.87
Gemma2-2B	Last Layer	62.16	26.89	80.79	68.02	35.67
	Best Single Layer	72.31	31.27	87.09	75.36	42.79
	DWAtt	88.82	61.22	98.39	90.93	67.79
	FC-Encoder	92.39	79.90	99.07	95.34	77.12
	Cayley-Encoder	92.58	69.60	99.16	96.43	83.27
Llama3-8B	Last Layer	68.25	34.23	84.42	73.39	40.96
	Best Single Layer	71.93	38.42	89.01	78.17	47.02
	DWAtt	90.04	66.55	97.97	92.41	75.00
	FC-Encoder	92.38	71.64	98.99	96.19	77.98
	Cayley-Encoder	92.85	73.43	99.03	96.46	79.04

在 8 项 STS 任务上（Spearman 相关系数 $\times 100$），Cayley-Encoder 在 24 个配置中 17 项拔得头筹（例如 Llama3-8B 上 STS14 达到 70.17，相比 Last Layer 的 42.88 提升超 27 个绝对点；STS12 从 32.65 跃升至 65.31）。

在与参数高效微调方法 LoRA 的直接对比中，Cayley-Encoder 在冻结 LLM 前提下，在 15 个分类配置中 13 个优于 LoRA，在 24 个 STS 配置中 21 个胜过 LoRA，而可训练参数量比 rank=16 的 LoRA 减少了 6~12 倍。

消融实验与机制分析¶

1. 序列级 Token 级全交互扩展消融（Gemma2-2B） 当突破层级池化限制，将每个 Token 视为节点构建 $L \times T$ 规模的超大交互图时，Cayley 图的稀疏稀释优势体现得淋漓尽致：

评测维度	方法配置	Banking77	Emotion	MTOP-Domain	MTOP-Intent	Poem Sentiment
准确率 (%)	FC-Encoder	OOM (显存溢出)	90.40	99.20	93.20	80.20
	Cayley-Encoder	92.60	89.70	99.10	96.30	86.40
GPU 显存 (MB)	FC-Encoder	OOM	8260	2978	1583	923
	Cayley-Encoder	269 (降 30.7×)	241 (降 34.3×)	141 (降 21.1×)	97 (降 16.3×)	77 (降 12.0×)

2. 层位置编码（Positional Encoding, PE）消融（Pythia-410M） 测试在节点中引入层物理深度位置编码的影响：

拓扑架构	是否加入层位置编码	分类平均准确率 (%)	STS 跨域平均相关系数 ($\times 100$)
FC-Encoder	无位置编码（默认）	85.75	56.40
	显式注入层深度 PE	83.86 (降 1.89)	56.55 (微升 0.15)
Cayley-Encoder	无位置编码（默认）	85.52	58.30
	显式注入层深度 PE	85.66 (微升 0.14)	54.76 (降 3.54)

关键发现¶

去中心化协同胜过单层主导：GNNExplainer 可解释性归因显示，所有层的贡献度权值均匀分布在 $0.035 \sim 0.040$ 区间内，没有任何一层出现注意力权重坍缩；这证实最终表征的飞跃来源于全层协同信息的互补涌现。
层到节点的映射具有置换鲁棒性：在固定超参下进行 10 次随机层-节点映射置换测试，标准差全部在 1.5% 以内，表明性能增益源于图拓扑本身的扩展与连通性质，而非特定层之间的邻接巧合。
极佳的小样本学习能力：在 Banking77 与 MTOP 任务上，Cayley-Encoder 仅需每个类别 8~32 个标注样本，准确率便可一举超越在全量数据上训练的 DWAtt 与 Last Layer 基线。

亮点与洞察¶

将网络层间关系严格几何化：跳出了深度学习常见的“加层堆注意力”惯性思维，引入代数拓扑中具有良好扩展特性的 Cayley 图，巧妙利用其 4-正则与对数直径解决了 GNN 在小图上的过拟合与过挤压难题。
极低算力开销下的表征榨取：不需要解冻庞大的 LLM 权重，在至多新增 0.1% 额外参数量的极轻量代价下，性能大幅超越参数量高出数倍的 DWAtt 甚至是直接修改权重的 LoRA 微调。
可平滑迁移至序列级稠密交互：Token-level 扩展实验展示了 Cayley 图在 $O(N)$ 稀疏边复杂度下的显存优势（相比全连接节省 10~30 倍显存且避免 OOM），为长上下文层-词联合表征建模提供了极具潜力的架构范式。

局限与展望¶

生成式任务迁移尚未验证：当前评测主要集中在辨别式分类（Classification）与语义匹配检索（STS）等表征提取任务上，如何将 Cayley 跨层拓扑有机整合进自回归解码（Autoregressive Generation）仍待探索。
静态预定义图结构的适应性：Cayley 图是数据无关的先验数学图，尽管正则性防止了结构过拟合，但缺乏针对不同下游数据分布动态调整连接强度的自适应图重构机制。
更大规模巨型模型的普适性：实验覆盖至 8B 模型，尚未在 70B+ 等具有更深网络层数（如 80+ 层）的超大模型或 MoE 混合专家模型上进一步验证层间通信模式的鲁棒性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将群论 Cayley 图与层间表征几何学习深度结合，打破了传统的最后一层与注意力堆叠范式。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 9 款模型、13 个基准任务，包含少样本、Token 级拓展、与 LoRA 横向对比以及细致的 GNN 归因分析。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密，拓扑动机与实验现象相互印证，行文逻辑流畅。
价值: ⭐⭐⭐⭐⭐ 为轻量级 LLM 表征利用、冻结特征高效调优及低算力下游适配提供了极佳的通用工具与理论启示。