Princeton University / Research talk

Linear Attention,
State Spaces, and RLT

Memory updates and computation across tokens

Yifan ZhangPrinceton ECE / October 2026

Based on A Note on Linear Attention and State Space Models
and Recurrent Looped Transformer

Arrow keys to navigate / O overview / N speaker notes

01 / 14

01 / A common language

The state-space view

\[h_t=\mathcal F_\theta(h_{t-1};u_t),\qquad y_t=\mathcal G_\theta(h_t;u_t)\]

Object	Meaning	Question to ask
State \(h_t\)	Information carried into the next update	What must persist?
Input \(u_t\)	New data and conditioning features	What is available before the update?
Transition \(\mathcal F_\theta\)	The rule that transforms the state	Affine or nonlinear in the old state?

The recurrence index can represent token position or network depth.

Note: Yifan Zhang and Sun-Yuan Kung
RLT: Yifan Zhang, Jichen Feng, Shihan Qin, Sun-Yuan Kung, and Mengdi Wang

Note / State-space representations / RLT / Sequence and state

02 / 14

02 / Linear attention

Linear attention as recurrent memory

Key features \(a_t=\phi(k_t)\), query features \(b_t=\phi(q_t)\). Read after writing the current token.

\[S_t=S_{t-1}+a_tv_t^\top,\qquad z_t=z_{t-1}+a_t,\qquad y_t=\frac{S_t^\top b_t}{z_t^\top b_t+\varepsilon}\]

State size \(O(md_v)\) stays fixed as the sequence grows. The vector \(z_t\) supplies the normalizer.

Note / Linear attention and original figure / Katharopoulos et al., 2020

03 / 14

03 / Exact dual views

The same computation in attention form

\[S_t=\sum_{j\le t}a_jv_j^\top\quad\Longrightarrow\quad S_t^\top b_t=\sum_{j\le t}(b_t^\top a_j)v_j\]

\[Y=(M\odot\Phi_Q\Phi_K^\top)V,\qquad M_{tj}=\mathbf 1\{j\le t\}\]

Same outputs, different evaluation: a running state or a causal \(T\times T\) score matrix.

Unnormalized outputs shown. Normalized attention also divides each row by its kernel sum.

Note / Original recurrent and parallel views / Katharopoulos et al., 2020

04 / 14

04 / Decay and selectivity

Selective SSMs and state-space duality

The SSM core

\[h_t=\bar A_t h_{t-1}+\bar B_t x_t\]

\[y_t=C_t^\top h_t\]

Mamba selects \(\Delta_t,B_t,C_t\) from layer inputs.
\(\bar A_t=\exp(\Delta_t A)\).

Scalar-transition SSD, with zero initial state

\[\begin{aligned} S_t&=\rho_t S_{t-1}+B_t x_t^\top,\qquad y_t=S_t^\top C_t\\[3pt] Y&=(L\odot CB^\top)X,\qquad L_{tj}=\mathbf 1\{j\le t\}\,\textstyle\prod_{r=j+1}^{t}\rho_r \end{aligned}\]

\(C\): queries, \(B\): keys, \(X\): values. Setting \(\rho_t=1\) recovers causal linear attention.

Note / Selective SSMs and SSD / Gu and Dao / Mamba, Fig. 1 / Dao and Gu / SSD

05 / 14

The figure is reproduced from Mamba Figure 1, https://arxiv.org/html/2312.00752v2/selection.svg. It illustrates selection and the expanded recurrent state. The equations on the right give the discrete SSM core. Mamba selects coefficients from the current layer input, which allows them to be computed before the recurrent scan of that layer. Once those inputs are fixed, the update is affine in the previous recurrent state. This is more specific than a general model whose coefficients depend on that same previous state.

For the scalar-transition SSD case, S is an m by d-v matrix, B and C are m-dimensional vectors, and x is d-v dimensional. The scalar rho is the decay at a position. Unrolling the recurrence weights each key–query interaction by the product of intervening decays. The empty product for j=t equals one.

The duality is the core connection used by Mamba-2. It concerns this structured mixer and does not equate every Transformer block with every SSM architecture.

Note / Selective SSMs and SSD / Gu and Dao / Mamba, Fig. 1 / Dao and Gu / SSD

05 / Corrective memory

DeltaNet writes the retrieval error

\[\begin{aligned} S_t&=S_{t-1}+\eta_ta_t\bigl(v_t-S_{t-1}^\top a_t\bigr)^\top\\[3pt] &=\bigl(I-\eta_ta_ta_t^\top\bigr)S_{t-1}+\eta_ta_tv_t^\top \end{aligned}\]

Write the retrieval error. The update is affine in the old state when the inputs are fixed.

Green arrows: writes. Dashed red arrows: key-dependent edits to the old memory.

Note / Delta Networks and original figure / Yang et al., 2024

06 / 14

The squared loss asks the memory matrix to retrieve v from key a. Taking one gradient step yields an error-correcting write. Expanding the error exposes an affine state transition with a rank-one edit.

The scalar example uses zero initial state and epsilon zero for normalized attention. After the first value, each memory holds or reads 1. The second additive write produces a sum of 4, and its normalized read divides by two. DeltaNet instead adds the error 3 minus 1, giving 3.

The edit matrix is an orthogonal projector only when eta equals the inverse squared key norm. The general update is a rank-one edit. The note also derives a triangular solve for the residual writes, followed by a masked read. That structure supports the chunkwise parallel algorithms studied by Yang and colleagues.

Note / Delta Networks and original figure / Yang et al., 2024

06 / The algebra of parallel evaluation

Why affine updates admit a scan

Fix a layer’s input sequence. Represent one transition by \((A_t,c_t)\):

\[h_t=A_t h_{t-1}+c_t\]

\[\underbrace{(A_2,c_2)\circ(A_1,c_1)}_{\text{two consecutive updates}}=(A_2A_1,\ A_2c_1+c_2)\]

Closed under composition

A segment of updates is still one affine map. Prefix composition is associative.

Structure controls cost

Additive, diagonal, scalar, and low-rank forms support efficient implementations.

A general nonlinear recurrence has no comparable fixed-size affine summary.

Note / Structured transitions / Mamba / Parallel scan / DeltaNet / Parallelization

07 / 14

Substitute the first update into the second to derive the pair-composition rule. Its associativity enables prefix scans. This is an algebraic statement about the affine core after its inputs and coefficients are fixed.

Associativity alone does not guarantee speed: dense matrix products can be expensive, and products of rank-one edits need not remain rank one. SSMs exploit structured transitions, while DeltaNet uses specialized factorizations and triangular solves.

Function composition is associative even for nonlinear functions. The issue for a general nonlinear cell is that composing many cells does not leave a compact affine pair with cheap composition. RLT therefore cannot directly use this affine-scan algorithm.

Note / Structured transitions / Mamba / Parallel scan / DeltaNet / Parallelization

07 / Recurrent Looped Transformer

RLT: a Transformer decoder as the recurrent cell

Causal encoder

Known tokens produce \(e_t\) and prefix memory \(M_{\le t}\).

Gated feedback

Merge \(e_t\) with the previous final decoder output \(s_{t-1}\).

Full decoder update

SWA, cross-attention, and an FFN transform the merged input.

SWA = sliding-window attention.
Every output predicts the next token.

RLT / Encoder, gated merge, and decoder blocks

08 / 14

Read the figure from bottom to top. The causal encoder produces a representation at each position and projected key–value memory. For known tokens, each encoder layer can process positions in parallel with a causal mask.

The decoder receives the current encoder representation and the previous token's final decoder output through the gated merge. Each decoder layer has local sliding-window self-attention, cross-attention to the causal prefix of encoder memory, and a feed-forward network. The orange feedback path returns the final decoder output to the next token's merge.

The paper authors are Yifan Zhang, Jichen Feng, Shihan Qin, Sun-Yuan Kung, and Mengdi Wang. The architecture figure is reused from the supplied general-exam slides, whose source is the RLT manuscript.

RLT / Encoder, gated merge, and decoder blocks

08 / A nonlinear state-space system

RLT’s explicit state-transition function

Complete decoder state: \(H=(s,C^D)\), including the feedback vector and local KV cache.

\[\boxed{\mathcal F_\Theta(H;e,M,t):=D_\phi\bigl(\operatorname{Merge}(e,s);M,C^D,t\bigr)}\]

\[\begin{aligned} r&=\operatorname{Norm}_s(s),\qquad g=\sigma\bigl(W_g[e;r]+b_g\bigr),\\ \operatorname{Merge}(e,s)&=e+\alpha\,g\odot W_s r \end{aligned}\]

\[H_t=\mathcal F_\Theta(H_{t-1};e_t,M_{\le t},t),\qquad H_0=(s_\star,\varnothing)\]

\[\mathcal G_\Theta(H)=\operatorname{softmax}\bigl(W_o\operatorname{Norm}_o(s)\bigr),\qquad p_\Theta(\cdot\mid x_{1:t})=\mathcal G_\Theta(H_t)\]

One transition updates both \(s_t\) and \(C_t^D\), conditioned on encoder features and memory.

RLT / Method, with the current manuscript’s explicit state-space notation

09 / 14

The function definition specifies the state transformation rather than merely calling RLT an SSM. Given the old state H, extract s for the merge and C-D for decoder self-attention. The full decoder returns the new feedback vector and the updated cache.

The conditioning inputs e, M, and the absolute token position t remain explicit. The output function uses the feedback vector component of the new state to predict the next token. The initial feedback vector is learned, and the initial decoder cache is empty.

The norm, gate, attention, and feed-forward operations make this transition nonlinear in the previous decoder state, even after the encoder inputs are fixed. The same cell parameters are reused across tokens, while different decoder layers may have different parameters.

RLT / Method, with the current manuscript’s explicit state-space notation

09 / RNN interpretation

Feedback extends the computation across tokens

\[H_T=F_T\circ\cdots\circ F_1(H_0),\qquad F_t(H)=\mathcal F_\Theta(H;e_t,M_{\le t},t)\]

Per token

\(L_E+L_D\) Transformer blocks.
The merge connects successive updates.

Across \(T\) tokens

The feedback path traverses \(T L_D\) decoder blocks with parameters reused.

Known-token encoding is parallel. RLT-1 decoder updates follow token order.

RLT / Prompt-state composition and depth accounting

10 / 14

A running parity computation is a simple example of a repeated state update: a one toggles parity and a zero preserves it. The picture illustrates the task, not a decoded internal neuron or a proof of the learned representation.

RLT is an RNN whose cell is a gated merge followed by a full Transformer decoder. It gains a computation path through multiple decoder stacks as more tokens arrive. Each token still runs the stated number of blocks. The T times L-D count refers to this feedback path, not to a claim of lower total attention cost.

For L-D equal to zero, there are no decoder blocks but the gated merge still recurs. The experiments use several encoder–decoder allocations, so decoder depth and recurrence remain distinct design choices.

RLT / Prompt-state composition and depth accounting

10 / RLT evidence

Length generalization depends on the depth split

Eight-layer models. Test accuracy at length 256, mean ± sample SD over three seeds.

Task	RLT-1 allocation	GPT Layer 8	RLT-1
Parity Train: 3–40 bits	Layer 5 + 3 or Layer 7 + 1	50.07 ± 1.63%	100.00 ± 0.00%
Swaps-\(S_5\) Train: 32 operations	Layer 4 + 4	0.85 ± 0.30%	97.30 ± 2.76%

Parity favors a shallow decoder. Swap tracking benefits from more decoder layers.

2,000 training steps, seeds 42–44, 1,024 shared test examples per length.
Each run selects its checkpoint using minimum in-distribution validation loss.

Layer counts match. Parameters and compute differ: RLT-1 has 26.10–28.73M parameters, versus GPT’s 25.31M.

RLT / Algorithmic experiments and length extrapolation

11 / 14

These are selected length-generalization results from the manuscript's six-task study. Parity scores the final answer, while swaps-S5 scores the final permutation state. The swaps task uses the identity and the ten single transpositions of five elements. Its uniform final-state reference is 1/120, about 0.83 percent.

The best encoder–decoder allocation depends on the task. Both parity allocations listed achieve 100 percent in each of the three seeds. For swaps, increasing decoder depth up to four layers improves the reported 256-operation result. The seeds vary initialization on a shared training stream and held-out examples.

All comparisons use width 512 and four attention heads. Equal logical depth is not a parameter- or compute-matched comparison. Addition extrapolation and standard S5 remain difficult in this study. These experiments compare RLT with a Transformer baseline and do not measure an SSM or DeltaNet baseline.

RLT / Algorithmic experiments and length extrapolation

11 / RLT-1 and RLT-2

Feedback frequency and parallelism

At Layer 4 + 4	Feedback interval	Parity / 64 bits	Swaps-\(S_5\) / 64 ops
RLT-1	1 token	100.00 ± 0.00%	100.00 ± 0.00%
RLT-2	4 tokens	98.99 ± 1.66%	19.60 ± 6.54%

RLT-2 shares the previous boundary state within each chunk.
Known positions in a chunk run in parallel, then its last output updates feedback.

Final-answer/state accuracy, mean ± sample SD, seeds 42–44. Same checkpoint protocol as the previous slide.

RLT / Feedback frequency / RLT-2 / Chunk semantics

12 / 14

Both models have the same encoder–decoder allocation and parameters. RLT-2 changes how often feedback updates. Within a chunk, each position receives the same previous boundary feedback vector, while causal attention and KV updates remain part of the model.

Chunk size four preserves nearly all 64-bit parity accuracy at this split but sharply reduces 64-operation swap tracking. Removing feedback entirely gives 50.23 plus or minus 3.80 percent parity accuracy and near-uniform swap tracking at this length. The RLT-0 control also removes the merge parameters, making it 787,968 parameters smaller.

RLT chunking changes the recurrence and therefore the model. This differs from an SSD chunk algorithm that reorganizes evaluation of the same affine recurrence. Known-token chunks can run in parallel, while autoregressive generation still waits for newly sampled tokens. The schedule diagram describes dependencies, not measured speedups.

RLT / Feedback frequency / RLT-2 / Chunk semantics

12 / Comparing the mechanisms

State size and parallel evaluation

Model core	Update in the previous state	Persistent memory*	Known-token evaluation
Linear attention	Additive / affine	Fixed matrix and normalizer	Prefix sums or chunked contraction
Selective SSM	Input-conditioned affine	Fixed state	Structured scan or SSD
DeltaNet	Affine rank-one edit	Fixed matrix	Structured chunk algorithm
RLT-1	Nonlinear decoder cell	Bounded decoder state, growing encoder memory	Parallel encoder, sequential decoder

*At fixed width, depth, and local window. State sizes refer to inference, excluding training activations.

Note / Linear attention, SSD, and Delta Networks / RLT / State and execution

13 / 14

This table compares the recurrent cores, with each layer's inputs held fixed when classifying affine transitions. Nonlinear feature projections, gates outside the core, and stacked layers can make the overall models nonlinear.

The fixed-state claims for linear attention, SSMs, and DeltaNet hold at fixed architectural dimensions. Their training activations can still grow with sequence length. RLT's decoder state contains a fixed-width vector and at most W minus one cached positions per SWA layer.

RLT also stores encoder memory over the entire prefix, as well as the encoder caches needed during incremental generation. It therefore does not have constant total inference memory. The full-prefix attention operations also prevent inferring linear total sequence cost from the bounded decoder state.

Note / Linear attention, SSD, and Delta Networks / RLT / State and execution

13 / Connection to RLT

The transition determines the computation

Structured recurrent memory

\[h_t=A_t h_{t-1}+c_t\]

Linear attention, SSMs, and DeltaNet compress history through structured affine updates.

Nonlinear recurrent computation

\[H_t=\mathcal F_\Theta(H_{t-1};e_t,M_{\le t},t)\]

RLT reuses a complete decoder cell across tokens and retains explicit context memory.

The state-space view connects the models.
The transition’s structure determines what can be composed efficiently.

Open question: can a structured transition retain RLT’s state-tracking gains?