【第633期】Lighthouse Attention重构长文本AI底层逻辑

Seventy3：借助NotebookLM的能力进行论文解读，专注人工智能、大模型、机器人算法、crypto方向，让大家跟着AI一起进步。

如果你想要解读自己的论文，获得更多曝光度。请联系小助手微信：seventy3_podcast 加群。
合作邮箱：zhiwudazhanjiangshi#gmail.com

今天的主题是：

Long Context Pre-Training with Lighthouse Attention

Summary

在极长序列长度下训练因果 Transformer（causal transformers）受到缩放点积注意力（SDPA）平方级时间和内存复杂度的瓶颈制约。在这项工作中，我们提出了 Lighthouse Attention，这是一种仅在训练阶段使用的、基于对称选择的分层注意力算法。它包裹在普通的 SDPA 周围，并且可以在训练接近尾声时轻松移除。我们的分层选择还是无梯度的，这使我们无需处理复杂且可能效率低下的反向传播算子（backward pass kernel）。

我们的贡献主要有三个方面：

1. 一个亚平方（subquadratic）级的分层预处理和后处理步骤，用于对序列进行自适应压缩和解压缩。

2. 一种对称压缩策略，在保持自左向右因果性的同时，同步对 query、key 和 value 进行池化（pooling），这大大提高了并行度。

3. 一种两阶段训练方法：在绝大部分时间内使用 Lighthouse Attention 进行预训练，最后通过简短的训练恢复为全注意力（full attention）模型。

我们进行了初步的小规模大语言模型（LLM）预训练实验，结果表明，在所有其他设置均匹配的情况下，与全注意力训练相比，我们的方法非常有效，不仅实现了更快的总训练时间，而且在恢复阶段后获得了更低的最终损失（loss）。

原文链接：arxiv.org