Inference and learning in sparse autoencoders as natural gradient flow

Vafaii, Hadi, Rao, Tejas, Chanin, David, Fel, Thomas, Yates, Jacob, Olshausen, Bruno, Klindt, David, George, Dileep, Lazaro-Gredilla, Miguel (October 2026) Inference and learning in sparse autoencoders as natural gradient flow. arXiv. ISSN 2331-8422 (Submitted)

PDF 10.48550.arXiv.2610.07389.pdf - Submitted Version Available under License Creative Commons Attribution. Download (2MB)

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.

Item Type:	Paper
Subjects:	bioinformatics bioinformatics > computational biology > algorithms bioinformatics > computational biology bioinformatics > computational biology > algorithms > machine learning
CSHL Authors:	Klindt, David
Communities:	CSHL labs > Klindt lab
SWORD Depositor:	CSHL Elements
Depositing User:	CSHL Elements
Date:	5 October 2026
Date Deposited:	08 Oct 2026 12:30
Last Modified:	08 Oct 2026 12:30
URI:	https://repository.cshl.edu/id/eprint/42335

Actions (login required)

Administrator's edit/view item