Get Started
Home
Topics
Search
Library
LLM Training · Reasoning · Aug 27, 2026

Fast Weight Attention for Continual Learning

Source: research paper via Hugging Face Daily Papers
Linear-recurrence states (Mamba-2, DeltaNet) are silently doing online learning, but paired off-by-one: they train on same-step (key, value) when the state gets read one token later. Falcon shifts to the prefix key with NLMS-normalized steps, lifting 33-48 digit addition extrapolation from 75% to 87%.
TL;DR
Falcon reframes the recurrent state in linear-attention and State Space Model models as an online regressor whose training pair must be prefix-shifted: at step t, the state should learn to predict value v_t from the previous key feature ϕ(k_{t-1}), not the same-step ϕ(k_t) that Delta Network and Mamba-2 use.
Why It Matters
You’re building a long-context assistant on top of a linear-recurrence backbone (Mamba-2, RWKV, DeltaNet) because you can’t afford the quadratic attention bill. These architectures compress the whole prefix into a fixed-size matrix state that gets updated one token at a time. That update rule is silently doing online learning inside your forward pass, and nobody tells you what loss it’s minimizing or whether the training pair even respects causality correctly.
The dominant baseline here is DeltaNet and its descendants, which treat the state update as one gradient step on a squared-error loss that pairs the current key with the current value: (ϕ(k_t), v_t). Falcon’s argument is that this pair, while causal for next-token prediction, mistrains the internal memory.
How It Works
Start with the intuition. Under read-after-write ordering, the state S_t is written after seeing token t, then read to predict token t+1. So when v_t arrives, the honest supervised example for the state is: “given the prefix feature that was available before v_t showed up, predict v_t.” That prefix feature is ϕ(k_{t-1}), not ϕ(k_t). The same-step pair is still causal, but it’s optimizing a different internal objective, one that’s misaligned with what the state actually gets asked to do at read time.
Once you accept that alignment, the update rule falls out of choosing a loss:
•
Squared-error regression gives you the Falcon-1/2/3 family. The state is a linear predictor; you take one online gradient step on the instantaneous ridge loss. With a normalized step size, this is exactly NLMS (Normalized Least Mean Squares) adapted to fast weights.
•
Negative inner product gives you Falcon-1A/2A/3A, which recover the additive Hebbian writes used by Linear Attention and Mamba-2, but with the shifted pair and an energy-normalized write gain.
The numeric suffix picks the dynamics: 1 is a scalar step size shared across value channels, 2 is per-column step sizes, 3 is a sliding minibatch over the last B causal pairs. Ridge shrinkage λ_t acts as forgetting, and the paper works out chunk-parallel kernels (WY transform for Falcon-2, ParallelFlow for Falcon-3) so training stays sequence-parallel.
# Falcon-1 step (regression, scalar rate), read-after-write x_t = phi(k[t-1]) # prefix write feature; x_1 := 0 r_t = v[t] - S.T @ x_t # residual against pre-update state L_t = (x_t @ x_t) + lam[t] + eps # local smoothness eta_t = beta[t] / L_t # NLMS-normalized step size S = (1 - eta_t * lam[t]) * S + eta_t * outer(x_t, r_t) o_t = S.T @ phi(q[t]) # read AFTER write
Core Insight
The prevailing view treats the recurrent state update as an architectural choice: pick your gate, pick your decay, ship it. This paper reframes it as a supervised learning problem hiding in plain sight, and then points out the standard pairing is off by one token. The state should be trained on (prefix-feature, newly-revealed-target), not (same-step-feature, same-step-target); the one-step shift is what makes the internal objective match what the state gets read for. The cleanest evidence isn’t the language-modeling table (where Falcon is competitive, not dominant); it’s the variable-digit addition task, where the shifted, normalized update extrapolates dramatically better than same-step baselines.
What They Found
•
On variable-length addition, trained on 1-32 digit widths and tested on 33-48 digits, Falcon-3A.3 reaches 87.2 mean out-of-distribution accuracy and Falcon-1A.3 reaches 85.9, versus 75.2 for Mamba-2 and 65.8 for a RoPE Transformer. This is the load-bearing result: it isolates the state’s ability to store and carry information causally, which is exactly what the alignment argument predicts should improve.
•
On language modeling at 124M-130M parameters trained on 50B tokens of FineWeb-Edu, the story is muted. The regression variant Falcon-1.3 gets the best FineWeb-Edu validation perplexity at 17.10, edging Gated DeltaNet at 17.32. Downstream zero-shot and one-shot averages across eight tasks are within ~1 point across all recurrent models, including the baselines. The authors call this out honestly: not a uniform win, but the aligned updates preserve LM quality while delivering the arithmetic gains.
•
Ablation within the Falcon-1A family: RMSNorm on queries and keys beats ℓ2-normalization on perplexity, and context-conditioned step size η beats context-conditioned gain β on downstream averages.
What’s Useful
Reach for this if you maintain a linear-recurrence or SSM-style backbone and you care about tasks where the state has to reliably carry structured information across many tokens: arithmetic, symbol manipulation, structured decoding, long-horizon state tracking. The concrete swap is small. In your existing DeltaNet-style block, change the write feature from ϕ(k_t) to ϕ(k_{t-1}) with x_1 := 0 as the boundary, and switch the raw step size to the NLMS form β_t / (‖x_t‖² + λ_t + ε) with β_t ∈ (0,2). You keep chunk-parallel training; the paper’s Algorithms 1-4 give explicit WY/ParallelFlow kernels.
Code and configs are at GitHub. No new datasets are released; language-modeling runs use FineWeb-Edu and the addition task follows Kaiser and Sutskever’s setup. The per-column variants (Falcon-2, Falcon-2A) and the sliding regression rule (Falcon-3) are derived and implemented but not separately benchmarked in the main tables, so treat them as available primitives rather than validated recipes.
Takeaway
When your recurrent state is silently doing online learning, make sure it’s learning from causally-aligned examples. The one-token shift between “the feature the state was read at” and “the feature the state gets trained on” sounds like bookkeeping, but it’s the difference between a memory that extrapolates on structured tasks and one that doesn’t.
Caveats
•
The headline evidence is a synthetic arithmetic task. On real language modeling at 130M scale, Falcon is competitive with Gated DeltaNet and Mamba-2, not clearly better. Whether the alignment story pays off at 7B+ scale or on natural long-context benchmarks is untested here.
•
The sliding-window variants (Falcon-3, Falcon-3A) are not Markov in the matrix state alone: exact continuation across a segment boundary needs the last B-1 causal pairs as a tail buffer. If you truncate or reset that tail (e.g., serving with KV-style eviction), near-boundary updates change and the descent guarantees no longer apply cleanly.
•
The per-step descent lemma is pointwise and applies to the unclamped recurrence. In practice, the log-space implementation clamps the decay fraction to keep γ_t > 0, and under that clamp the implemented shrinkage corresponds to an effective ridge coefficient, not exactly the one you parameterized. In mixed precision at long context, this surrogate is what actually runs.
Topics
Don't miss new content
Log in to follow topics and personalize your feed.
By content type
Research Paper171 episodes
AI171 episodes