🤖 machine learning

ATMA: Length-Invariant Language Modeling via Polar Attention and Gated-Delta Compression Memory

The paper introduces ATMA, a hybrid architecture combining a novel three-channel Polar Attention mechanism with gated-delta compression memory to overcome softmax-based length constraints, achieving monotonic perplexity reduction and high-fidelity long-range retrieval up to 64K tokens.

Original authors: Habibullah Akbar

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Habibullah Akbar

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

The Big Problem: The "Long Story" Dilemma

Imagine you are trying to remember a story that is 64,000 words long. You have two main ways to handle this, but both have a fatal flaw:

The "Sliding Window" Approach: You only look at the last few pages of the book.
- Pros: You remember the immediate details perfectly.
- Cons: You are completely blind to anything that happened 10 pages ago. If the answer to your question is in the first chapter, you miss it entirely.
The "Full Context" Approach: You try to hold the entire 64,000-word story in your head at once.
- Pros: You can see the beginning and the end.
- Cons: Your brain gets overwhelmed. The "signal" (the important facts) gets drowned out by the "noise" (all the other words). It's like trying to hear a whisper in a stadium full of people shouting; eventually, you can't hear anything clearly, and your performance collapses.

Current AI models (Large Language Models) mostly use the second approach, but as the story gets longer, they start to "forget" or get confused because the math they use (called Softmax Attention) spreads their attention too thin.

The Solution: ATMA

The authors created a new model called ATMA. Think of ATMA as a super-smart librarian who uses a special three-part system to manage a massive library without getting overwhelmed.

Instead of just one way of looking at the text, ATMA breaks the job down into three distinct channels:

1. The "What" Channel (Polar Attention - Direction)

Imagine you are looking for a specific book. Instead of counting how many people are in the room, you just focus on what the book looks like (its color, shape, title).

How it works: This part of ATMA ignores how many times a word appears. It only cares about the direction or the type of information.
The Magic: Even if the library grows from 100 books to 100,000 books, this channel stays calm. It doesn't get confused by the crowd size. It just points to the right "feature" of the answer.

2. The "How Much" Channel (Polar Attention - Magnitude)

Now, imagine you need to know how strong the signal is. Did one person shout the answer, or did a whole choir?

How it works: This channel counts the "effective matches." But here's the trick: it has a volume limiter.
The Magic: If 1,000 people shout the same thing, a normal brain might get overwhelmed. ATMA's volume limiter caps the volume so it doesn't explode. It says, "Okay, that's a lot of matches, but I'll treat it as a 'very loud' signal, not a 'broken speaker' signal." This keeps the math stable even for huge stories.

3. The "Long-Term Memory" Channel (Gated-Delta Compression)

Even with the "What" and "How Much" channels, you still need a place to store the gist of the story so you don't have to re-read the whole thing every time.

How it works: This is a special "sticky note" system. It constantly updates a summary of the story as it reads.
The Magic: Most summary systems are "lossy"—they forget details. But ATMA's memory is "gated." It decides what to keep and what to overwrite carefully. It acts like a high-quality compression algorithm that keeps the main plot points (the "needle") safe, even if the story is 64,000 words long.

Why the Combination is Key

The paper found that you can't just use one of these tools; you need all three working together:

If you only use the "What" channel (Polar Attention), you lose the ability to find specific facts in a huge haystack.
If you only use the Memory, you get a good summary but can't find specific, pinpoint details (like a 5-digit code hidden in the text).
ATMA combines them: The "What" channel finds the exact needle, and the "Memory" channel keeps the context stable so the needle doesn't get lost.

The Results: The "Needle in a Haystack" Test

The authors tested this by hiding a specific "needle" (a secret code) inside a massive "haystack" (a long document).

Old Models: When the document got very long (32 times longer than what they were trained on), the old models failed completely. They found the needle less than 20% of the time.
ATMA: Even at 64,000 words (32x the training length), ATMA found the needle 90%+ of the time.
Bonus: Not only did it find the needle, but it also understood the rest of the story better (lower "perplexity"), meaning it was less confused by the long text than any other model.

The Technical "Secret Sauce"

To make this work on real computers, the authors had to build custom software "engines" (kernels).

They created a special way to do the math that saves memory, like a Flash Drive that only loads the pages you need right now, rather than trying to load the whole book into RAM at once.
They optimized the "memory update" so it doesn't slow down the computer, keeping the speed high even while doing complex calculations.

Summary

ATMA is a new way for AI to read long books. It solves the problem of getting overwhelmed by long text by splitting the job into:

Direction: What are we looking for? (Stays calm regardless of size).
Magnitude: How strong is the signal? (Keeps the volume from exploding).
Memory: A smart, compressed summary that holds the big picture.

By combining these, ATMA can read a 64,000-word document and still find a tiny detail hidden inside it, something previous models couldn't do.

Technical Summary: ATMA – Length-Invariant Language Modeling via Polar Attention and Gated-Delta Compression Memory

1. Problem Statement

Modern Large Language Models (LLMs) relying on standard softmax scaled-dot-product attention (SDPA) face critical limitations when generalizing to sequence lengths exceeding their training distribution. The paper identifies two primary failure modes:

Softmax Dilution and Activation Shift: As the key-value sequence grows, the probability mass of softmax spreads across a wider distribution. This "dilution" causes the maximum coordinate of the softmax vector to approach zero, flattening attention and inducing activation shifts that push the model out-of-distribution.
The Structural Tension: Existing solutions force a trade-off between local coherence and global recall.
- Sliding Window Attention (SWA) maintains low perplexity but is blind to long-range dependencies beyond the window.
- Full Softmax Attention preserves global recall but suffers from perplexity explosion and retrieval collapse at extreme lengths.
- Recurrent/Linear Memories (e.g., Titans, DeltaNet) offer efficient long-term compression but act as "lossy" systems, failing to retrieve specific, pinpoint facts (the "needle-in-a-haystack" problem) from massive contexts.

2. Methodology: The ATMA Architecture

ATMA (Attention with Memory and Attention) is a hybrid convolutional-attention architecture designed to resolve the window-vs-retrieval trade-off. It utilizes a 16-layer decoder-only structure with a 3:1 ratio of LFM2 gated convolutional layers to Polar Attention layers. The core innovation lies in the Polar Attention mechanism and its integration with a Gated-Delta Titans Memory.

2.1 Polar Attention: A Three-Channel Decomposition

Unlike standard attention which entangles magnitude and direction, ATMA factorizes the attention mixing step into an additive combination of three streams:
$\text{out} = \text{content} + \text{count} + \text{memory}$

Direction Channel ("What"):
- Projects the matched value-sum onto the unit sphere.
- This creates a count-blind, size-invariant vector that represents purely what features were attended to, regardless of the number of matches or sequence length.
- It utilizes a Length Temperature parameter to sharpen distributions as context grows and an Extreme-Value-Corrected Null Floor to calibrate logits against random noise, preventing background noise from overwhelming the signal.
Magnitude Channel ("How Much"):
- Calculates the participation ratio (inverse Simpson index) of the attention distribution to determine the effective number of matches ( $n_{eff}$ ).
- This count is gated by the confidence of signal detection and mapped through a bounded, saturating monotonic function (tanh of log). This ensures the magnitude input remains in-distribution at any context length, preventing the unbounded growth seen in raw log-counts.
Recurrent Memory Channel:
- An additive third stream driven by a Gated-Delta Titans Memory (inspired by Behrouz et al., 2025).
- It acts as a lossy "gist" capturing long-term perplexity trends, complementing the high-fidelity, bounded readout of Polar Attention.
- Stability Mechanism: The paper notes that standard RMS-normalization causes exponential divergence in delta-rule updates. ATMA resolves this by applying L2-normalization to keys and queries, bounding the eigenvalues of the update rule to ensure absolute stability.

2.2 Hardware-Software Codesign

To ensure efficiency, the authors implemented custom Triton kernels:

Fused Polar Attention Kernel: Computes participation ratios online with $O(\text{block})$ memory using streaming accumulators and squared correction factors, avoiding the $O(T^2)$ memory cost of materializing intermediate scores.
GQA-Grouped Paged Decode Kernel: Optimizes inference by reading directly from paged KV cache blocks, reducing HBM traffic by 4x.
In-Place Gated-Delta Step Kernel: Performs recurrent memory updates in-place, eliminating gather-scatter cycles and reducing memory traffic by 3x.

3. Experimental Results

The authors conducted a 100-run factorial ablation sweep training 370M-parameter models on 1 billion tokens of FineWeb-Edu (training length 2K), evaluating performance up to 64K tokens (32x training length).

3.1 Key Findings

Retrieval Accuracy: The combination of Polar Attention + Titans Memory is the only configuration that maintains induction needle-in-a-haystack retrieval accuracy above 90% across the entire 2K to 64K range.
- Softmax + Memory: Collapses to 16% accuracy at 64K due to activation drift corrupting memory readout.
- Polar Only: Collapses to 0% at 64K, proving that attention alone cannot generalize without the memory branch.
Perplexity: The ATMA model demonstrates monotonic perplexity reduction, improving from 2.70 nats at 2K to 1.96 nats at 64K. This outperforms softmax-based memory baselines, which show a perplexity increase (collapse) at extreme lengths.
Ablation Insights:
- Adding Sliding Windows or Distractor Losses (auxiliary objectives to calibrate the null floor) actually harmed performance when the memory branch was active. The window restricted training on long-context patterns, and the distractor loss over-sharpened the null floor, suppressing necessary weak signals from the memory.
- The simplest configuration (Full Polar + Memory, no window, no distractor) was strictly superior.

3.2 Efficiency

The recurrent memory branch incurs a 4.5 percentage-point drop in Model FLOPs Utilization (MFU) on NVIDIA L4 GPUs compared to Polar-only training, a manageable overhead for the significant gains in retrieval and perplexity.

4. Significance and Claims

The paper claims that ATMA successfully resolves the structural tension between local perplexity and long-range retrieval in language modeling. Its primary significance lies in demonstrating that:

Length Invariance is Achievable: By decoupling direction and magnitude and correcting for extreme-value noise, attention mechanisms can remain stable and effective at lengths far exceeding training data.
Hybridization is Necessary: Neither Polar Attention nor Recurrent Memory is sufficient in isolation. Their combination allows for simultaneous high-fidelity retrieval of specific facts and continuous improvement in document perplexity.
Practical Scalability: The architecture, supported by custom kernel engineering, scales to 64K tokens (32x training length) without the "collapse" phenomena observed in standard softmax or linear attention baselines.

The authors conclude that this approach offers a viable path toward models capable of ingesting and reasoning over massive context windows while maintaining coherent representations and precise recall.