Lighthouse Attention: Cut Long-Context Pretraining Costs and Keep Dense Inference

Lighthouse Attention: Cut Long‑Context Pretraining Costs, Keep Dense Inference

TL;DR — Attention’s quadratic scaling makes training long‑context models expensive. Lighthouse Attention speeds up pretraining for 100K–1M token contexts by selecting a compact dense subsequence during training, running highly optimized FlashAttention on that subset, then briefly resuming with full dense attention so the final checkpoint behaves like any standard model at inference. The practical payoff: big GPU‑hour savings for teams that need long contexts without building a special sparse serving stack.

For execs: bottom line and quick decision criteria

What you get: Reported end‑to‑end pretraining speedups of ~1.4×–1.69× on long‑context runs, plus layer/kernel level forward gains (e.g., ~21× forward at 512K tokens) — while preserving a dense model for inference.
Engineering effort: Expect a small POC (2–4 engineer‑weeks using the public repo), and several months to integrate in production (KV cache, orchestration, multi‑backends) depending on your stack.
Good fit if: Your roadmap includes models trained on 100K–1M contexts (legal, medical, enterprise search, multi‑document reasoning) and you want to accelerate iteration without changing inference runtimes.

The business problem

Attention’s compute and memory costs grow roughly with the square of the sequence length. For teams trying to train models that read huge documents or stitch many documents together, that quadratic cost is the choke point: experiments take too long and GPU bills balloon. Many sparse or linear attention methods reduce asymptotic cost but force you into specialized kernels and inference stacks. Lighthouse takes a pragmatic route: make training fast, keep inference standard.

What Lighthouse Attention does, in plain English

Think of Lighthouse as using a cheap sieve during training to pick the most useful tokens, running the same fast dense attention you already trust on that compact set, then stitching results back into the full sequence. After the bulk of training uses this selection‑based, hierarchical attention, you perform a short “dense resumption” phase with regular scaled dot‑product attention (SDPA). The final checkpoint is a normal dense‑attention model suitable for standard inference and deployment.

“Pool queries, keys and values symmetrically into a pyramid, select a compact dense subsequence, run stock FlashAttention on it, then scatter results back.”

How it works — stepwise, without the heavy notation

Pyramid pooling: Q, K, V are pooled symmetrically into a small multi‑level pyramid so you have anchors that summarize local neighborhoods. This is linear work in sequence length.
Cheap scoring & selection: Anchors are scored with a fast, parameter‑free metric (e.g., an L2 or projection‑norm scorer). A chunked, stratified top‑K selects the most promising anchors. This selection is executed outside the attention kernel and is non‑differentiable (indices are discrete).
Dense hotspot: Gather the selected dense subsequence of S tokens and run FlashAttention (the highly optimized dense attention kernel) on those S tokens.
Scatter results: Use integer‑atomic scatter kernels to write the computed outputs back into the full sequence representation.
Dense resumption (final polish): After most training steps use Lighthouse, switch the model to normal dense SDPA for a short resumption phase so the checkpoint becomes a fully dense attention model for inference.

Because selection happens outside the attention kernel, Lighthouse reuses existing, optimized FlashAttention rather than requiring a bespoke sparse attention kernel. That makes the approach operationally attractive: you get runtime wins during training while keeping inference unchanged.

Why the non‑differentiable selection doesn’t break learning

The top‑K indices are discrete and carry no gradient, but gradients still flow through the gathered Q/K/V and the projection weights that produce them. In practice, the model learns to produce representations that are useful when those discrete indices pick them. The final dense resumption step acts as a recovery and polishing phase, ironing out transient gaps introduced by selection so the final loss and retrieval performance match or beat dense‑from‑scratch baselines.

“Selection indices are discrete and carry no gradient; projections learn to produce values that are useful when selected.”

Benchmarks & evidence — real numbers to anchor expectations

Model: Llama‑3‑style 530M decoder (dmodel=1024, 30 layers, 8 heads, head dim 128).
Hardware: single‑device B200 GPU (per‑device context experiments up to 512K), and scaled runs to 1M tokens across 32 Blackwell GPUs using context parallelism.
Layer/kernel gains: ~21× faster forward pass at 512K tokens vs cuDNN SDPA; ~17.3× faster forward+back.
End‑to‑end pretraining: observed wall‑clock speedups of ~1.40×–1.69× depending on split and config (same token budget comparisons).
Stage‑1 throughput: ~84K–126K tokens/s/GPU with Lighthouse vs ~46K tokens/s/GPU for dense SDPA. Projection‑norm scorers pushed the high end (~126K tokens/s/GPU).
Training recipe & recovery: short dense resumption (~1K–1.5K SDPA steps) stabilizes a transient loss spike (1.12–1.57 nats reported) and recovers or improves final loss. Example matched token budget (≈50.3B tokens): Lighthouse resumed runs 0.6980–0.7102 vs dense baseline 0.7237 (lower is better).
Retrieval tradeoffs (Needle‑in‑a‑Haystack): dense baseline mean retrieval 0.72; Lighthouse configurations reached 0.76 (k=2048, dilated scorer) and showed that larger k helps retrieval while the cheaper norm scorer can slightly reduce retrieval quality.

Tradeoffs, limitations and open questions

Training‑only: Lighthouse speeds training but does not reduce inference latency for single‑step autoregressive decoding—final checkpoints are dense models.
Tuning required: Hyperparameters such as k (selected tokens), pyramid depth L, pooling factor p, and scorer choice materially affect loss vs retrieval. Some experiments even showed smaller k helped final loss, suggesting implicit regularization effects that are not fully understood.
Serving integration: Production autoregressive serving (KV cache, continuous batching) and partial online selection are future work; Lighthouse doesn’t remove the need to design efficient inference systems for low‑latency serving.
Engineering effort: Implementing chunked‑bitonic top‑K and robust integer‑atomic scatter kernels across PyTorch/XLA/TPU stacks requires nontrivial engineering and testing.
Scaling rules: Regimes where k must grow with N are not fully characterized; adaptive per‑layer or per‑head k strategies are promising but unexplored at scale.

Proof‑of‑concept checklist (2–4 week POC using public repo)

Pick a representative workload: 500M–1B parameter model, 100K context target, in‑domain dataset (e.g., curated C4 subset or internal corpus).
Baseline run: train with dense SDPA for a fixed token budget; record wall‑clock, GPU‑hours, final loss and retrieval metrics.
Lighthouse run: same token budget, use Lighthouse for majority of steps and switch to dense resumption for ~1K–2K steps; record the same metrics.
Compare: end‑to‑end pretraining time, forward/back throughput, final loss delta, retrieval recall, and required engineering hours to integrate the repo.
Success criteria: ≥1.3× end‑to‑end speedup, final loss within 0.02 (absolute) of baseline or better, stable resumption within 2K steps.

Example ROI sketch (illustrative): Suppose a baseline run costs 1,000 GPU‑hours at $3/hr = $3,000. A 1.5× speedup reduces usage to ~667 GPU‑hours, saving ~333 GPU‑hours or $1,000. If the POC costs 4 engineer‑weeks (~640 engineer‑hours), and engineers bill at $80/hr, the projected engineering cost is high, but for repeated experiments or large projects the GPU‑hour savings compound across multiple runs. Use a concrete run mix to estimate real ROI for your team.

Practical recommendations

Start small: validate Lighthouse on a 500M–1B scale before committing it to larger budgets.
Keep a short dense resumption in every workflow—it’s the operational insurance policy that converts a training‑only sparse checkpoint into a standard dense model.
For retrieval or “needle‑in‑a‑haystack” tasks, bias toward larger k and the dilated scorer; for raw throughput, projection‑norm scorers are cheaper.
Retain boundary layers as dense (early and late layers) to preserve stability during selection‑based training.
Plan engineering time for selection kernels and multi‑node context parallelism if you target 1M contexts.

Key questions (quick Q&A)

Can you train for extremely long contexts faster without changing the final model’s inference behavior?

Yes. Use Lighthouse for the majority of training and perform a short dense resumption to convert the checkpoint to dense SDPA for inference.
How large are the runtime gains at extreme contexts?

Huge at the kernel/layer level (e.g., ~21× forward at 512K) and meaningful end‑to‑end: ~1.4×–1.69× pretraining speedups in reported experiments.
Do discrete, non‑differentiable selection steps break optimization?

No. Gradients flow through the gathered Q/K/V and projection weights adapt; a brief dense resumption stabilizes any transient effects.
Will Lighthouse reduce my single‑step autoregressive inference latency?

No. Lighthouse is a training‑time acceleration; final inference runs remain dense.

Resources & next steps

Paper: arXiv:2605.06554 — Long Context Pre‑Training with Lighthouse Attention.
Implementation: ighoshsubho/lighthouse-attention (GitHub).
Related reads: Saipien pieces on long‑context models, AI automation, and retrieval‑augmented agents (internal links recommended).

Lighthouse Attention is a pragmatic engineering tradeoff: aggressive, selection‑based compression during training for big iteration and cost wins, followed by a short dense resumption to deliver a conventional dense model. For teams investing in long‑context capabilities, it’s a practical lever to accelerate experiments and lower GPU bills without demanding a forever‑changed inference stack.