带有 Lighthouse Attention 的长上下文 预训练

在极端序列长度下训练因果 Transformer 受到缩放点积注意力 (SDPA) 的二次时间和内存的瓶颈。在这项工作中，我们提出了 Lighthouse Attention，这是一种仅训练的基于对称选择的分层注意力算法，它包裹着普通的 SDPA，并且可以在训练结束时轻松删除。我们的分层选择也是无梯度的，这使我们免于处理复杂且可能低效的后向传递内核。我们的贡献有三个方面：（i）次二次分层预处理和后处理步骤，对序列进行自适应压缩和解压缩。 (ii) 对称压缩策略，同时池化查询、键和值，同时保留从左到右的因果关系，这大大提高了并行性。 (iii) 两阶段训练方法，我们大部分时间都使用 Lighthouse Attention 进行预训练，并在最后通过短期训练恢复完整的注意力模型。我们进行了初步的小规模 LLM 预训练 实验，显示了我们的方法与所有其他设置匹配的全注意力训练相比的有效性，其中我们在恢复阶段后实现了更快的总训练时间和更低的最终损失。完整代码位于：https://github.com/ighoshsubho/lighthouse-attention