Research map / Large language models / Language, safety and interpretability
Jacobian Sparse Autoencoders: Sparsify Computations, Not Just Activations
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Jacobian SAEs, which train an input/output SAE pair around an MLP with an L1 penalty on the Jacobian of the latent-to-latent map so the computation, not just activations, is sparse.
Problem
Standard sparse autoencoders sparsify LLM activations but not the computation connecting them, so the computational graph between SAE latents across an MLP remains dense and hard to interpret.
Approach
Train a pair of k-sparse SAEs on an MLP's input and output and add an L1 penalty on the Jacobian of the function mapping input SAE latents to output SAE latents; restricting to the k active latents and deriving a closed-form Jacobian makes this tractable at roughly 2x SAE training cost.
Key results
- At lambda=0.5 on Pythia-410m (k=32), JSAEs substantially reduce the proportion of large Jacobian elements versus standard SAEs while keeping cross-entropy close to traditional SAEs.
- Jacobians are substantially sparser for trained Pythia models than for randomly re-initialized ones, and the JSAE sparsity gain is larger for trained models.
- About 88% of input-output scalar functions between SAE latents are linear and 10% are JumpReLU-like; 97.7% of samples satisfy |Delta s_y| approx |J| under a unit intervention.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add an L1 Jacobian sparsity term to the SAE objective | Traditional SAE pair · Pythia-410m MLPs | Jacobian sparsity | Helped |
| Add an L1 Jacobian sparsity term while retaining reasonable reconstruction quality | Traditional SAE pair · Pythia-410m MLPs | Cross-entropy preservation | Mixed |
| Add an L1 Jacobian sparsity term | Traditional SAE pair · Pythia-410m odd-numbered layers | Auto-interpretability | No clear effect |
| Increase the Jacobian loss coefficient above the useful range (ablation) | Jacobian SAE · Pythia MLPs | Reconstruction quality and model performance preservation | Hurt |
| Use λ=0.5 rather than a larger Jacobian loss coefficient (ablation) | Jacobian SAE · Pythia-410m layer 15 | Cross-entropy versus Jacobian sparsity | Helped |
| Train the Jacobian SAE on a pre-trained transformer rather than a randomly reinitialized transformer (ablation) | Jacobian SAE · Pythia-410m layer 15 | Jacobian sparsity | Helped |
| Train a traditional SAE on a pre-trained transformer rather than a randomly reinitialized transformer (ablation) | Traditional SAE · Pythia-410m layer 15 | Jacobian sparsity | Helped |
| Use Jacobian SAEs instead of traditional SAEs (ablation) | SAE input-output latent mappings · Pythia MLPs | Proportion of scalar functions that are linear | Helped |
| Use Jacobian SAE rather than traditional SAE (ablation) | Pythia-410m layer 15 · Jacobian sparsity | Jacobian sparsity | Helped |
Surprising findings
- Degradation of reconstruction and CE metrics under strong Jacobian penalty occurs almost entirely in the output SAE rather than the input SAE.
- Traditional SAEs already show somewhat sparser Jacobians on trained than random transformers, but the JSAE improvement is much larger for trained models.
Limitations
- Only applies to a single GPT-2-style MLP layer at a time; not GLU MLPs in the main experiments and no whole-model circuits.
- Jacobian is a local measure and an approximate proxy for computational sparsity; about 2% of scalar functions are neither linear nor JumpReLU.
- Relies on TopK activation, which can encourage high-density features.
Experimental setup
Datasets: C4 English subset
Benchmarks: Pythia model cross-entropy loss preservation fuzzing auto-interpretability score
Metrics: Jacobian sparsity (fraction of elements above 0.01) variance explained CE loss increase auto-interpretability fuzzing score
Baselines: standard TopK SAEs (lambda=0) randomly re-initialized Pythia models
Builds on: TopK SAE sparse autoencoders ACDC fuzzing auto-interpretability SAELens
Closest papers by meaning
- Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders (NeurIPS 2024 · similarity 0.92)
- Scaling and evaluating sparse autoencoders (ICLR 2025 · similarity 0.90)
- Are Sparse Autoencoders Useful? A Case Study in Sparse Probing (ICML 2025 · similarity 0.90)
- SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability (ICML 2025 · similarity 0.90)
- Low-Rank Adapting Models for Sparse Autoencoders (ICML 2025 · similarity 0.89)
- Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders (ICML 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.