💬 NLP

Long Context Pre-Training with Lighthouse Attention

This paper introduces Lighthouse Attention, a training-only, gradient-free hierarchical selection algorithm that wraps standard scaled dot-product attention to enable efficient subquadratic pre-training of causal transformers at extreme sequence lengths, which can be seamlessly removed via a short recovery phase to yield a full attention model with faster training and lower final loss.

Original authors: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a massive library of books (a very long text) all at once to understand a story. In the world of AI, this is called "long-context training."

The problem is that standard AI models (Transformers) try to read every single word against every other word to find connections. If you have 100,000 words, the model has to do 10 billion comparisons. It's like trying to find a specific sentence in a library by shouting every word to every other word simultaneously. It takes forever, costs a fortune in electricity, and the computer runs out of memory.

This paper introduces a new method called Lighthouse Attention to solve this. Here is how it works, using simple analogies:

1. The Problem: The "All-Seeing Eye" is Too Heavy

Standard AI attention is like a security guard who insists on looking at every single person in a stadium at the same time to see who is talking to whom. As the stadium gets bigger, the guard gets overwhelmed.

2. The Solution: The "Lighthouse" Strategy

Instead of looking at everything at full resolution, Lighthouse Attention acts like a lighthouse scanning a dark ocean. It doesn't ignore the ocean; it just scans it in layers to find the most important parts quickly.

Here is the three-step process the paper describes:

Step A: The "Pyramid" (Summarizing the Crowd)

Imagine you have a huge crowd of people (the text). Instead of looking at every individual face, you group them into small clusters (like families or teams).

The Trick: The paper does something unique here. Most other methods only summarize the "people being talked about" (the keys and values) but leave the "people doing the talking" (the queries) in high detail.
Lighthouse's Move: It summarizes everyone equally. It creates a pyramid of summaries:
- Level 0: Every single word.
- Level 1: Groups of 4 words summarized into one.
- Level 2: Groups of 16 words summarized into one.
- And so on.
  This creates a multi-level map of the text, from "super detailed" to "big picture."

Step B: The "Spotlight" (Picking the Best Bits)

Now, the model needs to decide which parts of this pyramid are important enough to read in full detail.

The Selection: It uses a simple, free rule (no extra learning required) to score every group. It asks, "Which groups have the most 'energy' or importance?"
The Cut: It picks the top $K$ most important groups from all levels of the pyramid.
The Result: Instead of reading 100,000 words, the model now only needs to read a tiny, dense list of the most important "chunks" (maybe 5,000 words).

Step C: The "Flash" (Reading the Selected Bits)

Once the model has picked its top 5,000 words, it runs them through the standard, super-fast "FlashAttention" engine. Because the list is now short, this step is incredibly fast and fits easily in the computer's memory.

3. The "Magic" Recovery (The Two-Stage Training)

This is the most critical part of the paper. The authors were worried: "If we train the AI to only look at summaries, will it forget how to read full sentences later?"

To fix this, they use a Two-Stage Training Recipe:

Stage 1 (The Lighthouse Phase): They train the model for most of the time using the Lighthouse method. The model learns to be efficient and pick the right highlights.
Stage 2 (The Recovery Phase): For the last little bit of training, they turn off the "summarizing" and "picking" parts. They force the model to read the full text again using the standard method.

The Result: The model "wakes up" from its efficient training and remembers how to use full attention perfectly. The paper claims that after this short recovery, the model performs just as well (or better) than a model that was trained on the full text the entire time, but it got there much faster and cheaper.

Why is this a big deal?

Speed: The paper shows that for very long texts (like 100,000+ words), this method is 17 to 21 times faster than standard methods.
No "Junk" Code: Unlike other methods that require building custom, complicated computer chips (kernels) to handle the "picking" process, Lighthouse uses standard, off-the-shelf computer parts for the actual reading. It just rearranges the data before handing it over.
Symmetry: It treats the "question" and the "answer" parts of the sentence equally, which makes the math cleaner and more stable.

The Bottom Line

Lighthouse Attention is like hiring a smart research assistant. Instead of reading 1,000 pages of a report word-for-word, the assistant quickly scans the whole document, highlights the 50 most important paragraphs, and then reads only those 50 paragraphs in deep detail.

The paper proves that if you train an AI this way, and then give it a quick "refresher course" on reading the whole thing at the end, it becomes a super-fast reader that doesn't lose any of its intelligence.

Technical Summary: Long Context Pre-Training with Lighthouse Attention

1. Problem Statement

Training causal transformers at extreme sequence lengths (128K to 1M+ tokens) is currently bottlenecked by the quadratic time and memory complexity ( $\Theta(N^2)$ ) of scaled dot-product attention (SDPA). While methods like FlashAttention optimize constants, they do not alter the asymptotic complexity. Existing sparse attention approaches attempt to mitigate this by replacing dense attention with selection mechanisms where queries attend only to a subset of keys. However, these methods suffer from two primary design flaws when applied to pre-training:

Asymmetry: Most methods (e.g., NSA, HISA) pool Keys and Values but keep Queries at full resolution. This treats the hierarchy as a compressed memory rather than a multi-scale representation, limiting the model's ability to learn coherent interactions across scales.
Architectural Entanglement: Selection logic is often embedded within custom sparse attention kernels. This prevents the reuse of highly optimized dense attention kernels (like FlashAttention) and requires custom backward passes, complicating training and inference deployment.

Furthermore, a critical concern for training-time sparse methods is correctness: unlike inference-only sparsity (which assumes a competent dense backbone), a training-time sparse method must prove that the resulting weights are capable of functioning as a competent dense-attention model once the sparsity is removed.

2. Methodology: Lighthouse Attention

Lighthouse Attention is a training-only, symmetrical, selection-based hierarchical attention algorithm designed to wrap around ordinary SDPA without modifying the attention kernel itself. It operates as a four-stage pipeline:

A. Symmetric Pyramid Construction

Unlike prior works that pool only the context (K/V), Lighthouse applies symmetric average-pooling to Queries, Keys, and Values ( $Q, K, V$ ) simultaneously.

The sequence is organized into an $L$ -level pyramid with a pooling factor $p$ .
Level 0 is the full-resolution sequence. Each subsequent level $\ell$ summarizes $p$ consecutive entries from level $\ell-1$ .
This produces coherent triples $(Q^{(\ell)}, K^{(\ell)}, V^{(\ell)})$ at every level, ensuring that pooled queries and keys reside in the same representation space.

B. Parameter-Free Scoring and Selection

A hierarchical selector ranks pyramid entries to determine which tokens to attend to.

Scoring: The method uses a parameter-free scoring functional. For each entry, it computes the $\ell_2$ norm of the projected $Q$ and $K$ vectors. Coarser levels inherit scores from finer levels via max-pooling, ensuring that a coarse span retains the importance of its strongest token.
Selection: A fused chunked-bitonic top-K kernel selects the top- $K$ entries across all levels jointly.
Non-Differentiability: The top-K selection step is discrete and non-differentiable. No straight-through estimator or auxiliary loss is used. Gradients flow only through the gather, FlashAttention, and scatter-back operations into the projection matrices ( $W_Q, W_K, W_V$ ). Consequently, the projections learn to produce values that are useful when selected, rather than learning to game a scorer.

C. Dense Sub-Sequence Attention

The selected entries are gathered into a contiguous, causally consistent sub-sequence of length $S \ll N$ .

Gather: The selected triples are topologically sorted to form a dense sequence.
Attention: Standard FlashAttention (or stock SDPA) is executed on this sub-sequence. Because the input is dense and contiguous, no custom sparse kernel is required.
Scatter-Back: A deterministic scatter kernel redistributes the attention outputs back to their original positions in the full sequence.

D. Two-Stage Training Recipe

To ensure the model retains full attention capabilities, the authors propose a two-stage training approach:

Stage 1: Pre-train the majority of the time using Lighthouse Attention.
Stage 2: Resume training with the same checkpoint using full dense SDPA for a short period. This "recovery phase" allows the model to adapt its weights to the full attention mechanism without starting from scratch.

3. Key Contributions

Symmetrical Hierarchical Attention: A novel design that pools $Q, K, V$ symmetrically, creating a true multi-scale representation rather than a compressed memory, while preserving left-to-right causality.
Kernel Decoupling: By placing selection outside the attention path, the method utilizes stock FlashAttention on a dense sub-sequence. This allows the reuse of highly optimized dense kernels and simplifies the backward pass.
Training-Time Correctness: The paper introduces a rigorous empirical criterion: a model pre-trained with Lighthouse must match or beat a dense-from-scratch baseline after a brief dense-SDPA resumption. The authors demonstrate that Lighthouse-trained models satisfy this criterion, proving that hierarchical training does not "hollow out" the model's ability to use full attention.
Efficient Implementation: The use of fused chunked-bitonic top-K and custom scatter-back kernels enables efficient training at extreme context lengths (up to 1M tokens) on multi-node Blackwell GPUs.

4. Experimental Results

The authors conducted experiments using a 530M-parameter Llama-3-style decoder trained on the C4 dataset with a context length of 98,304 tokens.

Recoverability: Models pre-trained with Lighthouse and then resumed with dense SDPA for 4,000 steps achieved lower final training loss (0.6980–0.7102) compared to a dense-from-scratch baseline (0.7237) trained on the same token budget. This confirms that the hierarchical training signal does not degrade the model's dense attention capabilities.
Throughput and Speedup:
- Lighthouse achieved a 2x throughput advantage (tokens/sec/GPU) during the pre-training stage compared to dense SDPA.
- End-to-end, the two-stage recipe achieved a 1.40x to 1.69x wall-clock speedup over dense-from-scratch training for the same token budget, with matched or lower final loss.
- At a context length of 512K, Lighthouse was 21x faster on the forward pass and 17.3x faster on the forward+backward pass compared to dense SDPA.
Scaling: The method successfully scaled to 1M tokens using 32 Blackwell GPUs with context parallelism, maintaining the speedup advantage over dense baselines.
Ablations:
- Scorer: The parameter-free $\ell_2$ norm scorer performed comparably to a more expensive dilated softmax scorer but was ~9% faster.
- Hyperparameters: Smaller pooling factors ( $p=2$ ), shallower pyramids ( $L=3$ ), and smaller top-K budgets ( $k=1536$ ) generally yielded the best loss/performance trade-offs.
Retrieval (Needle-in-a-Haystack): After the dense-SDPA resumption, Lighthouse-trained models achieved retrieval rates comparable to or better than the dense baseline (e.g., 0.76 vs 0.72 mean retrieval rate), validating long-context retention.

5. Significance and Claims

The paper claims that Lighthouse Attention addresses the fundamental bottleneck of long-context pre-training by reducing the attention complexity from quadratic to sub-quadratic (effectively linear in $N$ for bounded $K$ ) while maintaining the architectural integrity of standard Transformers.

The authors emphasize that their method is training-only; the selection mechanism is removed (or rather, the model is transitioned) to a full dense model for inference. This distinguishes Lighthouse from inference-only sparse methods, which inherit a "correctness floor" from their dense backbones but never prove their sparse training process yields competent dense weights.

Limitations acknowledged by the authors:

The symmetric pooling assumes all queries co-occur in one forward pass, which is violated during autoregressive decoding. Therefore, the method relies on the dense-SDPA resumption to produce an inference-ready model.
The inner attention is $\Theta(S^2 d)$ on the gathered sub-sequence. While sub-quadratic in $N$ , it is not strictly linear if the selection budget $K$ must scale with $N$ .
The method currently requires a two-stage training recipe; a native sparse inference target (like DSA or NSA) is not yet integrated, though the authors suggest this as a future direction.

In summary, Lighthouse Attention provides a pathway to train models on million-token contexts with significantly reduced compute costs, while guaranteeing that the resulting model is fully capable of utilizing dense attention mechanisms at inference time.