Paper deep dive
Inference and learning in sparse autoencoders as natural gradient flow
Hadi Vafaii, Tejas Rao, David Chanin, Thomas Fel, Jacob L. Yates, Bruno Olshausen, David Klindt, Dileep George, Miguel Lázaro-Gredilla
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 90%
Last extracted: 10/7/2026, 9:17:05 AM
Summary
The paper introduces BeFOND, an encoder-free sparse autoencoder (SAE) model that unifies inference and dictionary learning as natural-gradient flows on a shared variational free energy. BeFOND uses recurrent explaining away to reduce interference between overlapping features and Fisher preconditioning to handle rare features. It demonstrates superior dictionary recovery, rare-feature detection, and sample efficiency compared to amortized SAE baselines on both synthetic data and language model activations.
Entities (10)
Relation Signals (8)
BeFOND → isa → Sparse Autoencoder
confidence 95% · We instantiate this framework as BeFOND, an encoder-free sparse coding model
FOND → extendedto → BeFOND
confidence 92% · We instantiate this extended FOND framework as BeFOND
BeFOND → evaluatedon → Gemma-2-2B
confidence 90% · On Gemma-2-2B layer-12 activations, BeFOND outperforms all evaluated baselines
BeFOND → uses → Natural Gradient Flow
confidence 90% · unify inference and dictionary learning as natural-gradient flows
BeFOND → uses → Explaining Away
confidence 88% · recurrent explaining away reduces interference between overlapping features
BeFOND → uses → Fisher Preconditioning
confidence 88% · Fisher preconditioning can compensate for the slow learning of rare features
BeFOND → outperforms → BatchTopK
confidence 85% · BeFOND improves... with a growing advantage over amortized baselines
BeFOND → outperforms → JumpReLU
confidence 85% · outperforming pretrained reference SAEs... JumpReLU
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Sparse autoencoders are widely used to uncover interpretable features in neural networks, yet reliable recovery remains difficult when features overlap or activate infrequently. These challenges involve both inferring which features explain an input and learning the dictionary that represents them. Here, we unify inference and dictionary learning as natural-gradient flows on a shared variational free energy. We instantiate this framework as BeFOND, an encoder-free sparse coding model with closed-form inference and learning dynamics. We show how recurrent explaining away reduces interference between overlapping features, while Fisher preconditioning can compensate for the slow learning of rare features. On synthetic data, BeFOND improves dictionary recovery and rare-feature detection, with a growing advantage over amortized baselines as superposition increases. On language-model activations, it improves single-feature concept detection and selective intervention, outperforming pretrained reference SAEs with substantially less training data. Its feature quality continues to improve with dictionary width, whereas the evaluated baselines largely plateau. Together, these results show how improving inference and learning within a unified probabilistic framework can make better use of data and dictionary capacity to interpret and intervene on neural representations.
Tags
Links
- Source: https://arxiv.org/abs/2610.07389v1
- Canonical: https://arxiv.org/abs/2610.07389v1
Trouble viewing inline? Open PDF directly →
Full Text
283,057 characters extracted from source content.
Expand or collapse full text
Inference and learning in sparse autoencoders as natural gradient flow Hadi Vafaii Tejas Rao David Chanin Thomas Fel Jacob L. Yates Affiliation: Astera Institute Affiliation: UC Berkeley Affiliation: UCL Affiliation: Goodfire AI Bruno Olshausen David Klindt Dileep George Miguel Lázaro-Gredilla Affiliation: Astera Institute Affiliation: UC Berkeley Affiliation: CSHL*Co-first authors, correspondence to: hadi@astera.org tejasrao@berkeley.edu Abstract Sparse autoencoders are widely used to uncover interpretable features in neural networks, yet reliable recovery remains difficult when features overlap or activate infrequently. These challenges involve both inferring which features explain an input and learning the dictionary that represents them. Here, we unify inference and dictionary learning as natural-gradient flows on a shared variational free energy. We instantiate this framework as BeBeFOND, an encoder-free sparse coding model with closed-form inference and learning dynamics. We show how recurrent explaining away reduces interference between overlapping features, while Fisher preconditioning can compensate for the slow learning of rare features. On synthetic data, BeBeFOND improves dictionary recovery and rare-feature detection, with a growing advantage over amortized baselines as superposition increases. On language-model activations, it improves single-feature concept detection and selective intervention, outperforming pretrained reference SAEs with substantially less training data. Its feature quality continues to improve with dictionary width, whereas the evaluated baselines largely plateau. Together, these results show how improving inference and learning within a unified probabilistic framework can make better use of data and dictionary capacity to interpret and intervene on neural representations. 1 Introduction Sparse autoencoders (SAEs) infer sparse representations of data in terms of a learned, overcomplete dictionary of features. Their roots lie in sparse coding (Olshausen and Field, 1996; Lewicki and Sejnowski, 2000), motivated by the organization of sensory representations in biological systems (Olshausen and Field, 1997; Olshausen and Field, 2004). SAEs have since become central in interpreting neural network representations (Klindt et al., 2026, reviewed in), especially those of language models (Yun et al., 2021; Bricken et al., 2023; Cunningham et al., 2024; Templeton et al., 2024; Gao et al., 2025). Despite these successes, SAEs face challenges in both inference and dictionary learning. Standard SAEs use a trained neural network to perform one-pass inference—known as amortized inference—which can struggle to disentangle overlapping features in superposition (O’Neill et al., 2025). But improving inference alone is insufficient. In controlled experiments, Barin-Pacela et al. (2026) identified dictionary learning as the bottleneck: replacing an SAE’s encoder with iterative inference on the same learned dictionary failed to close the feature-recovery gap. These problems are coupled during training: the dictionary determines which explanations of the data are available, while inferred feature activations shape the learning signal used to update the dictionary. Therefore, inference errors can compromise dictionary learning, and a poor dictionary can limit what inference can recover. Unlike standard SAEs, classical sparse coding methods typically perform inference through iterative, per-sample optimization (Rozell et al., 2008; Gregor and LeCun, 2010). More recently in this lineage, Vafaii et al. (2025) introduced the Free-energy Online Natural-gradient Dynamics (FOND) framework, which formulates inference as natural-gradient flow on variational free energy. FOND’s guiding principle is to respect the intrinsic geometry of the statistical manifold: updates are measured by their effect on the represented probability distribution through the Fisher information metric (Amari, 1998; Martens, 2020; Khan and Rue, 2023). However, FOND did not address learning through the same geometric principle. Its generative model parameters were learned by backpropagating through the inference trajectory, with the associated computational and memory costs. Here, we extend this principle to learning. Just as inference follows natural-gradient flow on free energy with respect to variational parameters, we define learning as natural-gradient flow on the same objective with respect to generative parameters. Both processes thus respect the statistical geometry of the distributions they update. We instantiate this extended FOND framework as BeBeFOND: an encoder-free iterative SAE with Bernoulli latents, a factorized prior, and a linear Gaussian decoder. We derive closed-form expressions for its inference and learning dynamics, allowing the dictionary to be learned without backpropagating through the inference trajectory. Figure 1: Feature quality and sample efficiency on Gemma-2-2B activations. a,c: Training dynamics at 524524k width. b,d: Final performance across dictionary widths. Details in section 5. On the inference side, BeBeFOND implements explaining away through recurrent interactions between overlapping features, enabling lateral competition leading to sparse, disentangled representations. We isolate the contribution of these interactions and show that they become increasingly important as feature superposition increases. On the learning side, we show how infrequent activation slows Euclidean dictionary learning. We prove that Fisher preconditioning removes this frequency dependence, giving rare and common features the same local learning rate. Motivated by this result, we hypothesize that Fisher preconditioning would improve rare-feature recovery in BeBeFOND. Empirically, BeBeFOND learns more accurate dictionaries and detects rare features with few misses or false positives, while using substantially fewer training samples than SAE baselines and Euclidean ablations. We evaluate BeBeFOND on synthetic data with known ground-truth features using SynthSAEBench (Chanin and Garriga-Alonso, 2026), and on language-model representations using SAEBench (Karvonen et al., 2025). Baselines include amortized SAEs—BatchTopK (Bussmann et al., 2024), Matryoshka BatchTopK (Bussmann et al., 2025), and JumpReLU (Rajamanoharan et al., 2024b)—and an iterative mean-field variational baseline with Adam decoder learning (Kingma and Ba, 2015). On Gemma-2-2B layer-12 activations, BeBeFOND outperforms all evaluated baselines on both sparse probing (Kantamneni et al., 2025) and RAVEL disentanglement (Huang et al., 2024). BeBeFOND’s performance continues to improve as we increase dictionary width from 1616k to 524524k latents, whereas the baselines plateau (Fig. 1b,d). These gains also emerge with substantially less training data: at a dictionary width of 524524k, BeBeFOND matches the evaluated Gemma Scope JumpReLU checkpoints (Lieberum et al., 2024) in sparse probing after just ∼ 6M training tokens, compared with the 8B tokens used to train those checkpoints, achieving over 1,000× reduction in training tokens (Fig. 1a). Summary of our contributions and paper outline: (1) We review necessary background (section 2). We then extend FOND (Vafaii et al., 2025) to unify inference and learning as natural-gradient flows on variational free energy, instantiate it as Bernoulli-distributed BeBeFOND, and derive its dynamics in closed-form (section 3 and appendix B). (2) We analyze BeBeFOND’s recurrent “explaining away” (section 4.1) and use ablations to establish its contribution to feature recovery under increasing superposition (section C.3.1). (3) We characterize how Fisher preconditioning in BeBeFOND’s learning rule compensates for feature rarity (section 4.2), and test its consequences for dictionary learning through theoretical analysis (theorem 4.2), synthetic experiments, and optimizer ablations (sections C.4, C.3.2 and C.5). (4) We evaluate BeBeFOND on language-model representations, examining feature quality, sample efficiency, and scaling with dictionary width (section 5). 2 Background and Related Work Notation. Bold lowercase and uppercase symbols denote vectors and matrices, such as a and A. Inference and learning quantities appear in red and blue, respectively. Sparse coding. Classical sparse coding (Olshausen and Field, 1996; Olshausen and Field, 1997; Lewicki and Sejnowski, 2000; Aharon et al., 2006; Rozell et al., 2008; Gregor and LeCun, 2010; Mairal et al., 2010) models each observation ∈ℝM x ^M as a linear combination of dictionary atoms with sparse coefficients [rgb]0.9297,0.1094,0.1406 z: =+ϵ,ϵ∼(,M), x\;=\; [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 z+ ε~, ε ( 0,I_M)~, (1) The dictionary ∈ℝM×K [rgb]0.1602,0.6719,0.8867 ^M× K has atoms ⋅i≡ϕ(i) [rgb]0.1602,0.6719,0.8867 _· i≡ [rgb]0.1602,0.6719,0.8867 φ^(i) as its columns. We consider overcomplete dictionaries, whose atoms outnumber the input dimensions (K>MK>M), and therefore cannot all be mutually orthogonal, making some overlap between atoms unavoidable. Learning estimates the dictionary from data, while inference finds sparse codes that explain each input. These processes are coupled because inference uses the current dictionary, and dictionary updates depend on the inferred codes. In the overcomplete setting, superposition—overlap between atoms—makes sparse code inference particularly challenging (see section A.1). Amortized versus iterative inference. Standard SAEs learn an encoder =fϕ() [rgb]0.9297,0.1094,0.1406 z=f_ [rgb]0.9297,0.1094,0.1406 φ( x) that predicts a sparse code in one forward pass. These include ReLU (Bricken et al., 2023; Cunningham et al., 2024; Templeton et al., 2024), JumpReLU (Rajamanoharan et al., 2024b), gated (Rajamanoharan et al., 2024a), TopK (Makhzani and Frey, 2014; Gao et al., 2025), BatchTopK (Bussmann et al., 2024), and Matryoshka BatchTopK (Bussmann et al., 2025) SAEs. Training this mapping across samples amortizes inference, replacing a separate test-time optimization for each input. Iterative sparse coding instead refines the code for each observation (Rozell et al., 2008; Beck and Teboulle, 2009; Fang et al., 2022). In continuous time, this takes the schematic form, ˙=g(,,) [rgb]0.9297,0.1094,0.1406 z=g( [rgb]0.9297,0.1094,0.1406 z, [rgb]0.1602,0.6719,0.8867 , x). Iterative inference has yielded mixed results on language-model representations (Chanin, 2025; Sun et al., 2026). See section A.2 for further discussion. Variational inference and free energy. A generative model combines a prior with a likelihood: p(,)⏟joint distribution=p(∣)⏟conditional likelihoodp()⏟prior,e.g., p(∣)=(,,M)⏟probabilistic form of eq. 1. [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0p_ [rgb]0.1602,0.6719,0.8867 θ( x, z)_ [rgb]0,0,0joint distribution\;= [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0p_ [rgb]0.1602,0.6719,0.8867 θ( x z)_ [rgb]0,0,0conditional likelihood\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0p_ [rgb]0.1602,0.6719,0.8867 θ( z)_ [rgb]0,0,0prior~, .g., \;\;p_ [rgb]0.1602,0.6719,0.8867 θ( x z)\;\;= [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0N( x; [rgb]0.1602,0.6719,0.8867 z,I_M)_ [rgb]0,0,0probabilistic form of @crefcreftype~refnumeq:linear-generative-model. (2) Bayes’ rule gives p(∣)=p(,)/p()p_ [rgb]0.1602,0.6719,0.8867 θ( z x)=p_ [rgb]0.1602,0.6719,0.8867 θ( x, z)/p_ [rgb]0.1602,0.6719,0.8867 θ( x). Computing the evidence p()p_ [rgb]0.1602,0.6719,0.8867 θ( x) requires summing or integrating over latent configurations, which is generally intractable for large overcomplete models. Variational inference (Jordan et al., 1999; Blei et al., 2017, section A.3;) turns this problem into optimization: choose an approximate posterior q(∣)q_ [rgb]0.9297,0.1094,0.1406 u( z x) and optimize its parameters [rgb]0.9297,0.1094,0.1406 u. For example, [rgb]0.9297,0.1094,0.1406 u can be Gaussian means and variances, or Bernoulli logits. Define the variational free energy: ℱ(;,)≔q[−logp(,)q(∣)]=−q[logp(∣)]+(q(∣)∥p()).F( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)\; \;E_q_ [rgb]0.9297,0.1094,0.1406 u[- p_ [rgb]0.1602,0.6719,0.8867 θ( x, z)q_ [rgb]0.9297,0.1094,0.1406 u( z x)]\;=\;-E_q_ [rgb]0.9297,0.1094,0.1406 u\! [ p_ [rgb]0.1602,0.6719,0.8867 θ( x z) ]+D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z) )~. (3) Free energy is the same quantity as the negative evidence lower bound, ℱ=−ELBOF=-ELBO. Its terms balance explaining the observation against departing from the prior. The connection between free-energy minimization and posterior inference follows from the folloing identity: ℱ(,,)=(q(∣)∥p(∣))−logp().F( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)\;=\;D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z x) )\;-\; p_ [rgb]0.1602,0.6719,0.8867 θ( x)~. (4) This shows that, with x and [rgb]0.1602,0.6719,0.8867 θ fixed, reducing ℱF reduces the KL divergence to the true posterior because the evidence is independent of [rgb]0.9297,0.1094,0.1406 u (section A.3). An amortized encoder can predict =fϕ() [rgb]0.9297,0.1094,0.1406 u=f_ [rgb]0.9297,0.1094,0.1406φ( x), as in VAEs (Kingma and Welling, 2014), or inference can optimize it separately for each input: ∗()=argminℱ(,,). [rgb]0.9297,0.1094,0.1406 u^*( x)\;=\; [rgb]0.9297,0.1094,0.1406 uargmin\ F( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)~. (5) Holding x and [rgb]0.1602,0.6719,0.8867 θ fixed, write the objective as ℱ()F( [rgb]0.9297,0.1094,0.1406 u). Starting from some initialization 0 [rgb]0.9297,0.1094,0.1406 u_0, gradient descent with step size η>0η>0 gives: t+1=t−η∇ℱ(t). [rgb]0.9297,0.1094,0.1406 u_t+1\;=\; [rgb]0.9297,0.1094,0.1406 u_t\;-\;η\, _ [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u_t)~. (6) This update measures movement using Euclidean distance in parameter space (section A.4). However, equal-sized parameter changes can produce different changes in beliefs q_ [rgb]0.9297,0.1094,0.1406 u, and reparameterizing the same distributions can change the descent trajectory. Measuring movement between distributions instead leads to the geometry of their statistical manifold (Amari and Nagaoka, 2000). From gradient descent to natural gradients. Swapping the Euclidean geometry of gradient descent with the Fisher metric gives the natural-gradient update: t+1=t−ηu−1(t)∇ℱ(t). $ [rgb]0.9297,0.1094,0.1406 u_t+1\;=\; [rgb]0.9297,0.1094,0.1406 u_t\;-\;η\,G_u^-1( [rgb]0.9297,0.1094,0.1406 u_t) _ [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u_t)~.$ (7) where uG_u is the Fisher information matrix of the variational family: u()=∼q(⋅∣)[(∇logq(∣))(∇logq(∣))⊤].G_u( [rgb]0.9297,0.1094,0.1406 u)\;=\;E_ z q_ [rgb]0.9297,0.1094,0.1406 u(· x) [ ( _ [rgb]0.9297,0.1094,0.1406 u q_ [rgb]0.9297,0.1094,0.1406 u( z x) ) ( _ [rgb]0.9297,0.1094,0.1406 u q_ [rgb]0.9297,0.1094,0.1406 u( z x) ) ]~. (8) The Fisher matrix measures local statistical distinguishability, so the update accounts for how parameter changes affect the represented distribution (Amari, 1998; Amari and Nagaoka, 2000; Amari, 2016). Natural gradients have a long history in variational inference, underpinning stochastic variational inference (Hoffman et al., 2013), the Bayesian learning rule (Khan and Rue, 2023), and Bayesian Online Natural Gradient (Jones et al., 2024). Closed-form natural gradients need not be available, and computing the Fisher and applying its inverse can be expensive, motivating approximations (Martens and Grosse, 2015; Martens, 2020). Natural Reweighted Wake–Sleep uses natural gradients for generative and recognition networks (Várady et al., 2022), and recent work preconditions TopK SAE decoder updates using a damped empirical code Gram matrix (Picozzi, 2026). Both approaches retain learned encoders, whereas BeBeFOND directly optimizes variational posterior parameters for each input through recurrent inference. Building on this geometric perspective, Vafaii et al. (2025) introduced FOND by prescribing natural-gradient flow on ℱF with respect to [rgb]0.9297,0.1094,0.1406 u. A forward Euler discretization yields eq. 7. In the appendix, we motivate natural gradients using the standard trust region argument (Martens, 2020; Conn et al., 2000; Ollivier et al., 2017, section A.4;), as well as a broader “physics of adaptation” perspective (section A.5), inspired by the principle of general covariance (Einstein, 1916). 3 FOND unifies inference and learning as natural gradient flow FOND (Vafaii et al., 2025) formulated inference as natural gradient descent on variational free energy (eq. 7). But it stopped short of prescribing learning dynamics through the same principle. The geometric argument applies equally to learning: just as [rgb]0.9297,0.1094,0.1406 u changes the approximate posterior q_ [rgb]0.9297,0.1094,0.1406 u, the generative parameters [rgb]0.1602,0.6719,0.8867 θ change the joint distribution p(,)p_ [rgb]0.1602,0.6719,0.8867 θ( x, z). This motivates extending FOND by defining learning as natural-gradient flow on the same free energy with respect to [rgb]0.1602,0.6719,0.8867 θ. Thus, inference and learning are both treated as dynamical processes within a unified framework: FOND: Free-energy Online Natural-gradient Dynamics Flexible choices: Choose a generative model p(,)p_ [rgb]0.1602,0.6719,0.8867 θ( x, z), an inference model q(∣)q_ [rgb]0.9297,0.1094,0.1406 u( z x), and calculate the variational free energy ℱ(,)F( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ). FOND’s prescription: Define inference and learning as natural gradient flows on ℱF: Inference: ˙≔−u−1()∂ℱ(,)∂, [rgb]0.9297,0.1094,0.1406 u\, \,-G_u^-1( [rgb]0.9297,0.1094,0.1406 u)\, ( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)∂ [rgb]0.9297,0.1094,0.1406 u~, Learning: ˙≔−θ−1()∂ℱ(,)∂. [rgb]0.1602,0.6719,0.8867 θ\, \,-G_θ^-1( [rgb]0.1602,0.6719,0.8867 θ)\, ( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)∂ [rgb]0.1602,0.6719,0.8867 θ~. 3.1 BeBeFOND: a Bernoulli-distributed iterative sparse coding model We instantiate our extended FOND by following its three steps: choose the distributions, compute the free energy, and derive the dynamics. Distribution specifications. A generative model requires a prior and a conditional likelihood (eq. 2). We choose a factorized Bernoulli prior and a linear Gaussian decoder, together with a factorized Bernoulli variational posterior: Likelihood p(∣)p_ [rgb]0.1602,0.6719,0.8867 θ( x z) Prior p()p_ [rgb]0.1602,0.6719,0.8867 θ( z) Posterior q(∣)q_ [rgb]0.9297,0.1094,0.1406 u( z x) (,,) ( x; [rgb]0.1602,0.6719,0.8867 z, [rgb]0.1602,0.6719,0.8867 ) (9) ∏i=1KBernoulli(zi;π0,i) _i=1^KBernoulli(z_i; [rgb]0.1602,0.6719,0.8867π_0,i) (10) ∏i=1KBernoulli(zi;ri) _i=1^KBernoulli(z_i; [rgb]0.9297,0.1094,0.1406r_i) (11) [rgb]0.1602,0.6719,0.8867 diagonal and positive 0=σ(0) [rgb]0.1602,0.6719,0.8867 π_0=σ( [rgb]0.1602,0.6719,0.8867 u_0) =σ() [rgb]0.9297,0.1094,0.1406 r=σ( [rgb]0.9297,0.1094,0.1406 u) Here ∈0,1K z∈\0,1\^K, and σ(⋅)σ(·) is the logistic sigmoid. The generative parameters are ≡(,,0) [rgb]0.1602,0.6719,0.8867 θ≡( [rgb]0.1602,0.6719,0.8867 , [rgb]0.1602,0.6719,0.8867 , [rgb]0.1602,0.6719,0.8867 u_0), and the variational parameters [rgb]0.9297,0.1094,0.1406 u are the posterior logits. Free energy. Define ≔− e x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r as the mean reconstruction residual, ≔⊤−1 [rgb]0.1602,0.6719,0.8867W [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 as the precision-weighted Gram matrix, and ≔diag[⊙(−)] [rgb]0.9297,0.1094,0.1406C [ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)] as the posterior covariance. We weight the prior KL by β>0β>0, with β=1β=1 recovering the usual variational free energy. Because of the linear Gaussian assumption (section 3.1), the log-likelihood expectation can be evaluated exactly (section B.1): ℱ(,,)= ( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)\;\;=\;\; 12⊤−1+12Tr()+12logdet 12 e [rgb]0.1602,0.6719,0.8867 ^-1 e\;+\; 12 ( [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406C)\;+\; 12 [rgb]0.1602,0.6719,0.8867 +β∑i=1K[rilogriπ0,i+(1−ri)log1−ri1−π0,i]. +\;β _i=1^K [ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_i [rgb]0.1602,0.6719,0.8867π_0,i+(1- [rgb]0.9297,0.1094,0.1406r_i) 1- [rgb]0.9297,0.1094,0.1406r_i1- [rgb]0.1602,0.6719,0.8867π_0,i ]. (12) Dynamics. Applying FOND’s natural-gradient prescription, ˙≔−u−1()∇ℱ(,) [rgb]0.9297,0.1094,0.1406 u -G_u^-1( [rgb]0.9297,0.1094,0.1406 u) _ [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ) and ˙≔−θ−1()∇ℱ(,) [rgb]0.1602,0.6719,0.8867 θ -G_θ^-1( [rgb]0.1602,0.6719,0.8867 θ) _ [rgb]0.1602,0.6719,0.8867 θF( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ), with the Fisher metrics of the variational posterior and generative joint distribution, yields the following closed-form inference and learning dynamics (section B.2.1): Inference ˙=⊤−1⏟feedforward drive−⏟recurrentinteractions−12diag()⊙(−2)⏟variance penalty−β(−0)⏟homeostatic leak [rgb]0.9297,0.1094,0.1406 u\;\;=\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 x_ [rgb]0,0,0feedforward drive\,-\, [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406 r_ [rgb]0,0,0 subarraycrecurrent\\ interactions subarray\,-\, [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 12 ( [rgb]0.1602,0.6719,0.8867W) ( 1-2 [rgb]0.9297,0.1094,0.1406 r)_ [rgb]0,0,0variance penalty\,-\, [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0β( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_0)_ [rgb]0,0,0homeostatic leak (13) Learning ˙=ηΦ(⊤−)p−1 [rgb]0.1602,0.6719,0.8867 = _ ( x [rgb]0.9297,0.1094,0.1406 r - [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406E ) [rgb]0.1602,0.6719,0.8867E_p^-1 (14) ˙=ηΣ(−) [rgb]0.1602,0.6719,0.8867 = _ ( [rgb]0.9297,0.1094,0.1406V- [rgb]0.1602,0.6719,0.8867 ) (15) ˙0=η0β(−0) [rgb]0.1602,0.6719,0.8867 π_0= _0β( [rgb]0.9297,0.1094,0.1406 r- [rgb]0.1602,0.6719,0.8867 π_0) (16) Here ηΦ,ηΣ,η0>0 _ , _ , _0>0 set the learning timescales, and prior learning is expressed in mean coordinates 0=σ(0) [rgb]0.1602,0.6719,0.8867 π_0=σ( [rgb]0.1602,0.6719,0.8867 u_0). The posterior and prior second moments are: =+⊤≡diag[⊙(−)]+⊤,p=diag[0⊙(−0)]+00⊤, [rgb]0.9297,0.1094,0.1406E\;=\; [rgb]0.9297,0.1094,0.1406C\,+\, [rgb]0.9297,0.1094,0.1406 r [rgb]0.9297,0.1094,0.1406 r \;≡\; [ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)]\,+\, [rgb]0.9297,0.1094,0.1406 r [rgb]0.9297,0.1094,0.1406 r ~, [rgb]0.1602,0.6719,0.8867E_p\;=\; [ [rgb]0.1602,0.6719,0.8867 π_0 ( 1- [rgb]0.1602,0.6719,0.8867 π_0)]\,+\, [rgb]0.1602,0.6719,0.8867 π_0 [rgb]0.1602,0.6719,0.8867 π_0 ~, (17) and [rgb]0.9297,0.1094,0.1406V is diagonal, with Vmm=em2+(⊤)mm [rgb]0.9297,0.1094,0.1406V_\!m=e_m^2+( [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406C [rgb]0.1602,0.6719,0.8867 )_m, the expected squared reconstruction error in each observed dimension. Learning holds the inferred posterior fixed. For minibatches, we average ⊤ x [rgb]0.9297,0.1094,0.1406 r , [rgb]0.9297,0.1094,0.1406E, [rgb]0.9297,0.1094,0.1406V, and [rgb]0.9297,0.1094,0.1406 r before applying the learning rules. These equations keep the prior fixed during inference. Our implementation uses a rolling-reference variant (section B.1.1). Time discretization. Turning continuous flows into an implementable algorithm requires time discretization, which can affect stability. We use exponential integration for both inference (Cox and Matthews, 2002; Hochbruck and Ostermann, 2010, ETD1;) and learning (including matrix-free decoder updates; Al-Mohy and Higham (2011); Trefethen (2019)) to preserve linear relaxation within each step. See section B.6 for derivations and Alg. 1 for the full discretized algorithm. 4 Theoretical analyses, synthetic experiments, and ablations We study inference and learning using theory and synthetic data with known features and activations. Ablations test lateral interactions and compare separately tuned natural and Euclidean recipes. 4.1 Inference: BeBeFOND’s recurrent term enables explaining away Superposition mixes evidence from overlapping features (Elhage et al., 2022; Klindt et al., 2026), making inference of sparse codes difficult (section A.1). Here, we show how BeBeFOND’s iterative inference mitigates this interference through a mechanism known as explaining away. As an example, wet grass could be explained by rain or a sprinkler. Observing a wet roof supports rain without changing the sprinkler’s direct evidence from the grass. As rain explains more of the wetness, the sprinkler becomes less likely. This reduction in belief in one explanation through evidence for another is known as explaining away (Murphy, 2012). The challenge is that the influence of one piece of evidence depends on the others. An encoder that compresses the input to one fixed projection cannot generally capture the full dependence, as the following proposition establishes. Proposition 4.1 (A limitation of single-layer inference). Suppose the data are generated by the Bernoulli–Gaussian model: =+ϵ,zj∼indBernoulli(πj),ϵ∼(0,I), x\;=\; [rgb]0.1602,0.6719,0.8867 z\,+\, ε~, z_j ind Bernoulli( [rgb]0.1602,0.6719,0.8867π_j)~, ε (0,I)~, where 0<πj<10< [rgb]0.1602,0.6719,0.8867π_j<1 for every j, and the noise is independent of z. Suppose a dictionary atom ϕ(i) [rgb]0.1602,0.6719,0.8867 φ^(i) is neither orthogonal nor parallel to another atom ϕ(j) [rgb]0.1602,0.6719,0.8867 φ^(j). Then there exist no weights i [rgb]0.9297,0.1094,0.1406 w_i, bias bi [rgb]0.9297,0.1094,0.1406b_i, and scalar function fif_i such that: fi(i⊤+bi)=p(zi=1∣)=[zi∣],for every ∈ℝd.f_i( [rgb]0.9297,0.1094,0.1406 w_i x+ [rgb]0.9297,0.1094,0.1406b_i)\,\;=\;\,p(z_i=1 x)\,\;=\;\,E[z_i x]~, every x ^d~. For almost every dictionary with K,M≥2K,M≥ 2, no latent falls under the exception, so for generic dictionaries a feature’s exact posterior probability cannot be computed from one fixed input projection followed by a scalar nonlinearity. The proof shows that overlapping, nonparallel atoms make the posterior log-odds gradient change direction with the input, which is impossible for a function of one fixed projection (section C.2.1). This result generalizes O’Neill et al. (2025) and covers ReLU (Bricken et al., 2023; Cunningham et al., 2024), JumpReLU (including Gemma Scope) (Rajamanoharan et al., 2024b; Lieberum et al., 2024), weight-sharing gated (Rajamanoharan et al., 2024a), and fixed-threshold BatchTopK and Matryoshka BatchTopK SAEs (Bussmann et al., 2024; Bussmann et al., 2025), all of which compute each latent from a single projection of the input (section C.2.1). TopK SAEs fall outside this class, since whether a latent fires depends on every other pre-activation. BeBeFOND’s inference dynamics (section 3.1) couple explanations through the off-diagonal entries of −(⊤)-( [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ) [rgb]0.9297,0.1094,0.1406 r: positively overlapping features suppress one another’s activation drive. A wet roof strengthens rain, which explains the lawn’s wetness and suppresses the sprinkler. Repeated updates make this competition depend on the full observation (Fig. 2). Figure 2: Heatmaps show sprinkler probabilities under exact inference, converged BeBeFOND with a fixed prior, and a numerically optimized single-projection encoder. The marked observations have identical lawn wetness but different roof evidence: neutral (white, halfway dry / wet) or wet (blue). Fixed-prior inference and encoder fitting are detailed in section C.2.3. Section C.2.2 gives an approximate-inference counterpart to proposition 4.1 and sufficient conditions under which BeBeFOND can achieve strictly lower expected posterior KL than the best possible single-projection amortized encoder. We now empirically quantify the effects of explaining away. Explaining away under superposition. SynthSAEBench (Chanin and Garriga-Alonso, 2026) makes this problem measurable by providing synthetic data with known atoms and activations, together with feature correlations. Superposition is measured as the mean maximum (m) absolute cosine similarity: ρm=1K⋆∑i=1K⋆maxj≠i|(ϕ⋆(i))⊤ϕ⋆(j)|‖ϕ⋆(i)‖‖ϕ⋆(j)‖, _m\;=\; 1K_ _i=1^K_ _j≠ i |( [rgb]0.1602,0.6719,0.8867 φ_ ^(i)) [rgb]0.1602,0.6719,0.8867 φ_ ^(j) |\| [rgb]0.1602,0.6719,0.8867 φ_ ^(i)\|\,\| [rgb]0.1602,0.6719,0.8867 φ_ ^(j)\|~, (18) where ϕ⋆(i) [rgb]0.1602,0.6719,0.8867 φ_ ^(i) are the K⋆K_ ground-truth feature directions. Larger ρm _m means that each feature has, on average, a more strongly overlapping alternative. Reducing the input dimension at fixed feature count increases this overlap, modeling the interference created when many features share a real network’s activation space. Figure 3: BeBeFOND maintains the highest recovery, and without dead latents (K=4,096K=4,096). Means and 95% confidence intervals over 3 seeds. Exact values are reported in section C.3.1 and Table 4. As overlap increases, a strong detection score becomes less diagnostic of a feature’s presence: another feature may explain the same evidence. Explaining away addresses this ambiguity by discounting evidence already accounted for by competing causes. We measure dictionary recovery using the mean correlation coefficient (MCC), which quantifies alignment between learned and ground-truth feature directions. BeBeFOND achieves the highest MCC throughout the sweep, declining only modestly as superposition increases, while maintaining zero dead latents (Fig. 3). Ablating lateral interactions. In section C.3.1, we investigate the impact of ablating lateral interactions on dictionary recovery. Figure 6 shows that removing lateral interactions reduces dictionary recovery as feature overlap increases. Restoring suppressive interactions substantially improves recovery, whereas excitatory-only inference activates nearly the entire dictionary without comparable recovery. 4.2 Learning: BeBeFOND’s natural gradient flow improves rare feature learning Features extracted from language models vary widely in activation frequency (Gao et al., 2025), and larger SAE dictionaries recover increasingly rare concepts (Bricken et al., 2023; Templeton et al., 2024). Scaling dictionary width therefore places greater demands on learning from rare examples: additional neurons must capture infrequent features to expand the model’s coverage. Learning an infrequently active feature. For an isolated feature active with probability ρ, averaging over inputs scales the decoder gradient by ρ. The same decoder error therefore produces a smaller update for a rarer feature. Matched Fisher preconditioning divides out this frequency factor. Theorem 4.2 (Exact rarity compensation). For p(∣)=(,,)p( x z)=N( x; [rgb]0.1602,0.6719,0.8867 z,I), fix a training minibatch and hold its inferred distributions and prior references fixed. Let ∗ [rgb]0.1602,0.6719,0.8867 ^* minimize ℱF. If the averaged second moments satisfy p=≻0 [rgb]0.1602,0.6719,0.8867E_p= [rgb]0.9297,0.1094,0.1406E 0, then: Euclidean gradient :˙ gradient :\; [rgb]0.1602,0.6719,0.8867 =−ηΦ∇ℱ=−ηΦ(−∗), =\;\,- _ _ [rgb]0.1602,0.6719,0.8867 F\,\;=\;\,- _ ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E~, Natural gradient :˙ gradient :\; [rgb]0.1602,0.6719,0.8867 =−ηΦΦ−1∇ℱ=−ηΦ(−∗), =\;\,- _ G_ ^-1 _ [rgb]0.1602,0.6719,0.8867 F\,\;=\;\,- _ ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*)~, where Φ=p⊗G_ = [rgb]0.1602,0.6719,0.8867E_p and ηΦ>0 _ >0. Natural flow reduces every decoder error by e−ηΦΔte^- _ t over an interval Δt≥0 t≥ 0. For the same fractional error reduction, Euclidean flow takes κ()≥ρmax/ρminκ( [rgb]0.9297,0.1094,0.1406E)≥ _ / _ times as long in its slowest eigendirection as in its fastest, where ρi=ii _i= [rgb]0.9297,0.1094,0.1406E_i is the minibatch-averaged posterior activation probability. The rate difference becomes more pronounced as the frequency range widens. If ρi=ci−α _i=c\,i^-α, with α>0α>0 and i=1,…,Ki=1,…,K, the slowest Euclidean direction takes at least KαK^α times as long as the fastest to achieve the same fractional error reduction. Under the theorem’s matched-moment assumptions, natural flow gives every direction the same rate. For the proof and the extension to unequal prior and posterior second moments, see section C.4. Section C.4.2 shows that, with known codes and frozen Adam normalization, an optimer used by many common SAE baselines (Bussmann et al., 2024; Rajamanoharan et al., 2024b), Zipfian feature frequencies leave a local condition number of at least Kα/2K^α/2, whereas matched Fisher preconditioning reduces it to one. Recovering rare features. The full-width SynthSAEBench setting (Chanin and Garriga-Alonso, 2026) tests whether decoder capacity recovers the rare tail, with 16,384 learned latents for 16,384 generating features. This supplies one decoder slot per true feature, so complete recovery requires representing rare features alongside common ones. MCC measures recovery of the feature directions, while Rare F1 measures detection of the 25%25\% least frequent generating features, averaging their activation F1 scores. BeBeFOND leads all metrics, exceeding the evaluated baselines (Fig. 4; Table 5). Figure 4: BeBeFOND outperforms all evaluated baselines on dictionary recover (MCC) and feature detection (Macro F1, Rare F1) in SynthSAEBench (d=768d=768, K=K⋆=16,384K=K_ =16,384). Rare F1 covers the least frequent quarter of features. Bars show means over 3 initialization seeds with 95% confidence intervals, evaluated on 10610^6 held-out samples. Table 5 gives individual-seed results. Ablating inference and learning geometry. The natural-inference/natural-learning recipe reaches stronger recovery with fewer decoder updates than the separately tuned Euclidean variants (Fig. 7; section C.3.2). Section C.5.2 reports learning curves against the other baselines, and Table 5 reports dead-latent counts alongside recovery. 5 Evaluating BeBeFOND on language-model activations We evaluate whether BeBeFOND’s inference and learning benefits studied in section 4 translate to improved feature quality in language-model representations. We focus on three questions: (1) whether BeBeFOND outperforms widely used SAE baselines on Gemma 2B scale activations, (2) whether these benefits increase with dictionary width, and (3) how early they emerge during training. Experimental setup. We evaluate feature quality under a 5050M-token training budget, alongside released Gemma Scope JumpReLU checkpoints trained on 44B tokens at 1616k width and 88B tokens at widths 3232k–524524k (Lieberum et al., 2024; Rajamanoharan et al., 2024b). Our primary comparison trains BeBeFOND, BatchTopK (Bussmann et al., 2024), Matryoshka BatchTopK (Bussmann et al., 2025), and a mean-field variational baseline with Adam decoder learning (Kingma and Ba, 2015). These models use whitened Gemma-2-2B (Gemma Team, 2024) layer-12 residual-stream activations (d=2,304d=2,304) extracted from Pile Uncopyrighted (monology, 2023), with dictionary widths from 16,38416,384 to 524,288524,288 latents. We report means and sample standard deviations over three BeBeFOND training seeds and five baseline training seeds, using T=10T=10 inference steps for BeBeFOND and T=20T=20 for mean-field at evaluation. Amortized baselines use the selected target sparsities in Table 6. Gemma Scope checkpoints retain their original preprocessing. Appendix D provides further details. Evaluation metrics. Reconstruction quality alone does not establish whether an SAE has learned useful features (Gao et al., 2025; Karvonen et al., 2025). We therefore evaluate two complementary capabilities using SAEBench (Karvonen et al., 2025). Sparse probing (Kantamneni et al., 2025) tests concept detection using a logistic probe restricted to k=1k=1 or k=5k=5 selected SAE features. We average task-level accuracy, ROC-AUC, and F1 across 113113 classification tasks. Accuracy measures correct predictions, ROC-AUC measures the ranking of positive versus negative examples, and F1 summarizes the balance between precision and recall. RAVEL (Huang et al., 2024, Resolving Attribute–Value Entanglements in Language Models;) tests selective intervention on represented attributes. Its cause score measures successful changes to a target attribute, while isolation measures preservation of other attributes, and disentanglement is the mean of these two scores. Higher scores are better for all reported metrics. We report full numerical results in Table 6. Scaling with dictionary width. BeBeFOND achieves the highest mean RAVEL disentanglement at every tested width, including comparisons with the evaluated Gemma Scope checkpoints (Fig. 1d). At 524524k latents, its score is 0.827±0.0010.827± 0.001, compared with 0.748±0.0040.748± 0.004 for BatchTopK, the strongest amortized baseline trained on 5050M tokens (Table 6). This 7.97.9 percentage-point advantage reflects improvements in both cause and isolation: BeBeFOND scores 0.8000.800 and 0.8540.854, respectively, versus 0.7190.719 and 0.7760.776 for BatchTopK. Single-feature probing also favors BeBeFOND at large widths. At 524524k, its accuracy is 0.789±0.0020.789± 0.002, versus 0.774±0.0020.774± 0.002 for Matryoshka BatchTopK, an improvement of 1.51.5 percentage points. Across widths from 1616k to 524524k, BeBeFOND’s single-feature accuracy increases from 0.7710.771 to 0.7890.789, and its RAVEL score increases from 0.7860.786 to 0.8270.827 (Fig. 1b,d). Over the same range, BatchTopK and Matryoshka each improve by approximately 0.60.6 percentage points in single-feature accuracy. Thus, increasing dictionary width produces larger gains for BeBeFOND than for either amortized baseline. Differences between methods are smaller in five-feature probing accuracy, and the AUC and F1 rankings vary across widths (Table 9). High feature quality emerges early in training. At 524524k latents, BeBeFOND reaches the evaluated Gemma Scope single-feature probing score by ∼ 6M training tokens (Fig. 1a). The corresponding Gemma Scope checkpoints were trained on 8B tokens, a training-budget ratio of over 1,000×. By 50M tokens, BeBeFOND exceeds their single-feature accuracy (0.7890.789 versus 0.77350.7735) and RAVEL disentanglement (0.8270.827 versus 0.75550.7555), using 160× fewer training tokens. Comparison with iterative mean-field inference. The mean-field baseline performs iterative variational inference and learns its decoder using Adam. At 524524k, it reaches single-feature accuracy of 0.751±0.0070.751± 0.007 and RAVEL disentanglement of 0.685±0.0290.685± 0.029, significantly underperforming BeBeFOND on both metrics (Table 6). 6 Conclusion and Discussion We formulate inference and dictionary learning as natural-gradient flows on a shared variational free energy (the FOND framework). Our theory establishes limitations of single-projection encoders and shows how Fisher preconditioning compensates for infrequent activation under frozen, matched moments. We introduce BeBeFOND, which improves synthetic feature recovery and supports increasing feature quality at larger dictionary widths on Gemma. Limitations. Repeated inference steps and numerical decoder updates increase BeBeFOND’s computational cost (sections B.8 and D.3). Future work will focus on improving computational efficiency, speeding up inference convergence by adding momentum, and extending the factorized variational family with pairwise interactions to represent dependencies between latent features. On Gemma-2-2B layer-12 activations, BeBeFOND improves concept detection and selective intervention. Its gains increase with dictionary width while the evaluated baselines largely plateau. Together with recurrent explaining away and improved rare-feature recovery, these results show the practical value of taking the statistical geometry of inference and learning seriously. The broader lesson is how adaptation dynamics should be derived: Physics of Adaptation: deriving covariant adaptation dynamics Inference and learning dynamics should respect the intrinsic geometry and symmetries of the underlying model (section A.5). Code and data Our code, data, and model checkpoints are available here: https://github.com/hadivafaii/BeFOND AI use statement We used generative AI tools to assist with mathematical derivations, software implementation and debugging, experimental analysis and visualization, literature searches, and manuscript drafting and editing. The authors are responsible for the final text, results, and claims. References Aharon et al. (2006) M. Aharon, M. Elad, and A. Bruckstein rmKrmK-SVD: An Algorithm for Designing Overcomplete Dictionaries for Sparse Representation. IEEE Transactions on Signal Processing 54 (11), p. 4311–4322. External Links: Document, Link Cited by: §2. Al-Mohy and Higham (2011) A. H. Al-Mohy and N. J. Higham Computing the Action of the Matrix Exponential, with an Application to Exponential Integrators. SIAM Journal on Scientific Computing 33 (2), p. 488–511. External Links: Document, Link Cited by: §B.6, §3.1. Alemi et al. (2018) A. Alemi, B. Poole, I. Fischer, J. Dillon, R. A. Saurous, and K. Murphy Fixing a Broken ELBO. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 80, p. 159–168. External Links: Link Cited by: §A.3. Amari and Nagaoka (2000) S. Amari and H. Nagaoka Methods of Information Geometry. Translations of Mathematical Monographs, Vol. 191, American Mathematical Society. External Links: ISBN 9780821805312, Document, Link Cited by: §A.5, §2, §2. Amari (1998) S. Amari Natural Gradient Works Efficiently in Learning. Neural Computation 10 (2), p. 251–276. External Links: Document, Link Cited by: §A.4, §A.4, §A.4, §A.5, §A.5, Table 1, §B.1, §B.2, §1, §2. Amari (2016) S. Amari Information Geometry and Its Applications. Springer Japan. External Links: Document, Link, ISBN 9784431559771 Cited by: §A.5, §2. Attwell and Laughlin (2001) D. Attwell and S. B. Laughlin An Energy Budget for Signaling in the Grey Matter of the Brain. Journal of Cerebral Blood Flow & Metabolism 21 (10), p. 1133–1145. External Links: Document, Link Cited by: §A.5. Auconi and Ito (2026) A. Auconi and S. Ito Information-Geometric Signatures of Nonconservative Driving. External Links: 2605.03757, Document, Link Cited by: §A.5. Balles and Hennig (2018) L. Balles and P. Hennig Dissecting Adam: the sign, magnitude and variance of stochastic gradients. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 80, p. 404–413. External Links: Link Cited by: §C.4.2, §C.4.2. Barin-Pacela et al. (2026) V. Barin-Pacela, S. Joshi, I. Camacho, S. Lacoste-Julien, and D. Klindt Stop Probing, Start Coding: Why Linear Probes and Sparse Autoencoders Fail at Compositional Generalisation. External Links: 2603.28744, Link Cited by: §A.2, §1. Beck and Teboulle (2009) A. Beck and M. Teboulle A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems. SIAM Journal on Imaging Sciences 2 (1), p. 183–202. External Links: Document, Link Cited by: §A.2, §2. Bissiri et al. (2016) P. G. Bissiri, C. C. Holmes, and S. G. Walker A General Framework for Updating Belief Distributions. Journal of the Royal Statistical Society Series B: Statistical Methodology 78 (5), p. 1103–1130. External Links: Document, Link Cited by: §B.1.1. Blei et al. (2017) D. M. Blei, A. Kucukelbir, and J. D. McAuliffe Variational Inference: A Review for Statisticians. Journal of the American Statistical Association 112 (518), p. 859–877. External Links: Document, Link Cited by: §A.3, §A.3, §A.3, §A.3, §2. Boyd and Vandenberghe (2004) S. Boyd and L. Vandenberghe Convex optimization. Cambridge University Press. External Links: Link Cited by: §A.4, §A.4. Bricken et al. (2023) T. Bricken, A. Templeton, J. Batson, B. Chen, A. Jermyn, T. Conerly, N. L. Turner, C. Anil, C. Denison, A. Askell, R. Lasenby, Y. Wu, S. Kravec, N. Schiefer, T. Maxwell, N. Joseph, A. Tamkin, K. Nguyen, B. McLean, J. E. Burke, T. Hume, S. Carter, T. Henighan, and C. Olah Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits Thread. External Links: Link Cited by: §C.5.1, §1, §2, §4.1, §4.2. Bussmann et al. (2024) B. Bussmann, P. Leask, and N. Nanda BatchTopK Sparse Autoencoders. External Links: 2412.06410, Link Cited by: §1, §2, §4.1, §4.2, §5. Bussmann et al. (2025) B. Bussmann, N. Nabeshima, A. Karvonen, and N. Nanda Learning Multi-Level Features with Matryoshka Sparse Autoencoders. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 6077–6101. External Links: Link Cited by: §1, §2, §4.1, §5. Cauchy (1847) A. Cauchy Méthode générale pour la résolution des systèmes d’équations simultanées. Comptes rendus hebdomadaires des séances de l’Académie des sciences. External Links: Link Cited by: Table 1. Chanin and Garriga-Alonso (2026) D. Chanin and A. Garriga-Alonso SynthSAEBench: Evaluating Sparse Autoencoders on Scalable Realistic Synthetic Data. External Links: 2602.14687, Link Cited by: §C.1.1, §1, §4.1, §4.2. Chanin (2025) D. Chanin Training Matching Pursuit SAEs on LLMs. Note: LessWrong External Links: Link Cited by: §A.2, §2. chanind-goodfire (2026) chanind-goodfire SynthSAEBench Full-Width SAE Sweep. Note: Hugging Face model and training-recipe release External Links: Link Cited by: §C.5.1. Conn et al. (2000) A. R. Conn, N. I. M. Gould, and P. L. Toint Trust region methods. Society for Industrial and Applied Mathematics. External Links: Document Cited by: §2. Costa et al. (2025) V. Costa, T. Fel, E. S. Lubana, B. Tolooshams, and D. Ba From Flat to Hierarchical: Extracting Sparse Representations with Matching Pursuit. In Advances in Neural Information Processing Systems 38, p. 38625–38668. External Links: Document, Link Cited by: §A.2. Cox and Matthews (2002) S.M. Cox and P.C. Matthews Exponential Time Differencing for Stiff Systems. Journal of Computational Physics 176 (2), p. 430–455. External Links: Document, Link Cited by: §B.6, §B.6, §3.1. Cunningham et al. (2024) H. Cunningham, A. Ewart, L. Riggs, R. Huben, and L. Sharkey Sparse Autoencoders Find Highly Interpretable Features in Language Models. In International Conference on Learning Representations, Vol. 2024, p. 7827–7845. External Links: Link Cited by: §1, §2, §4.1. Daubechies et al. (2004) I. Daubechies, M. Defrise, and C. De Mol An iterative thresholding algorithm for linear inverse problems with a sparsity constraint. Communications on Pure and Applied Mathematics 57 (11), p. 1413–1457. External Links: Document, Link Cited by: §A.2. Dowty (2018) J. G. Dowty Chentsov’s theorem for exponential families. Information Geometry 1, p. 117–135. External Links: Document, Link Cited by: §A.4. Einstein (1916) A. Einstein Die grundlage der allgemeinen relativitätstheorie. Annalen der Physik 354 (7), p. 769–822. External Links: Document Cited by: §A.5, §A.5, §2. Elhage et al. (2022) N. Elhage, T. Hume, C. Olsson, N. Schiefer, T. Henighan, S. Kravec, Z. Hatfield-Dodds, R. Lasenby, D. Drain, C. Chen, R. Grosse, S. McCandlish, J. Kaplan, D. Amodei, M. Wattenberg, and C. Olah Toy models of superposition. Transformer Circuits Thread. External Links: Link Cited by: §A.1, §4.1. Fang et al. (2022) M. Y.-S. Fang, M. Mudigonda, R. Zarcone, A. Khosrowshahi, and B. A. Olshausen Learning and Inference in Sparse Coding Models With Langevin Dynamics. Neural Computation 34 (8), p. 1676–1700. External Links: Document, Link Cited by: §2. Gao et al. (2025) L. Gao, T. Dupre la Tour, H. Tillman, G. Goh, R. Troll, A. Radford, I. Sutskever, J. Leike, and J. Wu Scaling and evaluating sparse autoencoders. In International Conference on Learning Representations, Vol. 2025, p. 26721–26754. External Links: Link Cited by: §1, §2, §4.2, §5. Gemma Team (2024) Gemma Team Gemma 2: Improving Open Language Models at a Practical Size. arXiv preprint arXiv:2408.00118. External Links: Link Cited by: §D.1, §5. Gregor and LeCun (2010) K. Gregor and Y. LeCun Learning fast approximations of sparse coding. In Proceedings of the 27th International Conference on Machine Learning, p. 399–406. External Links: Link Cited by: §A.2, §1, §2. Hager (1989) W. W. Hager Updating the Inverse of a Matrix. SIAM Review 31 (2), p. 221–239. External Links: Link Cited by: §B.6. Hobson (1969) A. Hobson A new theorem of information theory. Journal of Statistical Physics 1 (3), p. 383–391. External Links: Document, Link Cited by: §A.4, §A.5, §A.5. Hochbruck and Ostermann (2010) M. Hochbruck and A. Ostermann Exponential integrators. Acta Numerica 19, p. 209–286. External Links: Document, Link Cited by: §B.6, §B.6, §B.6, §3.1. Hoffman et al. (2013) M. D. Hoffman, D. M. Blei, C. Wang, and J. Paisley Stochastic Variational Inference. Journal of Machine Learning Research 14 (40), p. 1303–1347. External Links: Link Cited by: §2. Huang et al. (2024) J. Huang, Z. Wu, C. Potts, M. Geva, and A. Geiger RAVEL: Evaluating Interpretability Methods on Disentangling Language Model Representations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Bangkok, Thailand, p. 8669–8687. External Links: Link, Document Cited by: §D.1, Table 8, §1, §5. Jermyn and Templeton (2024) A. Jermyn and A. Templeton Ghost Grads: An Improvement on Resampling. Transformer Circuits Thread. Note: Circuits Updates, January 2024 External Links: Link Cited by: §C.5.1. Jones et al. (2024) M. Jones, P. Chang, and K. Murphy Bayesian Online Natural Gradient (BONG). In Advances in Neural Information Processing Systems 37, p. 131104–131153. External Links: Document, Link Cited by: §2. Jordan et al. (1999) M. I. Jordan, Z. Ghahramani, T. S. Jaakkola, and L. K. Saul An Introduction to Variational Methods for Graphical Models. Machine Learning 37 (2), p. 183–233. External Links: Document, Link Cited by: §A.3, §2. Kantamneni et al. (2025) S. Kantamneni, J. Engels, S. Rajamanoharan, M. Tegmark, and N. Nanda Are Sparse Autoencoders Useful? A Case Study in Sparse Probing. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 29018–29049. External Links: Link Cited by: Table 8, §1, §5. Karvonen et al. (2025) A. Karvonen, C. Rager, J. Lin, C. Tigges, J. I. Bloom, D. Chanin, Y. Lau, E. Farrell, C. S. Mcdougall, K. Ayonrinde, D. Till, M. Wearden, A. Conmy, S. Marks, and N. Nanda SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 29223–29264. External Links: Link Cited by: §D.1, Table 9, §1, §5. Khan and Rue (2023) M. E. Khan and H. Rue The Bayesian Learning Rule. Journal of Machine Learning Research 24 (281), p. 1–46. External Links: Link Cited by: §A.5, §A.5, Table 1, §1, §2. Kingma and Ba (2015) D. P. Kingma and J. Ba Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, External Links: Link Cited by: §C.4.2, §1, §5. Kingma and Welling (2014) D. P. Kingma and M. Welling Auto-Encoding Variational Bayes. In International Conference on Learning Representations, External Links: Link Cited by: §A.3, §A.3, §2. Klindt et al. (2026) D. Klindt, C. O’Neill, P. Reizinger, H. Maurer, and N. Miolane A unifying framework from neural superposition to sparse interpretable codes. Nature Machine Intelligence 8 (7), p. 1025–1037. External Links: Document, Link Cited by: §1, §4.1. Landau and Lifshitz (1975) L.D. Landau and E.M. Lifshitz The Classical Theory of Fields. 4th edition, Course of Theoretical Physics, Vol. 2, Pergamon Press. External Links: Link Cited by: §A.5. Lewicki and Sejnowski (2000) M. S. Lewicki and T. J. Sejnowski Learning Overcomplete Representations. Neural Computation 12 (2), p. 337–365. External Links: Document, Link Cited by: §1, §2. Lieberum et al. (2024) T. Lieberum, S. Rajamanoharan, A. Conmy, L. Smith, N. Sonnerat, V. Varma, J. Kramar, A. Dragan, R. Shah, and N. Nanda Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2. In Proceedings of the 7th BlackboxNLP Workshop: Analyzing and Interpreting Neural Networks for NLP, p. 278–300. External Links: Document, Link Cited by: §C.5.1, Table 7, Table 9, §1, §4.1, §5. Lu et al. (2018) H. Lu, R. M. Freund, and Y. Nesterov Relatively Smooth Convex Optimization by First-Order Methods, and Applications. SIAM Journal on Optimization 28 (1), p. 333–354. External Links: Document, Link Cited by: §C.2.2. Mairal et al. (2010) J. Mairal, F. Bach, J. Ponce, and G. Sapiro Online Learning for Matrix Factorization and Sparse Coding. Journal of Machine Learning Research 11 (2), p. 19–60. External Links: Link Cited by: §2. Makhzani and Frey (2014) A. Makhzani and B. Frey k-Sparse Autoencoders. In International Conference on Learning Representations, External Links: Link Cited by: §2. Mallat and Zhang (1993) S.G. Mallat and Z. Zhang Matching pursuits with time-frequency dictionaries. IEEE Transactions on Signal Processing 41 (12), p. 3397–3415. External Links: Document, Link Cited by: §A.2. Martens and Grosse (2015) J. Martens and R. Grosse Optimizing Neural Networks with Kronecker-factored Approximate Curvature. In Proceedings of the 32nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 37, Lille, France, p. 2408–2417. External Links: Link Cited by: §2. Martens (2020) J. Martens New Insights and Perspectives on the Natural Gradient Method. Journal of Machine Learning Research 21 (146), p. 1–76. External Links: Link Cited by: §A.4, §A.4, §A.4, §A.4, §A.4, §A.5, §A.5, §B.1, §B.2, §C.4.2, §1, §2, §2. monology (2023) monology Pile Uncopyrighted. Note: Hugging Face dataset External Links: Link Cited by: §D.1, §5. Murphy (2012) K. P. Murphy Machine Learning: A Probabilistic Perspective. MIT Press. External Links: Link Cited by: §A.1, §4.1. Neal and Hinton (1998) R. M. Neal and G. E. Hinton A View of the EM Algorithm that Justifies Incremental, Sparse, and Other Variants. In Learning in Graphical Models, M. I. Jordan (Ed.), p. 355–368. External Links: Document, Link, ISBN 9789401061049 Cited by: §A.3. Ollivier et al. (2017) Y. Ollivier, L. Arnold, A. Auger, and N. Hansen Information-geometric optimization algorithms: a unifying picture via invariance principles. Journal of Machine Learning Research 18 (18), p. 1–65. External Links: Link Cited by: §2. Olshausen and Field (2004) B. A. Olshausen and D. J. Field Sparse coding of sensory inputs. Current Opinion in Neurobiology 14 (4), p. 481–487. External Links: Document, Link Cited by: §1. Olshausen and Field (1996) B. A. Olshausen and D. J. Field Emergence of simple-cell receptive field properties by learning a sparse code for natural images. Nature 381 (6583), p. 607–609. External Links: Document, Link Cited by: §1, §2. Olshausen and Field (1997) B. A. Olshausen and D. J. Field Sparse coding with an overcomplete basis set: A strategy employed by V1?. Vision Research 37 (23), p. 3311–3325. External Links: Document, Link Cited by: §1, §2. O’Neill et al. (2025) C. O’Neill, A. Gumran, and D. Klindt Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 46877–46896. External Links: Link Cited by: §1, §4.1. Picozzi (2026) D. Picozzi The information geometry of large language models is shared, learned, and controllable. External Links: 2609.11063, Link Cited by: §2. Rajamanoharan et al. (2024a) S. Rajamanoharan, A. Conmy, L. Smith, T. Lieberum, V. Varma, J. Kramár, R. Shah, and N. Nanda Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders. In Advances in Neural Information Processing Systems, Vol. 37, p. 775–818. External Links: Document, Link Cited by: §2, §4.1. Rajamanoharan et al. (2024b) S. Rajamanoharan, T. Lieberum, N. Sonnerat, A. Conmy, V. Varma, J. Kramár, and N. Nanda Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders. External Links: 2407.14435, Link Cited by: §1, §2, §4.1, §4.2, §5. Rangaraju et al. (2019) V. Rangaraju, M. Lauterbach, and E. M. Schuman Spatially Stable Mitochondrial Compartments Fuel Local Translation during Plasticity. Cell 176 (1–2), p. 73–84.e15. External Links: Document, Link Cited by: §A.5. Rozell et al. (2008) C. J. Rozell, D. H. Johnson, R. G. Baraniuk, and B. A. Olshausen Sparse Coding via Thresholding and Local Competition in Neural Circuits. Neural Computation 20 (10), p. 2526–2563. External Links: Document, Link Cited by: §A.1, §A.2, §1, §2, §2. Saad (2003) Y. Saad Iterative Methods for Sparse Linear Systems. 2 edition, Society for Industrial and Applied Mathematics. External Links: Link Cited by: §B.6. Särkkä and Svensson (2023) S. Särkkä and L. Svensson Bayesian Filtering and Smoothing. 2 edition, Cambridge University Press. External Links: Document, Link Cited by: §B.1.1. Sun et al. (2026) W. Sun, D. Wang, and L. Hu The Price of Amortized inference in Sparse Autoencoders. In International Conference on Learning Representations, Vol. 2026, p. 72686–72710. External Links: Link Cited by: §A.2, §A.2, §2. Templeton et al. (2024) A. Templeton, T. Conerly, J. Marcus, J. Lindsey, T. Bricken, B. Chen, A. Pearce, C. Citro, E. Ameisen, A. Jones, H. Cunningham, N. L. Turner, C. McDougall, M. MacDiarmid, A. Tamkin, E. Durmus, T. Hume, F. Mosconi, C. D. Freeman, T. R. Sumers, E. Rees, J. Batson, A. Jermyn, S. Carter, C. Olah, and T. Henighan Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet. Transformer Circuits Thread. External Links: Link Cited by: §1, §2, §4.2. Trefethen (2019) L. N. Trefethen Approximation Theory and Approximation Practice, Extended Edition. Society for Industrial and Applied Mathematics. External Links: Link Cited by: §B.6, §3.1. Vafaii et al. (2025) H. Vafaii, D. Galor, and J. Yates Brain-like Variational Inference. In Advances in Neural Information Processing Systems, Vol. 38, Main Conference, p. 163700–163778. External Links: Document, Link Cited by: §A.4, §B.1, §B.1.1, §1, §1, §2, §3. Várady et al. (2022) C. Várady, R. Volpi, L. Malagò, and N. Ay Natural Reweighted Wake–Sleep. Neural Networks 155, p. 574–591. External Links: Document, Link Cited by: §2. Wainwright and Jordan (2008) M. J. Wainwright and M. I. Jordan Graphical Models, Exponential Families, and Variational Inference. Foundations and Trends® in Machine Learning 1 (1-2), p. 1–305. External Links: Document, Link Cited by: §A.3, §B.1. Yun et al. (2021) Z. Yun, Y. Chen, B. Olshausen, and Y. LeCun Transformer visualization via dictionary learning: contextualized embedding as a linear superposition of transformer factors. In Proceedings of Deep Learning Inside Out (DeeLIO): The 2nd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, Online, p. 1–10. External Links: Link, Document Cited by: §1. Zellner (1988) A. Zellner Optimal Information Processing and Bayes’s Theorem. The American Statistician 42 (4), p. 278–280. External Links: Document, Link Cited by: §A.5. Zellner (2002) A. Zellner Information processing and Bayesian analysis. Journal of Econometrics 107 (1-2), p. 41–50. External Links: Document, Link Cited by: §A.5. Appendix: Inference and learning in sparse autoencoders as natural gradient flow 1 Introduction 2 Background and Related Work 3 FOND unifies inference and learning as natural gradient flow 3.1 BeBeFOND: a Bernoulli-distributed iterative sparse coding model 4 Theoretical analyses, synthetic experiments, and ablations 4.1 Inference: BeBeFOND’s recurrent term enables explaining away 4.2 Learning: BeBeFOND’s natural gradient flow improves rare feature learning 5 Evaluating BeBeFOND on language-model activations 6 Conclusion and Discussion References A Extended background A.1 Superposition and inference A.2 Iterative sparse encoders. A.3 Variational inference and the free energy objective A.4 Trust regions and natural-gradient inference A.5 “Physics of Adaptation”: deriving covariant inference and learning rules B Extended FOND model and dynamics B.1 Model and objective B.1.1 Online prior and working memory B.2 Natural-gradient dynamics B.2.1 Model updates B.3 Decoder bias B.4 Bernoulli uncertainty and sparse coding B.5 Activation threshold under a rolling prior B.6 Numerical integration and matrix-free updates B.7 Algorithm B.8 Computational complexity C Extended theoretical analyses, synthetic experiments, and ablations C.1 Experimental setup C.1.1 Datasets C.1.2 Training C.1.3 Evaluation C.1.4 Implementation C.2 Inference proofs C.2.1 Single-layer inference C.2.2 Iterative inference versus single projection encoders C.2.3 Rain–sprinkler example C.3 Inference and learning ablations C.3.1 Inference ablation: lateral interactions C.3.2 Learning ablation: inference and learning geometry C.3.3 Rolling prior C.4 Exact and approximate rarity compensation C.4.1 Activity-normalized inference C.4.2 Local conditioning under Adam normalization C.5 Learning and rare-feature recovery C.5.1 Resampling and auxiliary-loss variants C.5.2 Training sample efficiency D Extended Gemma experiments D.1 Experimental setup D.2 Results D.3 Wall-clock training times on H200 GPUs Appendix A Extended background A.1 Superposition and inference Inferring a sparse code can be difficult even with a perfectly learned dictionary. To isolate the inference problem, consider the noiseless generative process =∗ x= ^* z, where the ground truth dictionary ∗∈ℝM×K ^* ^M× K has unit-norm columns, and suppose we have learned the true dictionary =∗ [rgb]0.1602,0.6719,0.8867 = ^*. Projecting the observation onto each dictionary atom gives: (⊤)i=zi+∑j≠i(ϕ(i))⊤ϕ(j)zj⏟interference from other features. ( [rgb]0.1602,0.6719,0.8867 x )_i\;=\;z_i\,+\, _j≠ i ( [rgb]0.1602,0.6719,0.8867 φ^(i) ) [rgb]0.1602,0.6719,0.8867 φ^(j)z_j_interference from other features~. (19) The projection combines the feature’s own activity with contributions from other active features, weighted by their overlap with its dictionary atom. Even when only one feature j is active, an inactive feature i can receive a nonzero projection. When several features are active, their contributions can reinforce or cancel one another, making projection scores less reliable indicators of individual feature activity (Elhage et al., 2022). For an orthonormal dictionary, all interference terms vanish, and inference reduces to a single projection: =⊤ [rgb]0.9297,0.1094,0.1406 z= [rgb]0.1602,0.6719,0.8867 x. However, an overcomplete dictionary (K>MK>M) cannot have mutually orthogonal columns, so this interference cannot be eliminated for all feature pairs. Therefore, inference should distinguish a feature’s own contribution from evidence attributable to other active features. This motivates explaining away: support for a candidate feature decreases when other features already explain the evidence that initially supported it (Murphy, 2012). Such interactions allow inference to account for overlap between dictionary atoms rather than interpreting each projection independently (Rozell et al., 2008). A.2 Iterative sparse encoders. Iterative encoders refine sparse representations through successive updates rather than a single feed-forward prediction. ISTA, FISTA, and soft-thresholding LCA optimize ℓ1 _1-regularized reconstruction (Daubechies et al., 2004; Beck and Teboulle, 2009; Rozell et al., 2008), with LCA implementing lateral competition. LISTA learns a finite-step approximation (Gregor and LeCun, 2010), while matching pursuit (MP) greedily selects atoms to explain the remaining residual (Mallat and Zhang, 1993). With a fixed pretrained dictionary, better inference remains limited by the learned feature directions. Additional ISTA refinement improves reconstruction with little change in BatchTopK probe scores; non-amortized ISTA worsens splitting for Gemma TopK and increases dead latents for Pythia Standard SAEs (Sun et al., 2026). FISTA applied to fixed SAE dictionaries also fails to close the feature-recovery gap in controlled experiments, implicating the learned dictionary as a bottleneck (Barin-Pacela et al., 2026). On Pythia-160M, however, Sun et al. report improved targeted-probe selectivity and generally higher generative intervention scores with semi- and non-amortized inference (Sun et al., 2026, Appendix C). Approaches that also learn the dictionary have shown promising results: MP-SAE captures hierarchical features through residual-guided selection and supports different inference-time sparsity levels (Costa et al., 2025), while LocA-SAE seeks similar feature-quality benefits through grouped amortized encoders rather than explicit iterative refinement (Sun et al., 2026). But better reconstruction has not consistently translated into better interpretability tools. On Gemma-2-2B, MP-SAEs improve reconstruction but underperform the tested feed-forward baselines on sparse probing, with substantially longer training. The evaluator therefore recommends against them as a first choice for practical interpretability (Chanin, 2025). LocA-SAE eliminates dead latents on Gemma but scores below the Standard SAE on the same proxy probes, so improved utilization need not yield more predictive features (Sun et al., 2026). Across these evaluations, reconstruction, probing, and intervention metrics need not improve together, and the practical benefit of iterative refinement in these prior work remains task-dependent. A.3 Variational inference and the free energy objective We provide a pedagogical introduction to variational inference and derive the free energy objective. In this derivation, we distinguish two KL divergences: one measures how well we approximate the true posterior, while the other measures departure from the prior. They play different roles in the same exact identity. Finally, we explain why minimizing free energy improves inference using a few lines of algebra. Variational posterior inference. Consider an observation ∈ℝM x ^M and K-dimensional latent variables z. For Bernoulli latents, ∈0,1K z∈\0,1\^K. A generative model specifies a prior p()p_ [rgb]0.1602,0.6719,0.8867 θ( z) and a likelihood p(∣)p_ [rgb]0.1602,0.6719,0.8867 θ( x z), whose product is the joint distribution. Bayes’ rule gives: p(∣)=p(,)p(),p()=∑∈0,1Kp(∣)p().p_ [rgb]0.1602,0.6719,0.8867 θ( z x)\;=\; p_ [rgb]0.1602,0.6719,0.8867 θ( x, z)p_ [rgb]0.1602,0.6719,0.8867 θ( x)~, p_ [rgb]0.1602,0.6719,0.8867 θ( x)\;=\; _ z∈\0,1\^Kp_ [rgb]0.1602,0.6719,0.8867 θ( x z)\,p_ [rgb]0.1602,0.6719,0.8867 θ( z)~. (20) The denominator is the model evidence: the probability density of the observation after summing over all latent explanations. This sum contains 2K2^K terms, making direct evaluation generally infeasible for large K. The “true posterior” here means the exact posterior under the chosen generative model. Variational inference is a powerful approximation method that replaces this computation with an optimization problem (Jordan et al., 1999; Blei et al., 2017). In variational inference, we introduce a tractable approximate posterior q(∣)q_ [rgb]0.9297,0.1094,0.1406 u( z x) and adjust its variational parameters [rgb]0.9297,0.1094,0.1406 u to minimize (q(∣)∥p(∣))≔∑q(∣)logq(∣)p(∣).D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z x) )\; \; _ zq_ [rgb]0.9297,0.1094,0.1406 u( z x) q_ [rgb]0.9297,0.1094,0.1406 u( z x)p_ [rgb]0.1602,0.6719,0.8867 θ( z x)~. (21) This KL divergence is nonnegative and vanishes exactly when the two distributions agree (Blei et al., 2017). But its value still depends on the intractable posterior, so we can’t directly compute and optimize it. The next step isolates that dependence, leaving an objective we can optimize without evaluating the evidence. Free-energy derivation. Fix x and [rgb]0.1602,0.6719,0.8867 θ throughout the derivation. To keep the algebra readable, write q()≡q(∣) [rgb]0.9297,0.1094,0.1406q( z)≡ q_ [rgb]0.9297,0.1094,0.1406 u( z x) and suppress the parameter subscript on p. All expectations below are over ∼q z [rgb]0.9297,0.1094,0.1406q. For finite Bernoulli logits and a positive Gaussian covariance, the distributions are strictly positive, so the following ratios are well defined. Since q [rgb]0.9297,0.1094,0.1406q is normalized, ∑q()=1 _ z [rgb]0.9297,0.1094,0.1406q( z)=1. Multiply the log evidence by this sum, apply Bayes’ rule, and insert q()/q()=1 [rgb]0.9297,0.1094,0.1406q( z)/ [rgb]0.9297,0.1094,0.1406q( z)=1: logp() p( x) =logp()∑q()⏟= 1=q[logp()] =\; p( x) [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 _ z [rgb]0.9297,0.1094,0.1406q( z)_ [rgb]0,0,0=\,1\;=\;E_q[ p( x)] (22) =q[logp(,)p(∣)] =\;E_q[ p( x, z)p( z x)] =q[log[p(,)p(∣)q()q()]] =\;E_q[ [ p( x, z)p( z x) [rgb]0.9297,0.1094,0.1406q( z) [rgb]0.9297,0.1094,0.1406q( z) ]] =q[logp(,)q()]+q[logq()p(∣)] =\;E_q[ p( x, z) [rgb]0.9297,0.1094,0.1406q( z)]\;+\;E_q[ [rgb]0.9297,0.1094,0.1406q( z)p( z x)] =q[logp(,)q()]⏟ELBO=−ℱ+(q()∥p(∣)). =\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0E_q[ p( x, z) [rgb]0.9297,0.1094,0.1406q( z)]_ [rgb]0,0,0ELBO\;=\;-F\;+\;D_ KL ( [rgb]0.9297,0.1094,0.1406q( z)\, \|\,p( z x) )~. All we did was multiply by one twice, then rearrange logarithms. We didn’t have to invoke Jensen’s inequality as it is typically done. Instead, we have an exact decomposition of the log evidence. Nonnegativity of the final KL gives the evidence lower bound, ELBO≤logp()ELBO≤ p( x). Its negative is the variational free energy, ℱ≡−ELBOF≡-ELBO (Neal and Hinton, 1998; Kingma and Welling, 2014). Posterior and prior KL divergences. Expanding p(,)=p(∣)p()p_ [rgb]0.1602,0.6719,0.8867 θ( x, z)=p_ [rgb]0.1602,0.6719,0.8867 θ( x z)p_ [rgb]0.1602,0.6719,0.8867 θ( z) inside the ELBO reveals a second KL (Kingma and Welling, 2014). With ℱ≡ℱ(,,)F ( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ), we have: logp()⏟log evidence=−ℱ⏟ELBO+(q(∣)∥p(∣))⏟inference KL: original objective(gap to the true posterior)ℱ=q[−logp(∣)]⏟reconstruction term (accuracy)+(q(∣)∥p())⏟prior KL (coding rate). $ aligned [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 p_ [rgb]0.1602,0.6719,0.8867 θ( x)_ [rgb]0,0,0log evidence&\;\;=\;\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0-F_ [rgb]0,0,0ELBO\;+\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z x) )_ [rgb]0,0,0 subarraycinference KL: original objective\\ (gap to the true posterior) subarray\\[11.38109pt] F&\;\;=\;\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0E_q_ [rgb]0.9297,0.1094,0.1406 u[- p_ [rgb]0.1602,0.6719,0.8867 θ( x z)]_ [rgb]0,0,0reconstruction term (accuracy)\;+\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z) )_ [rgb]0,0,0prior KL (coding rate)~. aligned$ (23) The inference KL compares the approximate posterior with the true posterior, p(∣)p_ [rgb]0.1602,0.6719,0.8867 θ( z x). This is the original objective in eq. 21. The prior KL, by contrast, compares the approximate posterior with p()p_ [rgb]0.1602,0.6719,0.8867 θ( z), before conditioning on the observation. It is a component of free energy, and can be interpreted as a coding rate: the cost of representing the input-dependent distribution relative to the prior (Alemi et al., 2018). We use natural logarithms, so this cost is measured in nats. These quantities need not decrease together. If inference starts at the prior, the prior KL is initially zero. Explaining an informative observation generally requires moving away from that prior, increasing the coding rate while reducing the inference KL. Free energy balances this cost against the improvement in expected log likelihood. Minimizing the prior KL alone would simply return the prior and ignore the observation. Free energy and posterior approximation. Rearranging the first line of eq. 23 yields ℱ(,,)=(q(∣)∥p(∣))−logp().F( x; [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)\;=\;D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z x) )\;-\; p_ [rgb]0.1602,0.6719,0.8867 θ( x)~. (24) The evidence depends on the generative model and observation, but not on the variational parameters. Therefore, with x and [rgb]0.1602,0.6719,0.8867 θ fixed, every decrease in ℱF is an equal decrease in the inference KL (Blei et al., 2017). Equivalently, we have: ∇ℱ=∇(q(∣)∥p(∣)). _ [rgb]0.9297,0.1094,0.1406 uF\;=\; _ [rgb]0.9297,0.1094,0.1406 uD_ KL (q_ [rgb]0.9297,0.1094,0.1406 u( z x)\, \|\,p_ [rgb]0.1602,0.6719,0.8867 θ( z x) )~. (25) This is the sense in which free energy minimization improves inference. It holds regardless of whether the variational parameters are produced by an amortized encoder or refined iteratively. The bound becomes tight, ℱ=−logp()F=- p_ [rgb]0.1602,0.6719,0.8867 θ( x), only when the approximate posterior equals the true posterior. A restricted variational family may not contain that posterior (Wainwright and Jordan, 2008; Blei et al., 2017); for example, a factorized Bernoulli distribution cannot represent all the dependencies induced by overlapping dictionary atoms. Optimization can reduce the KL without eliminating this approximation gap. A.4 Trust regions and natural-gradient inference Inference moves the approximate posterior q_ [rgb]0.9297,0.1094,0.1406 u to reduce free energy. But specifying the objective is only part of specifying the dynamics. We also need to decide how far beliefs should move at each step, and what counts as a small change. This choice of distance is what takes us from ordinary gradient descent to the Fisher geometry used by FOND (Vafaii et al., 2025). Throughout this derivation, we hold the observation x, generative parameters [rgb]0.1602,0.6719,0.8867 θ, and prior reference fixed. We write the objective as ℱ()F( [rgb]0.9297,0.1094,0.1406 u) and suppress conditioning on x in q_ [rgb]0.9297,0.1094,0.1406 u. From a gradient to a step. At the current iterate t [rgb]0.9297,0.1094,0.1406 u_t, define t≔∇ℱ(t) g_t _ [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u_t). For a small displacement Δ [rgb]0.9297,0.1094,0.1406 u, the gradient predicts the change in free energy: ℱ(t+Δ)=ℱ(t)+t⊤Δ+o(‖Δ‖2).F( [rgb]0.9297,0.1094,0.1406 u_t+ [rgb]0.9297,0.1094,0.1406 u)\;=\;F( [rgb]0.9297,0.1094,0.1406 u_t)\;+\; g_t [rgb]0.9297,0.1094,0.1406 u\;+\;o\! ( [rgb]0.9297,0.1094,0.1406 u _2 )~. (26) There is a catch: if t≠ g_t≠ 0, the linear term keeps decreasing as we move farther along −t- g_t. Taken literally, it would tell us to take an infinitely large step. Yet, the prediction in eq. 26 is only useful nearby. We can keep the step finite by attaching a quadratic cost to movement: Δ=argmint⊤⏟predictedchange in ℱ+12η‖22⏟cost of the step=−ηt,η> 0. [rgb]0.9297,0.1094,0.1406 u\;=\; δargmin\ \ [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 g_t δ_ [rgb]0,0,0 subarraycpredicted\\ change in F subarray\;+\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 12η δ _2^2_ [rgb]0,0,0cost of the step \\;=\;-η\, g_t~, η\;>\;0~. (27) The last equality follows by differentiating with respect to δ and setting t+/η= g_t+ δ/η= 0. This is ordinary gradient descent: a balance between the predicted improvement and the cost of moving away from the current iterate (Martens, 2020). Increasing η makes movement cheaper, allowing a longer step. The quadratic term belongs to the update rule; the approximation to ℱF itself remains first order. Trust-region formulation. We can express the same tradeoff as a movement budget (Boyd and Vandenberghe, 2004, Sec. 9.4). Allow a step only within a ball around the current iterate, and choose the one that most reduces the linearized objective: minimizet⊤subject to‖22≤ϵ. δminimize\ \; g_t δ to δ _2^2\;≤\;ε~. (28) The radius ϵ ε sets how far we are willing to trust the local prediction. For ϵ>0ε>0 and t≠ g_t≠ 0, the optimum uses the full budget, giving Δ=−ϵt/‖t‖2 [rgb]0.9297,0.1094,0.1406 u=- ε\, g_t/ g_t _2. It agrees with eq. 27 when ϵ=η2‖t‖22ε=η^2 g_t _2^2. Thus, a fixed step size η lets a large gradient spend more of the movement budget, while a fixed ϵε enforces the same radius at each step. Changing the geometry of the step. Euclidean distance measures every direction with the same ruler. More generally, a positive-definite matrix tM_t assigns the squared length ⊤t δ M_t δ to a displacement. We evaluate this metric at the current iterate and hold it fixed while choosing the step. Replacing the Euclidean penalty gives: Δ=argmint⊤+12η⊤t=−ηt−1t. [rgb]0.9297,0.1094,0.1406 u\;=\; δargmin\ \ g_t δ\;+\; 12η δ M_t δ \\;=\;-η\,M_t^-1 g_t~. (29) The corresponding constraint, ⊤t≤ϵ δ M_t δ≤ε, has solution Δ=−ϵt⊤t−1tt−1t,t≠ 0. [rgb]0.9297,0.1094,0.1406 u\;=\;- ε g_t M_t^-1 g_t\,M_t^-1 g_t~, g_t\;≠\; 0~. (30) The ball has become an ellipsoid. Along an eigenvector with a large metric eigenvalue, movement costs more, so the same gradient component produces a smaller update. This is why “steepest descent” needs a geometry: the best direction depends on how we measure the cost of getting there (Amari, 1998; Boyd and Vandenberghe, 2004). Measuring changes in beliefs. So far, we have measured movement in parameter space. But [rgb]0.9297,0.1094,0.1406 u is a description of a belief, and equal parameter changes can mean very different belief changes. For example, a small change in a Bernoulli logit shifts the activation probability most strongly near 1/21/2, and much less near 00 or 11. A unit of logit displacement therefore has no fixed meaning as a change in belief. We should price the step by how much it changes the distribution it represents. KL divergence gives us such a measure: (q+Δ∥q)=q+Δ[logq+Δ()q()].D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u+ [rgb]0.9297,0.1094,0.1406 u\, \|\,q_ [rgb]0.9297,0.1094,0.1406 u )\;=\;E_q_ [rgb]0.9297,0.1094,0.1406 u+ [rgb]0.9297,0.1094,0.1406 u[ q_ [rgb]0.9297,0.1094,0.1406 u+ [rgb]0.9297,0.1094,0.1406 u( z)q_ [rgb]0.9297,0.1094,0.1406 u( z)]~. (31) Hobson (1969) characterized KL as an information measure for finite-outcome belief revision; section A.5 develops this motivation further. Here, it measures the change between successive approximate posteriors. The prior KL in free energy has a different job: it helps define the objective by comparing the posterior with its prior reference (section A.3). The stepwise KL tells us how much we change our mind in one update. Many small updates can still carry beliefs far from the prior. From local KL divergence to the Fisher metric. For a smooth family with common positive support and sufficient regularity to differentiate under the expectation, define the score ()≔∇logq() s_ [rgb]0.9297,0.1094,0.1406 u( z) _ [rgb]0.9297,0.1094,0.1406 u q_ [rgb]0.9297,0.1094,0.1406 u( z). Normalization gives q[]=E_q_ [rgb]0.9297,0.1094,0.1406 u[ s_ [rgb]0.9297,0.1094,0.1406 u]= 0. When the two distributions coincide, the KL divergence and its first derivative vanish. Its leading contribution is therefore quadratic, and its Hessian is the Fisher information: u() _u( [rgb]0.9297,0.1094,0.1406 u)\; =q[()()⊤], =\;E_q_ [rgb]0.9297,0.1094,0.1406 u\! [ s_ [rgb]0.9297,0.1094,0.1406 u( z) s_ [rgb]0.9297,0.1094,0.1406 u( z) ]~, (32) (q+Δ∥q) _ KL (q_ [rgb]0.9297,0.1094,0.1406 u+ [rgb]0.9297,0.1094,0.1406 u\, \|\,q_ [rgb]0.9297,0.1094,0.1406 u )\; =12Δ⊤u()Δ+o(‖Δ‖22). =\; 12\, [rgb]0.9297,0.1094,0.1406 u G_u( [rgb]0.9297,0.1094,0.1406 u) [rgb]0.9297,0.1094,0.1406 u\;+\;o\! ( [rgb]0.9297,0.1094,0.1406 u _2^2 )~. (33) The Fisher metric measures how sensitive the distribution is to movement in each direction. A large Fisher length means the step makes the old and new beliefs easier to distinguish, spending more of the same KL budget. Reversing the KL arguments gives the same quadratic term (Martens, 2020): KL is asymmetric beyond this local approximation, but both orientations supply the same local ruler. There is also a statistical reason to favor this ruler. A sufficient statistic retains all the information about the parameter; replacing the full observation by that statistic should preserve statistical lengths. For regular minimal exponential families, Dowty’s Chentsov-type theorem makes this requirement precise. Together with the appropriate scaling for independent replication and additional compatibility and continuity assumptions, preservation by canonical sufficient statistics characterizes the Fisher metric up to a positive constant (Dowty, 2018). The natural-gradient update. Set t=u(t)M_t=G_u( [rgb]0.9297,0.1094,0.1406 u_t) in eq. 29. To leading order, the movement cost is now the KL divergence divided by η. Balancing this cost against the predicted decrease in free energy yields: t+1=t−ηu(t)−1∇ℱ(t). $ [rgb]0.9297,0.1094,0.1406 u_t+1\;=\; [rgb]0.9297,0.1094,0.1406 u_t\;-\;η\,G_u( [rgb]0.9297,0.1094,0.1406 u_t)^-1 _ [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u_t)~.$ (34) This is natural-gradient descent wherever the Fisher is nonsingular (Amari, 1998; Martens, 2020). Constraining (q+Δ∥q)≤ϵ/2D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u+ [rgb]0.9297,0.1094,0.1406 u\, \|\,q_ [rgb]0.9297,0.1094,0.1406 u )≤ε/2 and using its quadratic approximation gives the same direction, with the budget-dependent normalization in eq. 30. For the continuous flow ˙=−u−1∇ℱ [rgb]0.9297,0.1094,0.1406 u=-G_u^-1 _ [rgb]0.9297,0.1094,0.1406 uF, the rate of change of free energy is: dℱdt=−(∇ℱ)⊤u−1∇ℱ≤ 0. dFdt\;=\;- ( _ [rgb]0.9297,0.1094,0.1406 uF ) G_u^-1 _ [rgb]0.9297,0.1094,0.1406 uF\;≤\;0~. (35) Wherever the gradient is nonzero, free energy strictly decreases along the flow. After discretization, the step size still matters: a large jump can leave the region where the local prediction is accurate, so an arbitrary finite update need not decrease ℱF. A Gaussian inference example. To see what changing the ruler buys us, consider a Gaussian target posterior p(∣)=(,,)p( z x)=N( z; μ,S) and a variational family q()=(,,)q_ [rgb]0.9297,0.1094,0.1406 u( z)=N( z; [rgb]0.9297,0.1094,0.1406 u,S), with fixed ≻0S 0. In this illustration, [rgb]0.9297,0.1094,0.1406 u denotes the posterior mean, rather than the Bernoulli logits used in BeBeFOND. For the ordinary variational free energy (β=1β=1), eq. 24 gives: ℱ()−ℱ()=(q∥p(⋅∣))=12(−)⊤−1(−).F( [rgb]0.9297,0.1094,0.1406 u)-F( μ)\;=\;D_ KL (q_ [rgb]0.9297,0.1094,0.1406 u\, \|\,p(· x) )\;=\; 12\,( [rgb]0.9297,0.1094,0.1406 u- μ) S^-1( [rgb]0.9297,0.1094,0.1406 u- μ)~. (36) The score is −1(−)S^-1( z- [rgb]0.9297,0.1094,0.1406 u), giving u=−1G_u=S^-1. Hence: ∇ℱ _ [rgb]0.9297,0.1094,0.1406 uF =−1(−), =\;S^-1( [rgb]0.9297,0.1094,0.1406 u- μ)~, (37) u−1∇ℱ _u^-1 _ [rgb]0.9297,0.1094,0.1406 uF =−. =\; [rgb]0.9297,0.1094,0.1406 u- μ~. Suppose the target posterior is very certain along one direction, but uncertain along another. Along an eigenvector of S with eigenvalue λj _j, a Euclidean step multiplies the mean error by 1−η/λj1-η/ _j. The narrow, certain direction forces a small η to avoid overshooting; with that same step size, the broad direction barely moves. Convergence in all directions requires 0<η<2λmin()0<η<2 _ (S). The natural-gradient step instead removes the precision factor −1S^-1: every component of the mean error contracts by 1−η1-η. It corrects the same fraction of the remaining error in each direction, with convergence for 0<η<20<η<2. For this mean-only quadratic problem, η=1η=1 reaches the target mean in one step. Reparameterization invariance. The belief itself should not care whether we describe it by a probability, a logit, or some other smooth coordinate. Let =f() [rgb]0.9297,0.1094,0.1406 v=f( [rgb]0.9297,0.1094,0.1406 u) be a smooth invertible reparameterization, with Jacobian =∂J= ∂ [rgb]0.9297,0.1094,0.1406 v∂ [rgb]0.9297,0.1094,0.1406 u. The chain rule and the score definition give: ∇ℱ _ [rgb]0.9297,0.1094,0.1406 vF =−⊤∇ℱ, =\;J^- _ [rgb]0.9297,0.1094,0.1406 uF~, (38) v _v =−⊤u−1, =\;J^- G_uJ^-1~, v−1∇ℱ _v^-1 _ [rgb]0.9297,0.1094,0.1406 vF =u−1∇ℱ. =\;JG_u^-1 _ [rgb]0.9297,0.1094,0.1406 uF~. The gradient and metric change together, so the natural-gradient velocity obeys ˙=˙ [rgb]0.9297,0.1094,0.1406 v=J [rgb]0.9297,0.1094,0.1406 u, exactly as a velocity should. Wherever the flow is well defined, both coordinate systems describe the same evolution of distributions (Amari, 1998; Martens, 2020). This is the optimization counterpart of the covariance principle developed in section A.5. This invariance is a property of the continuous dynamics. A finite Euler step preserves it exactly under affine coordinate changes, but generally only to first order under nonlinear ones (Martens, 2020). The geometry fixes the direction of motion, and discretization determines how faithfully our finite steps follow it. A.5 “Physics of Adaptation”: deriving covariant inference and learning rules In this section, we take a physics-inspired approach to dynamics and motivate the use of natural gradients through this angle. Specifically, we argue that the dynamics of inference and learning can be derived from stationary-action principles (Landau and Lifshitz, 1975), supplemented by dissipative forces. We emphasize how the underlying geometry of the system must constrain our choice of the kinetic Lagrangian. For beliefs existing on statistical manifolds, we will introduce the kinetic energy of beliefs as a natural choice, and use it to derive natural gradient descent in the overdamped (massless) limit of belief dynamics. The principle of stationary action. Physics is fundamentally about dynamics. The most successful theories in physics—including Newtonian mechanics, special and general relativity, electromagnetism, and quantum field theory—share two core attributes. They all: (1) identify some dynamical objects of interest, and (2) derive their dynamics via the principle of stationary action: Principle of stationary action: The closest physics has come to a theory of everything Physical systems evolve to extremize action S (time integral of the Lagrangian L): S=∫Lt,→real physical systems followtrajectories that extremize actionδS=0S\;= \!L\,dt, \, subarraycreal physical systems follow\\ trajectories that extremize action subarray\; $ δ S=0$ (39) Euler–Lagrange equations. Extremizing the action leads directly to the Euler–Lagrange equations: δS=0⟹dt(∂L∂˙)−∂L∂= 0.δ S=0 ddt ( ∂ L∂ x )- ∂ L∂ x\;=\; 0~. (40) Here, x is position, ˙=d/dt x=d x/dt is velocity, and L is the Lagrangian. In classical mechanics, the Lagrangian is simply kinetic minus potential energy, L=T−VL=T-V, and the kinetic energy is given by Tclassical=12m‖˙‖22T_classical= 12m x _2^2, where m is the rest mass. As an illustrative example, apply eq. 40 to Lclassical(,˙)=Tclassical(˙)−V()L_classical( x, x)=T_classical( x)-V( x): dt(∂Lclassical∂˙)−∂Lclassical∂ ddt ( ∂ L_classical∂ x )- ∂ L_classical∂ x =dt(∂Tclassical∂˙)+∂V()∂ =\;\; ddt ( ∂ T_classical∂ x )+ ∂ V( x)∂ x (41) =dt(m˙)+∇V=m¨+∇V= 0. =\;\; ddt (m x )+∇ V\;\,=\,\;m x+∇ V\;\,=\;\, 0~. Defining force ≔−∇VF -∇ V and acceleration ≔¨ a x, we recover Newton’s second law, =mF=m a. Generalized Euler–Lagrange equations. Equation 40 describes the dynamics when all forces are represented in the Lagrangian. To include dissipation and other non-conservative forces not represented in L, we add a generalized force vector Q. This gives: dt(∂L∂˙)−∂L∂=. ddt ( ∂ L∂ x )- ∂ L∂ x\;=\;Q~. (42) Defining the canonical momentum ≔∂L/∂˙ p ∂ L/∂ x, the generalized Euler-Lagrange equation is: dt−∂L∂=. $ ddt\, p\;-\; ∂ L∂ x\;=\;Q~.$ (43) Adaptation dynamics. The dynamics of various physical systems can be derived from eq. 43, by making different choices for the Lagrangian L and the non-conservative forces Q. We argue that the dynamics of adaptation is no different. We treat the belief state (t) [rgb]0.9297,0.1094,0.1406 u(t) as the position of a fictitious particle moving in a potential V()V( [rgb]0.9297,0.1094,0.1406 u), which we identify as the loss function. The same picture applies to the generative parameters (t) [rgb]0.1602,0.6719,0.8867 θ(t) during learning with potential V()V( [rgb]0.1602,0.6719,0.8867 θ). FOND unifies both processes as variational free energy minimization, giving ℱ(,)≡V(,)F( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)≡ V( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ) as one shared potential energy: inference changes [rgb]0.9297,0.1094,0.1406 u with [rgb]0.1602,0.6719,0.8867 θ held fixed, while learning changes [rgb]0.1602,0.6719,0.8867 θ with [rgb]0.9297,0.1094,0.1406 u held fixed. This leaves two modeling choices open: the kinetic part of the Lagrangian T, and the dissipation force Q. Different choices lead to different inference and learning dynamics and, after discretization, different optimization algorithms, as we demonstrate below. The principle of general covariance. With the potential energy fixed by the loss, let us turn to the kinetic part of the Lagrangian. How should we choose it? Should we just do a bunch of trial and error? Or is there a principled way to approach this? Physics gives us a basic requirement: The principle of general covariance “The general laws of nature are to be expressed by equations which hold good for all systems of co-ordinates, that is, are co-variant with respect to any substitutions whatever (generally co-variant).” — Einstein (1916) Imagine tossing a ball across a room. If two people track its motion—one standing upright and one tilted sideways—they will record different numbers for its horizontal and vertical motion. Yet, both of them watch the ball land in the exact same spot. The ball has no idea about the observers watching it; it follows a physical path that does not change simply because different observers choose to describe it in different ways. More concretely, consider Newton’s second law, =mF=m a. When you change your coordinate system, the actual numerical values of force and acceleration change, but the mathematical form of the law does not: =mF=m a becomes ′=m′F =m a because both sides “co-vary” in lockstep under the transformation. In this sense, Newton’s law describes a “co-variant” dynamical rule under Galilean transformations. If rotating or shifting our reference frame altered where the ball landed, the equation wouldn’t be describing the ball at all—it would merely be describing our choice of coordinates. Descartes, geometry, and algebra. This idea is not new. It goes back to Descartes’ link between geometry and algebra: choose a coordinate system, and a point becomes a tuple of numbers. We have done this so often that it is easy to forget what comes with it: redundancy in the description. The origin, the orientation of the axes, and the units are all choices of the modeler. The same physical trajectory admits many different numerical descriptions. Covariance requires the equations to transform consistently between them, so that they all predict the same physical motion. The principle of general covariance for inference and learning. Once again, we argue that the dynamics of adaptation are no different. A Bernoulli belief can be described by its activation probability or by its logit. These describe exactly the same distribution. Why should the learner’s response to evidence change just because we switched between these two descriptions? We require the continuous dynamics of inference and learning to describe the same evolution of distributions under smooth, invertible reparameterizations (Amari, 1998; Martens, 2020). There is also a concrete computational reason for this requirement. Take a scalar belief parameter u [rgb]0.9297,0.1094,0.1406u and write it instead as v=10u [rgb]0.9297,0.1094,0.1406v=10 [rgb]0.9297,0.1094,0.1406u. If we apply ordinary Euclidean gradient descent in v [rgb]0.9297,0.1094,0.1406v, the gradient becomes ten times smaller, and converting the step back to u [rgb]0.9297,0.1094,0.1406u introduces another factor of ten. Therefore, with the same learning rate, the update in u [rgb]0.9297,0.1094,0.1406u becomes 100× smaller. Same belief, same evidence, same objective, but a different pace of adaptation, just because we changed the units. We can fix this scalar example by re-tuning the learning rate. But when different coordinates are rescaled differently, a single learning rate cannot generally restore all their relative update speeds. Such artificial conditioning can force small steps for stability in some directions while leaving others barely moving (Martens, 2020). With a limited inference or training budget, this can determine which beliefs or features we have time to learn. We may then spend extra computation, and extra stochastic training samples, compensating for arbitrary scales chosen by the engineer. A coordinate choice should not get to decide how well we can adapt. Therefore, we draw inspiration from Einstein (1916) and postulate the same thing for physics of adaptation: The principle of general covariance, applied to inference and learning The physics of adaptation must not depend on your choice of coordinate system. To implement this principle, we need a measure of motion that remain invariant under coordinate transformations. This is what motivates the kinetic energy of beliefs below. The kinetic energy of beliefs. We can arrive at a suitable kinetic term from two directions. The first is geometric. Suppose some parameters λ live on a manifold with Riemannian metric ()G( λ). What is the simplest scalar, quadratic in velocity, that measures how fast they move? The metric gives us the squared speed, so for the kinetic energy, the natural choice is: T(,˙)=12m˙⊤()˙. $ T( λ, λ)\;=\; 12\,m\, λ G( λ) λ~.$ (44) Here, ˙=d/dt λ=d λ/dt, and m>0m>0 is a scalar that sets the relative scale of kinetic and potential energy. This expression is nonnegative, vanishes at rest, and treats motion in either direction equally. Crucially, it is invariant under a change of coordinates: the velocity and metric transform together, leaving their contraction unchanged. In Euclidean space, =G=I, and we recover the familiar 12m∥˙∥22 12m λ _2^2. This gives us the correct form of a kinetic energy on Riemannian manifolds, but we still need to choose the geometry of belief space. The second route starts with a more physical question: what does it cost to change one’s mind? Zellner (1988) derived Bayes’ rule as an optimal information processor, with no information lost between its inputs and outputs. Later, Zellner (2002) showed that introducing costs of changing beliefs can alter the optimal update, producing dynamic partial adjustment rather than an immediate revision. This is the distinction we need: specifying the posterior we ought to reach does not tell us how a physical learner gets there, or what the journey costs. We take the next step by giving beliefs a kinetic term, representing their resistance to change as inertia. For a brain, this cost has a concrete physical basis. Brains are made of matter, and changing their state takes energy. Fast changes in neural activity involve ion flows across cell membranes, and restoring the ion gradients costs ATP (Attwell and Laughlin, 2001). Slower changes in synaptic strength can involve protein synthesis and physical remodeling of dendritic spines, supported by local metabolic resources (Rangaraju et al., 2019). Inference and learning take place through these physical processes. Why should the beliefs they implement be free to move? To quantify this motion, let q(t)q_ λ(t) represent the learner’s current belief. How much information separates it from the belief a short time h later? Hobson (1969) asked how to measure the information in a message that changes probabilities from P to Q. He made a list of intuitive requirements: the measure should be continuous, independent of how outcomes are labeled, zero when nothing changes, and larger when a message rules out more equally likely possibilities. It should also obey a composition rule: describing an outcome in stages should give the same information as describing it all at once. For finite outcome spaces, these requirements uniquely select (Q∥P)D_ KL (Q\, \|\,P ), up to a positive constant. Its continuous counterpart retains coordinate invariance because the density Jacobians cancel in the ratio. Following (Hobson, 1969), we measure the information in a local belief change by: ℐt(h)≔(q(t+h)∥q(t)).I_t(h)\; \;D_ KL (q_ λ(t+h)\, \|\,q_ λ(t) )~. (45) For a regular, smooth family with common support, the KL divergence and its first derivative vanish when the two distributions coincide. Its leading contribution is therefore quadratic, with the Fisher information as its Hessian (Amari and Nagaoka, 2000; Amari, 2016): (q+Δ∥q) _ KL (q_ λ+ λ\, \|\,q_ λ ) =12Δ⊤()Δ+o(‖Δ‖2), =\; 12\, λ G( λ) λ\;+\;o\! ( λ ^2 ), (46) () ( λ) =q[(∇logq)(∇logq)⊤]. =\;E_q_ λ\! [ ( _\! λ q_ λ ) ( _\! λ q_ λ ) ]~. Along a smooth trajectory, Δ=h˙+O(h2) λ=h λ+O(h^2). Dividing the local information gain by h2h^2 and taking the limit gives what we call the kinetic energy of beliefs: Tbelief(,˙)≔limh→0ℐt(h)h2=12d2ℐt(h)dh2|h=0=12˙⊤()˙. $ aligned T_belief( λ, λ)\;& \; _h→ 0 I_t(h)h^2\\[5.0pt] \;&=\; 12 . d^2I_t(h)dh^2 |_h=0\;=\; 12\, λ G( λ) λ~. aligned$ (47) The two arguments yield the same quadratic form (eqs. 47 and 44). From the geometry argument, we derived a squared speed, and information processing supplied the Fisher metric that measures it. This speed tracks changes in the distribution itself, so it is invariant to rescaling the parameters. Doubling it quadruples TbeliefT_belief, just as doubling a classical particle’s speed quadruples TclassicalT_classical. At this stage, TbeliefT_belief has units of information per time squared. Multiplying by m gives the kinetic term T=mTbeliefT=mT_belief in eq. 44, with m()mG( λ) acting as a position-dependent mass matrix. We interpret m as setting the price of belief motion relative to reducing the loss. For a biological learner, this suggests a metabolic interpretation: moving beliefs faster requires more physical resources. Whether the brain’s ATP expenditure follows this precise quadratic form, and whether it shares the Fisher metric’s directional dependence, are open questions. We adopt this as a speculative interpretation of the kinetic term. The geometric derivation (eq. 44) stands on its own. Before we move on, we briefly mention a direct connection to recent stochastic thermodynamics. Auconi and Ito (2026) study the intrinsic speed of an evolving distribution, whose square is the temporal Fisher information, vinfo2=qt[(∂tlogqt)2]=2Tbeliefv_info^2=E_q_t\! [( _t q_t)^2 ]=2T_belief. For Markov jump and Fokker–Planck dynamics near a steady state, they relate 2vinfo22v_info^2 to the second time derivative of the KL divergence toward that state, and use their difference to bound steady-state entropy production. Thus, the same geometric measure of motion appears in physical constraints on dissipation. Finally, this construction (eqs. 47 and 44) applies to both sides of FOND: inference uses the Fisher metric of the variational posterior, uG_u, while learning uses that of the generative joint distribution, θG_θ. With the kinetic term chosen, we can now turn to the remaining ingredient: dissipation. Dissipative equations of motion for beliefs. We now have both parts of the Lagrangian, L=T−VL=T-V. Below, we will address the last missing piece: dissipation force Q. But first, we make the Lagrangian and the canonical momentum explicit. Let gijg_ij denote the entries of ()G( λ) and gijg^ij those of its inverse, with repeated indices summed. During each inference or learning phase, we hold the other parameter block fixed and take V and G to have no explicit time dependence. The Lagrangian and canonical momentum are then: L(,˙) L( λ, λ) =m2gij()λ˙iλ˙j−V(), =\; m2\,g_ij( λ) λ^i λ^j-V( λ), (48) pi p_i ≔∂L∂λ˙i=mgij()λ˙j,λ˙i=1mgij()pj. \; ∂ L∂ λ^i\;=\;m\,g_ij( λ) λ^j, λ^i= 1m\,g^ij( λ)p_j~. For dissipation, a reasonable standard choice is linear damping proportional to canonical momentum, Qi≔−pi/τdampQ_i -p_i/ _damp, where τdamp>0 _damp>0 is a constant damping time. This is isotropic with respect to the belief-space metric: every direction has the same damping rate when measured in that geometry. It introduces no preferred direction beyond the metric we already chose. Other dissipation laws are possible, but we use this one as a simple, covariant starting point. This force follows from the quadratic Rayleigh dissipation function: ℛ(,˙) ( λ, λ) =Tτdamp=m2τdamp˙⊤()˙, =\; T _damp\;=\; m2 _damp\, λ G( λ) λ, (49) Qi Q_i =−∂ℛ∂λ˙i=−piτdamp. =- ∂ λ^i=- p_i _damp~. It removes energy at a rate d(T+V)/dt=−2T/τdamp≤0d(T+V)/dt=-2T/ _damp≤ 0. Thus, the learner can lose momentum and settle, rather than oscillating or orbiting indefinitely under conservative forces. Substituting L and QiQ_i into the generalized Euler–Lagrange equation (eq. 43) gives: dpidt−m2∂igjkλ˙jλ˙k+∂iV=−piτdamp, dp_idt\,-\, m2\, _ig_jk\, λ^j λ^k\,+\, _iV\;=\;- p_i _damp~, (50) where ∂i≔∂/∂λi _i ∂/∂λ^i. Using λ˙i=gijpj/m λ^i=g^ijp_j/m and the inverse-metric identity ∂igjk=−gja(∂igab)gbk _ig^jk=-g^ja( _ig_ab)g^bk, we obtain the full, dissipative equations of motion: dpidt⏟inertia+12m∂igjkpjpk⏟geometric corrections+piτdamp⏟dissipation=−∂iV⏟task force. $ [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 dp_idt_ [rgb]0,0,0inertia\;+\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 12m\, _ig^jk\,p_jp_k_ [rgb]0,0,0geometric corrections\;+\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0 p_i _damp_ [rgb]0,0,0dissipation\;\;=\;\; [rgb]0.6016,0.6016,0.6016 [rgb]0,0,0- _iV_ [rgb]0,0,0task force~.$ (51) The geometric corrections arise because the metric changes as the learner moves. They are needed for the inertial dynamics to transform consistently between coordinate systems (i.e., to remain covariant). Equivalently, writing the equations in terms of velocity gives: m(λ¨i+Γjkiλ˙jλ˙k)+mτdampλ˙i=−gij∂jV,m ( λ^i+ ^i_jk λ^j λ^k )\,+\, m _damp\, λ^i\;=\;-g^ij _jV~, (52) where Γjki=12giℓ(∂jgℓk+∂kgℓj−∂ℓgjk) ^i_jk= 12g^i ( _jg_ k+ _kg_ j- _ g_jk) are the Christoffel symbols of the metric. The acceleration and geometric correction together describe acceleration on the statistical manifold. The overdamped (massless) limit. To recover first-order adaptation dynamics, we assume that momentum relaxes much faster than beliefs change. We make this precise by taking: m→ 0,τdamp→ 0,mτdamp→ζ∈(0,∞). $ m\,→\,0, _damp\,→\,0, m _damp\,→\,ζ∈(0,∞)~.$ (53) The last condition matters: inertia vanishes while friction remains finite. Heuristically, we can think of this as a “massless” learner. Imagine pushing a feather through thick honey. Its motion responds almost immediately to the balance between your push and the viscous drag, with very little coasting from its own momentum. The same intuition applies here: beliefs keep moving under the task force, but retain negligible momentum on the timescale of adaptation. Consider the slower dynamics after the initial rapid relaxation of momentum, in a smooth region where the metric is nonsingular and the velocity, acceleration, and metric derivatives remain bounded as m decreases. Then pi=O(m)p_i=O(m), so both dpi/dtdp_i/dt and the geometric correction in eq. 51 are O(m)O(m). The damping term, however, approaches ζgijλ˙jζ g_ij λ^j and remains finite. The surviving force balance is: ζ()˙=−∇V().ζ\,G( λ) λ\;=\;- _\! λV( λ)~. (54) Natural-gradient dynamics and FOND’s dual flows. Solving this balance for velocity gives: ˙=−1ζ()−1∇V(). $ λ\;=\;- 1ζ\,G( λ)^-1 _\! λV( λ)~.$ (55) This is natural-gradient flow. The inverse metric appears because friction converts force into velocity through the geometry of belief space. The geometric corrections have vanished with inertia, but the geometry itself remains. Setting V=ℱV=F and applying this result separately to inference and learning gives the two FOND natural-gradient flows: ˙=−1ζuu()−1∇ℱ(,),˙=−1ζθθ()−1∇ℱ(,). $ aligned [rgb]0.9297,0.1094,0.1406 u\;&=\;- 1 _u\,G_u( [rgb]0.9297,0.1094,0.1406 u)^-1 _\! [rgb]0.9297,0.1094,0.1406 uF( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)~,\\[4.0pt] [rgb]0.1602,0.6719,0.8867 θ\;&=\;- 1 _θ\,G_θ( [rgb]0.1602,0.6719,0.8867 θ)^-1 _\! [rgb]0.1602,0.6719,0.8867 θF( [rgb]0.9297,0.1094,0.1406 u, [rgb]0.1602,0.6719,0.8867 θ)~. aligned$ (56) Here, ζu,ζθ>0 _u, _θ>0 set the pace of the respective processes and can be absorbed into their inference and learning timescales. Each flow uses its own Fisher metric, with the other parameter block held fixed. Deriving the Bayesian learning rule (Khan and Rue, 2023, BLR,). An Euler discretization of eq. 55, with η≔Δt/ζη t/ζ, gives the familiar update: k+1=k−η(k)−1∇V(k). $ λ_k+1\;=\; λ_k\;-\;η\,G( λ_k)^-1 _\! λV( λ_k)~.$ (57) This is just natural gradient descent, and replacing the gradient by a minibatch estimate gives its stochastic version. For exponential-family variational beliefs, writing eq. 57 in natural parameters and setting V=ℱV=F recovers the Bayesian learning rule (BLR) of Khan and Rue (2023). By varying the candidate distributions and the approximations to their natural gradients, BLR unified learning rules ranging from SGD and Newton’s method to Kalman filtering. Here, we have derived the BLR from the physics of adaptation: it emerges as the overdamped limit of covariant belief dynamics. Table 1 summarizes how the same update gives familiar learning rules under different choices of metric and loss. Table 1: Special cases of eq. 57 under different choices of metric G and loss V. The BLR row uses exponential-family variational beliefs in natural parameters. Algorithm Metric G Loss V Gradient Descent (Cauchy, 1847, GD,) I (Euclidean) Generic task loss Natural Gradient Descent (Amari, 1998, NGD,) Fisher–Rao Generic task loss Bayesian Learning Rule (Khan and Rue, 2023, BLR,) Fisher–Rao ℱF (free energy) Summary. Let us take stock of what we have done. We started from a physics-centered view of “dynamics is everything,” and demonstrated how dynamics can be derived from applying the principle of stationary action plus dissipation. We then identified the loss as a potential, required adaptation to be independent of arbitrary coordinates (principle of general covariance), and used local information gain to construct the kinetic energy of beliefs. A standard choice of friction then gave us the full dissipative equations of motion, including inertia and geometric corrections. Natural-gradient inference and learning followed in the overdamped (massless) limit, where that inertia is negligible. Physics of adaptation as a broad framework for deriving inference and learning dynamics. We could have motivated the same practical update through a KL trust region (section A.4), or appealed directly to Amari’s natural gradient (Amari, 1998). Those arguments are enough to justify the rule. We hope this route also reveals the richer framework underneath it: beliefs moving through a statistical manifold, driven by a task and constrained by the physical cost of adaptation. Future work should consider deriving other dynamics within that same framework by keeping inertia, changing the dissipation law, or considering more generalized belief states with different, perhaps adaptive, geometries that evolve in time in response to experience. Appendix B Extended FOND model and dynamics B.1 Model and objective Bernoulli instantiation of FOND. Recovering a feature dictionary requires two coupled operations: infer which features explain each observation, and use those explanations to improve the dictionary. FOND derives the first operation as recurrent natural-gradient descent on variational free energy (Vafaii et al., 2025). Its original Poisson instantiation learns the decoder by backpropagating through the inference trajectory. Here we extend the natural-gradient prescription to learning itself. Inference changes the approximate posterior q_ [rgb]0.9297,0.1094,0.1406 u, and learning changes the generative distribution p_ [rgb]0.1602,0.6719,0.8867 θ. The two natural-gradient flows use their respective Fisher metrics to measure local changes in these distributions (Amari, 1998; Martens, 2020). Section A.4 develops this geometric motivation. Bernoulli features and a Gaussian decoder. In BeBeFOND, each latent indicates whether a feature is present, and its decoder atom specifies that feature’s contribution to the observation. The posterior mean ri [rgb]0.9297,0.1094,0.1406r_i therefore expresses confidence in feature presence. For K binary latents and an M-dimensional observation, the local inference model is p() p( z) =∏i=1KBernoulli(zi;πref,i), = _i=1^KBernoulli(z_i; [rgb]0.1602,0.6719,0.8867π_ref,i), q(∣) q_ [rgb]0.9297,0.1094,0.1406 u( z x) =∏i=1KBernoulli(zi;ri), = _i=1^KBernoulli(z_i; [rgb]0.9297,0.1094,0.1406r_i), [rgb]0.9297,0.1094,0.1406 r =σ(), =σ( [rgb]0.9297,0.1094,0.1406 u), p(∣) p( x z) =(,,), =N( x; [rgb]0.1602,0.6719,0.8867 z, [rgb]0.1602,0.6719,0.8867 ), (58) where [rgb]0.1602,0.6719,0.8867 is diagonal and positive. Equivalently, =+ϵ x= [rgb]0.1602,0.6719,0.8867 z+ ε with ϵ∼(0,) ε (0, [rgb]0.1602,0.6719,0.8867 ): the noise describes variation not explained by the selected atoms, while the atom itself sets the amplitude and direction of a present feature. The learned initial prior 0 [rgb]0.1602,0.6719,0.8867 π_0 supplies the first reference, and the rolling rule in section B.1.1 updates that reference as inference proceeds. A learned decoder bias accounts for a nonzero observation mean by replacing x with − x- [rgb]0.1602,0.6719,0.8867 b. Section B.3 gives its update and inference-response correction. Conditioning on the observation couples otherwise independent latent features: the likelihood contains cross-terms Wijzizj [rgb]0.1602,0.6719,0.8867W_ijz_iz_j, where Wij=ϕi⊤−1ϕj [rgb]0.1602,0.6719,0.8867W_ij= [rgb]0.1602,0.6719,0.8867 φ_i [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 φ_j. When atoms overlap, the evidence for one feature depends on which others explain the input. Factorizing q restricts the represented posterior dependencies, but each optimized mean can still depend on the other means through the shared reconstruction (section B.2). The Poisson and Bernoulli choices both admit analytic moments and Fisher metrics (Wainwright and Jordan, 2008). In canonical coordinates, latent familysupportriGu,ii=Varq(zi)Poisson0,1,2,…euiriBernoulli0,1σ(ui)ri(1−ri) array[]c|c|c|clatent family&support& [rgb]0.9297,0.1094,0.1406r_i&G_u,i=Var_q(z_i)\\ &\0,1,2,…\&e [rgb]0.9297,0.1094,0.1406u_i& [rgb]0.9297,0.1094,0.1406r_i\\ Bernoulli&\0,1\&σ( [rgb]0.9297,0.1094,0.1406u_i)& [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) array (59) The variance in the last column enters both the expected reconstruction loss and the geometry of inference. This is the connection that allows the distributional assumptions to determine the dynamics. Free-energy decomposition. Section A.3 derives the fixed-prior ELBO identity. Here we evaluate its two terms for the Bernoulli–Gaussian model, holding the current reference fixed. A distribution over codes assigns probability to several possible reconstructions. Write =− δ= z- [rgb]0.9297,0.1094,0.1406 r and =− e= x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r. Since q[]=0E_q[ δ]=0 and q[⊤]=E_q[ δ δ ]= [rgb]0.9297,0.1094,0.1406C, expanding the square gives: q[‖−‖−12] _q\! [\| x- [rgb]0.1602,0.6719,0.8867 z\|_ [rgb]0.1602,0.6719,0.8867 ^-1^2 ] =⊤−1−2⊤−1q[]+q[⊤] \;=\; e [rgb]0.1602,0.6719,0.8867 ^-1 e-2 e [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 E_q[ δ]+E_q[ δ [rgb]0.1602,0.6719,0.8867W δ] =‖−12+Tr(), =\| e\|_ [rgb]0.1602,0.6719,0.8867 ^-1^2+ ( [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406C), (60) where =⊤−1 [rgb]0.1602,0.6719,0.8867W= [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 and =diag[⊙(−)] [rgb]0.9297,0.1094,0.1406C=diag[ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)]. The cross-term vanishes because the centered code has zero mean. For the KL term, factorization turns the log ratio of products into a sum: DKL(q∥p) D_KL(q\|p) =∑iq[zilogriπref,i+(1−zi)log1−ri1−πref,i] = _iE_q\! [z_i [rgb]0.9297,0.1094,0.1406r_i [rgb]0.1602,0.6719,0.8867π_ref,i+(1-z_i) 1- [rgb]0.9297,0.1094,0.1406r_i1- [rgb]0.1602,0.6719,0.8867π_ref,i ] =∑i[rilogriπref,i+(1−ri)log1−ri1−πref,i]. = _i [ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_i [rgb]0.1602,0.6719,0.8867π_ref,i+(1- [rgb]0.9297,0.1094,0.1406r_i) 1- [rgb]0.9297,0.1094,0.1406r_i1- [rgb]0.1602,0.6719,0.8867π_ref,i ]. (61) A decoder that reconstructs x from [rgb]0.9297,0.1094,0.1406 r can therefore still have a large expected error if the individual latent configurations produce different observations. Including this uncertainty cost and the KL penalty gives ℱ= = 12(−)⊤−1(−)+12∑iWiiri(1−ri)+12logdet 12( x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r) [rgb]0.1602,0.6719,0.8867 ^-1( x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r)+ 12 _i [rgb]0.1602,0.6719,0.8867W_i [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)+ 12 [rgb]0.1602,0.6719,0.8867 +β∑i[rilogriπref,i+(1−ri)log1−ri1−πref,i]+const, +β _i [ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_i [rgb]0.1602,0.6719,0.8867π_ref,i+(1- [rgb]0.9297,0.1094,0.1406r_i) 1- [rgb]0.9297,0.1094,0.1406r_i1- [rgb]0.1602,0.6719,0.8867π_ref,i ]+const, (62) with =⊤−1 [rgb]0.1602,0.6719,0.8867W= [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 . The log determinant normalizes the Gaussian likelihood, and β weights departure from the current prior reference. At β=1β=1 and a fixed generative prior, this is the negative evidence lower bound. The Bernoulli uncertainty term also cancels the quadratic self-interaction in each coordinate, leaving reconstruction affine in ri [rgb]0.9297,0.1094,0.1406r_i when the other means are fixed (section B.4). Overlapping atoms still couple different coordinates, so a factorized posterior does not imply independent inference updates. The moments needed for inference and learning are available without sampling latent codes: [rgb]0.9297,0.1094,0.1406C =diag[⊙(−)], =diag[ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)], [rgb]0.9297,0.1094,0.1406E =+⊤, = [rgb]0.9297,0.1094,0.1406C+ [rgb]0.9297,0.1094,0.1406 r [rgb]0.9297,0.1094,0.1406 r , p [rgb]0.1602,0.6719,0.8867E_p =diag[ref⊙(−ref)]+refref⊤, =diag[ [rgb]0.1602,0.6719,0.8867 π_ref ( 1- [rgb]0.1602,0.6719,0.8867 π_ref)]+ [rgb]0.1602,0.6719,0.8867 π_ref [rgb]0.1602,0.6719,0.8867 π_ref , Vmm [rgb]0.9297,0.1094,0.1406V_m =(−)m2+(⊤)mm. =( x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r)_m^2+( [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406C [rgb]0.1602,0.6719,0.8867 )_m. (63) Here [rgb]0.9297,0.1094,0.1406E and p [rgb]0.1602,0.6719,0.8867E_p are the posterior and prior second moments, and Vmm [rgb]0.9297,0.1094,0.1406V_m is the posterior-expected squared residual in observed dimension m. Bernoulli inference in BeBeFOND. For our factorized Bernoulli posterior, [rgb]0.9297,0.1094,0.1406 u contains logits and ri=σ(ui) [rgb]0.9297,0.1094,0.1406r_i=σ( [rgb]0.9297,0.1094,0.1406u_i). The score is − z- [rgb]0.9297,0.1094,0.1406 r, so u=diag(⊙(−)),Δ⊤uΔ=∑iri(1−ri)(Δui)2.G_u=diag\! ( [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r) ), [rgb]0.9297,0.1094,0.1406 u G_u [rgb]0.9297,0.1094,0.1406 u= _i [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)( [rgb]0.9297,0.1094,0.1406u_i)^2. (64) For a fixed small logit displacement, the distribution changes less when a Bernoulli probability is near zero or one. In probability coordinates, the same infinitesimal length is ∑i(dri)2/[ri(1−ri)] _i(d [rgb]0.9297,0.1094,0.1406r_i)^2/[ [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)]: a fixed probability displacement is more costly near the boundary. For finite logits, all these variances are positive. Applying the inverse Fisher cancels the variance factor in the Euclidean logit gradient, yielding the inference dynamics derived in section B.2. The update therefore accounts for how much a parameter movement changes the inferred beliefs. Natural-gradient inference and learning. For Bernoulli logits, the Fisher is u=G_u= [rgb]0.9297,0.1094,0.1406C. The same factor ri(1−ri) [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) multiplies the ordinary logit gradient, so applying u−1G_u^-1 removes this attenuation from the inference drive. The remaining terms combine evidence from the input, recurrent interactions through the dictionary, and relaxation toward the reference. Section B.2 derives these terms and their cancellation explicitly. Learning then holds the inferred moments fixed. The decoder gradient compares the observed input–latent moment ⟨⊤⟩ x [rgb]0.9297,0.1094,0.1406 r with the model prediction [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406E. Its Fisher uses the prior second moment, Φ=p⊗−1G_ = [rgb]0.1602,0.6719,0.8867E_p [rgb]0.1602,0.6719,0.8867 ^-1, giving a correction by p−1 [rgb]0.1602,0.6719,0.8867E_p^-1 rather than the posterior’s diagonal Fisher. This correction is what can compensate for weak learning signals in rarely active directions. Exact compensation requires matching the minibatch-averaged second moments (section C.4). The covariance update tracks expected squared residuals, while the learned initial prior tracks detached posterior means across inputs. These decoder, covariance, and initial-prior updates use the inferred statistics directly, without backpropagating through the inference trajectory. Their explicit flows and derivations appear in sections B.2 and B.2.1. The decoder evolves in learning time τ over an interval Δt t, and h denotes the separate inference integration interval. B.1.1 Online prior and working memory Consider observing a scene whose underlying features persist. After the first observation, the inferred distribution contains information about that particular scene that the population prior did not contain. The next update should start from this informed belief. FOND’s rolling rule makes it the new prior reference (Vafaii et al., 2025): the KL penalty now measures departure from the current prediction, while the likelihood supplies evidence for revising it. The learned initial prior starts the sequence, and the rolling reference carries its evolving, input-specific belief forward. Figure 5 shows this predict–update sequence and its fixed-input specialization. Section B.5 derives its activation threshold, and section C.3.3 reports the rolling-prior experiments varying β and the initial-prior cap. GENERATIVE MODEL PARAMETERS≡(,,0) [rgb]0.1602,0.6719,0.8867 θ≡( , , π_0)ROLLING UPDATEpref,t=qt−1 [rgb]0.1602,0.6719,0.8867p_ref,t= [rgb]0.9297,0.1094,0.1406q_t-10z_01z_1t−1z_t-1tz_t0x_01x_1t−1x_t-1tx_t0 π_0pref,1=q0 [rgb]0.1602,0.6719,0.8867p_ref,1= [rgb]0.9297,0.1094,0.1406q_0⋯·s⋯·s, , q0q_0, , q1q_1, , qt−1q_t-1, , qtq_t0=1=⋯=T=x_0=x_1=·s=x_T=xFixed inputℱt(t,t,ref,t)=ℒrec(t,t)+βDKL(qt∥pref,t) aligned F_t (x_t; [rgb]0.9297,0.1094,0.1406u_t, [rgb]0.1602,0.6719,0.8867u_ref,t )&=L_rec (x_t; [rgb]0.9297,0.1094,0.1406r_t )\\[8.0pt] & +β D_KL ( [rgb]0.9297,0.1094,0.1406q_t\,\|\, [rgb]0.1602,0.6719,0.8867p_ref,t ) aligned Figure 5: Online inference with a rolling prior. The learned initial prior initializes the first belief. Under identity state dynamics, prediction carries each posterior into the next reference (green arrows). The observation model is shown by black arrows and inference by red dashed arrows. The fixed-input case presents the same observation at each step. Prior prediction. Let qtq_t be the inferred distribution after observation t x_t. To form the next reference, a state-transition model predicts which latent configurations remain plausible. For a transition kernel (′∣)K( z z), prediction averages over the preceding belief (Särkkä and Svensson, 2023, Chapter 6): pref,t(′)=∑(′∣)qt−1().p_ref,t( z )= _ zK( z z)q_t-1( z). (65) If the latent state persists without process noise, then (′∣)=[′=]K( z z)=1[ z = z], so prediction simply copies the belief: pref,0=p0,pref,t=qt−1(t≥1),p0()=∏iBernoulli(zi;π0,i).p_ref,0=p_0, p_ref,t=q_t-1 (t≥ 1), p_0( z)= _iBernoulli(z_i; [rgb]0.1602,0.6719,0.8867π_0,i). (66) The identity transition means that the causes persist, not that we know which causes are present. Copying qt−1q_t-1 therefore retains both its feature preferences and its uncertainty. With no new likelihood contribution, the belief is preserved unchanged. This is the working memory supplied by the rolling reference. A new independent input starts again from the learned initial prior. Posterior update. The observation can support the prediction or require it to change. Expected reconstruction error rewards explanations of the current input, while the KL penalty controls departure from the prediction. For a Bernoulli reference with means ref,t [rgb]0.1602,0.6719,0.8867 π_ref,t and logits ref,t=logit(ref,t) [rgb]0.1602,0.6719,0.8867 u_ref,t=logit( [rgb]0.1602,0.6719,0.8867 π_ref,t), the local objective is ℱt()=−qlogp(t∣)⏟ℒrec(t,)+βDKL(q∥pref,t),=σ().F_t( [rgb]0.9297,0.1094,0.1406 u)= -E_q_ [rgb]0.9297,0.1094,0.1406 u p_ [rgb]0.1602,0.6719,0.8867 θ( x_t z)_L_rec( x_t; [rgb]0.9297,0.1094,0.1406 r)+β D_KL\! (q_ [rgb]0.9297,0.1094,0.1406 u\, \|\,p_ref,t ), [rgb]0.9297,0.1094,0.1406 r=σ( [rgb]0.9297,0.1094,0.1406 u). (67) The decoder and reference remain fixed during this local update. At β=1β=1, minimizing over unrestricted distributions gives the Bayesian posterior for the reference and observation. For general β>0β>0, it gives a distribution proportional to pref,t()p(t∣)1/βp_ref,t( z)p_ [rgb]0.1602,0.6719,0.8867 θ( x_t z)^1/β (Bissiri et al., 2016). The factorized Bernoulli family gives tractable expectations and local updates while retaining marginal uncertainty and approximating posterior dependencies. Recurrent belief revision. Write the negative mean-coordinate reconstruction gradient as t() a_t( [rgb]0.9297,0.1094,0.1406 r). Including the Bernoulli uncertainty term, its iith coordinate is [t()]i=(ϕ(i))⊤−1t−12Wii−∑j≠iWijrj.[ a_t( [rgb]0.9297,0.1094,0.1406 r)]_i=( [rgb]0.1602,0.6719,0.8867 φ^(i)) [rgb]0.1602,0.6719,0.8867 ^-1 x_t- 12 [rgb]0.1602,0.6719,0.8867W_i- _j≠ i [rgb]0.1602,0.6719,0.8867W_ij [rgb]0.9297,0.1094,0.1406r_j. (68) The first term supplies evidence for feature i, while the last adjusts it for other active explanations. When Wij>0 [rgb]0.1602,0.6719,0.8867W_ij>0, increasing belief in feature j reduces the drive for i. Repeated updates thus let explanations revise one another even though q is factorized. Applying the posterior Fisher inverse gives: ˙=t()−β(−ref,t). [rgb]0.9297,0.1094,0.1406 u= a_t( [rgb]0.9297,0.1094,0.1406 r)-β ( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref,t ). (69) Initialize at the predictive logits. The KL force is initially zero, so evidence can change the belief before the fixed reference exerts a restoring force. One ETD1 step with interval h integrates this relaxation while freezing the reconstruction drive (section B.6): t=ref,t+1−e−βhβt(ref,t),β>0. [rgb]0.9297,0.1094,0.1406 u_t= [rgb]0.1602,0.6719,0.8867 u_ref,t+ 1-e^-β hβ a_t\! ( [rgb]0.1602,0.6719,0.8867 π_ref,t ), β>0. (70) The resulting posterior becomes the reference for the next outer step. The next update therefore retains the inferred changes instead of pulling them back toward the original population prior. If several inner steps are taken, their reference stays fixed and the restoring force generally becomes nonzero. The parameter β controls resistance to local belief change, and h controls how far the local flow is integrated. Evidence accumulation and sparsity. Rolling also changes the influence of the initial prior as a sequence proceeds. To see this, suppose the initial prior and likelihood are positive and each local objective is minimized exactly over unrestricted distributions. Repeated substitution of the posterior for the next reference gives qt()∝p0()∏j=0tp(j∣)1/β.q_t( z)\ \ p_0( z) _j=0^tp_ [rgb]0.1602,0.6719,0.8867 θ( x_j z)^1/β. (71) Each observation contributes a likelihood factor, while the initial prior enters once. For a fixed repeated input, the likelihood power is (t+1)/β(t+1)/β: the initial prior becomes weaker relative to reconstruction evidence. Replaying one stored observation supplies no new measurement, so this computation is not ordinary conditioning on one observation. Finite mean-field ETD1 steps approximate the local updates and need not attain this unrestricted endpoint. This explains why a sparse initial prior need not enforce the same sparsity throughout a rollout. Its low probabilities set a high initial activation threshold, but the rolling reference can retain increases supported by reconstruction evidence. A cap on mean initial-prior activity therefore constrains the starting belief, not final posterior activity. B.2 Natural-gradient dynamics For the free energy in Section B.1, inference changes the posterior while holding the model fixed. Learning then changes the model with the inferred statistics fixed: ˙ [rgb]0.9297,0.1094,0.1406 u =⊤−1−−12diag()⊙(−2)−β(−ref), = [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 x- [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406 r- 12diag( [rgb]0.1602,0.6719,0.8867W) ( 1-2 [rgb]0.9297,0.1094,0.1406 r)-β( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref), (72) ˙ [rgb]0.1602,0.6719,0.8867 =ηΦ(⊤−)p−1, = _ ( x [rgb]0.9297,0.1094,0.1406 r - [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406E) [rgb]0.1602,0.6719,0.8867E_p^-1, (73) ˙ [rgb]0.1602,0.6719,0.8867 =ηΣ(−), = _ ( [rgb]0.9297,0.1094,0.1406V- [rgb]0.1602,0.6719,0.8867 ), (74) ˙0 [rgb]0.1602,0.6719,0.8867 π_0 =η0β(−0). = _0β( [rgb]0.9297,0.1094,0.1406 r- [rgb]0.1602,0.6719,0.8867 π_0). (75) Here ref=logit(ref) [rgb]0.1602,0.6719,0.8867 u_ref=logit( [rgb]0.1602,0.6719,0.8867 π_ref). The initial prior 0 [rgb]0.1602,0.6719,0.8867 π_0 tracks detached posterior means, and the rolling reference is refreshed within inference (section B.1.1). For minibatches, average ⊤ x [rgb]0.9297,0.1094,0.1406 r , [rgb]0.9297,0.1094,0.1406E, p [rgb]0.1602,0.6719,0.8867E_p, and [rgb]0.9297,0.1094,0.1406V before applying the updates. Derivations appear in Section B.2.1. Fisher preconditioning for inference. For the Bernoulli posterior, ∂logq/∂ui=zi−ri∂ q/∂ [rgb]0.9297,0.1094,0.1406u_i=z_i- [rgb]0.9297,0.1094,0.1406r_i. Independence under q gives u=q[(−)(−)⊤] _u=E_q[( z- [rgb]0.9297,0.1094,0.1406 r)( z- [rgb]0.9297,0.1094,0.1406 r) ] =, = [rgb]0.9297,0.1094,0.1406C, (76) ∇ℱ _ [rgb]0.9297,0.1094,0.1406 uF =[−⊤−1+12diag()⊙(−2)] = [rgb]0.9297,0.1094,0.1406C [ [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406 r- [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 x+ 12diag( [rgb]0.1602,0.6719,0.8867W) ( 1-2 [rgb]0.9297,0.1094,0.1406 r) ] +β(−ref). + [rgb]0.9297,0.1094,0.1406Cβ( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref). (77) Multiplication by −u−1-G_u^-1 cancels [rgb]0.9297,0.1094,0.1406C and yields eq. 72 whenever 0<ri<10< [rgb]0.9297,0.1094,0.1406r_i<1. Ordinary logit-gradient descent instead multiplies the same drive by ri(1−ri) [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i), making it small near either boundary. The natural flow removes this factor from the logit update, while the probability still changes according to ˙=˙ [rgb]0.9297,0.1094,0.1406 r= [rgb]0.9297,0.1094,0.1406C [rgb]0.9297,0.1094,0.1406 u. This distinction matters when interpreting the revival of a nearly inactive latent. Decoder Fisher information. Let ⟨⋅⟩ · denote the minibatch average. The Gaussian decoder score is −1(−)⊤ [rgb]0.1602,0.6719,0.8867 ^-1( x- [rgb]0.1602,0.6719,0.8867 z) z . Averaging its outer product over the complete-data model, with the chosen prior references held fixed, gives Φ=p⊗−1,∇ℱ=−1(−⟨⊤⟩),G_ = [rgb]0.1602,0.6719,0.8867E_p [rgb]0.1602,0.6719,0.8867 ^-1, _ [rgb]0.1602,0.6719,0.8867 F= [rgb]0.1602,0.6719,0.8867 ^-1 ( [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406E- x [rgb]0.9297,0.1094,0.1406 r ), (78) where both second moments are averaged across the minibatch before the Fisher action. Vectorization of decoder entries is implicit. The resulting matrix-form flow is eq. 73, with ⊤ x [rgb]0.9297,0.1094,0.1406 r replaced by its minibatch average. This is the complete-data model Fisher, rather than the observed-data Fisher or an outer product of empirical loss gradients (Amari, 1998; Martens, 2020). The distinction between p [rgb]0.1602,0.6719,0.8867E_p and [rgb]0.9297,0.1094,0.1406E is essential: the former supplies the metric, and the latter the curvature of the frozen reconstruction objective. Averaging per-example inverses would give a different update. Likewise, matching average firing rates alone does not match second moments. For example, a shared factorized prior with mean ⟨⟩ [rgb]0.9297,0.1094,0.1406 r has off-diagonal moments ⟨ri⟩⟨rj⟩ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_j , whereas the averaged posterior has ⟨rirj⟩ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_j . Their difference is the across-input covariance of the inferred means. Section C.4 states the exact matching condition and the unequal-moment extension. B.2.1 Model updates Proof of Equations 72 and 75. For Bernoulli inference, differentiating eq. 62 gives ∇ℱ=−⊤−1++12diag()⊙(−2)+β(−ref). _ [rgb]0.9297,0.1094,0.1406 rF=- [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 x+ [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406 r+ 12diag( [rgb]0.1602,0.6719,0.8867W) ( 1-2 [rgb]0.9297,0.1094,0.1406 r)+β( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref). The posterior Fisher [rgb]0.9297,0.1094,0.1406C cancels the logit Jacobian, yielding eq. 72. For the decoder, the complete-data Fisher block is p⊗−1 [rgb]0.1602,0.6719,0.8867E_p [rgb]0.1602,0.6719,0.8867 ^-1, which cancels the precision in the gradient and gives eq. 73. For each diagonal variance and each initial-prior probability, ∂ℱ∂Σmm ∂ [rgb]0.1602,0.6719,0.8867 _m =Σmm−Vmm2Σmm2, = [rgb]0.1602,0.6719,0.8867 _m- [rgb]0.9297,0.1094,0.1406V_m2 [rgb]0.1602,0.6719,0.8867 _m^2, GΣmm G_ [rgb]0.1602,0.6719,0.8867 _m =12Σmm2, = 12 [rgb]0.1602,0.6719,0.8867 _m^2, ∂[β(q∥p0)]∂π0,i ∂[ _ KL (q\, \|\,p_0 )]∂ [rgb]0.1602,0.6719,0.8867π_0,i =βπ0,i−riπ0,i(1−π0,i), =β [rgb]0.1602,0.6719,0.8867π_0,i- [rgb]0.9297,0.1094,0.1406r_i [rgb]0.1602,0.6719,0.8867π_0,i(1- [rgb]0.1602,0.6719,0.8867π_0,i), Gπ0,i G_ [rgb]0.1602,0.6719,0.8867π_0,i =1π0,i(1−π0,i). = 1 [rgb]0.1602,0.6719,0.8867π_0,i(1- [rgb]0.1602,0.6719,0.8867π_0,i). (79) Preconditioning proves eqs. 74 and 75. With frozen targets, integration yields: (Δt) [rgb]0.1602,0.6719,0.8867 ( t) =+[(0)−]e−ηΣΔt, = [rgb]0.9297,0.1094,0.1406V+[ [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.9297,0.1094,0.1406V]e^- _ t, 0(Δt) [rgb]0.1602,0.6719,0.8867 π_0( t) =+[0(0)−]e−η0βΔt. = [rgb]0.9297,0.1094,0.1406 r+[ [rgb]0.1602,0.6719,0.8867 π_0(0)- [rgb]0.9297,0.1094,0.1406 r]e^- _0β t. (80) ∎ Decoder-bias correction. Proof of Equation 81. Hold the incoming reference and covariance fixed during one unclamped inference step. Only the feedforward drive depends on the bias, so ∂new∂=−1−e−βhβ⊤−1,∇ℱ=−˙. ∂ [rgb]0.9297,0.1094,0.1406 u^new∂ [rgb]0.1602,0.6719,0.8867 b=- 1-e^-β hβ [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1, _ [rgb]0.9297,0.1094,0.1406 uF=- [rgb]0.9297,0.1094,0.1406C [rgb]0.9297,0.1094,0.1406 u. Applying the chain rule at the updated posterior gives ∇ℱ=−1(−)+1−e−βhβ−1˙. _ [rgb]0.1602,0.6719,0.8867 bF=- [rgb]0.1602,0.6719,0.8867 ^-1( x- [rgb]0.1602,0.6719,0.8867 b- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r)+ 1-e^-β hβ [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406C [rgb]0.9297,0.1094,0.1406 u. Multiplication by the inverse bias Fisher block [rgb]0.1602,0.6719,0.8867 and minibatch averaging gives eq. 81. ∎ B.3 Decoder bias A learned bias gives p(∣)=(,+,)p( x z)=N( x; [rgb]0.1602,0.6719,0.8867 b+ [rgb]0.1602,0.6719,0.8867 z, [rgb]0.1602,0.6719,0.8867 ), and inference uses − x- [rgb]0.1602,0.6719,0.8867 b. With a fixed posterior, the bias fits the mean residual. Differentiating through one unclamped ETD1 inference step adds a response term: ˙=ηΦbatch[−⏟reconstruction residual−1−e−βhβ˙⏟inference response]. [rgb]0.1602,0.6719,0.8867 b= _ E_batch\! [ x- [rgb]0.1602,0.6719,0.8867 b- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r_reconstruction residual- 1-e^-β hβ [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406C [rgb]0.9297,0.1094,0.1406 u_inference response ]. (81) The incoming reference is fixed, and [rgb]0.9297,0.1094,0.1406C and ˙ [rgb]0.9297,0.1094,0.1406 u are evaluated after that step. This expression assumes one shared covariance across the batch. Section B.2.1 derives the correction. The full-width synthetic experiment freezes the bias. B.4 Bernoulli uncertainty and sparse coding The Bernoulli posterior variance is Varq(zi)=ri(1−ri)Var_q(z_i)= [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i), which vanishes at 0 and 1 and is maximal at 1/21/2. In isolation, its cost favors confident absence or presence. In the full expected reconstruction loss, it cancels the quadratic self-term: 12Wiiri2⏟mean reconstruction+12Wiiri(1−ri)⏟posterior uncertainty=12Wiiri. 12 [rgb]0.1602,0.6719,0.8867W_i [rgb]0.9297,0.1094,0.1406r_i^2_mean reconstruction+ 12 [rgb]0.1602,0.6719,0.8867W_i [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)_posterior uncertainty= 12 [rgb]0.1602,0.6719,0.8867W_i [rgb]0.9297,0.1094,0.1406r_i. (82) Consequently, with the other coordinates fixed, reconstruction is affine in ri [rgb]0.9297,0.1094,0.1406r_i and the KL term is strictly convex for β>0β>0. The full objective has an interior coordinate optimum at finite evidence. B.5 Activation threshold under a rolling prior Section B.1.1 motivates the rolling rule. Here we make its sparsity effect explicit. Starting one ETD1 step at the reference, the gain in eq. 70 tends to 1/β1/β as h→∞h→∞. Since σ(u)>1/2σ(u)>1/2 exactly when u>0u>0, for β>0β>0 we obtain rinew>12⟺(ϕ(i))⊤−1−12Wii−∑j≠iWijπref,j>βlog1−πref,iπref,i. [rgb]0.9297,0.1094,0.1406r_i^new> 12 ( [rgb]0.1602,0.6719,0.8867 φ^(i)) [rgb]0.1602,0.6719,0.8867 ^-1 x- 12 [rgb]0.1602,0.6719,0.8867W_i- _j≠ i [rgb]0.1602,0.6719,0.8867W_ij [rgb]0.1602,0.6719,0.8867π_ref,j>β 1- [rgb]0.1602,0.6719,0.8867π_ref,i [rgb]0.1602,0.6719,0.8867π_ref,i. (83) The right side is positive for a reference probability below 1/21/2 and grows as that probability decreases. It sets the reconstruction evidence needed to activate the feature, after accounting for the other features’ contributions on the left. Updating the reference changes this threshold at the next outer step. The prior cap and inference depth can therefore affect posterior sparsity through different parts of the rollout (section C.3.3). B.6 Numerical integration and matrix-free updates Stiffness and exponential integration. Natural gradients specify the continuous-time evolution of beliefs and model parameters, but implementing these dynamics requires a numerical update. A difficulty with explicit Euler is that stable continuous dynamics can become unstable after discretization. For the simple relaxation v˙=−βv v=-β v, Euler multiplies v by 1−βh1-β h, whereas exact integration multiplies it by e−βhe^-β h. Euler therefore requires 0<βh<20<β h<2 for decay. When fast relaxation accompanies slower dynamics, this stability restriction can force small steps even when the solution changes slowly. Exponential integrators address this source of stiffness by integrating the linear component exactly (Cox and Matthews, 2002; Hochbruck and Ostermann, 2010). Inference integration. Write eq. 72 as ˙=()−β(−ref) [rgb]0.9297,0.1094,0.1406 u= a( [rgb]0.9297,0.1094,0.1406 r)-β( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref), where a contains the reconstruction and Bernoulli uncertainty terms. Holding the reference fixed and freezing a during one interval gives the first-order exponential time-differencing update (Cox and Matthews, 2002; Hochbruck and Ostermann, 2010, ETD1;): new=ref+e−βh(−ref)+1−e−βhβ(),β>0. $ [rgb]0.9297,0.1094,0.1406 u_new\;=\; [rgb]0.1602,0.6719,0.8867 u_ref\,+\,e^-β h( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref)\,+\, 1-e^-β hβ a( [rgb]0.9297,0.1094,0.1406 r)~,$ β>0~. (84) The reference relaxation is thus preserved exactly. Here ETD1 follows the same natural-gradient direction as Euler, with effective step (1−e−βh)/β(1-e^-β h)/β, bounded by 1/β1/β. Starting from the reference removes the middle term. Section B.1.1 describes how that reference is passed between inference steps. At h→∞h→∞, ETD1 solves the subproblem with the reconstruction drive held fixed. We apply a finite number of recurrent inference steps before each decoder update. Frozen-statistics learning. The same principle applies to learning. With minibatch statistics, prior references, and covariance fixed, and ,p≻0 [rgb]0.9297,0.1094,0.1406E, [rgb]0.1602,0.6719,0.8867E_p 0, the decoder flow has the exact solution (section C.4) (Δt)−∗=((0)−∗)exp[−ηΦΔtp−1]. [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*= ( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^* ) \! [- _ t\, [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1 ]. (85) Euler must accommodate the fastest relaxation mode, but exponential integration instead applies each mode’s exact decay over the chosen interval. When =p [rgb]0.9297,0.1094,0.1406E= [rgb]0.1602,0.6719,0.8867E_p, every mode contracts by the same factor, recovering exact rarity compensation. The configured decoder solver, matrix_free_phi1_auto, evaluates a matrix-function action on the decoder residual. Let s=ηΦΔts= _ t, p,λ=p+λ [rgb]0.1602,0.6719,0.8867E_p,λ= [rgb]0.1602,0.6719,0.8867E_p+ with Fisher damping λ≥0λ≥ 0, and =p,λ−1A= [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p,λ^-1. The corresponding frozen update is: (Δt) [rgb]0.1602,0.6719,0.8867 ( t) =(0) = [rgb]0.1602,0.6719,0.8867 (0) (86) +s(⟨⊤⟩−(0))p,λ−1φ1(−s), +s ( x [rgb]0.9297,0.1094,0.1406 r - [rgb]0.1602,0.6719,0.8867 (0) [rgb]0.9297,0.1094,0.1406E ) [rgb]0.1602,0.6719,0.8867E_p,λ^-1 _1(-sA), where brackets denote minibatch averages and φ1()=∑m=0∞m/(m+1)! _1(Z)= _m=0^∞Z^m/(m+1)! is the entire matrix function used in exponential integration (Hochbruck and Ostermann, 2010). For λ=0λ=0, this recovers the exact frozen-statistics solution above. The residual form avoids explicitly computing ∗ [rgb]0.1602,0.6719,0.8867 ^* or subtracting a near-identity exponential from the identity. The implementation applies this function without constructing a dense K×K× K matrix. For a matrix R whose rows are minibatch posterior means, =diag(⟨⊙(−)⟩)+1B⊤(), [rgb]0.9297,0.1094,0.1406E v=diag( [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r) ) v+ 1BR (R v), (87) and the same representation applies to p [rgb]0.1602,0.6719,0.8867E_p. Low-rank solves use the Woodbury identity (Hager, 1989) when the rank is at most 1,024. Larger ranks use Jacobi-preconditioned conjugate gradients (Saad, 2003). Thus the minibatch rank and conditioning also affect cost (section B.8). The solver combines scaled Taylor evaluation (Al-Mohy and Higham, 2011) with Chebyshev polynomial approximation (Trefethen, 2019). It selects Taylor evaluation when the predicted work is within 32 operator applications, and otherwise uses a Chebyshev recurrence with degree chosen from the coefficient tail. Each operator application requires a posterior-moment product and a solve with the damped prior moment, so the polynomial degree and linear-solve iterations contribute to decoder-update cost. The selected synthetic recipes and audited whitened Gemma configurations request Chebyshev tolerance 10−510^-5 and maximum degree 2,048 (512 for the raw 16K Gemma run). These control the matrix-function approximation. The variance and initial-prior flows likewise integrate to exponential averages of their frozen targets (section B.2.1). These updates exploit solvable components of the derived dynamics while leaving approximation to the remaining nonlinear terms and numerical matrix-function evaluation. Parameter floors, a mean initial-prior cap, and decoder constraints are applied as specified by each recipe. All inferred statistics are detached during learning, so the decoder update does not backpropagate through the inference trajectory. B.7 Algorithm Algorithm 1 gives the Bernoulli inference and moment updates. During inference, the generative parameters remain fixed. Learning uses the inferred minibatch moments. Algorithm 1 FOND with a rolling inference reference 1: Decoder [rgb]0.1602,0.6719,0.8867 , diagonal covariance ≻0 [rgb]0.1602,0.6719,0.8867 0, initial prior 0∈(0,1)K [rgb]0.1602,0.6719,0.8867 π_0∈(0,1)^K, inference depths, β>0β>0, integration parameter h, and learning schedules. 2: for each training minibatch do 3: Center observations by the decoder bias, if learned. 4: ←logit(0) [rgb]0.9297,0.1094,0.1406 u ( [rgb]0.1602,0.6719,0.8867 π_0), independently for each observation. 5: for each outer step up to the sampled training horizon do 6: ref← [rgb]0.1602,0.6719,0.8867 u_ref← [rgb]0.9297,0.1094,0.1406 u, ref←σ(ref) [rgb]0.1602,0.6719,0.8867 π_ref←σ( [rgb]0.1602,0.6719,0.8867 u_ref). 7: p←batch[diag(ref⊙(−ref))+refref⊤] [rgb]0.1602,0.6719,0.8867E_p _batch[diag( [rgb]0.1602,0.6719,0.8867 π_ref ( 1- [rgb]0.1602,0.6719,0.8867 π_ref))+ [rgb]0.1602,0.6719,0.8867 π_ref [rgb]0.1602,0.6719,0.8867 π_ref ]. 8: for each inner inference step do 9: ←σ() [rgb]0.9297,0.1094,0.1406 r←σ( [rgb]0.9297,0.1094,0.1406 u), ←⊤−1 [rgb]0.1602,0.6719,0.8867W← [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 . 10: ˙←⊤−1−−12diag()⊙(−2)−β(−ref) [rgb]0.9297,0.1094,0.1406 u← [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 x- [rgb]0.1602,0.6719,0.8867W [rgb]0.9297,0.1094,0.1406 r- 12diag( [rgb]0.1602,0.6719,0.8867W) ( 1-2 [rgb]0.9297,0.1094,0.1406 r)-β( [rgb]0.9297,0.1094,0.1406 u- [rgb]0.1602,0.6719,0.8867 u_ref). 11: ←+1−e−βhβ˙ [rgb]0.9297,0.1094,0.1406 u← [rgb]0.9297,0.1094,0.1406 u+ 1-e^-β hβ [rgb]0.9297,0.1094,0.1406 u. 12: end for 13: if this is the final outer step then 14: ←σ() [rgb]0.9297,0.1094,0.1406 r←σ( [rgb]0.9297,0.1094,0.1406 u), ←diag[⊙(−)] [rgb]0.9297,0.1094,0.1406C [ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)], ←batch[+⊤] [rgb]0.9297,0.1094,0.1406E _batch[ [rgb]0.9297,0.1094,0.1406C+ [rgb]0.9297,0.1094,0.1406 r [rgb]0.9297,0.1094,0.1406 r ]. 15: Vmm←batch[(−)m2+(⊤)mm] [rgb]0.9297,0.1094,0.1406V_m _batch[( x- [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r)_m^2+( [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406C [rgb]0.1602,0.6719,0.8867 )_m]. 16: ∗←batch[⊤]−1 [rgb]0.1602,0.6719,0.8867 ^* _batch[ x [rgb]0.9297,0.1094,0.1406 r ] [rgb]0.9297,0.1094,0.1406E^-1. 17: ←∗+(−∗)exp[−ηΦΔtp−1] [rgb]0.1602,0.6719,0.8867 ← [rgb]0.1602,0.6719,0.8867 ^*+( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [- _ t\, [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1]. 18: Update enabled variance and bias parameters by eqs. 80 and 81, then apply the stated decoder constraints. 19: end if 20: end for 21: if an initial-prior update is scheduled then 22: Infer [rgb]0.9297,0.1094,0.1406 r on the prior-update batch, then detach it. 23: 0←batch[]+(0−batch[])e−η0βΔt [rgb]0.1602,0.6719,0.8867 π_0 _batch[ [rgb]0.9297,0.1094,0.1406 r]+( [rgb]0.1602,0.6719,0.8867 π_0-E_batch[ [rgb]0.9297,0.1094,0.1406 r])e^- _0β t. 24: Apply the specified prior floor and mean-probability cap. 25: end if 26: end for The decoder line is the frozen-moment matrix-exponential update. Implementations can apply its matrix functions without forming dense inverses. ETD1 uses gain 1/β1/β at h=∞h=∞. Evaluation resets the initial prior for each input and performs inference without learning. B.8 Computational complexity Let B be minibatch size, d input dimension, K dictionary width, and T the total number of inner inference steps for an input. A dense SAE encoder costs O(BdK)O(BdK); its training step has the same order, with a larger constant for decoding and backpropagation. BeBeFOND applies the lateral operator as ⊤[] [rgb]0.1602,0.6719,0.8867 [ [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406 r] (with diagonal observation precision inserted when used), so it does not form a dense K×K× K Gram matrix. Its costs are: BeFOND inference: $ Be$FOND inference: O(TBdK), O(TBdK)~, (88) BeFOND training step: $ Be$FOND training step: O(TBdK+CNGD+Cprior), O(TBdK+C_NGD+C_prior)~, where CNGDC_NGD includes the moment products and linear solves used by the decoder update, and CpriorC_prior is the amortized cost of the separate initial-prior fits. Solver work depends on minibatch rank, conditioning, and numerical tolerance. Section B.6 describes the matrix-free evaluation. No gradient is propagated through the T inference steps, avoiding storage of that trajectory for backpropagation. If a conventional SAE update uses c units of dense matrix-multiplication work, its cost is O(cBdK)O(cBdK). With T=O(c)T=O(c), BeBeFOND’s inference contribution has that same order. Our current implementation takes approximately 1212–14×14× as long as BatchTopK for the measured training prefix (section D.3). Appendix C Extended theoretical analyses, synthetic experiments, and ablations C.1 Experimental setup C.1.1 Datasets SynthSAEBench. We measure dictionary and support recovery using the known generating features of SynthSAEBench (Chanin and Garriga-Alonso, 2026). The full-width experiment uses d=768d=768 and K=K⋆=16,384K=K_ =16,384. The superposition experiment fixes K=4,096K=4,096, K⋆=16,384K_ =16,384, and varies d from 256 to 1,536. Inputs are unwhitened, with parent scaling disabled (scale_children_by_parent=false). Feature probabilities follow a Zipf law with exponent 0.5 over [0.0005,0.4][0.0005,0.4]. The generator uses 128 roots, branching factor 4, depth 3, and probability compensation. C.1.2 Training BeBeFOND alternates inference and learning as in Alg. 1. Table 2 groups the synthetic settings. Prior-update examples are counted separately from decoder-update examples. Geometry arms are separately tuned. Their early checkpoints include recipe-matched reproductions. Table 2: SynthSAEBench training settings. Decoder and prior examples are counted separately. Full-width runs use d=768d=768, K=K⋆=16,384K=K_ =16,384; superposition runs use K=4,096K=4,096. Experiment Settings BeBeFOND best / continuous Seed 0; batch 2,048; 88,779 updates; 181.82M decoder and 18.18M prior examples; linear LR 0.02→10−50.02→ 10^-5, no warmup; β=1.6β=1.6 Inference / decoder Bernoulli; ETD1 h=∞h=∞; one inner step; uniform training depth 1–10, test depth 200. Prior Fisher, current moments, no damping; Δt=1 t=1; bias and variance fixed Initialization / constraints Gaussian scale 0.1443375673; prior logits −4.7538364662-4.7538364662; soft decoder norm 3.5 with coefficient 100; prior mean cap 3×10−43× 10^-4 and floor 10−510^-5 Prior / matrix functions Prior update every 10 batches on 2,048 separate examples; depth ≤200≤ 200, KL threshold 0.01, rejection threshold 1, window 10. Taylor limit 32 applications; Chebyshev tolerance 10−510^-5, maximum degree 2,048 Original baselines BatchTopK/JumpReLU/Matryoshka: target L0=35/45/20L_0=35/45/20, batch 1,024, LR 3×10−43× 10^-4, 200M decoder examples. MP/ReLU-L1: target L0=35L_0=35, 200,000,512 decoder plus 19M monitoring examples. MF: 20 inference steps in training and testing, batch 1,024, 195,313 updates (200,000,512 examples), cosine schedule with dictionary / other-parameter LR 0.001/0.0030.001/0.003, damping 1 on the first step then 0.2, decoder norm cap 4 Modified baselines Target L0=45/40/35L_0=45/40/35 for BatchTopK/JumpReLU/Matryoshka; MinFire or resampling and learning-rate decay Training streams Model seed 0; BeBeFOND data seed 1; plain amortized baseline data seed 100101; native MP/ReLU-L1 use the global PyTorch and NumPy streams initialized with seed 0, without a separate training-stream seed; MF uses data seed 0 Geometry comparison Batch 2,048; 100,000 updates; 150,000-step cosine schedule stopped at 100,000; bias and variance fixed. Tuples below are (β,norm,prior cap,peak LR)(β,norm,prior cap,peak LR) N/E; E/E; N/N (1.125,2.5,0.001,0.1)(1.125,2.5,0.001,0.1); (0.8,3.5,0.01,0.1)(0.8,3.5,0.01,0.1); (1.125,3.5,0.0002,0.006)(1.125,3.5,0.0002,0.006). Natural inference: ETD1 h=∞h=∞; Euclidean inference: Euler h=128h=128 Superposition / lateral Baseline comparison: initialization seeds 0, 1, 2. Lateral ablations: seed 0. Batch 1,024. Recipe differences are specified in Section C.3.1 C.1.3 Evaluation Synthetic recovery. MCC is dictionary mean correlation coefficient, feature-macro F1 uses positive-cosine one-to-one matching to the generating dictionary. and rare F1 averages over its least frequent quartile. A latent is dead when it is inactive on every test example. Activity is ri>0.5 [rgb]0.9297,0.1094,0.1406r_i>0.5 for Bernoulli BeBeFOND/MF and positive native output for the SAE baselines. Mean L0L_0 counts active latents per example. Tests use 10610^6 examples with seed 49100, including the MF reevaluation reported in Table 5. The independent frequency reference uses 5×1055× 10^5 examples with seed 315031. BeBeFOND/MF use 200/20 test inference steps. Geometry validation uses 32,768 examples, seed 813011, and 100 inference steps. C.1.4 Implementation Training horizon. Each optimizer step draws one inference horizon for the whole minibatch, resets the posterior, runs that many outer steps, and updates the decoder only at the final step. The exponential sampler assigns: Pr(t=j)=e−λ(j−1)∑s=1toutere−λ(s−1),j=1,…,touter, (t=j)= e^-λ(j-1) _s=1^t_outere^-λ(s-1), j=1,…,t_outer, (89) where λ is solved numerically so the truncated discrete distribution has the configured mean. Probabilities use FP64. Brent’s root solver uses tolerance 10−1410^-14. At mean 1 the distribution is a point mass at 1, and at mean (touter+1)/2(t_outer+1)/2, it is uniform. Sampling uses NumPy’s default_rng seeded by the model seed and zero-based optimizer step, identically on every device and after resumption. The draw does not advance the data or latent-sampling generators. K1 uses a fixed horizon of one and one inner step. Its initial-prior fit and test inference can still use longer horizons. Uniform sampling on 1–10 has mean 5.5, Pr(t=1)=0.1 (t=1)=0.1, and 95th percentile 10. Exponential sampling on 1–50 with the same mean has Pr(t=1)≈0.182 (t=1)≈ 0.182 and 95th percentile 1515. Synthetic baseline implementations use SAE Lens 6.51.0 (BatchTopK/JumpReLU/Matryoshka) or 6.37.0 (native MP/ReLU-L1). C.2 Inference proofs C.2.1 Single-layer inference Proposition C.1 (Two-atom case). Let zi∈0,1z_i∈\0,1\ with zi∼Bernoulli(πi)z_i ( [rgb]0.1602,0.6719,0.8867π_i) independently, with 0<πi<10< [rgb]0.1602,0.6719,0.8867π_i<1, and let =ϕ(1)z1+ϕ(2)z2+ϵ x= [rgb]0.1602,0.6719,0.8867 φ^(1)z_1+ [rgb]0.1602,0.6719,0.8867 φ^(2)z_2+ ε, where ϵ∼(,M) ε ( 0,I_M) independently of z. If the dictionary atoms are linearly independent and W12=(ϕ(1))⊤ϕ(2)≠0 [rgb]0.1602,0.6719,0.8867W_12=( [rgb]0.1602,0.6719,0.8867 φ^(1)) [rgb]0.1602,0.6719,0.8867 φ^(2)≠ 0, then no encoder z^1()=f(⊤+b) [rgb]0.9297,0.1094,0.1406 z_1( x)=f( [rgb]0.9297,0.1094,0.1406 w x+ [rgb]0.9297,0.1094,0.1406b), for any scalar function f, can equal p(z1=1∣)p(z_1=1 x) for every input x. Assume independent Bernoulli latents with interior prior probabilities and unit observation covariance. Set: ai()=logit(πi)+(ϕ(i))⊤−12‖ϕ(i)‖2,Wij=(ϕ(i))⊤ϕ(j).a_i( x)\;=\;logit( [rgb]0.1602,0.6719,0.8867π_i)+( [rgb]0.1602,0.6719,0.8867 φ^(i)) x- 12\| [rgb]0.1602,0.6719,0.8867 φ^(i)\|^2, [rgb]0.1602,0.6719,0.8867W_ij\;=\;( [rgb]0.1602,0.6719,0.8867 φ^(i)) [rgb]0.1602,0.6719,0.8867 φ^(j). Proof of Proposition C.1. Marginalizing z2z_2 gives logitp(z1=1∣) (z_1=1 x)\; =a1()+log1+ea2()−W121+ea2(), =\;a_1( x)+ 1+e^a_2( x)- [rgb]0.1602,0.6719,0.8867W_121+e^a_2( x), ∇logitp(z1=1∣) _ xlogitp(z_1=1 x)\; =ϕ(1)+[σ(a2−W12)−σ(a2)]ϕ(2). =\; [rgb]0.1602,0.6719,0.8867 φ^(1)+[σ(a_2- [rgb]0.1602,0.6719,0.8867W_12)-σ(a_2)] [rgb]0.1602,0.6719,0.8867 φ^(2). (90) When W12≠0 [rgb]0.1602,0.6719,0.8867W_12≠ 0, the bracket is nonconstant. Linear independence of the atoms therefore makes the gradient direction vary with x. A posterior of the form f(⊤+b)f( [rgb]0.9297,0.1094,0.1406 w x+ [rgb]0.9297,0.1094,0.1406b) is constant along every direction orthogonal to [rgb]0.9297,0.1094,0.1406 w, so its log-odds gradient must be parallel to [rgb]0.9297,0.1094,0.1406 w everywhere. This is a contradiction, even if f is not assumed differentiable. ∎ Multiple atoms. By assumption, some atom ϕ(i) [rgb]0.1602,0.6719,0.8867 φ^(i), i≥2i≥ 2, is neither orthogonal to nor a multiple of ϕ(1) [rgb]0.1602,0.6719,0.8867 φ^(1). Proof of Proposition 4.1. Suppose f(⊤+b)=p(z1=1∣)f( [rgb]0.9297,0.1094,0.1406 w x+ [rgb]0.9297,0.1094,0.1406b)=p(z_1=1 x) for every x, and write L()=logitp(z1=1∣)L( x)=logitp(z_1=1 x). As in the proof of proposition C.1, ∇L∇ L is then parallel to [rgb]0.9297,0.1094,0.1406 w everywhere. Differentiation of the finite partition functions gives: ∇L=ϕ(1)+∑k≥2Δkϕ(k),Δk=[zk∣z1=1,]−[zk∣z1=0,].∇ L\;=\; [rgb]0.1602,0.6719,0.8867 φ^(1)\,+\, _k≥ 2 _k [rgb]0.1602,0.6719,0.8867 φ^(k)~, _k\;=\;E[z_k z_1=1, x]\,-\,E[z_k z_1=0, x]~. (91) Along =t x=t d, with d orthogonal to no atom, the fields of all other bits diverge while their couplings remain bounded. Each bit becomes certain independently of z1z_1. Hence, Δk→0 _k→ 0 and ∇L→ϕ(1)∇ L→ [rgb]0.1602,0.6719,0.8867 φ^(1). Therefore, |ϕ(1) [rgb]0.9297,0.1094,0.1406 w [rgb]0.1602,0.6719,0.8867 φ^(1), and ∇L⊤=0∇ L v=0 everywhere for every ⟂ϕ(1) v [rgb]0.1602,0.6719,0.8867 φ^(1). Let Q contain the atoms parallel to the chosen atom i, writing ϕ(j)=λjϕ(i) [rgb]0.1602,0.6719,0.8867 φ^(j)= _j [rgb]0.1602,0.6719,0.8867 φ^(i). Choose ⟂ϕ(1) v [rgb]0.1602,0.6719,0.8867 φ^(1) with (ϕ(i))⊤≠0( [rgb]0.1602,0.6719,0.8867 φ^(i)) v≠ 0, and choose ⟂ϕ(i) d [rgb]0.1602,0.6719,0.8867 φ^(i) orthogonal to no atom outside Q. Such a d exists because only finitely many proper hyperplanes are excluded. Along =0+t x= x_0+t d, all bits outside Q∪1Q∪\1\ freeze, leaving finite conditional partition functions Z0(s),Z1(s)Z_0(s),Z_1(s) with s=(ϕ(i))⊤0s=( [rgb]0.1602,0.6719,0.8867 φ^(i)) x_0. Thus ∇L⊤⟶((ϕ(i))⊤)D′(s),D(s)=logZ1(s)Z0(s).∇ L v (( [rgb]0.1602,0.6719,0.8867 φ^(i)) v)D (s), D(s)= Z_1(s)Z_0(s). As s moves from −∞-∞ to +∞+∞, every bit in Q switches according to the sign of λj _j. Since W1j=λjW1i [rgb]0.1602,0.6719,0.8867W_1j= _j [rgb]0.1602,0.6719,0.8867W_1i, D(+∞)−D(−∞)=−W1i∑j∈Q|λj|≠0.D(+∞)-D(-∞)=- [rgb]0.1602,0.6719,0.8867W_1i _j∈ Q| _j|≠ 0. Therefore D′D is nonzero somewhere, contradicting ∇L⊤=0∇ L v=0 for a single-projection encoder. ∎ Orthogonality and parallelism of two atoms are polynomial conditions on [rgb]0.1602,0.6719,0.8867 that fail for almost every dictionary when M≥2M≥ 2, so for almost every dictionary with K≥2K≥ 2 no latent falls under the exception. If no two atoms are parallel, exceptions must be orthogonal to every other atom. For K>MK>M, at most M−1M-1 atoms can have that property. Whitening extends both results to any fixed positive-definite [rgb]0.1602,0.6719,0.8867 , with =⊤−1 [rgb]0.1602,0.6719,0.8867W= [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 . C.2.2 Iterative inference versus single projection encoders Proposition C.2 (Fixed-prior inference versus single-projection encoders). Consider the Bernoulli–Gaussian model of proposition 4.1 with unit observation covariance. Let =⊤,=−diag(), [rgb]0.1602,0.6719,0.8867W= [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 , [rgb]0.1602,0.6719,0.8867 J= [rgb]0.1602,0.6719,0.8867W-diag( [rgb]0.1602,0.6719,0.8867W), and suppose that λmin()>−4 _ ( [rgb]0.1602,0.6719,0.8867 J)>-4 and that some pair of atoms is neither orthogonal nor parallel. Let A contain all encoders qf(∣)=∏iBernoulli(zi;fi(i⊤+bi))q_f( z x)= _iBernoulli (z_i;\,f_i( [rgb]0.9297,0.1094,0.1406 w_i x+ [rgb]0.9297,0.1094,0.1406b_i) ) with arbitrary weights, biases, and measurable fi:ℝ→[0,1]f_i:R→[0,1], and let R(q)=∼p()[(q(⋅∣)∥p(⋅∣))].R(q)=E_ x p( x) [D_ KL (q(· x)\, \|\,p(· x) ) ]. Let qtq_t be the factorized posterior with means (t)=σ((t)), [rgb]0.9297,0.1094,0.1406 r(t)=σ( [rgb]0.9297,0.1094,0.1406 u(t)), where (t) [rgb]0.9297,0.1094,0.1406 u(t) follows section 3.1 with β=1β=1 and the prior fixed at the generating one, 0=logit,(0)=0. [rgb]0.1602,0.6719,0.8867 u_0=logit [rgb]0.1602,0.6719,0.8867 π, [rgb]0.9297,0.1094,0.1406 u(0)= [rgb]0.1602,0.6719,0.8867 u_0. Then the optimal factorized posterior q∗(⋅∣)=argminqfactorized(q∥p(⋅∣))q^*(· x)= *arg\,min_q\ factorizedD_ KL (q\, \|\,p(· x) ) is unique and limt→∞R(qt)=R(q∗)<infqf∈R(qf). _t→∞R(q_t)=R(q^*)< _q_f R(q_f). Consequently, for all sufficiently large finite t, R(qt)<infqf∈R(qf).R(q_t)< _q_f R(q_f). Proof. Zero atoms decouple, since their rows of [rgb]0.1602,0.6719,0.8867 J vanish, so we drop them. Write ai()=logit(πi)+(ϕ(i))⊤−12Wiia_i( x)=logit( [rgb]0.1602,0.6719,0.8867π_i)+( [rgb]0.1602,0.6719,0.8867 φ^(i)) x- 12 [rgb]0.1602,0.6719,0.8867W_i and ψ()=∑i[rilogri+(1−ri)log(1−ri)]ψ( [rgb]0.9297,0.1094,0.1406 r)= _i [ [rgb]0.9297,0.1094,0.1406r_i [rgb]0.9297,0.1094,0.1406r_i+(1- [rgb]0.9297,0.1094,0.1406r_i) (1- [rgb]0.9297,0.1094,0.1406r_i) ] for the negative Bernoulli entropy. For a factorized q with means ∈[0,1]K [rgb]0.9297,0.1094,0.1406 r∈[0,1]^K, using zi2=ziz_i^2=z_i, the divergence (q∥p(⋅∣))D_ KL (q\, \|\,p(· x) ) equals, up to a term independent of [rgb]0.9297,0.1094,0.1406 r, F()=ψ()−()⊤+12⊤.F_ x( [rgb]0.9297,0.1094,0.1406 r)=ψ( [rgb]0.9297,0.1094,0.1406 r)- a( x) [rgb]0.9297,0.1094,0.1406 r+ 12 [rgb]0.9297,0.1094,0.1406 r [rgb]0.1602,0.6719,0.8867 J [rgb]0.9297,0.1094,0.1406 r. Its Hessian diag(1/(ri(1−ri)))+diag (1/( [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)) )+ [rgb]0.1602,0.6719,0.8867 J is smallest at =12 [rgb]0.9297,0.1094,0.1406 r= 12 1, where it equals 4+4I+ [rgb]0.1602,0.6719,0.8867 J, so ∇2F⪰4m∇^2F_ x 4m\,I with m=1+λmin()/4>0m=1+ _ ( [rgb]0.1602,0.6719,0.8867 J)/4>0. Hence F_ x has a unique minimizer ∗()∈(0,1)K [rgb]0.9297,0.1094,0.1406 r^*( x)∈(0,1)^K, which defines q∗q^* and satisfies ∗:=logit∗=−∗ [rgb]0.9297,0.1094,0.1406 u^*:=logit [rgb]0.9297,0.1094,0.1406 r^*= a- [rgb]0.1602,0.6719,0.8867 J [rgb]0.9297,0.1094,0.1406 r^*. Convergence. With β=1β=1 and 0=logit [rgb]0.1602,0.6719,0.8867 u_0=logit [rgb]0.1602,0.6719,0.8867 π, section 3.1 reads ˙=−−=−∇F() [rgb]0.9297,0.1094,0.1406 u= a- [rgb]0.1602,0.6719,0.8867 J [rgb]0.9297,0.1094,0.1406 r- [rgb]0.9297,0.1094,0.1406 u=-∇ F_ x( [rgb]0.9297,0.1094,0.1406 r), a smooth, globally Lipschitz flow in [rgb]0.9297,0.1094,0.1406 u. Let V(t)=(q∗∥q(t))V(t)=D_ KL (q_ [rgb]0.9297,0.1094,0.1406 r^*\, \|\,q_ [rgb]0.9297,0.1094,0.1406 r(t) ), the Bregman divergence of ψ. Then V˙=−⟨∇F(),−∗⟩≤−[F()−F(∗)] V=- ∇ F_ x( [rgb]0.9297,0.1094,0.1406 r), [rgb]0.9297,0.1094,0.1406 r- [rgb]0.9297,0.1094,0.1406 r^* ≤- [F_ x( [rgb]0.9297,0.1094,0.1406 r)-F_ x( [rgb]0.9297,0.1094,0.1406 r^*) ] by convexity, while dtF()=−∇F⊤diag(ri(1−ri))∇F≤0 ddtF_ x( [rgb]0.9297,0.1094,0.1406 r)=-∇ F_ x diag ( [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) )∇ F_ x≤ 0. Hence t[F((t))−F(∗)]≤∫0t[F((s))−F(∗)]ds≤V(0)≤∑i−logminπi,1−πit [F_ x( [rgb]0.9297,0.1094,0.1406 r(t))-F_ x( [rgb]0.9297,0.1094,0.1406 r^*) ]≤ _0^t [F_ x( [rgb]0.9297,0.1094,0.1406 r(s))-F_ x( [rgb]0.9297,0.1094,0.1406 r^*) ]ds≤ V(0)≤ _i- \ [rgb]0.1602,0.6719,0.8867π_i,1- [rgb]0.1602,0.6719,0.8867π_i\ for every x. Since R(qt)−R(q∗)=[F((t))−F(∗)]R(q_t)-R(q^*)=E [F_ x( [rgb]0.9297,0.1094,0.1406 r(t))-F_ x( [rgb]0.9297,0.1094,0.1406 r^*) ], this gives R(qt)→R(q∗)R(q_t)→ R(q^*) at rate O(1/t)O(1/t). ri∗ [rgb]0.9297,0.1094,0.1406r_i^* is not a single-projection function. Choose i,ji,j with Wij≠0 [rgb]0.1602,0.6719,0.8867W_ij≠ 0 and ϕ(i),ϕ(j) [rgb]0.1602,0.6719,0.8867 φ^(i), [rgb]0.1602,0.6719,0.8867 φ^(j) not parallel. Implicit differentiation of ∇F(∗)=∇ F_ x( [rgb]0.9297,0.1094,0.1406 r^*)= 0 gives ∂∗/∂=−1⊤∂ [rgb]0.9297,0.1094,0.1406 r^*/∂ x= H_ x^-1 [rgb]0.1602,0.6719,0.8867 with =∇2F(∗)⪰4m H_ x=∇^2F_ x( [rgb]0.9297,0.1094,0.1406 r^*) 4m\,I, and ∂∗/∂=⊤−−1⊤∂ [rgb]0.9297,0.1094,0.1406 u^*/∂ x= [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 J H_ x^-1 [rgb]0.1602,0.6719,0.8867 . Both are bounded. Along a ray =t x=t d with d orthogonal to no atom, ∗=−∗ [rgb]0.9297,0.1094,0.1406 u^*= a- [rgb]0.1602,0.6719,0.8867 J [rgb]0.9297,0.1094,0.1406 r^* forces every coordinate to saturate, so ∂∗/∂=diag(rk∗(1−rk∗))∂∗/∂→∂ [rgb]0.9297,0.1094,0.1406 r^*/∂ x=diag ( [rgb]0.9297,0.1094,0.1406r_k^*(1- [rgb]0.9297,0.1094,0.1406r_k^*) )\,∂ [rgb]0.9297,0.1094,0.1406 u^*/∂ x→ 0 and ∇ui∗→ϕ(i) _ x [rgb]0.9297,0.1094,0.1406u_i^*→ [rgb]0.1602,0.6719,0.8867 φ^(i). Now let Q=k:ϕ(k)=λkϕ(j)Q=\k: [rgb]0.1602,0.6719,0.8867 φ^(k)= _k [rgb]0.1602,0.6719,0.8867 φ^(j)\ and take a ray =t x=t v with ⟂ϕ(j) v [rgb]0.1602,0.6719,0.8867 φ^(j) orthogonal to no atom outside Q. The coordinates outside Q saturate, while Q∗ [rgb]0.9297,0.1094,0.1406 r^*_Q converges to the minimizer ∈(0,1)Q s∈(0,1)^Q of the reduced problem, which is strongly convex because λmin(QQ)≥λmin() _ ( [rgb]0.1602,0.6719,0.8867 J_Q)≥ _ ( [rgb]0.1602,0.6719,0.8867 J). With Q=diag(1/(sk(1−sk)))+QQ≻0 H_Q=diag (1/(s_k(1-s_k)) )+ [rgb]0.1602,0.6719,0.8867 J_Q 0, the limiting equations give ∂Q∗/∂→Q−1(ϕ(j))⊤∂ [rgb]0.9297,0.1094,0.1406 r^*_Q/∂ x→ H_Q^-1 λ( [rgb]0.1602,0.6719,0.8867 φ^(j)) , and since Wik=λkWij [rgb]0.1602,0.6719,0.8867W_ik= _k [rgb]0.1602,0.6719,0.8867W_ij for k∈Qk∈ Q, ∇ui∗(t)⟶ϕ(i)−Wij(⊤Q−1)ϕ(j). _ x [rgb]0.9297,0.1094,0.1406u_i^*(t v) [rgb]0.1602,0.6719,0.8867 φ^(i)- [rgb]0.1602,0.6719,0.8867W_ij\, ( λ H_Q^-1 λ )\, [rgb]0.1602,0.6719,0.8867 φ^(j). The coefficient is nonzero, so this limit is not parallel to ϕ(i) [rgb]0.1602,0.6719,0.8867 φ^(i). The gradient of ui∗ [rgb]0.9297,0.1094,0.1406u_i^* therefore has no fixed direction, and neither ui∗ [rgb]0.9297,0.1094,0.1406u_i^* nor ri∗ [rgb]0.9297,0.1094,0.1406r_i^* is a function of a single projection. Positive approximation gap. Let δ=inff,,b[(f(⊤+b)−ri∗())2]δ= _f, w,bE [(f( w x+b)- [rgb]0.9297,0.1094,0.1406r_i^*( x))^2 ] over measurable f:ℝ→[0,1]f:R→[0,1]. If δ=0δ=0, take gn=fn(n⊤+bn)→ri∗g_n=f_n( w_n x+b_n)→ [rgb]0.9297,0.1094,0.1406r_i^* in L2(p)L^2(p), with unit n w_n (absorbing scale into fnf_n. Constant functions can take any direction) and, along a subsequence, n→ w_n→ w. Since p()p( x) is continuous and positive, the convergence also holds in Lloc2L^2_loc. Each gng_n is invariant along n⟂ w_n , so ∫gn∂nψ0=0 g_n\, _ v_n _0=0 for every test function ψ0 _0 and n⟂n v_n w_n. Taking n=(−nn⊤)→ v_n=(I- w_n w_n ) v→ v for ⟂ v w and passing to the limit gives ∫ri∗∂ψ0=0 [rgb]0.9297,0.1094,0.1406r_i^*\, _ v _0=0, so the smooth function ri∗ [rgb]0.9297,0.1094,0.1406r_i^* is constant along ⟂ w , contradicting the previous paragraph. Hence δ>0δ>0. Gap in expected divergence. For qf∈q_f with means f [rgb]0.9297,0.1094,0.1406 r_f, 4m4m-strong convexity and the optimality of ∗ [rgb]0.9297,0.1094,0.1406 r^* give R(qf)−R(q∗)=[F(f())−F(∗())]≥2m∥f−∗∥22≥2mδ>0.R(q_f)-R(q^*)=E [F_ x( [rgb]0.9297,0.1094,0.1406 r_f( x))-F_ x( [rgb]0.9297,0.1094,0.1406 r^*( x)) ]≥ 2m\,E [rgb]0.9297,0.1094,0.1406 r_f- [rgb]0.9297,0.1094,0.1406 r^* _2^2≥ 2m\,δ>0. Together with R(qt)→R(q∗)R(q_t)→ R(q^*), this proves the claim. ∎ The fixed-reference specialization of Alg. 1 discretizes this flow. With β=1β=1 and the reference at the prior, its inner step is logit←logit−(1−e−h)∇F()logit [rgb]0.9297,0.1094,0.1406 r [rgb]0.9297,0.1094,0.1406 r-(1-e^-h)∇ F_ x( [rgb]0.9297,0.1094,0.1406 r), that is, mirror descent with mirror map ψ. Since ∇2F⪯(1+λmax()/4)∇2ψ∇^2F_ x (1+ _ ( [rgb]0.1602,0.6719,0.8867 J)/4 )∇^2ψ, the same conclusion holds after finitely many inner steps, at a geometric rate, whenever 1−e−h≤4/(4+λmax())1-e^-h≤ 4/(4+ _ ( [rgb]0.1602,0.6719,0.8867 J)) (Lu et al., 2018). This discrete guarantee requires both β=1β=1 and the stated step-size bound. The condition λmin()>−4 _ ( [rgb]0.1602,0.6719,0.8867 J)>-4 is exactly the condition for the mean-field objective to be strongly convex, so it has a single optimum that BeBeFOND approaches. A width-independent sufficient condition follows for atoms of equal norm c under unit covariance: λmin()≥−c2 _ ( [rgb]0.1602,0.6719,0.8867 J)≥-c^2, with equality when K>MK>M, so it reduces to c<2c<2 in units of the noise standard deviation. For stronger coupling the objective can have several local minima, one per competing explanation, and the proposition does not apply. Rolling-reference inference (section B.1.1) allows the inference budget to be allocated across reference updates. It includes fixed-reference inference by assigning the entire budget to the first inner block and returning its output before refreshing the reference. With a general positive-definite [rgb]0.1602,0.6719,0.8867 , whitening gives the same statement with =⊤−1 [rgb]0.1602,0.6719,0.8867W= [rgb]0.1602,0.6719,0.8867 [rgb]0.1602,0.6719,0.8867 ^-1 [rgb]0.1602,0.6719,0.8867 and =−diag() [rgb]0.1602,0.6719,0.8867 J= [rgb]0.1602,0.6719,0.8867W-diag( [rgb]0.1602,0.6719,0.8867W). The strong-convexity condition and discrete step-size bound use this [rgb]0.1602,0.6719,0.8867 J, while the sufficient norm bound applies to whitened atoms −1/2ϕ(i) [rgb]0.1602,0.6719,0.8867 ^-1/2 [rgb]0.1602,0.6719,0.8867 φ^(i). C.2.3 Rain–sprinkler example Fig. 2 illustrates the two-atom case with independent rain and sprinkler latents, each with prior probability 0.10.1. The coordinates of x are lawn and roof wetness, with: =(2230),=(1.05)2I2. [rgb]0.1602,0.6719,0.8867 = pmatrix2&2\\ 3&0 pmatrix, [rgb]0.1602,0.6719,0.8867 =(1.05)^2I_2. Rain contributes to both coordinates, while the sprinkler contributes only to lawn wetness. The atoms are neither orthogonal nor parallel, so proposition 4.1 applies after whitening. Moreover, λmin()=−4/1.052>−4 _ ( [rgb]0.1602,0.6719,0.8867 J)=-4/1.05^2>-4, satisfying proposition C.2. The exact posterior is computed by enumerating the four latent configurations. Fixed-prior inference. The BeBeFOND panel shows the unique optimal factorized Bernoulli posterior. We solve the stationary equations from proposition C.2 by fixed-point iteration, initialized at =(0.1,0.1) [rgb]0.9297,0.1094,0.1406 r=(0.1,0.1), with β=1β=1 and the generating prior held fixed. The update map has contraction factor 1/1.052<11/1.05^2<1. Iteration continues until the fixed-point residual divided by 1−1/1.0521-1/1.05^2 is below 2×10−132× 10^-13. The panel shows the converged solution of section 3.1 with the decoder and reference fixed. Single-projection encoder. We numerically fit ri()=fi(i⊤) [rgb]0.9297,0.1094,0.1406r_i( x)=f_i( [rgb]0.9297,0.1094,0.1406 w_i x) to minimize the same model-averaged reverse KL to the exact posterior. Each unit direction is parameterized by an angle. A 12×1212× 12 angular grid over [0,π)2[0,π)^2 supplies five starting points for local refinement. For each direction pair, the scalar response functions are solved from the conditional stationarity equations of this KL objective and represented by piecewise-linear probabilities. Final fits use 2,0012,001 grid points per function, extending eight noise standard deviations beyond the projected component means, with constant extrapolation. Conditional expectations during fitting use 160-point Gauss–Hermite quadrature. The fitted rain and sprinkler directions are 65.80∘65.80 and −20.02∘-20.02 , respectively. We evaluate both methods using 160-point Gauss–Hermite quadrature per input coordinate and mixture component. Averaging DKL(q(∣)∥p(∣))D_KL(q( z x)\|p( z x)) over ∼p() x p( x) gives 0.0085080.008508 nats for converged fixed-prior BeBeFOND and 0.0215540.021554 nats for the fitted single-projection encoder. The white point has roof wetness 1.51.5, halfway between the dry and wet conditional means, and therefore equal roof likelihoods with and without rain. Raising roof wetness to 33 while holding lawn wetness fixed strengthens rain and suppresses the sprinkler. Table 3 reports the exact and factorized probabilities at these points. Table 3: Posterior activation probabilities at the marked observations in Figure 2, rounded to six decimal places. Neutral and wet correspond to =(2,1.5)⊤ x=(2,1.5) and =(2,3)⊤ x=(2,3) , respectively. Neutral roof Wet roof Method Rain Sprinkler Rain Sprinkler Exact posterior 0.292138 0.292138 0.960707 0.033017 BeBeFOND 0.228996 0.228996 0.974111 0.019506 Single-projection encoder 0.217514 0.139225 0.958493 0.064003 C.3 Inference and learning ablations C.3.1 Inference ablation: lateral interactions Figures 3 and 6 use K=4,096K=4,096. The baseline comparison uses initialization seeds 0, 1, and 2, with individual results in Table 4. The lateral ablations use seed 0. Metric definitions are in Section C.1.3. The lateral ablations use matched learning rates, β, and decoder norms. Their full BeBeFOND control differs from the selected BeBeFOND recipe in the baseline comparison at d=512,384,256d=512,384,256: the ablations use (LR,β)=(0.001,1),(0.001,1),(0.001,2)(LR,β)=(0.001,1),(0.001,1),(0.001,2), whereas the baseline comparison uses (0.0003,0.5),(0.0003,1),(0.0003,1.5)(0.0003,0.5),(0.0003,1),(0.0003,1.5). Excitatory-only runs at d=1,024d=1,024 and 1,5361,536 stopped early. At d=256d=256, suppressive interactions raise MCC from 0.28100.2810 to 0.68620.6862 within the matched ablation recipes. Excitatory-only inference activates almost all 4,096 latents but reaches MCC 0.28250.2825. Suppressive-only inference has one dead latent at d=1,536d=1,536 and zero at the other dimensions. Both full and zero-off-diagonal inference have zero throughout. Table 4: Individual-seed results for Figure 3 (K=4,096K=4,096, K⋆=16,384K_ =16,384). MCC values retain the recorded precision. Dead latents are counted on 10610^6 held-out examples using each model’s native activity rule (ri>0.5r_i>0.5 for BeBeFOND and Mean-field + Adam). MCC Dead latents Method d Seed 0 Seed 1 Seed 2 0 1 2 BatchTopK 1536 0.6447949409484863 0.6418958902359009 0.6367465853691101 178 217 206 1024 0.7172077894210815 0.7101552486419678 0.715070366859436 78 80 66 768 0.7237797379493713 0.7247610688209534 0.7268125414848328 49 47 39 512 0.6935576796531677 0.6925675868988037 0.6897946000099182 37 38 55 384 0.6238665580749512 0.6256771087646484 0.6227874159812927 56 82 58 256 0.4887199401855469 0.48809993267059326 0.4880616068840027 152 176 157 JumpReLU 1536 0.7402234673500061 0.7397291660308838 0.7378072738647461 0 0 0 1024 0.741813063621521 0.739402711391449 0.7388435006141663 3 2 2 768 0.7435916662216187 0.7454241514205933 0.7452358603477478 2 0 1 512 0.7406731843948364 0.7403227090835571 0.7406840324401855 3 2 1 384 0.7123968005180359 0.7126669883728027 0.712004542350769 78 90 96 256 0.5576815605163574 0.55924391746521 0.5596640110015869 1249 1232 1224 Matryoshka 1536 0.6665480136871338 0.6840691566467285 0.6623464822769165 266 223 279 1024 0.7525583505630493 0.7481846213340759 0.7474654912948608 88 86 85 768 0.7555773258209229 0.7561160326004028 0.7561805248260498 54 51 49 512 0.6980507969856262 0.6989551186561584 0.6928752660751343 54 87 84 384 0.6256262063980103 0.6234501600265503 0.6201913356781006 93 87 85 256 0.5091139078140259 0.5098989605903625 0.5087519288063049 170 153 144 Matching Pursuit 1536 0.7041427493095398 0.7060015201568604 0.7078589200973511 0 0 0 1024 0.7029024362564087 0.7033072113990784 0.6972198486328125 0 0 0 768 0.7092912197113037 0.7096635699272156 0.705768346786499 0 0 0 512 0.7090048789978027 0.7061803340911865 0.7071677446365356 0 0 0 384 0.6933020353317261 0.6994995474815369 0.6963708400726318 0 0 0 256 0.6802412867546082 0.679160475730896 0.6853885054588318 0 0 0 ReLU-L1 1536 0.597491443157196 0.5992552042007446 0.598943829536438 0 0 0 1024 0.59971684217453 0.6019572019577026 0.6006413698196411 0 0 0 768 0.6030566692352295 0.6011125445365906 0.6026996374130249 0 0 0 512 0.5990475416183472 0.5983516573905945 0.5974198579788208 0 0 0 384 0.5921348333358765 0.5912116765975952 0.590924084186554 0 0 0 256 0.5765519738197327 0.5788543224334717 0.5773016214370728 0 0 0 Mean-field + Adam 1536 0.6974799036979675 0.7041385769844055 0.584987998008728 0 0 291 1024 0.7075880169868469 0.7071225643157959 0.7132526636123657 3 9 2 768 0.7050071358680725 0.7113231420516968 0.7108327746391296 8 12 11 512 0.719290018081665 0.7244423627853394 0.7237823009490967 10 3 7 384 0.7211331129074097 0.7238309979438782 0.7178806662559509 14 5 20 256 0.46739497780799866 0.4664973020553589 0.47486913204193115 1645 1655 1664 BeBeFOND 1536 0.8128377795219421 0.8148884773254395 0.8158183097839355 0 0 0 1024 0.8047415018081665 0.8055942058563232 0.8029000759124756 0 0 0 768 0.8140117526054382 0.8091920614242554 0.8121296167373657 0 0 0 512 0.812289297580719 0.8088412880897522 0.8094125986099243 0 0 0 384 0.799635648727417 0.8007169365882874 0.7975541353225708 0 0 0 256 0.7783401012420654 0.7805794477462769 0.7804576754570007 0 0 0 Figure 6: Lateral suppression supports feature recovery under superposition. Hollow markers denote early-stopped excitatory runs (K=4,096K=4,096). C.3.2 Learning ablation: inference and learning geometry The Euclidean ablations also reveal a substantial gain in training efficiency (Fig. 7). Natural inference and learning (N/N) reach MCC 0.5000.500 and Rare F1 0.2800.280 after just 1,000 updates, exceeding both Euclidean variants after 100,000 updates on their respectively tuned runs. Figure 7: The natural-inference/natural-learning recipe reaches higher recovery with fewer samples. Panels show MCC, Macro F1, and Rare F1 on SynthSAEBench (d=768d=768, K=K⋆=16,384K=K_ =16,384). N/E denotes inference/learning geometry: natural (N) or Euclidean (E). Recipes are tuned separately (section C.1.2). The natural-inference/natural-learning recipe reaches stronger recovery in fewer decoder updates. Each geometry uses a separately tuned recipe, with settings given in Table 2. Section C.4 derives equal contraction for frozen, matched second moments and extends the result to unequal moments. C.3.3 Rolling prior Fig. 8 varies β and the cap on the mean initial prior in 72 completed fits: six values of each hyperparameter and seeds 0 and 1. Increasing β reduces final activity in this sweep, but also reduces recovery at the larger values. Several looser caps produce overlapping curves because the learned prior lies below those caps. The initial-prior cap constrains the starting activity, while the rolling-reference dynamics determine activity later in the rollout. Figure 8: Sparsity and recovery versus β on SynthSAEBench (d=768d=768, K=4,096K=4,096, K⋆=16,384K_ =16,384). Colors indicate the mean initial-prior cap, lines show the mean over two seeds, and shading shows their range. Each fit uses 12,000 updates and is evaluated on 8,192 examples after 5,000 inference steps. All fits use raw, unwhitened inputs, batch size 1,024, a learning rate decaying from 0.001 to 10−510^-5, one inner ETD1 step with h=∞h=∞, and an exponential training horizon with mean 5.5 capped at 50. The decoder dictionary starts at norm 2 and is clipped at norm 4. The decoder covariance is fixed, decoder bias is learned, and Fisher damping is 10−510^-5. Initial-prior fits use 100 inference steps on 1,024 examples every 500 updates, with no positive prior floor. Total exposure is 12.3M main plus 24,576 prior examples per fit. C.4 Exact and approximate rarity compensation Fix a training minibatch and hold its inferred distributions, prior references, and covariance ≻0 [rgb]0.1602,0.6719,0.8867 0 fixed during the decoder update. Throughout this subsection, ≻0 [rgb]0.9297,0.1094,0.1406E 0 and p≻0 [rgb]0.1602,0.6719,0.8867E_p 0 denote the minibatch averages of the posterior and prior second moments in eq. 63, and ℱF denotes the resulting minibatch-averaged free energy. In Fisher-preconditioned decoder updates, the matrix gradient is implicitly vectorized columnwise before multiplication by Φ−1G_ ^-1, and the result is reshaped to the decoder matrix. Proof of Theorem 4.2. For the unit-covariance likelihood in the theorem, the decoder-dependent terms are ℱ()=12batchq[‖−‖22]+const.F( [rgb]0.1602,0.6719,0.8867 )= 12E_batchE_q[\| x- [rgb]0.1602,0.6719,0.8867 z\|_2^2]+const. The posterior and prior references are fixed, so the KL term is constant in [rgb]0.1602,0.6719,0.8867 . Expanding the quadratic gives ∇ℱ _ [rgb]0.1602,0.6719,0.8867 F =−batch[⊤], = [rgb]0.1602,0.6719,0.8867 [rgb]0.9297,0.1094,0.1406E-E_batch[ x [rgb]0.9297,0.1094,0.1406 r ], ∇vec()2ℱ ∇^2_vec( [rgb]0.1602,0.6719,0.8867 )F =⊗≻0. = [rgb]0.9297,0.1094,0.1406E 0. (92) Hence the minimizer is unique, ∗=batch[⊤]−1 [rgb]0.1602,0.6719,0.8867 ^*=E_batch[ x [rgb]0.9297,0.1094,0.1406 r ] [rgb]0.9297,0.1094,0.1406E^-1, and ∇ℱ=(−∗) _ [rgb]0.1602,0.6719,0.8867 F=( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E. For a model draw ~|∼(,) x z ( [rgb]0.1602,0.6719,0.8867 z,I), the decoder score is ∇logp(~∣)=(~−)⊤. _ [rgb]0.1602,0.6719,0.8867 p_ [rgb]0.1602,0.6719,0.8867 ( x z)=( x- [rgb]0.1602,0.6719,0.8867 z) z . The conditional residual has second moment I. Averaging the vectorized score’s outer product over these model draws and the fixed prior references therefore gives Φ=p⊗G_ = [rgb]0.1602,0.6719,0.8867E_p . Its inverse acts on a matrix gradient by right multiplication with p−1 [rgb]0.1602,0.6719,0.8867E_p^-1. Thus, when p= [rgb]0.1602,0.6719,0.8867E_p= [rgb]0.9297,0.1094,0.1406E, Euclidean:dτ(−∗) : ddτ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) =−ηΦ(−∗), =- _ ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E, Natural:dτ(−∗) : ddτ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) =−ηΦ(−∗). =- _ ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*). Integrating the natural flow yields (Δt)−∗=e−ηΦΔt((0)−∗). [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*=e^- _ t( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^*). (93) The matching condition is p= [rgb]0.1602,0.6719,0.8867E_p= [rgb]0.9297,0.1094,0.1406E for the minibatch-averaged second moments. To compare relaxation times, let j=λjj [rgb]0.9297,0.1094,0.1406E v_j= _j v_j. Since ≻0 [rgb]0.9297,0.1094,0.1406E 0, its eigenvectors form an orthonormal basis and λj>0 _j>0. Along the Euclidean trajectory, ((Δt)−∗)j=e−ηΦλjΔt((0)−∗)j.( [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*) v_j=e^- _ _j t( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^*) v_j. Its modal relaxation times are 1/(ηΦλj)1/( _ _j), so the ratio of the slowest to the fastest is λmax()/λmin()=κ() _ ( [rgb]0.9297,0.1094,0.1406E)/ _ ( [rgb]0.9297,0.1094,0.1406E)=κ( [rgb]0.9297,0.1094,0.1406E). Natural flow instead has relaxation time 1/ηΦ1/ _ in every direction. Since zi2=ziz_i^2=z_i for Bernoulli latents, ρi=ii=batchq[zi]=batch[ri]. _i= [rgb]0.9297,0.1094,0.1406E_i=E_batchE_q[z_i]=E_batch[ [rgb]0.9297,0.1094,0.1406r_i]. The Rayleigh quotient bounds each diagonal entry between the smallest and largest eigenvalues, giving λmax()≥ρmax,λmin()≤ρmin,κ()≥ρmaxρmin. _ ( [rgb]0.9297,0.1094,0.1406E)≥ _ , _ ( [rgb]0.9297,0.1094,0.1406E)≤ _ , κ( [rgb]0.9297,0.1094,0.1406E)≥ _ _ . For one feature with known activity, =p [rgb]0.9297,0.1094,0.1406E=p, recovering the Euclidean rate ηΦp _ p and natural rate ηΦ _ used in the scalar example. ∎ Corollary C.3 (Power-law feature frequencies). Under theorem 4.2, suppose the averaged posterior activation probabilities satisfy ρi=ci−α _i=c\,i^-α for i=1,…,Ki=1,…,K, with c>0c>0 and α>0α>0. The ratio of the slowest to the fastest Euclidean relaxation time is at least KαK^α; whereas, for the matched natural-gradient flow it is one. Proof. The largest and smallest frequencies are c and cK−αcK^-α. Substituting their ratio into the theorem gives κ()≥Kακ( [rgb]0.9297,0.1094,0.1406E)≥ K^α. eq. 93 gives the common natural-gradient relaxation time. ∎ With the inferred codes fixed and the averaged second moments matched, natural-gradient flow has the same local contraction rate in every decoder direction. General observation covariance. For fixed ≻0 [rgb]0.1602,0.6719,0.8867 0, the corresponding gradient, Hessian, and decoder Fisher are ∇ℱ _ [rgb]0.1602,0.6719,0.8867 F =−1(−∗), = [rgb]0.1602,0.6719,0.8867 ^-1( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E, ∇vec()2ℱ ∇^2_vec( [rgb]0.1602,0.6719,0.8867 )F =⊗−1,Φ=p⊗−1. = [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867 ^-1, _ = [rgb]0.1602,0.6719,0.8867E_p [rgb]0.1602,0.6719,0.8867 ^-1. (94) The optimum is unchanged. Euclidean flow now obeys ˙=−ηΦ−1(−∗) [rgb]0.1602,0.6719,0.8867 =- _ [rgb]0.1602,0.6719,0.8867 ^-1( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E; its Hessian has condition number κ()κ()κ( [rgb]0.9297,0.1094,0.1406E)κ( [rgb]0.1602,0.6719,0.8867 ) because the Kronecker eigenvalues are products. Natural preconditioning cancels the observation precision and gives dτ(−∗)=−ηΦ(−∗)p−1. ddτ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*)=- _ ( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1. The coefficients are fixed during the update, hence (Δt)−∗=((0)−∗)exp[−ηΦΔtp−1]. [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*=( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^*) [- _ t\, [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1]. (95) When p= [rgb]0.1602,0.6719,0.8867E_p= [rgb]0.9297,0.1094,0.1406E, this again reduces to eq. 93, independently of the fixed observation covariance. The following extension treats unequal moments. Proposition C.4 (Rarity compensation with unequal moments). Under the same frozen-substep assumptions, let ≻0 [rgb]0.9297,0.1094,0.1406E 0, p≻0 [rgb]0.1602,0.6719,0.8867E_p 0, and ηΦ>0 _ >0, without requiring =p [rgb]0.9297,0.1094,0.1406E= [rgb]0.1602,0.6719,0.8867E_p. The decoder error obeys eq. 95. Define λ∗:=λmin(p−1/2p−1/2)>0. _*:= _ \! ( [rgb]0.1602,0.6719,0.8867E_p^-1/2 [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1/2 )>0. Then, for Δt≥0 t≥ 0, ‖((Δt)−∗)p1/2‖F≤e−ηΦλ∗Δt‖((0)−∗)p1/2‖F.\|( [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2\|_F≤ e^- _ _* t\|( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2\|_F. (96) In particular, if (1−ϵ)p⪯(1+ϵ)p,0≤ϵ<1,(1-ε) [rgb]0.1602,0.6719,0.8867E_p [rgb]0.9297,0.1094,0.1406E (1+ε) [rgb]0.1602,0.6719,0.8867E_p, 0≤ε<1, then ‖((Δt)−∗)p1/2‖F≤e−ηΦ(1−ϵ)Δt‖((0)−∗)p1/2‖F.\|( [rgb]0.1602,0.6719,0.8867 ( t)- [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2\|_F≤ e^- _ (1-ε) t\|( [rgb]0.1602,0.6719,0.8867 (0)- [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2\|_F. (97) Proof. eq. 95 follows from the quadratic decoder gradient derived above. Multiplying its error dynamics on the right by p1/2 [rgb]0.1602,0.6719,0.8867E_p^1/2 gives: dτ[(−∗)p1/2]=−ηΦ[(−∗)p1/2](p−1/2p−1/2). ddτ [( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2 ]=- _ [( [rgb]0.1602,0.6719,0.8867 - [rgb]0.1602,0.6719,0.8867 ^*) [rgb]0.1602,0.6719,0.8867E_p^1/2 ] ( [rgb]0.1602,0.6719,0.8867E_p^-1/2 [rgb]0.9297,0.1094,0.1406E [rgb]0.1602,0.6719,0.8867E_p^-1/2 ). The matrix in the final parentheses is symmetric positive definite, so its decay exponential has spectral norm e−ηΦλ∗Δte^- _ _* t. Applying ‖AB‖F≤‖A‖F‖B‖2\|AB\|_F≤\|A\|_F\|B\|_2 proves eq. 96. Under the stated moment bounds, this matrix has eigenvalues in [1−ϵ,1+ϵ][1-ε,1+ε], hence λ∗≥1−ϵ _*≥ 1-ε, which proves eq. 97. ∎ C.4.1 Activity-normalized inference Theorem C.5 (Activity-normalized inference). Fix the observation, decoder, covariance, and current prior reference during inference. For the factorized Bernoulli posterior, let ri=σ(ui)∈(0,1) [rgb]0.9297,0.1094,0.1406r_i=σ( [rgb]0.9297,0.1094,0.1406u_i)∈(0,1). At the same posterior state, the unit-rate Euclidean and natural-gradient flows in membrane-potential (logit) coordinates satisfy u˙iNGD=−1ri(1−ri)∂ℱ∂ui,u˙iGD=ri(1−ri)u˙iNGD. $ [rgb]0.9297,0.1094,0.1406u_i^NGD=- 1 [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) ∂ [rgb]0.9297,0.1094,0.1406u_i, [rgb]0.9297,0.1094,0.1406u_i^GD= [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i)\, [rgb]0.9297,0.1094,0.1406u_i^NGD.$ Thus natural-gradient inference removes the multiplicative Bernoulli variance factor from each membrane-potential update. Proof of Theorem C.5. For factorized Bernoulli latents, ri=σ(ui)∈(0,1) [rgb]0.9297,0.1094,0.1406r_i=σ( [rgb]0.9297,0.1094,0.1406u_i)∈(0,1), so ∂ri∂ui=ri(1−ri),=diag[⊙(−)]. ∂ [rgb]0.9297,0.1094,0.1406r_i∂ [rgb]0.9297,0.1094,0.1406u_i= [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i), [rgb]0.9297,0.1094,0.1406C=diag[ [rgb]0.9297,0.1094,0.1406 r ( 1- [rgb]0.9297,0.1094,0.1406 r)]. The score is ∂logq(∣)/∂ui=zi−ri∂ q_ [rgb]0.9297,0.1094,0.1406 u( z x)/∂ [rgb]0.9297,0.1094,0.1406u_i=z_i- [rgb]0.9297,0.1094,0.1406r_i. Independence therefore gives the posterior Fisher ()= G( [rgb]0.9297,0.1094,0.1406 u)= [rgb]0.9297,0.1094,0.1406C. The chain rule yields ∇ℱ=∇ℱ _ [rgb]0.9297,0.1094,0.1406 uF= [rgb]0.9297,0.1094,0.1406C _ [rgb]0.9297,0.1094,0.1406 rF. Natural preconditioning cancels this factor: u˙iNGD=−1ri(1−ri)∂ℱ∂ui=−∂ℱ∂ri,u˙iGD=ri(1−ri)u˙iNGD. [rgb]0.9297,0.1094,0.1406u_i^NGD=- 1 [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) ∂ [rgb]0.9297,0.1094,0.1406u_i=- ∂ [rgb]0.9297,0.1094,0.1406r_i, [rgb]0.9297,0.1094,0.1406u_i^GD= [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) [rgb]0.9297,0.1094,0.1406u_i^NGD. ∎ When the inference drive favors activation, this cancellation removes an additional suppression of weakly active latents’ logit updates. When the inference drive favors activation, this cancellation removes the variance factor from weakly active latents’ logit updates. Activation probabilities evolve according to r˙i=ri(1−ri)u˙i [rgb]0.9297,0.1094,0.1406r_i= [rgb]0.9297,0.1094,0.1406r_i(1- [rgb]0.9297,0.1094,0.1406r_i) [rgb]0.9297,0.1094,0.1406u_i. C.4.2 Local conditioning under Adam normalization Adam normalizes updates using coordinatewise estimates of gradient second moments (Kingma and Ba, 2015; Balles and Hennig, 2018). The following calculation isolates the effect of this normalization on decoder conditioning. Proposition C.6 (Local conditioning under frozen Adam normalization). Let =∗+ x= [rgb]0.1602,0.6719,0.8867 _* z+ ξ, where the known codes have independent coordinates zi∼Bernoulli(ρi)z_i ( _i) with 0<ρi<10< _i<1, and ∼(,d) ξ ( 0,I_d) is independent of z. For the unconstrained population loss L()=12‖−‖22,=[⊤],=⊗d,L( [rgb]0.1602,0.6719,0.8867 )= 12E\| x- [rgb]0.1602,0.6719,0.8867 z\|_2^2, =E[ z z ], H=E _d, consider gradients averaged over B independent samples. Freeze the Adam denominator at the square root of its expected second-moment target at ∗ [rgb]0.1602,0.6719,0.8867 _*, with stabilizer ϵ>0ε>0. The resulting diagonal preconditioner is ϵ=diagi=1K[(ϵ+ρi/B)−1]⊗d.D_ε=diag_i=1^K\! [(ε+ _i/B)^-1 ] _d. Writing ϵ=ϵ1/2ϵ1/2A_ε=D_ε^1/2 HD_ε^1/2, its condition number satisfies κ(ϵ)≥ρmax/(ϵ+ρmax/B)ρmin/(ϵ+ρmin/B)≥ρmaxρmin.κ(A_ε)≥ _ /(ε+ _ /B) _ /(ε+ _ /B)≥ _ _ . (98) Thus ρi=ci−α _i=c\,i^-α, with 0<c<10<c<1 and α>0α>0, gives κ(ϵ)≥Kα/2κ(A_ε)≥ K^α/2. Matched complete-data Fisher preconditioning instead gives −1/2−1/2= H^-1/2 H H^-1/2=I and condition number one. Proof. The Hessian is H, and =diag(ρi(1−ρi))+⊤≻0E=diag( _i(1- _i))+ ρ ρ 0. At the optimum, a single-sample gradient is −⊤- ξ z , with mean zero and vectorized second moment H. Independence gives second moment /B H/B for the averaged gradient, yielding ϵD_ε above. Each diagonal entry of ϵA_ε is ai=ρi/(ϵ+ρi/B)a_i= _i/(ε+ _i/B), repeated d times. Rayleigh quotients give κ(ϵ)≥maxiai/miniaiκ(A_ε)≥ _ia_i/ _ia_i. Since aia_i increases with ρi _i, the first bound follows. For s=ρmax/ρmin≥1s= _ / _ ≥ 1 and t=ρmin/Bt= _ /B, the ratio is s2(ϵ+t)/(ϵ+st)≥s^2(ε+t)/(ε+st)≥ s. Substitution gives the width bound. The complete-data Fisher equals H for this model, proving the final claim. ∎ This local conditioning analysis uses known latent codes and fixes Adam’s gradient second-moment estimate at its expected value at the decoder optimum, before taking the square root to form the normalization. Adam’s normalization depends on gradient variance (Balles and Hennig, 2018), while the model Fisher measures distributional sensitivity (Martens, 2020). Section B.8 discusses the computational cost of the implemented Fisher updates. C.5 Learning and rare-feature recovery C.5.1 Resampling and auxiliary-loss variants Resampling and auxiliary losses aim to revive inactive features (Bricken et al., 2023; Jermyn and Templeton, 2024). Fig. 9 and Table 5 compare the original baselines with modified recipes using MinFire (BatchTopK), resampling (JumpReLU and Matryoshka), learning-rate decay, and retuned sparsity. These modifications substantially improve baseline recovery, but BeBeFOND retains the highest MCC, Macro F1, and Rare F1. The comparison evaluates complete recipes with different target sparsities. Measured L0L_0 is reported in the table. The modified recipes follow the released full-width sweep (chanind-goodfire, 2026), using seed 0, 200M samples, batch size 1,024, Adam with learning rate 3×10−43× 10^-4 and moments (0.9,0.999)(0.9,0.999), BF16 autocast, gradient-norm clipping at 1, and learning-rate decay during the final 20% of training. Target L0L_0 is 45 for BatchTopK, 40 for JumpReLU, and 35 for Matryoshka. MinFire (the sweep implementation’s MinFireGate in gates.py) reserves each latent’s strongest positive activation before filling the remaining batch-wide TopK budget. The selected BatchTopK recipe applies this gate every batch for the first 90% of training, then uses ordinary BatchTopK. JumpReLU uses an input straight-through estimator and a base L0L_0 coefficient of 0.1 with an adaptive multiplier. JumpReLU and Matryoshka resample every 2,500 updates during the first 80% of training, with inactivity windows of 1,000 and 100 updates, respectively. Resampled decoder atoms follow reconstruction-residual directions, and their optimizer moments are reset. BatchTopK and Matryoshka retain auxiliary-loss coefficient 1. Matryoshka uses prefixes 128, 512, 2,048, and 16,384 with per-level auxiliary losses. Gemma Scope uses reconstruction and L0L_0 sparsity losses, without a separate dead-latent auxiliary loss, alongside straight-through estimators, tied initialization, learning-rate and sparsity warmup, and decoder normalization (Lieberum et al., 2024, Sections 2.2 and 3.2). Figure 9: BeBeFOND retains the highest MCC, Macro F1, and Rare F1 after baseline modifications aimed at preventing inactive features. Hatched bars denote modified recipes (d=768d=768, K=K⋆=16,384K=K_ =16,384, seed 0); numerical results appear in Table 5. Table 5: Full-width SynthSAEBench results (d=768d=768, K=K⋆=16,384K=K_ =16,384). Individual rows report initialization seeds 0–2; mean rows give pointwise 95% Student-t confidence intervals for the three recovery metrics, as in Figure 4. Each evaluation uses 10610^6 held-out examples. L0L_0 and dead counts in mean rows are arithmetic means. The lower block retains the modified recipes (++) and continuous-run checkpoint, each at seed 0. The BeBeFOND runs use 181.82M decoder-update plus 18.18M prior examples. Method Seed MCC ↑ Macro F1 ↑ Rare F1 ↑ L0L_0 Dead ↓ BatchTopK 0 0.534523 0.445173 0.299436 34.833798 1,577 1 0.535311 0.444452 0.299666 35.031285 1,733 2 0.537442 0.446642 0.296203 34.948331 1,707 Mean 0.5358±0.00380.5358± 0.0038 0.4454±0.00280.4454± 0.0028 0.2984±0.00480.2984± 0.0048 34.94 1672.33 JumpReLU 0 0.483953 0.326567 0.190439 45.148172 8,069 1 0.481728 0.326672 0.184864 44.826161 8,090 2 0.482102 0.325518 0.184310 45.098772 8,117 Mean 0.4826±0.00300.4826± 0.0030 0.3263±0.00160.3263± 0.0016 0.1865±0.00840.1865± 0.0084 45.02 8092.00 Matryoshka 0 0.590345 0.556314 0.433018 19.978349 4,923 1 0.574629 0.530946 0.366127 19.991146 5,505 2 0.583862 0.552021 0.415913 19.989287 5,049 Mean 0.5829±0.01960.5829± 0.0196 0.5464±0.03370.5464± 0.0337 0.4050±0.08630.4050± 0.0863 19.99 5159.00 Matching Pursuit 0 0.812752 0.782551 0.822595 29.646241 2 1 0.811610 0.781264 0.821854 29.648383 1 2 0.815241 0.784698 0.827665 29.688899 3 Mean 0.8132±0.00460.8132± 0.0046 0.7828±0.00430.7828± 0.0043 0.8240±0.00790.8240± 0.0079 29.66 2.00 ReLU-L1 0 0.651464 0.641205 0.522534 35.089073 174 1 0.650650 0.640259 0.524647 34.972258 135 2 0.650945 0.641409 0.524073 34.976428 137 Mean 0.6510±0.00100.6510± 0.0010 0.6410±0.00150.6410± 0.0015 0.5238±0.00270.5238± 0.0027 35.01 148.67 Mean-field + Adam 0 0.779715 0.841050 0.704307 45.107146 514 1 0.783435 0.846995 0.743286 48.116016 640 2 0.790151 0.855224 0.729628 68.019993 556 Mean 0.7844±0.01310.7844± 0.0131 0.8478±0.01770.8478± 0.0177 0.7257±0.04910.7257± 0.0491 53.75 570.00 BeFOND 0 0.898559 0.970345 0.955392 36.033549 1 1 0.898064 0.970567 0.949700 35.679190 0 2 0.897350 0.969745 0.950726 35.485085 0 Mean 0.8980±0.00150.8980± 0.0015 0.9702±0.00110.9702± 0.0011 0.9519±0.00750.9519± 0.0075 35.73 0.33 Additional seed-0 checkpoints BatchTopK+ 0 0.8123 0.6826 0.5851 45.01 504 JumpReLU+ 0 0.8612 0.8487 0.7288 39.96 367 Matryoshka+ 0 0.8405 0.8348 0.8457 34.95 495 BeBeFOND 0 0.8974 0.9701 0.9482 35.70 0 Recovery across feature frequencies. Fig. 10 resolves the comparison into frequency quintiles. The BeBeFOND checkpoint is the best run in Table 5 (MCC 0.8986), and the original and modified baseline checkpoints are those in Fig. 9. Improvements from the modified recipes extend into the rare tail, while BeBeFOND retains higher coverage and F1 in every quintile. Coverage measures dictionary alignment. F1 measures detection of the matched feature’s activity, so the two need not coincide. Figure 10: Recovery across feature-frequency quintiles (d=768d=768, K=K⋆=16,384K=K_ =16,384, seed 0). Dashed lines show original baselines, and “+” denotes their modified recipes. Coverage is the fraction of features matched above cosine 0.8. Quintiles use each run’s 500K-example frequency reference. F1 uses 1M held-out examples. Random draws differ between some runs. These fifths differ from Rare F1’s least-frequent quarter. Dictionary capacity. Table 5 combines recovery, sparsity, and dead-latent counts. Main-text Fig. 4 reports three-seed means: MCC 0.8980 and Rare F1 0.9519, with 0.33 dead latents on average. Its seed-0 run reaches MCC 0.8986 and Rare F1 0.9554, with one dead latent. The separate continuous run below reaches MCC 0.8974 and Rare F1 0.9482, with zero. MP has two dead latents but lower recovery, so keeping latents active is not sufficient for accurate feature recovery. C.5.2 Training sample efficiency Fig. 11 reports MCC, Macro F1, and Rare F1 throughout training in the full-width setting. Its BeBeFOND curve follows one continuous run, whose endpoint is identified separately in Table 5. Figure 11: Feature recovery throughout training on unwhitened SynthSAEBench (d=768d=768; 16,384 learned and ground-truth features), using plain baseline recipes and seed 0. (a) Dictionary mean correlation coefficient (MCC). (b) Macro F1 over ground-truth features after positive-cosine one-to-one matching. (c) Rare F1 over the least frequent quartile, defined by a shared independent 500K-example frequency reference (seed 315031). Every checkpoint is evaluated on 1M held-out examples (seed 49100). BeFOND uses 200 inference steps with activity p>0.5p>0.5; baselines use positive native activations. The x-axis counts decoder-update examples and is piecewise linear: 1M examples at 0–10M, 5M at 10–50M, and 25M beyond 50M occupy equal horizontal distances. Cross-region slopes therefore do not have a common per-sample interpretation. Baseline points through 10M, after 10M through 50M, and beyond 50M come from separate matching-recipe runs with their original full-budget learning-rate schedules; they do not form a single continuous trajectory. Overlapping checkpoints are selected by these fixed windows, not by score. BeFOND uses one continuous run. Samples are rounded up to full batches. Baselines finish at approximately 200M decoder-update examples and BeFOND at 181.82M plus 18.18M separate prior-update examples, which are excluded from the x-axis. Baseline auxiliary exposure is not counted on this axis. Final measured L0L_0 values are BatchTopK: 34.83, JumpReLU: 45.15, Matryoshka: 19.98, Matching Pursuit: 29.65, ReLU-L1: 35.09, BeFOND: 35.70. One seed is shown; no uncertainty intervals are estimated. Appendix D Extended Gemma experiments Table 6: SAEBench results on whitened Gemma-2-2B layer-12 activations after 49,999,87249,999,872 training tokens. Entries are mean ± sample standard deviation over three training seeds for BeBeFOND and five for each baseline. For SAE baselines, we selected target L0L_0 values by mean top-1 accuracy, and then used it for all probing columns, and the model selected by mean disentanglement supplies all RAVEL columns. Parentheses give both values when they differ. Sparse probing measures how well k=1k=1 or k=5k=5 features predict concepts, averaged equally across 113 tasks: accuracy is the fraction of correct predictions; AUC measures how well positives rank above negatives across thresholds; F1 balances precision (correctness of positive predictions) and recall (coverage of positives). RAVEL measures whether interventions change a targeted attribute (cause) while preserving other attributes (isolation), and disentanglement is their mean. Higher is better throughout. Within each width and metric, bold marks the highest mean, selected before rounding. Sparse probing RAVEL Model k=1k=1 k=5k=5 Acc. ↑ AUC ↑ F1 ↑ Acc. ↑ AUC ↑ F1 ↑ Disentangle- ment ↑ Cause ↑ Isolation ↑ Dictionary width: 16k BatchTopK (target L0=400L_0=400) .767±.002.767 [rgb]0.3984,0.3984,0.3984±.002 .798±.003.798 [rgb]0.3984,0.3984,0.3984±.003 .761±.004.761 [rgb]0.3984,0.3984,0.3984±.004 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .862±.002.862 [rgb]0.3984,0.3984,0.3984±.002 .812±.002.812 [rgb]0.3984,0.3984,0.3984±.002 .729±.005.729 [rgb]0.3984,0.3984,0.3984±.005 .692±.019.692 [rgb]0.3984,0.3984,0.3984±.019 .767±.017.767 [rgb]0.3984,0.3984,0.3984±.017 Matryoshka BatchTopK (L0=200L_0=200 probing; 400400 RAVEL) .768±.003.768 [rgb]0.3984,0.3984,0.3984±.003 .794±.003.794 [rgb]0.3984,0.3984,0.3984±.003 .762±.002.762 [rgb]0.3984,0.3984,0.3984±.002 .810±.002.810 [rgb]0.3984,0.3984,0.3984±.002 .858±.002.858 [rgb]0.3984,0.3984,0.3984±.002 .808±.002.808 [rgb]0.3984,0.3984,0.3984±.002 .731±.003.731 [rgb]0.3984,0.3984,0.3984±.003 .695±.013.695 [rgb]0.3984,0.3984,0.3984±.013 .767±.015.767 [rgb]0.3984,0.3984,0.3984±.015 Mean-field + Adam (T=20T=20 inference steps) .761±.003.761 [rgb]0.3984,0.3984,0.3984±.003 .761±.003.761 [rgb]0.3984,0.3984,0.3984±.003 .749±.004.749 [rgb]0.3984,0.3984,0.3984±.004 .801±.002.801 [rgb]0.3984,0.3984,0.3984±.002 .842±.001.842 [rgb]0.3984,0.3984,0.3984±.001 .798±.002.798 [rgb]0.3984,0.3984,0.3984±.002 .696±.003.696 [rgb]0.3984,0.3984,0.3984±.003 .639±.006.639 [rgb]0.3984,0.3984,0.3984±.006 .753±.006.753 [rgb]0.3984,0.3984,0.3984±.006 BeBeFOND (T=10T=10 inference steps) .771±.004.771 [rgb]0.3984,0.3984,0.3984±.004 .781±.004.781 [rgb]0.3984,0.3984,0.3984±.004 .765±.005.765 [rgb]0.3984,0.3984,0.3984±.005 .812±.002.812 [rgb]0.3984,0.3984,0.3984±.002 .856±.001.856 [rgb]0.3984,0.3984,0.3984±.001 .811±.002.811 [rgb]0.3984,0.3984,0.3984±.002 .786±.002.786 [rgb]0.3984,0.3984,0.3984±.002 .759±.002.759 [rgb]0.3984,0.3984,0.3984±.002 .813±.003.813 [rgb]0.3984,0.3984,0.3984±.003 Dictionary width: 65k BatchTopK (L0=400L_0=400 probing; 800800 RAVEL) .770±.003.770 [rgb]0.3984,0.3984,0.3984±.003 .800±.002.800 [rgb]0.3984,0.3984,0.3984±.002 .765±.003.765 [rgb]0.3984,0.3984,0.3984±.003 .814±.003.814 [rgb]0.3984,0.3984,0.3984±.003 .862±.002.862 [rgb]0.3984,0.3984,0.3984±.002 .812±.003.812 [rgb]0.3984,0.3984,0.3984±.003 .745±.005.745 [rgb]0.3984,0.3984,0.3984±.005 .712±.006.712 [rgb]0.3984,0.3984,0.3984±.006 .777±.008.777 [rgb]0.3984,0.3984,0.3984±.008 Matryoshka BatchTopK (target L0=400L_0=400) .772±.001.772 [rgb]0.3984,0.3984,0.3984±.001 .801±.001.801 [rgb]0.3984,0.3984,0.3984±.001 .766±.002.766 [rgb]0.3984,0.3984,0.3984±.002 .811±.001.811 [rgb]0.3984,0.3984,0.3984±.001 .858±.001.858 [rgb]0.3984,0.3984,0.3984±.001 .809±.002.809 [rgb]0.3984,0.3984,0.3984±.002 .737±.004.737 [rgb]0.3984,0.3984,0.3984±.004 .704±.007.704 [rgb]0.3984,0.3984,0.3984±.007 .771±.006.771 [rgb]0.3984,0.3984,0.3984±.006 Mean-field + Adam (T=20T=20 inference steps) .757±.002.757 [rgb]0.3984,0.3984,0.3984±.002 .757±.002.757 [rgb]0.3984,0.3984,0.3984±.002 .744±.002.744 [rgb]0.3984,0.3984,0.3984±.002 .794±.001.794 [rgb]0.3984,0.3984,0.3984±.001 .832±.001.832 [rgb]0.3984,0.3984,0.3984±.001 .790±.002.790 [rgb]0.3984,0.3984,0.3984±.002 .733±.003.733 [rgb]0.3984,0.3984,0.3984±.003 .697±.007.697 [rgb]0.3984,0.3984,0.3984±.007 .769±.006.769 [rgb]0.3984,0.3984,0.3984±.006 BeBeFOND (T=10T=10 inference steps) .776±.004.776 [rgb]0.3984,0.3984,0.3984±.004 .794±.004.794 [rgb]0.3984,0.3984,0.3984±.004 .769±.006.769 [rgb]0.3984,0.3984,0.3984±.006 .812±.002.812 [rgb]0.3984,0.3984,0.3984±.002 .856±.001.856 [rgb]0.3984,0.3984,0.3984±.001 .811±.002.811 [rgb]0.3984,0.3984,0.3984±.002 .814±.000.814 [rgb]0.3984,0.3984,0.3984±.000 .789±.001.789 [rgb]0.3984,0.3984,0.3984±.001 .840±.001.840 [rgb]0.3984,0.3984,0.3984±.001 Dictionary width: 131k BatchTopK (target L0=600L_0=600) .773±.002.773 [rgb]0.3984,0.3984,0.3984±.002 .804±.002.804 [rgb]0.3984,0.3984,0.3984±.002 .768±.003.768 [rgb]0.3984,0.3984,0.3984±.003 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .862±.001.862 [rgb]0.3984,0.3984,0.3984±.001 .813±.002.813 [rgb]0.3984,0.3984,0.3984±.002 .751±.003.751 [rgb]0.3984,0.3984,0.3984±.003 .720±.008.720 [rgb]0.3984,0.3984,0.3984±.008 .782±.013.782 [rgb]0.3984,0.3984,0.3984±.013 Matryoshka BatchTopK (L0=200L_0=200 probing; 400400 RAVEL) .773±.004.773 [rgb]0.3984,0.3984,0.3984±.004 .799±.004.799 [rgb]0.3984,0.3984,0.3984±.004 .768±.004.768 [rgb]0.3984,0.3984,0.3984±.004 .814±.001.814 [rgb]0.3984,0.3984,0.3984±.001 .860±.002.860 [rgb]0.3984,0.3984,0.3984±.002 .812±.001.812 [rgb]0.3984,0.3984,0.3984±.001 .739±.002.739 [rgb]0.3984,0.3984,0.3984±.002 .713±.008.713 [rgb]0.3984,0.3984,0.3984±.008 .764±.006.764 [rgb]0.3984,0.3984,0.3984±.006 Mean-field + Adam (T=20T=20 inference steps) .754±.003.754 [rgb]0.3984,0.3984,0.3984±.003 .754±.003.754 [rgb]0.3984,0.3984,0.3984±.003 .740±.005.740 [rgb]0.3984,0.3984,0.3984±.005 .791±.002.791 [rgb]0.3984,0.3984,0.3984±.002 .830±.003.830 [rgb]0.3984,0.3984,0.3984±.003 .787±.002.787 [rgb]0.3984,0.3984,0.3984±.002 .734±.005.734 [rgb]0.3984,0.3984,0.3984±.005 .694±.007.694 [rgb]0.3984,0.3984,0.3984±.007 .774±.004.774 [rgb]0.3984,0.3984,0.3984±.004 BeBeFOND (T=10T=10 inference steps) .785±.002.785 [rgb]0.3984,0.3984,0.3984±.002 .802±.003.802 [rgb]0.3984,0.3984,0.3984±.003 .778±.003.778 [rgb]0.3984,0.3984,0.3984±.003 .816±.001.816 [rgb]0.3984,0.3984,0.3984±.001 .857±.002.857 [rgb]0.3984,0.3984,0.3984±.002 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .816±.001.816 [rgb]0.3984,0.3984,0.3984±.001 .790±.002.790 [rgb]0.3984,0.3984,0.3984±.002 .842±.003.842 [rgb]0.3984,0.3984,0.3984±.003 Dictionary width: 524k BatchTopK (target L0=800L_0=800) .773±.004.773 [rgb]0.3984,0.3984,0.3984±.004 .803±.004.803 [rgb]0.3984,0.3984,0.3984±.004 .768±.006.768 [rgb]0.3984,0.3984,0.3984±.006 .814±.003.814 [rgb]0.3984,0.3984,0.3984±.003 .860±.002.860 [rgb]0.3984,0.3984,0.3984±.002 .813±.003.813 [rgb]0.3984,0.3984,0.3984±.003 .748±.004.748 [rgb]0.3984,0.3984,0.3984±.004 .719±.010.719 [rgb]0.3984,0.3984,0.3984±.010 .776±.010.776 [rgb]0.3984,0.3984,0.3984±.010 Matryoshka BatchTopK (L0=200L_0=200 probing; 400400 RAVEL) .774±.002.774 [rgb]0.3984,0.3984,0.3984±.002 .799±.002.799 [rgb]0.3984,0.3984,0.3984±.002 .768±.002.768 [rgb]0.3984,0.3984,0.3984±.002 .813±.001.813 [rgb]0.3984,0.3984,0.3984±.001 .859±.001.859 [rgb]0.3984,0.3984,0.3984±.001 .812±.000.812 [rgb]0.3984,0.3984,0.3984±.000 .743±.003.743 [rgb]0.3984,0.3984,0.3984±.003 .708±.013.708 [rgb]0.3984,0.3984,0.3984±.013 .778±.009.778 [rgb]0.3984,0.3984,0.3984±.009 Mean-field + Adam (T=20T=20 inference steps) .751±.007.751 [rgb]0.3984,0.3984,0.3984±.007 .751±.007.751 [rgb]0.3984,0.3984,0.3984±.007 .735±.009.735 [rgb]0.3984,0.3984,0.3984±.009 .786±.009.786 [rgb]0.3984,0.3984,0.3984±.009 .825±.010.825 [rgb]0.3984,0.3984,0.3984±.010 .782±.009.782 [rgb]0.3984,0.3984,0.3984±.009 .685±.029.685 [rgb]0.3984,0.3984,0.3984±.029 .622±.047.622 [rgb]0.3984,0.3984,0.3984±.047 .748±.013.748 [rgb]0.3984,0.3984,0.3984±.013 BeBeFOND (T=10T=10 inference steps) .789±.002.789 [rgb]0.3984,0.3984,0.3984±.002 .805±.004.805 [rgb]0.3984,0.3984,0.3984±.004 .784±.004.784 [rgb]0.3984,0.3984,0.3984±.004 .817±.001.817 [rgb]0.3984,0.3984,0.3984±.001 .859±.001.859 [rgb]0.3984,0.3984,0.3984±.001 .816±.001.816 [rgb]0.3984,0.3984,0.3984±.001 .827±.001.827 [rgb]0.3984,0.3984,0.3984±.001 .800±.001.800 [rgb]0.3984,0.3984,0.3984±.001 .854±.002.854 [rgb]0.3984,0.3984,0.3984±.002 Figure 12: Performance as a function of inference depth. Sparse probing top-1 accuracy (left), top-5 accuracy (middle), and RAVEL disentanglement (right) for BeBeFOND on Gemma-2-2B layer-12 activations, evaluated over T=1,…,40T=1,…,40 inference steps at dictionary widths of 1616k, 6565k, 131131k, and 524524k. All evaluations use checkpoint 12,20712,207, corresponding to 49,999,87249,999,872 training tokens. Lines show means and shaded regions indicate ±1± 1 sample standard deviation across three training seeds. Sparse probing peaks within the first few steps, whereas RAVEL generally benefits from longer inference. D.1 Experimental setup Dataset. We extract layer-12 post-residual activations from Gemma-2-2B (Gemma Team, 2024) (d=2,304d=2,304). Documents from Pile Uncopyrighted (monology, 2023) are tokenized separately with BOS and truncated to 1,024 tokens. BOS and padding positions are excluded. Validation, calibration, and training use disjoint document ranges. Raw and whitened activations are evaluated separately. Validation and calibration each use 65,536 tokens, with up to 500M training tokens. Extraction uses BF16, batch 4, and no document packing. Whitening statistics were fitted on 500M training tokens. Training. Table 7 summarizes the training and evaluation settings. The selected BeBeFOND results in Table 6 use seeds 0–2 and checkpoint 12,207, corresponding to 49,999,872 decoder-training tokens. Evaluation uses ten outer inference steps with one inner step per reference update. The readout retains posterior means above 0.5 and sets the remainder to zero, followed by decoder-norm rescaling. At widths 16K, 65K, and 131K, the recipes use β=1β=1, ETD1 step h=1.5h=1.5, reference retention 0.9990.999, and decoder-norm targets 44, 1.51.5, and 0.750.75, respectively. At 524K they use β=1.5β=1.5, h=∞h=∞, full reference retention, and norm target 0.50.5. All use a global batch of 4,096 and a linear learning-rate schedule with peak 10−310^-3 and minimum 10−510^-5. Warmup occupies 5% of training at 16K–131K and 1% at 524K. Table 7: Gemma training and evaluation settings. Preprocessing, checkpoint, inference depth, and readout identify an evaluation configuration. Method / metric Settings BeBeFOND Widths 16,384–524,288; exponentially sampled training depth (mean 16, cap 100; Appendix C.1); FP32 inference/parameters; prior Fisher with current moments and no damping; Δt=1 t=1; decoder solver matrix_free_phi1_auto Mean-field + Adam Whitened inputs, 50M tokens, 20 inference steps, readout p>0.5p>0.5. After 2,000 updates, clip the dictionary gradient at three times its running-mean norm and skip batches whose negative ELBO exceeds five times its running mean, capped at approximately 20% of recent batches BatchTopK / Matryoshka Whitened inputs; batch 4,096; 12,207 updates (49,999,872 tokens); constant LR 3×10−43× 10^-4; FP32 SAEs, BF16 Gemma; SAELens 6.51.1. ZCA fitted on 500M training tokens. Target L0L_0 values in Table 8 were selected on a separate pilot Baseline initialization / losses Decoder norm initialized to 0.1; decoder bias subtracted before encoding; activations rescaled by decoder norm; auxiliary-loss coefficient 1; threshold EMA rate 0.01 Matryoshka levels 4,096, 16,384, and the full width (16,384 appears once for the smallest model); Matryoshka auxiliary loss enabled Gemma Scope Release google/gemma-scope-2b-pt-res, checkpoint paths layer_12/width_w/average_l0_s, with (w,s)(w,s) given by each row of Table 9. Training: 4B tokens at 16K, 8B at 32K–524K, and 16B at 1M (Lieberum et al., 2024) Sparse probing 113 tasks; final non-padding token; context 1,024; 1/2/5 selected features; L1 logistic regression; equal task averaging Probe selection Rank absolute normalized class-mean difference; 10 inverse-penalty values from 10510^5 to 10−510^-5; validation ROC-AUC; refit on training data RAVEL Average cause and isolation over attributes within each entity, then equally over entities; disentanglement is their mean. Aggregation requires completed results for every configured entity Evaluation. Sparse probing reports mean test accuracy across 113 tasks, using the best 1 or 5 features selected on training data. RAVEL reports cause, isolation, and their mean (Huang et al., 2024). Table 7 specifies the evaluation settings (Karvonen et al., 2025). Implementation. The Gemma reference environment uses Python 3.11.14, PyTorch 2.7.1+cu126, TransformerLens 2.16.1, Transformers 4.57.3, and SAE Lens 6.24.6. D.2 Results The whitened baseline comparison reports sparse probing and RAVEL in Table 8, and Table 9 reports the Gemma Scope checkpoints. Table 8: SAEBench results for SAEs trained on 50M tokens of whitened Gemma-2-2B layer-12 activations. Entries report mean ± sample standard deviation over five training seeds (1–5), evaluated at the final checkpoint (49,999,872 tokens). Sparse-probing metrics are macro-averaged over 113 tasks; RAVEL overall is the disentanglement score. All scores are fractions, and higher is better. L0L_0 denotes target training sparsity, fixed using a separate seed-0 pilot. Bold indicates the highest mean across L0L_0 values within each width–architecture group, selected before rounding. Sparse probing (Kantamneni et al., 2025) RAVEL (Huang et al., 2024) Width Target L0L_0 k=1k=1 k=5k=5 Acc. ↑ AUC ↑ F1 ↑ Acc. ↑ AUC ↑ F1 ↑ Overall ↑ Cause ↑ Isolation ↑ BatchTopK (16k) 16k 200 .765±.004.765 [rgb]0.3984,0.3984,0.3984±.004 .791±.004.791 [rgb]0.3984,0.3984,0.3984±.004 .758±.005.758 [rgb]0.3984,0.3984,0.3984±.005 .810±.002.810 [rgb]0.3984,0.3984,0.3984±.002 .859±.002.859 [rgb]0.3984,0.3984,0.3984±.002 .808±.003.808 [rgb]0.3984,0.3984,0.3984±.003 .717±.009.717 [rgb]0.3984,0.3984,0.3984±.009 .698±.018.698 [rgb]0.3984,0.3984,0.3984±.018 .736±.021.736 [rgb]0.3984,0.3984,0.3984±.021 16k 400 .767±.002.767 [rgb]0.3984,0.3984,0.3984±.002 .798±.003.798 [rgb]0.3984,0.3984,0.3984±.003 .761±.004.761 [rgb]0.3984,0.3984,0.3984±.004 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .862±.002.862 [rgb]0.3984,0.3984,0.3984±.002 .812±.002.812 [rgb]0.3984,0.3984,0.3984±.002 .729±.005.729 [rgb]0.3984,0.3984,0.3984±.005 .692±.019.692 [rgb]0.3984,0.3984,0.3984±.019 .767±.017.767 [rgb]0.3984,0.3984,0.3984±.017 Matryoshka BatchTopK (16k) 16k 200 .768±.003.768 [rgb]0.3984,0.3984,0.3984±.003 .794±.003.794 [rgb]0.3984,0.3984,0.3984±.003 .762±.002.762 [rgb]0.3984,0.3984,0.3984±.002 .810±.002.810 [rgb]0.3984,0.3984,0.3984±.002 .858±.002.858 [rgb]0.3984,0.3984,0.3984±.002 .808±.002.808 [rgb]0.3984,0.3984,0.3984±.002 .718±.004.718 [rgb]0.3984,0.3984,0.3984±.004 .688±.018.688 [rgb]0.3984,0.3984,0.3984±.018 .747±.013.747 [rgb]0.3984,0.3984,0.3984±.013 16k 400 .767±.001.767 [rgb]0.3984,0.3984,0.3984±.001 .797±.001.797 [rgb]0.3984,0.3984,0.3984±.001 .761±.002.761 [rgb]0.3984,0.3984,0.3984±.002 .809±.002.809 [rgb]0.3984,0.3984,0.3984±.002 .857±.001.857 [rgb]0.3984,0.3984,0.3984±.001 .808±.002.808 [rgb]0.3984,0.3984,0.3984±.002 .731±.003.731 [rgb]0.3984,0.3984,0.3984±.003 .695±.013.695 [rgb]0.3984,0.3984,0.3984±.013 .767±.015.767 [rgb]0.3984,0.3984,0.3984±.015 BatchTopK (65k) 65k 400 .770±.003.770 [rgb]0.3984,0.3984,0.3984±.003 .800±.002.800 [rgb]0.3984,0.3984,0.3984±.002 .765±.003.765 [rgb]0.3984,0.3984,0.3984±.003 .814±.003.814 [rgb]0.3984,0.3984,0.3984±.003 .862±.002.862 [rgb]0.3984,0.3984,0.3984±.002 .812±.003.812 [rgb]0.3984,0.3984,0.3984±.003 .740±.003.740 [rgb]0.3984,0.3984,0.3984±.003 .700±.015.700 [rgb]0.3984,0.3984,0.3984±.015 .781±.012.781 [rgb]0.3984,0.3984,0.3984±.012 65k 800 .768±.002.768 [rgb]0.3984,0.3984,0.3984±.002 .803±.002.803 [rgb]0.3984,0.3984,0.3984±.002 .764±.002.764 [rgb]0.3984,0.3984,0.3984±.002 .813±.001.813 [rgb]0.3984,0.3984,0.3984±.001 .860±.001.860 [rgb]0.3984,0.3984,0.3984±.001 .812±.001.812 [rgb]0.3984,0.3984,0.3984±.001 .745±.005.745 [rgb]0.3984,0.3984,0.3984±.005 .712±.006.712 [rgb]0.3984,0.3984,0.3984±.006 .777±.008.777 [rgb]0.3984,0.3984,0.3984±.008 Matryoshka BatchTopK (65k) 65k 200 .772±.002.772 [rgb]0.3984,0.3984,0.3984±.002 .798±.002.798 [rgb]0.3984,0.3984,0.3984±.002 .766±.002.766 [rgb]0.3984,0.3984,0.3984±.002 .813±.001.813 [rgb]0.3984,0.3984,0.3984±.001 .861±.002.861 [rgb]0.3984,0.3984,0.3984±.002 .811±.001.811 [rgb]0.3984,0.3984,0.3984±.001 .734±.009.734 [rgb]0.3984,0.3984,0.3984±.009 .707±.015.707 [rgb]0.3984,0.3984,0.3984±.015 .761±.010.761 [rgb]0.3984,0.3984,0.3984±.010 65k 400 .772±.001.772 [rgb]0.3984,0.3984,0.3984±.001 .801±.001.801 [rgb]0.3984,0.3984,0.3984±.001 .766±.002.766 [rgb]0.3984,0.3984,0.3984±.002 .811±.001.811 [rgb]0.3984,0.3984,0.3984±.001 .858±.001.858 [rgb]0.3984,0.3984,0.3984±.001 .809±.002.809 [rgb]0.3984,0.3984,0.3984±.002 .737±.004.737 [rgb]0.3984,0.3984,0.3984±.004 .704±.007.704 [rgb]0.3984,0.3984,0.3984±.007 .771±.006.771 [rgb]0.3984,0.3984,0.3984±.006 BatchTopK (131k) 131k 600 .773±.002.773 [rgb]0.3984,0.3984,0.3984±.002 .804±.002.804 [rgb]0.3984,0.3984,0.3984±.002 .768±.003.768 [rgb]0.3984,0.3984,0.3984±.003 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .862±.001.862 [rgb]0.3984,0.3984,0.3984±.001 .813±.002.813 [rgb]0.3984,0.3984,0.3984±.002 .751±.003.751 [rgb]0.3984,0.3984,0.3984±.003 .720±.008.720 [rgb]0.3984,0.3984,0.3984±.008 .782±.013.782 [rgb]0.3984,0.3984,0.3984±.013 131k 800 .770±.003.770 [rgb]0.3984,0.3984,0.3984±.003 .803±.003.803 [rgb]0.3984,0.3984,0.3984±.003 .765±.003.765 [rgb]0.3984,0.3984,0.3984±.003 .814±.002.814 [rgb]0.3984,0.3984,0.3984±.002 .859±.002.859 [rgb]0.3984,0.3984,0.3984±.002 .812±.002.812 [rgb]0.3984,0.3984,0.3984±.002 .746±.004.746 [rgb]0.3984,0.3984,0.3984±.004 .720±.011.720 [rgb]0.3984,0.3984,0.3984±.011 .771±.005.771 [rgb]0.3984,0.3984,0.3984±.005 Matryoshka BatchTopK (131k) 131k 200 .773±.004.773 [rgb]0.3984,0.3984,0.3984±.004 .799±.004.799 [rgb]0.3984,0.3984,0.3984±.004 .768±.004.768 [rgb]0.3984,0.3984,0.3984±.004 .814±.001.814 [rgb]0.3984,0.3984,0.3984±.001 .860±.002.860 [rgb]0.3984,0.3984,0.3984±.002 .812±.001.812 [rgb]0.3984,0.3984,0.3984±.001 .739±.008.739 [rgb]0.3984,0.3984,0.3984±.008 .723±.019.723 [rgb]0.3984,0.3984,0.3984±.019 .754±.006.754 [rgb]0.3984,0.3984,0.3984±.006 131k 400 .773±.001.773 [rgb]0.3984,0.3984,0.3984±.001 .801±.001.801 [rgb]0.3984,0.3984,0.3984±.001 .767±.002.767 [rgb]0.3984,0.3984,0.3984±.002 .812±.001.812 [rgb]0.3984,0.3984,0.3984±.001 .858±.001.858 [rgb]0.3984,0.3984,0.3984±.001 .810±.002.810 [rgb]0.3984,0.3984,0.3984±.002 .739±.002.739 [rgb]0.3984,0.3984,0.3984±.002 .713±.008.713 [rgb]0.3984,0.3984,0.3984±.008 .764±.006.764 [rgb]0.3984,0.3984,0.3984±.006 BatchTopK (524k) 524k 600 .772±.004.772 [rgb]0.3984,0.3984,0.3984±.004 .801±.004.801 [rgb]0.3984,0.3984,0.3984±.004 .766±.004.766 [rgb]0.3984,0.3984,0.3984±.004 .813±.002.813 [rgb]0.3984,0.3984,0.3984±.002 .861±.002.861 [rgb]0.3984,0.3984,0.3984±.002 .812±.003.812 [rgb]0.3984,0.3984,0.3984±.003 .747±.006.747 [rgb]0.3984,0.3984,0.3984±.006 .719±.011.719 [rgb]0.3984,0.3984,0.3984±.011 .775±.007.775 [rgb]0.3984,0.3984,0.3984±.007 524k 800 .773±.004.773 [rgb]0.3984,0.3984,0.3984±.004 .803±.004.803 [rgb]0.3984,0.3984,0.3984±.004 .768±.006.768 [rgb]0.3984,0.3984,0.3984±.006 .814±.003.814 [rgb]0.3984,0.3984,0.3984±.003 .860±.002.860 [rgb]0.3984,0.3984,0.3984±.002 .813±.003.813 [rgb]0.3984,0.3984,0.3984±.003 .748±.004.748 [rgb]0.3984,0.3984,0.3984±.004 .719±.010.719 [rgb]0.3984,0.3984,0.3984±.010 .776±.010.776 [rgb]0.3984,0.3984,0.3984±.010 Matryoshka BatchTopK (524k) 524k 200 .774±.002.774 [rgb]0.3984,0.3984,0.3984±.002 .799±.002.799 [rgb]0.3984,0.3984,0.3984±.002 .768±.002.768 [rgb]0.3984,0.3984,0.3984±.002 .813±.001.813 [rgb]0.3984,0.3984,0.3984±.001 .859±.001.859 [rgb]0.3984,0.3984,0.3984±.001 .812±.000.812 [rgb]0.3984,0.3984,0.3984±.000 .734±.005.734 [rgb]0.3984,0.3984,0.3984±.005 .705±.011.705 [rgb]0.3984,0.3984,0.3984±.011 .764±.011.764 [rgb]0.3984,0.3984,0.3984±.011 524k 400 .771±.002.771 [rgb]0.3984,0.3984,0.3984±.002 .800±.002.800 [rgb]0.3984,0.3984,0.3984±.002 .767±.003.767 [rgb]0.3984,0.3984,0.3984±.003 .810±.001.810 [rgb]0.3984,0.3984,0.3984±.001 .858±.001.858 [rgb]0.3984,0.3984,0.3984±.001 .809±.002.809 [rgb]0.3984,0.3984,0.3984±.002 .743±.003.743 [rgb]0.3984,0.3984,0.3984±.003 .708±.013.708 [rgb]0.3984,0.3984,0.3984±.013 .778±.009.778 [rgb]0.3984,0.3984,0.3984±.009 Table 9: SAEBench evaluations of released Gemma Scope layer-12 checkpoints (Lieberum et al., 2024; Karvonen et al., 2025), organized by dictionary width and average sparsity L0L_0. Sparse probing RAVEL Width L0L_0 k=1k=1 k=2k=2 k=5k=5 Disent- angle ↑ Cause ↑ Isola- tion ↑ Acc. ↑ AUC ↑ F1 ↑ Acc. ↑ AUC ↑ F1 ↑ Acc. ↑ AUC ↑ F1 ↑ 16k 22 0.7337 0.7476 0.7159 0.7529 0.7861 0.7441 0.7771 0.8243 0.7717 0.5989 0.5121 0.6857 16k 41 0.7519 0.7710 0.7366 0.7686 0.8056 0.7601 0.7893 0.8370 0.7835 0.6644 0.6163 0.7126 16k 82 0.7627 0.7857 0.7545 0.7821 0.8207 0.7768 0.8047 0.8523 0.8016 0.6960 0.6489 0.7431 16k 176 0.7679 0.7938 0.7602 0.7890 0.8262 0.7851 0.8129 0.8587 0.8119 0.7206 0.6508 0.7904 16k 445 0.7725 0.8033 0.7687 0.7942 0.8363 0.7921 0.8171 0.8632 0.8162 0.7488 0.7205 0.7771 32k 12 0.7179 0.7264 0.6960 0.7350 0.7665 0.7181 0.7570 0.8046 0.7489 0.6050 0.5461 0.6639 32k 22 0.7379 0.7498 0.7221 0.7535 0.7824 0.7418 0.7758 0.8187 0.7711 0.6416 0.6251 0.6582 32k 40 0.7520 0.7695 0.7422 0.7662 0.7995 0.7590 0.7872 0.8327 0.7842 0.6755 0.6200 0.7311 32k 76 0.7607 0.7815 0.7534 0.7761 0.8138 0.7721 0.7990 0.8466 0.7958 0.7075 0.6714 0.7436 32k 155 0.7675 0.7907 0.7609 0.7865 0.8255 0.7833 0.8054 0.8528 0.8025 0.7249 0.6881 0.7616 32k 360 0.7749 0.8049 0.7713 0.7897 0.8306 0.7866 0.8140 0.8601 0.8129 0.7450 0.7105 0.7794 65k 21 0.7361 0.7459 0.7139 0.7507 0.7785 0.7359 0.7743 0.8186 0.7686 0.6469 0.6053 0.6885 65k 38 0.7501 0.7656 0.7399 0.7660 0.7962 0.7589 0.7874 0.8290 0.7834 0.6767 0.6345 0.7188 65k 72 0.7647 0.7816 0.7552 0.7810 0.8128 0.7751 0.7987 0.8451 0.7965 0.7041 0.6596 0.7487 65k 141 0.7717 0.7930 0.7653 0.7858 0.8231 0.7814 0.8087 0.8541 0.8069 0.7346 0.7158 0.7535 65k 297 0.7788 0.8032 0.7731 0.7947 0.8343 0.7925 0.8178 0.8611 0.8169 0.7487 0.7016 0.7959 131k 12 0.7108 0.7198 0.6870 0.7311 0.7585 0.7146 0.7525 0.7952 0.7443 0.6004 0.5764 0.6245 131k 20 0.7279 0.7362 0.7074 0.7483 0.7752 0.7335 0.7698 0.8132 0.7621 0.6408 0.5950 0.6866 131k 36 0.7379 0.7513 0.7222 0.7565 0.7843 0.7471 0.7810 0.8226 0.7767 0.6828 0.6451 0.7205 131k 67 0.7554 0.7713 0.7444 0.7693 0.8019 0.7622 0.7913 0.8347 0.7876 0.7083 0.6734 0.7432 131k 129 0.7677 0.7865 0.7608 0.7837 0.8173 0.7794 0.8027 0.8475 0.8011 0.7265 0.6561 0.7969 131k 264 0.7814 0.8033 0.7773 0.7941 0.8290 0.7914 0.8138 0.8575 0.8125 0.7486 0.7281 0.7691 262k 11 0.7068 0.7168 0.6786 0.7232 0.7503 0.7054 0.7430 0.7872 0.7310 0.6166 0.5987 0.6345 262k 21 0.7240 0.7355 0.7037 0.7443 0.7732 0.7318 0.7657 0.8078 0.7607 0.6355 0.6169 0.6540 262k 36 0.7375 0.7533 0.7252 0.7573 0.7879 0.7479 0.7772 0.8195 0.7723 0.6795 0.6567 0.7023 262k 67 0.7510 0.7684 0.7427 0.7666 0.7960 0.7596 0.7911 0.8325 0.7883 0.7039 0.6652 0.7426 262k 121 0.7656 0.7854 0.7572 0.7818 0.8155 0.7772 0.8029 0.8480 0.8008 0.7158 0.6737 0.7579 262k 243 0.7733 0.7951 0.7655 0.7928 0.8291 0.7906 0.8119 0.8561 0.8105 0.7523 0.7219 0.7827 524k 22 0.7129 0.7276 0.6976 0.7218 0.7529 0.7104 0.7453 0.7893 0.7384 0.6138 0.5800 0.6475 524k 29 0.7258 0.7421 0.7139 0.7400 0.7725 0.7297 0.7629 0.8076 0.7573 0.6451 0.6002 0.6900 524k 46 0.7387 0.7552 0.7264 0.7535 0.7872 0.7457 0.7770 0.8211 0.7731 0.6803 0.6602 0.7005 524k 65 0.7523 0.7688 0.7412 0.7663 0.7981 0.7589 0.7879 0.8303 0.7856 0.6958 0.6583 0.7332 524k 115 0.7637 0.7815 0.7550 0.7791 0.8110 0.7734 0.8004 0.8419 0.7981 0.7251 0.6760 0.7741 524k 227 0.7735 0.7955 0.7682 0.7927 0.8268 0.7890 0.8106 0.8520 0.8094 0.7555 0.7330 0.7781 1m 19 0.7028 0.7182 0.6875 0.7165 0.7497 0.7081 0.7396 0.7847 0.7350 0.5910 0.5570 0.6250 1m 26 0.7234 0.7388 0.7047 0.7357 0.7693 0.7273 0.7587 0.8084 0.7523 0.6371 0.6120 0.6622 1m 58 0.7354 0.7548 0.7239 0.7550 0.7876 0.7499 0.7755 0.8193 0.7737 0.6722 0.6560 0.6884 1m 73 0.7526 0.7714 0.7435 0.7671 0.8013 0.7633 0.7887 0.8295 0.7866 0.6933 0.6533 0.7334 1m 107 0.7633 0.7797 0.7545 0.7737 0.8030 0.7684 0.7959 0.8385 0.7937 0.7201 0.6857 0.7545 1m 207 0.7788 0.7983 0.7726 0.7919 0.8264 0.7881 0.8095 0.8514 0.8081 0.7420 0.6860 0.7979 D.3 Wall-clock training times on H200 GPUs Table 10 compares wall-clock training times on eight NVIDIA H200 GPUs, with one model measured at a time and seed 0 throughout. Each measurement covers the first 999,424 main-training tokens of the original 50M-token schedule. Global batch size is 4,096 for BeBeFOND, BatchTopK, and Matryoshka, and 2,048 for MF20+clip. Table 10: Elapsed training time in seconds for 999,424 training tokens on eight NVIDIA H200 GPUs (lower is better). Each model is measured once at seed 0. Timings include BeBeFOND’s prior updates and exclude evaluation, initialization, warmup, data preparation, checkpointing, and logging. Matryoshka denotes Matryoshka BatchTopK. Mean-field + Adam uses 20 inference steps, gradient clipping, and batch skipping. Width K BeBeFOND BatchTopK Matryoshka Mean-field + Adam 16,384 33.125 2.719 2.991 2.714 65,536 118.959 9.202 10.737 10.157 131,072 234.665 17.935 20.075 19.914 524,288 989.321 70.121 77.382 78.331 Inputs are loaded, shuffled, whitened, and placed in GPU memory before timing. Initialization and CUDA/JIT warmup are excluded, and the complete initial training state is restored after warmup. The timed interval includes inference, losses, parameter updates, normalization, required communication, and BeBeFOND’s prior updates with convergence checks. Evaluation, checkpoint writing, and logging are excluded. PyTorch timings use the maximum CUDA-synchronized elapsed time across ranks, while JAX waits for all output shards on all devices. BeBeFOND retains its original stochastic training horizons, capped at 100 with one inner step, and warms the complete measured prefix to cover adaptive solver paths. MF uses 20 inference steps. Its clipping and loss-based skipping begin after 2,000 updates and are inactive during this prefix. Model tensors are FP32, with native FP64 intermediates retained in BeBeFOND’s moment reductions and Woodbury calculations. BeBeFOND and the amortized baselines use PyTorch, while MF uses JAX. The distributed baseline implementations preserve the global TopK budget and synchronize gradients, thresholds, and dead-feature histories. BeBeFOND takes approximately 1212–14×14× as long as BatchTopK across these widths.