SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Benjamin L. Edelman<sup>1</sup> Surbhi Goel<sup>2</sup> Sham Kakade<sup>1</sup> Eran Malach<sup>3</sup> Cyril Zhang<sup>4</sup>
$^{1}$ Harvard University  $^{2}$ University of Pennsylvania  $^{3}$ Hebrew University of Jerusalem  $^{4}$ Microsoft Research NYC
bedelman@g.harvard.edu, surbhig@cis.upenn.edu, sham@seas.harvard.edu, eran.malach@mail.huji.ac.il, cyrilzhang@microsoft.com
In modern deep learning, algorithmic choices (such as width, depth, and learning rate) are known to modulate nuanced resource tradeoffs. This work investigates how these complexities necessarily arise for feature learning in the presence of computational-statistical gaps. We begin by considering offline sparse parity learning, a supervised classification problem which admits a statistical query lower bound for gradient-based training of a multilayer perceptron. This lower bound can be interpreted as a multi-resource tradeoff frontier: successful learning can only occur if one is sufficiently rich (large model), knowledgeable (large dataset), patient (many training iterations), or lucky (many random guesses). We show, theoretically and experimentally, that sparse initialization and increasing network width yield significant improvements in sample efficiency in this setting. Here, width plays the role of parallel search: it amplifies the probability of finding "lottery ticket" neurons, which learn sparse features more sample-efficiently. Finally, we show that the synthetic sparse parity task can be useful as a proxy for real problems requiring axis-aligned feature learning. We demonstrate improved sample efficiency on tabular classification benchmarks by using wide, sparsely-initialized MLP models; these networks sometimes outperform tuned random forests.
Algorithm design in deep learning can appear to be more like "hacking" than an engineering practice. Numerous architectural choices and training heuristics can affect various performance criteria and resource costs in unpredictable ways. Moreover, it is understood that these multifarious hyperparameters all interact with each other; as a result, the task of finding the "best" deep learning algorithm for a particular scenario is foremost empirically-driven. When this delicate balance of considerations is achieved (i.e. when deep learning works well), learning is enabled by phenomena which cannot be explained by statistics or optimization in isolation. It is natural to ask: is this heterogeneity of methods and mechanisms necessary?
This work studies a single synthetic binary classification task in which the above complications are recognizable, and, in fact, provable. This is the problem of offline (i.e. small-sample) sparse parity learning: identify a  $k$ -way multiplicative interaction between  $n$  Boolean variables, given  $m$  random examples. We begin by interpreting the standard statistical query lower bound for this problem as a multi-resource tradeoff frontier for deep learning, balancing between the heterogeneous resources of dataset size, network size, number of iterations, and success probability. We show that in different regimes of simultaneous resource constraints (data, parallel computation, sequential computation,
Figure 1: "More data or larger network?" Effects of jointly scaling the dataset and model sizes, for 2-layer MLPs trained to learn sparse parities. Left: A success frontier, where the computational challenge of feature learning can be surmounted by increasing sample size or model size. Each point shows the percentage of successful training runs (out of 50 trials) for each (dataset size  $m$ , width  $r$ ); (✓) denotes  $100\%$  successful trials. Right: 10 example training curves (gray = training error; colored = test error) for each of the boxed  $(m, r)$  pairs. See main text and Section 4 for further details.
and random trials), the standard algorithmic choices in deep learning can succeed by diverse and entangled mechanisms. Specifically, our contributions are as follows:
Multi-resource lower and upper bounds. We formulate a "data  $\times$  width  $\times$  time  $\times$  luck" lower bound for offline sparse parity learning with feedforward neural nets (Theorem 3). This barrier arises from the classic statistical query (SQ) lower bound for this problem. We show that under different resource constraints, the tractability of learning can be "bought" with varied mixtures of these resources. In particular, in Theorems 4 and 5 we prove that by tuning the width and initialization scheme, we can populate this frontier with a spectrum of successful models ranging from narrow networks requiring many training samples to sample-efficient networks requiring many neurons, as summarized by the following informal theorem statement:
Informal Theorem 1. Consider the problem of learning  $(n, k)$ -parities from  $m$  i.i.d. samples with a 2-layer width- $r$  ReLU MLP, whose first-layer neurons are initialized with sparsity  $s$ . After  $O_{n}(1)$  steps of gradient descent, the  $k$  relevant coordinates are identified with probability 0.99 when (1)  $s > \Omega(k)$ ,  $r = \Theta((n/s)^k)$  and  $m = \Theta(n^2(s/k)^{k-1})$ , and when (2)  $s < k$ ,  $r = \Theta((n/k)^s)$  and  $m = \Theta((n/k)^{k-s-1})$ .
Intuitively, this analysis reveals a feature learning mechanism by which overparameterization (i.e. large network width) plays a role of parallel search over randomized subnetworks. Each individual hidden-layer neuron has its own sample complexity for identifying the relevant coordinates, based on its Fourier gap (Barak et al., 2022) at initialization. Trained with parallel gradient updates, the full network implicitly acts as an ensemble model over these neurons, whose overall sample complexity is determined by the "winning lottery tickets" (Frankle and Carbin, 2018) (i.e. the lucky neurons initialized to have the lowest sample complexity). This departs significantly from the neural tangent kernel (Jacot et al., 2018) regime of function approximation with wide networks, in which overparameterization removes data-dependent feature selection (rather than parallelizing it across subnetworks).
Empirical study of neural nets' statistical thresholds for sparse parity learning. We corroborate the theoretical analysis with a systematic empirical study of offline sparse parity learning using SGD on MLPs, demonstrating some of the (perhaps) counterintuitive effects of width, data, and initialization. For example, Figure 1 highlights our empirical investigation into the interactions between data, width, and success probability. The left figure shows the fractions of successful training runs as a function of dataset size (x-axis) and width (y-axis). Roughly, we see a "success frontier", where having a larger width can be traded off with smaller sample sizes. The right figure depicts some training curves (for various widths and sample sizes). Grokking (Power et al., 2022; Liu et al., 2022) (discontinuous and delayed generalization behavior induced by optimization dynamics) is evident in some of these figures.
"Parity2real" transfer of algorithmic improvements. It is often observed that deep learning methods underperform tree-based methods (e.g. random forests) on tabular datasets, particularly those where the target function depends on a few of the input features in a potentially non-smooth manner; see Grinsztajn et al. (2022) for a recent discussion. Motivated by our findings in the synthetic parity setting (and the observation that the sparse parity task possesses a number of the same properties of these problematic real-world datasets), we then turn to experimentally determine the extent to which our findings also hold for real tabular data. We evaluate MLPs of various depths and initialization sparsities on 16 tabular classification tasks, which were standardized by Grinsztajn et al. (2022) to compare neural vs. tree-based methods. Figure 3 shows that wider networks and sparse initialization yield improved performance, as in the parity setting. In some cases, our MLPs outperform tuned random forests.
In the nascent empirical science of large-scale deep learning, scaling laws (Kaplan et al., 2020; Hoffmann et al., 2022) have been shown to extrapolate model performance with remarkable consistency, revealing flexible tradeoffs and Pareto frontiers between the heterogeneous resources of data and computation. The present work reveals that in the simple synthetic setting of parity learning, the same intricacies can be studied theoretically and experimentally. In particular, viewing model size  $\times$  training iterations  $\times$  random restarts as a single "total FLOPs" resource, our study explains why data  $\times$  compute can be a necessary and sufficient resource, through the lens of SQ complexity.
Analyses of deep feature learning. Formally characterizing the representation learning mechanisms of neural networks is a core research program of deep learning theory. Many recent works have analyzed gradient-based feature learning (Wei et al., 2019; Barak et al., 2022; Zhenmei et al., 2022; Abbe et al., 2022; Damian et al., 2022; Telgarsky, 2022), escaping the "lazy" neural tangent kernel (NTK) regime (Jacot et al., 2018; Chizat et al., 2019), in which features are fixed at initialization.
Learning parities with neural networks. The XOR function has been studied as an elementary challenging example since the dawn of artificial neural networks (Minsky and Papert, 1969), and has been revisited at various times: e.g. neural cryptography (Rosen-Zvi et al., 2002); learning interactions via hints in the input distribution (Daniely and Malach, 2020); a hard case for self-attention architectures (Hahn, 2020). Closest to this work, Barak et al. (2022) find that in the case of online (infinite-data) parity learning, SGD provides a feature learning signal for a single neuron, and can thus converge at a near-optimal computational rate for non-overparameterized networks. They note computational-statistical tradeoffs and grokking in the offline setting, which we address systematically. Merrill et al. (2023) examine the same problem setting empirically, investigating a mechanism of competing sparse and dense sub-networks. Abbe et al. (2023) provide evidence that the time complexity required by an MLP to learn an arbitrary sparse Boolean function is governed by the largest "leap" in the staircase of its monomials, each of which is a sparse parity. Telgarsky (2022) gives a margin-based analysis of gradient flow on a two-layer neural network that achieves improved sample complexity  $(\tilde{O}(n/\epsilon))$  for the 2-sparse parity problem, at the cost of exponential width.
Neural nets and axis-aligned tabular data. Decision tree ensembles such as random forests (Breiman, 2001) and XGBoost (Chen and Guestrin, 2016) remain more popular among practitioners than neural networks on tabular data (Kaggle, 2021), despite many recent attempts to design specialized deep learning methods (Borisov et al., 2022). Some of these employ sparse networks (Yang et al., 2022b; Lutz et al., 2022) similar to those considered in our theory and experiments.
We refer the reader to Appendix A for additional related work.
2.1 Parities, and parity learning algorithms
A light bulb is controlled by  $k$  out of  $n$  binary switches; each of the  $k$  influential switches can toggle the light bulb's state from any configuration. The task in parity learning is to identify the subset of  $k$  important switches, given access to  $m$  i.i.d. uniform samples of the state of all  $n$  switches. Formally, for any  $1 \leq k \leq n$  and  $S \subseteq [n]$  such that  $|S| = k$ , the parity function  $\chi_S: \{\pm 1\}^n \to \pm 1$  is
defined as  $\chi_S(x_{1:n})\coloneqq \prod_{i\in S}x_i$ . The  $(n,k)$ -parity learning problem is to identify  $S$  from samples  $(x\sim \mathrm{Unif}(\{\pm 1\} ^n,y = \chi_S(x)))$ ; i.e. output a classifier with  $100\%$  accuracy on this distribution, without prior knowledge of  $S$ .
This problem has a rich history in theoretical computer science, information theory, and cryptography. There are several pertinent ways to think about the fundamental role of parity learning:
(i) Monomial basis elements: A  $k$ -sparse parity is a degree- $k$  monomial, and thus the analogue of a Fourier coefficient with "frequency"  $k$  in the harmonic analysis of Boolean functions (O'Donnell, 2014). Parities form an important basis for polynomial learning algorithms (e.g. Andoni et al. (2014)).
(ii) Computational hardness: There is a widely conjectured computational-statistical gap for this problem (Applebaum et al., 2009, 2010), which has been proven in restricted models such as SQ (Kearns, 1998) and streaming (Kol et al., 2017) algorithms. The statistical limit is  $\Theta (\log (\text{number of possibilities for } S)) = \Theta (k\log n)$  samples, but the amount of computation needed (for an algorithm that can tolerate an  $O(1)$  fraction of noise) is believed to scale as  $\Omega (n^{ck})$  for some constant  $c > 0$  - i.e., there is no significant improvement over trying all subsets.
(iii) Feature learning: This setting captures the learning of a concept which depends jointly on multiple attributes of the inputs, where the lower-order interactions (i.e. correlations with degree- $k' < k$  parities) give no information about  $S$ . Intuitively, samples from the sparse parity distribution look like random noise until the learner has identified the sparse interaction. This notion of feature learning complexity is also captured by the more general information exponent (Ben Arous et al., 2021).
2.2 Notation for neural networks
For single-output regression, a 2-layer multi-layer perceptron (MLP) is a function class, parameterized by a matrix  $W \in \mathbb{R}^{r \times n}$ , vectors  $v, b \in \mathbb{R}^r$  and scalar  $\beta \in \mathbb{R}$ , defined by
$$
\widehat {y}: x \mapsto v ^ {\top} \sigma (W x + b) + \beta ,
$$
where  $\sigma$  is an activation function, usually a scalar function applied identically to each input. For rows  $w_{i}$  of  $W$ , the intermediate variables  $\sigma (w_i^\top x + b_i)$  are thought of as hidden layer activations or neurons. The number of parallel neurons is often called the width.
In this section we theoretically study the interactions between data, width, time, and luck on the sparse parity problem. We begin by rehashing Statistical Query (SQ) lower bounds for parity learning in the context of gradient-based optimization, showing that without sufficient resources sparse parities cannot be learned. Then, we prove upper bounds showing that parity learning is possible with correctly scaling either width (keeping sample size small), sample size (keeping width small), or a mixture of the two.
Statistical query algorithms. A seminal work by Kearns (1998) introduced the statistical query (SQ) algorithm framework, which provides a means of analyzing noise-tolerant algorithms. Unlike traditional learning algorithms, SQ algorithms lack access to individual examples but can instead make queries to an  $SQ$  oracle, which responds with noisy estimates of the queries over the population. Notably, many common learning algorithms, including noisy variants of gradient descent, can be implemented within this framework. While SQ learning offers robust guarantees for learning in the presence of noise, there exist certain problems that are learnable from examples but not efficiently learnable from statistical queries (Blum et al., 1994, 2003). One notable example of such a problem is the parity learning problem, which possesses an SQ lower bound. This lower bound can be leveraged to demonstrate the computational hardness of learning parities with gradient-based algorithms (e.g., Shalev-Shwartz et al. (2017)).
Equivalently, parity can be represented as a function of a bit string  $b \in \{0,1\}^n$ , which computes the XOR of the influential subset of indices  $S$ :  $\chi_S(b_{1:n}) \coloneqq \oplus_{i \in S} b_i$ .
3.1 Lower bound: a multi-resource hardness frontier
We show a version of the SQ lower bound in this section. Assume we optimize some model  $h_\theta$ , where  $\theta \in \mathbb{R}^r$  (i.e.,  $r$  parameters). Let  $\ell$  be some loss function satisfying:  $\ell'(\hat{y}, y) = -y + \ell_0(\hat{y})$ ; for example, one can choose  $\ell_2(\hat{y}, y) = \frac{1}{2}(y - \hat{y})^2$ . Fix some hypothesis  $h$ , target  $f$ , a sample  $S \subseteq \{\pm 1\}^n$  and distribution  $\mathcal{D}$  over  $\{\pm 1\}^n$ . We denote the empirical loss by  $L_S(h, f) = \frac{1}{|S|} \sum_{\mathbf{x} \in S} \ell(h(\mathbf{x}), f(\mathbf{x}))$  and the population loss by  $L_{\mathcal{D}}(h, f) := \mathbb{E}_{\mathbf{x} \sim \mathcal{D}}[\ell(h(\mathbf{x}), f(\mathbf{x}))]$ . We consider SGD updates of the form:
$$
\theta_ {t + 1} = \theta_ {t} - \eta_ {t} \left(\nabla_ {\theta} \left(L _ {\mathcal {S} _ {t}} \left(h _ {\theta_ {t}}, f\right) + R \left(\theta_ {t}\right)\right) + \xi_ {t}\right),
$$
for some sample  $S_{t}$ , step size  $\eta_t$ , regularizer  $R(\cdot)$ , and adversarial noise  $\xi_t \in [-\tau, \tau]^r$ . For normalization, we assume that  $\|\nabla h_\theta(\mathbf{x})\|_\infty \leq 1$  for all  $\theta$  and  $\mathbf{x} \in \mathcal{X}$ .
Denote by  $\mathbf{0}$  the constant function, mapping all inputs to 0. Let  $\theta_t^\star$  be the following trajectory of SGD that is independent of the target. Define  $\theta_t^\star$  recursively s.t.  $\theta_0^\star = \theta_0$  and
$$
\theta_ {t + 1} ^ {\star} = \theta_ {t} ^ {\star} - \eta_ {t} \nabla_ {\theta} \left(L _ {\mathcal {D}} \left(h _ {\theta_ {t} ^ {\star}}, \mathbf {0}\right) + R \left(\theta_ {t} ^ {\star}\right)\right)
$$
Assumption 2 (Bounded error for gradient estimator). For all  $t$ , suppose that
$$
\left\| \nabla L _ {\mathcal {S} _ {t}} \left(h _ {\theta_ {t} ^ {\star}}, \mathbf {0}\right) - \nabla L _ {\mathcal {D}} \left(h _ {\theta_ {t} ^ {\star}}, \mathbf {0}\right) \right\| _ {\infty} \leq \tau / 2.
$$
Remark: If  $m = \tilde{O}(1/\tau^2)$ , Assumption 2 is satisfied w.h.p. for: 1)  $\mathcal{S}_t \sim \mathcal{D}^m$  ("online" SGD), 2)  $\mathcal{S}_t = \mathcal{S}$  for some  $\mathcal{S} \sim \mathcal{D}^m$  ("offline" GD) and 3)  $\mathcal{S}_t$  is a batch of size  $m$  sampled uniformly at random from  $\mathcal{S}$ , where  $\mathcal{S} \sim \mathcal{D}^M$  and  $M \geq m$  ("offline" SGD). Indeed, in all these cases we have  $\mathbb{E}_{\mathcal{S}_t}\left[\nabla L_{\mathcal{S}_t}(h_{\theta_t^\star}, \mathbf{0})\right] = \nabla L_{\mathcal{D}}(h_{\theta_t^\star}, \mathbf{0})$ , and the above follows from standard concentration bounds.
The lower bound in this section uses standard statistical query arguments to show that gradient-based algorithms, without sufficient resources, will fail to learn  $(n,k)$ -parities. We therefore start by stating the four types of resources that impact learnability:
- The number of parameters  $r$  - equivalently, the number of parallel "queries".
- The number of gradient updates  $T$  - i.e. the serial running time of the training algorithm.
- The gradient precision  $\tau$  - i.e. how close the empirical gradient is to the population gradient. As discussed above,  $\tilde{O}(1/\tau^2)$  samples suffice to obtain such an estimator.
- The probability of success  $\delta$ .
The following theorem ties these four resources together, showing that without a sufficient allocation of these resources, gradient descent will fail to learn:
Proposition 3. Assume that  $\theta_0$  is randomly drawn from some distribution. For every  $r, T, \delta, \tau > 0$ , if  $\frac{rT}{\tau^2\delta} \leq \frac{1}{2}\binom{n}{k}$ , then there exists some  $(n, k)$ -parity s.t. with probability at least  $1 - \delta$  over the choice of  $\theta_0$ , the first  $T$  iterates of SGD are statistically independent of the target function.
The proof follows standard SQ lower bound arguments (Kearns, 1998; Feldman, 2008), and for completeness is given in Appendix B. The core idea of the proof is the observation that any gradient step has very little correlation (roughly  $n^{-k}$ ) with most sparse-parity functions, and so there exists a parity function that has small correlation with all steps. In this case, the noise can force the gradient iterates to follow the trajectory  $\theta_1^{\star}, \ldots, \theta_T^{\star}$ , which is independent of the true target. We note that, while the focus of this paper is on sparse parities, similar analysis applies for a broader class of functions that are characterized by large SQ dimension (see Blum et al. (1994)).
Observe that there are various ways for an algorithm to escape the lower bound of Theorem 3. We can scale a single resource with  $\binom{n}{k}$ , keeping the rest small, e.g. by training a network of size  $n^k$ , or using a sample size of size  $n^k$ . Crucially, we note the possibility of interpolating between these extremes: one can spread this homogenized "cost" across multiple resources, by (e.g.) training a network of size  $n^{k/2}$  using  $n^{k/2}$  samples. In the next section, we show how neural networks can be tailored to solve the task in these interpolated resource scaling regimes.
2We use the notation  $\tilde{O}$  to hide logarithmic factors.
3.2 Upper bounds: many ways to trade off the terms
As a warmup, we discuss some simple SQ algorithms that succeed in learning parities by properly scaling the different resources. First, consider deterministic exhaustive search, which computes the error of all possible  $(n,k)$ -parities, choosing the one with smallest error. This can be done with constant  $\tau$  (thus, sample complexity logarithmic in  $n$ ), but takes  $T = \Theta(n^k)$  time. Since querying different parities can be done in parallel, we can reduce the number of steps  $T$  by increasing the number of parallel queries  $r$ . Alternatively, it is possible to query only a randomly selected subset of parities, which reduces the overall number of queries but increases the probability of failure.
The above algorithms give a rough understanding of the frontier of algorithms that succeed at learning parities. However, at first glance, they do not seem to reflect algorithms used in deep learning, and are specialized to the parity problem. In this section, we will explore the ability of neural networks to achieve similar tradeoffs between the different resources. In particular, we focus on the interaction between sample complexity and network size, establishing learning guarantees with interpolatable mixtures of these resources.
Before introducing our main positive theoretical results, we discuss some prior theoretical results on learning with neural networks, and their limitations in the context of learning parities. Positive results on learning with neural networks can generally be classified into two categories: those that reduce the problem to convex learning of linear predictors over a predefined set of features (e.g. the NTK), and those that involve neural networks departing from the kernel regime by modifying the fixed features of the initialization, known as the feature learning regime.
Kernel regime. When neural networks trained with gradient descent stay close to their initial weights, optimization behaves like kernel regression on the neural tangent kernel (Jacot et al., 2018; Du et al., 2018): the resulting function is approximately of the form  $\mathbf{x} \mapsto \langle \psi(\mathbf{x}), \mathbf{w} \rangle$ , where  $\psi$  is a data-independent infinite-dimensional embedding of the input, and  $\mathbf{w}$  is some weighting of the features. However, it has been shown that the NTK (more generally, any fixed kernel) cannot achieve low  $\ell_2$  error on the  $(n, k)$ -parity problem, unless the sample complexity grows as  $\Omega(n^k)$  (see Kamath et al. (2020)). Thus, no matter how we scale the network size or training time, neural networks trained in the NTK regime cannot learn parities with low sample complexity, and thus do not enjoy the flexibility of resource allocation discussed above.