Looking at Section 4 and Appendix B of the paper, the expectation is taken over the data distribution. Specifically:

𝔼_𝐱[ℒ_θ(𝐱)] where 𝐱 ~ 𝒟

Here, 𝐱 represents language model activations sampled from the training dataset 𝒟.

The key insight is that while the loss ℒ_θ(𝐱) is piecewise constant in θ for any fixed input 𝐱 (because the JumpReLU and Heaviside functions create discontinuities), when you take the expectation over many different inputs 𝐱, the expected loss becomes smooth and differentiable.

This happens because:

• For a fixed 𝐱, the discontinuities occur when pre-activations π_i(𝐱) equal the threshold θ_i
• Different inputs 𝐱 have different pre-activation values
• So the discontinuity locations vary across the data distribution
• Averaging over the distribution smooths out these discontinuities

In Equation 13, they show the derivative of this expected loss:

$\frac{\partial\mathbb{E}_{\mathbf{x}}\left[\mathcal{L}_{\boldsymbol{\theta}}(\mathbf{x})\right]}{\partial\theta_{i}} = \left(\mathbb{E}_{\mathbf{x}}\left[I_{i}(\mathbf{x})|\pi_{i}(\mathbf{x})=\theta_{i}\right]-\lambda\right)p_{i}(\theta_{i})$

where $p_i$ is the probability density of pre-activations $\pi_i(\mathbf{x})$ across the data distribution.

[deleted]

The expectation is taken over the distribution of language model activations that the SAE is being trained on.

Specifically, in the paper's notation:

• $\mathbf{x} \sim \mathcal{D}$ represents language model activations sampled from the training dataset
• The loss function is $\mathcal{L}_{\boldsymbol{\theta}}(\mathbf{x})$ which depends on both the input activations $\mathbf{x}$ and the threshold parameters $\boldsymbol{\theta}$
• The expected loss is $\mathbb{E}_{\mathbf{x}}\left[\mathcal{L}_{\boldsymbol{\theta}}(\mathbf{x})\right]$

The key insight is that while the loss for any single input $\mathbf{x}$ is piecewise constant with respect to the threshold $\theta_i$ (and thus has zero gradient almost everywhere), the expected loss over the distribution of all possible activations is smooth.

This happens because different activation samples will have their pre-activations $\pi_i(\mathbf{x})$ distributed around the threshold. As you vary the threshold slightly, you change which samples fall above or below it, creating a non-zero gradient in expectation.

In Appendix B, they show this gradient is:

$\frac{\partial\mathbb{E}_{\mathbf{x}}\left[\mathcal{L}_{\boldsymbol{\theta}}(\mathbf{x})\right]}{\partial\theta_{i}} = \left(\mathbb{E}_{\mathbf{x}}\left[I_{i}(\mathbf{x})|\pi_{i}(\mathbf{x})=\theta_{i}\right]-\lambda\right)p_{i}(\theta_{i})$

where $p_i$ is the probability density function of the pre-activation distribution $\pi_i(\mathbf{x})$. This density needs to be estimated from data, which is where the kernel density estimation techniques come in.

The training efficiency advantages of JumpReLU SAEs stem from their simpler computational requirements compared to Gated and TopK architectures:

JumpReLU vs Gated SAEs

Gated SAEs require:

• Two encoder paths: Computing both $\boldsymbol{\pi}_{\text{gate}}(\mathbf{x})$ and $\boldsymbol{\pi}_{\text{mag}}(\mathbf{x})$ means double the encoder computations
• Auxiliary loss with frozen decoder: The auxiliary term $\|\mathbf{x} - \hat{\mathbf{x}}_{\text{frozen}}(\text{ReLU}(\boldsymbol{\pi}_{\text{gate}}(\mathbf{x})))\|_2^2$ requires maintaining a frozen copy of the decoder and computing an additional reconstruction
• Element-wise multiplication: The final features are computed as $\mathbf{f}_{\text{gate}}(\mathbf{x}) \odot \mathbf{f}_{\text{mag}}(\mathbf{x})$

JumpReLU SAEs only need:

• Single encoder path: Just compute $\boldsymbol{\pi}(\mathbf{x}) = \mathbf{W}_{\text{enc}}\mathbf{x} + \mathbf{b}_{\text{enc}}$ once
• No auxiliary loss: The L0 penalty provides sufficient gradient signal without auxiliary terms
• Simple thresholding: Apply JumpReLU$_\theta(z) = z \cdot H(z-\theta)$ elementwise

JumpReLU vs TopK SAEs

TopK SAEs require:

• Partial sorting: Finding the top K activations has complexity $O(M \log K)$ where $M$ is the dictionary size (e.g., 131k features)
• Even approximate TopK (like jax.lax.approx_max_k) is more expensive than elementwise operations
• AuxK auxiliary loss: Adds computational overhead to reduce dead features

JumpReLU's elementwise operation is just:

feature_magnitudes = x * (x > threshold)  # O(M) elementwise

Memory Efficiency

JumpReLU also has lower memory requirements:

• No need to store frozen decoder copies (unlike Gated)
• No need to maintain sorting indices or auxiliary computations (unlike TopK)
• Straight-through estimators add minimal memory overhead

The paper notes these efficiency gains are particularly important when training at scale - with 131k-width SAEs on billions of tokens, even small per-step efficiency improvements compound significantly.

The training efficiency advantage of JumpReLU over TopK SAEs comes down to the fundamental difference in their activation mechanisms:

TopK's computational overhead:

• Partial sorting requirement: For each input, TopK must identify and retain only the K largest pre-activations out of M total features (where M is the dictionary size, e.g., 131k in the paper's experiments)
• Complexity: Even with efficient algorithms, this requires O(M log K) operations for exact TopK, or O(M) for approximate methods
• Non-local operation: The activation of any feature depends on the values of all other features, requiring global information

JumpReLU's efficiency:

• Simple elementwise comparison: Each feature activation is determined by a single comparison: if pre_activation > threshold
• Complexity: O(M) with a very low constant factor - just M comparisons
• Local operation: Each feature's activation is independent of others

Practical implications: The paper mentions using jax.lax.approx_max_k (based on Chern et al., 2022) for TopK implementation, which is faster than exact sorting but still involves:

• Multiple passes through the data to approximate the top-K selection
• Additional memory access patterns that are less cache-friendly
• More complex gradient computation in the backward pass

For a 131k-width SAE trained on billions of tokens, these differences compound significantly. The elementwise nature of JumpReLU also makes it more amenable to parallelization and hardware acceleration compared to the inherently sequential nature of sorting operations.

This efficiency gain is particularly important when training multiple SAEs with different sparsity levels or when scaling to larger models and dictionary sizes.

In the context of Figure 1's caption, "direction" refers to a feature direction in the Sparse Autoencoder (SAE) - essentially one of the learned dictionary elements.

More precisely, each direction corresponds to:

• A row in the encoder weight matrix $\mathbf{W}_{\text{enc}}$
• A column in the decoder weight matrix $\mathbf{W}_{\text{dec}}$ (denoted as $\mathbf{d}_i$ in the paper)

So when the caption says "Consider a direction in which the encoder pre-activation is high...", it's referring to one specific feature in the SAE's dictionary. The figure shows how the encoder pre-activation $\pi_i(\mathbf{x}) = (\mathbf{W}_{\text{enc}}\mathbf{x} + \mathbf{b}_{\text{enc}})_i$ behaves for that particular feature $i$.

The key insight illustrated is that for this feature:

• When the feature is truly "active" (should be detected), the pre-activation is high
• When the feature is "inactive" (shouldn't be detected), the pre-activation is low but sometimes still positive

This creates the challenge that ReLU addresses poorly but JumpReLU handles well by introducing an adjustable threshold.

Yes, you've got it exactly right! Figure 1 illustrates the key advantage of JumpReLU over standard ReLU for sparse autoencoders.

The core insight is that with a learned threshold vector $\boldsymbol{\theta}$, JumpReLU can:

1. Filter out false positives: Set the threshold above the "noise floor" of pre-activations that occur when the feature is actually inactive (but happen to be slightly positive)

2. Preserve true activation magnitudes: Unlike solutions that shift the bias to make more pre-activations negative, JumpReLU doesn't systematically shrink the magnitudes of true activations - it just passes them through unchanged if they're above the threshold

The figure shows this nicely:

• ReLU (center-left): Lets through all positive values, including false positives
• ReLU with shifted bias (center-right): Eliminates false positives but systematically underestimates all magnitudes
• JumpReLU (far-right): Clean separation - zeros out everything below $\theta$, preserves everything above

This is why the paper emphasizes that JumpReLU (like TopK and Gated SAEs) can "separate out deciding which features are active from estimating active features' magnitudes." The threshold handles the binary active/inactive decision, while the magnitude above threshold captures the strength of activation.

Sure! Here's a simple toy example to illustrate how the L1 penalty is not invariant under reparameterizations:

Original SAE Setup

Consider a simple SAE with:

• Input dimension: 2
• Dictionary size: 2
• Input vector: $\mathbf{x} = [1, 0]$

Original parameters:

• Encoder weights: $\mathbf{W}_{\text{enc}} = \begin{bmatrix} 2 & 0 \\ 0 & 2 \end{bmatrix}$
• Encoder bias: $\mathbf{b}_{\text{enc}} = [-1, -1]$
• Decoder weights: $\mathbf{W}_{\text{dec}} = \begin{bmatrix} 0.5 & 0 \\ 0 & 0.5 \end{bmatrix}$
• Decoder bias: $\mathbf{b}_{\text{dec}} = [0, 0]$

Forward pass:

1. Pre-activations: $\boldsymbol{\pi} = \mathbf{W}_{\text{enc}}\mathbf{x} + \mathbf{b}_{\text{enc}} = [2 \cdot 1 - 1, 0 \cdot 1 - 1] = [1, -1]$
2. Feature activations: $\mathbf{f} = \text{ReLU}([1, -1]) = [1, 0]$
3. Reconstruction: $\hat{\mathbf{x}} = \mathbf{W}_{\text{dec}}\mathbf{f} = [0.5 \cdot 1, 0 \cdot 1] = [0.5, 0]$
4. L1 penalty: $\|\mathbf{f}\|_1 = 1$

Reparameterized SAE

Now, let's scale encoder weights down by factor 10 and decoder weights up by factor 10:

Reparameterized parameters:

• $\mathbf{W}'_{\text{enc}} = \frac{1}{10}\mathbf{W}_{\text{enc}} = \begin{bmatrix} 0.2 & 0 \\ 0 & 0.2 \end{bmatrix}$
• $\mathbf{b}'_{\text{enc}} = \frac{1}{10}\mathbf{b}_{\text{enc}} = [-0.1, -0.1]$
• $\mathbf{W}'_{\text{dec}} = 10 \cdot \mathbf{W}_{\text{dec}} = \begin{bmatrix} 5 & 0 \\ 0 & 5 \end{bmatrix}$
• $\mathbf{b}'_{\text{dec}} = \mathbf{b}_{\text{dec}} = [0, 0]$

Forward pass:

1. Pre-activations: $\boldsymbol{\pi}' = [0.2 \cdot 1 - 0.1, 0 \cdot 1 - 0.1] = [0.1, -0.1]$
2. Feature activations: $\mathbf{f}' = \text{ReLU}([0.1, -0.1]) = [0.1, 0]$
3. Reconstruction: $\hat{\mathbf{x}}' = [5 \cdot 0.1, 0 \cdot 0.1] = [0.5, 0]$
4. L1 penalty: $\|\mathbf{f}'\|_1 = 0.1$

Key Observations

1. Same reconstruction: $\hat{\mathbf{x}} = \hat{\mathbf{x}}' = [0.5, 0]$
2. Same sparsity pattern: Both have only the first feature active
3. Different L1 penalties: Original = 1, Reparameterized = 0.1

By scaling down the encoder by any factor $\alpha$ and scaling up the decoder by $1/\alpha$, we can make the L1 penalty arbitrarily small while keeping everything else the same. This is why the L1 penalty needs to be paired with decoder weight normalization constraints (typically $\|\mathbf{d}_i\|_2 = 1$) to be meaningful.

The RI-L1 penalty $\sum_i f_i \|\mathbf{d}_i\|_2$ fixes this issue because when we scale decoder weights up, the penalty increases proportionally, canceling out the decrease from smaller feature magnitudes.

The paper identifies two key limitations of L1 sparsity penalties for training SAEs:

1. Feature Magnitude Suppression

The L1 penalty penalizes feature magnitudes in addition to sparsity, which harms reconstruction fidelity. This is because L1 measures the sum of absolute values, so it incentivizes making active features smaller even when they should have larger values for accurate reconstruction.

2. Lack of Reparameterization Invariance

The L1 penalty is not invariant under reparameterizations of the SAE. Specifically:

• You can scale down encoder parameters and scale up decoder parameters proportionally
• This arbitrarily shrinks feature magnitudes (and thus the L1 penalty)
• But it doesn't change the number of active features or the SAE's output reconstructions
• This creates a "cheat" where the model can reduce the loss without actually becoming more sparse

Due to this issue, L1-trained SAEs require additional constraints during training, typically constraining decoder weight matrix columns (𝐝ᵢ) to have unit norm.

The paper introduces the RI-L1 (Reparameterization-Invariant L1) penalty as an alternative:

$S_{\text{RI-L1}}(\mathbf{f}) = \sum_{i=1}^{M} f_i \|\mathbf{d}_i\|_2$

This weights feature coefficients by the norms of corresponding dictionary directions, making it invariant to the reparameterization problem.

By using L0 instead, JumpReLU SAEs avoid both issues - L0 only counts the number of non-zero features, not their magnitudes, and is inherently invariant to reparameterization.

Yes, exactly! Equation 10 is showing that when you expand the L0 norm of the feature vector $\mathbf{f}(\mathbf{x})$ using the JumpReLU definition, you get a sum of Heaviside step functions.

Here's the reasoning:

• Each feature $f_i(\mathbf{x}) = \text{JumpReLU}_{\theta_i}(\pi_i(\mathbf{x})) = \pi_i(\mathbf{x}) \cdot H(\pi_i(\mathbf{x}) - \theta_i)$
• The L0 norm counts non-zero entries
• $f_i(\mathbf{x})$ is non-zero exactly when $H(\pi_i(\mathbf{x}) - \theta_i) = 1$ (since pre-activations are continuous, $\pi_i(\mathbf{x}) = 0$ happens with probability 0)
• Therefore $\|f_i(\mathbf{x})\|_0 = H(\pi_i(\mathbf{x}) - \theta_i)$

So the sparsity penalty becomes: $\mathcal{L}_{\text{sparsity}} = \lambda\|\mathbf{f}(\mathbf{x})\|_0 = \lambda\sum_{i=1}^{M}H(\pi_i(\mathbf{x})-\theta_i)$

This reformulation is crucial for the paper's approach because it shows that the threshold parameters $\theta_i$ only appear inside Heaviside functions in the loss, which is what motivates their use of straight-through estimators to get gradient signals for training these thresholds.

Great observation! The ð (eth) symbol is not a standard mathematical convention - it's a notation choice specific to this paper.

The authors explain in footnote 4:

"We use the notation ð/ð⁢z to denote pseudo-derivatives, to avoid conflating them with actual partial derivatives for these functions."

They're using ð instead of the standard ∂ (partial derivative) because the JumpReLU and Heaviside functions aren't differentiable in the classical sense - they have discontinuities at the threshold. The "pseudo-derivatives" are custom gradient definitions they create for the backward pass during training.

This is actually quite thoughtful notation design! Since they're defining straight-through estimators (STEs) that approximate gradients where true gradients don't exist or are zero, using a visually similar but distinct symbol helps avoid confusion. It makes it immediately clear when they're talking about:

• ∂: True partial derivatives (for differentiable functions)
• ð: Their custom-defined gradients for training (for non-differentiable points)

The eth symbol (ð) comes from Old English and Icelandic alphabets. While it's occasionally used in some mathematical contexts, this specific usage as "pseudo-derivative" appears to be novel to this paper. It's a clever way to maintain mathematical clarity while introducing new gradient approximations for discontinuous functions.

Looking at the Gao et al. paper, the definition of "fraction of variance unexplained" (FVU) can be found in Section 2.1 under "Evaluation":

Evaluation: After training, we evaluate autoencoders on sparsity L0, and reconstruction mean-squared error (MSE). We report a normalized version of all MSE numbers, where we divide by a baseline reconstruction error of always predicting the mean activations.

So FVU is calculated as:

$\text{FVU} = \frac{\text{MSE of SAE reconstruction}}{\text{MSE of always predicting the mean}}$

This is also referred to as the "normalized loss" in the paper. It's a standard normalization that measures what fraction of the variance in the data remains unexplained by the model. A perfect reconstruction would have FVU = 0, while a model that does no better than predicting the mean would have FVU = 1.

The JumpReLU paper uses this as a secondary metric alongside their primary metric of delta LM loss (the increase in cross-entropy loss when the SAE is spliced into the LM's forward pass).

Figure 5 is plotting two important characteristics of Sparse Autoencoders (SAEs) against each other to reveal a key trade-off:

Delta LM Loss (X-axis)

• What it measures: The increase in cross-entropy loss when the SAE is spliced into the language model's forward pass
• Why it matters: This is a measure of reconstruction fidelity - lower values mean the SAE is better at faithfully reconstructing the original activations without harming the model's performance
• Interpretation: Moving left on the x-axis means better reconstruction quality

Feature Activation Frequency (Y-axis)

• What it measures: The proportion of SAE features that fire very frequently (>10% of tokens in top plot, >1% in bottom plot)
• Why it matters: Very high frequency features tend to be:
  • Less interpretable (they activate so often they may not represent specific, meaningful concepts)
  • Less sparse (defeating the purpose of "sparse" autoencoders)
  • Potentially acting more like bias terms than meaningful feature detectors

The Significance of This Plot

1. Different architectures have different characteristics: The plot reveals that TopK and JumpReLU SAEs tend to learn more high-frequency features than Gated SAEs, especially in the low delta LM loss (high fidelity) regime.

2. Trade-off visualization: As you improve reconstruction fidelity (move left), you tend to get more high-frequency features. This suggests there's a fundamental tension between:

   • Making accurate reconstructions (low delta LM loss)
   • Learning sparse, interpretable features (few high-frequency features)

3. Architecture comparison at fixed performance: By plotting against delta LM loss rather than just comparing raw frequency counts, the authors ensure they're comparing architectures at equivalent levels of reconstruction quality. This is a fairer comparison than just looking at frequency distributions alone.

4. Practical implications: For practitioners, this helps in choosing architectures - if you care more about avoiding high-frequency features for interpretability, you might prefer Gated SAEs despite slightly worse reconstruction fidelity.

The paper notes that while JumpReLU has more very high frequency features (>10%) like TopK, the absolute numbers remain small (<0.06% of all features even in the worst case), suggesting this may not be a major practical concern.