AdaptiveK Sparse Autoencoders

• The paper demonstrates that adapting latent support size based on input complexity yields more faithful and interpretable reconstructions than fixed-K autoencoders.
• It employs methods like a linear-probe complexity estimator and differentiable soft TopK selection to dynamically allocate sparsity, balancing reconstruction fidelity and efficiency.
• Empirical evaluations across language and vision tasks indicate that AdaptiveK models reduce reconstruction error and improve metrics such as NMSE and cosine similarity compared to fixed sparsity baselines.

Adaptive Top K Sparse Autoencoders, commonly abbreviated AdaptiveK, are sparse autoencoders in which the latent support size is not a fixed global constant but a quantity that varies with the input. In contrast to standard TopK sparse autoencoders, which allocate exactly $K$ active latents to every sample, AdaptiveK methods are motivated by the claim that real activations differ in semantic complexity and local intrinsic dimensionality, so the number of explanatory factors required for faithful and interpretable reconstruction should also vary across samples (Yao et al., 24 Aug 2025). In current arXiv usage, the term covers both explicit per-input $k$-prediction mechanisms and closely related designs that achieve variable support size through batch-level competition or data-dependent sparse selection; among the clearest explicit instances are the complexity-driven AdaptiveK Sparse Autoencoders (Yao et al., 24 Aug 2025) and SoftSAE, which is described as “essentially a concrete AdaptiveK SAE instantiation” (Stępień et al., 7 May 2026).

1. Fixed-$K$ origins and the motivation for adaptation

The immediate background to AdaptiveK is the fixed-cardinality k-sparse autoencoder. In the original k-Sparse Autoencoder, the encoder computes

$$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$

then keeps only the $k$ largest hidden activations and sets the rest to zero, yielding exact per-example sparsity with linear reconstruction and gradient flow only through the active units (Makhzani et al., 2013). Modern TopK SAEs in mechanistic interpretability retain this basic hard-support idea, but apply it to overcomplete dictionaries trained on residual-stream or related activations (Gao et al., 2024).

The central objection to fixed $K$ is that it imposes the same representational budget on all inputs. The AdaptiveK paper states that some contexts are “semantically simple” while others are “denser, more abstract, more domain-specific, or more logically entangled,” and therefore should not be forced to use the same number of SAE features (Yao et al., 24 Aug 2025). SoftSAE makes the same point in geometric language: natural data lies on manifolds with varying local intrinsic dimensionality, so a fixed $K$ gives simple examples noisy, overly long explanations and complex examples insufficient capacity (Stępień et al., 7 May 2026).

This critique is also visible inside the large-scale fixed-TopK literature. “Scaling and evaluating sparse autoencoders” explicitly states that forcing every token to use exactly $k$ latents is “likely suboptimal” and that, ideally, one would constrain $\mathbb E[L_0]$ rather than $L_0$ itself (Gao et al., 2024). SplInterp reaches a related conclusion from a different direction: its optimal piecewise-affine baseline is a local-PCA-style autoencoder with region-specific rank $k$0, which suggests that variable local rank is theoretically natural even though the paper itself analyzes fixed-$k$1 TopK partitions (Budd et al., 17 May 2025).

2. Canonical AdaptiveK formulations

A standard SAE baseline in this literature uses

$k$2

AdaptiveK augments this with a complexity estimator and replaces the fixed TopK budget by an input-dependent one (Yao et al., 24 Aug 2025).

In “AdaptiveK Sparse Autoencoders: Dynamic Sparsity Allocation for Interpretable LLM Representations,” the adaptive signal is a linear probe trained to predict a scalar complexity score from the activation: $k$3 The probe is trained by ridge regression with objective

$k$4

and closed-form solution

$k$5

The predicted complexity $k$6 is then mapped by a sigmoid-based rule to an adaptive support size $k$7 bounded between $k$8 and $k$9, and TopK is applied with that budget (Yao et al., 24 Aug 2025). The paper reports $K$0, $K$1, $K$2, and sigmoid steepness $K$3.

SoftSAE implements the same general idea by making the support size itself a learned latent decision. It starts from

$K$4

then predicts a continuous budget

$K$5

with a single-hidden-layer Dynamic Sparsity MLP whose first layer is initialized from $K$6 (Stępień et al., 7 May 2026). During training it applies a differentiable soft top-$K$7 operator,

$K$8

and uses

$K$9

During inference, and in the late hardening phase of training, it switches to hard TopK with rounded $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$0 (Stępień et al., 7 May 2026).

The distinction between these two formulations is architectural rather than conceptual. AdaptiveK uses an externally pretrained complexity signal to set $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$1, whereas SoftSAE predicts $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$2 directly inside the SAE and trains the whole mechanism end to end. Both are explicit AdaptiveK designs in the sense that the model makes a sample-dependent decision about how many features should be active, not merely which ones.

3. The broader design space of variable-support SAEs

Not every variable-support sparse autoencoder is an explicit AdaptiveK model. The literature now contains several mechanisms that relax fixed TopK in different ways.

Method	Adaptive mechanism	Relation to AdaptiveK
AdaptiveK SAE (Yao et al., 24 Aug 2025)	Linear-probe complexity prediction mapped to $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$3	Explicit learned per-context budget
SoftSAE (Stępień et al., 7 May 2026)	Dynamic Sparsity MLP predicts continuous $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$4 with differentiable SoftTopK	Explicit per-input AdaptiveK instantiation
BatchTopK SAE (Bussmann et al., 2024)	Top $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$5 activations kept across a batch, so samplewise $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$6 varies while average stays fixed	Adaptive per sample only in a batch-average sense
Sparsemax attention SAE (Wang et al., 16 Apr 2026)	Sparsemax infers support size from a simplex projection over concept scores	AdaptiveK-like, but attention-based rather than TopK-with-$$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$7
Sampled-SAE (Oozeer et al., 29 Aug 2025)	Batch-level feature scoring restricts the candidate pool before BatchTopK	Distribution-aware feature eligibility, not explicit tokenwise $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$8

BatchTopK is the simplest nontrivial relaxation of fixed TopK. For a batch of $$\boldsymbol{z}=W^\top \boldsymbol{x}+\boldsymbol{b},$$9 samples, it keeps the top $k$0 latent activations across the entire batch, so

$k$1

and the average per-sample $k$2 is exactly $k$3, while the individual support sizes $k$4 vary (Bussmann et al., 2024). The paper explicitly interprets this as allocating more latents to more complex samples and fewer to simpler ones, but there is no separate $k$5 predictor.

“Improving Sparse Autoencoder with Dynamic Attention” replaces TopK entirely by a sparsemax cross-attention encoder. With concept-score vector $k$6, sparsemax computes

$k$7

so the support size is data-dependent through the threshold $k$8 (Wang et al., 16 Apr 2026). The paper explicitly states that sparsemax can be viewed as “a more precise version of BatchTopK, where we set $k$9 at the sample level rather than the batch level,” but its architecture is a cross-attention SAE rather than a standard encoder–TopK–decoder pipeline.

By contrast, some influential TopK-adjacent papers are best read as motivation rather than as AdaptiveK methods. “Unlocking the Address Book” introduces a Dual-Budget Strategy for KV-cache interpretation, with $K$0 for Keys and shallow layers and $K$1 for deep Values, but this is a static, representation-type-dependent policy rather than learned per-token adaptation (Ma et al., 11 Dec 2025). “TopK LLMs” moves fixed TopK sparsity into the transformer itself with $K$2, but $K$3 remains a global hyperparameter and does not vary by token, layer, or context (Takahashi et al., 26 Jun 2025).

4. Training objectives, regularization, and optimization

AdaptiveK training introduces a second problem beyond ordinary reconstruction: learning or stabilizing the mechanism that determines support size.

The complexity-driven AdaptiveK paper uses a three-phase procedure. First, it trains the ridge probe. Second, it freezes the probe and trains the SAE with

$K$4

where

$K$5

Third, it jointly fine-tunes probe and SAE with

$K$6

and

$K$7

to prevent the probe from drifting too far from the pretrained complexity estimator (Yao et al., 24 Aug 2025). The paper reports $K$8, $K$9, $K$0, and an adaptive $K$1 between $K$2 and $K$3.

SoftSAE trains adaptation end to end by replacing discrete cardinality with a differentiable relaxation. Its loss is

$K$4

where the key regularizer constrains only the expected budget,

$K$5

This means $K$6 is not a fixed per-sample support size but a target mean sparsity level, while $K$7 is free to vary under a global average constraint (Stępień et al., 7 May 2026). The auxiliary term is the standard dead-feature prevention loss based on underused neurons reconstructing residual error.

SoftSAE also identifies a training pathology specific to soft adaptive selection. As $K$8, the soft selector becomes nearly binary, but the model may “hide” information in tiny nonzero weights and thereby bypass the intended sparsity constraint. Its remedy is operationally important: freeze the Dynamic Sparsity MLP near the end of training and replace soft selection with hard TopK (Stępień et al., 7 May 2026).

BatchTopK exposes a different optimization issue: the training rule depends on the other samples in the minibatch. To remove this dependency at inference, the paper estimates a global threshold

$K$9

and replaces BatchTopK by JumpReLU-like thresholding at test time (Bussmann et al., 2024). The paper is explicit that this is a heuristic approximation rather than an exact match to the training mechanism.

Several later papers suggest complementary, rather than competing, solutions. “Beyond the Hard Budget” adds pre-TopK sparsity regularizers to fixed-$k$0 SAEs, including an off-support $k$1 penalty and a scale-invariant $k$2 ratio penalty, both restricted to batch-active units (Jacquier et al., 25 Jun 2026). SplInterp, from a spline-theoretic perspective, introduces PAM-SGD, an alternating method with exact decoder updates for TopK SAEs and argues that it works well “provided that the sparsity can adapt to the data,” which suggests a natural optimization interface for future AdaptiveK variants (Budd et al., 17 May 2025).

5. Empirical behavior across language and vision

The most direct empirical support for AdaptiveK comes from the complexity-probe results and from the reconstruction frontiers reported by explicit adaptive models.

In the AdaptiveK paper, the linear complexity probe performs comparably to more expressive regressors. The reported scores are RMSE $k$3, Pearson $k$4, and Spearman $k$5 for the linear model, versus RMSE $k$6, Pearson $k$7, and Spearman $k$8 for a one-hidden-layer MLP, and RMSE $k$9, Pearson $\mathbb E[L_0]$0, and Spearman $\mathbb E[L_0]$1 for XGBoost (Yao et al., 24 Aug 2025). The paper interprets this as evidence that context complexity is largely linearly encoded in LLM activations. It also gives concrete adaptive budgets: predicted complexity $\mathbb E[L_0]$2, $\mathbb E[L_0]$3, $\mathbb E[L_0]$4, $\mathbb E[L_0]$5, $\mathbb E[L_0]$6, and $\mathbb E[L_0]$7 (Yao et al., 24 Aug 2025). Across Pythia-70M, Pythia-160M, and Gemma-2-2B, the paper reports improved reconstruction fidelity, explained variance, and cosine similarity relative to fixed-sparsity baselines, while reducing the need for wide $\mathbb E[L_0]$8-sweeps (Yao et al., 24 Aug 2025).

SoftSAE provides the most detailed explicit AdaptiveK benchmark numbers. On CLIP embeddings at target $\mathbb E[L_0]$9, TopK achieves FVE $L_0$0, BatchTopK $L_0$1, Matryoshka $L_0$2, and SoftSAE $L_0$3; SoftSAE’s realized $L_0$4 is $L_0$5, closer to the target than BatchTopK’s $L_0$6 or Matryoshka’s $L_0$7 (Stępień et al., 7 May 2026). At $L_0$8, SoftSAE reaches FVE $L_0$9 with $k$00, versus TopK FVE $k$01 at exact $k$02 and BatchTopK FVE $k$03 at $k$04 (Stępień et al., 7 May 2026). On Gemma-2-2B activations at target $k$05, SoftSAE reports $k$06, FVE $k$07, absorption fraction $k$08, SCR top-2 $k$09, and TPP top-10 $k$10 (Stępień et al., 7 May 2026). Its strongest claim, however, is qualitative rather than purely reconstructive: $k$11 tracks intuitive semantic complexity in images and text.

BatchTopK offers a simpler form of adaptivity and consistently improves over fixed TopK in reconstruction quality at the same average sparsity. The paper states that at $k$12, BatchTopK has lower NMSE and less CE degradation than TopK on both GPT-2 Small and Gemma 2 2B, and that the per-sample active-count distribution is genuinely broad: some samples activate only a single latent, while some activate more than 80 (Bussmann et al., 2024). This is strong evidence that adaptive per-sample support size need not collapse back to near-constant $k$13.

The sparsemax attention SAE shows similar qualitative behavior in a different architecture. On OpenWeb with dictionary size $k$14, it reports NMSE $k$15 versus TopK $k$16 and cross-entropy degradation $k$17 versus TopK $k$18; on WikiText-103 at the same $k$19, it reports NMSE $k$20 versus TopK $k$21 and cross-entropy degradation $k$22 versus TopK $k$23 (Wang et al., 16 Apr 2026). In vision, the same paper reports that sparsemax explains class information with far fewer concepts than TopK or BatchTopK in top-$k$24 concept classification, which it interprets as evidence that support size should depend on sample complexity rather than be globally fixed (Wang et al., 16 Apr 2026).

6. Misconceptions, limitations, and open problems

A persistent misconception is to equate any nonconstant support size with AdaptiveK. The literature is more differentiated than that. BatchTopK is adaptive only because it enforces a batch-level top-$k$25 budget, not because it learns a dedicated per-example cardinality predictor (Bussmann et al., 2024). Sparsemax attention infers a variable support set, but does so through simplex projection in a cross-attention SAE rather than through explicit TopK-with-$k$26 gating (Wang et al., 16 Apr 2026). The KV-cache Dual-Budget Strategy uses different fixed budgets for Keys and Values, but it is “fixed Top-K with manually chosen asymmetric budgets,” not AdaptiveK in the strict sense (Ma et al., 11 Dec 2025).

A second misconception is that explicit AdaptiveK is always the simplest route to better sparse representations. “Beyond the Hard Budget” shows that fixed-$k$27 TopK SAEs can behave more adaptively in effect when pre-TopK activations are regularized, especially with the $k$28 ratio penalty that concentrates information into fewer effective units and improves robustness to inference-time changes in $k$29 (Jacquier et al., 25 Jun 2026). Conversely, “AbsTopK” argues that the more fundamental limitation of standard TopK may be nonnegativity rather than fixed support size: by selecting the $k$30 largest-magnitude activations and preserving sign, AbsTopK improves reconstruction and several steering/probing metrics while still using a fixed $k$31 (Zhu et al., 1 Oct 2025).

Adaptive methods also introduce new engineering costs. SoftSAE notes that differentiable SoftTopK becomes a bottleneck for large dictionaries and reports training times of about 2 hours versus 1.5 for baselines on CLIP and about 6 hours versus 2.5 on Gemma, together with the need to manage temperature annealing, budget annealing, and the soft-to-hard transition (Stępień et al., 7 May 2026). BatchTopK inherits batch dependence and an inference mismatch due to threshold approximation (Bussmann et al., 2024). ATM, a time-adaptive masking method designed to reduce feature absorption, achieves an absorption score of $k$32 versus TopK’s $k$33, but does so with reported $k$34 sparsity $k$35 versus TopK’s $k$36, so its stability improvements are not obtained under the same hard sparsity regime (Li et al., 9 Oct 2025).

Finally, there is a theoretical caution. “Toward Identifiable Sparse Autoencoders” does not study AdaptiveK directly, but it argues that even fixed-$k$37 TopK SAEs can be unstable because of one-sided activations, poor dictionary conditioning, and weak amortized encoders (Nelson et al., 29 May 2026). A plausible implication is that adaptive support size could improve fit while requiring stronger monitoring of support conditioning and code stability, especially if support cardinality itself becomes variable. SplInterp adds a complementary observation: fixed-$k$38 TopK partitions have a clean $k$39-th order power-diagram geometry, whereas a true AdaptiveK system would likely mix multiple support orders and therefore lose some of that geometric simplicity (Budd et al., 17 May 2025).

Across these works, AdaptiveK emerges less as a single algorithm than as a design principle: sparse autoencoders should allocate latent support in a way that reflects the heterogeneity of the data. Explicit per-input $k$40-prediction, differentiable soft top-$k$41, batch-average budgeting, sparsemax support inference, and effective-sparsity regularization are all different realizations of that principle, but they differ sharply in whether they learn how many features to use, merely vary the support as a side effect, or retain fixed cardinality while reshaping the code geometry.