JumpReLU SAEs optimize L0 directly and match or beat Gated and TopK

measured in 1 paper

Rajamanoharan et al. train JumpReLU SAEs with a per-feature learnable threshold on Gemma-2-9B residual, attention-output, and MLP-output sites at layers 9/20/31 [rajamanoharan-etal-2024-jumprelu] Because the step gate and L0 penalty are piecewise-constant in the threshold, they train via a kernel-density straight-through estimator proven to estimate the true expected-loss gradient [rajamanoharan-etal-2024-jumprelu] On delta-LM-loss versus L0, JumpReLU SAEs consistently beat Gated and match or slightly beat TopK at matched sparsity [rajamanoharan-etal-2024-jumprelu] Manual and automated interpretability ratings are similar across the three, so the fidelity gain costs no interpretability [rajamanoharan-etal-2024-jumprelu] The gain is attributable to directly optimizing L0 rather than an L1 proxy, since a weight-tied Gated encoder is already equivalent in form to JumpReLU [rajamanoharan-etal-2024-jumprelu]

Structure

Linear Direction

Context

straight-through estimator via kernel-density-estimate surrogate gradient for a discontinuous threshold gate, direct L0 optimization instead of an L1 sparsity proxy, Pareto-equal-or-better reconstruction fidelity vs Gated and TopK at matched sparsity, with matched interpretability, fidelity gain attributed to training objective, not functional form (Gated encoder is already equivalent to JumpReLU in form)

Confirmed in models

Gemma-2-9B · Gemma

Method

JumpReLU Sparse Autoencoders

Papers

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders — Rajamanoharan, Senthooran, Lieberum, Tom, Sonnerat, Nicolas, Conmy, Arthur, Varma, Vikrant, Kramár, János, Nanda, Neel2024 · arXiv:2407.14435