What survives the bottleneck

An autoencoder is a network trained to give back its input. An encoder $f_\phi : \mathbb{R}^d \to \mathbb{R}^k$ maps $x$ to a code $z$, a decoder $g_\theta : \mathbb{R}^k \to \mathbb{R}^d$ maps the code back, and training minimises

$$\mathcal{L}(\phi, \theta) = \mathbb{E}_x\, \lVert x - g_\theta(f_\phi(x)) \rVert^2 .$$

Left alone, the task is trivial: with $k \ge d$ both maps can be the identity. Everything interesting comes from the constraint placed on the code, and the constraint decides what the model is for. A narrow code compresses. A code that is a distribution can be sampled, so the decoder becomes a generator. A sparse code, wider than the input, can be read: each coordinate comes to stand for one thing. This piece takes the three in turn, and trains each one live in front of you.

A narrow code

Start with the linear case, $f(x) = Ex$ and $g(z) = Dz$ with $k \lt d$. Their product $DE$ has rank at most $k$, so the autoencoder is searching for the best rank-$k$ map from the data to itself. That is the Eckart–Young problem from the piece on low-rank updates, and its answer is the same: projection onto the top $k$ principal directions. For centred data with covariance eigenvalues $\lambda_1 \ge \dots \ge \lambda_d$,

$$\min_{D,\,E}\; \mathbb{E}\,\lVert x - DEx\rVert^2 = \sum_{j > k} \lambda_j .$$

Baldi and Hornik showed more: this loss has no spurious local minima. Every local minimum is global, and every other critical point is a saddle, so gradient descent finds PCA.Baldi and Hornik, Neural networks and principal component analysis, 1989. It finds the principal subspace, though, not the principal directions.For any invertible $k\times k$ matrix $M$, the pair $(DM^{-1}, ME)$ reconstructs exactly as well. The code’s coordinates are fixed only up to $M$, which is why nothing forces them to mean anything. The sparse autoencoder at the end of this piece is one answer to that.

A nonlinear autoencoder can do better when the data lie near a curved surface instead of a flat one. The figure trains one with a single number as its code, on points near a curve in the plane. It trains while it is on the page; change any setting and it starts again.

Turn off the nonlinear layers first. The blue line straightens onto PCA’s dashed line, and the error trace settles exactly on PCA’s level: a linear autoencoder is PCA, found slowly. Turn them back on and the decoder’s curve bends to follow the data, and the error falls by close to two orders of magnitude.

The strip of codes is the more telling panel. On the arc, the shading runs light to dark from one end to the other. The code has recovered position along the curve, up to direction and scale, though nothing in the loss asked for it. The readout measures this as rank agreement between code and true position.

Now choose the spiral. The error can still get low while the order is lost: the decoder’s curve visits the arms in a different sequence, jumping between them. A low reconstruction error says the code is sufficient. It does not say the code is the coordinate you had in mind. Raise the noise to 0.05 and beyond, and once the arms blur into each other, no single number can tell them apart.

What a denoising autoencoder learns

A different way to stop an autoencoder copying is to corrupt its input and ask for the original. Train $g(f(\tilde x))$ to output $x$ where $\tilde x = x + \sigma\varepsilon$.Vincent et al., Extracting and composing robust features with denoising autoencoders, 2008. The optimal reconstruction is the posterior mean $r(\tilde x) = \mathbb{E}[x \mid \tilde x]$, and for Gaussian noise that mean has a closed form, Tweedie’s formula:

$$\mathbb{E}[x \mid \tilde x] = \tilde x + \sigma^2\, \nabla_{\tilde x} \log p_\sigma(\tilde x),$$

where $p_\sigma$ is the density of the noisy data. The reconstruction error $r(\tilde x) - \tilde x$ is therefore a scaled estimate of the score, the gradient of the log density.Alain and Bengio, What regularized auto-encoders learn from the data-generating distribution, 2014. A denoising autoencoder trained at many noise levels is, in effect, the network inside a diffusion model.

A code you can sample from

A plain autoencoder’s code space has no reason to be filled. Decode a random $z$ and you may land anywhere. To generate, we need a code with a known distribution.

The variational autoencoder starts from a latent variable model: a prior $p(z) = \mathcal{N}(0, I)$ and a decoder $p_\theta(x \mid z)$.Kingma and Welling, Auto-encoding variational Bayes, 2014; Rezende, Mohamed and Wierstra, 2014, independently. The likelihood $p_\theta(x) = \int p_\theta(x\mid z)\,p(z)\,dz$ is intractable, so an encoder $q_\phi(z \mid x)$ approximates the posterior. For any $q$, the log likelihood splits exactly:

$$\begin{aligned} \log p_\theta(x) &= \underbrace{\mathbb{E}_{q_\phi(z\mid x)}\!\left[\log p_\theta(x \mid z)\right] - \mathrm{KL}\!\left(q_\phi(z\mid x)\,\Vert\, p(z)\right)}_{\text{ELBO}} \\ &\quad + \mathrm{KL}\!\left(q_\phi(z\mid x)\,\Vert\, p_\theta(z\mid x)\right). \end{aligned}$$

The last term is non-negative, so the first two form a lower bound, the ELBO, and maximising it both raises the likelihood and pulls $q$ towards the true posterior.

Two details make it trainable. The sample $z \sim q_\phi(z\mid x)$ is written as $z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, I)$, the reparameterisation, so gradients flow through $\mu$ and $\sigma$. And for a diagonal Gaussian $q$, the KL term has a closed form:

$$\mathrm{KL}\!\left(\mathcal{N}(\mu, \sigma^2)\,\Vert\,\mathcal{N}(0, 1)\right) = \tfrac12\left(\mu^2 + \sigma^2 - 1 - \log \sigma^2\right).$$

With a Gaussian decoder of variance $s^2$, $-\log p_\theta(x\mid z) = \lVert x - g_\theta(z)\rVert^2 / 2s^2$ plus a constant. Multiply the negative ELBO by $2s^2$ and it becomes

$$\mathcal{L} = \underbrace{\mathbb{E}\,\lVert x - g_\theta(z)\rVert^2}_{D,\ \text{distortion}} + \beta\, \underbrace{\mathbb{E}_x\,\mathrm{KL}\!\left(q_\phi(z\mid x)\,\Vert\, p(z)\right)}_{R,\ \text{rate}}, \qquad \beta = 2s^2 .$$

So the β of the β-VAE is not a hack added to the objective.Higgins et al., β-VAE, 2017. The rate–distortion reading is from Alemi et al., Fixing a broken ELBO, 2018. It is the decoder’s assumed noise level, and it sets the price of information. The rate $R$ is how many nats the code spends describing each point, and it splits in two:

$$R = I_q(x; z) + \mathrm{KL}\!\left(q(z)\,\Vert\, p(z)\right), \qquad q(z) = \mathbb{E}_x\, q_\phi(z \mid x).$$

The first part is information the code genuinely carries about $x$. The second is waste: the mismatch between the prior and the aggregate posterior $q(z)$, the distribution of codes the encoder actually uses. Where $q(z)$ leaves holes that the prior covers, samples from the prior decode to places the data never were.

At β = 0.010 the posteriors are narrow bars laid out in order, the aggregate posterior roughly covers the prior, and the generated rings sit on the data. Drag β down towards $10^{-4}$. The bars shrink to slivers and the rate climbs past five nats: the code becomes nearly deterministic, an ordinary autoencoder with a little noise. A rough rule is that a 1-d code pinning a point to within $\sigma$ of a range of width $w$ costs about $\log(w/\sigma)$ nats.

Now drag β up. The posteriors widen and drift towards zero, the rate falls, and the distortion in the readout rises. Near β = 1 the rate drops to zero. Every posterior has become the prior, the decoder ignores $z$, and every generated ring lands on the same point, the mean of the data. This is posterior collapse. Here we forced it by pricing information too high. With a decoder powerful enough to model $x$ on its own, such as an autoregressive one, it happens even at β = 1, because the decoder can do without the code.Bowman et al., Generating sentences from a continuous space, 2016, which also introduced KL annealing as a remedy. “Free bits” (Kingma et al., 2016) is another. Watch the training path in the bottom panel: at high β it runs left, towards zero rate, and stays there.

A code you can read

The last constraint goes the other way: make the code wider than the input, but insist that only a few of its coordinates are active at once.

The motivation comes from looking inside trained networks. A layer with $d$ neurons can represent far more than $d$ features, provided the features are rarely active together, by giving each feature a direction and letting the directions overlap. This is superposition.Elhage et al., Toy models of superposition, 2022. Its cost is that single neurons respond to several unrelated things, so reading the network neuron by neuron fails. If an activation is a sparse sum $x \approx \sum_i a_i\, d_i$ over $n \gt d$ directions, recovering the $d_i$ and $a_i$ is sparse dictionary learning, an old problem from vision.Olshausen and Field, Emergence of simple-cell receptive field properties by learning a sparse code for natural images, 1996.

A sparse autoencoder learns the dictionary and the encoding together:Bricken et al., Towards monosemanticity, 2023; Cunningham et al., Sparse autoencoders find highly interpretable features in language models, 2023.

$$f(x) = \mathrm{ReLU}\!\left(W_e (x - b_d) + b_e\right), \qquad \hat x = W_d f(x) + b_d,$$ $$\mathcal{L} = \lVert x - \hat x\rVert^2 + \lambda \lVert f(x)\rVert_1, \qquad \lVert W_d[:, j]\rVert = 1 \text{ for every } j.$$

The columns of $W_d$ are the dictionary, one atom per latent. The unit-norm constraint matters: without it, the model could shrink $f$ by any factor $c$ and grow $W_d$ by $c$, and the penalty would fall to nothing at no cost.

The figure works in two dimensions so that everything can be seen. A handful of true features point in different directions in the plane, and each is on, with a random strength, only some of the time. Two dimensions hold five features only because they are rarely on together.

With the starting settings, a dictionary of 8 latents finds most of the features, each with an atom lying along its dashed line, and leaves a couple of latents dead. Lower λ. Reconstruction stays perfect, but the atoms stop tracking the features and more latents fire per point: in two dimensions, any two atoms can rebuild any point, so without a price on activity nothing pulls the dictionary towards the truth. Raise λ and the codes get sparser but the reconstruction suffers, because the L1 penalty charges for the size of an activation as well as its presence. The encoder learns to under-report.

That shrinkage motivated replacing the penalty with a hard constraint: keep only the $k$ largest activations and zero the rest, a TopK autoencoder.Gao et al., Scaling and evaluating sparse autoencoders, 2024. JumpReLU (Rajamanoharan et al., 2024) is a related fix. Set keep only the top k to 1 and every feature is found, nearly exactly. Set it to 2 and recovery breaks down even though the variance left unexplained drops to zero. Two active atoms span the plane, so any dictionary reconstructs perfectly and the objective no longer cares where the atoms point. The general lesson: a dictionary is identifiable only when the code is much sparser than the dimension it lives in. Real sparse autoencoders work with $d$ in the thousands, tens of thousands to millions of latents, and a few dozen active at a time.

The dictionary size has its own failure modes. With fewer latents than features, atoms settle between features and each latent stands for a blend. With many more, latents die, or two atoms split one feature between them. Note also what the readout can do here that a real study cannot: count the features found. Inside a trained model there is no ground truth, and judging whether a learned feature is real is the hardest part of the method.

A minimal sparse autoencoder in PyTorch:

class SAE(nn.Module):
    def __init__(self, d, n_latents, lam=1e-3):
        super().__init__()
        self.W_e = nn.Parameter(torch.randn(n_latents, d) / d**0.5)
        self.b_e = nn.Parameter(torch.zeros(n_latents))
        self.W_d = nn.Parameter(self.W_e.detach().T.clone())   # start tied
        self.b_d = nn.Parameter(torch.zeros(d))
        self.lam = lam

    def forward(self, x):
        f = torch.relu((x - self.b_d) @ self.W_e.T + self.b_e)
        x_hat = f @ self.W_d.T + self.b_d
        loss = ((x - x_hat) ** 2).sum(-1).mean() + self.lam * f.abs().sum(-1).mean()
        return x_hat, f, loss

    @torch.no_grad()
    def renorm(self):                                         # after every optimiser step
        self.W_d /= self.W_d.norm(dim=0, keepdim=True)