GitHub - jonsmirl/ssa: Subquadratic Sparse Attention (SSA): content-routed exact attention for long contexts — paper, theory, and reference implementations

SSA is a research implementation of content-routed sparse attention for long contexts. For each query it selects a small set of key blocks, adds a local causal window, and computes ordinary softmax attention over the selected keys. The implementation avoids materializing the dense n × n attention matrix.

This repository is a working research artifact, not a production inference stack. Its strongest result is a complete frozen Qwen2.5-0.5B prefill over 10,000,128 tokens on one RTX Pro 6000. Its quality evidence at that length is one semantic needle-in-a-haystack ranking, so the run establishes execution at scale rather than broad quality preservation.

The full technical account is available as the compiled paper, LaTeX source, and Markdown source. Detailed experiment records live in RESULTS.md.

Current status

The partial-coordinate experiment turns partial key reads into attention-logit intervals and feeds them to the existing mass/output certificate. On the cached Qwen head, high top-key recall alone still leaves too much mass; adaptive reads certify the requested bound. All 880 numerical trials pass the audit. This measures logical data reads, not runtime savings.

The GPU experiment tests actual latency on fresh Qwen geometry. Fused partial-score/bound and sparse-value kernels preserve the checked certificate, but the adaptive implementation is slower than dense attention in every measured group. It is not a deployable acceleration: reducing requested K/V bytes did not offset its routing/control work. The device-resident follow-up removes host acceptance decisions and charges conditional dense fallback. It still loses against matched dense graph replay through 32K. This partial-coordinate certified GPU acceleration research line is closed; its code and negative results are preserved. This is not a mathematical impossibility claim about all certified sparse attention.

The summary-recovery experiment separates information loss from fitting instability. On 16 cached Qwen-8K head queries and three seeds, persistent tanh ridge reduces mean head-output L2 error from sparse 2.3583 to 1.9266/1.8589 at 8K/16K scalar state caps, including Gram/cross statistics. However, the true-weight selected-only oracle with no tail estimate achieves 1.0661; stability alone has not recovered a useful tail. Median cell decoders achieve 0.9934/0.9822, with most improvement over sparse already explained by oracle normalization. Their worst-case cube radii remain 0.7366/0.7346, versus 0.7491 omitted mass with no summary. These are not measured accuracy floors or deployable certificates. See results and artifact.

Sparse reads plus bounded tail state

hybrid_tail_attention.py combines exact selected attention with an approximate tail computed from 16 cell counts/value sums per KV head. Selected contributions are subtracted from the summary before their exact weights are inserted. All incoming values update the state; it is not reconstructing arbitrary unseen values from sparse output alone. Raw next-token CE trains 336 per-head gains while all base Qwen weights remain frozen. No filler tokens are needed.

On the entire 298,938-token official WikiText-2 test split on an RTX Pro 6000, with the gains frozen after training on the official train split and development selection on validation:

These use batched causal tree routing and token-weighted loss, including final partial windows. Training uses only 512-token windows. The exact read budget stays at two past 64-key blocks plus the causal current block. The summary holds 49,920 scalars across all 24 layers, plus centers; inference prefix snapshots now advance in bounded chunks. The KV archive still grows with context. The learned estimate is not a deterministic mass/output certificate, and this reference remains slower than dense attention.

The cap was selected on development validation before the second full-corpus run; no gains were refitted. It improves sparse-only perplexity at all four lengths but gives up some original-tail benefit at 512. These are repeated regression tests, not a fresh holdout after the first run was inspected.

Nine fixed NIAH probes through 128K give dense 7/9, sparse 4/9, capped tail 2/9 strict wins plus one tie; the original tail gives 1/9. At 32K the capped model loses all three probes, versus sparse 2/3. The model uses native, unscaled RoPE (configured for 32K), so the 128K failures include a positional extrapolation confound; retrieval degradation already occurs at 32K. The correction is not a retrieval-safe long-context solution, despite its CE improvement. See the capped full-scale results, original run, and Kaggle reproduction. Peak allocation remains 5.81 GB at 128K. Earlier eight-window results remain recorded, but do not override the full-corpus and retrieval evidence.

Reproduce training or evaluate the portable learned gains without retraining:

python -m ssa.qwen_tail_demo --context 512 --steps 100 --test-examples 8 \
  --official-splits --eval-contexts 1024,4096 --out runs/qwen_tail_final
python -m ssa.qwen_tail_demo --context 512 --test-examples 8 --official-splits \
  --eval-contexts 1024,4096 --router tree \
  --load-gains runs/qwen_tail_final/results.json --out runs/qwen_tail_tree

The demo requires locally cached Qwen2.5-0.5B weights/tokenizer and WikiText-2, and runs offline on CUDA. results.json includes all 336 gain values, per-window losses, source hashes, and the base commit. Detailed results include the negative untrained and MLP controls. The new tail architecture has not been evaluated at 10M.

An experimental correction in tail_correction.py distinguishes assignment prototypes from the actual means of unopened cells. Actual means give a Jensen lower mass bound, not the upper bound needed by the deterministic certified reader. A separate mixture-share limit bounds how far the estimated tail can move the selected output; it does not guarantee improved accuracy. The formal follow-up landed in Substrate commit 206193290: exact reverse-KL projection, unique cell-uniform optimum, refinement, selected-key streaming identities, conditional output transport, and the movement/alignment criteria. The self-contained public specification is included here; no private repository is needed. These results do not imply lower model CE, and the capped prototype variant does not inherit the actual-mean estimator's variational optimum. The route-aware follow-up landed at 68968fc7a. Its public specification connects exact route-switch costs, full comparison traces, nonlinear multi-layer propagation, and a strict prediction-margin test under one hypothesis set. This preserves a supplied reference winner, not its correctness. SSA now implements a numerical path checker and a complete small causal transformer with analytic uniform RMSNorm, attention, and feed-forward bounds. Across 28 seeded correction trials, 18 pass the algebraic margin test, including 14 nonzero corrections; both prediction-changing proposals are rejected. There are zero observed interval violations. These are float64 checks, not an IEEE proof; the tiny model uses bound-friendly weights and normalization, and charges dense Jacobians and routing scans. It does not instantiate useful Qwen path bounds or establish subquadratic certification.

A separate Qwen endpoint policy evaluates both the sparse reference and capped-tail candidate. It accepts a candidate row only when both have the same strict winner, falling back on disagreement, ties, or a nonfinite candidate. No target labels enter this decision. On two development-validation windows per context on the local RTX 4080:

This full-vocabulary preservation is by construction, not evidence that the reference is correct. The measured CE improvement is empirical, and the policy cannot repair a wrong reference argmax. Both full model forwards and both logit projections are charged; this is paired teacher-forced prefill, not dual-cache generation, preserved sampling behavior, a Qwen Taylor certificate, or a speedup. When both endpoints are already computed, their direct comparison avoids the extra path-bound work. The run artifact records source hashes, costs, and separate four-candidate retrieval probes; this small validation run does not replace the full-corpus results above.

OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 python -m ssa.routed_certificate_experiment
python -m ssa.routed_acceptance_demo

The Qwen command uses the existing saved gains and offline model/data caches. It fits the local RTX 4080; Kaggle is not required for this experiment.

The storage comparison tests now check joint batch fitting, sequential delta updates, and cell sums on the cached 8K Qwen head. At two representation-state caps and three seeds, none of the tested linear/random-tanh fitting variants improves sparse attention's mean output error, even when supplied oracle attention weights; cell summaries do. This is a negative result for these particular updates, not a no-go theorem for fixed-state memory. Frozen unseen-value prediction and causal online reconstruction are reported separately. The results also contain an exact-arithmetic witness ruling out perfect linear key-to-value reconstruction on the stored fixture. No additional full-model training or Kaggle run is justified by this comparison.

Complete 10M transformer result

ssa/streaming_qwen.py executes a memory-bounded Qwen prefill using layerwise K/V lifetime, token-chunked projections and MLPs, native grouped-query attention, and a bounded hierarchical router. ssa/kaggle_10m_runner.py supplies equivalence and quality gates and is packaged for the ARC3-attached Kaggle RTX Pro 6000 by kaggle_10m/.

The recorded run in runs/kaggle_10m_v5/ssa_10m_result.json reports:

Routing uses pre-RoPE query/key content geometry to propose blocks and post-RoPE query/key vectors for the actual attention scores. Layer 1 forms a cross-head consensus route; a bounded 128-block reservoir carries high-vote evidence through later layers. Each query/head attends to at most 198 routed and local blocks. This policy is an empirical retrieval mechanism, not a certificate of dense-equivalent output.

The static YaRN factor is approximately 306× the model's training range. That makes the successful probe useful execution evidence but prevents treating it as a general long-context model-quality result.

Algorithm and complexity

For block c, SSA uses the second-cumulant routing score

r_c(q) = <q, μ_c> + (β/2) qᵀΣ_cq

and then evaluates standard softmax attention on the selected keys. If κ keys are selected per query, the attention work is O(nκ). A flat block router gives an O(n√n) operating point when block size and budget scale as √n. A hierarchical or IVF router with fixed block size and fixed κ approaches linear work on benign geometry. These are algorithmic or measured typical-case statements; worst-case exact routing can still require a full scan.

A different subquadratic construction factors positions as (block, slot) and composes a within-block pass with a same-slot cross-block pass. Every source can then reach every target in two hops, and a supplied cost model proportional to n(w + n/w) is minimized at w = √n. This is a useful comparison, not SSA's mechanism: complete structural reach does not imply dense-softmax equivalence. With either pass fixed, the other ranges over a strict subspace of all token-mixing maps in the nondegenerate dimensions; no rank bound or two-free-pass impossibility follows. Likewise, the scalar tree cost f log_f B is minimized continuously at e and over integer fanouts at 3, but that model omits GPU utilization and beam quality. The measured 10M tree therefore remains at fanout 16 pending a controlled benchmark.

The retrieval law

explains why selection can make recall insensitive to context length: it reduces the effective distractor count μ from n to κ. The selector must still include the target. Coherent spans provide block-level signal; isolated unit-norm needles can disappear inside summaries and are the principal adversarial case.

Mathematical status

The repository contains two distinct guarantee levels:

• ssa/certified_attention.py adaptively opens blocks until it certifies an omitted-mass and/or attention-output tolerance, with a full-scan fallback.
• ssa/hierarchical_certified_attention.py applies the same target through a tree and can evaluate O(log B) summaries on concentrated geometry. Its worst case is O(B).
• ssa/score_tail_certificate.py accepts CCC/IVF blocks only as routing seeds, derives separate admissible attention-logit caps, and sends a multi-level unopened-score profile through the same mass, KL, and value-output certificate. It is a float64 CPU reference.
• The fast budgeted routes in streaming_qwen.py, ivf_kernel.py, and the FlexAttention path are empirical; exact softmax over their selected set does not imply equality with dense attention.

The current ellipsoidal upper bound is

<q, μ_c> + ρ_c √(qᵀΣ_cq + ε||q||²),    S_c = Σ_c + εI.

The ε||q||² term is required because ρ_c is measured in the metric of S_c. Samuelson's inequality gives the tightest upper bound available from (μ_c, Σ_c, b) alone. On real transformer keys the exact covariance bound can cost at least a full key scan, which is why the scalable implementation uses budgeted lossy routing.

Related algebraic results are machine-checked in the separate Substrate Lean development. The audit includes the potential-store precursor 88a4c7012, the six subsequent SSA commits through 77cd1b8c6, the trace/covariance/outlier-peeled mass-tree results through 9c6b1ad35, their Inference-facing consumer fad55ff82, and the cell-mean projection/streaming/output results at 206193290, and the routed nonlinear comparison-path certificate at 68968fc7a. The broader mapping between Substrate results and this code is documented in docs/substrate_math_imports.md. The checked results now include recursive real-valued ball containment and its pairing cap, monotone expansion of the certified drop set as a tree bound tightens, causal selection by cutting every routed set at the query's original position, retention of every item with at least nine votes by a highest-count reservoir of capacity 128 when 14 selectors each contribute at most 70 items, and survival plus a uniform W + C cardinality bound for union with a fixed carrier. The radial pairing cap is attained when an extremal centre is aligned with the query and the corresponding boundary member is realizable; these are sufficient hypotheses, not a necessary characterization. The per-centre cap is never larger than the reach cap, and the plane witness exhibits a strict gap as large as the whole cap. No converse states that nonalignment forces strictness or that equality forces alignment.

The same audit now covers the complete abstract restricted-read certificate: total variation equals omitted share, selected-to-dense KL equals minus the log kept share, the reverse smoothed divergence is unbounded as the omitted support vanishes, and both value-output certificate arms follow from an exact residual identity. It also proves that admissible bounds plus strict skip/truncation conditions return the strict top set; without a margin, a deterministic index order is required to name one top set. SSA now uses larger parent indices to break exact routing-score ties. A separate adaptive-read theorem gives the b/n recall ceiling only for selectors that return positions they actually probed, and extends it to finite randomized mixtures. Finally, the reservoir retention, roundWidth + 128 cap, past bound, and causal cut are checked together as one composed plan. Retention remains conditional on nine pre-reservoir votes, and none of these facts proves attention-output quality or runtime.

The newer audit also puts the read and router on a common finite-potential object. For scores $a_i(q)$ and values $v_i$, the log-partition, softmax weights, and weighted read are three views of one construction; for scalar values the read is the directional derivative of the potential. Appending a fresh item of exponential mass $a$ to an old mass $Z$ changes the potential by $\log(1+a/Z)$ and moves the read by exactly $a/(Z+a)$ times the fresh value's displacement from the old read. This is an access law, not deletion or rewriting of append-only stored content, and changing scores adds a separate score-only term.

For a block whose logits lie in $[\ell,h]$, the normalized log-sum-exp change is now certified within $|\beta|^3(h-\ell)^3/6$ of its mean-plus-variance Taylor center. The bound controls the third cumulant over the whole tilt interval; the third cumulant measured only at zero is insufficient. Across sequential routed reads, local output errors obey $e_{t+1}\le\epsilon_t+L_t e_t$ when the exact next read is $L_t$-Lipschitz. Errors add only in the separately assumed nonexpansive case; without gain control, even a unit first-step error can be amplified beyond any proposed two-step bound.

For query groups, the pointwise maximum pairing is the least shared safe bound, adding a query can only reduce the certified drop set, and applying one common linear isometry to every query and region preserves the bound and retained set exactly. This does not extend to position-dependent transports. A complementary finite objective weights tree-node reads by a supplied query distribution: tighter safe bounds cannot increase this cost, and a least-cost member exists in any supplied finite family of certified partitions. That is a selection theorem, not a tree-construction, optimization, generalization, asymptotic, or latency result.

Finally, the read-budget boundary now includes finite preprocessing capacity. An index with $K$ states, per-state adaptive depth at most $b$, and at most $a$ unprobed reference-output positions can reach at most $K(b+a)$ single-spike placements. A finite randomized family therefore has some fixed placement with recall at most $K(b+a)/n$; grounded outputs remove the $a$ term. The identity index with $K=n$ reaches every spike at depth zero, so there is still no index-independent lower bound. State count is not equated with bits, index-build work, lookup work, or a general cell-probe cost.

The audit also sharpens two boundaries used by this project. A routed set reused wholesale at every query is generally non-causal; SSA is causal because the final attention relation masks by original token position, including partially visible blocks. And RoPE's relative-offset algebra does not itself prove unbounded length generalization: a phase band cannot acquire a nonzero winding inside less than one wavelength, while changing its winding under a reschedule requires crossing the anti-aliasing margin somewhere. The 306× YaRN run is therefore reported as measured execution and one retrieval outcome, not as a theorem of positional transfer.

The public paper does not require access to that separate repository: Appendix B of the Markdown source and the compiled paper reproduce the definitions, theorem statements, proofs, counterexamples, and scope for these routing invariants. The private Lean commit is corroborating audit provenance, not the only available mathematical argument.

These theorems verify the abstract exact-arithmetic components, not the complete Python/CUDA execution. The production CausalTree now conservatively inflates every recursive radius and query/node score cap and rounds the result toward +∞. On an RTX 4080, ssa/float_tree_verification.py compared fanouts 2, 4, and 16 over ordinary, scale-separated, large-offset, cancellation-heavy, and axis-aligned float32 geometries against float64 descendant oracles. The former formulas underestimated 8,701 of 90,105 radii and 367,480 of 1,081,260 caps; the guarded production formulas had zero observed underestimates. A 65,536-leaf, 64-dimensional fanout-16 build took 0.433 ms guarded versus 0.223 ms raw. These are adversarial numerical measurements, not an IEEE/CUDA proof, and they do not make fixed-beam routing an exact selector. Fixed-beam retrieval quality and the premise that the measured 10M target received nine pre-consensus base-route votes remain empirical.

Other measured evidence

• A flat block-sparse kernel is 20.6× faster than dense exact attention at n=262,144 on synthetic keys.
• The FAISS-GPU IVF kernel reaches 12M tokens in 139.5 ms single-head and uses 6.55 GB; its decode path remains about 0.6 ms per step from 1M to 12M at fixed κ.
• A fused Qwen2.5-0.5B swap reaches 128K under YaRN. NIAH is preserved in the measured settings while the two-hop task degrades at tight budgets.
• Routability co-training reduces lossless branch-and-bound key reads from 26.5% to 4.2% on the controlled task.
• Staged adaptation reaches 32× the trained length at 0.979 SSA recall in a toy MQAR task.
• The 124M construction experiment ends within +1.2 perplexity of an equal-training dense control while attending to about 38% of keys. This constant fraction is still asymptotically quadratic.
• The Certified Causal Cascade's per-query routing certificate has zero observed violations in its clustered and random tests; it certifies top-κ under the routing metric, not omitted attention mass.
• Trace, full covariance, and deterministic outlier peeling make Bennett mass trees dramatically tighter on a concentrated synthetic control (3,884 to 556 keys at a 10% mass certificate). On the measured Qwen-8K layer-18 head, even four peeled vectors plus covariance leaves 23.39 log units of median cap slack and opens every visible block. Force-keeping the peeled vectors' containing blocks likewise ends in a full read; it only reduces bound evaluations. The exact certificate remains a diagnostic, not a production router.
• A 16-level score-tail profile improves the Qwen-8K fixed-10%-block certificate margin from 34.54 (one residual threshold) to 31.81, with zero mass/output oracle violations, but certifies none of the queries and still opens all 6,247 visible keys at $\eta=0.10$ and $0.01$. The oracle contrast is sharp: the median largest 3.38% of visible keys holds 90% of true mass, while exact top keys plus one residual maximum need 19.82%. Sparse mass exists; the current block partition and caps do not expose it cheaply.
• Cross-layer route sharing from a middle donor layer reduces measured routing overhead from about 59% to 6% while preserving the single-needle probe. Sharing from layer 0 does not.

These results use different fixtures to isolate different claims. Only the complete 10M run combines a real pretrained model, every layer and head, long context, subquadratic routing, sparse attention execution, and an observed quality outcome in one experiment.

Known limitations

• Cheap summary routing is not lossless for arbitrary keys.
• Complete graph reach in a factorized attention pattern is not dense-attention functional equivalence.
• A routed subset is causal only after it is cut or masked at each query's original position.
• Multi-hop, multi-needle, and low-margin retrieval are harder than single high-margin NIAH.
• The 10M run has one semantic ranking and no dense 10M reference.
• Qwen2.5-0.5B is frozen and extended far beyond its trained positional range.
• The 12M IVF speed result is single-head and uses synthetic geometry.
• The real-model 128K and 10M results do not establish production throughput or frontier-model behavior.
• The certified attention implementations are CPU references and may open every block.

Quickstart

pip install -r requirements.txt
python -m ssa.ssa_demo
python -m ssa.niah_analysis
python -m ssa.staged_extension
python -m ssa.float_tree_verification
pytest ssa/tests

The core demos run on CPU. CUDA is required for the fused kernels and materially shortens the training demos. Scripts using pretrained models also require transformers, datasets, and either a local model cache or network access. The Kaggle 10M package is designed for internet-off execution with attached model and wheel datasets; see kaggle_10m/README.md.

Repository map

Building the paper

cd paper
latexmk -pdf subquadratic_attention.tex

License

Apache License 2.0 — see LICENSE.