Hugging Face Infini-Attention 复现分析

Infini-Attention 机制

Infini-Attention 旨在通过用固定大小的压缩记忆缓冲区取代标准自注意力的二次记忆增长，从而实现理论上的无限上下文长度。该过程遵循以下技术步骤：

  1. 分段：将输入序列划分为固定大小的片段。
  2. 局部注意力：在当前片段内计算标准的因果点积注意力。
  3. 记忆检索：模型使用当前片段的查询向量 ($Q$)，通过非线性激活函数（ELU + 1）从压缩记忆矩阵 ($M_{s-1}$) 中检索长期上下文。
  4. 整合：可学习的标量参数 ($\beta$) 充当门控机制，使用 sigmoid 函数平衡检索到的长期记忆 ($A_{\text{mem}}$) 与局部点积注意力 ($A_{\text{dot}}$) 的贡献。
  5. 记忆更新：通过将当前片段的键值状态加入现有缓冲区来更新压缩记忆。
  6. 状态转移：丢弃前一片段的注意力状态，仅将更新后的压缩记忆传递给下一片段。

复现挑战与收敛问题

在使用 200M Llama 模型和 Llama 3 8B 进行复现的过程中，Hugging Face 发现了若干关键的收敛障碍：

门控收敛

初始实验表明约 95% 的平衡权重集中在 0.5 附近，说明门控机制未收敛。分析发现标准的 Llama 3 8B 超参数（学习率 $3.0 \times 10^{-4}$）不足以让平衡因子显著偏离其初始化值。

为了解决此问题，团队为门控函数实现了单独的更高学习率（0.01），同时保持全局学习率为 $3.0 \times 10^{-4}$。这使得平衡因子能够达到更理想的范围，但在 200M 模型中在处理 200 亿 token 后最初出现了 NaN 损失。

权重衰减与 rollout

进一步检查发现，权重衰减会促使平衡因子的 L2 范数变小，使 sigmoid 值集中在 0.5 附近。为此，Hugging Face 采取了以下改动：

  • 移除 平衡因子的权重衰减。
  • 将 rollout 数量提升 至 16（片段长度为 64），以更强地激励模型使用压缩记忆。

这些调整使全局权重分布在 0 到 1 的完整范围内，且 10% 的头部权重达到 0.9 到 1.0 之间。

实验结果与评估

评估使用“密码检索任务”进行，其中特定的“针”（例如密码）隐藏在上下文中不同位置的无关文本中。

  • 小规模（200M 模型）：早期信号表明模型能够生成与前段相关的内容，但在检索精确的针时表现不佳。
  • Llama 3 8B：当针位于前段时，模型最初未通过针的评估。
  • 优化设置：在调整学习率并移除权重衰减后，模型在继续前段精确内容方面表现出改进，并在特定提示下通过了一些针密码测试，但仍不稳定。

结论与技术要点

尽管门控收敛有所改进，Hugging Face 仍然认为 Infini-Attention 尚未达到可用于生产的可靠性。团队坚持认为 Ring Attention、YaRN 和 RoPE scaling 仍是扩展预训练模型至更长上下文的首选方案。

实验的关键技术教训包括：

  • 门控敏感性：门控函数的正确收敛需要特定的学习率调节，并且要移除权重衰减以避免权重集中在 0.5 附近。
  • 损失 vs. 实用性：即使模型条件不佳或注意力输出存在维度错误，梯度下降仍可能降低训练损失，这意味着仅凭损失不足以衡量连贯性；必须进行持续评估。
  • 压缩权衡：随着记忆压缩次数的增加，性能必然下降。

Sources