PaperlayerView exact revision on arXiv

Source-linked AI summary

BatchTopK Sparse Autoencoders

Bart Bussmann, Patrick Leask, Neel Nanda

arXiv:2412.06410v1cs.LGcs.AIstat.ML

TL;DR

Sparse autoencoders face a trade-off between reconstruction and sparsity. BatchTopK relaxes TopK’s fixed per-sample constraint to a batch-level constraint, allowing variable latent counts. Across GPT-2 Small and Gemma 2 2B, it consistently outperforms TopK and JumpReLU in reconstruction, with a downstream cross-entropy caveat at high active-latent counts.

  • Problem

    Improving SAE reconstruction at a given dictionary size can require reduced sparsity, motivating architectures that address this trade-off.

    Source: §1 · p. 1

  • Method

    BatchTopK selects the top activations across a batch, allowing variable active-latent counts per sample while preserving average sparsity.

    Source: §1 · p. 1

  • Results

    Across GPT-2 Small and Gemma 2 2B, BatchTopK consistently outperforms TopK and JumpReLU in reconstruction, although JumpReLU has less downstream CE degradation in large models with many active latents.

    Source: §1 · p. 1

  • Takeaways & Limitations

    BatchTopK allows direct specification of average sparsity rather than costly hyperparameter tuning to target a level.

    Source: §1 · p. 2

    JumpReLU sparsity cannot be set explicitly, so targeting a specific level requires costly hyperparameter tuning.

    Source: §2 · p. 3

Abstract

from arXiv · show

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting language model activations by decomposing them into sparse, interpretable features. A popular approach is the TopK SAE, that uses a fixed number of the most active latents per sample to reconstruct the model activations. We introduce BatchTopK SAEs, a training method that improves upon TopK SAEs by relaxing the top-k constraint to the batch-level, allowing for a variable number of latents to be active per sample. As a result, BatchTopK adaptively allocates more or fewer latents depending on the sample, improving reconstruction without sacrificing average sparsity. We show that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, and achieve comparable performance to state-of-the-art JumpReLU SAEs. However, an advantage of BatchTopK is that the average number of latents can be directly specified, rather than approximately tuned through a costly hyperparameter sweep. We provide code for training and evaluating BatchTopK SAEs at https://github.com/bartbussmann/BatchTopK

1 Introduction

Sparse autoencoders balance reconstruction quality against sparsity, motivating architectures that improve reconstruction at a given sparsity. BatchTopK relaxes TopK’s per-sample constraint to the batch level, and experiments report stronger reconstruction than TopK and JumpReLU across two models, with a downstream cross-entropy caveat at high sparsity.

1 Introduction

  • SAE architectures address the tension between sparsity and reconstruction, since improving reconstruction for a given dictionary can require reduced sparsity.Source: §1 · p. 1
  • BatchTopK relaxes TopK’s fixed per-sample top-k limit to a batch-level constraint, allowing variable active-latent counts while improving reconstruction without sacrificing average sparsity.A single global threshold parameter removes batch dependence during inference.Source: §1 · p. 1
  • Across GPT-2 Small and Gemma 2 2B, BatchTopK consistently outperforms TopK and JumpReLU in reconstruction, although JumpReLU has less downstream CE degradation in large models with many active latents.The introduction reports this pattern across varied dictionary sizes and sparsity levels.Source: §1 · p. 1
  • The method allows direct specification of average sparsity, avoiding costly hyperparameter tuning to target a sparsity level.Source: §1 · p. 2

2 Background: Sparse Autoencoder Architectures

Sparse autoencoders reconstruct language-model activations from sparse combinations of learned directions. Their architectures differ in how they balance reconstruction against sparsity and constrain latent activations.

2 Background: Sparse Autoencoder Architectures

  • SAEs decompose model activations into sparse linear combinations of learned directions and reconstruct the input from those latents.An SAE has an encoder and decoder; its latent representation is sparse, and its reconstructed input approximates the original activation.Source: §2 · p. 2
  • SAE training balances reconstruction and sparsity; increasing λ encourages sparsity at the cost of reconstruction error.Some architectures also use an auxiliary loss, for example to recycle inactive latents in TopK SAEs.Source: §2 · pp. 2–3
  • ReLU SAEs use an L1 penalty as an L0 proxy, but it suppresses latent activations and can be reduced by decoder scaling unless decoder directions are unit-normalized.The L1 penalty provides a gradient for training, unlike the L0 norm.Source: §2 · p. 3
  • TopK SAEs retain only the k largest activations per sample, avoiding an explicit sparsity penalty but fixing the number of active latents.An auxiliary loss based on dead latents is used to avoid dead latents.Source: §2 · p. 3
  • JumpReLU SAEs use learned per-latent thresholds and an L0 sparsity penalty, but targeting a desired sparsity requires costly hyperparameter tuning.The thresholded activation sets values below the learned threshold to zero; training uses straight-through estimators for the discontinuous function.Source: §2 · p. 3

3 BatchTopK Sparse Autoencoders

BatchTopK selects activations across a batch rather than imposing a fixed per-sample limit, allowing samples to use different numbers of latents. At inference, a threshold removes the batch dependency.

3 BatchTopK Sparse Autoencoders

  • Selecting the top n × k activations across n samples allows some samples to use more than k latents and others fewer.The method replaces sample-level TopK with a batch-level operation.Source: §3 · p. 3
  • The training objective includes an auxiliary loss scaled by α to prevent dead latents, as in TopK SAEs.Source: §3 · p. 4
  • At inference, a threshold estimated from minimum positive activations across batches removes the dependency between samples.Activations below θ are zeroed using JumpReLU instead of BatchTopK.Source: §3 · p. 4

4 Experiments

Across GPT-2 Small and Gemma 2 2B, BatchTopK generally improves reconstruction over TopK, with mixed downstream cross-entropy results against JumpReLU. Its active-latent counts vary widely across samples.

4 Experiments

  • At fixed L0=32, BatchTopK has lower normalized MSE and less cross-entropy degradation than TopK on both models.It also outperforms TopK across different k values at fixed dictionary size 12288.Source: §4 · p. 4
  • BatchTopK beats JumpReLU on both measures across all sparsity levels for GPT-2 Small, but on Gemma 2 2B its CE advantage is limited to k=16.On Gemma 2 2B, BatchTopK has better reconstruction than JumpReLU for every tested k.Source: §4 · p. 4
  • The active-latent distribution spans from one latent on some samples to more than 80 on others.The distribution’s left peak corresponds to <BOS>-token activations; the authors note that TopK would use the same number of latents regardless of input.Source: §4 · p. 4, Figure 3 · p. 5
  • The authors interpret the variable counts as corroborating their hypothesis that fixed TopK is too restrictive and samples contain different numbers of active dictionary latents.Source: §4 · p. 5

5 Conclusion

The conclusion presents BatchTopK as a batch-level relaxation of TopK’s fixed per-sample sparsity constraint. The authors report reconstruction gains without sacrificing average sparsity, while noting that human interpretability was not directly assessed.

5 Conclusion

  • BatchTopK allocates variable numbers of latents across samples and improves reconstruction without sacrificing average sparsity.The conclusion describes the method as selecting top activations across a batch rather than imposing a strict per-sample limit.Source: §5 · p. 5
  • The evaluation used reconstruction loss and sparsity, but did not directly assess human interpretability.The authors suggest architectural similarity to TopK means the latents would remain comparably interpretable.Source: §5 · p. 5
  • The authors expect future work to continue improving SAEs’ approximation of model-activation latent structure.Source: §5 · p. 5

A.1 Experimental Details

The experiments used OpenWebText sequences of up to 128 tokens and trained SAEs on residual-stream activations from GPT-2 Small and Gemma 2 2B. The appendix reports optimizer settings, dictionary sizes, sparsity levels, and JumpReLU penalties.

A.1 Experimental Details

  • SAEs were trained on OpenWebText sequences capped at 128 tokens, using Adam with learning rate 3 × 10^-4 and batch size 4096.Training processed 1 × 10^9 tokens; β1=0.9 and β2=0.99.Source: §A.1 · p. 6
  • GPT-2 Small dictionary sizes were 3072, 6144, 12288, and 24576; Gemma 2 2B used 16384, with k values 16, 32, and 64.Source: §A.1 · p. 6
  • JumpReLU sparsity penalties were 0.004, 0.0018, and 0.0008 for GPT-2 Small and 0.02, 0.005, and 0.001 for Gemma 2 2B.The bandwidth parameter was 0.001 for both experiments.Source: §A.1 · p. 6

Related papers

Loading 2412.06410v1…