Rational Sparse Autoencoder

2026-06-12

AI Summary

This work addresses the rigidity of sparsity mechanisms in existing sparse autoencoders, which stems from their reliance on fixed nonlinear activation functions and hinders the trade-off between reconstruction fidelity and sparsity. To overcome this limitation, the authors propose a trainable rational function activation, optimized via a two-stage procedure: first, scale parameters are initialized and calibrated on synthetic data using a relaxed Remez exchange algorithm; second, the model is fine-tuned with a standard sparsity-regularized objective. This approach unifies the approximation of diverse sparse activation mechanisms, adapts to the geometry of pre-activations, and introduces only a few scalar parameters. Experiments on residual streams of three open-source language models demonstrate that the proposed Rational Sparse Autoencoder (RSAE) consistently outperforms existing activation baselines in both reconstruction accuracy and downstream task performance, while preserving feature-level interpretability and enabling training within minutes on a single GPU.

Abstract

Sparse autoencoders (SAEs) are standard tools for mechanistic interpretability, but current SAE families are constrained by fixed encoder nonlinearities such as ReLU, JumpReLU, and TopK. This hard-codes a particular sparsity mechanism into the model and can distort the reconstruction-versus-sparsity trade-off. We introduce the Rational Sparse Autoencoder (RSAE), which replaces the fixed encoder activation with a trainable rational function. Rational activations are flexible enough to uniformly approximate the activation primitives used by existing SAE families on compact domains (for TopK, the thresholded gate obtained after a separating top-k threshold is supplied), while also providing a richer function class for adapting to the observed pre-activation geometry. We realise this idea through a two-stage pipeline: an initialisation procedure that copies the pre-trained baseline SAE weights, plugs in rational coefficients obtained by the relaxed Remez exchange on synthetic data, and calibrates the scale parameters along with the rational coefficients; followed by a fine-tuning step under the standard sparsity-regularised reconstruction objective. Empirically, on residual-stream activations of three open-weight language models and across all three baseline activation families, the RSAE strictly improves on it after the fine-tuning step, both on reconstruction-side metrics and on downstream-behaviour metrics, without sacrificing feature-level interpretability under sparse probing. These gains are consistent across host language models, across baseline activation families, and across the full range of baseline sparsity we tested, while the upgrade itself adds only a handful of scalar parameters per autoencoder and runs in minutes on a single consumer GPU.

Problem

• Sparse Autoencoder
• Nonlinear Activation
• Sparsity Mechanism
• Reconstruction-Sparsity Trade-off
• Interpretability

Innovation

• Rational Sparse Autoencoder
• trainable rational activation
• mechanistic interpretability
• sparsity-reconstruction trade-off
• Remez exchange algorithm