🤖 AI Summary
This work addresses the quadratic computational complexity of standard dot-product attention, which hinders the training of causal Transformers on extremely long sequences. The authors propose Lighthouse Attention—a symmetric, selection-based hierarchical attention mechanism used exclusively during training—that adaptively compresses sequences by jointly pooling queries, keys, and values through a gradient-agnostic hierarchical selection and symmetric compression strategy. This approach enhances parallelism while strictly preserving causality. Integrated within a two-stage training paradigm—initial pretraining with Lighthouse Attention followed by fine-tuning with full attention—the method substantially reduces training costs. Experiments on small-scale language models demonstrate that, under identical settings, the proposed approach achieves shorter total training time and yields lower final loss after the recovery phase.
📝 Abstract
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