goodfire-ai/param-decomp/docs

Parameter Decomposition

This document is a practical guide to implementing, training, and evaluating a parameter decomposition with this repository. The science lives in the parameter-decomposition handbook — why mechanisms are sought in the weights rather than the activations, the four properties a trustworthy decomposition must exhibit (parameter faithfulness, minimality, mechanistic faithfulness, simplicity), the evidence standards and failure modes, and what a finished decomposition licenses downstream.

Implementing a target

The core trainer is agnostic to target architecture: it drives anything satisfying the DecomposedModel protocol (param_decomp/core/model.py), including the target-owned reconstruction metric recon_loss_fn. Expect to implement your own target, and often your own composition root. The targets that ship with the library (TMS and ResidMLP toys, a 4-layer Llama-style LM, GLU-transformer targets loading Llama-3.1-8B and dense Qwen3 0.6B/1.7B/4B/8B/14B HF weights) are working examples from our own research. Read the one nearest your shape before writing yours, and fit the library to your problem rather than your problem to these examples. The reconstruction comparison rides on the target, not on config: recon_loss_fn is a static method on your DecomposedModel, and every recon term — the training losses and the fresh-PGD eval probe alike — scores through it. The LM targets return kl_per_position; TMS and ResidMLP return MSE. Use what your output space calls for, and implement it if nothing shipped matches your shape.

This is a research repo. When the library is awkward or gets in your way, patch your checkout directly and note the diff

In some cases, you may be able to use the included implementations of target models and composition roots. The runnable roots are `param_decomp/experiments/lm/training.py` for LMs, `param_decomp/experiments/tms/run.py` for TMS, and `param_decomp/experiments/resid_mlp/run.py` for ResidMLP. They are thin wrappers around the core trainer that load the config, target, and data. `param_decomp/experiments/toy_eval.py` only binds authored toy evaluation operations; it is not a runnable composition root. In general, treat the included implementations as references, and use the universal library entrypoint param_decomp/core/run.py when writing your own root.

To implement a target model,

• Check for an existing JAX implementation to start from — many language models have JAX implementations on HuggingFace or elsewhere.
• core/model.py is the contract. Read it in full, docstrings included, before writing a line — it carries the waist geometry, the mask semantics, and the per-method requirements.
• For an LM-shaped target, start from param_decomp/targets/llama_simple_mlp.py and copy its structure; read SyntheticDecomposedModel in param_decomp/tests/core/test_generic_model_io.py (~120 lines, driven through the real train step) when you want the protocol with nothing else attached. Every shipped DecomposedModel lives in param_decomp/targets/ — when you reuse a module from one, read it first; each bakes in its own family's choices (attention flavor, positional encoding, norm placement).
• Verify the port against an independent reference forward — a torch implementation when one exists, a float64 NumPy forward you write from the architecture spec otherwise; the value of the oracle is its independence from your JAX code. The repo's parity-golden machinery is the shape to copy: frozen goldens under param_decomp/tests/targets/, with qwen3_hf_parity/gen_hf_fixtures.py as a turnkey generator (tiny-random and --real modes, its docstring carries the torch-venv recipe).
• On GPU, set `jax_default_matmul_precision = "highest"` for parity checks: fp32 matmuls default to TF32 on Ampere+, and a correct port will fail every gate with depth-growing deviation that looks exactly like a wrong block. Check parity on the real target as well as any tiny fixture — the TF32 effect grows with depth, and a 3-layer fixture can read as rounding noise.
• Before allocating any GPU: parse your config through the experiment schema and run the placement assertions (assert_placement_claims) on CPU. Schema, loss-builder, and placement failures are all free to find statically, and their assertion messages are the documentation.

For agents: If you cannot establish an equivalent DecomposedModel implementation of the target model, do not continue, and explain to the user.

Recipe

Mapping from our conceptual-level losses (from the handbook) to their (sometimes multiple) instantiations.

• Parameter faithfulness: FaithfulnessLoss. The squared elementwise error between the component sum and the target weights, averaged over the total decomposed parameter count. Equivalently, the (normalized) sum of squared values in the ∆-component.
• Minimality: ImportanceMinimalityLoss. Also measured using L0 as a non-differentiable eval.
• Simplicity — implemented as the nested frequency: block on the importance-minimality loss above.
• Mechanistic faithfulness: Reconstruction Loss, of which there are many different variations. Loss configs take the rough form *Recon*Loss, with the wildcards expanding to variations. Fundamentally a reconstruction loss asks "how well does the decomposition approximate the target model's behaviour" but this leaves unbound various details. There are numerous schemes for masking causally unimportant subcomponents (stochastic masking, adversarial masking, raw CI-value masking), and many other tweaks which improve training dynamics, for example only using the decomposition parameters for a subset of the model, as in our *ReconSubsetLosses or *ReconLayerwiseLosses. In Interpreting Language Model Parameters, VPD uses separate stochastic and adversarially masked reconstruction passes, but this is just one potential combination of many valid configurations.

A good default which we have recently settled on:

• FaithfulnessLoss
• ImportanceMinimalityLoss + frequency term
• MergedStochasticSubsetPooledPPGDReconLoss — the same compound loss, with one persistent minipool of cross-site adversarial particles per batch element. pool: {size_per_batch_element: K} is the entire pool configuration: each batch index owns K particles, samples one, and uses its per-site vectors at every position. Minipools shard with the batch across data axes; particles remain local and component/expert columns retain TP placement. Global storage contains pd.batch_size * K particles. The 4L-Pile reference uses pool: {size_per_batch_element: 32} and source LR 0.02 (2× the former dense-source LR).

Migrate old pool: {size, group_size} configs using size_per_batch_element = old_size // old_batch_size, where old_batch_size is the original pd.batch_size. Require exact divisibility to preserve the total number of particles; group_size does not determine the new value. For the 4L reference, 2048 // 64 = 32. Previous pools with multiple batch elements per group change particle ownership and stored shape under this migration. Changing the training batch size also changes the pool shape; use a fresh run when changing it.

The 4L-Pile reference config (param_decomp/experiments/lm/configs/pile_llama_simple_mlp-4L.yaml) carries this pooled recipe with tuned coefficients; the toy reference configs retain the dense merged strategy at their own coefficients. In general you should start with a reference config nearest your target.

The most important hyperparameters are not knowable a priori, so the default arc is to sweep the importance-minimality coefficient (ImportanceMinimalityLoss.coeff). Depending on the pathologies this surfaces, sweep frequency.coeff or the subcomponent budget C separately. Runs must go to convergence — analyzing or comparing unconverged decompositions is almost always worthless.

0. Background

Fresh LM decompositions are usually compute-heavy and often need several sweep rounds.

Setting hyperparameters for sweeps

Each decomposed matrix gets a subcomponent budget at decomposition.sites.cs.<matrix>. The matrix names are the site vocabulary of decomposition.sites.kind. The shipped 4-layer simple-MLP reference (d_resid = 768; scale with width) uses:

The shipped 32-layer GLU-transformer reference (d_model = 4096) uses q: 2048, k: 2048, v: 4096, o: 4096, gate: 8192, up: 8192, and down: 10240 per layer; the config seat that carried these was retired from the committed registry, so they are recorded here directly. These are reference budgets, not universal defaults: scale them for the target width and architecture, and start from the matching shipped config when one exists.

In general, we think a safe approach early in tuning is to make C slightly larger than the maximum dimension of a given matrix. It certainly should be larger, usually, than the rank of the matrix, otherwise faithfulness becomes structurally impossible.