Overview
Lighthouse Attention is a training-time attention replacement from Nous Research (Bowen Peng, Subho Ghosh, Jeffrey Quesnelle) targeting the Θ(N²d) cost of scaled dot-product attention (SDPA) in long-context pre-training. Unlike inference-only sparsity methods (H2O, SnapKV, HISA), Lighthouse guarantees that a model trained with it can be recovered to full dense SDPA inference — and the recovered models outperform dense-from-scratch baselines.
Key points
- Architecture — four-stage pipeline that keeps the attention kernel untouched: 1. Pyramid Pool: builds an L-level pyramid by non-overlapping mean pooling (factor p) of *Q, K, V symmetrically* — unlike NSA/HISA/InfLLM-V2 which pool only K/V. Total entries are O(N). 2. Hierarchical Selector: parameter-free scores (‖Q‖₂ / ‖K‖₂ norms, inherited via max-pool at coarser levels), fused chunked-bitonic top-k selection. A dilated-softmax scorer ablation shows ~0.01 loss difference at ~9% extra cost. 3. Gather + FlashAttention: selected entries form a contiguous subsequence of length S = N/p^{L-1} + (L−1)·p·k (≈65K at N=10⁶, L=4, p=4, k=4096), processed with stock FlashAttention and a standard causal mask; symmetric pooling means no gradient-free holes. 4. Scatter-Back: outputs are written to offset ranges R(l,i) = [i·p^l + p^l − 1, i·p^l + 2·p^l − 2], preserving causality; final output sequence is fully dense.
- Complexity: pyramid pooling, scoring, and scatter are Θ(Nd) / Θ(N log k); attention is Θ(S²d). Balancing with L = log_p(N/k) yields multi-logarithmic cost in N — versus Θ(N²d) for dense SDPA.
- Gradient flow: the top-k selector is non-differentiable by design — no straight-through estimator, no Gumbel softmax, no auxiliary scorer loss. Projections learn implicitly to be useful when selected.
- Two-stage training: pre-train a 530M Llama-3 model for 10k–12k steps with Lighthouse, then recover with dense SDPA for 4k–6k steps. Loss spikes (1.12–1.57) recover within ~1–1.5k steps and drop below the dense-from-scratch baseline by step 16,000 (0.6980–0.7102 vs 0.7237). Recovery is insensitive to the exact switch point.
- Performance (single B200, single layer): 21× faster forward, 17.3× forward+backward at 512K context; 1.4–1.7× faster full-model training vs cuDNN SDPA (530M, 98K context, 8×B200); advantage preserved at 1M context / 32 GPUs via context parallelism.
- Symmetric Q/K/V pooling assumes all queries coexist in one forward pass — broken in autoregressive decoding, so Lighthouse is a training-time method, not an inference method.
- Gathered attention is Θ(S²d): sub-quadratic but not strictly linear if k must grow with N.
- Experiments are at ablation-friendly scale (530M params, 50B tokens); frontier-scale validation is pending.
- Open directions: recovery via asymmetric sparse methods, per-layer/per-head adaptive k, multi-scale extension to vision/audio/video, serving integration.
- Paper: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle. "Long Context Pre-Training with Lighthouse Attention". arXiv:2605.06554
- Code: https://github.com/ighoshsubho/lighthouse-attention
- Nous Research write-up: https://nousresearch.com/lighthouse-attention/
- Podcast discussion: https://podcast.do-not-panic.com/episodes/long-context-pre-training-with-lighthouse-attention/
Ablations
| Dimension | Finding | |---|---| | Scorer | Parameter-free ≈ dilated-softmax (~0.01 gap), ~9% cheaper | | Pooling factor p | p=2 slightly better than p=4 | | Levels L | L=3 most balanced | | Top-k budget | Counter-intuitively, smaller k → lower loss (k=1536 best at 50B tokens: 0.6825, Pareto-optimal), suggesting hierarchical selection acts as regularization |
All Lighthouse configs matched or beat the 0.7237 dense-from-scratch baseline.
Positioning vs prior work
Selection logic lives entirely *outside* the attention kernel, unlike MoBA, NSA, and DSA which modify kernels. Benefits: inherits all FlashAttention optimizations, bit-identical forward/backward semantics to dense Transformers, standard ring-attention context parallelism without sparse-aware collectives.
Limitations and open questions
Significance
Lighthouse reframes long-context training economics: use O(N log N) attention for most of pre-training, briefly recover dense capability, and end with both lower cost and better final loss than training dense throughout. If it holds at frontier scale, it could substantially reduce the cost ceiling of 1M-context model training.