Nous Research's Lighthouse Attention Trains 1.69x Faster on Million-Token Contexts

Nous Research's Lighthouse Attention cuts long-context LLM pre-training time by 1.4–1.7x, running 17x faster than standard attention at 512K context — with no custom sparse kernel and no quality loss at inference.

·
·
·
Nous Research's Lighthouse Attention Trains 1.69x Faster on Million-Token ContextsPRO
  • Lighthouse Attention is a training-only hierarchical sparse attention from Nous Research that wraps standard FlashAttention and is removed before inference.
  • It achieves a 1.4–1.7x wall-clock speedup over dense SDPA pre-training at 98K context, saving 75–106 B200-hours per 50B-token run.
  • At 512K context on a single B200, the forward+backward pass runs 17.3x faster than standard attention with no custom sparse kernel.
  • A two-stage recipe (Lighthouse pre-training + brief dense SDPA resume) produces models that match or beat a dense-from-scratch baseline on training loss and NIAH retrieval.
  • The method scales to 1M-token context across 32 Blackwell GPUs using standard ring attention with no sparse-aware collectives.
  • Full code is released as a patch on torchtitan with configs covering the paper's complete ablation grid.

Training a transformer on sequences longer than 100K tokens is expensive in a very specific way: attention cost grows quadratically with sequence length. Double the context, quadruple the compute. Lighthouse Attention, released by Nous Research, attacks this problem at the training stage with a hierarchical selection mechanism that wraps around standard attention , and can be completely removed before inference, leaving behind a normal dense model.

The wall that FlashAttention can't move

Scaled dot-product attention (SDPA) has O(N²) time and memory complexity. FlashAttention reduces constants but not asymptotics; at sequence lengths above 100K, the attention term dominates all other compute in the model. The frontier of language modeling has moved to 128K, 1M, and longer contexts , pushed by agentic reasoning, long-document understanding, and multimodal inputs , making this bottleneck increasingly costly.

Prior work tried to fix this with sparse attention: instead of attending to every token, each query attends to a small subset. Methods like MoBA and Native Sparse Attention (NSA) select contiguous blocks; DeepSeek Sparse Attention (DSA) uses a learned indexer to pick the top-k tokens. These methods produce meaningful inference speedups but inherit two design decisions that fit long-context pretraining poorly: asymmetry (queries stay at full resolution while keys and values are pooled) and architectural entanglement (selection lives inside the attention kernel, so carefully optimized dense-attention kernels cannot be reused).

A pyramid that scores itself

Lighthouse replaces a standard attention layer with a four-stage pipeline that surrounds, but does not modify, the attention kernel: a pre-attention selection stage drives a contiguous gather, stock FlashAttention runs on the gathered sub-sequence, and a post-attention scatter writes the result back to the original positions. The key insight is that the selection logic is entirely decoupled from the attention kernel itself.

Needle-in-a-Haystack retrieval heatmaps comparing Lighthouse variants against the dense SDPA baseline across context lengths up to 96K

Here's how the four stages work:

  1. Pyramid Pool: Q, K, and V are all average-pooled symmetrically into an L-level pyramid with pooling factor p. Each level summarizes p times more tokens than the level below it. Unlike all prior hierarchical sparse methods (NSA, HISA, InfLLM-V2), which only pool the key-value side, Lighthouse pools queries too.

Pro article

This story is for Pro members

You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.

Trending
  • No trending articles

Comments

avatar

Next Reads