Sparse Autoencoders

The problem, from first principles

Individual neurons in a language model are hard to read because the model packs many more concepts into its activations than it has dimensions (superposition), so each neuron responds to a mixture. A sparse autoencoder (SAE) tries to unpack them. It maps an activation vector up into a much wider layer of feature activations, keeps only a few of them active through a sparsity penalty, and reconstructs the original activation from those few features. If it works, each feature fires on one recognizable concept. The training objective is a trade-off: reconstruct well, but explain each activation with as few features as possible.

From theory to code

Implement sae_forward and sae_loss.

Constraints

• x has shape (n, d). The encoder is f = relu((x - b_dec) @ W_enc + b_enc) with W_enc of shape (d, m). The decoder is x_hat = f @ W_dec + b_dec with W_dec of shape (m, d).
• sae_forward(x, W_enc, b_enc, W_dec, b_dec) returns (x_hat, f).
• sae_loss(x, x_hat, f, l1_coeff) returns mse + l1_coeff * l1 where mse is the mean over examples of the sum over dimensions of (x - x_hat) ** 2, and l1 is the mean over examples of the sum of |f| over features. Return a Python float.
• m > d is expected (an overcomplete dictionary).

Hints