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

Source-linked AI summary

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

arXiv:2608.27763v1Aug 27, 2026cs.LGcs.CLstat.ML

TL;DR

Long-context recurrent models must learn fast-memory updates online while respecting which information was available at prediction time. The paper makes this alignment explicit, derives normalized regression and inner-product updates, and evaluates them across language modeling and arithmetic extrapolation. Representative variants remain competitive in language modeling, while the best inner-product variants improve arithmetic length extrapolation.

Problem
Long-context models need online fast-memory updates without confusing the causal same-step association with the prefix-aligned example available when a target was predicted.
Method
The paper derives normalized first-order regression and inner-product fast-weight updates under prefix alignment, including scalar, per-column, and sliding-window variants with recurrent and parallel forms.
Results
Representative variants remain competitive in language modeling, and the best inner-product variants improve arithmetic length extrapolation.
Takeaways & Limitations
The framework separates temporal alignment, plasticity, forgetting, and bounded rehearsal in recurrent sequence models.

Abstract

from arXiv · show

Recurrent fast-weight memories and selective state-space models compress an expanding context into a fixed-size recurrent state, making the state transition an online learning rule. We study this rule under read-after-write autoregressive semantics. For the prefix-prediction objective considered here, the local fast-memory example revealed at step $t$ is the prefix-aligned pair $(\mathbf{x}_t,\mathbf{y}_t)=(φ(\mathbf{k}_{t-1}),\mathbf{v}_t)$. The common same-step association $(φ(\mathbf{k}_t),\mathbf{v}_t)$ remains causal, but optimizes a different internal objective. We derive normalized first-order updates for squared-error regression and negative inner-product objectives. The regression family comprises Falcon-1 (a scalar NLMS update), Falcon-2 (its per-column extension), and Falcon-3 (a sliding-window mini-batch update); Falcon-1A/Falcon-2A/Falcon-3A are the corresponding inner-product variants. We provide recurrent, masked-parallel, and chunk-parallel forms, together with numerically stable positive-decay renormalization. Representative variants remain competitive in language modeling and improve length extrapolation on variable-digit addition. This framework separates temporal alignment, plasticity, forgetting, and bounded rehearsal in recurrent sequence models.

1 Introduction

Long-context sequence modeling faces both quadratic attention costs and continual-learning demands: new evidence must be bound online without catastrophic interference. Recurrent alternatives compress fast memory into fixed-size states, but their local objectives and temporal alignment require explicit treatment.

Standard self-attention scales as O(N^2) in sequence length, making attention matrices and KV-cache traffic bottlenecks for long contexts.
Continual long-context modeling requires online binding of new evidence without catastrophic interference.
Transformers externalize fast memory in growing KV caches, whereas SSMs and fast-weight models compress it into fixed-size recurrent states.
In recurrent alternatives, the state-update rule acts as a local learning rule whose temporal alignment determines whether training information was available at prediction time.
Linear attention, fast-weight methods, RWKV, and Mamba demonstrate competitive recurrent alternatives with O(N) training and O(1) per-step inference.

C. Falcon-3: sliding-window mini-batch

Falcon makes the fast-memory example explicit under read-after-write semantics and derives normalized update families, including sliding-window mini-batch regression and inner-product variants with parallel implementations.

Falcon’s prefix-prediction pair is (ϕ(k_t−1), v_t), while the same-step pair remains causal but optimizes a different internal objective.
The recurrent state acts as a fast linear predictor from the prefix feature x_t := ϕ(k_t−1) to the newly revealed target v_t.
Falcon-1, Falcon-2, and Falcon-3 are normalized first-order regression writes with scalar or per-column plasticity and shrinkage controls.
Falcon-1A, Falcon-2A, and Falcon-3A are inner-product counterparts that retain the same plasticity and forgetting semantics while normalizing write magnitude.
The resulting recurrences receive chunk-parallel implementations and are evaluated in language modeling and arithmetic extrapolation.

2 Background

The background connects linear attention and selective SSMs through compressed recurrent states, then distinguishes normalized reads from denominator-free updates and motivates Falcon’s online objective.

Selective SSMs make transition coefficients input-dependent, while Mamba-2 connects simplified SSM transitions to causal linear attention and chunk-parallel training.
Linear attention replaces softmax with a kernel feature map, exploiting matrix associativity to reduce sequence-processing complexity.
The recurrent formulation compresses context into a matrix state S_t and a normalizer z_t.
Under read-after-write semantics, next-latent alignment writes v_t under the previous feature ϕ(k_t−1), shifting the standard write stream by one.
Denominator-free linear attention controls state magnitude and memory timescale through explicit decay or gain control, while its numerator update is gradient descent on an inner-product objective.
Delta Networks instead use an error-driven update derived from instantaneous squared error between the state’s reconstruction and the target.

A. Recurrent Form (causal scan)

Falcon’s recurrent view maintains a sequence of fixed-size matrix states, providing constant-state inference; Figure 2 relates this scan to masked and chunk-parallel forms.

The recurrent scan represents successive fast-memory states S_0, S_1, S_2, through S_t while processing key-value pairs.
The recurrent implementation offers O(N) time and O(1) state at inference.
Figure 2 presents recurrent, masked-attention, and chunk-wise parallel views of denominator-free linear attention.
The chunk-wise form splits the sequence into chunks, computes intra-chunk masked attention in parallel, and propagates a fixed-size state across chunks.

3 Autoregressive Next-Latent Prediction

The paper recasts recurrent fast-weight updates as online learning under read-after-write causality, pairing each newly revealed target with its prefix write feature. It derives normalized regression and inner-product updates whose gains, shrinkage, and normalization define distinct fast-memory behaviors.

Causal Alignment: Under read-after-write semantics, the causal fast-memory pair is (ϕ(k_t−1), v_t), whereas same-step pairing remains causal but optimizes a different local objective.The state predicts v_t from the prefix feature before incorporating the newly revealed target.
Online Objective: The recurrent state is modeled as an online linear predictor, and its instantaneous ridge objective is optimized with Online Gradient Descent.The fast-memory objective is internal to the forward pass rather than an additional supervised loss.
Normalization and Stability: Normalized step sizes match regression updates to local smoothness, while stabilization improves numerical robustness and recovers classical NLMS in the zero-ridge, zero-stabilizer case.The normalized update yields per-step descent in the instantaneous regularized objective, but not necessarily cumulative online-loss descent.
Connections: The framework connects fast-weight rules to Delta Networks and Linear Attention by identifying the local objectives that induce their update forms.The shifted regression assignment recovers the Delta Network functional form with a critical index shift.
Inner-Product Variants: Inner-product variants remove the residual term and produce additive writes, with normalization interpreted as write-magnitude control rather than curvature normalization.Falcon-1A, Falcon-2A, and Falcon-3A preserve the corresponding scalar, per-column, and sliding-window organization.

C. Chunk-wise Parallel Form (per-chunk WY/Gram)

The chunk-wise formulation expresses recurrent Falcon updates through within-chunk WY/Gram computations while preserving causal history and fixed-state sequential execution. Boundary handling is exact only when the required sliding-window overlap is retained.

C. Chunk-wise Parallel Form (per-chunk WY/Gram): The chunk-parallel form separates cross-chunk history from intra-chunk parallel computation while representing the same recurrent kernel.The recurrent and chunk-wise views are equivalent for the rank-one kernel described in Figure 4.
C.1 Scaled Fast Weight Recurrences: Positive-decay renormalization reduces positive-ridge recurrences to the displayed no-ridge WY form, while clamping preserves positive carry under normalized updates.The additive normalized case is recovered when η_t = 1 and λ_t = 0.
C. Chunk-wise Parallel Form (per-chunk WY/Gram): Falcon-2 chunk-parallel training builds a shared key Gram matrix and channel-dependent triangular systems independently within each head.A merged residual solve replaces the older two-solve presentation without changing asymptotic complexity.
A. Recurrent Form (per-channel rates): The recurrent implementation maintains a fixed-size state and supports O(N) time with O(1) sequential state during causal scanning.The per-channel WY representation provides an equivalent view of the no-ridge rank-one kernel.
C.1 Scaled Fast Weight Recurrences: Normalization is invariant to uniform feature rescaling in the non-sliding case, while sliding variants use window-dependent smoothness or write-energy statistics.For inner-product writes, the statistics serve as practical magnitude controls rather than curvature requirements.
C.2 Signed-Feature Normalizers: Signed-feature normalized reads require sign-stable denominators because adding a positive constant alone does not prevent zero or negative read denominators.The paper therefore treats denominator-free inner-product reads as the default signed-feature interpretation.
C.3 Step Size and Gating Conventions: Falcon expresses forgetting through derived decay α_t = η_tλ_t and carry γ_t = 1 − α_t, tying shrinkage to the local objective and normalized step size.This differs from gated Delta rules whose scalar carry is introduced independently.
C.7 Sliding-Window Boundary State: For sliding-window Falcon-3 variants, exact continuation requires the last B−1 causal pairs or an equivalent rolling-window state; matrix state alone is insufficient.Resetting the tail changes the near-boundary update, and aggregated statistics cannot indefinitely remove the oldest contribution exactly.

B. Parallel Form (batched per-channel TriSolve)

Falcon-2 extends scalar fast-weight regression with per-channel plasticity and equivalent recurrent, masked-parallel, and chunk-parallel formulations. The batched TriSolve implementation preserves the recurrence while reducing chunk-local computation.

B. Parallel Form (batched per-channel TriSolve): The batched TriSolve formulation solves per-channel triangular systems for the chunked recurrence.Each channel uses its own system while the batch shares the core computation.
B. Parallel Form (batched per-channel TriSolve): Chunking regroups the sequence into contiguous blocks, replacing O(N^2) work with O(NC) time across M batched solves.The chunk size is C and M chunks cover the length-N sequence.
A. Recurrent Form: Falcon-2 maintains a fixed-size matrix state while assigning each value channel its own step size and plasticity trajectory.All channels share the shifted write-feature direction and residual construction.
B. Parallel Form (batched per-channel TriSolve): Falcon-2 uses shifted write features and residual-based regression updates rather than same-step feature-value associations.The implementation retains the read-after-write alignment used by the paper’s fast-memory objective.
B. Parallel Form (batched per-channel TriSolve): The optimized implementation merges residual systems when they share a triangular factor, removing one batched TriSolve per chunk without changing recurrence or asymptotic complexity.This is presented as the single-inversion form for the chunk-wise forward pass.
B. Parallel Form (batched per-channel TriSolve): The section’s setup includes per-channel nonnegative step sizes, nonnegative ridge coefficients, a positive decay floor, divisible chunk size, and an initial state.These conditions define the inputs and stability conventions for the batched forward pass.
B. Parallel Form (batched per-channel TriSolve): The chunk-parallel algorithm initializes the state, processes each chunk, computes local prefixes and triangular solves, then carries the decayed output state forward.Positive-decay factors are accumulated in log space before the next chunk begins.

C. Falcon-3A: sliding-window inner-product write

Falcon-3A applies inner-product fast-weight updates over a finite sliding history, extending scalar and per-column additive writes to bounded rehearsal. Its recurrent and parallel views preserve shifted causal alignment while supporting chunk-parallel training.

Parallel forms: Unrolling produces decay-weighted causal masks, while chunking divides the sequence into fixed-size blocks for parallel computation.Per-column variants use channel-specific masks, and Falcon-3A admits a fully vectorized masked-attention form.
Scalar and per-column variants: Falcon-1A uses one scalar gain across value channels, while Falcon-2A assigns each value column its own energy-normalized gain.Falcon-2A is explicitly the per-column inner-product analogue of Falcon-2.
Scalar and per-column variants: Inner-product normalization controls additive write magnitude, unlike regression normalization, which is matched to local curvature or smoothness.The objective is λt-smooth independently of write-feature energy.
Recurrent form: The recurrent update combines additive window statistics with shrinkage, with the boundary convention η1 = 0 preventing an initial write.The implementation maintains the final state and FIFO window for continuation.
Falcon-3A sliding-window update: The mini-batch mechanism averages a finite history window to capture local dependencies and reduce noise accumulation.Exact continuation additionally requires a fixed-width tail of the last B−1 causal pairs.
Falcon-3A sliding-window update: The windowed update averages rank-one gradients evaluated at the pre-update state, rather than updating each residual sequentially.This is the direct mini-batch analogue of the instantaneous regression update.

5 Experiments

The experiments evaluate representative Falcon variants on language modeling and variable-digit addition, finding competitive language-model quality and improved arithmetic length extrapolation.

Experimental setup: 124M–130M-parameter models are evaluated on FineWeb-Edu and downstream accuracy with matched training budgets and recurrent baselines.The study also uses variable-length multi-digit addition to test causal storage and extrapolation.
Language modeling: Falcon-1.3 achieves the strongest FineWeb-Edu validation perplexity at 17.10, ahead of Gated DeltaNet at 17.32 and Falcon-1A.3 at 17.40.Lower perplexity is better in Table 1.
Language modeling: Falcon-1A.2 has the best listed small-model zero-shot average at 49.30, while Falcon-1.3 has the best recurrent one-shot average at 49.54.These are downstream-task averages across the evaluated tasks.
Variable-length arithmetic addition: Falcon-3A.3 achieves 87.2 mean accuracy on 33–48-digit out-of-distribution teacher-forced addition, followed by Falcon-1A.3 at 85.9.Both outperform the reported RetNet/LightningAttn and Transformer baselines.
Variable-length arithmetic addition: The arithmetic experiment is treated as supporting evidence that shifted, normalized updates extrapolate well when recurrent storage and carry propagation dominate.The diagnostic trains on widths uniformly sampled from 1–32 digits and predicts reversed sums.

C.6 Implementation Notes and Related Update Rules

The implementation notes explain how ridge decay, normalization, scale statistics, and optional local convolution shape stable fast-weight updates and relate them to contemporary delta-rule models.

Decay and stability: Ridge regularization contributes a scalar carry γt := 1 −ηtλt, with distinct eigenvalues along and orthogonal to the current write-feature direction.The write-feature direction has eigenvalue 1 −ηt(∥xt∥2 + λt), while orthogonal directions have eigenvalue γt.
Decay and stability: Positive-decay renormalization rescales step sizes, states, and write targets together; rescaling ηt alone is insufficient for exact equivalence.Chunked implementations apply the equivalent local transformation using within-chunk log decays.
Implementation choices: A short causal convolution is optional and may be applied before forming queries, keys, and values.This choice follows prior observations about local mixing and associative recall.
Sliding updates: Falcon-3 uses a small-Gram smoothness statistic for sliding regression, while sliding inner-product writes use window energy; the current shrinkage multiplier is statistics-only.The step-size denominator still uses the live statistic.
Related update rules: The formulation is complementary to RWKV-7 and Kimi Linear Attention because it grounds writes in causal next-latent regression and stabilizes them with NLMS normalization.Those related models add learned gating and in-context learning-rate control to delta-style fast weights.

D.1 Mask Identities and Chunk-Local Evaluation

The appendix derives mask identities and chunk-local algorithms for evaluating decayed, sliding-window fast-weight recurrences while preserving causal structure and numerical stability.

Mask identities: The fast-weight output decomposes into a masked-attention term from local writes and a decayed-history term from the incoming state.The mask is step-dependent and can remain nonzero for all causal positions.
Mask identities: Each mask coefficient sums at most B decayed write coefficients, and coefficients need not be monotone in lag because tokens are re-injected across the moving window.This produces a moving-sum followed by an exponential-tail structure in stationary settings.
Numerical stability: Chunk-local log-decay resets the prefix at boundaries, uses within-chunk differences, and propagates boundary states through chunk-exit decay.This avoids underflow from materializing global decay products in reduced precision.
Chunk-parallel evaluation: Chunk-parallel evaluation has three phases: parallel local precomputation, sequential inter-chunk boundary propagation, and parallel output materialization.The mask combines a causal decay kernel with a B-banded moving-average window operator.
Rank-one evaluation: For rank-one no-ridge updates, block transitions admit a WY representation, while a single residual solve removes one triangular solve per chunk without changing the recurrence.The accumulated transition uses update vectors collected in U and a triangular mixing matrix T.
Algorithm assumptions: The implementation assumes shifted write features xt := ϕ(kt−1), nonnegative step sizes, a defined initial state, and chunk sizes dividing the sequence length.The boundary step is hard-set to η1 = 0 and γ1 = 1, so it carries no gradient to β1 or λ1.

F Falcon-2 Parallel Implementation

Falcon-2 gives each value channel its own adaptive step size while sharing key-defined update geometry, enabling batched chunk-parallel computation. A merged residual solve removes one triangular solve per chunk, while shared-dynamics Falcon-1 further reduces system-construction overhead.

Channel-wise dynamics: Falcon-2 assigns each value channel a distinct trajectory and step size, while keys determine shared update geometry across channels.This structure permits vectorized chunk-wise recurrence with per-dimension adaptivity.
Chunk parallelization: The chunk implementation forms shared Gram and score structures, then solves channel-dependent triangular systems in batch over value channels.The shared Gram matrix is computed once per chunk, while step-size modulation differentiates each channel’s system.
One-TriSolve reduction: The two-solve decomposition is algebraically replaced by one merged residual solve because both paths use the same triangular factor.This preserves the recurrence and asymptotic complexity while removing one batched TriSolve per chunk.
Decay handling: Positive-decay renormalization handles ridge or weight decay by rescaling each channel’s state, reducing the decayed recurrence to a no-ridge rank-one form.The implemented path uses a clamped positive carry when the raw decay is not guaranteed positive.
Shared-dynamics variant: Falcon-1 shares the adaptive learning rate and transition matrix across value channels, reducing distinct chunk systems from dv to 1.The remaining solve still has O(dvC^2) cost but becomes a single high-throughput multi-right-hand-side solve.

G.4 Complexity Analysis

The complexity analysis contrasts full Falcon-2’s channel-specific systems with Falcon-1’s shared system and explains how Falcon-3 uses associative chunk scans for low-rank recurrences. These implementations retain the stated recurrences while improving parallel execution structure.

Falcon-1 complexity: Falcon-1 removes one forward TriSolve per chunk and shares system construction across value channels, yielding the larger practical gain over Falcon-2.Its remaining triangular solve remains O(dvC^2) but is a single optimized multi-right-hand-side solve.
Falcon-3 scan: Falcon-3 represents the sliding-window update as a rank-B affine recurrence and trains it with ParallelFlow without dense transition matrices or token-by-token scanning.The chunk maps compose associatively, allowing either sequential composition or an associative scan.
Causal masking: The tensorInv solve uses a block-strict-causal mask so same-time rank components do not interact, while chunk outputs use inclusive causality for read-after-write behavior.The structured solve avoids explicitly materializing the full inverse tensor.
ParallelFlow structure: ParallelFlow computes local chunk propagators and input injections in parallel, then links chunks through a global associative scan.The formulation separates temporal dynamics from their computation.
Complexity caveat: The compressed O(L^2_cR + d) complexity expression is misleading because it hides dominant dimension-dependent terms.The implementation also involves matrix multiplications involving feature and value dimensions.

H.3 Mapping Falcon-3 to ParallelFlow

Falcon-3 maps its sliding regression update exactly to a fixed-rank tensorInv recurrence by padding each active window to width B. This enables chunk-parallel scans without approximating the discrete update.

Sliding regression recurrence: Falcon-3 uses a sliding regression window of size B, converting the rank-one update into a rank-B affine recurrence.The active window contains at most B shifted write features and associated value blocks.
Fixed-rank mapping: The tensorInv mapping zero-pads every active window to fixed rank R = B, with padded columns vanishing identically.The resulting drivers use padded write-feature and value blocks.
Causal semantics: The block-strict-causal solve evaluates residuals at the pre-update state, while output materialization uses inclusive block causality for within-chunk read-after-write.For the first boundary step, the update is a strict no-op.
Exact parallelization: Falcon-3 can use tensorInv without approximation at the discrete level because its padded update is exactly a rank-B low-rank affine recurrence.The method avoids dense feature-by-feature state propagation.
Practical complexity: For small experimental windows such as B = 4, the rank-B overhead is modest and remains far cheaper than full O(L^2) attention.The chunk-parallel method retains the sliding-window regressor update exactly.