Inference and learning in sparse autoencoders as natural gradient flow

1 INTRODUCTION

• BeFOND extends FOND by making inference and dictionary learning natural-gradient flows on shared variational free energy, with closed-form dynamics that avoid backpropagating through inference trajectories.The model is an encoder-free iterative SAE with Bernoulli latents, a factorized prior, and a linear Gaussian decoder.Source: §1 · pp. 2–3
• Fisher preconditioning compensates for infrequent feature activation, which otherwise slows Euclidean dictionary learning, and gives rare and common features the same local learning rate.The paper tests the predicted benefit for rare-feature recovery through theory, synthetic experiments, and optimizer ablations.Source: §1 · pp. 2–3
• Evaluations cover synthetic data with known ground-truth features and language-model representations, comparing BeFOND with amortized SAEs and an iterative mean-field variational baseline.The language-model evaluation examines feature quality, sample efficiency, and dictionary-width scaling.Source: §1 · pp. 2–3
• On Gemma-2-2B layer-12 activations, BeFOND outperforms evaluated baselines on sparse probing and RAVEL disentanglement, with over 1,000× fewer training tokens than Gemma Scope JumpReLU checkpoints for sparse probing.At 524k latents, it matches those checkpoints after ∼6M tokens versus 8B training tokens; its performance also keeps improving with dictionary width while baselines plateau.Source: §1 · p. 3
• Recurrent explaining away creates competition between overlapping features, and ablations test its increasing contribution to feature recovery as superposition grows.The introduction also reports that BeFOND learns more accurate dictionaries and detects rare features with few misses or false positives.Source: §1 · pp. 2–3

2 BACKGROUND AND RELATED WORK

• Classical sparse coding represents each observation as a linear combination of dictionary atoms with sparse coefficients; overcomplete dictionaries have more atoms than input dimensions, so atom overlap is unavoidable.Inference and dictionary learning are coupled because codes depend on the dictionary and dictionary updates depend on inferred codes, making superposition challenging.Source: §2 · p. 3
• Standard SAEs amortize inference with a learned encoder that predicts codes in one pass, whereas iterative sparse coding refines codes separately for each observation.Iterative inference has yielded mixed results on language-model representations.Source: §2 · p. 3
• Variational inference optimizes an approximate posterior, and minimizing free energy reduces its KL divergence to the true posterior when the observation and generative parameters are fixed.Free energy is the negative ELBO, balancing explanation of the observation against departure from the prior.Source: §2 · p. 4
• Natural gradients replace Euclidean parameter-space movement with Fisher-metric updates that account for how parameter changes alter represented distributions.FOND applies this principle to inference; related natural-gradient approaches and approximations are discussed, while BeFOND directly optimizes posterior parameters through recurrent inference.Source: §2 · p. 4

3 FOND UNIFIES INFERENCE AND LEARNING AS NATURAL GRADIENT FLOW

• FOND extends natural-gradient inference on free energy to learning, treating both variational and generative parameters as dynamical variables.The extension applies the same objective to inference parameters u and generative parameters θ.Source: §3 · pp. 4–5

3.1 BeFOND: A BERNOULLI-DISTRIBUTED ITERATIVE SPARSE CODING MODEL

• BeFOND uses a factorized Bernoulli prior and posterior with a linear Gaussian decoder, and represents posterior beliefs with logits.Its latent variables are binary; the generative parameters include the dictionary, covariance, and prior parameters.Source: §3.1 · p. 5
• The linear Gaussian likelihood permits an exact free-energy expression using the mean residual, precision-weighted dictionary Gram matrix, and posterior covariance.The prior KL is weighted by β, with β = 1 recovering the usual variational free energy.Source: §3.1 · p. 5
• BeFOND derives closed-form inference and learning dynamics using Fisher metrics for the variational posterior and generative joint distribution.Inference includes recurrent interactions; learning updates the dictionary, noise covariance, and prior mean on their respective timescales.Source: §3.1 · p. 5
• The implementation averages minibatch statistics before learning, keeps the prior fixed during inference, and discretizes both flows with exponential integration.A rolling-reference variant is used; the discretization preserves linear relaxation within each step.Source: §3.1 · p. 6

4.1 INFERENCE: BeFOND’S RECURRENT TERM ENABLES EXPLAINING AWAY

• For generic dictionaries with overlapping, nonparallel atoms, no single fixed input projection followed by a scalar nonlinearity can compute exact posterior probabilities for every feature.The result covers several single-projection SAE encoders, while TopK SAEs fall outside this class.Source: §4.1 · p. 6
• BeFOND’s recurrent updates implement explaining away: evidence for one overlapping feature suppresses another feature’s activation drive.For example, roof evidence can strengthen the rain explanation for wet grass and thereby suppress the sprinkler.Source: §4.1 · p. 6
• An approximate-inference result gives sufficient conditions for BeFOND to achieve lower expected posterior KL than the best possible single-projection amortized encoder.This is the theoretical counterpart to the empirical investigation of explaining away.Source: §4.1 · p. 6
• Removing lateral interactions reduces recovery as feature overlap increases; restoring suppressive interactions improves recovery, unlike excitatory-only inference.Excitatory-only inference activates nearly the entire dictionary without comparable recovery.Source: §4.1 · p. 7
• At K = 4,096, BeFOND maintains the highest MCC throughout the superposition sweep, with no dead latents.MCC measures alignment between learned and ground-truth feature directions; recovery declines only modestly as superposition increases.Source: Figure 3 · p. 7, §4.1 · p. 7

4.2 LEARNING: BeFOND’S NATURAL GRADIENT FLOW IMPROVES RARE FEATURE LEARNING

• Under matched-moment assumptions, natural flow gives decoder directions the same learning rate, while Euclidean flow slows in directions associated with rarer features.For frequencies ρ_i = c i^-α, the slowest Euclidean direction takes at least K^α times as long as the fastest for the same fractional error reduction.Source: §4.2 · p. 8
• BeFOND leads all metrics over evaluated baselines in full-width recovery, including detection of the 25% least frequent generating features.The setting provides one learned latent per generating feature, so complete recovery requires representing rare features alongside common ones.Source: §4.2 · p. 8
• The natural-inference/natural-learning recipe reaches stronger recovery with fewer decoder updates than separately tuned Euclidean variants.The comparison is reported in the geometry ablation.Source: §4.2 · p. 8

Evaluation setup

• The evaluation tests feature quality under a 50M-token budget using sparse probing and RAVEL, which measures selective intervention on represented attributes.Sparse probing restricts logistic probes to k = 1 or k = 5 selected features; RAVEL disentanglement averages cause and isolation scores.Source: §5 · pp. 8–9

Scaling with dictionary width

• At 524k latents, BeFOND reaches 0.827 ± 0.001 RAVEL disentanglement versus 0.748 ± 0.004 for BatchTopK, a 7.9 percentage-point advantage.The advantage reflects both cause and isolation scores.Source: §5 · p. 9
• At 524k, BeFOND’s single-feature accuracy is 0.789 ± 0.002 versus 0.774 ± 0.002 for Matryoshka BatchTopK, a 1.5-point improvement.From 16k to 524k, BeFOND’s single-feature accuracy rises from 0.771 to 0.789 and RAVEL from 0.786 to 0.827; each amortized baseline gains about 0.6 points in single-feature accuracy.Source: §5 · p. 9
• Differences between methods are smaller in five-feature probing accuracy, and AUC and F1 rankings vary across dictionary widths.Source: §5 · p. 9

High feature quality emerges early in training

• At 524k latents, BeFOND reaches the evaluated Gemma Scope single-feature probing score after ∼6M training tokens, compared with 8B tokens for those checkpoints.By 50M tokens, BeFOND exceeds their single-feature accuracy (0.789 versus 0.7735) and RAVEL disentanglement (0.827 versus 0.7555), using 160× fewer tokens.Source: §5 · p. 10

Comparison with iterative mean-field inference

• At 524k, the mean-field baseline scores 0.751 ± 0.007 in single-feature accuracy and 0.685 ± 0.029 in RAVEL disentanglement, significantly underperforming BeFOND on both.The baseline uses iterative variational inference and Adam decoder learning.Source: §5 · p. 10

6 CONCLUSION AND DISCUSSION

• Inference and dictionary learning are formulated as natural-gradient flows on shared variational free energy; Fisher preconditioning compensates for infrequent activation under frozen, matched moments.BeFOND improves synthetic feature recovery and supports increasing feature quality at larger dictionary widths on Gemma.Source: §6 · p. 10

Limitations

• Repeated inference steps and numerical decoder updates increase BeFOND’s computational cost.The paper identifies improving efficiency and speeding inference convergence as future work.Source: §6 · p. 10

6 CONCLUSION AND DISCUSSION

• BeFOND improves concept detection and selective intervention on Gemma-2-2B layer-12 activations, with gains increasing as dictionary width grows while evaluated baselines largely plateau.The paper connects these results with recurrent explaining away and improved rare-feature recovery.Source: §6 · p. 10
• The discussion frames adaptation dynamics as needing to respect the underlying model’s intrinsic geometry and symmetries.Source: §6 · p. 10

CODE AND DATA

• Code, data, and model checkpoints are stated to be available, and generative AI tools assisted with derivations, implementation, analysis, literature searches, and manuscript work.The authors state that they are responsible for the final text, results, and claims.Source: CODE AND DATA · p. 10, AI USE STATEMENT · p. 10

A.1 SUPERPOSITION AND INFERENCE

• With an overcomplete dictionary, projecting an input onto an atom mixes that feature’s activity with contributions from overlapping active features, making projection scores unreliable.Explaining away motivates reducing support for a candidate feature when other features already explain the evidence.Source: §A.1 · p. 18

A.2 ITERATIVE SPARSE ENCODERS.

• Iterative sparse encoders refine codes, but prior results show that improved reconstruction or latent utilization need not improve interpretability metrics.Reported effects vary by task and method; fixed-dictionary inference can remain limited by learned feature directions.Source: §A.2 · p. 18

A.3 VARIATIONAL INFERENCE AND THE FREE ENERGY OBJECTIVE

• Variational free energy decomposes into reconstruction and prior-KL terms, and minimizing it improves inference by reducing the KL to the true posterior.A restricted variational family may not represent all posterior dependencies, so optimization can reduce but not eliminate the approximation gap.Source: §A.3 · pp. 19–20

A.4 TRUST REGIONS AND NATURAL-GRADIENT INFERENCE

• A trust-region step uses the full movement budget and matches the fixed-step update when ϵ = η^2∥g_t∥_2^2.A fixed step size lets larger gradients spend more of the movement budget; a fixed ϵ enforces the same radius at each step.Source: §A.4 · p. 21
• Using the Fisher metric makes the local movement cost proportional to KL divergence and yields natural-gradient descent when the Fisher is nonsingular.The local Fisher metric is the quadratic approximation to KL between successive approximate posteriors.Source: §A.4 · pp. 21–22
• Free energy strictly decreases along the continuous flow wherever its gradient is nonzero, but an arbitrary finite step need not decrease it.A large discretized jump can leave the region where the local prediction is accurate.Source: §A.4 · p. 22
• In the Gaussian mean example, natural-gradient steps contract error equally in every direction and reach the target mean in one step when η = 1.The corresponding Euclidean update is constrained by the narrowest direction, requiring 0 < η < 2λ_min(S).Source: §A.4 · p. 22
• Natural-gradient flow describes the same distribution evolution under smooth coordinate changes, but finite Euler steps preserve this exactly only for affine changes.For nonlinear changes, finite Euler steps preserve invariance only to first order.Source: §A.4 · pp. 22–23

A.5 “PHYSICS OF ADAPTATION”: DERIVING COVARIANT INFERENCE AND LEARNING RULES

• The framework treats variational free energy as a shared potential, with inference changing beliefs and learning changing generative parameters while holding the other block fixed.Choosing the kinetic term and dissipation force determines the resulting dynamics and, after discretization, the optimization algorithm.Source: §A.5 · p. 24
• A coordinate-invariant kinetic energy follows from the Riemannian metric on belief space and from local information gain, whose quadratic form is measured by the Fisher metric.The resulting motion tracks changes in distributions rather than parameter rescalings; its biological metabolic interpretation is speculative.Source: §A.5 · pp. 25–26
• For exponential-family variational beliefs in natural parameters, the overdamped natural-gradient update recovers the Bayesian learning rule of Khan & Rue (2023).The framework also motivates exploring dynamics that retain inertia, use other dissipation laws, or adapt the belief geometry over time.Source: §A.5 · p. 28
• In the overdamped limit, dissipative covariant belief dynamics reduce to natural-gradient flows for both inference and learning, each using its own Fisher metric.The full dynamics include inertia and metric-dependent geometric corrections; these vanish in the massless limit while the geometry remains.Source: §A.5 · p. 28

B.1 MODEL AND OBJECTIVE

• BeFOND extends FOND’s recurrent natural-gradient inference by applying natural-gradient learning to the generative distribution, using each distribution’s Fisher metric.In the Bernoulli instantiation, binary latents indicate feature presence, decoder atoms specify feature contributions, and posterior means express confidence in feature presence.Source: §B.1 · p. 29

A.5 “PHYSICS OF ADAPTATION”: DERIVING COVARIANT INFERENCE AND LEARNING RULES

• For the latent families described, analytic moments and Fisher metrics make posterior variance part of both expected reconstruction loss and inference geometry.Source: §B.1 · p. 29
• The Bernoulli free energy includes expected reconstruction uncertainty and a KL penalty; its uncertainty term cancels self-interaction, while overlapping atoms still couple coordinates.Source: §B.1 · pp. 29–30
• For Bernoulli inference, natural-gradient preconditioning cancels probability-variance attenuation in logit gradients, leaving input evidence, recurrent dictionary interactions, and relaxation toward the reference.Source: §B.1 · p. 30
• The decoder update matches observed input–latent moments to model predictions and preconditions them with the inverse prior second moment, potentially compensating for weak learning signals in rare directions.Exact compensation requires matching minibatch-averaged second moments.Source: §B.1 · p. 30
• The covariance update tracks expected squared residuals, and the learned initial prior tracks detached posterior means across inputs; these updates use inferred statistics without backpropagating through inference.Source: §B.1 · p. 30

B.1.1 ONLINE PRIOR AND WORKING MEMORY

• For an exactly minimized repeated-input rollout, each observation contributes another likelihood factor while the initial prior enters once, so its influence wanes and its activity cap constrains the starting belief, not final activity.Replaying one stored observation is not ordinary conditioning on a new measurement; finite mean-field ETD1 steps may not reach the unrestricted endpoint.Source: §B.1.1 · p. 32

B.2 NATURAL-GRADIENT DYNAMICS

• Posterior-Fisher preconditioning removes the r_i(1 − r_i) shrinkage from logit updates, while decoder learning uses the complete-data model Fisher and distinguishes metric moments E_p from reconstruction curvature E.The passage emphasizes that matching average firing rates alone does not match second moments when inferred means covary across inputs.Source: §B.2 · p. 33

B.3 DECODER BIAS

• With a fixed posterior, the learned decoder bias fits the mean residual; differentiating through one unclamped ETD1 step adds a response term.The stated expression assumes one shared covariance across the batch.Source: §B.3 · p. 34

B.4 BERNOULLI UNCERTAINTY AND SPARSE CODING

• Bernoulli uncertainty cancels the quadratic self-term in expected reconstruction loss, leaving an affine reconstruction objective in each coordinate and an interior optimum at finite evidence.The KL term is strictly convex for β > 0 when the other coordinates are fixed.Source: §B.4 · p. 34

B.5 ACTIVATION THRESHOLD UNDER A ROLLING PRIOR

• With a rolling reference below 1/2, the activation threshold requires positive reconstruction evidence after accounting for other features, and updating the reference changes that threshold.The prior cap and inference depth affect posterior sparsity through different parts of the rollout.Source: §B.5 · p. 34

B.6 NUMERICAL INTEGRATION AND MATRIX-FREE UPDATES

• ETD1 integrates reference relaxation exactly with gain bounded by 1/β, while the algorithm alternates rolling inference with moment-based learning and supports matrix-free decoder updates.When E = E_p, each decoder mode contracts by the same factor, recovering exact rarity compensation; evaluation resets the initial prior for each input.Source: §B.6 · p. 35, §B.7 · pp. 35–36

B.8 COMPUTATIONAL COMPLEXITY

• 12–14×: BeFOND takes approximately this much longer than BatchTopK on the measured training prefix; its inference cost is O(TBdK), with additional decoder-learning and prior-fit costs.BeFOND applies the lateral operator without forming a dense K × K Gram matrix.Source: §B.8 · p. 36

C.1.1 DATASETS

• Synthetic evaluation uses unwhitened SynthSAEBench data with Zipf-distributed feature probabilities and measures dictionary and support recovery, rare-feature F1, dead latents, and mean L0.The evaluation passage specifies test and frequency-reference sample sizes and inference steps; the implementation passage notes K1 uses a fixed horizon of one, while prior fitting and test inference may use longer horizons.Source: §C.1.1 · p. 37, §C.1.3 · pp. 37–38, §C.1.4 · p. 38

C.2.1 SINGLE-LAYER INFERENCE

• For overlapping, linearly independent atoms, no scalar function of a single projection can equal the exact posterior; this holds for almost every dictionary when M ≥ 2.The generic-dictionary result allows exceptions only for atoms orthogonal to every other atom when no pair is parallel.Source: §C.2.1 · pp. 38–39

C.2.2 ITERATIVE INFERENCE VERSUS SINGLE PROJECTION ENCODERS

• Under the stated conditions, the optimal factorized posterior is not a single-projection function, leaving every such encoder with a positive expected-divergence gap.The result assumes λmin(J) > −4 and a pair of atoms that are neither orthogonal nor parallel.Source: §C.2.2 · p. 39, Then the optimal factorized posterior · p. 40

Then the optimal factorized posterior

• When λmin(J) > −4, the mean-field objective is strongly convex and has a single optimum that BeFOND approaches; for stronger coupling, it can instead have several local minima, one per competing explanation.Source: Then the optimal factorized posterior · p. 41

C.2.3 RAIN–SPRINKLER EXAMPLE

• 0.008508 nats for converged fixed-prior BeFOND versus 0.021554 nats for the fitted single-projection encoder in average KL divergence on the rain–sprinkler example.The example uses independent rain and sprinkler latents with prior probability 0.1 each.Source: §C.2.3 · p. 41

C.3.1 INFERENCE ABLATION: LATERAL INTERACTIONS

• Lateral ablations match learning rates, β, and decoder norms, while excitatory-only runs at d = 1,024 and 1,536 stopped early.The ablation control uses different learning-rate and β settings from the selected baseline-comparison recipe at d = 512, 384, and 256.Source: §C.3.1 · p. 42
• At d = 256, suppressive interactions raise MCC from 0.2810 to 0.6862 within matched ablation recipes.Figure 6 uses K = 4,096; hollow markers denote early-stopped excitatory runs.Source: §C.3.1 · p. 42, Figure 6 · p. 42
• Excitatory-only inference activates almost all 4,096 latents but reaches MCC 0.2825.Source: §C.3.1 · p. 42
• Suppressive-only inference has one dead latent at d = 1,536 and none at the other tested dimensions; full and zero-off-diagonal inference have zero throughout.Source: §C.3.1 · p. 42

C.3.2 LEARNING ABLATION: INFERENCE AND LEARNING GEOMETRY

• MCC 0.500 and Rare F1 0.280 after 1,000 updates exceed both Euclidean variants after 100,000 updates.The geometry recipes were tuned separately.Source: §C.3.2 · pp. 42, 44, Figure 7 · p. 44

C.3.3 ROLLING PRIOR

• Increasing β reduces final activity in the 72-fit sweep, but larger values also reduce recovery.Several looser initial-prior caps yield overlapping curves because the learned prior falls below those caps; the cap constrains starting activity while rolling-reference dynamics determine later activity.Source: §C.3.3 · p. 44, Figure 8 · p. 44

C.4 EXACT AND APPROXIMATE RARITY COMPENSATION

• For power-law feature frequencies, the slowest-to-fastest Euclidean relaxation-time ratio is at least K^α, versus one for matched natural-gradient flow.With inferred codes fixed and averaged second moments matched, natural-gradient flow has the same local contraction rate in every decoder direction.Source: §C.4 · p. 46

C.4.1 ACTIVITY-NORMALIZED INFERENCE

• Natural-gradient inference removes the Bernoulli variance factor from membrane-potential updates, including an extra suppression of weakly active latents when activation is favored.Source: §C.4.1 · p. 47

C.4.2 LOCAL CONDITIONING UNDER ADAM NORMALIZATION

• Under the local Adam-normalization analysis, power-law frequencies give κ(Aϵ) ≥ K^α/2; the calculation assumes known codes and a frozen expected second-moment estimate at the decoder optimum.Source: §C.4.2 · pp. 47–48

• BeFOND retains the highest MCC, Macro F1, and Rare F1 after modifications that substantially improve baseline recovery.The comparison is between complete recipes with different target sparsities.Source: Figure 9 · p. 48, §C.5.1 · p. 48
• BeFOND has higher coverage and F1 in every frequency quintile, including the rare tail.Coverage measures dictionary alignment, while F1 measures detection of the matched feature’s activity; they need not coincide.Source: Figure 10 · p. 50, §C.5.1 · p. 50
• The three-seed full-width mean is MCC 0.8980 and Rare F1 0.9519, with 0.33 dead latents on average.The seed-0 run reaches MCC 0.8986 and Rare F1 0.9554 with one dead latent; a separate continuous run has zero dead latents but lower recovery.Source: §C.5.1 · p. 50

C.5.2 TRAINING SAMPLE EFFICIENCY

• Figure 11 tracks MCC, Macro F1, and Rare F1 through training, with BeFOND on one continuous run and baseline checkpoints drawn from separate matching-recipe runs.The figure shows one seed and estimates no uncertainty intervals; its x-axis excludes BeFOND’s separate prior-update examples.Source: §C.5.2 · p. 50, Figure 11 · p. 51
• Sparse probing peaks within the first few inference steps, whereas RAVEL generally benefits from longer inference.Figure 12 evaluates 1–40 steps at widths from 16k to 524k using checkpoint 12,207, corresponding to 49,999,872 training tokens.Source: Figure 12 · p. 51

D.1 EXPERIMENTAL SETUP

• Gemma-2-2B layer-12 evaluations use selected BeFOND results at 49,999,872 decoder-training tokens and report sparse probing across 113 tasks alongside RAVEL metrics.The activations are evaluated in raw and whitened forms, with disjoint document ranges for training, validation, and calibration.Source: §D.1 · pp. 52–53

D.2 RESULTS

• The whitened baseline comparison reports sparse probing and RAVEL in Table 8, while Table 9 reports released Gemma Scope checkpoint evaluations.Table 8 summarizes five-seed baseline results at 50M training tokens; the supplied passages do not state comparative winners.Source: §D.2 · p. 53, Table 8 · p. 54

• The timed interval includes inference, updates, required communication, and BeFOND prior updates, but excludes evaluation, initialization, warmup, data preparation, checkpointing, and logging.Source: §D.3 · p. 54, Table 10 · p. 55
• BeFOND takes approximately 12–14× as long as BatchTopK across the measured widths.Table 10 covers 999,424 training tokens on eight NVIDIA H200 GPUs.Source: §D.3 · p. 54, Table 10 · p. 55