用分层选择机制突破长序列训练瓶颈，训练更快更省内存。
Long Context Pre-Training with Lighthouse Attention
- 通过分层压缩与解压，实现次二次复杂度的注意力计算。
- 对查询、键、值同步压缩，保持因果关系，提升并行效率。
- 两阶段训练：前期用新方法，后期快速恢复完整注意力模型。
在极长序列下训练因果Transformer受限于缩放点积注意力（SDPA）的二次时间与内存开销。本文提出Lighthouse Attention，一种仅用于训练的对称选择式分层注意力算法，可嵌入普通SDPA且训练末期轻松移除。其分层选择无需梯度，避免了复杂低效的反向传播。贡献有三：(i) 增量式预处理与后处理步骤，实现序列的自适应压缩与解压；(ii) 对称压缩策略，同时处理查询、键、值，保持左到右因果性，显著提升并行性；(iii) 两阶段训练：大部分时间使用Lighthouse Attention预训练，最后以短时训练恢复完整注意力模型。小规模大语言模型预训练实验表明，相较全注意力训练，在其他条件相同的情况下，该方法总训练时间更短，最终损失更低。完整代码已开源。
原文摘要 · Abstract (English)
Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。