🤖 machine learning

WriteSAE: Sparse Autoencoders for Recurrent State

This paper introduces WriteSAE, the first sparse autoencoder designed to decompose and edit the matrix cache writes of state-space and hybrid recurrent language models by factoring decoder atoms into their native rank-1 update shape, enabling precise behavioral interventions and highly accurate effect prediction.

Original authors: Jack Young

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jack Young

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a large language model (like the ones that write stories or answer questions) as a massive, high-speed factory. Inside this factory, there's a special "memory board" where the machine writes down its current thoughts before moving to the next step.

For a long time, scientists have tried to understand what's written on this board using tools called Sparse Autoencoders (SAEs). Think of these tools as a pair of glasses that let researchers see the individual "ideas" or "features" the machine is thinking about. However, there was a problem: these glasses only worked on the output of the factory, not on the specific way the machine writes to its memory board in certain advanced models.

This paper introduces a new tool called WriteSAE. Here is how it works, explained simply:

1. The Problem: The Wrong Shape of the Glasses

In older models, the machine wrote its thoughts as a simple list of numbers. The old glasses (standard SAEs) were designed to read lists.
But in newer, faster models (like Gated DeltaNet, Mamba-2, and RWKV-7), the machine doesn't write a list. Instead, it writes a matrix (a grid of numbers) using a specific mathematical trick called a "rank-1 update."

The Analogy: Imagine the machine is painting a picture. The old tools tried to describe the painting by looking at the finished canvas. But the new machines paint by adding one specific brushstroke (a single line of color) at a time to a complex grid. The old tools couldn't see that single brushstroke because they were looking for a whole list of colors, not a single stroke.

2. The Solution: WriteSAE

The authors built WriteSAE, a new set of glasses designed specifically to see that single brushstroke.

The Shape Match: Instead of trying to read a list, WriteSAE's "atoms" (the pieces it looks for) are shaped exactly like the brushstroke the machine uses. They are tiny, single-stroke patterns.
The Result: Because the shape matches perfectly, WriteSAE can isolate exactly what the machine is writing to its memory at any given moment.

3. The Experiments: Swapping the Brushstrokes

To prove this works, the researchers did a "surgery" on the machine's memory.

The Swap: They found a moment where the machine wrote a specific brushstroke. They erased that stroke and replaced it with a "learned" stroke from their new dictionary (WriteSAE).
The Test: They compared this to two other scenarios:
1. Erasing it completely (leaving a blank spot).
2. Putting in a random scribble.
The Outcome: When they swapped in the WriteSAE stroke, the machine kept working almost exactly as it did before (92.4% of the time). When they erased it or used a random scribble, the machine got confused and made mistakes.
The Metaphor: It's like replacing a specific gear in a clock with a custom-made gear that fits perfectly. The clock keeps ticking. If you remove the gear or put in a random rock, the clock stops.

4. What They Found

Two Types of Strokes: They discovered the machine uses two main types of brushstrokes:
- Registers: These are precise, focused strokes that do specific jobs (like marking the start of a sentence or a proper name).
- Bundles: These are more scattered strokes that seem to do a mix of things.
Predicting the Future: They found a mathematical formula that can predict exactly how changing a specific stroke will change the machine's next word. It's like knowing that if you tweak this one specific gear, the clock will chime one second earlier.
Editing the Machine: They successfully used this tool to "install" new behaviors. For example, they could force the machine to keep talking about a specific topic (like a "mid-rank" word) that it wouldn't normally choose, simply by inserting the right brushstroke into the memory board.

5. The Limits

The paper is very careful to say this works best on models that write in this specific "single-stroke" (rank-1) way.

If a model writes in a more complex way (like using two strokes at once or a diagonal pattern), the tool doesn't work as perfectly, though it still finds some patterns.
The tool works great on the Qwen3.5 model (a specific type of AI), and they showed it works on other similar models too, but the "magic formula" for predicting the future changes depending on the model's architecture.

Summary

WriteSAE is a new tool that lets us see and edit the specific way advanced AI models write to their internal memory. By matching the shape of the tool to the shape of the machine's writing, the researchers can swap out specific thoughts, predict how the machine will react, and even steer the machine to say things it wouldn't normally say, all without breaking the machine. It's the first time scientists have successfully performed this kind of "surgery" directly on the memory-writing site of these specific types of models.

Technical Summary: WriteSAE for Recurrent State Decomposition

Problem Statement

State-space and hybrid recurrent language models (e.g., Gated DeltaNet, Mamba-2, RWKV-7, Qwen3.5) utilize a matrix cache for state propagation, where each token writes a rank-1 outer product ( $k_t v_t^\top$ ) into a $d_k \times d_v$ matrix. Existing Sparse Autoencoders (SAEs) operate on residual streams, reading the state after emission. Consequently, standard SAEs cannot directly access or decompose the native write mechanism, which occurs upstream in the recurrence. Furthermore, because the state is updated via rank-1 outer products, a standard vector-based SAE decoder atom cannot replace a single cache write without breaking the structural correspondence required for causal intervention (cache patching).

Methodology

The paper introduces WriteSAE, the first sparse autoencoder designed to decompose and edit the matrix cache write of recurrent models.

Architecture and Training

Native Shape Matching: WriteSAE decoder atoms are constrained to be rank-1 outer products ( $v_i w_i^\top$ ), matching the native write primitive ( $k_t v_t^\top$ ) of the host architecture. This allows a single atom to replace a specific cache slot update.
Bilinear Encoder: To efficiently encode the matrix state $S_t$ , the model uses a bilinear encoder ( $a_i = v_i^\top S_t w_i$ ) rather than a flattened dense encoder. This reduces parameter count per atom from $d_k d_v$ (16,384 for $128 \times 128$ ) to $d_k + d_v$ (256), a 64x reduction.
Training Objective: The model is trained on mean-centered states ( $x = \text{vec}(S_t - M)$ ) using a TopK sparsity constraint and an auxiliary loss to revive dead features. The training corpus consists of OpenWebText sequences processed through Qwen3.5-0.8B (and other models) at specific layers and heads.

Feature Classification

The trained dictionary is partitioned into two classes based on the median cosine similarity between the atom's write direction and the native write direction:

Registers: Atoms with high cosine alignment to native writes (recoverable from the cache).
Bundles: Atoms with dispersed write directions across the cache.

Causal Intervention Tests

The paper proposes three primary validation mechanisms:

Cache-Slot Substitution: Replacing the native write ( $k_t v_t^\top$ ) with a matched-Frobenius-norm SAE atom ( $\hat{S}_t = S_t - \Delta_{nat} + \Delta_{atom}$ ) and measuring the downstream KL divergence.
Closed-Form Logit Shift: Deriving a three-factor expression to predict per-token logit shifts: $\Delta \ell \approx G_{t_0 \to t} \langle w_i, q_t \rangle \langle v_i, W_U[\text{tok}] \rangle$ , where $G$ is the gate product, $q_t$ is the read query, and $W_U$ is the unembedding row.
Direct Cache Interventions: Erasing or installing atoms at specific positions to observe changes in token probabilities and generation behavior.

Key Contributions

Architecture-Matched Dictionary: A dictionary where atoms are rank-1 outer products, enabling direct substitution of native cache writes, a capability previously unavailable to residual-stream SAEs.
Substitution Validation: Demonstration that learned atoms can replace native writes with high fidelity. At Qwen3.5-0.8B (Layer 9, Head 4), atoms beat matched-norm ablation on 92.4% of 4,851 firings.
Closed-Form Prediction: A parameter-free, three-factor formula that predicts per-firing logit shifts with a median $R^2$ of 0.98 across 200 atom-by-epsilon cells.
Cross-Architecture Transfer: Validation that the substitution success and feature partitioning (Register vs. Bundle) hold across different matrix-recurrent architectures (Gated DeltaNet, RWKV-7, Mamba-2), though success rates vary by the "write rank" of the substrate.

Experimental Results

Substitution Performance

Qwen3.5-0.8B (Gated DeltaNet): At L9 H4, atoms beat ablation on 92.4% of firings. The population test over 87 atoms holds at 89.8%. The strict ordering $KL_{atom} < KL_{ablate} < KL_{random}$ holds on 89.5% of firings.
Mamba-2-370M: Substitution succeeds at 88.1% over 2,500 firings, despite Mamba-2 using a diagonal state update rather than a full rank-1 outer product.
RWKV-7: Shows intermediate performance, consistent with its rank-2 write rule.

Logit Shift Prediction

The closed-form expression accurately tracks measured effects. At L9 H4, the median $R^2$ between predicted and measured logit shifts is 0.98. However, this closed form fails (negative $R^2$ ) on Mamba-2 and Qwen3.5-4B, indicating the gate factor $G$ is substrate-specific.

Behavioral Interventions

Erasure: Erasing a specific register atom (F412) at its natural firing positions reduced the log-probability of the target token ("space") by a median of -0.116 nats ( $p=1.07 \times 10^{-6}$ ).
Generation Steering: Sustained installation of a closed-form direction at three consecutive positions increased the "target-in-continuation" rate for midrank tokens from 33.3% to 100% under greedy decoding.
Passage-Level Amplification: Amplifying boundary-correlated features in a 4B model reduced newline generation by 33% (from 16.8 to 11.2 per 400 tokens) at a 5x dose, demonstrating control over document structure.

Architectural Scope

The paper establishes a hierarchy of "write rank" effectiveness:

Gated DeltaNet (Rank-1): Highest register cosine separation (0.262) and substitution fidelity.
RWKV-7 (Rank-2): Intermediate separation (0.180).
Mamba-2 (Diagonal/Scalar): Lowest separation (0.0575), yet still achieves ~88% substitution success.

Significance and Claims

The paper claims to provide the first causal dictionary for the matrix-recurrent write site, moving beyond residual stream analysis. It demonstrates that:

Write Rank is Determinative: The ability to decompose the state into interpretable "registers" correlates with the rank of the native write rule (Rank-1 > Rank-2 > Diagonal).
Causal Substitution is Possible: Learned atoms can replace native writes with minimal downstream KL divergence, validating the "register" concept in recurrent models.
Closed-Form Interpretability: For Gated DeltaNet, the effect of a feature on logits can be predicted analytically without fitting parameters, linking the dictionary directly to the model's gating and reading mechanisms.

The authors modestly note limitations: the closed-form coefficient does not transfer across substrates (e.g., fails on Mamba-2), and per-atom identity is seed-specific (though class-level properties are stable). The work establishes a new benchmark for mechanistic interpretability in state-space models, showing that the "write" primitive is a viable target for dictionary learning and causal editing.