Technical Summary: Rational Sparse Autoencoder (RSAE)

Problem Statement

Sparse Autoencoders (SAEs) are the standard tool for mechanistic interpretability in large language models (LLMs), decomposing internal activations into sparse linear combinations of overcomplete feature dictionaries. However, existing SAE families rely on fixed, non-smooth encoder nonlinearities—specifically ReLU (with $\ell_1$ regularization), TopK, and JumpReLU. These fixed activation primitives introduce specific pathologies:

• ReLU SAEs: Suffer from magnitude shrinkage of active features and a persistent population of "dead" latents.
• TopK SAEs: Replace soft penalties with hard cardinality constraints, breaking gradient flow through inactive features and requiring auxiliary revival losses.
• JumpReLU SAEs: Introduce learnable thresholds but require continuous-relaxation surrogates for back-propagation through their indicator gates.

These hard-coded mechanisms constrain the model's ability to adapt to the observed pre-activation geometry, potentially distorting the trade-off between reconstruction fidelity and sparsity.

Methodology

The authors propose the Rational Sparse Autoencoder (RSAE), which replaces the fixed encoder activation with a trainable rational function. The RSAE retains the standard SAE skeleton (an affine pre-activation layer followed by a sparse activation block and a linear decoder) but modifies the activation function $\phi(\cdot)$ to be a learnable rational function applied element-wise.

Architecture

The RSAE activation is defined as:
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
where $h$ is the pre-activation vector, $C_{in}$ and $C_{out}$ are learnable scaling parameters, and $r_{(a,b)}(t) = P(t)/Q(t)$ is a rational function with trainable coefficients $(a, b)$. The scaling parameters map the pre-activation distribution into the rational function's design interval $[-1, 1]$ and map the output back to the expected feature magnitude.

Theoretical Foundation

The paper establishes an expressive asymmetry between rational functions and the piecewise-affine gates used in current SAEs:

1. Upper Bound (Rational Approximation): Using Zolotarev sign functions, the authors prove that trainable rational activations can approximate ReLU, JumpReLU, and supplied-threshold TopK gates with polylogarithmic size in the target accuracy $\epsilon$ (specifically $O(\log(1/\epsilon)\log(1/\delta))$ for discontinuous gates with margin $\delta$).
2. Lower Bound (Piecewise-Affine Limit): Conversely, they demonstrate that approximating certain $O(1)$-parameter rational target maps using scalar-output single-layer piecewise-affine encoders (ReLU/JumpReLU/TopK) requires $\Omega(\epsilon^{-1/2})$ activated coordinates. This suggests that rational activations offer a more efficient representation class for the SAE encoder.

Training Pipeline

The RSAE is constructed via a two-stage pipeline designed to upgrade pre-trained "teacher" SAEs:

1. Initialization:
   • Remez Exchange: A rational function is fitted to the teacher's activation primitive (ReLU, JumpReLU, or TopK) on a synthetic dense grid using the relaxed Remez exchange algorithm to minimize the $L_\infty$ error.
   • Calibration: The rational coefficients and scaling parameters ($C_{in}, C_{out}$) are adapted to the teacher's actual pre-activation distribution by minimizing the squared error between the rational output and the teacher's output.
   • Weight Inheritance: The encoder and decoder weights are copied verbatim from the pre-trained teacher.
2. Fine-Tuning: All parameters (weights, biases, and rational coefficients) are unfrozen and jointly optimized under the standard $\ell_1$-regularized reconstruction objective.

Key Contributions

1. RSAE Architecture: Introduction of a sparse autoencoder with a trainable rational activation, grounded in approximation theory tailored to the shallow SAE encoder.
2. Theoretical Expressivity Gap: Proof that rational activations provide compact approximations to fixed SAE gates, whereas piecewise-affine encoders require significantly more activated coordinates to approximate simple rational targets.
3. Two-Step Training Algorithm: A practical initialization and fine-tuning procedure that allows RSAEs to be "dropped in" to replace pre-trained SAEs with negligible overhead.
4. Empirical Validation: Comprehensive evaluation across three open-weight language models (GPT-2 small, Pythia-160m, Gemma-2-2B) and all three baseline activation families.

Results

The authors evaluate RSAE against pre-trained baselines on residual-stream activations. The results indicate strict improvements after fine-tuning:

• Reconstruction Metrics: RSAE consistently achieves lower Mean Squared Error (MSE), lower $\ell_0$ (sparsity), and a higher fraction of "alive" features compared to the teacher models across all tested configurations.
• Downstream Metrics: RSAE improves downstream behavior, showing reduced cross-entropy degradation ($\Delta CE$) and higher loss recovery ($LR$) when intercepting the residual stream.
• Interpretability: Sparse probing experiments on Pythia-160m show that RSAE preserves feature-level interpretability. The full-dictionary probe accuracy on eight binary classification tasks remains comparable to the teacher, with no significant degradation.
• Efficiency: The upgrade adds only a handful of scalar parameters per autoencoder. The entire pipeline (initialization + 2,000 fine-tuning steps) runs in minutes on a single consumer GPU (NVIDIA RTX 5090).
• Robustness: Gains are consistent across different host models, baseline activation families, and a wide range of baseline sparsity levels ($\ell_0 \in [51, 683]$).

Significance and Claims

The paper claims that the RSAE offers a "drop-in" modification for pre-trained SAEs that strictly improves model fidelity at matched sparsity without sacrificing interpretability. By replacing fixed nonlinearities with trainable rational functions, the RSAE adapts to the specific geometry of the pre-activations in the host model, overcoming the limitations of hard-coded activation primitives.

The authors emphasize that these improvements are achieved with negligible computational cost (minutes on a single GPU) and minimal parameter overhead. They position RSAE as a universal upgrade strategy for existing SAEs, capable of enhancing both reconstruction quality and downstream task performance while maintaining the monosemantic features required for mechanistic interpretability. The paper notes that while the evaluation covers three models and three activation families, the approach is general and could potentially be extended to other frontier models or architectural variants (e.g., Gated, BatchTopK, Matryoshka SAEs).