Sparse Autoencoders

The superposition hypothesis holds that a network represents more features than it has dimensions by placing them in almost-orthogonal directions and relying on their sparsity to keep interference low. If true, the activations are sparse combinations of an overcomplete dictionary, and that dictionary is recoverable.

Weighting the L1 by the decoder column norm is not cosmetic: without it the penalty is trivially gamed by shrinking $f$ and growing $d$ proportionally. The alternative is to constrain the columns to unit norm and project after every step.

Shrinkage

An L1 penalty biases every active feature toward zero, so reconstructions are systematically too small. Top-k and JumpReLU encoders replace the penalty with a hard sparsity constraint, which removes the bias at the cost of a non-differentiable activation that needs a straight-through estimator.

import torch
from torch import Tensor, nn


class TopKSAE(nn.Module):
    """Top-k encoder: sparsity is a constraint, not a penalty,
    so there is no shrinkage to correct for."""

    def __init__(self, d_model: int, d_hidden: int, k: int = 32):
        super().__init__()
        self.enc = nn.Linear(d_model, d_hidden)
        self.dec = nn.Linear(d_hidden, d_model, bias=False)
        self.pre_bias = nn.Parameter(torch.zeros(d_model))
        self.k = k

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
        centred = x - self.pre_bias
        acts = self.enc(centred)
        vals, idx = acts.topk(self.k, dim=-1)
        f = torch.zeros_like(acts).scatter_(-1, idx, vals.relu())
        return self.dec(f) + self.pre_bias, f

    @torch.no_grad()
    def normalise_decoder_(self) -> None:
        self.dec.weight /= self.dec.weight.norm(dim=0, keepdim=True)

Subtracting a learned pre-bias before encoding and adding it back after decoding lets the dictionary model deviations from the mean activation rather than the mean itself — without it the first few features are spent representing a constant. Call normalise_decoder_ after each optimiser step, and remove the component of the gradient parallel to each column beforehand if you want the step to be exact rather than approximately tangent.