1. 问题与动机
Transformer 架构中的标准全注意力机制（Full Attention）的计算量与序列长度呈二次增长关系，这是长上下文语言建模的核心瓶颈。滑动窗口注意力（Sliding Window Attention, SWA）通过限制上下文范围提升效率，但代价是缩小了感受野，导致模型无法捕获远距离依赖。
当前工业界的混合 Transformer 方案（如 Mistral、Gemma-2）采用静态的、启发式设计的交替模式（例如每 3 层 SWA 插入 1 层 Full Attention），但这种静态分配无法适应不同输入上下文的信息需求差异。作者观察到一个关键直觉：当滑动窗口之外的上下文与当前生成内容相关性较低时，应当降低昂贵的全注意力比例——而这种相关性是动态变化的，因此需要 token 级别的动态路由决策。
该工作的核心动机在于：现有混合 Transformer 在效率和效果之间的 trade-off 是静态的、粗粒度的，而信息检索需求在不同 token、不同层之间存在显著差异，需要一种动态、细粒度的注意力分配机制。
2. 方法详解
SwiAttn 的核心架构由三部分组成：轻量路由器、双分支注意力计算和自适应正则化。
路由器设计：在每一层 l，对每个 token 通过一个线性映射计算 soft gate σ̃ = sigmoid(R(H))，其中 R 是线性映射函数，H 是该层的隐状态。然后通过直通估计器（Straight-Through Estimator, STE）得到硬门控 σ = 𝟙(σ̃ > 0.5) + σ̃ - sg(σ̃)，其中 sg(·) 为 stop-gradient 操作。STE 的关键在于前向传播使用离散的 0/1 决策，反向传播时梯度直接通过连续的 sigmoid 值流回，使得离散路由可微训练。
双分支计算：两个分支（Full Attention 和 SWA）共享相同的 Q、K、V 投影，最终输出为门控加权：O = σ ⊙ O_FULL + (1-σ) ⊙ O_SWA。在推理的解码阶段，两个分支共享 KV cache 以提升效率。
自适应正则化损失：这是本文最关键的创新。正则化损失为 L_t(l) = γ_t(l) · softplus(R(h_t(l)))，其中自适应权重 γ_t(l) = γ_base / (ε + NLL(t) + α·||o_FULL - o_SWA||²₂)。这个设计的精妙之处在于：当某个 token 预测难度大（NLL 高）或两个分支输出差异大（MSE 高）时，正则化权重自动降低，鼓励路由器选择 Full Attention；反之对简单 token 施加更强正则化压力推向 SWA。超参数配置：γ_base = 1×10⁻³, ε = 0.1, α = 100。
训练流程：采用持续预训练策略（Continual Pretraining, CPT），两个注意力分支均从预训练的全注意力模型初始化。模型规模 1.5B 参数、22 层 Transformer，SWA 窗口大小 1024 tokens。初始预训练 240B tokens（4K 序列），SwiAttn 持续预训练 20B tokens。长上下文变体额外通过 YaRN 位置编码扩展到 32K，再训练 8B tokens。
3. 核心技术贡献
贡献一：Token 级动态注意力路由。将注意力类型选择从层级静态规则提升为 token×层 的二维动态决策。实验表明，平均全注意力比例约为 0.13（远低于静态混合的 0.25），但在关键位置（如 needle-in-haystack 任务中 needle 在窗口外时）全注意力比例可达 0.30，在窗口内时仅约 0.05，展现了精准的信息需求感知。
贡献二：自适应正则化机制。通过 NLL 和分支间 MSE 双信号动态调节正则化强度，解决了「简单 token 浪费全注意力 vs 困难 token 被迫使用 SWA」的矛盾。消融实验显示，去除自适应正则化后 GovReport 等长文任务性能从与全注意力持平下降至明显劣势（平均 24.9 vs 26.2）。
贡献三：与全注意力匹配的效果 + 显著效率提升。在 28 个评测任务上，SwiAttn 在常识推理（53.7% vs 53.4%）、上下文内检索（52.9% vs 52.9%）和长文理解（37.4% vs 37.4%）三个维度均匹配全注意力，同时在 32K 上下文推理时实现 4× 加速（标准分布）至 12× 加速（简单分布）。
贡献四：层间路由模式的可解释发现。分析显示浅层（1-3 层）和最后一层倾向使用 SWA，而中间层（8、11、14 层）更频繁使用全注意力。这与 Transformer 中浅层提取局部特征、深层聚合全局信息的认知一致。
4. 实验结果分析
实验在 28 个数据集上进行，涵盖常识推理（11 个）、上下文内检索（6 个）和长文理解（7 个）三大类。
| 模型 | 常识推理 (avg%) | 上下文检索 (avg%) | 长文理解 (avg) |
|---|
| SWA-ZS | 52.5 | 30.1 | 33.3 |
| SWA-CPT | 52.5 | 41.5 | 34.9 |
| StaticHybrid | 53.6 | 50.6 | 35.2 |
| SwiAttn | 53.7 | 52.9 | 37.4 |
| FullAttn | 53.4 | 52.9 | 37.4 |
关键发现：(1) 常识推理任务上各方法差距极小，因为大部分实例的上下文长度在 SWA 窗口（1024）内；(2) 上下文检索任务上 SwiAttn 比 StaticHybrid 提升 4.5%（相对 8.9%），比 SWA-CPT 提升 27.5%，验证了动态路由在信息检索场景的核心价值；(3) 长文理解上 SwiAttn 超越 StaticHybrid 6.3%，在 GovReport（25.0 vs 22.6）和 MultiFieldQA（35.3 vs 29.6）上优势最为明显。
消融实验（Table 4）在 GovReport、MultiNews、SAMSum 上进行：移除自适应正则化（V1）导致平均从 26.2 降至 24.9；α 参数敏感性较高（α=100 最优，α=1 或 α=0 显著下降），而 ε 参数较不敏感（ε=0.01~0.2 间变化有限）。这表明分支间 MSE 信号是自适应机制中最关键的组件。
Needle-in-Haystack 测试中，SwiAttn 在 32K 上下文、35 个深度百分比 × 35 个上下文长度配置下实现完美检索准确率，与全注意力一致。
5. 局限与未来方向
作者明确指出两个局限：(1) 缺乏硬件高效实现——在 prefill 阶段计算双分支的实现尚未针对 GPU 做硬件级优化，当前效率收益主要体现在 decode 阶段；(2) 规模验证不足——受限于计算资源，仅在 1.5B 模型上验证，未在更大模型（如 7B、13B）上确认方法的可扩展性。
另一个作者未提及但值得关注的局限是：路由器的训练依赖持续预训练（需 20B tokens），这一成本对于频繁更新的搜索排序模型来说可能偏高。此外，STE 在大规模训练中可能存在梯度估计不准确的问题，论文未讨论这一风险。
6. 工业实践意义
对搜索/推荐系统的直接价值：
(1) 检索模型效率优化：搜索排序中的 cross-encoder（如 BERT-based reranker）和 late-interaction 模型（如 ColBERT）面临长文档处理的效率瓶颈。SwiAttn 提供了一个即插即用的注意力替换方案——在平均仅使用 13% 全注意力的情况下匹配全注意力效果，这对在线 reranking 的 latency 约束极为有利。
(2) 推理阶段的实际加速：在 32K 上下文下 decode 阶段实现 4-12× 加速。对于搜索系统中需要处理长文档（如长帖子拼接、多文档联合排序）的场景，这一加速可直接转化为 QPS 提升和成本节约。
(3) 路由模式的可解释性启示：层间路由分析揭示了浅层偏好局部特征、中间层需要全局信息的模式。这一发现可指导搜索模型的结构设计——例如在知道浅层不需要全局注意力的前提下，可以安全地对浅层使用更高效的局部注意力，而将计算预算集中在中间层。
(4) 部署注意事项：当前方案的主要限制在 prefill 阶段缺乏硬件优化，因此在 batch 推理（离线 reranking）中的收益可能不如 decode-heavy 的在线生成场景显著。建议在实际部署前针对具体场景（prefill vs decode 比例）评估实际加速效果。训练成本方面，20B tokens 的持续预训练对基础设施团队来说是可接受的一次性投入。
Peking University