bartbussmann / batchtopk

BatchTopK Sparse Autoencoders

BatchTopK is a novel approach to Sparse Autoencoders (SAEs) that offers an alternative to TopK SAEs as introduced by OpenAI. This repository contains the implementation and experiments for BatchTopK SAEs, as described in our preliminary findings.

What is the BatchTopK activation function?

BatchTopK modifies the TopK activation in SAEs in the following way:

1. Instead of applying TopK to each sample independently, it flattens all feature activations across the batch.
2. It then takes the top (K * batch_size) activations.
3. Finally, it reshapes the result back to the original batch shape.

Usage

git clone https://github.com/bartbussmann/BatchTopK.git
cd BatchTopK
pip install transformer_lens
python main.py

Acknowledgments

The training code is heavily inspired and basically a stripped-down version of SAELens. Thanks to the SAELens team for their foundational work in this area!

batchtopk's Issues

aux and l1 loss may not be worthwhile for topksae

Just a quick comment: I did a little bit of profiling with this code (thanks by the way, great to have a clean reference) and it seems that for the TopKSAE at least the aux_loss and the l1_loss do not significantly alter training performance. The l2_loss alone seems to work just as well.

      l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
      variance = ((x - x.mean(0)) ** 2).mean()
      l1_norm = acts_topk.float().abs().sum(-1).mean()
      l1_loss = self.cfg["l1_coeff"] * l1_norm
      l0_norm = (acts_topk > 0).float().sum(-1).mean()
      aux_loss = self.get_auxiliary_loss(x, x_reconstruct, acts)
      loss = l2_loss + l1_loss + aux_loss