Post

multi_doc_qa 从 0.086 涨到 0.200: SWA + LowRank Gated + Hybrid RoPE 的长文配方

multi_doc_qa 从 0.086 涨到 0.200: SWA + LowRank Gated + Hybrid RoPE 的长文配方

某 3B 模型的 baseline 在 multi_doc_qa 上只有 0.086(几乎是 random)。加上 SWA + LowRank Gated + Hybrid RoPE 三件套后涨到 0.200——2.3 倍提升。最关键的组件不是 SWA(它只解决 OOM),而是 Hybrid RoPE(它让 attention 能做跨文档的内容匹配)。


1. 背景:长文 Fine-tuning 阶段的架构选择

模型预训练阶段已经引入了 SWA + Gated Attention,但 16K+ 的长上下文能力需要专门的 fine-tuning 阶段(FD stage)来激活。问题是:attention 机制 + position encoding 的哪种组合在 FD 阶段表现最好?

我们测试了 4 个变体,逐步叠加组件。


2. 四个变体与评测结果

ModelAveragelongbench_16kmulti_doc_qapassagesingle_doc_qasummary
baseline0.1670.1390.0860.0270.2080.206
SWA+ElemGated0.1870.1550.1100.0270.2390.212
SWA+ElemGated+HybridRoPE0.2090.1740.1290.0300.2950.204
SWA512+LowRank128+HybridRoPE0.2160.1790.2000.0300.2300.217

关键数字:

  • Hybrid RoPE 单独贡献 +2.2% precision(0.187 → 0.209)
  • LowRank128 Gated 再加 +0.7%(0.209 → 0.216)
  • multi_doc_qa 的绝对增益最大:0.086 → 0.200,2.3x

3. 各组件的作用

graph TD
    A[Long-Context FD Recipe] --> B[SWA w=512]
    A --> C[Gated Attention]
    A --> D[Hybrid RoPE]
    
    B --> B1[解决 OOM]
    B --> B2["O(n*w) local + full attention on global layers"]
    
    C --> C1[Elementwise Gate]
    C --> C2[LowRank128 Gate]
    C1 --> C1a["per-dim sigmoid, H params/head"]
    C2 --> C2a["bottleneck projection dim=128"]
    C2 --> C2b["学习高层次信息选择"]
    
    D --> D1["部分维度使用 position encoding"]
    D --> D2["其余维度 position-free"]
    D2 --> D2a["纯 content-based attention"]
    D2 --> D2b["长距离依赖不受位置约束"]

SWA (Sliding Window Attention, w=512)

强制局部层只在窗口内做 attention,节省长序列显存:

  • 无 SWA:O(n^2) attention → 16K+ 直接 OOM
  • 有 SWA:局部层 O(n·w),仅全局层做 full attention

SWA 是必要条件(没它跑不起来),但不是性能增益的主要来源。

Gated Attention

两种实现:

  • Elementwise:每个 attention head 对输出做 per-dimension sigmoid gate,参数量 = H per head
  • LowRank (dim=128):attention 输出先投影到 128 维 bottleneck,再做 gating

LowRank gate 在更高抽象层次上学习”传递哪些信息”——对 multi-document reasoning 特别有效,因为模型需要在多个文档间选择相关信息。

Hybrid RoPE

标准 RoPE 对所有 attention 维度施加 position encoding。Hybrid RoPE 只对部分维度使用位置编码,其余维度完全 position-free。

position-free 维度的 attention 完全基于内容匹配,不受相对位置影响。这对跨文档的长距离依赖至关重要——不同文档中的相关信息,位置可能相距数千 token,但内容语义是对齐的。


4. 为什么 Hybrid RoPE 是 multi_doc_qa 的关键

Multi-document QA 的核心需求:

  1. 跨文档信息检索:答案可能在第 3 个文档的某段,而问题在最后——纯位置编码会让模型偏好”近”的 token
  2. 跨文档信息对比/综合:需要比较不同位置的内容片段
  3. Position 不是有效信号:哪个文档排第几跟答案无关

Hybrid RoPE 的 position-free 维度天然适配这三个需求。模型可以在这些维度上做纯 content-based matching,不被位置距离”惩罚”。

数据验证:multi_doc_qa 从加入 Hybrid RoPE 开始跳涨(0.110 → 0.129 → 0.200),而 passage retrieval(强位置依赖)几乎不变(0.027 → 0.030)。


5. LowRank vs Elementwise Gating:什么时候 bottleneck 更好

指标ElementwiseLowRank128Delta
multi_doc_qa0.1290.200+0.071
single_doc_qa0.2950.230-0.065
average0.2090.216+0.007

LowRank gate 的 bottleneck 强制模型学习”粗粒度的信息选择策略”——在多文档场景下,这等价于学习”哪个文档的信息值得传递”。

Elementwise gate 保留了更细粒度的 per-dimension 控制,在 single_doc_qa(局部推理、位置相关)中表现更好。

直觉:multi_doc_qa 是一个”选择哪个文档”的高层决策问题,LowRank 的 abstraction 恰好匹配这个粒度;single_doc_qa 是”精确定位文档内位置”的低层问题,Elementwise 的细粒度控制更合适。


6. Trade-off 与工程决策

最终选择 SWA512 + LowRank128 + Hybrid RoPE,接受了一个 trade-off:

  • multi_doc_qa: +0.071(0.129 → 0.200)
  • single_doc_qa: -0.065(0.295 → 0.230)

为什么接受这个 trade-off:

  • multi_doc_qa 更代表真实长上下文使用场景(用户不会只给一个文档然后问问题)
  • single_doc_qa 的 0.230 仍然远高于 baseline 的 0.208
  • 总 average 仍然是最高的(0.216 vs 0.209)

7. Loss 排序确认层级

训练 loss 的排序:

SWA+LowRank+HybridRoPE < SWA+Elem+HybridRoPE < SWA+Elem < baseline

(lower = better)

Loss 排序与下游评测完全一致。这说明这些组件的增益不是”碰巧在某个 benchmark 上有效”,而是genuinely 提升了模型的长文本建模能力。Loss 作为更稳定的信号,给了我们更强的信心。


8. 长文配方的工程经验

组件叠加顺序有逻辑

  1. SWA 解决”能不能跑”——没它 16K 直接 OOM
  2. Gated Attention 解决”信息选择”——不是所有 token 都值得传递
  3. Hybrid RoPE 解决”跨位置匹配”——长距离内容对齐不应被位置惩罚

不要只看 average

如果只看 average(0.209 vs 0.216),会觉得 LowRank gate 的提升微不足道。但拆开看 multi_doc_qa(+0.071),才能理解它的真正价值——以及代价(single_doc_qa -0.065)。

Hybrid RoPE 的适用条件

  • 任务越依赖”跨位置内容匹配”,Hybrid RoPE 收益越大
  • 任务越依赖”位置本身就是信号”,Hybrid RoPE 收益越小甚至有害
  • multi_doc_qa 是前者的典型;passage retrieval 接近后者

LowRank gate 的适用条件

  • 信息选择粒度越粗(”选哪个文档”),LowRank 越好
  • 信息选择粒度越细(”选哪个 token”),Elementwise 越好

9. References

  1. Beltagy et al., “Longformer: The Long-Document Transformer,” arXiv:2004.05150 — SWA 的经典实现
  2. Su et al., “RoFormer: Enhanced Transformer with Rotary Position Embedding,” arXiv:2104.09864 — RoPE 原始论文
  3. Child et al., “Generating Long Sequences with Sparse Transformers,” arXiv:1904.10509 — local + global attention pattern
  4. Press et al., “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation,” arXiv:2108.12409 — position encoding 对长度外推的影响
This post is licensed under CC BY 4.0 by the author.