Wall Attention: Length Generalization With Diagonal Gates

Positional encodings have become a critical bottleneck for long-context generalization. RoPE, the dominant approach used by nearly every major frontier model, is data-independent → the same two positions always receive the same bias, regardless of which tokens actually appear there. This blocks length generalization, preventing models from accessing long context without training on increasingly long data.

Meanwhile, a parallel line of work in linear RNNs has quietly developed a more powerful idea. Diagonal forget gates (e.g. GLA, RWKV-7, and Kimi Linear) allow different feature channels to decay at different rates, giving models a granular, learned, content-dependent sense of what to remember and what to forget. This has become one of the key ingredients behind the latest generation of linear RNN architectures.

While studying how to bring diagonal gating into softmax attention, we found a fundamental obstruction. A finite-dimensional diagonal gate cannot directly act on the infinite-dimensional feature space of the exponential kernel → we need a lift.

Most interestingly, this obstruction has a clean resolution through the induced action framework, which lifts any input-space linear operator to the full feature space by applying before embedding rather than after. This turns out to unify FoX, PaTH, and Wall as special cases of a single construction.

Today, we’re releasing Wall Attention, which generalizes diagonal forget gates to softmax attention. Wall achieves state-of-the-art length extrapolation, outperforms RoPE and FoX in pretraining, generalizes from 4k to 160k+ context zero-shot, and is a drop-in replacement with open-source Triton kernels for training and decoding → where WallDecode is comparable to FA3 decode.

Code: github.com/tilde-research/wall-attention-release

The Problem: Softmax Attention Can’t Forget

Standard softmax attention is permutation-equivariant. To the pure attention operation, context isn’t a timeline; it’s a static bag of tokens. To fix this, we inject positional embeddings (PEs), such as RoPE or ALiBi, to introduce a recency bias.

Regardless of the order in which tokens are introduced, vanilla softmax attention has the same output.

Regardless of the order in which tokens are introduced, vanilla softmax attention has the same output.

But methods like RoPE are data-independent. The same two positions receive the exact same rotational penalty regardless of what text they represent. Natural language, however, has highly variable information density. The model should be able to look at a token and decide, “This is the end of a semantic thought, I should clear my working memory,” or “This is a critical entity, I need to remember it forever.”

Modern linear RNNs (such as Mamba, GLA, and KDA) address this with data-dependent diagonal forget gates. They maintain a recurrent state, and apply a learned, per-channel decay before writing new information. This is what gives them selective memory.

Architectures with increasingly sophisticated gating & update strategies.

Architectures with increasingly sophisticated gating & update strategies.

The question is: how do we bring diagonal gating to softmax attention? Prior work, such as FoX, tried scalar gates by absorbing decay into an ALiBi-like distance penalty. We seek to bring the diagonal gate over.

FoX formulation.

FoX formulation.

Meet Wall Attention - a full generalization of diagonal gates to softmax attention, a complete replacement for RoPE, and a new SOTA for zero-shot length extrapolation.

Wall formulation.

Wall formulation.

Diagonal Forget Gates on Infinite Dimensions

To understand why deriving Wall is nontrivial, look to the unrolled gated linear RNN:

St=AtSt−1+ϕ(kt)vt⊤S_t = A_t S_{t-1} + \phi(k_t) v_t^\top

Unrolling this yields unnormalized attention weights

wi,j=⟨ϕ(qi),(∏r=j+1iAr)ϕ(kj)⟩w_{i,j} = \langle \phi(q_i), (\prod_{r=j+1}^i A_r) \phi(k_j) \rangle

To recover standard softmax attention, the feature map φ is infinite-dimensional.

For a scalar gate, the cumulative product is just a number. It factors right out of the inner product, giving

exp⁡(q⊤k+log⁡ Fij)\exp(q^\top k + \log \ F_{ij})

This is the FoX formulation: a simple additive bias.

But we want a **diagonal gate **to control forgetting per-channel.

Fundamentally, you cannot directly parametrize an infinite-dimensional diagonal matrix. As a result, the scalar absorption trick completely breaks down.

The Induced Action Framework

Instead of trying to lift the diagonal matrix to the feature space directly, we let it act on the input space and induce the action through φ. We define the induced action such that gating the input before embedding equals gating the embedded state:

A~ ϕ(x)  :=  ϕ(A x)\tilde{A}, \phi(x) ;:=; \phi(A, x)

Crucially, this induced action is a group homomorphism. This means the cumulative induced transition equals the induced action of the cumulative product! When we pass this through the exponential kernel trick, the infinite-dimensional math collapses into something incredibly beautiful and simple:

wij=exp⁡ ⁣(qi⊤Aj→i kj)w_{ij} = \exp!\left(q_i^\top A_{j \to i}, k_j\right)

When specialized to diagonal gates,

Aj→i=diag(∏r=j+1igr)A_{j \to i} = \mathrm{diag}(\prod_{r=j+1}^i g_r)

Definition: Wall Attention

Wall attention allows the model to set dynamic per-channel forget rates within the attention mechanism, allowing for selective retention.

ot=∑jsoftmaxj ⁣(∑nFij,n qi,n kj,n)vj=∑jsoftmaxj ⁣(∑n(∏r=j+1igr,n)qi,n kj,n)vj\begin{aligned} o_t = \sum_j \mathrm{softmax}j !\left( \sum_n F{ij,n}, q_{i,n}, k_{j,n} \right) v_j \[4pt] = \sum_j \mathrm{softmax}j !\left( \sum_n \left( \prod_{r=j+1}^i g_{r,n} \right) q_{i,n}, k_{j,n} \right) v_j \end{aligned}

However, we can also rewrite Wall into an efficient factorized format. Define the log-space prefix sum of the gates as

Pt=∑u≤tlog⁡ guP_t = \sum_{u \leq t} \log \ g_u

We just rescale the queries and keys before standard attention:

q~i=exp⁡(Pi)⊙qi,k~j=exp⁡(−Pj)⊙kjot=Attn(q~, k~, v)\tilde{q}_i = \exp(P_i) \odot q_i, \qquad \tilde{k}_j = \exp(-P_j) \odot k_j \qquad o_t = \mathrm{Attn}(\tilde{q},, \tilde{k},, v)

That’s it. That is Wall Attention.

Visualization of sample per-channel selective retention in Wall.

Visualization of sample per-channel selective retention in Wall.

Fast Training & Inference for Wall

If you code the factorized Wall Attention naively - computing Q̃ and K̃ in HBM and passing them to FlashAttention - the model will explode.

Why? Because P_t is a monotonic cumulative sum. By sequence length 8192, the size of P_t exceeds 160. 2^(P_i) underflows and 2^(-P_j) overflows bf16 limits, causing catastrophic cancellation, even though the final reconstructed score is perfectly bounded.

To fix this, we wrote a custom Triton kernel.

1. Per-Tile Anchors: Instead of global rescaling, we introduce a local anchor R.

exp⁡2(Pi−Pj)=exp⁡2(Pi−R)⋅exp⁡2(R−Pj)\exp_2(P_i - P_j) = \exp_2(P_i - R) \cdot \exp_2(R - P_j)

This bounds the exponents to the maximum accumulated gate within a single sequence tile, mathematically guaranteeing stability.

2. The Fused Gate Gradient: In the backward pass, calculating the gradient for the gate prefix dP required an extra accumulator. This blew up our register pressure, forcing the Triton compiler to silently halve our block sizes, cratering SM occupancy.

The critical insight? We realized

dPj=−ln⁡2⋅Kj⊙dKjdP_j = -\ln{2} \cdot K_j \odot dK_j

Both K_j and dK_j are already resident in SRAM when the inner loop terminates. By fusing the gate gradient post-loop, we freed up the registers, restored block sizes to 128x64, and hit tensor core saturation.

3. Decoding Competitively With Flash Attention 3

For inference, storing P_t for all tokens would inflate the KV cache. Instead, we absorb the gate into the cached keys using chunk-anchored references. The inner loop requires zero extra gate arithmetic. Our WallDecode kernel is comparable to FA3 decode across sequence lengths.

Across sequence lengths, WallDecode achieves competitive throughput to FA3 decode.

Across sequence lengths, WallDecode achieves competitive throughput to FA3 decode.

Empirical Results

Small Scale Pretraining

Wall cleanly beats RoPE and FoX across the scales we tested. Wall sets a new SOTA on our pre-training tokens vs performance curve over our previous record Aurora 1.1B, matching strong open-source SLMs on key benchmarks.

Top: downstream evaluation across PE strategies. Wall (NoPE) achieves the strongest performance. Bottom: token efficiency vs. publicly available models trained on 30-500x more data and previous generation Aurora models.

Top: downstream evaluation across PE strategies. Wall (NoPE) achieves the strongest performance. Bottom: token efficiency vs. publicly available models trained on 30-500x more data and previous generation Aurora models.

Zero-Shot Length Extrapolation

Pretraining gains are great, but Wall was built for length. We took our 1B models - trained with a maximum context window of exactly 4,096 tokens - and evaluated them on extreme context extrapolation.

Wall (NoPE) stably generalizes to 160,000+ tokens zero-shot without performance degradation.

Length generalization benchmarks.

Length generalization benchmarks.

On the Needle-in-a-Haystack (NIAH) benchmark, every RoPE and FoX model collapsed at 8k tokens. Wall maintained strong retrieval at 16k and continued performing beyond. Wall also achieved a ~9% relative improvement over RoPE on LongBench v1.

The Discovery of Bimodal Channels

When we opened up the wall to see how it was routing information, we found something fascinating.

Because the Wall gates per-channel, we tracked the retention scores across 65k-token documents. At initialization, all channels start roughly identical (fully open). But as training progresses, Wall learns to split its head dimensions into two distinct populations:

• Static Memory Channels: Some channels lock their retention to exactly 1.0 (zero variance). They never forget. They serve as unconditional long-range memory, acting exactly like vanilla attention.

• Dynamic Forgetting Channels: Other channels become highly reactive. Their retention fluctuates wildly from step to step, depending solely on the text’s semantic content. At paragraph breaks or semantic shifts, they snap shut, clearing the channel’s working memory before opening again to absorb the next thought.

Each dot is one channel, plotted by its mean and standard deviation of retention. The colors correspond to different layers in the model.

Each dot is one channel, plotted by its mean and standard deviation of retention. The colors correspond to different layers in the model.

Wall learns a multi-timescale memory hierarchy end-to-end, offloading the structural prior of positional embeddings directly into the sequence’s latent dynamics.

Climbing the Wall

We started by identifying a mathematical pathology (how do you project a diagonal matrix onto an infinite-dimensional feature space?), built a rigorous framework to solve it (the induced action on the symmetric algebra), formulated a new attention mechanism, and engineered a stable, hardware-aware kernel to make it scale.

• Sufficiency. Wall outperforms every positional variant we tested. It does not need RoPE stacked on top of it. It proves that diagonal gating is a sufficient framework for temporal understanding in language models.

• **Hardware-aligned.**Wall retains the embarrassingly parallel structure of vanilla attention. It drops seamlessly into GQA and MLA architectures and has an efficient decoding formulation.

• **Induced Action Bridges RNN and Softmax Attention. **The induced action framework formally connects the mechanism powering modern linear RNNs with standard softmax attention in a way that, to our knowledge, hasn’t been formalized before.

The future of Wall is wide open. From WallMLA, to kernel improvements, to uptraining open source RoPE models into Wall models, we are incredibly excited to see what the open-source community does with Wall Attention.

Read the paper: https://blog.tilderesearch.com/blog/wall-attn

Get the kernels: github.com/tilde-research/wall-attention-release

Dhruv Pai, Timor Averbuch, Ashley Zhang, Ben Keigwin, and Alec Dewulf - Tilde Research