Falcon realigns fast-weight writes: learn v_t from the previous key, not the current one

Fast Weight Attention for Continual Learning

Yifan Zhang, Steve Ta, Jasper Zhang, Jichen Feng, Shuzhen Li, Yongxin Zhang, Yifeng Liu, Huizhuo Yuan, Mengdi Wang, Quanquan Gu, Andrew Chi-Chih Yao

cs.LG, cs.CL, stat.ML

2026-08-28

Falcon trains fast weights on (phi(k_{t-1}), v_t). At 130M it matches Gated DeltaNet on FineWeb-Edu and reaches 87.2 mean accuracy on longer addition.

What problem this solves

A Transformer KV cache grows with length. Mamba, linear attention, and DeltaNet compress history into a fixed recurrent state, so each transition is an online learning rule. The pair that rule writes is often the same-step association: current key with current value. Under read-after-write next-token prediction, that pair is the wrong local example. When the model reads the updated state to predict the next token, the write feature that was actually available is ϕ(k{t-1}), and the newly observed target is vt. Binding (ϕ(kt), vt) is still causal. It optimizes a different internal memory objective.

This is continual learning with a bounded fast memory: bind new evidence online without wiping what is already stored. If the alignment is off, plasticity, forgetting, and windowed rehearsal all take gradient steps on the wrong sample.

Method

Falcon treats the state S as a linear predictor from xt=ϕ(k{t-1}) to yt=vt, then takes one normalized gradient step on an instantaneous ridge loss. The step size is ηt = βt / (||xt||² + λt + ε), the NLMS choice from adaptive filtering, which is less brittle to feature scale. The residual rt = vt − S{t-1}^T xt is the error of predicting the current value from the previous write feature, not the usual DeltaNet reconstruction of vt from kt.

Six variants sit on two axes:

• indices 1/2/3: one scalar gain for all value channels, a per-column gain, or a sliding mini-batch over the last B causal pairs
• suffix A: drop the residual and write an inner-product update, i.e. Hebbian accumulation plus shrinkage

Falcon-1 is scalar NLMS. Falcon-2 promotes η to a vector. Falcon-3 averages residuals over a window and uses the largest eigenvalue of the window covariance as the smoothness scale. The A family writes η xt yt^T. Reads are ot = St^T ϕ(qt). Q/K default to RMSNorm, which the paper prefers to ℓ2 in mixed precision. Recurrent, masked-parallel, and SSD-style chunk-parallel forms are all given, with positive-decay renormalization so log-space unrolling does not go negative.

Results

Models sit at 124M to 130M parameters, trained for about 49.2B FineWeb-Edu tokens at length 1,024. Baselines are a LLaMA-style Transformer, RetNet, Mamba-2, DeltaNet, and Gated DeltaNet. The main tables report scalar regression Falcon-1.3 and several inner-product variants. Falcon-2, Falcon-2A, and Falcon-3 are defined and not separately scored.

Language modeling is competitive, not a sweep. Falcon-1.3 wins FineEdu at 17.10 versus Gated DeltaNet at 17.32, then loses Wiki at 33.00 versus 30.99. The best listed zero-shot average is Falcon-1A.2 at 49.30. Variable-digit addition isolates storage and carry: train on 1 to 32 digits, test 33 to 48. Falcon-3A.3 reaches a mean of 87.2; the Transformer sits at 65.8. The regression variant Falcon-1.3 falls to 68.8, back near the Transformer. The paper treats this as supporting evidence, not the headline result.

Why it matters

SSM and fast-weight work gets a construction drawing: temporal alignment, plasticity, forgetting, and bounded rehearsal can be tuned as separate knobs. Changing same-step binding to prefix-feature predicts current-target leaves a DeltaNet-shaped update with one index shifted. Chunk-parallel training remains, so the cost class does not change.

Before inventing a new write rule for long-context recurrences, check that the pair being stored was visible at prediction time. On language modeling this is a small move. On addition extrapolation, inner-product plus a sliding window is a real gap. Do not sell those two plots as the same story.

Limitations

The study stays at about 130M parameters and 50B tokens. There is no 1B-plus comparison. Falcon-2, Falcon-2A, and Falcon-3 exist on paper and not in the main tables. Addition accuracy is teacher-forced on the target suffix, not free generation. Gated DeltaNet still has a clear Wiki perplexity lead, so the framework does not show that realignment wins generic language modeling. Same-step pairing remains a legal causal rule; the paper does not lock every other choice and compare the two alignments at large scale. Positive-decay renormalization and the x1=0 boundary are easy to get wrong. The appendix is long for a reason.

Terms

• fast weight: a short-term memory matrix updated token by token in the forward pass.
• NLMS: least mean squares with a step size divided by input energy, less explosive under rescaling.
• read-after-write: write the current token, then read the updated state to predict the next one.
• next-latent alignment: train the fast memory on the pair (ϕ(k{t-1}), vt).
• SSD: Mamba-2's duality that writes selective SSMs and linear attention as the same chunk scan.

Source

• arXiv 2608.27763v1
• PDF
• Hugging Face
• GitHub · https://github.com/yifanzhang-pro/fast-weight-attention

What people are saying

• Falcon Framework: Adaptive Memory Updates for Efficient Continual Learning burny_tech
• Yao Class Team Proposes Fast Weight Attention for Continual Learning QuanquanGu

Related papers

• Yale: a coverage-SNR threshold decides SFT vs ICL, and GPU congestion can reverse it
• Penn's SSV Cuts Strong Verifier Calls 46% on Sudoku While Matching Oracle Accuracy
• Train grokking 10 million steps further and test accuracy falls back to 50%
• A 109M orbital transformer hits R² 0.9999, then recovers five fake gravity laws
• CoT monitors cannot be sound: Harvard on linguistic illegibility in LLM security
• Adam adapts per-parameter steps from first and second moments, α=0.001

All paper explainers