Needle-in-a-haystack retrieval accuracy as evaluation context grows to 32× the training length (2K→64K tokens). Polar Attention with the Titans memory stays flat above 90%; every other variant collapses toward chance. Identical 370M-parameter model, identical 1B-token training budget — only the attention recipe differs.
Abstract
Softmax scaled-dot-product attention is numerically stable but not length-invariant: as the key set grows, softmax mass spreads thinner, so a model trained at length T behaves differently at 32·T — it never learned to scale its computation with length, and "how many things matched" is entangled with "what matched" in a single normalized output. Borrowing the size-invariance of convolution, Polar Attention derives two channels from one temperature-sharpened softmax with a learned, extreme-value-corrected null sink: a unit-vector direction (size-invariant, count-blind) and a bounded magnitude read from the participation ratio of effective matches. The output stays bounded with infinite dynamic range at any context length, where softmax both dilutes and explodes.
Attention alone forces a tradeoff — a sliding window wins perplexity but is retrieval-blind, full attention preserves recall but leaks perplexity. We resolve it by attaching a Titans-style compression memory (a per-head gated-delta fast-weight store, in the Memory-as-Gate configuration) as an additive third channel: out = content + count + memory. The memory supplies the diffuse long-context perplexity gain; full polar plus a distractor objective supplies exact pinpoint retrieval. The combined stack is implemented across three numerically cross-verified pipelines — reference, training, and a paged inference engine — with a FlashAttention-style Triton kernel and a fused gated-delta kernel. We validate the design with a 120-cell factorial ablation over the attention core, regularizer, distractor, memory, and window axes.
Method
Each attention layer writes out = content + count + memory. The first two come from one shared softmax; the third is a recurrent associative memory trained in-loop.
A unit vector over values: what was attended to. Projected to the sphere, it is invariant to how many keys matched and to total length. c = s / ‖s‖
The participation ratio of effective matches, gated by a null-sink confidence and squashed by tanh(β·log1p) into [0,1) — bounded, ordered, length-invariant. 1K vs 1M matches stay distinguishable but bounded.
A per-head linear memory updated by a gated delta rule (Titans / Gated DeltaNet). Self-stabilizing in norm across length; γ is a per-head horizon dial. A lossy gist that carries long-context perplexity.
The extreme-value-corrected null floor (+ slope·√log n) tracks the noise maximum so pure noise never reads as a match at long context; a length temperature (1 + softplus·log n) sharpens the softmax as context grows. A small distractor objective pushes random keys below the floor, widening the signal/noise margin — this is what buys exact retrieval far beyond the training length. The full derivation follows.
Intuition
Two animations for the two ideas: why polar stays bounded as context grows, and how the memory writes a token once and recalls it later.
As context grows toward 64×, softmax weight dilutes and its output norm blows up (red); polar locks the direction to a unit vector and keeps the magnitude bounded in [0,1) (green).
Each token: the state M decays, then takes a rank-1 write. A needle is written once; a later query reads it back — while the state norm stays flat in length (self-stabilizing).
Schematic, built to convey the mechanism — not a live model. Animations pause off-screen and respect prefers-reduced-motion.
Derivation
A single query \(i\) attends over keys \(j<n_i\) (causal, \(n_i=i+1\) visible keys). With QK-norm on \(q,k\), scores are \(\sigma_{ij}=(q_i\cdot k_j)/\sqrt{d_k}\). Two length-aware scalars per head do the work.
Sharpens the softmax as context grows (Scalable-Softmax style), so weight does not spread thinner with more keys:
The logit of an off-by-one "null sink" key. The \(\sqrt{\log n}\) growth is load-bearing — the maximum of \(n\) noisy scores grows like \(\sqrt{2\ln n}\), so any fixed floor is eventually overtaken by noise:
Append the null logit, softmax over real keys + sink, mix the values (plus a learned default \(v_{\text{null}}\)), project to the sphere:
Reuse the same weights. Renormalize over real keys, take the participation ratio (effective match count), gate by confidence, then bound with a saturating map:
The count is invariant in length when \(\text{len\_gain}\cdot\Delta \ge 1\) (\(\Delta\) = the QK margin between signal and noise): the temperature then suppresses each noise weight faster than \(n\) grows. End-to-end, the same key population at 1× and 64× yields identical \(c\) (cosine 1.0) and identical \(\mathrm{mag}\).
Random keys, pushed through the same \(k\) projection, must lose to the sink — this widens \(\Delta\) and is what buys retrieval far past train length:
Designed by falsification. An additive soft-count \(\sum\sigma(\sigma-\theta)\) leaks \(\mathcal{O}(N)\) (89× drift over a 1000× sweep); the participation ratio is flat. Raw \(\log(m)\) pushes the projection out of distribution — the bounded \(\tanh\!\cdot\operatorname{log1p}\) keeps it in range. A fixed null floor is overtaken by the noise maximum. The distractor alone collapses the count (floor \(\to\infty\)); the task loss is its counterweight.
Each head carries a matrix memory \(M\in\mathbb{R}^{d_v\times d_k}\) — a fast-weight key→value store — reusing the layer's \(q,k,v\) plus two data-dependent scalar gates per head.
The Titans neural memory (momentum \(\eta=0\)) reparametrizes exactly to the Gated DeltaNet recurrence. FLA's convention is canonical — decay-first, undecayed write, self-inclusive readout:
Pulled back into polar's discipline before it enters the residual (zero-init \(\mathrm{proj}\) → enabling the branch is a safe no-op at step 0):
The delta rule's per-step eigenvalue along the key is \(\gamma\,(1-\beta\lVert k\rVert^2)\). Polar's RMS-norm gives \(\lVert k\rVert^2=d_k\) → eigenvalue \(\approx-7\) → the state diverges (oracle blew up to \(\sim\!10^{57}\)). L2-normalized keys give \(\gamma(1-\beta)\in(0,1)\) — stable.
The delta memory self-stabilizes: state-norm is flat in \(N\) (256→16384) even at \(\gamma=1\), because the \((I-\beta k k^{\top})\) key-replacement converges toward the least-squares solution \(M\approx V K^{+}\) (capacity \(\sim d_k\)). So \(\gamma\) is learned per-head: some heads hold \(\gamma\approx1\) for distant recall, others \(\gamma<1\) for recency.
Division of labor. The linear memory is a lossy gist (capacity \(\sim d_k\)) — it carries diffuse long-context perplexity but cannot exact-recall a planted needle. Full polar + the distractor does the pinpoint retrieval. The recurrence has a closed chunkwise-parallel form so it trains in-loop at ~6–9% MFU overhead (fused FLA / Triton kernel).
Full write-ups with verification (float64 gradcheck, bit-exact train↔reference↔inference parity): POLAR_ATTENTION.md · TITANS_MEMORY.md.
Systems
The architecture lives in three implementations that are kept numerically identical: a pure-PyTorch reference oracle, an optimized training path (FP16/FP8 matmuls, Muon), and a paged inference engine. Every layer is cross-checked, so a new variant is safe to ship the moment parity holds.
A fused Triton kernel streams the polar reduction with query- and key-blocking — fully O(block) memory on tensor cores. The count channel needs one extra streamed accumulator Q2 = Σexp(·)² for the participation ratio, rescaled by α² on every online-softmax max update. The backward keeps a cheap per-query preamble in PyTorch and runs only the two O(T²) matmul loops as Triton.
Paged KV cache with prefix caching; centralized per-sequence state tables hold the conv states and the Titans memory (per-head M, fp32, FLA [K,V] layout). Prefill runs FLA's fused chunk_gated_delta_rule; decode runs a GQA-grouped paged kernel plus an in-place gated-delta step — one state read+write per token, CUDA-graph capturable.
Every reduction is established against a token-by-token oracle: float64 gradcheck on the streaming polar autograd and the gated-delta recurrence, then train == reference bit-exact in the default config. The same per-layer and end-to-end checks route through the Triton kernels on GPU. Nothing merges until parity holds.
verify.py 30/30 CPU + 30/30 CUDA
test_polar_kernel.py 104/104
test_integration.py 21/21
verify_titans / mag / window fp64 gradcheck
train ↔ reference bit-exact
The torch.compile saga. FLA's gated-delta kernel is a custom autograd function that graph-breaks at every memory layer under torch.compile — a ~2× time / ~2× peak-RAM regression, since the memory's FLOPs are only ~5% of the model (overhead, not compute). Wrapping it as two opaque custom ops (FLA_CUSTOM_OP=1) with register_fake shapes restores fusion and recomputes activations in the backward — bringing the real cost down to ~6–9% MFU.
Results
370M parameters (16 layers, 12 gated-conv + 4 attention), ~1B tokens of FineWeb-Edu, sequence length 2K. Every cell evaluates at full context. The winning recipe is polar + memory (no window, no distractor, baseline regularizer).
| Context | Polar+Titans | Softmax (NoPE)+Titans | RoPE+Titans | Polar only | Softmax (NoPE) only | RoPE only |
|---|---|---|---|---|---|---|
| 2K (1×) | 91% | 98% | 74% | 96% | 98% | 43% |
| 8K (4×) | 93% | 94% | 29% | 15% | 28% | 4% |
| 16K (8×) | 98% | 85% | 9% | 8% | 8% | 0% |
| 32K (16×) | 96% | 48% | 0% | 3% | 3% | 0% |
| 64K (32×) | 93% | 16% | 0% | 0% | 1% | 0% |
| Context | Polar+Titans | Softmax (NoPE)+Titans | RoPE+Titans | Polar only | Softmax (NoPE) only | RoPE only |
|---|---|---|---|---|---|---|
| 2K (1×) | 2.70 | 2.66 | 2.81 | 2.83 | 2.76 | 2.85 |
| 8K (4×) | 2.44 | 2.39 | 2.74 | 3.32 | 2.67 | 2.78 |
| 16K (8×) | 2.28 | 2.29 | 2.79 | 3.53 | 2.97 | 2.85 |
| 32K (16×) | 2.14 | 2.29 | 2.82 | 3.60 | 3.38 | 3.07 |
| 64K (32×) | 1.96 | 2.34 | 2.86 | 3.60 | 3.71 | 3.36 |
Polar + Titans is the only recipe whose perplexity improves monotonically as context grows (2.70→1.96 across a 32× sweep) while retrieval stays flat. Softmax holds early then collapses past ~16×; RoPE collapses even faster. Polar without the memory loses long-context perplexity. Convergence and quality rank Polar+Titans > Softmax+Titans > RoPE+Titans > Polar-only at ~6–9% MFU overhead.
The Ablation
Five regularizers × distractor on/off × memory on/off × window on/off × three attention cores (polar, softmax/nope, rope) — with a second batch adding wall attention. Every cell: clean & junk perplexity and needle accuracy at six lengths, plus MFU and wall-clock — in one interactive dashboard.
Cite
@misc{kreasof2026polar,
title = {Polar Attention: Length-Invariant Sequence Mixing with Compression Memory},
author = {Kreasof AI Research},
year = {2026},
note = {Research preview},
url = {https://github.com/kreasof-ai/atma}
}
A full technical report is in preparation. Documentation: Polar Attention · Titans Memory · Atma overview.