Top-AFA: Adaptive Activation for SAEs

• Top-AFA is an input-adaptive activation function for sparse autoencoders that adjusts active feature counts based on embedding norms.
• It replaces the fixed-k hyperparameter with a rule that matches the sparse code norm to the dense embedding norm under quasi-orthogonality.
• Empirical evaluations on GPT-2 embeddings show that Top-AFA yields lower NMSE and improved quasi-orthogonality metrics compared to traditional top-k methods.

Top-AFA is an input-adaptive activation function for sparse autoencoders (SAEs) introduced in “Evaluating and Designing Sparse Autoencoders by Approximating Quasi-Orthogonality” (Lee et al., 31 Mar 2025). In mechanistic interpretability, where SAEs are used to decompose dense language-model embeddings into sparse feature vectors, Top-AFA replaces the fixed hyperparameter $k$ of top-$k$ SAEs with a rule that selects, for each input, however many features are needed so that the sparse feature norm matches the norm predicted by the input embedding under approximate quasi-orthogonality. In the same framework, “AFA” denotes approximate feature activation: a theoretical link between the norm of the dense embedding and the norm of the sparse code, used both to design the activation rule and to evaluate pretrained SAEs (Lee et al., 31 Mar 2025).

1. Context and motivation

Sparse autoencoders are widely used in mechanistic interpretability research for LLMs. In the basic formulation, a dense hidden embedding $\mathbf{z}\in\mathbb{R}^d$ is mapped to a sparse feature vector $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$, typically with $h>d$, and then reconstructed linearly as

$$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$

The paper situates Top-AFA against the state of the art in $k$-sparse autoencoders, where sparsity is imposed by keeping exactly $k$ nonzero activations. Its central critique is that the choice of $k$, often treated as an $\ell_0$ hyperparameter, lacks a theoretical grounding. A fixed $k$0 also assumes that every input should activate the same number of features, despite possible variation in embedding complexity across tokens and layers (Lee et al., 31 Mar 2025).

Top-AFA is proposed as a replacement for fixed-$k$1 sparsification. Rather than prescribing a constant activation count, it uses the input embedding itself to infer how much sparse activation should occur. This suggests an input-conditioned sparsity regime in which activation count is an emergent quantity rather than a manually specified constant.

2. Approximate feature activation and quasi-orthogonality

The theoretical basis of Top-AFA is the approximate feature activation formulation. The paper assumes the linear representation hypothesis for SAEs with linear decoders: $k$2 where $k$3 is the decoder dictionary and $k$4 is the sparse code. This is paired with the superposition setting $k$5.

The relevant geometric notion is $k$6-quasi-orthogonality. For a column-normalized dictionary $k$7, the paper defines

$k$8

If $k$9 is approximately quasi-orthogonal, then $\mathbf{z}\in\mathbb{R}^d$0 is approximately identity, and the dense embedding norm and sparse-code norm become tightly linked. The paper’s central bound is

$\mathbf{z}\in\mathbb{R}^d$1

Equivalently,

$\mathbf{z}\in\mathbb{R}^d$2

In the limit $\mathbf{z}\in\mathbb{R}^d$3, this reduces to $\mathbf{z}\in\mathbb{R}^d$4. The paper uses that approximation as its operative target: the input embedding norm is treated as the theoretically expected sparse activation norm. A related corollary is that when $\mathbf{z}\in\mathbb{R}^d$5,

$\mathbf{z}\in\mathbb{R}^d$6

so $\mathbf{z}\in\mathbb{R}^d$7 can be read as an approximation to the number of activated features. This is the paper’s route from an $\mathbf{z}\in\mathbb{R}^d$8-norm statement to adaptive control of $\mathbf{z}\in\mathbb{R}^d$9-style sparsity (Lee et al., 31 Mar 2025).

3. Definition of the Top-AFA activation rule

Top-AFA operationalizes the AFA theory as an activation rule. For a batch input $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$0, the algorithm first centers the input,

$\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$1

then computes the target norm

$\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$2

Positive pre-activations are obtained by

$\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$3

Each candidate feature is scored by its squared activation weighted by the norm of the corresponding decoder column: $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$4 The score vector is sorted in descending order, cumulative sums are taken, and the selected activation count is the prefix whose cumulative norm is closest to the target: $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$5 after converting both $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$6 and the cumulative sums to norms. The final sparse code keeps exactly those top $\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$7 activations for that input. Reconstruction then proceeds as

$\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$8

This differs from standard top-$\mathbf{f}(\mathbf{z})\in\mathbb{R}^h$9 SAEs in that $h>d$0 is no longer fixed globally. It also differs from batch top-$h>d$1, which constrains average sparsity over a batch rather than choosing an input-specific activation count. The paper notes that decoder columns need not have unit norm; weighting by $h>d$2 compensates for that by absorbing decoder scaling into the feature contribution. The associated norm-matching penalty is

$h>d$3

and the full objective is

$h>d$4

In the reported experiments, $h>d$5, while $h>d$6 is tuned over $h>d$7 (Lee et al., 31 Mar 2025).

4. Evaluation methodology derived from AFA

The same theory is used to evaluate pretrained SAEs. The paper argues that standard SAE evaluation typically measures either reconstruction quality, through quantities such as MSE or NMSE, or sparsity in isolation, but does not assess whether the magnitude of feature activation is justified by the input embedding itself. To address that gap, it defines a lower-bound estimate on the degree of quasi-orthogonality required to explain the observed mismatch between dense and sparse norms: $h>d$8 This quantity is interpreted as the minimum quasi-orthogonality level any dictionary would need in order to explain the activation norm actually produced by the SAE.

The paper also introduces ZF plots, which place $h>d$9 and $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$0 on opposing axes. In that visualization, points near the diagonal indicate agreement with the AFA prediction, while large deviations indicate over- or under-activation relative to the input-conditioned norm target. The paper presents $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$1 as complementary to NMSE rather than a replacement for it: a model can reconstruct accurately while still using feature activations whose norm is difficult to reconcile with quasi-orthogonal sparse decomposition. The paper interprets lower $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$2 as evidence of better-justified feature activations (Lee et al., 31 Mar 2025).

5. Empirical performance on GPT-2 SAEs

The empirical study trains SAEs on GPT-2 Small residual-stream embeddings from layers 6, 7, and 8, using OpenWebText with sequence length 128. Training runs for 20,000 iterations with 4,096 tokens per iteration, for a total of 81,920,960 tokens, and uses width $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$3 on a single NVIDIA A100 GPU. The paper compares Top-AFA against fixed top-$$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$4 and batch top-$$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$5 baselines (Lee et al., 31 Mar 2025).

The main reported result is that Top-AFA achieves the best NMSE on all three tested layers. The best settings reported in the paper are summarized below.

Layer	Best Top-AFA result	Best fixed-sparsity baselines
6	$$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$6, $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$7, NMSE $$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$8	top-$$\mathbf{f}(\mathbf{z})=\sigma(W_{\text{enc}}\mathbf{z}+\mathbf{b}_{\text{enc}}), \qquad \hat{\mathbf{z}}=W_{\text{dec}}\mathbf{f}(\mathbf{z})+\mathbf{b}_{\text{dec}}.$$9: $k$0; batch top-$k$1: $k$2
7	$k$3, $k$4, NMSE $k$5	top-$k$6: $k$7; batch top-$k$8: $k$9
8	$k$0, $k$1, NMSE $k$2	top-$k$3: $k$4; batch top-$k$5: $k$6

The reported average $k$7 values are notably larger than the best fixed-$k$8 baselines, which is consistent with the paper’s claim that the correct activation count should vary with the input rather than be set in advance. The paper also reports that Top-AFA achieves better $k$9 distributions than top-$k$0 and batch top-$k$1, suggesting that the gain is not confined to reconstruction alone.

Computationally, the paper states that Top-AFA has the same asymptotic top-selection cost as top-$k$2, with leading complexity

$k$3

and is therefore not asymptotically worse than the compared activation rules. This suggests that its adaptive sparsity is not obtained by a qualitatively heavier activation mechanism.

6. Limitations, scope, and interpretation

The paper is explicit that the AFA theory is derived under the linear representation hypothesis with a one-layer linear decoder. Extensions to more general superposition mappings, two-layer decoders, or architectures such as Matryoshka SAEs are left open. It also notes that the expectation argument used to identify the embedding norm with the expected sparse-feature norm is simplifying, and that a more rigorous treatment could yield looser but more principled bounds (Lee et al., 31 Mar 2025).

A second limitation is that Top-AFA removes the inference-time hyperparameter $k$4 but still introduces a training-time regularization coefficient $k$5. The paper treats this as fundamentally different from choosing a fixed activation count, because $k$6 shapes optimization rather than directly prescribing $k$7, but the method is not hyperparameter-free. The paper also acknowledges that real pretrained decoder dictionaries are often less quasi-orthogonal than ideal high-dimensional geometry would permit, so the approximation $k$8 is not exact in practice.

Finally, the empirical case made for Top-AFA is primarily geometric and reconstruction-centered. The paper does not claim a direct proof that lower $k$9 yields uniquely canonical or monosemantic features. A plausible implication is that Top-AFA improves the alignment between input magnitude and sparse activation magnitude, which may be favorable for interpretability, but the paper stops short of proving that this alone resolves broader questions about feature semantics.

Top-AFA is therefore best understood as a theoretically motivated, input-adaptive alternative to fixed top-$\ell_0$0 sparsification in SAEs. Its distinctive contribution is not merely higher reconstruction quality, but a reformulation of sparsity control around approximate feature activation: the sparse code should activate to the extent predicted by the input embedding norm, not to the extent imposed by a global constant.