Overview
Lighthouse Attention, from Bowen Peng, Subho Ghosh, and Jeffrey Quesnelle at Nous Research, is not another inference-time sparse attention trick. It is a training-phase attention replacement: most of pre-training uses hierarchical sparse attention, followed by a brief switch back to dense SDPA, yielding a model fully usable with full-attention inference.
This "recoverable-to-dense after training" correctness criterion is the fundamental distinction from prior sparse methods. Inference-only approaches (H2O, SnapKV, HISA) rely on a dense backbone as a quality floor. Training-time sparse methods must answer a harder question: after training, can the model still run with dense attention? Lighthouse's answer: yes—and the recovered model outperforms dense-from-scratch training.
The Problem: O(N²) Long-Context Training
Standard scaled dot-product attention costs Θ(N²d) in compute and memory. FlashAttention improves constants, not asymptotics. For N ≥ 10⁵, attention dominates training cost. Doubling context quadruples attention compute: going from 128K to 1M context is 8× the length but 64× the compute—directly capping frontier-model training budgets.
Architecture: A Four-Stage Pipeline
Lighthouse replaces the attention layer with a four-stage pipeline, leaving the attention kernel itself untouched—selection logic lives outside FlashAttention.
1. Pyramid Pool
Given Q, K, V ∈ R^{N×d}, build an L-level pyramid by non-overlapping mean pooling with factor p per level: level 0 is the full-resolution sequence, level 1 pools every p tokens, level 2 every p², and so on.
Key design: symmetric pooling—Q, K, and V are all pooled (unlike NSA, HISA, InfLLM-V2, which pool only K/V). This means pooled Q(l) and K(l) share one representation space and each pyramid entry is a coherent (Q, K, V) triple summarizing p^l tokens. Total entries are Σ N/p^l ≤ N·p/(p−1), i.e., O(N), with linear construction cost.
2. Hierarchical Selector
Each entry gets two scalar scores (as query, as key):
- Level 0: sQK_{0,i} = ||Q_i||₂, sKQ_{0,i} = ||K_i||₂
- Coarser levels inherit scores via max-pool from level 0 rather than recomputing
- 512K context: 21× faster forward, 17.3× forward+backward vs SDPA
- Equivalently, SDPA needs ~113K context to match Lighthouse at 512K
- Full-model training (530M params, 98K context, 8×B200): 1.4–1.7× faster than cuDNN SDPA
- 1M context / 32 GPUs: advantage cleanly preserved via context parallelism
- Recovery is insensitive to the switch point (10k/11k/12k all work)—no magic schedule needed
- A longer dense recovery tail yields lower final loss
- Hierarchical training signal does not erode full-attention capability
- Paper: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle. "Long Context Pre-Training with Lighthouse Attention". arXiv:2605.06554, 2026-05-07
- 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/
Max-pooling lets a coarse span inherit importance if it contains an important token. All levels' scores are concatenated and top-k entries are selected with a fused chunked-bitonic top-K kernel.
The scorer has no learnable parameters—deliberately: any positive result above this is a lower bound for richer scorers. Ablations show a dilated-softmax scorer performs within ~0.01 loss but costs ~9% more.
3. Gather + FlashAttention
Selected entries from different levels are gathered into a contiguous subsequence of length S = N/p^{L-1} + (L−1)·p·k, then processed with stock FlashAttention. At N=10⁶, L=4, p=4, k=4096: S ≈ 6.5×10⁴, far below N.
The causal mask derives from pyramid coordinates (entries only attend to entries with smaller base positions). Because Q/K/V are pooled symmetrically, the gather has no "holes"—no tokens are cut off from gradients, which asymmetric pooling cannot guarantee.
4. Scatter-Back
Outputs Ō ∈ R^{S×d} are redistributed to N original positions; entry (l, i) writes to offset range R(l,i) = [i·p^l + p^l − 1, i·p^l + 2·p^l − 2]. The p^l − 1 offset guarantees causality. Per-position fan-in is bounded by L, independent of k. The final output sequence is fully dense—a compressed approximation of full attention.
Gradient Flow: Non-Differentiable Selection
Top-K selection is non-differentiable, yet Lighthouse uses no straight-through estimator, no Gumbel softmax, no auxiliary scorer loss. The selector branch receives no gradient; gradients flow only loss → scatter → FlashAttention → gather → pyramid pool → W_Q/W_K/W_V.
The projections thus learn not to "raise scores" but to make selected Q, K, V useful when selected—an implicit optimization that sidesteps the optimization fragility of learnable scorers.
Complexity: O(N²) → O(N log N)
| Stage | Complexity | |---|---| | Pyramid Pool | Θ(Nd) | | Scoring | Θ(N log k) | | Top-K Selection | Θ(N log k) | | Gather + FlashAttention | Θ(S²d) | | Scatter-Back | Θ(Nd) |
Balancing with L = log_p(N/k) gives S = Θ(k·log_p(N/k)) and attention cost Θ(k² log²N · d)—polylogarithmic in N for fixed k. Total compute is linear in N plus a log k factor.
Benchmarks (single B200, single attention layer):
Two-Stage Training: Hierarchical-to-Dense Recovery
Stage 1: pre-train with Lighthouse (most steps). A 530M Llama-3 model used Lighthouse for the first 10k–12k steps.
Stage 2: recover with dense SDPA for the final 4k–6k steps, same optimizer state and data stream. Loss spikes initially (1.12–1.57) as the model first sees untrained attention, but recovers within ~1–1.5k steps and by step 16,000 beats the dense-from-scratch baseline (0.6980–0.7102 vs 0.7237).
Key findings:
This is a strong empirical claim: training mostly on selected tokens, then exposing all tokens, lets the model exploit full attention better than training densely from scratch.
Ablations: Four Knobs
| Dimension | Range | Finding | |---|---|---| | Scorer | projection-norm vs dilated-softmax | gap ~0.01; parameter-free is ~9% cheaper | | Pooling factor p | 2, 4 | smaller p slightly better | | Levels L | 2, 3, 4 | L=3 most balanced | | Top-k budget k | 1536–6144 | counterintuitively, smaller k gives lower loss (down to 1536) |
At a 50B-token budget, k=1536 with the dilated scorer achieves the lowest loss (0.6825) and is Pareto-optimal. The authors conjecture hierarchical selection acts as regularization—forcing focus on the most important tokens aids generalization under limited budget; whether this reverses at larger budgets is future work. All configurations matched or beat the dense baseline (0.7237), so recoverability is not a hyperparameter fluke.
Comparison with Existing Methods
| Method | Mechanism | Training-time? | Symmetric pooling? | Selection in-kernel? | Learnable scorer? | |---|---|---|---|---|---| | Lighthouse | hierarchical + selection | ✅ | ✅ | ❌ | ❌ | | MoBA | block selection | ✅ | ❌ | ✅ | ❌ | | NSA | block + learned | ✅ | ❌ | ✅ | ✅ | | DSA | token-level + learned | ✅ | ❌ | ✅ | ✅ | | HISA | hierarchical index | ❌ (inference) | ❌ | ✅ | ❌ | | H2O/SnapKV | KV-cache eviction | ❌ (inference) | ❌ | ❌ | ❌ | | Linear Attention | state compression | ✅ | N/A | N/A | ❌ |
Because selection lives entirely outside the kernel, Lighthouse inherits all FlashAttention optimizations (including future ones), keeps forward/backward bit-identical to a dense transformer (no custom sparse backward), and supports standard ring-attention context parallelism without sparse-aware collectives.
Limitations and Open Questions
1. Symmetric Q/K/V pooling assumes all queries coexist in one forward pass. Autoregressive decoding breaks this—each new token produces a single query—so Lighthouse is a training-time technique, not an inference one; deployment requires dense-SDPA recovery first. 2. Gathered attention costs Θ(S²d)—sub-quadratic for fixed k, but not strictly linear. If k must grow with N to preserve recall, complexity is super-linear; these regimes are uncharacterized. 3. Scale of experiments. 530M parameters, 16K steps, 50B tokens—ablation-friendly scale. Frontier-scale (70B–1T) validation is explicitly pending.
Open questions include recovering with asymmetric sparse methods (DSA, NSA, MoBA) for natively servable checkpoints, per-layer/per-head adaptive k budgets, extending the pyramid to multi-scale vision/audio/video, and serving integration (continuous batching, speculative decoding, KV-cache management).
Significance
Lighthouse's contribution is a training-phase compute alternative: pre-train with O(N log N) attention for most of the run, spend a short tail recovering dense capability—less total time, better final results. If this reproduces at frontier scale, it changes the economics of long-context training: a 1M-context model might cost 1/1.7× or less of dense training cost, without sparse attention from scratch.
More deeply, it validates a "compress information during training, release it during recovery" strategy—build coarse understanding under limited information, then refine under full information—offering a possible explanation for why recovered models beat dense-from-scratch.