Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SplInterp: Improving our Understanding and Training of Sparse Autoencoders

Jeremy Budd, Javier Ideami, Benjamin Macdowall Rynne, Keith Duggar, Randall Balestriero

TL;DR

This work tackles the interpretability and data-efficiency of sparse autoencoders (SAEs) in large models by embedding SAEs in spline theory. It develops a unified SAE-spline framework, revealing the exact geometric structure of TopK (and related) SAEs as $K$-th order power diagrams, and introduces a proximal alternating training method (PAM-SGD) with convergence guarantees. The results show that SAEs generalize $k$-means with piecewise affine encodings while trading some accuracy for monosemantic sparsity, and PAM-SGD enhances sample efficiency and activation sparsity in MNIST and LLM settings. This spline-based perspective provides both theoretical grounding and practical tools for mechanistic interpretability in high-dimensional sparse representations.

Abstract

Sparse autoencoders (SAEs) have received considerable recent attention as tools for mechanistic interpretability, showing success at extracting interpretable features even from very large LLMs. However, this research has been largely empirical, and there have been recent doubts about the true utility of SAEs. In this work, we seek to enhance the theoretical understanding of SAEs, using the spline theory of deep learning. By situating SAEs in this framework: we discover that SAEs generalise ``$k$-means autoencoders'' to be piecewise affine, but sacrifice accuracy for interpretability vs. the optimal ``$k$-means-esque plus local principal component analysis (PCA)'' piecewise affine autoencoder. We characterise the underlying geometry of (TopK) SAEs using power diagrams. And we develop a novel proximal alternating method SGD (PAM-SGD) algorithm for training SAEs, with both solid theoretical foundations and promising empirical results in MNIST and LLM experiments, particularly in sample efficiency and (in the LLM setting) improved sparsity of codes. All code is available at: https://github.com/splInterp2025/splInterp

SplInterp: Improving our Understanding and Training of Sparse Autoencoders

TL;DR

This work tackles the interpretability and data-efficiency of sparse autoencoders (SAEs) in large models by embedding SAEs in spline theory. It develops a unified SAE-spline framework, revealing the exact geometric structure of TopK (and related) SAEs as -th order power diagrams, and introduces a proximal alternating training method (PAM-SGD) with convergence guarantees. The results show that SAEs generalize -means with piecewise affine encodings while trading some accuracy for monosemantic sparsity, and PAM-SGD enhances sample efficiency and activation sparsity in MNIST and LLM settings. This spline-based perspective provides both theoretical grounding and practical tools for mechanistic interpretability in high-dimensional sparse representations.

Abstract

Sparse autoencoders (SAEs) have received considerable recent attention as tools for mechanistic interpretability, showing success at extracting interpretable features even from very large LLMs. However, this research has been largely empirical, and there have been recent doubts about the true utility of SAEs. In this work, we seek to enhance the theoretical understanding of SAEs, using the spline theory of deep learning. By situating SAEs in this framework: we discover that SAEs generalise ``-means autoencoders'' to be piecewise affine, but sacrifice accuracy for interpretability vs. the optimal ``-means-esque plus local principal component analysis (PCA)'' piecewise affine autoencoder. We characterise the underlying geometry of (TopK) SAEs using power diagrams. And we develop a novel proximal alternating method SGD (PAM-SGD) algorithm for training SAEs, with both solid theoretical foundations and promising empirical results in MNIST and LLM experiments, particularly in sample efficiency and (in the LLM setting) improved sparsity of codes. All code is available at: https://github.com/splInterp2025/splInterp
Paper Structure (49 sections, 8 theorems, 57 equations, 38 figures, 1 algorithm)

This paper contains 49 sections, 8 theorems, 57 equations, 38 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. The spline geometry of sparse autoencoders (SAEs)
  3. A primer on SAEs
  4. SAEs are piecewise affine splines
  5. SAEs, $k$-means, and principal component analysis (PCA)
  6. Visualising the SAE bridge between $k$-means and PCA
  7. A proximal alternating method (PAM-SGD) for training SAEs
  8. Convergence theory and the PAM-SGD algorithm
  9. Sample-efficient sparse coding for MNIST and LLMs: PAM-SGD vs. SGD
  10. PAM-SGD generalised better than SGD in low-data regimes. (\ref{['fig:mnist02', 'fig:gemma_training_size', 'fig:mnistTopK', 'fig:ReLU_reruns', 'fig:gemma_K_values', 'fig:gemma_training_size_Top320', 'fig:gemma_loss_curves_Top640', 'fig:gemma_loss_curves_Top320', 'fig:gemma_training_size_Top640']})
  11. PAM-SGD was faster, more accurate, and more interpretable on MNIST. (\ref{['fig:mnist_reconstruction', 'fig:mnist_evolution', 'fig:mnist_filters', 'fig:mnist_reconstruction_TopK', 'fig:mnist_evolution_TopK', 'fig:mnist_filters_TopK']})
  12. Summary and practical implications.
  13. Conclusions and limitations
  14. Visualisations of Voronoi and power diagrams
  15. Spatial Partitioning Methods: From Voronoi to Power Diagrams
  16. ...and 34 more sections

Key Result

Theorem 2.3

The cells $\{\Omega^{\mathop{\mathrm{TopK}}\nolimits}_S\}$ form a $K^\text{th}$-order power diagram with $\binom{d}{K}$ cells. Conversely, for any $K^\text{th}$-order power diagram with $\binom{d}{K}$ cells with centroids $\{\mu_i\}_{i=1}^d$ and weights $\{\alpha_i\}_{i=1}^d$, there exist $W_{enc} \ for all $i$, where $e_i$ is the elementary basis vector with $1$ in coordinate $i$ and $0$ in every

Figures (38)

  • Figure 1: Key ideas in this work.Left: Sparse autoencoders (SAEs) vs. regular autoencoders; in an SAE, the code $z$ can have very high dimension, but most entries are zero (greyed out). Right: The superior sample efficiency (on MNIST and Gemma-2-2B) of our novel proximal alternating method SGD (PAM-SGD) algorithm.
  • Figure 2: Visualising the SAE bridge between $k$-means clustering and PCA. The cyan/brown boundary intersecting the point cloud (left and right figure) is a visualiser bug.
  • Figure 3: Training and test loss curves at different data sizes for MNIST, with ReLU activation. The chart highlights PAM-SGD's superior sample efficiency and convergence speed.
  • Figure 4: Training and test loss curves at different data sizes for Gemma-2-2B, with ReLU activation. PAM-SGD again has a huge advantage at low data, and remains superior throughout.
  • Figure 5: Spatial Partitioning Methods (a) Standard Voronoi diagram where space is partitioned based on the nearest generator point using Euclidean distance; (b) Nearest-neighbor ($k=1$) classification showing how Voronoi cells define decision boundaries for point classification (c) $k$-means clustering with $k=3$ demonstrating how cluster centroids generate Voronoi cells that define cluster boundaries (d) Power diagram (weighted Voronoi) where each generator has an associated weight, creating curved boundaries between regions. Power diagrams generalise Voronoi diagrams and provide additional flexibility for modeling spatial relationships.
  • ...and 33 more figures

Theorems & Definitions (20)

  • Definition 2.1: Power and Voronoi diagrams
  • Theorem 2.3
  • Definition 2.4: $k$-means clustering
  • Theorem 2.6
  • Theorem 3.1
  • Theorem B.1
  • proof : Proof of \ref{['thm:openpoly']}
  • Theorem B.2
  • proof : Proof of \ref{['thm:Kthpowerdiagram']}
  • proof : Proof of \ref{['nb:hexagon']}
  • ...and 10 more