[Paper] 具选择性访问早期表示的 Transformer

Source: arXiv - 2605.03953v1

概述

一个名为 Selective Access Transformers (SATFormer) 的新型 Transformer 模型族表明，让更深层在真正有帮助时“窥视”第一层的表示——仅在必要时——可以在不增加常见的内存和速度开销的情况下提升性能。通过将早期表示的复用视为上下文相关的检索问题，SATFormer 在各种模型规模上均优于普通 Transformer 和早期的静态残差技巧。

关键贡献

• 选择性门控机制，能够为每个 token、每个头和每层动态决定注入多少第一层值矩阵 (V_1)。
• 内存高效设计：门控是每个头/层的轻量标量，使整体占用与标准 Transformer 相当。
• 广泛的实证提升：在 130 M 到 1.3 B 参数范围内，验证损失和零样本准确率均有一致提升，在检索密集型基准上提升最大（约 +≈1.5 % 平均）。
• 可解释性洞察：对学习到的门控进行分析，揭示出稀疏、深度相关、头部特定以及任务类别的模式，证实模型学会了 何时 与 何处 使用早期信息。
• 开源实现（GitHub 链接），可最小改动地嵌入现有 Transformer 代码库。

方法论

1. 基线架构 – 从标准 Transformer（预归一化、多头自注意力、残差连接）开始。

2. 保留第一层的值通路 – 保持来自第一层的值投影 (V_1)，使其在所有后续层可用。

3. 上下文相关门控 – 对于每个下游层 (l) 和头 (h)，使用一个小型前馈网络，以当前隐藏状态为输入，计算标量门 (g_{l,h}\in[0,1])。

4. 选择性注入 – 层 (l) 的注意力输出变为

   $$\text{output}{l} = \text{Attention}(Q_l,K_l,V_l) ;+; g{l,h},\cdot, V_1$$

   其中门可以完全关闭早期值贡献（0），或完全通过（1），或介于两者之间的任意值。

5. 训练 – 整个系统使用常规语言模型或分类损失端到端训练；门控与其他参数一起联合学习。

6. 效率技巧 – 门控在 token 位置上进行广播，因此额外计算仅为少量逐元素乘法，保持吞吐量。

结果与发现

• Zero‑shot accuracy 在检索为中心的任务（如 MS‑MARCO、Natural Questions）上相比静态残差提升约 1.5 个百分点。
• Throughput 与原始 Transformer 相差仅 2‑3 %，且 GPU memory 开销可以忽略不计（< 5 %）。
• Gate analysis 表明，早期层的值在中层的前几个注意力头中被大量用于词汇密集的 token，但在处理更高层语义的更深层时逐渐衰减——正是作者所假设的选择性行为。

实际意义

• 更好的检索增强模型 – 如果你在构建依赖提取精确 token 级线索的搜索或问答系统，SATFormer 可以在不需要额外索引结构的情况下提供显著的准确率提升。
• 即插即用升级 – gating 模块只需几行代码；你可以对现有的 Transformer 堆栈（BERT、GPT、T5 等）进行改造，几乎无需工程开销即可获得收益。
• 性价比高的扩展 – 对于内存受限的大规模语言模型，SATFormer 在廉价的静态残差技巧和重量级密集检索层之间提供了折中方案。
• 可解释性用于调试 – 学习到的门控模式可以可视化，以了解模型哪些部分仍然依赖低层词汇信号，有助于模型内省和偏差分析。

限制与未来工作

• Gate granularity – 当前的门在层/头内的所有 token 位置之间共享；更细粒度（每个 token）的 gating 可能捕获更细微的复用，但会增加内存开销。
• Task scope – 论文聚焦于语言建模和检索密集型基准；尚需观察 SATFormer 在生成为中心的任务（例如摘要、代码合成）上的表现。
• Training stability – 某些非常深的配置出现偶发的门饱和（全零或全一），需要仔细调节学习率。
• Future directions suggested include exploring multi‑layer early‑representation pools (not just the first layer), hierarchical gating, and applying the selective‑access idea to vision Transformers where early‑level texture cues may be similarly valuable.

作者

• Skye Gunasekaran
• Téa Wright
• Rui‑Jie Zhu
• Jason Eshraghian

论文信息

• arXiv ID: 2605.03953v1
• 分类: cs.LG, cs.CL
• 出版日期: 2026年5月5日
• PDF: 下载 PDF