The two-teacher census — population loss landscapes of two-neuron teacher and student networks

Both networks read a planar Gaussian input \(x\sim\mathcal N(0,I_2)\) through two hidden neurons. The teacher network with parameters \((s,\beta)\) and the student network with parameters \((c,\theta)\) have the same shape, \[f_{s,\beta}(x)=\sum_{i=0}^{1} s_i\,\mathrm{ReLU}(\langle w_i,x\rangle)\;[\,+\langle b,x\rangle\,],\qquad w_i=(\cos\beta_i,\sin\beta_i),\] with \(\mathrm{ReLU}(z)=\max(z,0)\) and unit directions \(w_i\). For a neuron with general first-layer weight \(\tilde w_i\) and second-layer weight \(\tilde s_i\), the generator is \(g_i=\tilde s_i\,\tilde w_i\), and its mass is the second-layer weight times the norm of the first-layer weight, \(s_i=\tilde s_i\lVert\tilde w_i\rVert\) — positive or negative — with direction \(w_i=\tilde w_i/\lVert\tilde w_i\rVert\), so \(g_i=s_iw_i\). Plain network: the formula without the bracket. Skip network: with the learned linear skip \(\langle b,x\rangle\). The skip makes the sign of a direction irrelevant (see the paper), so the explorer draws a skip-network generator as a segment through the origin.

Scaling symmetry. ReLU is positively homogeneous, \(\mathrm{ReLU}(\lambda z)=\lambda\,\mathrm{ReLU}(z)\) for \(\lambda>0\), so \((\tilde s_i,\tilde w_i)\mapsto(\tilde s_i/\lambda,\lambda\tilde w_i)\) changes neither the generator nor the network function: the loss depends on a neuron only through its generator, and the critical points of the full parametrisation are those of the (mass, direction) parametrisation used here.

Population loss. The teacher network is fixed and the student network is scored by the mean squared difference of the outputs over the Gaussian input, \[\mathscr L(c,\theta)=\tfrac12\,\mathbb E_{x\sim\mathcal N(0,I_2)}\bigl(f_{c,\theta}(x)-f_{s,\beta}(x)\bigr)^2 .\] A student network with parameters \((c,\theta)\in\mathbb R^2\times\mathbb R^2\) is a critical point if \(\nabla\mathscr L(c,\theta)=0\). A critical point is a local minimum if \(\mathscr L\ge\mathscr L(c,\theta)\) on a neighbourhood — strict if the inequality is strict away from the point, flat otherwise; a local maximum if \(\mathscr L\le\mathscr L(c,\theta)\) on a neighbourhood; a saddle if it is neither. The global minimum is a critical point at which \(\mathscr L\) attains its smallest value; in this setting it always exists and the value is \(0\), because the student network can copy the teacher network. A spurious minimum is a local minimum that is not global.

Each spurious minimum gets a type: collided (co) if the two student neurons share a direction, separated (sep) otherwise; with same-sign (+) or opposite-sign (±) masses \(c_0c_1\gtrless0\). The classification gives, for every teacher network, the complete list of critical points in closed form and the type of each; the table of the explorer lists them.

For a teacher network \((s,\beta)\) write \(\mathsf M(s,\beta)\) for the multiset of types of its spurious minima, e.g. \(\{\mathrm{co}^{\pm}\}\) or \(\{2\times\mathrm{co}^{+},\mathrm{sep}^{\pm}\}\); the global minimum is present for every teacher network and not listed. A census is a set of teacher parameters \((s,\beta)\) with the same \(\mathsf M\).

Rotating, reflecting, swapping and jointly rescaling the teacher network's generators does not change \(\mathsf M\), so a teacher network is a point of a plane: horizontal the gap \(\widehat\beta=\beta_1-\beta_0\) (half a period suffices), vertical the mass ratio \(s_0/s_1\) on a compressed scale (\(0\) at the top and bottom edges: the first mass vanishes; \(\pm\infty\) in the middle: the second mass vanishes; \(\pm1\): equal masses up to sign). The census map colours this plane by \(\mathsf M\); saddles play no role in it. Degenerate teacher networks — coincident directions, a right angle or antipodal directions, a vanishing mass — have lists of critical points of their own, but their minima agree with a neighbouring region, so for simplicity the map does not single them out. Click any region to load one of its teacher networks and see the critical points of the student landscape in the table.

About