anthology-of-the-sota

SOTA-344: Train sparse autoencoders with a TopK activation instead of an L1 penalty, and prevent dead latents with transposed-decoder initialization and an auxiliary loss

2 min read

Properties5
StatusProposed
Consensusunassessed
Tagsanalysis-and-evaluation, concept-geometry
SourceLIT-571 — Scaling and evaluating sparse autoencoders
Introduced byLIT-571 — Scaling and evaluating sparse autoencoders

Graph View

Source

Gao et al. (2024), LIT-571. Read as NOTE-307.

The practice

If you train an SAE on a model’s activations:

  • Use TopK for sparsity. z = TopK(W_enc(x − b_pre)), loss = MSE. Choose k as the L0 you want. There is no L1 coefficient to sweep, and runs at different sizes compare at identical sparsity
  • Initialize the encoder to the decoder’s transpose
  • Add the auxiliary dead-latent loss: reconstruct the current error from the top-k_aux latents that have not fired recently. Together with the initialization, this is what took dead latents from up to 90% to 7%
  • Scale the learning rate as 1/√(latents), the largest that converged in the paper

The recipe is Gao et al.’s, LIT-571, which introduced TopK SAEs together with the dead-latent mitigations. At 32k latents TopK beat ReLU and ProLU SAEs and matched Gated on the MSE–L0 frontier, and at fixed L0 = 128 it scaled more steeply with size (Figure 2). Without mitigation up to 90% of latents died; with the transposed initialization and the auxiliary loss, 7% did in their 16M-latent SAE on GPT-4.

Conditions

  • This says how to train an SAE, not whether to use one. On probing, SAE latents do not beat raw activations (LIT-568, SOTA-345)
  • Fixed per-token L0 is a known weakness. The authors suggest constraining the expected L0 instead, which BatchTopK does

Backlinks