Sparse Weight Decomposition for Efficient Circuit Extraction
Chuanhao Yan, Xuhan Huang, Yawen Duan, Zhenfei Yin, Hang Zhao, Bryan Dai, Jie Fu
cs.LG, cs.CL
2026-08-05
SWD factors each weight matrix into two sparse matrices whose shared coordinates become ablatable circuit units, matching baselines with under 1% of the data and fewer edges.
Mechanistic interpretability tries to answer a specific question: when a model gets something right, which components are doing the work? Those components form a "circuit." To find them you need units you can keep or switch off individually. Dense pretrained weights offer no such units; every neuron connects to everything, so you cannot say which edge belongs to which function.
Existing answers all bolt something onto the model. Sparse autoencoders and Transcoders train a sparse dictionary or replacement network on lots of text, which costs compute and leaves a fidelity gap between the replacement and the original. Sparse pretraining trains a whole new weight-sparse model, which the paper notes takes 100 to 1000 times the training and inference compute. All of that, just to obtain units you can treat as circuit components.
The question here: can you pull those units straight out of an existing checkpoint, without training any auxiliary representation?
SWD factorizes the weight matrix W of a dense linear projection into two sparse factors, W ≈ AB. A and B are sparse, and they share a column of intermediate coordinates. Coordinate i is a "bottleneck unit": it reads a few input dimensions through column A:,i, computes a scalar zi, and writes a few output dimensions through row Bi,:. The path A:,i · Bi,: is rank one and touches only a handful of coordinates on each side.
The useful property falls out of one identity: hAB = Σ zi Bi,:. Replace zi and you change only one term of the sum; the other paths are untouched. Each unit is individually addressable, exactly what circuit extraction needs.
Why two sparse factors instead of pruning or SVD? Pruning removes scalar edges but creates no intermediate unit, so there is nothing to ablate. SVD gives rank-one components, but each one's read and write vectors are dense and touch every dimension, so units carry many edges. SWD bets on sparse read/write structure itself, not on factorization as such. A control confirms it: exact SVD and random-orthogonal factorizations both need more active edges than SWD at the same fidelity, so it is the sparsity doing the work.
The objective minimizes E[‖hW − hAB‖²] under a fixed nonzero budget K, where the calibration Gram matrix G = E[h⊤h] steers capacity toward input directions that carry mass under the model's activation distribution. The solver is a Double Sparse Factorization heuristic: split the budget between A and B, alternate block updates each solved by ADMM (regularized least squares, hard-threshold projection to the budget, dual update), then freeze the supports and refit the surviving B entries.
A zero-data variant sets G to the identity when no calibration activations are available, reducing the objective to ‖W − AB‖²F. It trades some behavioral fidelity for weight fidelity and can be applied at every training step.
The contrast with Transcoder is clean (Table 1). Transcoder trains a replacement network whose hidden features are the units, with activation sparsity and a need for text data. SWD fits sparse factors post hoc on a single matrix, with weight sparsity and almost no data.
On single-matrix replacement, data efficiency is the most visible result. SWD at 50% sparsity drives the cross-entropy delta low with a few thousand tokens; Transcoder and VPD, which train replacement networks, need roughly 10⁶ tokens. SWD uses under 1% of their data. This holds on GPT-2 and Qwen2.5 at 0.5B/1.5B/3B, and the single-matrix experiment extends to Qwen3.5-27B.
Circuit cost is compared at matched replacement fidelity (replacement CE within 0.001 of the dense model), reporting how many active edges are needed to reach given sufficiency (task behavior retained when only the selected units are kept) and necessity (behavior lost when they are ablated). On four GPT-2 tasks (greater-than, IOI, docstring, gendered-pronoun) SWD generally hits the thresholds with fewer active edges, and the conclusion holds under zero ablation. The same pattern appears on Qwen2.5-3B and Qwen3.5-27B.
Full-model replacement on GPT-2 Small replaces all 48 attention and MLP weight matrices across 12 transformer blocks:
| Method | Data | Active nonzeros | LM CE |
| SWD | 4.19M calibration tokens | 26.77M | 3.90 |
| SWD-FT | 4.19M calibration + 16.38M fine-tuning | 26.77M | 3.44 |
| Sparse pretrain | 2.884B tokens | 26.77M | 3.45 |
SWD-FT freezes the support and refits only nonzero entries, reaching CE 3.44, slightly below the sparse-pretrained checkpoint at 3.45 trained on 2.884B tokens. It uses about 20.6M tokens in total, under 1% of that budget. On circuit quality (Figure 7) the three are comparable on greater-than and gendered-pronoun, but sparse pretraining has near-zero or negative necessity drop on docstring and IOI, while SWD-FT stays effective on all four.
Two qualitative experiments probe whether units carry real semantics. In the GreaterThan circuit, 9,208 bottleneck units are ranked by attribution, and four of six displayed units relate to numbers, quantities, or measurement. A targeted edit is sharper: editing unit c205 on the prompt "the opposite of up is" raises the answer margin by 0.216 at the strongest positive setting, while the mean final-token KL on seven unrelated factual prompts is 4.02×10⁻⁵, lower than a random-unit control and rank-4 LoRA. In attention diagnostics, the SWD-reconstructed attention for layer 9, head 3 matches the dense version with KL 0.0405; ablating query unit q266 redirects attention sharply toward the first token, raising KL to 2.496 and producing a maximum probability change of 0.775.
For mechanistic-interpretability practitioners, getting addressable circuit units turns from training an auxiliary model on millions of tokens into one post-hoc matrix factorization on the checkpoint. The cost drops by an order of magnitude, and because SWD fits the actual checkpoint weights rather than a learned proxy, the fidelity gap that dogs SAE and Transcoder comparisons narrows. The zero-data variant opens per-step analysis during training.
A more concrete use is targeted intervention. The c205 example shows you can nudge one concept with almost no spillover, useful for steering and controllability work.
Be honest about scope. This is an efficiency improvement on a specific front, single-matrix and GPT-2-scale full-model replacement, not a new theory of how models work. The circuits you get are still local and non-unique.
The paper states its limitations more plainly than most.
Efficient extraction is not complete understanding. Making units cheap to obtain does not settle whether large nonlinear models admit a compact, human-readable account at all. Sufficiency and necessity show the selected units causally affect held-out task margins; they do not explain the whole model or guarantee coverage of rare, safety-relevant cases.
SWD circuits are local and non-unique. Identifiability of a two-factor sparse decomposition needs extra structural conditions, and even then only up to unavoidable scaling and permutation. SWD does not establish those conditions, so bottleneck units depend on the objective, calibration distribution, initialization, and sparsity budget. It intervenes on an approximate replacement AB, not the exact W; matched cross-entropy controls average drift, not agreement on every input. First-order ranking can miss behaviorally important edges, and mean-ablation conclusions shift with the protocol.
Scaling is still open. Qwen3.5-27B covers one matrix, full-model replacement is only on GPT-2 Small, and the semantic audit looks at six units.
One caveat to pull out: sufficiency and necessity are compared inside a controlled protocol where each method's own replacement is matched to similar pre-pruning CE. Fewer edges holds within that protocol; it can change with the graph, threshold, or query/key representation. The targeted edit is compelling but it is a single unit and a single prompt, more a sample of evidence than a statistical result.