Dictionary Learning for NN

The generative picture

Dictionary learning starts from a generative assumption: each observed vector is approximately a sparse linear mix of a shared set of atoms. Collect activations as columns of X: [d, N] (d dimensions, N samples). We posit a dictionary D: [d, m] whose m columns d_1, …, d_m are the atoms, and codes Z: [m, N], so that

X ≈ D Z,   with each column z_i sparse (few nonzeros)

Reading one sample: x_i ≈ Σ_k z_{ki} d_k, a weighted sum of a handful of atoms. The interpretive bet is that the underlying signal really is generated this way — a scene is a few objects, a token’s activation is a few active concepts — even though the raw vector looks dense. When m > d the dictionary is overcomplete: there are more atoms than dimensions, so the columns cannot be orthogonal and the representation of a generic dense vector is non-unique. Sparsity is precisely what selects one meaningful decomposition out of that infinite family.

The sparse-coding objective

We want a dictionary and codes that reconstruct well while keeping every code sparse. Written as a single objective over both unknowns:

min over D, Z   (1/2) ||X - D Z||_F^2  +  λ Σ_i ||z_i||_1
subject to     ||d_k||_2 ≤ 1   for every atom k

The first term is squared reconstruction error (Frobenius norm). The second is a sparsity penalty. Ideally we would penalize ||z_i||_0, the literal count of nonzeros, but the L0 norm is non-convex and combinatorial, so we relax it to the L1 norm ||z_i||_1 = Σ_k |z_{ki}|, whose sharp corners at zero still drive coefficients exactly to zero. The hyperparameter λ trades reconstruction against sparsity: large λ yields fewer, cleaner active atoms but blurrier reconstructions. The norm constraint on each atom is essential — without it the model cheats by scaling atoms up and codes down, making the L1 penalty meaninglessly small while nothing is truly sparse.

Why alternate: a biconvex problem

The joint objective is not convex in D and Z together — the product D Z couples them bilinearly. But it is convex in each one with the other held fixed. That structure invites alternating minimization, the workhorse of every classical dictionary-learning method:

repeat:
  (1) sparse coding:  fix D, solve for Z   (a Lasso per column)
  (2) dict update:    fix Z, update D      (least squares + renorm)

Each phase is a well-understood convex subproblem, so each step cannot increase the loss and the alternation converges to a stationary point — though, because the joint surface is non-convex, only a local one that depends on initialization. Phase 1 asks ‘given these atoms, what is the sparsest explanation of each sample?’ Phase 2 asks ‘given how samples are explained, what atoms reconstruct them best?’ Iterating lets the two co-adapt: atoms migrate toward the directions the data repeatedly uses.

Phase 1: solving for the codes

With D fixed, each column decouples into an independent Lasso problem: min_z (1/2)||x_i - D z||_2^2 + λ||z||_1. Two families dominate. Proximal / ISTA methods take a gradient step on the smooth reconstruction term, then apply the soft-threshold operator S_λ(v) = sign(v) · max(|v| - λ, 0), which shrinks small coefficients to exactly zero — that shrinkage is where sparsity comes from. FISTA adds momentum for faster convergence.

Greedy methods such as Orthogonal Matching Pursuit (OMP) instead target the L0 form directly: repeatedly pick the atom most correlated with the current residual, add it to the active set, re-solve least squares over those atoms, and stop at a chosen nonzero count. OMP gives you an exact sparsity budget (say, at most 8 active atoms); L1/ISTA gives you a soft penalty you tune via λ. Modern SAE variants echo both: Top-K SAEs are the greedy fixed-budget idea, plain L1 SAEs are the relaxed penalty.

Phase 2: updating the dictionary

With the codes Z fixed, updating D is a constrained least-squares problem: min_D ||X - D Z||_F^2 with each column renormalized to unit length. The simplest route is a gradient step, D ← D - η(D Z - X) Z^T, followed by rescaling every atom to ||d_k||_2 = 1. The classical K-SVD algorithm is sharper: it updates one atom at a time. For atom k it looks only at the samples that actually use k, forms the residual those samples would have if k were removed, and sets d_k (and their k-coefficients) to the best rank-one fit of that residual — the top singular vector from an SVD. Because it refits the coefficients simultaneously, K-SVD converges in far fewer outer iterations than gradient updates, at the cost of an SVD per atom. Either way the geometric effect is the same: atoms rotate to point along the directions the data keeps activating together.

A small worked example

Take d = 3, an overcomplete dictionary of m = 4 unit atoms, and a target activation:

d_1=(1,0,0)  d_2=(0,1,0)  d_3=(0,0,1)  d_4=(.58,.58,.58)
x = (0.60, 0.60, 0.02)

A dense least-squares fit would smear weight across all four atoms, including a little of d_3 and d_4, to shave the last hundredth of error. The sparse solution instead reads off the obvious structure: z ≈ (0.60, 0.60, 0, 0) — two active atoms, reconstruction (0.60, 0.60, 0), residual norm about 0.02. The L1 penalty happily accepts that tiny residual to buy two hard zeros. Notice the temptation the overcomplete atom d_4 creates: because it correlates with the x_1, x_2 pattern, a greedy coder could pick it first, and whether the ‘right’ two-atom answer wins depends on atom coherence — the quantity recovery guarantees turn on.

Overcompleteness and superposition

Why insist on m > d for neural networks? Because transformers appear to store more features than they have neurons. Under superposition, a d-dimensional activation space represents a much larger set of sparse, near-orthogonal feature directions by tolerating small interference — a lossy compression the network can afford precisely because only a few features fire on any given token. A complete (m = d) basis cannot untangle that; it can only rotate the same d axes. An overcomplete dictionary with m >> d gives you enough atoms to assign a distinct direction to each underlying feature, and the sparsity constraint matches the assumption that few features are active at once. This is the theoretical bridge: superposition says features live in an overcomplete, sparse code inside the residual stream, and dictionary learning is the tool designed to recover exactly that kind of code from data.

The SAE is dictionary learning as a network

A sparse autoencoder is dictionary learning re-cast as a single-hidden-layer network with an amortized encoder. Instead of solving a Lasso per input at inference time, the SAE trains an encoder to predict the sparse code in one forward pass:

z = ReLU(W_enc x + b_enc)          # amortized sparse coding (phase 1)
x_hat = W_dec z + b_dec           # W_dec columns = dictionary atoms
L = ||x - x_hat||_2^2 + λ ||z||_1

The decoder weight matrix W_dec: [d, m] is the dictionary — its columns are the atoms, conventionally unit-normalized like the classical constraint. Training by gradient descent jointly fits encoder, decoder, and codes, folding the two alternating phases into one differentiable loss. The ReLU plus L1 (or a Top-K gate) plays the role of the soft-threshold in ISTA. So the SAE is not a different idea from dictionary learning; it is the same objective, made cheap to evaluate by learning the sparse-coding map instead of re-optimizing it for every activation.

Feature recovery and identifiability

When does dictionary learning return the true generating atoms rather than some arbitrary basis that happens to reconstruct? This is the identifiability question, and the answer hinges on two conditions. First, the codes must be genuinely sparse — each sample uses few atoms — and there must be enough samples exercising every atom. Second, the atoms must be sufficiently incoherent: the mutual coherence μ = max_{j≠k} |d_j · d_k| must be small, so no two atoms are near-parallel. Classical results show that if the true sparsity per sample is below roughly 1/(2μ), sparse coding recovers the correct support and the dictionary is identifiable up to the unavoidable ambiguities: you can permute the atoms and flip their signs and get the same factorization. Those two symmetries are why recovered features come out unordered and why sign conventions vary. When coherence is high or codes are too dense, recovery degrades — atoms merge, split, or absorb each other, the interpretability analogue of dead and duplicated SAE features.

Complexity and CPU-SLM implications

The costs split cleanly. Amortized inference — running a trained SAE encoder — is a single matrix multiply, O(m d) per activation, cheap enough to run over an entire corpus on a CPU. Classical per-sample sparse coding is heavier: ISTA costs O(m d) per iteration times many iterations, and K-SVD training adds an SVD per atom per epoch. Memory is the real pressure point for small-model tooling: an overcomplete dictionary with m = 16d or 32d atoms is large, and storing codes for a big activation corpus is larger still — though the codes are sparse, so a sparse format pays off. For CPU-bound SLM interpretability the practical recipe is to train the dictionary/SAE offline once, then keep only the encoder and a sparse feature store for analysis. Top-K coding also bounds work per token to a fixed budget, making latency predictable on hardware without a GPU.

Common pitfalls

Several failure modes recur. Dead atoms: columns that never activate, wasting capacity — usually a sign that m is too large or the encoder collapsed; resampling or reinitializing them helps. Shrinkage bias: the L1 penalty not only zeros small coefficients but shrinks the surviving ones, biasing reconstruction low; Top-K and gated variants exist partly to decouple ‘which atoms’ from ‘how much.’ Forgetting to normalize atoms breaks the sparsity penalty entirely, as noted earlier. Reading too much into a single atom: incoherence is never perfect, so an atom can still be mildly polysemantic, and the permutation/sign ambiguity means atom indices carry no intrinsic meaning. Finally, sparsity is a dial, not a truth: too aggressive and you merge distinct features into one atom, too loose and you fail to disentangle superposition at all — tuning λ or K against the reconstruction-versus-sparsity frontier is the core empirical work.