PaperlayerView exact revision on arXiv

Source-linked AI summary

Long Context Pre-Training with Lighthouse Attention

Bowen Peng, Subho Ghosh, Jeffrey Quesnelle

arXiv:2605.06554v1cs.CL

TL;DR

Extreme-length training is limited by the quadratic compute and memory of SDPA. Lighthouse Attention adds gradient-free hierarchical selection around stock attention, then briefly resumes dense training; experiments report 1.4–1.7× speedups while matching or beating dense-from-scratch training loss.

  • Problem

    Training at extreme sequence lengths is bottlenecked by SDPA’s Θ(N^2) compute and memory, while training-time sparse methods must still produce competent dense-attention models.

  • Method

    Lighthouse Attention symmetrically pools Q, K, and V across a multi-resolution pyramid, selects a causal dense sub-sequence around stock FlashAttention, and trains without gradient-based top-K selection or auxiliary losses.

  • Results

    1.4–1.7× end-to-end speedups against cuDNN SDPA at ≥100K context accompany dense-SDPA resumption that matches or beats dense-from-scratch training loss at matched tokens.

  • Takeaways & Limitations

    A brief dense-SDPA resumption recovers a full-attention model while preserving the reported training-time efficiency gains.

  • Takeaways & Limitations

    Symmetric Q/K/V pooling assumes all queries co-occur in one forward pass, so autoregressive decoding requires dense-SDPA resumption; scaling regimes where k grows with N remain uncharacterised.

Abstract

from arXiv · show

Training causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention

1 Introduction

Lighthouse Attention targets the quadratic cost of long-context SDPA with hierarchical selection designed for training-time correctness. It uses symmetric multi-scale pooling and a brief dense resumption, with reported matching or better training loss than dense training.

  • Extreme-context training is bottlenecked by SDPA’s Θ(N^2) compute and memory, despite FlashAttention reducing constants.
  • Lighthouse Attention pools Q, K, and V symmetrically across a multi-level pyramid, scores entries bidirectionally, and selects a causal dense sub-sequence for stock FlashAttention.
  • The selection branch is non-differentiable and gradient-free, while gradients flow through scatter, FlashAttention, gather, and the projection matrices.
  • A brief dense-SDPA resumption lets Lighthouse-trained models match or beat dense-from-scratch models on training loss using the same token budget.
  • The contribution combines hierarchical attention, fused top-K and scatter kernels, and a training-correctness criterion based on dense-SDPA resumption.

2 Related works

Prior long-context methods compress states or select blocks or tokens, often using key-value hierarchies and custom sparse kernels. Lighthouse instead ranks symmetric Q/K/V pyramid representations while retaining stock FlashAttention on a gathered dense sequence.

  • Compression-based approaches improve asymptotics by replacing softmax attention with bounded states, but can compress the entire past and limit long-range recall.
  • Block-pruning methods make a single retain-or-discard decision per block, while token-level methods select smaller subsets of keys or past tokens.
  • Existing hierarchical methods generally apply hierarchies to keys and values and feed selected content into custom sparse-attention kernels.
  • Lighthouse differs by pooling queries symmetrically with keys and values, using the pyramid only for ranking, and applying stock FlashAttention to a dense sub-sequence.
  • Figure 1 depicts Lighthouse as a pipeline surrounding, rather than modifying, the attention kernel with selection, gather, stock attention, and scatter-back stages.

3 Method

Lighthouse Attention replaces dense attention with a four-stage hierarchical pipeline that symmetrically pools Q, K, and V, selects pyramid entries, applies stock FlashAttention to a causal gathered subsequence, and scatters outputs back. The design uses fixed pooling, parameter-free scoring, discrete top-k selection, and deterministic reconstruction while preserving dense outputs and causality.

  • Pipeline: Lighthouse surrounds, rather than modifies, the attention kernel with selection, contiguous gathering, stock FlashAttention, and postattention scatter-back.The pipeline replaces standard attention while retaining the same FlashAttention call used by the dense baseline.
  • Optimization: The top-k indices are discrete and non-differentiable, so gradients flow through gather, attention, and scatter into W_Q, W_K, and W_V without a straight-through estimator.Lighthouse adds no learnable parameters or auxiliary losses; projections learn representations useful when selected.
  • Pyramid: Symmetric average-pooling constructs an L-level pyramid of coherent Q, K, V triples, with each coarser entry summarizing p^ℓ consecutive tokens.Unlike asymmetric designs, Lighthouse pools queries, keys, and values together; pyramid construction costs Θ(N) time and memory.
  • Scoring and selection: Parameter-free query and key scores select the highest-ranked entries jointly across pyramid levels using a fused chunked-bitonic top-k kernel.Level-0 scores use per-head ℓ2 norms, while coarser levels inherit maxima from finer-level scores; the coarsest level is always retained.
  • Gathered-sequence attention: Selected entries form a contiguous length-S subsequence, where standard masked softmax attention uses a causal S×S mask derived from pyramid coordinates.The gather is topologically sorted, so the attention computation contains no sparse indexing; at N = 10^6, one configuration gives S ≈ 6.5 × 10^4 ≪ N.
  • Scatter-back reconstruction: Scatter-back redistributes each selected output to its represented base positions using causality-preserving shifts, disjoint within-level ranges, and bounded fan-in across levels.The final sequence is fully dense, and each position receives contributions from at most L levels.

4 Design Choices

Lighthouse’s design choices reduce training cost while keeping selection separate from attention and making dense recovery exact. The method uses symmetric query compression, inexpensive parameter-free scoring, stock attention, and gradient-free selection.

  • Symmetric compression: Symmetric Q/K/V pooling changes the dense kernel call from O(NSd) to O(S^2d) during training while keeping pooled queries and keys in the same representation space.This contrasts with methods that leave queries dense or pool only keys and values.
  • Parameter-free scoring: Parameter-free per-head ℓ2 norms replace learned scoring heads, providing a cheaper but weaker scorer than attention- or QK-interaction-based alternatives.The paper describes positive results with this weaker scorer as a lower bound on what richer scorers might extract.
  • Decoupled attention: Top-k selection is decoupled from attention by feeding a contiguous dense subsequence to stock SDPA or FlashAttention.The same attention kernel runs during training and inference, and disabling selection exactly recovers the dense baseline.
  • Gradient-free selection: The top-k operation is not differentiated, so gradients pass only through gathered Q, K, and V rather than through a learned selector.The design uses neither a straight-through estimator, Gumbel softmax, nor an auxiliary scorer loss.

5 Complexity Analysis and Kernel Design

Lighthouse confines custom GPU work to selection and scatter-back while reducing the attention subsequence to a polylogarithmic-size function of context length at fixed k. Consequently, bounded-k per-layer compute is linear in N up to a log k factor.

  • Kernel design: Only top-k selection and scatter-back use custom CUDA or Triton kernels; the remaining stages are PyTorch primitives that torch.compile can fuse.The custom selection kernel uses chunked bitonic merging, while scatter-back is the other bespoke stage.
  • Complexity analysis: S = N/p^(L−1) + (L−1)pk is the sole super-linear per-layer term, arising from dense attention on the gathered subsequence.The other stages contribute linear scoring and Θ(N log k) selection passes.
  • Complexity analysis: Choosing L = log_p(N/k) yields S = Θ(k log_p(N/k)) and attention cost Θ(k^2 log^2 N · d), polylogarithmic in N for fixed k.This balances the two terms in S.
  • Complexity analysis: For bounded k, total per-layer compute is linear in N up to a log k factor.The result combines the polylogarithmic gathered-attention cost with linear scoring and Θ(N log k) selection work.
  • Kernel design: Chunked bitonic top-k keeps an in-register top-m buffer per chunk and dispatches independent CTAs, avoiding textbook bitonic shared-memory blow-up at k = 4096.The kernel also produces stratified selection intended to resist span collapse.

6 Experiments

Experiments evaluate recoverability, scaling, ablations, throughput, and end-to-end training cost. Lighthouse recovers dense-attention performance while reducing attention and total training time at long contexts.

  • Recoverability: Dense-SDPA resumption recovered Lighthouse-trained models within ≈1–1.5k steps, reaching final losses of 0.6980–0.7102 versus 0.7237 for dense-from-scratch.All three resume schedules matched or beat the dense baseline at the same 16,000-step, ≈50B-token budget.
  • Scaling: At 512K context, Lighthouse was 21× faster than SDPA forward and 17.3× faster forward-plus-backward.The measured setting used a single B200 with L=3, p=4, and sparsity ≈1:64.
  • Ablations: Every Lighthouse ablation matched or beat the dense-SDPA-from-scratch final-loss baseline of 0.7237.The comparison used post-resume training loss at step 16,000 across independently varied design axes.
  • Ablations: The projection-norm scorer was roughly 9% cheaper in B200-hours than dilated softmax, with 179.6–180.9 versus 197.2–199.7 at L=3, p=4.Its loss was within ≈0.01 of dilated softmax in either direction, with no uniform winner.
  • Throughput: Stage-1 Lighthouse throughput reached 84–126k tok/s/GPU versus ≈46k for dense SDPA, providing a roughly 2× per-step advantage.The projection-norm scorer at L=3, p=4, k=1536 reached 126k tok/s/GPU.
  • Training cost: End-to-end runtime was 22.5–27.0 hours versus 37.9 hours for dense SDPA, a 1.40×–1.69× wall-clock speedup at matched or lower final loss.The budget was 16,000 steps and 50.3B tokens; savings came entirely from stage 1.

7 Conclusion

Lighthouse is a hierarchical selection method for long-context pretraining that surrounds stock FlashAttention with symmetric sequence compression and decompression. After brief dense-SDPA resumption, it matches or beats dense training while achieving 1.4–1.7× speedups, but its deployment scope remains limited.

  • Conclusion: Lighthouse pools Q, K, and V symmetrically across a multi-resolution pyramid and runs stock FlashAttention on the selected dense sub-sequence.Selection is placed outside the attention kernel and uses no learnable parameters, auxiliary losses, or straight-through estimators.
  • Conclusion: Brief dense-SDPA resumption matched or beat dense-from-scratch training loss and long-context retrieval at matched tokens.The method achieved 1.4–1.7× end-to-end speedups against cuDNN SDPA at ≥100K context and scaled to 1M tokens on multi-node Blackwell.
  • Limitations: Autoregressive decoding violates the assumption that all queries co-occur in one forward pass, so Lighthouse relies on dense-SDPA resumption for an inference-ready model.Downstream evaluations are run after resumption rather than directly on the hierarchical forward.
  • Limitations: The gathered attention remains Θ(S2d), so regimes requiring k to scale with N are not characterized.The method is subquadratic in N at fixed k but not strictly linear.

A Ablations

The ablation summary compares Lighthouse-stage training, dense-resume training, throughput, compute cost, and final loss across configurations on 530M-parameter Llama-3.

  • Ablation table: Table 2 reports LH Steps, SDPA Steps, B200-Hrs, Lighthouse-stage Tok/s (k), and Final Loss at step 16,000.The dense-from-scratch row has no LH stage and reports dense values; the final block adds context parallelism.
  • Ablation table: Bold values mark the per-block best on each metric, except throughput in the context-parallel block because context length varies there.Throughput is aggregated over 8 ranks, while B200-Hrs is combined wall-clock multiplied by 8 GPUs.

B Complexity Derivation

Lighthouse reduces the attention sub-sequence to S = Θ(k log T), making dense attention polylogarithmic in T and total per-layer compute linear in T for bounded k.

  • Sub-sequence size: S = N/p^(L−1) + (L −1)pk defines the gathered sub-sequence size for a Lighthouse layer.The setup uses sequence length T, pooling factor p, L pyramid levels, top-k budget k, and head dimension d.
  • Total compute: For bounded k, total per-layer compute is linear in T because projection, pooling, scoring, and scatter stages remain Θ(T · d).The logarithmic factor applies to the sub-sequence size, not the total per-layer compute.
  • Choice of L: Setting L = log_p(T/k) balances the pyramid terms and yields S = Θ(k log T).The coarsest level contributes the first term, while finer levels contribute at most pk entries each.
  • Attention cost: Θ(S^2 · d) is the dense FlashAttention cost on the gathered sub-sequence, becoming Θ(k^2 log^2 T · d).This dependence is polylogarithmic in T for bounded k.
  • Comparison: Table 4 compares Lighthouse with dense softmax, log-linear attention, and linear-attention / SSM families.Lighthouse and linear/SSM families share the same asymptotic class, while dense softmax is quadratic.

C.1 Symmetric Q/K/V Pooling

Lighthouse pools queries, keys, and values symmetrically, separates selection from stock FlashAttention, and uses gradient-free selection to preserve a recoverable dense model.

  • Symmetric Q/K/V Pooling: Symmetric Q/K/V pooling reduces the dense call from O(NSd) to O(S^2d) while preserving coherent representations across pyramid levels.Pooled queries route to pooled keys, enabling summary–summary interactions unavailable to asymmetric pyramids.
  • Scoring: Projection-norm scoring is parameter-free and avoids the extra compute and parameters of a learned scoring head.The dilated scorer remains an alternative, but projection norms are the throughput-sensitive default.
  • Kernel interface: Selection produces a contiguous sub-sequence that stock FlashAttention can process without a custom sparse-attention kernel.The same attention kernel runs during training and inference, and disabling selection recovers dense attention.
  • Gradient flow: The discrete top-K operation carries no gradient; gradients instead flow through gather, attention, scatter, and the projected Q, K, V values.This avoids straight-through estimators, auxiliary scorer losses, and scorer–attention optimization pathologies.
  • Chunked-Bitonic Selection: Chunked top-K is stratified rather than globally exact, replacing clustered global winners with lower-scoring entries from other chunks.The design guarantees broader regional coverage and empirically avoids selection collapse onto a narrow span.
  • Context-Parallel Execution: The contiguous output supports context-parallel ring attention and enables 1M-token pretraining across 32 Blackwell GPUs without changing the attention kernel.Lighthouse’s local preprocessing and contiguous tensors avoid custom sparse collectives.

E Design Ablations and Throughput (extended)

Ablations favor shallow pyramids and throughput-efficient scoring, while larger selection budgets do not consistently improve post-resume loss in the tested range.

  • Scorer variants: At k=1536, dilated scoring reaches 0.6881 while projection-norm scoring reaches 0.6946; at k=2048, the values are 0.6969 and 0.6921 respectively.Neither scorer is uniformly better, and their losses remain within ∼0.01 in both directions.
  • Scorer variants: Projection-norm scoring saves approximately 9% compute, costing 179.6–180.9 B200-h versus 197.2–199.7 for dilated scoring.It also uses no additional learnable parameters.
  • Pooling factor: At k=1536, pooling factors p=2, 4, and 8 yield final losses of 0.6825, 0.6881, and 0.6828 respectively.The study adopts p=2 as default; p=4 with projection-norm scoring is the wall-clock-favoured alternative.
  • Pyramid depth: With p=2 and k=1536, increasing L from 3 to 5 raises final loss from 0.6825 to 0.6991.L=3 is best across both tested selection budgets.
  • Selection budget: With L=3 and p=2, tested k values produce final losses of 0.6825, 0.6880, 0.6890, 0.6951, and 0.6831.A larger selection budget does not consistently lower post-resume loss within the tested range; reversal at larger budgets remains future work.

E.5 Throughput Decomposition

Lighthouse training reduces end-to-end runtime versus dense SDPA, while retrieval outcomes vary across selection budgets and scorers. The retrieval test compares four Lighthouse→SDPA configurations against a dense baseline at 98K context.

  • Throughput: 83.5–126.0k tok/s/GPU stage-1 throughput exceeds dense SDPA’s ∼46k across the 98K ablation grid.At k=2048, raising p from 2 to 4 increases throughput from 90.9 to 97.1k, while raising L from 3 to 4 increases it to 94.5k.
  • End-to-end runtime: 22.5–27.0h end-to-end runtime yields a 1.40× to 1.69× wall-clock speedup versus 37.9h for dense SDPA-from-scratch.The comparison uses the same 16,000-step, 50.3B-token budget and reports matched or lower final loss.
  • Retrieval: 0.76 is the best mean retrieval rate, achieved by k=2048 with the dilated scorer, versus 0.72 for the dense SDPA baseline.The test averages retrieval over 10 single-digit passkeys across context lengths and depths; random chance is 10%.
  • Retrieval: k=2048 outperforms k=1536 for both scorers, while switching from dilated to norm costs 0.04 at k=2048 and 0.08 at k=1536.The norm scorer’s retrieval penalty is larger than its effect on training loss, and the preferred default depends on whether the task is loss- or retrieval-driven.
Loading 2605.06554v1…