Learning sparse parities with SGD

Theorems (disjoint-PolyNet). Split $[n]$ into $k$ blocks of size $n'=n/k$, take $f(x;w)=\prod_{i\le k}\langle w_i,x_{P_i}\rangle$, the correlation loss, and one relevant coordinate per block. (i) Plateau: for $k\ge3$, on the event that the relevant initial weights have positive product (for $\pm1$, and with high probability for Gaussian initialization), with $T(\alpha)$ the first time the error is at most $\alpha$, $$\frac{T(1/2-\gamma)}{T(0)}=1-\tilde O\big((n')^{1-k/2}\gamma^{2/k-1}\big),$$ so gradient flow spends all but a vanishing fraction of its time with error above $49\%$. (ii) SGD at any batch size: from $\pm1$ initialization with an adaptive step, the error falls below $\varepsilon$ within $O\big(\log(nT/\delta)\log(k/\varepsilon)(3n'-2)^{2k-1}\big)$ steps, with probability $1/2$ over the initialization and $1-\delta$ over the samples.

The population gradient of one neuron is a vector of Fourier coefficients, and a Fourier gap reveals $S$. For a neuron $\sigma(w^\top x)$ and the correlation loss, coordinate $i$ of the population gradient is $-\mathbb E[\chi_S(x)x_i\,\sigma'(w^\top x)]$, and since $x_i^2=1$ multiplying by $x_i$ removes $i$ from $S$ or adds it: $$\bar g_i=-\hat h(S\setminus\{i\})\ \ (i\in S),\qquad \bar g_i=-\hat h(S\cup\{i\})\ \ (i\notin S),\qquad h(x)=\sigma'(w^\top x).$$ If $h$ has a $\gamma$-Fourier gap at $S$, any estimate within $\gamma/2$ in sup norm has its $k$ largest entries exactly on $S$, and Hoeffding with a union bound gives one from $O(\log n/\gamma^2)$ samples. The question becomes how large the gap is.

For a ReLU with sign weights, the gap comes from the explicit Fourier spectrum of majority. With $w\in\{\pm1\}^n$ and $|b|\lt1$, the ReLU derivative is a threshold, $\mathbf 1\{w^\top x+b>0\}=\frac12+\frac12\mathrm{Maj}(w\odot x)$ for odd $n$, whose coefficients are known in closed form, depend only on $|T|$, vanish for even $|T|$, and satisfy $$\frac{|\xi_{k+1}|}{|\xi_{k-1}|}=\frac{k-1}{n-k},\qquad |\xi_{k-1}|-|\xi_{k+1}|\ \ge\ 0.03\,(n-1)^{-(k-1)/2},$$ with $\xi_m=\widehat{\mathrm{Maj}}([m])$. So relevant gradients exceed irrelevant ones by a factor of about $n/k$, and the gap is of order $n^{-(k-1)/2}$, the per-sample signal SGD must accumulate. For even $k$, $k\pm1$ are odd and the gap is nonzero; for odd $k$ an unbiased threshold has no signal at all, which is why biases are needed in practice. (My brute-force check at $n=15$, $k=4$ gives gradient $0.0081$ on $S$ and $0.0022$ off it, ratio $3/11$.)

One large-batch step aligns neurons with $S$, and a convex fit of the second layer gives the parity. Neurons come in pairs with opposite output weights, so $f(x;\theta_0)\equiv0$ and the hinge loss has derivative $-y$; first-layer weights are random signs and biases sit on a grid. Weight decay $\lambda_0=1$ removes the initialization, so each neuron's new weights are its rescaled batch gradient, which by Step 2 points at $S$: a neuron whose initial signs on $S$ agree (probability $2^{-k+1}$) gets weights about $\frac1{2k}$ on $S$ and at most $\Delta\approx\frac1{2n}+\frac{\tau}{k|\xi_{k-1}|}$ off it, with $\tau$ the batch error. The off-$S$ errors add up inside the neuron with spread $\Delta\sqrt n$, which must stay below about $1/(8k^2)$ because the second layer adds $k$ such neurons with weights up to $8k$, so $\tau\lesssim|\xi_{k-1}|/\sqrt n$ and $B\approx n/\xi_{k-1}^2\approx n^k$ up to logs. An aligned neuron then computes a ramp $\sigma(\frac1{2k}\sum_{j\in S}x_j+b_i)$; with width $r\gtrsim k2^k\log(k/\delta)$ every bias level has one, and since $\chi_S$ is a function of $\sum_{j\in S}x_j$, alternating second-layer weights of size at most $8k$ interpolate $2\chi_S$. With the first layer frozen the second-layer problem is convex, and standard SGD analysis gives loss $\varepsilon$ in $O(k^3r^2n/\varepsilon^2)$ steps.

Why the small networks must be learning features. The paper cites the kernel bound from earlier work; for the hinge loss Bessel's inequality already gives it (this short argument is not in the paper). The $\binom nk$ parities are orthonormal, so $\sum_S\|\mathbb E[\chi_S\Psi]\|^2\le\mathbb E\|\Psi\|^2\le1$, some $S$ has $\|\mathbb E[\chi_S\Psi]\|\le\binom nk^{-1/2}$, and then $\mathbb E\,\ell\ge1-R\binom nk^{-1/2}\gt1-\varepsilon$ once $R^2\lt\varepsilon^2\binom nk$. The tangent kernel of a width-$2^{O(k)}$ network is such a small embedding. In the other direction the SQ bound $n^{\Omega(k)}$ applies to SGD with noisy gradients, so $n^{O(k)}$ is the right order.

On a disjoint-PolyNet, gradient flow is an explosion ODE, and the loss changes only at the end. Here the population gradient vanishes off the relevant coordinates, so irrelevant weights $u_i$ never move, and the relevant weights drive each other: $$\dot v_i=\prod_{j\ne i}v_j,\qquad \frac{d}{dt}v_i^2=2\prod_{j}|v_j|\ \text{ for every }i.$$ All $v_i^2$ move together, and Maclaurin's inequality sandwiches the common increment between two ODEs $\dot q=2(q+c)^{k/2}$, which for $k\ge3$ blow up in finite time: going from $|v|=b$ to $\infty$ takes at most $b^{2-k}/(k-2)$. The loss lags because the error depends on $v/\|u\|$, a parity of independent sign events, one per block, $$\mathrm{err}\ \ge\ \frac12-\frac12\prod_{i\le k}\mathrm{erf}\Big(\frac{|v_i|}{\sqrt2\,\|u_i\|_2}\Big)-\text{(Berry-Esseen error)},$$ so with $\|u_i\|_2\approx\sqrt{n'}$ it stays above $\frac12-\gamma$ until $|v_i|\approx\gamma^{1/k}\sqrt{n'}$, and the blow-up time from there gives the plateau fraction. Progress in $v$ is steady throughout, but $v/\|u\|$ stays near zero for almost the entire run.

SGD on the disjoint-PolyNet tracks the flow with a normalized step. The stochastic gradient differs from the population one by a martingale increment bounded by $\prod_{l\ne i}\|w_l\|_1$, so normalizing the step by that product makes each increment at most $1/\sqrt{2T\log(2nT/\delta)}$. Azuma-Hoeffding with a union bound keeps the irrelevant weights within $1/2$ of their $\pm1$ start, and on the positive-product event an induction shows each relevant weight keeps its sign and gains at least $(3n'-2)^{-(k-1)}/\sqrt{8T\log(2nT/\delta)}$ per step, because the normalized drift is a product of ratios $|v_j|/\|w_j\|_1\ge1/(3n'-2)$. The error is at most $\varepsilon$ once $|v_i|\approx\sqrt{n'\log(k/\varepsilon)}$, which takes $T\approx(3n')^{2k-1}\log(k/\varepsilon)$ steps at any batch size.

The experiments test random search against hidden progress. Success is robust across $n\le30$, $k\le4$, batch sizes from 1 to 1024, three initializations, three losses and fifteen architectures, from ReLU MLPs to single sinusoidal neurons and a small Transformer (trained with Adam), reaching full accuracy in at least 20% of runs within $c\,n^{\alpha k}\le10^5$ steps. The timing rules out random search: memoryless search would succeed at a geometric time, most likely at $t=0$, but over $10^6$ PolyNet runs no early successes appear; times are concentrated given the initialization; widening an MLP up to $3\times10^6$ neurons gives no $r$-fold speedup, only a thinner upper tail; and the exponent worsens with $n$. The hidden progress can be seen directly: $\|w_t-w_0\|_\infty$ separates training on parity from training on noise long before the loss moves, and measured Fourier gaps under other initializations and activations are comparable to majority's and grow during training, which no theorem covers. With a finite sample and weight decay the networks grok, and a deep quadratic network learns a 4-parity that greedy layer-wise training provably cannot.