ATMA: Long-Context Language Modeling via Polar Attention and Gated-Delta Compression Memory

1. Introduction

Long-context language modeling is often compressed into a single retrieval curve. That curve matters, but it is not the whole problem. A useful model must preserve relevant information as context grows without giving up document likelihood, ordinary short-context quality, reasoning performance, or practical inference cost.

These objectives compete. Full attention preserves direct access to tokens, but its compute and cache grow with sequence length and its probability mass can disperse over more keys. Recurrent models keep fixed-size state, but that state must compress the past and can lose detail. We therefore treat long-context modeling as a Pareto problem: an architecture is an operating point on a multi-objective frontier, not a winner by one number.

The problem with more keys

Standard softmax attention normalizes over every available key. As the sequence grows, irrelevant scores create a larger noise field. Positional schemes such as RoPE change how order enters the score, but they do not by themselves separate the direction of the retrieved value from the amount of evidence that participated in the match.

That distinction motivates Polar Attention. We want one channel to represent what matched, without its norm growing simply because more keys were present, and another bounded channel to express how much matched.

Our approach

ATMA is a 16-layer decoder with 12 gated convolutional layers and four global attention layers. The global layers use Polar Attention and also maintain gated-delta recurrent memory. The attention path can still read the full context; the recurrent path supplies a compressed history whose state size does not grow with the sequence.

We test this in two stages. First, a complete factorial sweep chooses the recipe across attention type, regularization, distractor alignment, memory, and local training windows. Second, matched NoPE, RoPE, and Polar models are trained for 9.816B tokens and evaluated from their 2K training length through 256K.

2. Architecture

The matched attention models contain 378.16–378.22M parameters and share hidden size, heads, tokenizer, data order, optimizer, schedule, and token budget. Only the attention core changes. This keeps the main comparison narrow enough to interpret.

Polar Attention

Polar Attention decomposes its output into a normalized direction and a bounded magnitude. The direction aggregates values and removes radial scale. The magnitude is derived from an effective participation ratio, then passed through an extreme-value-aware null sink and a bounded map. In plain language: the model records what the matching values point toward separately from how convincing the match appears to be.

This is a channel decomposition, not an independence claim. Changing the match set can still rotate the direction. The design instead prevents the raw number of participating keys from directly controlling output norm.

Gated-delta compression memory

Each global layer also updates a recurrent key-value state. A learned retention gate controls what remains; a delta update writes new information; a zero-initialized projection adds the retrieved memory back to the attention output. Unit-normalized keys keep the tested recurrence stable, while the gate bounds how much the state can change at each step.

3. Experimental design

The experiment is deliberately divided into recipe selection and matched comparison. Mixing these two stages would overstate what the ablation proves.

The evaluation contains 24,000 paired retrieval trials over synthetic and FinePDFs haystacks, eight context lengths, three needle depths, and 50 trials per cell. We additionally measure BABILong after adaptation, fixed-target bits per byte, eight short-context tasks, and single-sequence serving on one L40S.

4. Results

Memory consistently helps Polar in the factorial

At 64K, adding memory improves Polar in all 20 matched combinations of regularizer, distractor setting, and window setting. The mean gain is 47.8 percentage points. The corresponding effect on NoPE is small and heterogeneous, while RoPE remains at zero.

A local training window reduces Polar with memory by 25.5 points on average. Distractor alignment also has a smaller, nonuniform negative effect. We therefore promote full-context Polar with memory, without a local window or distractor alignment.

Browse every Stage I cell and validation trajectory in the interactive dashboard →

Polar retains more target signal at 256K

After matched training, NoPE and Polar both approach perfect teacher-forced target-token accuracy at 2K. Their extrapolation is different. At 256K, Polar retains 34.4% averaged across tasks, haystack suites, and depths; NoPE retains 0.6% and RoPE 0.0%.

Where retrieval fails

The aggregate masks the most important boundary in the result. At 256K, Polar retains 68.4% target-token accuracy on synthetic haystacks, but only 0.3% on FinePDFs. Exact five-token accuracy is 18.0% on synthetic contexts and zero on FinePDFs.

5. The multi-objective frontier

Polar improves retrieval and long-document likelihood among the matched attention models, but it pays elsewhere. Its mean accuracy over eight 2K controls is 43.29%, compared with 45.15% for NoPE. At 128K, Polar attention decodes at 16.3 ms/token on one L40S; Atma-Raven-Titans records 2.27 ms/token with fixed recurrent state.

Raven Native reaches 46% BABILong at 256K and keeps likelihood nearly flat. Because Raven changes both model family and optimizer, it is not an attention ablation. It represents another practical point on the frontier: stronger adapted reasoning and cheaper decode, with lower direct retrieval.

There is no single “long-context score.” Retrieval, likelihood, adapted reasoning, short-context quality, memory growth, and latency answer different questions.

6. What a checkpoint audit revealed

A post-hoc audit found that three NoPE checkpoints could have nearly identical 2K validation curves while differing by 6.70 nats at 256K. The divergence appeared around blocks 6–7 in sampled secant-gain probes.

The finding emerged during an infrastructure transition, not a randomized hardware experiment. Re-evaluation excluded an evaluation-version mismatch, and a microbatch-4 retrain excluded microbatch size alone in the tested runs. It did not distinguish initialization, data order, device-dependent kernels, or their interactions.

7. Limitations

• The study uses one model scale and one 2K pretraining regime.
• Stage I demonstrates robustness across nuisance-factor settings, not across random seeds.
• Polar has zero exact FinePDFs retrieval at 64K and 256K.
• Teacher-forced retrieval is not autoregressive generation.
• Raven is a separately optimized family comparison, not an optimizer-matched ablation.
• Serving measurements are single-sequence measurements on one L40S and should be treated as descriptive.

Conclusion

ATMA establishes a reproducible operating point for bounded full-context attention plus recurrent memory. Polar materially extends target-token signal, preserves exact synthetic retrieval, and stabilizes long-document likelihood. The result does not solve natural-text retrieval at extreme length, and it does not dominate recurrent alternatives on reasoning or serving.

That mixed result is the point. Long-context systems should be compared as trade-offs, with metrics kept distinct and failure boundaries made visible.