Hidden Progress in Deep Learning: SGD Learns Parities Near the Computational Limit

Introduction

In deep learning, performance improvements are frequently observed upon simply scaling up resources (such as data, model size, and training time). While these improvements are often continuous in terms of these resources, some of the most surprising recent advances in the field have been emergent capabilities: at a certain threshold, behavior changes qualitatively and discontinuously. Through a statistical lens, it is well-understood that larger models, trained with more data, can fit more complex and expressive functions. However, far less is known about the analogous computational question: how does the scaling of these resources influence the success of gradient-based optimization?

These phase transitions cannot be explained via statistical capacity alone: they can appear even when the amount of data remains fixed, with only model size or training time increasing. A timely example is the emergence of reasoning and few-shot learning capabilities when scaling up language models (Radford et al., 2019, Brown et al., 2020, Chowdhery et al., 2022, Hoffmann et al., 2022); Srivastava et al. (2022) identify various tasks which language models are only able to solve if they are larger than a critical scale. Power et al. (2022) give examples of discontinuous improvements in population accuracy (“grokking”) when running time increases, while dataset and model sizes remain fixed.

In this work, we analyze the computational aspects of scaling in deep learning, in an elementary synthetic setting which already exhibits discontinuous improvements. Specifically, we consider the supervised learning problem of learning a sparse parity: the label is the parity (XOR) of $k\ll n$ bits in a random length-$n$ binary string. This problem is computationally difficult for a range of algorithms, including gradient-based (Kearns, 1998) and streaming (Kol et al., 2017) algorithms. We focus on analyzing the resource measure of training time, and demonstrate that the loss curves for sparse parities display a phase transition across a variety of architectures and hyperparameters (see Figure 1, left). Strikingly, we observe that SGD finds the sparse subset (and hence, reaches 0 error) with a variety of activation functions and initialization schemes, even with no over-parameterization.

A natural hypothesis to explain SGD’s success in learning parities, with no visible progress in error and loss for most of training, would be that it simply “stumbles in the dark”, performing random search for the unknown target (e.g. via stochastic gradient Langevin dynamics). If that were the case, we might expect to observe a convergence time of $2^{\Omega(n)}$, like a naive search over parameters or subsets of indices. However, Figure 1 (right), already provides some evidence against this “random search” hypothesis: the convergence time adapts to the sparsity parameter $k$, with a scaling of $n^{O(k)}$ on small instances. Notably, such a convergence rate implies that SGD is closer to achieving the optimal computation time among a natural class of algorithms (namely, statistical query algorithms).

Through an extensive empirical analysis of the scaling behavior of a variety of models, as well as theoretical analysis, we give strong evidence against the “stumbling in the dark” viewpoint. Instead, there is a hidden progress measure under which SGD is steadily improving. Furthermore, and perhaps surprisingly, we show that SGD achieves a computational runtime much closer to the optimal SQ lower bound than simply doing (non-sparse) parameter search. More generally, our investigations reveal a number of notable phenomena regarding the dependence of SGD’s performance on resources: we identify phase transitions when varying data, model size, and training time.

It is known from SQ lower bounds that with a constant noise level, gradient descent on any architecture requires at least $n^{\Omega(k)}$ computational steps to learn $k$-sparse $n$-dimensional parities (for background, see Appendix A). We first show a wide variety of positive empirical results, in which neural networks successfully solve the parity problem in a number of iterations which scales near this computational limit:

For all small instances ($n\leq 30,k\leq 4$) of the sparse parity problem, architectures $\mathcal{A}\in\{$2-layer MLPs, TransformersWith a smaller range of hyperparameters., sinusoidal/oscillating neurons, PolyNetsA non-standard architecture introduced in this work; see Section 3 for the definition.$\}$, initializations in $\{$uniform, Gaussian, Bernoulli$\}$, and batch sizes $1\leq B\leq 1024$, SGD on $\mathcal{A}$ solves the $(n,k)$-sparse parity problem (w.p. $\geq 0.2$) within at most $c\cdot n^{\alpha k}$ steps, for small constants $c,\alpha$.

Our empirical results suggest that, in a number of computational steps matching the SQ limit, SGD is able to solve the parity problem and identify the influential coordinates, without an explicit sparse prior. We give a theoretical analysis which validates this claim.

We also present a stronger analysis for an idealized architecture (which we call the disjoint-PolyNet), which allows for any batch size, and captures the phase transitions observed in the error curves.

On disjoint-PolyNets, SGD (with any batch size $B\geq 1$) converges with high probability to a solution with at most $\epsilon$ error on the $(n,k)$-parity problem in at most $n^{O(k)}\cdot\log(1/\epsilon)$ iterations. Continuous-time gradient flow exhibits a phase transition: it spends a $1-o(1)$ fraction of its time before convergence with error $\geq 49\%$.

Our theoretical and empirical results hold in non-overparameterized regimes (including with a width-$1$ sinusoidal neuron), in which no fixed kernel, including the neural tangent kernel (NTK) (Jacot et al., 2018), is sufficiently expressive to fit all sparse parities with a large margin. Thus, our findings comprise an elementary example of combinatorial feature learning: SGD can only successfully converge by learning a low-width sparse representation.

Building upon our core positive results, we provide a wide variety of preliminary experiments, showing sparse parity learning to be a versatile testbed for understanding the challenges and surprises in solving combinatorial problems with neural networks. These include quantities which reveal the continual hidden progress behind uninformative training curves (as predicted by the theory), experiments at small sample sizes which exhibit grokking (Power et al., 2022), as well as an example where greedy layer-wise learning is impossible but end-to-end SGD can learn the layers jointly.

2 Related work

We present the most directly related work on feature learning, and learning parities with neural nets. A broader discussion can be found in Appendix A.3.

Theoretical analysis of gradient descent on neural networks is notoriously hard, due to the non-convex nature of the optimization problem. That said, it has been established that in some settings, the dynamics of GD keep the weights close to their initialization, thus behaving like convex optimization over the Neural Tangent Kernel (see, for example, (Jacot et al., 2018, Allen-Zhu et al., 2019, Du et al., 2018)). In contrast, it has been shown that in various tasks, moving away from the fixed features of the NTK is essential for the success of neural networks trained with GD (for example (Yehudai and Shamir, 2019, Allen-Zhu and Li, 2019, Wei et al., 2019) and the review in (Malach et al., 2021)). These results demonstrate that feature learning is an important part of the GD optimization process. Our work also focuses on a setting where feature learning is essential for the target task. In our theoretical analysis, we show that the initial population gradient encodes the relevant features for the problem. The importance of the first gradient step for feature learning has been recently studied in (Ba et al., 2022).

The problem of learning parities using neural networks has been investigated in prior works from various perspectives. It has been demonstrated that parities are hard for gradient-based algorithms, using similar arguments as in the SQ analysis (Shalev-Shwartz et al., 2017, Abbe and Sandon, 2020). One possible approach for overcoming the computational hardness is to make favorable assumptions on the input distribution. Indeed, recent works show that under various assumptions on the input distribution, neural networks can be efficiently trained to learn parities (XORs) (Daniely and Malach, 2020, Shi et al., 2021, Frei et al., 2022, Malach et al., 2021). In contrast to these results, this work takes the approach of intentionally focusing on a hard benchmark task, without assuming that the distribution has some favorable (namely, non-uniform) structure. This setting allows us to probe the performance of deep learning at a known computational limit. Notably, the work of Andoni et al. (2014) provides analysis for learning polynomials (and in particular, parities) under the uniform distribution. However, their main results require a network of size $n^{O(k)}$ (i.e., extremely overparameterized network), and provides only partial theoretical and empirical evidence for the success of smaller networks. Studying a related subject, some works have shown that neural networks display a spectral bias, learning to fit low-frequency coefficients before high-frequency ones (Rahaman et al., 2019, Cao et al., 2019).

Preliminaries

We provide an expanded discussion of background and related work in Appendix A.

A key fact about parities is that they are orthogonal under the correlation inner product: for $S^{\prime}\subseteq[n]$,

Our main results are presented in the online learning setting, with a stream of i.i.d. batches of examples. At each iteration $t=1,\ldots,T$, a learning algorithm $\mathcal{A}$ receives a batch of $B$ examples $\{(x_{t,i},y_{t,i})\}_{i=1}^{B}$ drawn i.i.d. from $\mathcal{D}_{S}$, then outputs a classifier $\widehat{y}_{t}:\{\pm 1\}^{n}\rightarrow\{\pm 1\}$. We say that $\mathcal{A}$ solves the parity task in $t$ steps (with error $\epsilon$) if

Empirical findings

The central phenomenon of study in this work is the empirical observation that neural networks, with standard initialization and training, can solve the $(n,k)$-parity problem in a number of iterations scaling as $n^{O(k)}$ on small instances. We observed robust positive results for randomly-initialized SGD on the following architectures, indexed by Roman numerals:

2-layer MLPs: ReLU $(\sigma(z)=(z)_{+})$ or polynomial ($\sigma(z)=z^{k}$) activation, in a wide variety of width regimes $r\geq k$. Settings (i), (ii), (iii) (resp. (iv), (v), (vi)) use $r=\{10,100,1000\}$ ReLU (resp. polynomial) activations. We also consider $r=k$ (exceptional settings (*i), (*ii) ), the minimum width for representing a $k$-wise parity for both activations.

1-neuron networks: Next, we consider non-standard activation functions $\sigma$ which allow a one-neuron architecture $f(x;w)=\sigma(w^{\top}x)$ to realize $k$-wise parities. The constructions stem from letting $w^{*}=\sum_{i\in S}e_{i}$, and constructing $\sigma(\cdot)$ to interpolate (the appropriate scaling of) $\frac{k-{w^{*}}^{\top}x}{2}\text{ mod }2$ with a piecewise linear $k$-zigzag activation (vii), or a degree-$k$ polynomial (viii). Going a step further, a single $\infty$-zigzag (ix) or sinusoidal (x) neuron can represent all $k$-wise parities. In settings (xi), (xii), (xiii), (xiv), we remove the second trainable layer (setting $u=1$). We find that wider architectures with these activations also train successfully.

Transformers: There is growing interest in using parity as a benchmark for combinatorial function learning, long-range dependency learning, and length generalization in Transformers (Lu et al., 2021, Edelman et al., 2021, Hahn, 2020, Anil et al., 2022, Liu et al., 2022). Motivated by these recent theoretical and empirical works, we consider a simplified specialization of the Transformer architecture to this sequence classification problem. This is the less-robust setting (*iii); the architecture and optimizer are described in Appendix D.1.3.

PolyNets: Our final setting (xv) is the PolyNet, a slightly modified version of the parity machine architecture. Parity machines have been studied extensively in the statistical mechanics of ML literature (see the related work section) as well as in a line of work on ‘neural cryptography’ (Rosen-Zvi et al., 2002). A parity machine outputs the sign of the product of $k$ linear functions of the input. A PolyNet simply outputs the product itself. Both architectures can clearly realize $k$-sparse parities. The PolyNet architecture was originally motivated by the search for an idealized setting where an end-to-end optimization trajectory analysis is tractable (see Section 4.1); we found in these experiments that this architecture trains very stably and sample-efficiently.

All of the networks listed above were observed to successfully learn sparse parities in a variety of settings. We summarize our findings as follows: for all combinations of $n\in\{10,20,30\}$, $k\in\{2,3,4\}$, batch sizes $B\in\{1,2,4,\ldots,1024\}$, initializations $\{$uniform, Gaussian, Bernoulli$\}$, loss functions $\{$hinge, square, cross entropy$\}$, and architecture configurations $\{\text{(i)},\text{(ii)},\ldots,\text{(xv)}\}$, SGD solved the parity problem (with $100\%$ accuracy, validated on a batch of $2^{13}$ samples) in at least $20\%$ of $25$ random trials, for at least one choice of learning rate $\eta\in\{0.001,0.01,0.1,1\}$. The models converged in $t_{c}\leq c\cdot n^{\alpha k}\leq 10^{5}$ steps, for small architecture-dependent constants $c,\alpha$ (see Appendix C). Figure 1 (left) shows some representative training curves.

Settings (*i) and (*ii), where the MLP just barely represents a $k$-sparse parity, and the Transformer setting (*iii), are less robust to small batch sizes. In these settings, the same positive results as above only held for sufficiently large batch sizes: $B\geq 16$. Also, setting (*iii) used the Adam optimizer (which is standard for Transformers); see Appendix D.1.3 for details.

For almost all of the architectures, we find that that the training curves exhibit phase transitions in terms of running time (and thus, in the online learning setting, dataset size as well): long durations of seemingly no progress, followed by periods of rapid decrease in the validation error. Strikingly, for architectures (v) and (vi), this plateau is absent: the error in the initial phase appears to decrease with a linear slope. See Appendix C.8 for more plots.

2 Random search or hidden progress?

The remainder of this paper seeks to answer the question: “By what mechanism does deep learning solve these emblematic computationally-hard optimization problems?”

A natural hypothesis would be that SGD somehow implicitly performs Monte Carlo random search, “bouncing around” the loss landscape in the absence of a useful gradient signal. Upon closer inspection, several empirical observations clash with this hypothesis:

Scaling of convergence times: Without an explicit sparsity prior in the architecture or initialization, it is unclear how to account for the runtimes observed in experiments, which adapt to the sparsity $k$. The initializations, which certainly do not prefer sparse functionsIndeed, under all standard architectures and initialization, the probability that a random network is $\Omega(1)$-correlated with a sparse parity would be $2^{-\Omega(n)}$, since with that probability $1-o(1)$ of the total influence would be accounted by the $n-k$ irrelevant features., are close to the correct solutions with probability $2^{-\Omega(n)}\ll n^{-k}$.

Sensitivity to initialization, not SGD samples: Running these training setups over multiple stochastic batches from a common initialization, we find that loss curves and convergence times are highly correlated with the architecture’s random initialization, and are quite concentrated conditioned on initialization; see Figure 2 (center).

Elbows in the scaling curves: For larger $n$, the power-law scaling ceases to hold: the exponent worsens (see Figure 2 (right), as well as the discussion in Appendix C.2). This would not be true for random exhaustive search.

Even these observations, which do not probe the internal state of the algorithm, suggest that exhaustive search is an insufficient picture of the training dynamics, and a different mechanism is at play.

Theoretical analyses

We now provide a theoretical account for the success of SGD in solving the $(n,k)$-parity problem. Our main theoretical observation is that, in many cases, the population gradient of the weights at initialization contains enough “information” for solving the parity problem. That is, given an accurate enough estimate of the initial gradient (by e.g. computing the gradient over a large enough batch size), the relevant subset $S$ can be found.

Carefully extending this insight, we obtain an end-to-end convergence result for ReLU-activation MLP networks with a particular symmetric choice of $\pm 1$ initialization, trained with the hinge loss:

This does not capture the full range of settings in which we empirically observe successful convergence. First, it requires a sign vector initialization, while we observe convergence with other random initialization schemes (namely, uniform and Gaussian). Second, it requires the batch size to scale with $n^{\Omega(k)}$In fact, at this batch size, the correct parity indices emerge in a single SGD step., while we also obtain positive results when $B$ is small (even $B=1$). Analogous statements for these cases (as well as other activations and losses) would require Fourier gaps for population gradient functions other than majority; lower bounds on the degree-$(k-1)$ coefficients (“Fourier anti-concentration”) are particularly elusive in the literature, and we leave it as an open challenge to establish them in more general settings. We provide preliminary empirics in Appendix C.1, suggesting that the Fourier gaps in our empirical settings are sufficiently large.Interestingly, we observe that the Fourier gap tends to increase over the course of training. This is not captured by our current theoretical analysis.

We note that in the low-width (non-overparameterized) regimes considered in this work, no fixed kernel (including the neural tangent kernel (Jacot et al., 2018), whose dimensionality is the network’s parameter count) can solve the sparse parity problem. The following is a consequence of results in (Kamath et al., 2020, Malach and Shalev-Shwartz, 2022):

Thus, our low-width results lie outside the NTK regime, which requires far larger models (size $n^{\Omega(k)}$) to express parities. However, we note that better sample complexity bounds are possible in the NTK regime, with an algorithm more similar to standard SGD (see (Telgarsky, 2022) and Appendix A.3).

2 Disjoint-PolyNet: exact trajectory analysis for an idealized architecture

In this section, we present an architecture (a version of PolyNets (xv)) which empirically exhibits similar behavior to MLPs and bypasses the difficulty of analyzing Fourier gaps. The disjoint-PolyNet takes a product over $k$ linear functions of an equal-sizedWe assume for simplicity that $n$ is divisible by $k$. partition $P_{1},\ldots,P_{k}$ of the input coordinates: $f(x;w_{1:k}):=\prod_{i=1}^{k}\langle w_{i},x_{P_{i}}\rangle$. As noted in the Section 1.2, this is equivalent to a tree parity machine, with real-valued (instead of $\pm 1$) outputs.

This architecture also requires us to assume that the set $S$ of size $k$ in the $(n,k)$-parity problem is selected such that exactly one index belongs to each disjoint partition, that is, for all $i\in[k]$, $S\cap P_{i}=1$. We refer to this problem as the $(n,k)$-disjoint parity problem. Note that there are still $(n^{\prime})^{k}=(n/k)^{k}$ different possibilities for set $S$ under this restriction. For fixed $k$, these represent a constant fraction of the $\binom{n}{k}\approx(ne/k)^{k}$ (by Stirling’s approximation) possibilities for $S$ in the general non-disjoint case.

Consider training a disjoint-PolyNet w.r.t. the correlation loss. Without loss of generality, assume that each relevant coordinate in $S$ is the first element $P_{i}$. Then, the population gradient is non-zero only at indices $i\in S$:

Suppose $k\geq 3$. Let $T(\epsilon)$ denote the smallest time at which the error is at most $\epsilon$. Then,

Informally, the network takes much longer to reach slightly-better-than-trivial accuracy than it takes to go from slightly better than trivial to perfect accuracy. Returning to discrete time, we also analyze the trajectory of disjoint-PolyNets trained with online SGD at any batch size, confirming that a neural network can learn $k$-sparse disjoint parities within $n^{O(k)}$ iterations.

Extended versions of these theorems, along with proofs, can be found in Appendix B.3.

Hidden progress: discussion and additional experiments

So far, we have shown that sparse parity learning provides an idealized setting in which neural networks successfully learn sparse combinatorial features, with a mechanism of continual progress hiding behind discontinuous training curves. In this section, we outline preliminary explorations on a broader range of interesting phenomena which arise in this setting. Details are provided in Appendix C, while more systematic investigations are deferred to future work.

The theoretical and (black-box) empirical results suggest that SGD does not learn parities via the memoryless process of random exhaustive search. This suggests the existence of progress measures: scalar quantities which are functions of the training algorithm’s state (i.e. the model weights $w_{t}$) and are predictive of the time to successful convergence. We provide some white-box investigations which further support the hypothesis of hidden progress, by examining the gradual improvement in quantities other than the training loss. In Appendix C.1, we directly plot the Fourier gaps of the population gradient, as a function of $t$, finding that they are large (within a small constant factor of those of majority) in practice. In Figure 3 and Appendix C.3, we examine the weight movement norm $\rho(w_{0:t}):=\left\lVert w_{t}-w_{0}\right\rVert_{\infty}$ to reveal hidden progress, motivated by the fact that $w_{t}-w_{0}$ is a linearized estimate for the initial population gradient.

An interesting consequence of our analysis is that it illuminates scaling behaviors with respect to a third fundamental resource parameter: model size, which we study in terms of network width $r$. If SGD operated by a “random search” mechanism, one would expect width to provide a parallel speedup. Instead, we find that SGD sequentially amplifies progress. The sharp lower tails in Figure 2 (left) imply that running $r$ identical copies of SGD does not give $(1/r)\times$ speedups; more directly, in Appendix C.4 (previewed in Figure 4 (left)), we find that convergence times for sparse parities empirically plateau at large model sizes.

Our main results are presented in the online learning setting (fresh minibatches from $\mathcal{D}_{S}$ at each iteration). While this mitigates the confounding factor of overfitting, it couples the resources of training time and independent samples in a suboptimal way, due to the computational-statistical gap for parity learning. In Appendix C.5, we find empirically that minibatch SGD (with weight decay) can learn sparse parities, even with smaller sample sizes $m\ll n^{k}$. We reliably observe the grokking phenomenon (Power et al., 2022): an initial overfitting phase, then a delayed phase transition in the generalization error; see the two center panels of Figure 4 (right). These results complement and corroborate the findings of Nanda and Lieberum (2022), who analyze the hidden progress of Transformers trained on arithmetic tasks (a setting which also exhibits grokking).

It is a significant challenge (and generally outside the scope of this work) to understand the interactions between network depth and computational/statistical efficiency. In Appendix C.7, we show that learning parities with deeper polynomial-activation MLPs comprises a simple counterexample to the “deep only works if shallow is good” principle of Malach and Shalev-Shwartz (2019): a deep network can get near-perfect accuracy, even when greedy layer-wise training (e.g. (Belilovsky et al., 2019)) cannot beat trivial performance. By providing positive theory and empirics which elude these simplified explanations of SGD, we hope to point the way to a more complete understanding of learning dynamics in the challenging cases where no apparent progress is made for extended periods of time.

Conclusion

This work puts forward sparse parity learning as an elementary test case to explore the puzzling features of the role of computational (as opposed to statistical) resources in deep learning. In particular, we have shown that a variety of neural architectures solve this combinatorial search problem, with a number of computational steps nearly matching the sparsity-dependent SQ lower bound. Furthermore, we have shown that despite abrupt phase transitions in the loss and accuracy curves, SGD works by gradually amplifying the sparse features “under the hood”.

Even in this simple setting, there are several open experimental and theoretical questions. This work largely focuses on the online learning case, which couples training iterations with fresh i.i.d. samples. We believe it would be instructive to investigate parity learning when the three resources of samples, time, and model size are scaled separately. Some very preliminary findings along these lines are presented in Section 3. It is an open problem to extend our theoretical results to the small-batch setting, as well as to the full range of architectures and losses in our experiments. Resolving these questions would require a better understanding of the anti-concentration behavior of Boolean Fourier coefficients, which is much less studied than the analogous concentration phenomena.

Another important follow-up direction is understanding the extent to which these insights extend from parity learning to more complex (including real-world) combinatorial problem settings, as well as the extent to which non-synthetic tasks (in, e.g., natural language processing and program synthesis) embed within them parity-like subtasks of exhaustive combinatorial search. We hope that our results will lead to further progress towards understanding and improving the optimization dynamics behind the recent slew of dramatic empirical successes of deep learning in these types of domains.

This work seeks to contribute to the foundational understanding of computational scaling behaviors in deep learning. Our theoretical and empirical analyses are in a heavily-idealized synthetic problem setting. Hence, we see no direct societal impacts of the results in this study.

We would like to thank Lenka Zdeborová for providing us with references to the statistical physics literature on phase transitions in the learning curves of neural networks, and Matus Telgarsky for bringing to our attention the better sample complexity guarantees of 2-sparse parity learning in the NTK regime. Sham Kakade acknowledges funding from the Office of Naval Research under award N00014-22-1-2377.

References

Appendix

This work leverages the parity problem as a “computationally hard case” for identifying the features $S$ which are relevant to the label. Observe that for any $S^{\prime}\subseteq[n]$, it holds that

That is, a learner who guesses indices $S^{\prime}$ cannot use correlations as feedback to reveal which (or how many) indices in $S^{\prime}$ are correct, unless $S^{\prime}$ is exactly the correct subset. In this sense, for the $(n,k)$-parity problem, the $\binom{n}{k}-1$ wrong answers are indistinguishable from each other. Thus, the structure of this problem forces this form of learner (but not necessarily all learning algorithms) to perform exhaustive search over subsets.

It should be mentioned that the $(n,k)$-parity problem can be solved efficiently by a learning algorithm that has access to examples (i.e., an algorithm that does not operate in the SQ framework). Specifically, this problem can be solved by the Gaussian elimination algorithm. Moreover, it has been shown that the (Stochastic) Gradient Descent algorithm, discussed in the next section, can also be utilized for solving parities, given accurate enough estimates of the gradient and a very particular choice of neural network architecture Abbe and Sandon (2020). That said, when the accuracy of the gradients is not sufficient, GD suffers from the same SQ lower bound mentioned above (i.e., GD is essentially an SQ algorithm Abbe et al. (2021)).

Learning sparse noisy parities, even at a very small noise level (i.e., $o(1)$ or $n^{-\delta}$) is believed to be computationally hard. This was first explicitly conjectured by Alekhnovich (2003), and has been the basis for several cryptographic schemes (e.g., (Ishai et al., 2008, Applebaum et al., 2009, 2010)). For noiseless sparse parities, Kol et al. (2017) show time-space hardness in the setting where $k=\omega(1)$. We present some experiments with noisy parities in Appendix C.6, finding that our empirical results (and theoretical analysis) are robust to $\Theta(1)$ noise.

A.2 Neural networks and standard training

where $\sigma(\cdot)$ is applied entrywise. It is standard to use GD to jointly update the network’s parameters. Our results include positive results about “single neurons”: MLPs with width $r=1$. We note that for our theoretical analysis, when training MLPs with GD, we allow for different learning rate and weight decay schedule for the different layers.

Finally, we will analyze randomized learning algorithms, such as GD with random initialization $\theta_{0}$, whose iterates $\theta_{t}$ (and thus classifiers $\widehat{y}$) are random variables even when the samples are not. A learning algorithm has permutation symmetry if, for all sequences of data $\{(x_{t,i},y_{t,i})\}$, the classifiers $\widehat{y}_{t}\circ\pi$ resulting from feeding $\{(\pi(x_{t,i}),y_{t,i})\}$ to the learner have identical distributions as $\pi$ ranges over all permutations of indices. The neural architectures and initializations (and thus, SGD) considered in this work are permutation-symmetric; for this reason, it is convenient for notation to choose $S=[k]$ as the canonical $(n,k)$-parity learning problem, without loss of generality.

A.3 Additional related work

A line of recent work has focused on understanding the feature learning ability of gradient descent dynamics on neural networks. These analyses go beyond the Neural Tangent Kernel (NTK) regime, where they show a separation between learning with fixed features versus GD on neural networks, for these problems. Several of these works assume structure (often "sparse") in the input data which is useful for the prediction task, and helps avoid computational hardness. In contrast, our work focuses on studying hard problems at their computational limit. Here we discuss the most relevant works in detail:

A line of work (Diakonikolas et al. (2020), Yehudai and Ohad (2020), Frei et al. (2020)) focuses on learning a single non-linearity $y=\sigma(w^{\top}x)$ (typically $\sigma(\cdot)$ is the ReLU or sigmoid) using gradient-based methods. These works obtain polynomial-time convergence guarantees when the distribution satisfies a spread condition. These results do not extend to the Boolean hypercube.

Daniely and Malach (2020) also study the problem of learning sparse parities using neural networks. One key difference from our work is that they assume a modified version of the problem, where the input distribution is not uniform over the hypercube, but instead leaks information about the label. In particular, the distribution ensures that the relevant parity bits always have the same value. Shi et al. (2021) generalize this setting by considering a setting where labels are generated based on certain class specific patterns and the data itself is generated using these patterns with some extra background patterns. This also embeds information in the data itself regarding the label, unlike our setting, where the labels are uncorrelated with the input features. Under these structural assumption, the papers study how GD on a two-layer network can learn useful features in polynomial time. Both these analysis also exploits the first gradient step to find useful features. Shi et al. (2021) additionally require a second step to refine the features.

Ba et al. (2022) show how the first gradient step is important for feature leaning. In particular, they show that first update is essentially rank-1 and aligns with the linear component of the underlying function. The functions we consider (parity) do not have a linear component.

Abbe et al. (2022) define a notion of initial alignment between the network at initialization and the target function and show that it is essential to get polynomial time learnability with noisy gradients on a fully connected network. Our MLP results also exploit the correlation between the gradient and the label to ensure that the gradient update gives us signal.

Frei et al. (2022) also study learnability of a parity-like function with $k=2$ under noisy labels. The paper analyzes early stopping GD for learning the underlying labeling function. Our setup is quite different from theirs and can handle $k>2$.

In concurrent work, Damian et al. (2022) consider the problem of learning polynomials which depend on few relevant directions using gradient descent on a two-layer network. They assume that the distributional Hessian of the ground truth function spans exactly the subspace of the relevant direction. Using this, they show that gradient descent can learn the relevant subspace with sample complexity scaling as $d^{2}$ and not $d^{p}$ where $p$ is the degree of the underlying polynomial as long as the number of relevant directions is much less than $d$. Their proof technique is similar to our two-layer MLP result which also exploits correlation in the first gradient step. However, for our setting, the distributional Hessian has rank 0 and does not satisfy their assumptions.

An extensive body of work originating in the statistical physics community has studied phase transitions in the learning curves of neural networks (Gardner and Derrida, 1989, Watkin et al., 1993, Engel and Van den Broeck, 2001). These works typically focus on student-teacher learning in the “thermodynamic limit”, in which the number of training examples is $\alpha$ times larger than the input dimension and both are taken to infinity. One of the classic toy architectures analyzed in this literature is the parity machine (Mitchison and Durbin, 1989, Hansel et al., 1992, Opper, 1994, Kabashima, 1994, Simonetti and Caticha, 1996). In our work, we introduce PolyNets, a variant of parity machines in which the output is real-valued rather than $\pm 1$; and we theoretically analyze disjoint-Polynets, which are the real-output analogue of the oft-considered parity machines with tree architecture. While much of the statistical mechanics of ML literature focuses on an idealized training limit in which the weights reach a Gibbs distribution equilibrium, there is a strand of the literature that aims to characterize the trajectory of SGD training in the high-dimensional limit with constant-sized sets of ordinary differential equations (Saad and Solla, 1995b, a, Goldt et al., 2019). These papers discuss cases, including problems that share aspects with 2-sparse parities Refinetti et al. (2021), where the network gets stuck in (and then escapes from) a plateau of suboptimal generalization error. Recently, Arous et al. (2021) studied (for rank-one parameter estimation problems) the relative amount of time spent by SGD in an initial high-error “search” phase versus a final “descent” phase, which is reminiscent of the framing of Theorem 6. However, to our knowledge prior work has not shown $k$-sparse parities can be learned with a number of iterations that nearly matches known lower bounds, nor has it specifically studied phase transitions in $k$-sparse parity learning during gradient descent.

Another relevant line of work studies learning the parity problem using the neural tangent kernel (NTK) (Jacot et al., 2018). Namely, in some settings, when the network’s weights stay close to their initialization throughout the training, SGD converges to a solution that is given by a linear function over the initial features of the NTK. As shown in Theorem 5, learning parities over a fixed set of features requires the size of the model to be $\Omega(n^{k})$. In contrast, the model size (number of hidden neurons) considered in this paper does not depend at all on the input dimension $n$. Nevertheless, the NTK analysis does give better sample complexity guarantees than the ones presented in this work, with a somewhat more natural version of SGD. For example, the work of Ji and Telgarsky (2019) demonstrates learning 2-sparse parities using NTK analysis with a sample complexity of $O(n^{2})$, which matches the sample complexity lower bound for learning this problem with NTK (see Wei et al. (2019)). Concurrent work by Telgarsky (2022) shows that this sample complexity can be improved to $O(n)$ once the optimization leaves the NTK regime. However, this analysis is given for networks of size $O(n^{n})$, much larger than the networks considered in this paper. We refer the reader to Table 1 in Telgarsky (2022) for a complete comparison of the sample-complexity, run-time and model-size bounds achieved by different works studying 2-sparse parities.

Appendix B Proofs

For some even number $r$, consider a ReLU MLP of size $r$:

For all $1\leq i\leq r/2$, randomly initialize

Fix some $k$, and assume that $n\geq 2k^{2}$. Then, for every $S\subseteq[n]$ s.t. $|S|=k$ it holds that:

for some universal constants $c_{1},c_{2}$. More precisely,

Therefore, by the previous lemma we get that for every even $k$ the following holds:

Fix some $k$ and assume that $n\geq 4k$. Then, Majority has a $\gamma_{k}$-Fourier gap at $S$ of size $k$ with $\gamma_{k}=0.03(n-1)^{-\frac{k-1}{2}}$.

First we establish a simple relationship between $|\xi_{k-1}|$ and $|\xi_{k+1}|$.

Here, the first equation follows from Lemma 1, and the second equation follows by simple algebra using the following equality: $\binom{m}{r}=\frac{m-r+1}{r}\binom{m}{r-1}$.

Here, the first equality holds from above, the second by Lemma 1, the third inequality holds from standard approximations of the binomial coefficients, and the last inequality follows from the following inequalities: $n-2k+1\geq(n-1)/2$ (by assumption on $n$) and $\frac{\sqrt{2\pi}k^{k/2}}{(k-1)e^{k+2}}>0.03$ (by standard calculus). This gives us the desired result.

Assume that $k$ is even and that $n\geq 2(k+1)^{2}$. Then, the following hold:

Let $\tau>0$ be some tolerance parameter, fix $\epsilon\in(0,1)$ and let $\eta=\frac{1}{k|\xi_{k-1}|}$. Assume that $k$ is an even number. Fix some $w_{1},\dots,w_{k}\in\{\pm 1\}^{n}$, $b_{1},\dots,b_{r}\in(-1,1)$ and $u_{1},\dots,u_{k}\in\{\pm 1\}$. Let $\widehat{w}_{i}=-\eta\widehat{g}_{i}$ and $\widehat{b}_{i}=b_{i}-\eta\widehat{\gamma}_{i}$ s.t. $\|\widehat{g}_{i}-g_{i}\|_{\infty}\leq\tau$ and $\|\widehat{\gamma}_{i}-\gamma_{i}\|\leq\tau$. Assume the following holds:

Additionally, for all $i$ and all $x$ it holds that $|\sigma(\widehat{w}_{i}\cdot x+\widehat{b}_{i})|\leq n+1$.

Claim 1: For all $i$ and for all $j\in[k]$ it holds that $\left|\widehat{w}_{i,j}-\frac{1}{2k}\right|\leq\frac{\tau}{|\xi_{k-1}|}$.

Proof: First, observe that by the assumption it holds that for all $i,j\in[k]$,

Claim 2: For all $i$ and for all $j>k$ it holds that $|\widehat{w}_{i,j}|\leq\frac{|\xi_{k+1}|+2\tau}{2k|\xi_{k-1}|}$

Claim 3: For all $i$ it holds that $|\widehat{b}_{i}-b_{i}|\leq\frac{\tau}{k|\xi_{k-1}|}$

Claim 4: Fix $\delta>0$. Let $h_{i}$ be a function s.t. $h_{i}(x)=\sigma(\frac{1}{2k}\sum_{j=1}^{k}x_{j}+b_{i})$ and $\widehat{h}_{i}$ a function s.t. $\widehat{h}_{i}(x)=\sigma(\widehat{w}_{i,j}\cdot x+\widehat{b}_{i})$. Then, if $\tau\leq\frac{\delta}{2}\frac{k|\xi_{k-1}|}{\sqrt{2n\log(2k/\epsilon)}}$ and $n\geq\frac{32c_{2}^{2}\log(2k/\epsilon)}{c_{1}^{2}\delta^{2}}$, the following holds:

where in the last inequality we use Eq. (6). So, choosing $\tau\leq\frac{\delta}{2}\frac{k|\xi_{k-1}|}{\sqrt{2n\log(2k/\epsilon)}}$ and $n\geq\frac{32c_{2}^{2}\log(2k/\epsilon)}{c_{1}^{2}\delta^{2}}$ gives the required.

Claim 5: Let $h_{1},\dots,h_{k}$ be the functions defined in the previous claim. Then, there exists weights $u^{*}$ with $\|u^{*}\|_{\infty}\leq 8k$ s.t. for $f^{*}(x)=\sum_{i=1}^{k}u^{*}_{i}h_{i}(x)$ it holds that $f^{*}(x)=2\chi_{[k]}(x)$ for all $x\in\{\pm 1\}^{n}$.

Proof: For $i\leq k-2$ define $u^{*}_{i}=8k(-1)^{i+1}$ and $u^{*}_{k-1}=6k$, $u^{*}_{k}=-2k$.

Proof of Lemma 4: Choose $\widehat{u}=u^{*}$. Using Claim 4 and the union bound, w.p. $1-\epsilon$ over $x\sim\{\pm 1\}^{n}$ it holds that for all $i\in[k]$, $|h_{i}(x)-\widehat{h}_{i}(x)|\leq\delta$. Therefore, w.p. $\geq 1-\epsilon$

so, choosing $\delta=\frac{1}{8k^{2}}$ we get that, w.p. at least $1-\epsilon$ over the choice of $x$ it holds that $f(x)\chi_{[k]}(x)\geq 1$. Additionally, for every $x$ it holds that

Assume we randomly initialize an MLP using the unbiased initialization defined previously. Consider the following update:

Additionally, it holds that $\|\sigma(W^{(1)}\cdot x+b^{(1)})\|_{\infty}\leq n+1$.

Claim: with probability at least $1-\delta$,

Proof: Fix some $i,j$ and note that by Hoeffding’s inequality,

Using the union bound, with probability at least $1-\delta$, there exists a set of $k$ neurons satisfying the conditions of Lemma 4, and therefore the required follows from the Lemma. ∎

We use the following well-known result on convergence of SGD (see for example Shalev-Shwartz and Ben-David (2014)):

Let $k$ be an even number. Assume we randomly initialize an MLP using the unbiased initialization defined previously. Fix $\epsilon\in(0,1/2)$ and let $T\geq\frac{2^{9}k^{3}rn^{2}}{\epsilon^{2}},r\geq k\cdot 2^{k}\log(8k/\epsilon),B\geq c_{1}^{-1}2^{8}k^{7/6}n\binom{n}{k-1}\log(128k^{3}n/\epsilon)\log(32nr/\epsilon),n\geq\frac{2^{11}k^{4}c_{2}^{2}\log(128k^{3}n/\epsilon)}{c_{1}^{2}}$. Choose the following learning rate and weight decay schedule:

At the first step, use $\eta_{0}=\frac{1}{k|\xi_{k-1}|}$, $\lambda_{0}=1$ for all weights.

After the first step, use $\eta_{t}=0$ for the first layers weights and biases and $\eta_{t}=\frac{4k^{1.5}}{n\sqrt{r(T-1)}}$ for the second layer weights, with $\lambda_{t}=0$ for both layers.

Then, the following holds, with expectation over the randomness of the initialization and the sampling of the batches:

Now, we will show that w.h.p. there exists $u^{*}$ with good loss. Let $\epsilon^{\prime}=\frac{\epsilon}{64k^{2}n},\delta^{\prime}=\frac{\epsilon}{8}$. Denote $\tau=\frac{|\xi_{k-1}|}{16k\sqrt{2n\log(128k^{3}n/\epsilon)}}=\frac{|\xi_{k-1}|}{16k\sqrt{2n\log(2k/\epsilon^{\prime})}}$. Observe that $r\geq k\cdot 2^{k}\log(k/\delta^{\prime})$, and using the fact that $|\xi_{k-1}|\geq c_{1}(k-1)^{-1/3}\binom{n}{k-1}^{-1}$ we get

and additionally $n\geq\frac{2^{11}k^{4}c_{2}^{2}\log(2k/\epsilon^{\prime})}{c_{1}^{2}}$.

So, we can apply Theorem 8 with $M=8k\sqrt{k}$ and $\rho=2\sqrt{r}n$ and get that, w.p. $1-\epsilon/4$ over the initialization and the first step, it holds that

B.2 Recoverability of the parity indices from Fourier gaps

Given a network architecture where some neuron has a $\gamma$-Fourier gap with respect to the target subset $S$, we quantify how the indices in $S$ can be determined by observing an estimate of the population gradient for a general activation function $\sigma$ and $w_{t}$:

Then, for every $w$ such that $\sigma^{\prime}(w^{\top}x)$ has a $\gamma$-Fourier gap at $S$, the $k$ indices at which $g(w)$ has the largest absolute values are exactly the indices in $S$.

Let $h(x):=\sigma^{\prime}(w^{\top}x)$. We compute the population gradient, we we call $\bar{g}(w)$:

where the inequality in the final $i\not\in S$ case is due to the Fourier gap property. Then, it holds that for all $i\in S$ we have $|g_{i}|>\gamma/2$ and for all $i\notin S$ we have $|g_{i}|<\gamma/2$. Thus, the largest entries of the estimate $g(w)$ occur at the indices in $S$, as claimed. ∎

B.3 Global convergence for disjoint-PolyNets

In this section we will develop theory for disjoint-PolyNets trained with correlation loss, as illustrated in Figure 5. Section B.3.1 will consider optimization with gradient flow, and section B.4 will consider optimization with SGD at any batch size $B\geq 1$.

In gradient flow, the relevant weights evolve according to the following differential equations:

Suppose disjoint-PolyNet for $k>2$ is initialized such that $\prod_{i}v_{i}(0)>0$, and optimized with gradient flow. Let $\bar{v}_{a}:=\frac{1}{k}\sum_{i=1}^{k}(v_{i}(0))^{2}$, and $\bar{v}_{g}:=\left(\prod_{i=1}^{k}(v_{i}(0))^{2}\right)^{1/k}$. For any $b\geq 0$ and $i\in[k]$, let $T_{i}(b):=\arg\sup_{t\geq 0}(|v_{i}(t)|\leq b)$. Then

Let $T_{i}(\infty):=\arg\sup_{t\geq 0}(|v_{i}(t)|<\infty)$. Then

First, observe that the product of the relevant weights is non-decreasing during gradient flow.

Thus, $\prod_{i}v_{i}(0)>0$ implies that $\prod_{i}v_{i}(t)>0$ for all $t$.

In other words, the squares of the relevant weights each follow the same trajectory, shifted according to their initializations. Let $q(t):=({v_{i}(t)})^{2}-({v_{i}(0)})^{2}$, for any $i$. This quantity evolves as follows:

Since $q(t)$ is strictly increasing, its inverse $q^{-1}$ is well-defined, and we can use the inverse function theorem to characterize $q^{-1}$ for all $t\geq 0$:

We can upper- and lower-bound the integrand by applying Maclaurin’s inequality (see page 52 in Hardy et al. (1952)):

The amount of time it takes for $q$ to reach a value of $c$ is thus:

Hence, for any $b\geq 0$, for each $i$, $|v_{i}(t)|\leq b$ as long as

Meanwhile, the amount of time after $q^{-1}(c)$ it takes for $q$ to explode to infinity is

Substituting $c=b^{2}-v_{i}(0)^{2}$, we obtain that the amount of time it takes for $|v_{i}|$ to grow from $b$ to $\infty$ is

We can upper- and lower-bound $\dot{q}$ by applying Maclaurin’s inequality (see page 52 in Hardy et al. (1952)):

When $k=2$, solving the LHS and RHS differential inequalities yields:

From the lower bound on $q(t)$, we can infer that the relevant weights all explode to infinity by the following time:

From the upper bound, we can infer that for any $c>0$, it is the case that $q(t)\leq c$ so long as

Hence, for each $i$, $|v_{i}(t)|\leq b$ for all

Now we analyze the relationship between the relevant weights and the accuracy of the disjoint-PolyNet.

Then define the error of $f$ with parameters $w_{1:k}$ as

Let $w$ be any setting of the weights of a disjoint-PolyNet such that $\prod_{i}v_{i}>0$. For ease of notation, let $u_{i}:=w_{i,2:n^{\prime}}$ be the irrelevant portion of $w_{i}$. There is a constant $c$ such that

where $\operatorname*{erf}$ is the Gauss error function $\operatorname*{erf}(y):=\frac{2}{\sqrt{\pi}}\int_{0}^{y}e^{-\tau^{2}}\,d\tau$.

Let $z_{i}=x_{(i-1)n^{\prime}+1}$ be the $i$th relevant coordinate of $x$, and let $z_{i}^{-}=x_{(i-1)n^{\prime}+2:in^{\prime}}$ be the irrelevant coordinates in $P_{i}$.

Then we have that the error of $f$ with parameters $w_{1:k}$ is

Line 10 follows because $u_{i}^{\top}z_{i}^{-}$ and $z_{i}$ are independent. In line 11 we use that $u_{i}^{\top}z_{i}^{-}$ is symmetric about 0. Finally, line 12 uses the assumption that $\prod_{i=1}^{k}v_{i}>0$.

We can bound $\Pr_{x}\left[\prod_{i=1}^{k}w_{i}^{\top}x_{P_{i}}=0\right]$ using Hoeffding’s inequality:

The indicator random variables $\operatorname*{1{\hskip-2.5pt}\hbox{I}}[u_{i}^{\top}z_{i}^{-}>|v_{i}|]$