BatchTopK relaxes per-sample sparsity to a per-batch budget

measured in 1 paper

Bussmann, Leask & Nanda replace TopK SAE's per-sample k-active constraint with a per-batch budget of n*k activations, so per-sample counts vary while the batch average stays k [bussmann-etal-2024-batchtopk] Single-sample inference is restored via a JumpReLU-style threshold estimated from training batches [bussmann-etal-2024-batchtopk] On GPT-2 Small and Gemma-2-2B, BatchTopK consistently beats TopK on reconstruction and cross-entropy across widths and sparsities [bussmann-etal-2024-batchtopk] Against JumpReLU the result is scale-dependent: BatchTopK wins on GPT-2 at every sparsity but on Gemma-2-2B only at k=16 for CE degradation [bussmann-etal-2024-batchtopk] Unlike JumpReLU, BatchTopK sets average sparsity directly as a hyperparameter rather than via an L0-penalty sweep [bussmann-etal-2024-batchtopk]

Structure

Linear Direction

Context

batch-level relaxation of a per-sample top-k sparsity constraint, variable per-sample active-latent count with a fixed batch-average, inference-time threshold estimated from training batch statistics, applied as a JumpReLU gate, scale-dependent comparison to JumpReLU (wins on a small model, only partially on a larger one)

Confirmed in models

GPT-2 Small · GPT-2, Gemma-2-2B · Gemma

Method

BatchTopK Sparse Autoencoders

Papers

BatchTopK Sparse Autoencoders — Bussmann, Bart, Leask, Patrick, Nanda, Neel2024 · arXiv:2412.06410