Train Once, Reuse Everywhere：Generalizable Implicit In-Context Learning by Routing Attention
（ICML 2026）
motivation
上下文学习（ICL）的局限性：虽然大语言模型可以通过在提示词中插入少量标注示例（即 ICDs）来适应新任务，但这会显著增加输入序列的长度，从而推高推理成本。此外，这种显式的方法对示例的顺序和格式非常敏感。
隐式 ICL 及其缺陷：为了解决成本问题，近期研究提出了“隐式 ICL”，即把示例压缩成稠密向量，并注入到模型的中间层（残差流）中。然而，现有的隐式方法本质上是“事后（post-hoc）”的残差调整，缺乏结构性，导致模型只是在“模仿”而不是真正“内化”ICL机制。因此，这些方法往往与特定任务高度绑定，在面对域外（OOD）任务或缺乏示例的任务时泛化能力极差。
方法
阶段一：提取主 ICL 方向 (Principal ICL Directions, PIDs)
第一步的目标是找到那些在不同任务之间共享的、能够代表“上下文学习（ICL）能力”的底层注意力结构。作者通过多领域数据来提取这些模式。
数据收集：假设有多个领域（Domains），对于每个领域，构建包含示例和查询的 ICL 提示词（Prompts）。
提取特征：将这些提示词输入到冻结的大语言模型（LLM）中，专门提取最后一个 token（即整合了所有上下文信息的那个词）在第 $l$ 层、第 $h$ 个注意力头的 Query (查询) 和 Key (键) 表征： $q _ {d,i}^{l,h}, k _ {d,i}^{l,h} \in \mathbb{R}^{1 \times d_k}$
拼接与堆叠：将所有注意力头（Heads）的表征拼接起来，得到层级的向量： $q _ {d,i}^{l} = \text{Concat}_h q _ {d,i}^{l,h} \in \mathbb{R}^{1 \times d}$ 随后，把所有提示词和所有领域的向量堆叠成宏观的 ICL 基矩阵（ICL bases）： $\tilde{Q}^{l}, \tilde{K}^{l} \in \mathbb{R}^{N \times d}$
主成分分析 (PCA)：这是最关键的一步。对上述两个矩阵进行 PCA 降维，提取出前 $r$ 个主成分方向，得到我们要的主 ICL 方向 (PIDs)： $U _ {q}^{l}, U _ {k}^{l} \in \mathbb{R}^{d \times r}$ 这里的 $r$ 通常很小（论文中消融实验最优选了 8），这些矩阵作为可复用的路由方向被保存下来。
PIDs如何捕获通用ICL模式（理论证明）:
作者利用混合尖峰协方差模型进行了数学证明：
根据尖峰协方差模型，这个矩阵可以被分解为三个独立的部分：
$$ \Sigma_{Q,d} = S_q \Lambda_q S_q^\top + \sigma^2 I + B_{q,d} \Gamma_{q,d} B_{q,d}^\top $$
- 共享部分 ($S_q \Lambda_q S_q^\top$)：这代表了跨领域共享的低维注意力结构，也就是所有任务在进行 ICL 时都需要的通用能力（例如将注意力分配给正确格式的上下文）。
- 各向同性噪声 ($\sigma^2 I$)：代表模型内部署基础的、均匀的白噪声。
- 领域特异性部分 ($B_{q,d} \Gamma_{q,d} B_{q,d}^\top$)：这代表了领域 $d$ 特有的变化（例如特定任务的词汇或格式）。
2. 多领域数据的混合与池化 (Pooling)
为了提取通用的 ICL 模式，作者并没有只看一个领域，而是将所有 $D$ 个不同领域的 ICL 提示词放在一起，计算总的“池化协方差矩阵” $\hat{\Sigma}_Q$。
当计算这个多领域混合协方差矩阵的期望值时，可以将其展开为：
$$ \mathbb{E}[\hat{\Sigma}*Q] = S_q \Lambda_q S_q^\top + \sigma^2 I + \frac{1}{N} \sum_{d=1}^D |\mathcal{D}d| B{q,d} \Gamma_{q,d} B_{q,d}^\top $$
在这个公式中，我们可以看到不同成分在混合后的表现：
- 第一项 (共享结构)：因为它是所有领域共有的，所以在加总过程中，这一部分会不断累积并被强化。
- 第三项 (领域特异性干扰)：这是 $D$ 个不同领域的特有特征的加权和。
如果选择的 $D$ 个领域足够多样化且缺乏一致的对齐（例如，既有新闻分类，又有常识问答和情感分析），那么这些领域特异性子空间 $B_{q,d}$ 的方向是杂乱无章的。
在求和的过程中，这些不一致的领域特异性方向会相互抵消（averages out），最终退化为增加背景方差的各向同性噪声，而不是形成主导的特征方向。
阶段二：查询条件路由器 (Query-conditioned Router)
得到了通用的 ICL 方向后，由于每个具体的输入问题（Query）不同，模型需要“因地制宜”地决定在推理时具体使用哪些方向。为此，论文引入了一个极轻量级的路由器。
对于一个新的零样本（Zero-shot）输入文本 $x$，先用一个较小且冻结的文本编码器（如 MiniLM）将其转换为语义特征向量 $E(x)$。
将 $E(x)$ 并行输入到两个多层感知机（MLP）分支中，生成两组动态参数：
- 路由矩阵：$\alpha(x) = \tanh(g _ {\theta _ {\alpha}}(E(x))) \in \mathbb{R}^{L \times r}$ （决定每层 $l$ 中 $r$ 个 PID 方向的正负和权重）。
- 门控矩阵：$\gamma(x) = \sigma(g _ {\theta _ {\gamma}}(E(x))) \in \mathbb{R}^{L \times H}$ （利用 Sigmoid 函数生成，控制第 $l$ 层中第 $h$ 个注意力头是否需要开启干预）。
在零样本推理时，模型会将上述生成的低秩偏置动态添加到原始的注意力 Logits（即 $A^{l,h}(x)$）上： $$ \tilde{A}^{l,h}(x) = A^{l,h}(x) + \gamma^{l,h}(x) (Q _ {zs}^{l} U _ {q}^{l}) \text{diag}(\alpha^{l}(x)) (K _ {zs}^{l} U _ {k}^{l})^{\top} $$ 在这个公式中：
- $Q _ {zs}^{l}$ 和 $K _ {zs}^{l}$ 是零样本推理时的 Query 和 Key。
- $(Q _ {zs}^{l} U _ {q}^{l})$ 和 $(K _ {zs}^{l} U _ {k}^{l})^{\top}$ 是将当前的 Q 和 K 投影到之前提取的通用 PID 空间中。
- $\text{diag}(\alpha^{l}(x))$ 则是根据当前输入，对这些通用方向进行个性化的加权。
看上去有点像线性注意力的式子。
阶段三：多目标训练 (Multi-objective Training)
了让这个路由器既能听懂指令，又不会过度干预 LLM 原本的能力，作者设计了一个复合的损失函数。在训练时，LLM 的参数是完全冻结的，只更新路由器的参数。
总体损失函数由三个部分组成： $$ \mathcal{L} = \mathcal{L}\times{CE} + \lambda\times{conf}\mathcal{L}\times{conf} + \lambda\times{spar}\mathcal{L}\times{spar} + \lambda\times{gate}\mathcal{L} _ {gate} $$
- 监督交叉熵损失 ($\mathcal{L} _ {CE}$)：最基础的分类/生成损失，确保添加了路由干预后，模型依然能输出正确的标签 $y_i$。 $$ \mathcal{L}{CE} = -\frac{1}{B} \sum{i=1}^{B} \log P^{ICR}(y_i | x_i) $$
- 置信度对齐损失 ($\mathcal{L} _ {conf}$)：为了防止路由器为了降低损失而采取“输出模棱两可结果”的捷径，引入了基于熵（Entropy）的惩罚，确保干预后的预测置信度不低于零样本时的置信度。
- 稀疏性正则化 ($\mathcal{L} _ {spar}$ 和 $\mathcal{L} _ {gate}$)：引入 $L_1$ 正则化，鼓励路由向量 $\alpha$ 和门控 $\gamma$ 尽可能稀疏。特别地，论文设计了一个随层数线性增加的权重 $w^l$，因为越靠后的层越接近最终预测，干预应当越谨慎、越聚焦于少数几个决定性的方向。
