PaperlayerView exact revision on arXiv

Source-linked AI summary

WriteSAE: Sparse Autoencoders for Recurrent State

Jack Young

arXiv:2605.12770v4cs.LGcs.AIcs.CL

TL;DR

Residual-stream SAEs cannot directly replace matrix-shaped recurrent cache updates, so the paper introduces rank-1 WriteSAE atoms matched to the model’s writes. Direct replacement generally improves final-output agreement over deleting writes, while a Gated DeltaNet formula predicts logit changes and supports cache interventions, with transfer to Mamba-2. The paper also bounds these results by architecture-specific failures and limited scaling evidence.

  • Problem

    Residual-stream SAEs have vector-shaped atoms and cannot directly replace the matrix-shaped updates written into recurrent caches.

  • Method

    WriteSAE learns rank-1 matrix atoms and tests them by replacing native recurrent writes with activation-scaled atoms during continued forward passes.

  • Results

    92.4% of evaluated positions favored atom replacement over deletion, the Gated DeltaNet logit formula reached R^2=0.98, and Mamba-2-370M reached 88.08%.

  • Takeaways & Limitations

    Rank-matched atoms can often replace individual recurrent cache writes, and the replacement test transfers across recurrent families even though the Gated DeltaNet gate coefficient does not.

  • Takeaways & Limitations

    The Gated DeltaNet logit formula does not transfer directly to Mamba-2, and the reported results remain bounded by limited larger-model and multi-feature evidence.

Abstract

from arXiv · show

We introduce WriteSAE, a sparse autoencoder for the matrix updates written into recurrent language-model state. In Gated DeltaNet, Mamba-2, and RWKV-7, each token writes a matrix-shaped update to a recurrent cache; a residual-stream SAE has vector-shaped atoms and cannot replace that update directly. WriteSAE learns rank-1 matrix atoms with the same shape as the model's own write. This lets us test a direct replacement: at positions where the SAE activates an atom, we remove the model's write, insert the atom scaled by its SAE activation, and continue the forward pass. The atom gives a closer final token distribution than deleting the write on 92.4% of evaluated positions; averaged per atom, the rate is 89.8%. For Gated DeltaNet, a formula using the forget gate, read query, and output embedding predicts the resulting logit change with $R^2 = 0.98$. The same replacement test transfers to Mamba-2-370M at 88.1%. In generation, the formula chooses a write direction; writing it into three consecutive cache positions at $3\times$ the norm of the model's write makes tokens initially ranked 100--1000 by the unmodified model appear in 100% of continuations, up from 33.3%. To our knowledge this is the first cache-level steering intervention reported in a state-space or hybrid recurrent layer.

1 Introduction

WriteSAE represents recurrent cache writes as rank-1 matrix atoms, enabling direct replacement of native updates that vector-shaped residual-stream SAEs cannot perform. In Gated DeltaNet, the replacement improves final-output agreement over deleting writes and supports logit and generation interventions.

  • Motivation: Recurrent models write matrix-shaped updates into a cache, whereas residual-stream SAEs produce vector-shaped atoms that cannot directly replace those updates.WriteSAE matches the rank-1 outer-product shape of the model’s write.
  • Method: WriteSAE decoder atoms are rank-1 matrices, preserving the cache shape read by later tokens.A firing is a position where the SAE assigns a nonzero coefficient to an atom.
  • Results: 92.4% of evaluated positions favored replacing the native Gated DeltaNet write with an atom over deleting the write.The per-atom average was 89.8%.
  • Results: R^2=0.98 was achieved by a formula using the forget gate, read query, and output embedding to predict logit changes.The formula is evaluated for Gated DeltaNet’s rank-1 perturbations.
  • Results: The same replacement test reached 88.08% on Mamba-2-370M.This result extends the replacement comparison beyond the primary Gated DeltaNet setting.
  • Results: Writing the formula-selected direction into three cache positions at 3× the model-write norm raised the appearance rate of initially rank-100–1000 tokens from 33.3% to 100%.The intervention used generation continuations.

2 Method

The method trains rank-1 WriteSAE atoms on recurrent cache states and evaluates them by replacing native writes during continued forward passes. A closed-form Gated DeltaNet expression maps cache perturbations to later-token logit changes, while KL divergence measures output-distribution effects.

  • Logit prediction: The analysis predicts a later token’s logit change from a rank-1 atom update written into the cached state.The target is the logit change for token tok at a later position t.
  • Logit prediction: The formula uses forget-gate products, the later read query, and the output-embedding row for the candidate token.Every quantity on the right is observable from a single forward pass.
  • Logit prediction: Subtracting perturbed and native trajectories isolates how later recurrence steps propagate the cache perturbation before the query reads it.When the atom direction is nearly orthogonal to later keys, the extra update term is small, matching the regime of high fit.
  • WriteSAE: A rank-1 decoder atom matches the model’s rank-1 write, allowing one native cache update to be replaced without changing the matrix shape.WriteSAE uses a TopK encoder and rank-1 decoder atoms.
  • Evaluation: Figure 2 compares cosine-based atom classes, held-out passage NLL under deletion or random writes, pooled KL, and a larger-model generation probe.The figure reports 316 alive atoms and evaluates KL over 4,851 positions.
  • Replacement test: At the largest TopK activation, the protocol swaps the native write for the atom and continues the forward pass.The final-output score is KL(p_patched ∥ p_baseline).

3 Experiments

Experiments test whether WriteSAE atoms can replace recurrent cache writes, validate the mechanism behind substitution quality, and assess transfer across architectures and encoder choices.

  • 3.1 Feature classes: WriteSAE trains 2,048 atoms, with 316 surviving validation and a descriptive split into 222 registers, 94 bundles, and 1,732 null atoms.The two-component Gaussian mixture used median cosine-to-native-write, with ΔBIC = −296 over a one-component null.
  • 3.1 Feature classes: Deleting every register firing raises NLL by +0.005 bits/token, versus +0.226 for matched-norm random rank-1 writes across 20 held-out passages.The 41.87× gap holds in 19/20 passages.
  • 3.2 Mechanism validation: 92.4% of 4,851 firings favor atom substitution over deletion, while the strict ordering KLatom < KLdelete < KLrandom holds on 89.5% of firings.The per-atom average win rate is 89.8% across 87 atoms; bundle and register atoms both show replacement success.
  • 3.2 Mechanism validation: Atom substitution remains strongest across all L9 heads, averaging 89.3% ± 2.6% with a range of 82.6%–93.2%.H12 is dead, so the test covers 15 of 16 L9 heads.
  • 3.2 Mechanism validation: The logit-change formula obtains median per-feature R2 = 0.98 across seven registers and bundle F87, using gates, the read query, and output embedding.The cosine factor accounts for the substitution gap, while output-embedding projection is not limiting.
  • 3.3 Architectural scope: Write-rank comparisons extend the replacement test beyond Gated DeltaNet: Mamba-2-370M reaches 88.08% atom-over-deletion wins, while Qwen3.5-4B is at chance under the same recipe.The cross-host averaged rates are 89.8% for GDN and 88.1% for Mamba-2.
  • 3.4 Ablations: At matched nf=2,048, sparsity k=32, and training budget, the bilinear encoder has 32% dead features versus 80% for FlatSAE across 720 runs.BatchTopK and JumpReLU recover the same register/bundle partition under the bilinear encoder.

4 Cache Intervention Probes

Cache-level interventions use formula-selected write directions and feature amplification to alter target inclusion and newline rates. Effects depend on target class, dose, and the host model.

  • Cache-slot erasure: Erasing F412 at 150 natural firing positions lowers the most-affected token’s log-probability by a median 0.116 nats, unlike non-firing positions.The natural-firing effect is significant, whereas the non-firing median change is +0.016 with p=0.15.
  • Single-position prediction: 84.6% of 2,000 trials have the predicted sign of the single-position logit change, but pooled magnitude prediction has R2 = −0.06.The median measured-to-predicted ratio is 1.08, while small noisy effects weaken the pooled correlation.
  • Generation edit: 100% target inclusion, up from 33.3%, occurs when the formula-selected direction is written at three cache positions for initially rank-100–1000 targets.This probe uses m=3× on Qwen3.5-0.8B L9 H4 with n=300 contexts; out-of-context targets remain at 0%.
  • Held-out 4B probe: At 5× amplification, boundary features reduce mean newlines from 16.8 to 11.2 per 400 tokens on a held-out 4B model.The reduction is 33% across 40 prompts; at 10×, newline counts rebound toward baseline.
  • Controls: Only BilinearSAE features produce the target newline reduction; FlatSAE amplification reduces word length without changing paragraph count.These pilots serve as negative controls with different feature pools and higher dead-feature rates.

5 Related Work

WriteSAE targets matrix-valued recurrent-state writes, unlike prior SAE and transcoder interventions that operate on vectors. This distinct intervention site reflects the matrix cache used by recurrent and hybrid models.

  • Intervention site: Residual-stream SAEs decode vector atoms, while WriteSAE uses atoms matching the matrix shape written into recurrent state.The paper positions this as a different intervention site from residual-stream analysis and MLP transcoders.
  • Recurrent-model context: Modern recurrent hybrids reuse recurrent-state arithmetic across architectures including RetNet, GLA, Gated DeltaNet, RWKV-7, and Mamba-2.Prior probes describe recurrent-state content, whereas this work intervenes directly on the matrix state later tokens read.

6 Discussion and Conclusion

The paper’s central claim is that rank-1 atoms can often replace rank-1 recurrent writes, with the clearest evidence and logit formula fit in Gated DeltaNet. Transfer across architectures is partial and bounded by several failures.

  • Discussion and conclusion: Rank-1 SAE atoms can often replace a single rank-1 matrix update written by the model.Gated DeltaNet provides the clearest register class and the best closed-form logit-shift fit.
  • Cross-architecture scope: The replacement test transfers across host architectures, but the Gated DeltaNet gate coefficient does not transfer directly.The GDN formula reaches R2=0.98, whereas the same formula yields negative R2 on Mamba-2 and Qwen3.5-4B.
  • Limitations: At Qwen3.5-4B, the same SAE recipe beats deleting the write only at chance despite better validation MSE than the 0.8B SAE.This separates state reconstruction quality from write-direction alignment.
  • Observed failures: Mamba-2 generation edits produce no target-inclusion lift, consistent with the failure of the GDN-derived logit formula outside GDN.The reported generation probe uses 3,600 trials and has approximately zero median first-step log-probability lift.

A Derivation of the Three-Factor Logit Factorization

The derivation factorizes a rank-1 cache perturbation’s eventual logit change into gate decay, cache read alignment, and output projection. Empirical tests support this approximation for Gated DeltaNet, while also identifying architecture- and sparsity-dependent scope conditions.

  • Three-factor factorization: The predicted logit change combines a prompt-specific gate product, read-query alignment, and output-embedding projection.The gate product absorbs propagation decay and Jacobian rotation or rescaling into one scalar factor.
  • Derivation: A rank-1 perturbation propagates through the remaining recurrence, is read by the later query, and then passes through downstream layers to produce a logit shift.The perturbation begins as εv_iw_i^T in the cached state and is propagated to the final logits.
  • Empirical validation: Median R2=0.98 across the tested feature population, with all 200 fitted cells exceeding R2=0.95.Top-ranked evaluation-token logits retain R2 in [0.983, 0.984], whereas rank-50 tail logits degrade to approximately 0.05.
  • Scope: The approximation is first-order in ε and its selectivity support is ε-invariant across ε ∈[0.1, 3].The reported scope is limited by higher-order Taylor effects in tail logits and by architecture-specific behavior in host-architecture analogs.
  • Sparsity robustness: JumpReLU reproduces the register/bundle partition, while gated SAE variants fail to provide a stable matched-L0 objective.JumpReLU achieves 105× within-SAE separation versus 29× for BatchTopK; gated variants either collapse, stall, or have worse MSE.

C Full Encoder-Swap MSE Comparison

The appendix compares architecture-matched rank-1 WriteSAE atoms with dense-state encoders and controls. Rank-1 structure preserves the model’s write primitive while remaining competitive in reconstruction and more interpretable for single-write replacement.

  • Architectural match: A Gated DeltaNet write is one rank-1 outer product, so a rank-1 atom patches the same cache event directly.Higher-rank atoms can instead compress multiple writes that the firing-level cache intervention cannot accept as one event.
  • MSE comparison: The flat upper bound beats rank-1 by only 8% at L1/L9 despite using 60× more decoder parameters per atom, and by 1.75× at diffuse L17.Flat atoms span the full state but may mix several native writes into one feature.
  • Training-budget control: At 5×10^6 states/head with 200 epochs, the matched-cohort ordering returns: BilinearSAE 85.8% < FlatSAE 93.4% < MatrixSAE 94.0%.The 5K-corpus crossover is attributed to training dynamics rather than coverage, and the parameter-matched BilinearSAE control remains worse downstream.
  • Selectivity controls: Selectivity remains high across the measured head sweep, with mean 0.9953 and 39/47 cells exceeding 0.99 at K=32.Across 592 feature-cell pairs, no measured control dips below 0.996 selectivity.

F Cross-Architecture Partition and Scaling

Cross-architecture and scaling analyses show that replacement success is broadly stable in gated Qwen3.5 Gated DeltaNet, but geometry alone does not guarantee success when cache contents integrate history without decay.

  • Scaling: The partition persists across the 34× Qwen3.5 scale range, with register counts near 220 at 0.8B and 4B and 147 at 27B.Median register cosine softens from 0.26 to 0.09 but remains above the threshold cos=0.05.
  • Head-level replication: Across 15 L9 heads, the mean win rate is 89.29% ± 2.63%, while L9 H4 reaches 90.84%.Fourteen of fifteen heads exceed 85%, and twelve exceed 88%.
  • Population replacement: 89.80% of firings across 87 alive atoms beat deletion, with 91.36% for registers and 88.98% for bundles.The overall 95% CI is [88.1, 91.3], and the class difference is not significant by the reported Mann-Whitney test.
  • Rank matching: The rank-2 substitution has downstream perplexity 20.360 versus 20.347 for rank-1 despite lower reconstruction error.This supports matching atom rank to the model’s rank-1 write for the substitution intervention.
  • Failure analysis: The 7.6% cases where deletion beats the atom concentrate among small-effect firings, falling from 12.3% in Q1 to 4.9% in Q4.The quartiles are ordered from smallest to largest deletion effect size.

G.2 Cosine-free classifier reproduces the substitution result

A classifier based on replacement success rather than decoder cosine reproduces the main register result, showing that the substitution pattern is not dependent on the cosine-based partition.

  • Classifier agreement: The replacement-success classifier agrees with the cosine classifier on 7 of 8 features that fired above threshold.The single disagreement is bundle exemplar F87, whose atom beats deletion on 94% of 300 firings.
  • Robustness: Removing cosine from the pipeline leaves the reported result unchanged at 91.25% for the replacement-success register class.The replacement-success class fires 1,851 times, while cosine-register and cosine-bundle rates are 90.72% and 94.00%, respectively.
  • Seed stability: The class-level claims replicate across SAE seeds even though individual atom identities do not.This distinguishes stability of the population-level partition from stability of particular learned features.
  • Interpretive connection: The no-decay DeltaNet failure shows that high cosine can coexist with poor replacement when an atom matches integrated history rather than a local write.This contrasts with gated Qwen3.5 DeltaNet, where per-position decay restores replacement success above 89%.

H Register vs bundle: systematic differences beyond population win rate

Register and bundle atoms differ in firing breadth and top-1 disruption, while showing comparable lesion magnitude and no significant population-scale substitution gap. Reader traces further indicate that selected register atoms are enriched in specific downstream attention heads rather than diffusely represented.

  • Population comparison: Population substitution rates are 91.4% for registers and 89.0% for bundles, with no significant gap.The population comparison reports p=0.24, so class membership does not predict whether an individual substitution beats deletion.
  • Top-1 disruption: Registers flip the model’s top-1 on 70.0% of firings, versus 63.3% for bundles.The difference is significant with Cliff’s δ=+0.29 and p=0.041.
  • Firing breadth: Registers fire on a median of 100 validation tokens, versus 70.5 for bundles.Firing breadth correlates strongly with cosine to the native write (Spearman ρ=+0.73, p<10^-5).
  • ΔKL effect size: Median ΔKL is 0.682 nats for registers and 0.579 for bundles, a nonsignificant difference.The reported effect size is Cliff’s δ=+0.11 with p=0.45.
  • Control analysis: Matched-norm random rank-1 controls score above 0.99 on top-K-overlap selectivity, so firing-level KL ordering remains the distinguishing evidence.The orthogonal control returns selectivity 0.998, and matched-norm random perturbations average 0.9953 across 47/48 cells.
  • Reader traces: Register atoms read into specific later attention heads at 3–7× baseline, while the signal does not diffuse across the residual stream.The reported exemplars are F53→L5 H5 at 6.9×, F63→L17 H4 at 5.2×, and F1335→L21 H10 at 7.5×; the authors do not claim generality.

I.6 Generation intervention: full results

The generation intervention writes a formula-selected rank-1 direction into three consecutive cache positions and evaluates greedy continuations across target-rank strata and write magnitudes. It is most effective for targets initially ranked 100–1000, while out-of-context targets show rank shifts without appearing and larger writes eventually degrade generation.

  • Intervention setup: The intervention applies the formula-selected direction across three consecutive cache positions before generating 20 tokens greedily.Targets are stratified by their initial unmodified-model rank into frequent, 100–1000, rare, and semantic classes.
  • Pooled result: Across all four target classes, the 3× edit lifts target appearance by +16.7pp, from 8.3% native to 25.0%.The pooled evaluation contains n=1,200 trials.
  • Rank 100–1000: At 3× write norm, targets initially ranked 100–1000 appear in 300/300 continuations, versus 33.3% under the unmodified model.This is a +66.7 percentage-point lift for the rank-100-to-1000 class.
  • Target-stratification boundary: Frequent, rare, and semantic out-of-context targets appear in 0% of continuations at every tested magnitude despite 4,039–17,526-position rank shifts.These targets begin at ranks of at least 17,000, which greedy decoding over three edited positions does not promote to top-1.
  • Magnitude sweep: Pooled lift is non-monotone across magnitudes: +8.3pp at 1.5×, +16.7pp at 3.0×, and +8.3pp at 6.0×.Beyond the best measured magnitude, the cache write dominates surrounding context and degrades the rest of generation.

J Reproducibility

The project provides code, checkpoints, cached states, and deletion-control outputs, while documenting the per-firing replacement procedure and implementation choices used for training and generation probes.

  • Artifacts: All scripts producing the reported numbers, tables, and figures are provided in the repository snapshot, with trained checkpoints and cached states available on HuggingFace.The listed artifacts include per-head deletion-control JSON outputs; code and checkpoints are under MIT.
  • Replacement procedure: Each firing selects the dominant TopK atom, rescales its rank-1 decoder atom to the native write’s Frobenius norm, and compares atom, deletion, and random conditions.The procedure uses k=32 nonzeros and matched-norm forward passes.
  • Training implementation: Dense-encoder WriteSAE training keeps 300–1,200 alive features only with both auxiliary dead-feature loss and periodic resampling.The auxiliary coefficient is λaux=10^-2, and the resampler fires every 250 steps.
  • Encoder and execution choices: The bilinear matched-filter encoder is used for the 4B generation probe because its firing coefficient matches projection onto the decoder’s rank-1 direction, despite 5–15% higher MSE.Skipping the cache deep-copy biases results toward whichever condition runs last and costs about 1.4× wall-clock per firing.
Loading 2605.12770v4…