SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Feature learning regime. Due to the limitation of neural networks trained in the kernel regime, some works study learning in the "rich" regime, quantifying how hidden-layer features adapt to the data. Among these, Barak et al. (2022) analyze a feature learning mechanism requiring exceptionally small network width: SGD on 2-layer MLPs can solve the online sparse parity learning problem with network width independent of  $n$  (dependent only on  $k$ ), at the expense of requiring a suboptimal  $(\approx n^{k})$  number of examples. This mechanism is Fourier gap amplification, by which SGD through a single neuron  $\sigma(w^{\top}x)$  can perform feature selection in this setting, via exploiting a small gap between the relevant and irrelevant coordinates in the population gradient. The proof of Theorem 4 below relies on a similar analysis, extended to the offline regime (i.e., multiple passes over a dataset of limited size).
3.2.1 "Data  $\times$  model size" success frontier for sparsely-initialized MLPs
In this section, we analyze a 2-layer MLP with ReLU  $(\sigma(x) = \max(0, x))$  activation, trained with batch ("offline") gradient-descent over a sample  $S$  with  $\ell_2$ -regularized updates:
$$
\theta^ {(t + 1)} = (1 - \lambda^ {(t)}) \theta^ {(t)} - \eta^ {(t)} \nabla L _ {\mathcal {S}} \left(h _ {\theta^ {(t)}}\right)
$$
We allow learning rates  $\eta$ , and weight decay coefficients  $\lambda$  to differ between layers and iterations. For simplicity, we analyze the case where no additional noise is added to each update; however, we believe that similar results can be obtained in the noisy case (e.g., using the techniques in Feldman et al. (2017)). Finally, we focus on ReLU networks with  $s$ -sparse initialization of the first layer: every weight  $\mathbf{w}_i$  has  $s$  randomly chosen coordinates set to 1, and the rest set to 0. Note that after initialization, all of the network's weights are allowed to move, so sparsity is not necessarily preserved during training.
For simplicity, we do not assume adversarial noise in the gradients as in the lower bound. However, similar results can be shown under bounded noise.
Over-sparse initialization ( $s > \Omega(k)$ ). The following theorem demonstrates a "data  $\times$  width" success frontier when learning  $(n, k)$ -parities with sparsely initialized ReLU MLPs at sparsity levels  $s > \Omega(k)$ .
Theorem 4. Let  $k$  be an even integer, and  $\epsilon \in (0,1/2)$ . Assume that  $n \geq \Omega(1/\epsilon^2)$ . For constants  $c_1, c_2, c_3, c_4$  depending only on  $k$ , choose the following: (1) sparsity level:  $s \geq c_1/\epsilon^2$ , for some odd  $s$ , (2) width of the network:  $r = c_2(n/s)^k$ , (3) sample size:  $m \geq c_3(s/k)^{k-1}n^2\log n$ , and (4) number of iterations:  $T \geq c_4/\epsilon^2$ . Then, for every  $(n,k)$ -parity distribution  $\mathcal{D}$ , with probability at least 0.99 over the random samples and initialization, gradient descent with these parameter settings returns a function  $h_T$  s.t.  $L_{\mathcal{D}}(h_T) \leq \epsilon$ .
Intuitively, by varying the sparsity parameter in Theorem 4, we obtain a family of algorithms which smoothly interpolate between the small-data/large-width and large-data/small-width regimes of tractability. First, consider a sparsity level linear in  $n$  (i.e.  $s = \alpha \cdot n$ ). In this case, a small network (with width  $r$  independent of the input dimension) is sufficient for solving the problem, but the sample size must be large  $(\Omega(n^{k+1}))$  for successful learning; this recovers the result of Barak et al. (2022). At the other extreme, if the sparsity is independent of  $n$ , the sample complexity grows only as  $O(n^2 \log n)^4$ , but the requisite width becomes  $\Omega(n^k)$ .
Proof sketch. The proof of Theorem 4 relies on establishing a Fourier anti-concentration condition, separating the relevant (i.e. indices in  $S$ ) and irrelevant weights in the initial population gradient, similarly as the main result in Barak et al. (2022). When we initialize an  $s$ -sparse neuron, there is a probability of  $\gtrsim (s / n)^k$  that the subset of activated weights contains the "correct" subset  $S$ . In this case, to detect the subset  $S$  via the Fourier gap, it is sufficient to observe  $s^{k-1}$  examples instead of  $n^{k-1}$ . Initializing the neurons more sparsely makes it less probable to draw a lucky neuron, but once we draw a lucky neuron, it requires fewer samples to find the right features. Thus, increasing the width reduces overall sample complexity, by sampling a large number of "lottery tickets".
Under-sparse initialization ( $s < k$ ). The sparsity parameter can modulate similar data vs. width tradeoffs for feature learning in the "under-sparse" regime. We provide a partial analysis for this more challenging case, showing that one step of gradient descent can recover a correct subnetwork. Appendix B.3 discusses the mathematical obstructions towards obtaining an end-to-end guarantee of global convergence.
Theorem 5. For even  $k, s$ , sparsity level  $s < k$ , network width  $r = O((n / k)^s)$ , and  $\varepsilon$ -perturbed  $s$ -sparse random initialization scheme s.t. for every  $(n, k)$ -parity distribution  $\mathcal{D}$ , with probability at least 0.99 over the choice of sample and initialization after one step of batch gradient descent (with gradient clipping) with sample size  $m = O((n / k)^{k - s - 1})$  and appropriate learning rate, there is a subnetwork in the ReLU MLP that approximately computes the parity function.
Here, the sample complexity can be improved by a factor of  $n^s$ , at the cost of requiring the width to be  $n^s$  times larger. The proof of Theorem 5 relies on a novel analysis of improved Fourier gaps with "partial progress": intuitively, if a neuron is randomly initialized with a subset of the relevant indices  $S$ , it only needs to identify  $k - s$  more coordinates, inheriting the improved sample complexity for the  $(n - s, k - s)$ -parity problem. Note that the probability of finding such a lucky neuron scales as  $(n / k)^{-s}$ , which governs how wide (number of lottery tickets) the network needs to be.
Remarks on the exact-sparsity regime. Our theoretical analyses do not extend straightforwardly to the case of  $s = \Theta(k)$ . Observe that if the value of  $k$  is known, initializing a network of size  $r = \widetilde{O}(n^k)$  with sparsity  $s = k$  gives w.h.p. a subnetwork with good first-layer features at initialization. We believe that with a proper choice of regularization and training scheme, it is possible to show that such a network learns to select this subnetwork with low sample complexity. We leave the exact details of this construction and end-to-end proofs for future work.
Analogous results for dense initialization schemes? We believe that the principle of "parallel search with randomized per-subnetwork sample complexities" extends to other initialization schemes
<sup>4</sup>We note that the additional  $n^2$  factor in the sample complexity can be removed if we apply gradient truncation, thus allowing only a logarithmic dependence on  $n$  in the small-sample case.
For ease of analysis, we use a close variant of the sparse initialization scheme:  $s$  coordinates out of the  $n$  coordinates are chosen randomly and set to 1, and the rest of the coordinates are set to  $\epsilon < (n - s)^{-1}$ . Without the small norm dense component in the initialization, the population gradient will be 0 at initialization.
Empirical statistical threshold vs. width+init
Computation time vs. width+init
Figure 2: Zoomed-in views of the interactions between width, data, time, and luck. Left: Locations of these runs in the larger parameter space. Center: Success probability vs. dataset size. In accordance with our theory, width buys luck, and improves end-to-end sample efficiency, despite increasing the network's capacity. Right: Number of iterations to convergence vs. dataset size. We observe a data vs. time tradeoff in the grokking regime (at the edge of feasibility), as well as a "sample-wise double descent" performance drop with more data (which can also be seen in Figure 1, and disappears with larger widths). Comparing dense vs. sparse initializations (upper vs. lower plots; sparse inits are colored more brightly), we see computational and statistical benefits of the sparse initialization scheme from Section 3.2.1.
(including those more commonly used in practice), and leads to analogous success frontiers. To support this, our experiments investigate both sparse and uniform initialization, with qualitatively similar findings. For dense initializations, the mathematical challenge lies in analyzing the Fourier anti-concentration of general halfspaces (see the discussion and experiments in Appendix C.1 of (Barak et al., 2022)). The axis-aligned inductive biases imparted by sparse initialization may also be of independent practical interest.
A high-level takeaway from Section 3 is that when a learning problem is computationally difficult but statistically easy, a complex frontier of resource tradeoffs can emerge; moreover, it is possible to interpolate between extremes along this frontier using ubiquitous algorithmic choices in deep learning, such as overparameterization, random initialization, and weight decay. In this section, we explore the nature of the frontier with an empirical lens—first with end-to-end sparse parity learning, then with natural tabular datasets.
4.1 Empirical Pareto frontiers for offline sparse parity learning
We launch a large-scale ( $\sim$ 200K GPU training runs) exploration of resource tradeoffs when training neural networks to solve the offline sparse parity problem. While Section 3 analyzes idealized variants of SGD on MLPs which interpolate along the problem's resource tradeoff frontier, in this section we ask whether the same can be observed end-to-end with standard training and regularization.
On various instances of the  $(n,k)$ -sparse parity learning problem, we train a 2-layer MLP with identical hyperparameters, varying the network width  $r\in \{10,30,100,\dots ,10^5\}$  and the dataset size  $m\in \{100,200,300,600,1000,\ldots ,10^{6}\}$ . Alongside standard algorithmic choices, we consider one non-standard augmentation of SGD: the under-sparse initialization scheme from Section 3.2.1; we have proven that these give rise to "lottery ticket" neurons which learn the influential coordinates more sample-efficiently. Figure 1 (in the introduction) and Figure 2 illustrate our findings at a high level; details and additional discussion are in Appendix C.1). We list our key findings below:
(1) A "success frontier": large width can compensate for small datasets. We observe convergence and perfect generalization when  $m \ll n^k$ . In such regimes, which are far outside the
online setting considered by Barak et al. (2022), high-probability sample-efficient learning is enabled by large width. This can be seen in Figure 1 (left), and analogous plots in Appendix C.1.
(2) Width is monotonically beneficial, and buys data, time, and luck. In this setting, increasing the model size yields exclusively positive effects on success probability, sample efficiency, and the number of serial steps to convergence (see Figure 2). This is a striking example where end-to-end generalization behavior runs opposite to uniform convergence-based upper bounds, which predict that enlarging the model's capacity worsens generalization.
(3) Sparse axis-aligned initialization buys data, time, and luck. Used in conjunction with a wide network, we observe that a sparse, axis-aligned initialization scheme yields strong improvements on all of these axes; see Figure 2 (bottom row). In smaller hyperparameter sweeps, we find that  $s = 2$  (i.e. initialize every hidden-layer neuron with a random 2-hot weight vector) works best.
(4) Intriguing effects of dataset size. As we vary the sample size  $m$ , we note two interesting phenomena; see Figure 2 (right). The first is grokking (Power et al., 2022), which has been previously documented in this setting (Barak et al., 2022; Merrill et al., 2023). This entails a data vs. time tradeoff: for small  $m$  where learning is marginally feasible, optimization requires significantly more training iterations  $T$ . Our second observation is a "sample-wise double descent" (Nakkiran et al., 2021): success probability and convergence times can worsen with increasing data. Both of these effects are also evident in Figure 1).
Lottery ticket neurons. The above findings are consistent with the viewpoint taken by the theoretical analysis, where randomly-initialized SGD plays the role of parallel search, and a large width increases the number of random subnetworks available for this process—in particular, the "winning lottery ticket" neurons, for which feature learning occurs more sample-efficiently. To provide further evidence that sparse subnetworks are responsible for learning the parities, we perform a smaller-scale study of network pruning ability in Appendix C.2.
4.2 Sample-efficient deep learning on natural tabular datasets
Sparse parity learning is a toy problem, in that it is defined by an idealized distribution, averting the ambiguities inherent in reasoning about real-life datasets. However, due to its provable hardness (Theorem 3, as well as the discussion in Section 2.1), it is a maximally hard toy problem in a rigorous sense<sup>6</sup>. In this section, we perform a preliminary investigation of how the empirical and algorithmic insights gleaned from Section 4.1 can be transferred to more realistic learning scenarios.
To this end, we use the benchmark assembled by Grinsztajn et al. (2022), a work which specifically investigates the performance gap between neural networks and tree-based classifiers (e.g. random forests, gradient-boosted trees), and includes a standardized suite of 16 classification benchmarks with numerical input features. The authors identify three common aspects of tabular<sup>7</sup> tasks which present difficulties for neural networks, especially vanilla MLPs:
(i) The feature spaces are not rotationally invariant. In state-of-the-art deep learning, MLPs are often tasked with function representation in rotation-invariant domains (token embedding spaces, convolutional channels, etc.).
(ii) Many of the features are uninformative. In order to generalize effectively, especially from a limited amount of data, it is essential to avoid overfitting to these features.
(iii) There are meaningful high-frequency/non-smooth patterns in the target function. Combined with property (i), decision tree-based methods (which typically split on axis-aligned features) can appear to have the ideal inductive bias for tabular modalities of data.
Noting that the sparse parity task possesses all three of the above qualities, we conduct a preliminary investigation on whether our empirical findings in the synthetic case carry over to natural tabular data. In order to study the impact of algorithmic choices (mainly width and sparse initialization) on sample efficiency, we create low-data problem instances by subsampling varying fractions of each dataset for training. Figure 3 provides a selection of our results. We note the following empirical findings, which are the tabular data counterparts of results (2) and (3) in Section 4.1:
$^{6}$ Namely, its SQ dimension is equal to the number of hypotheses, which is what leads to Theorem 3.  $^{7}$ Tabular data" refers to the catch-all term for data sources where each coordinate has a distinct semantic meaning which is consistent across points.
Figure 3: Analogous investigation of MLP width  $r$  and sparse initialization for real-world tabular datasets (OpenML benchmarks assembled by Grinsztajn et al. (2022)), varying the dataset size  $m$  via downsampling. Large width and sparse initialization tend to improve generalization, in accordance with the theory and experiments for synthetic parity tasks. In some settings, our best MLPs outperform tuned random forests. Dotted lines denote the test errors reported by Grinsztajn et al. (2022) of tuned MLPs and RFs on the full datasets. Full results on all 16 tasks are in Appendix C.3.
(2T) Wide networks generalize on small tabular datasets. Like in the synthetic experiments, width yields nearly monotonic end-to-end benefits for learning. This suggests that the "parallel search + pruning" mechanisms analyzed in our paper are also at play in these settings. In some (but not all) cases, these MLPs perform competitively with tuned tree-based classifiers.
(3T) Sparse axis-aligned initialization sometimes improves end-to-end performance. This effect is especially pronounced on datasets which are downsampled to be orders of magnitude smaller. We believe that this class of drop-in replacements for standard initialization merits further investigation, and may contribute to closing the remaining performance gap between deep learning and tree ensembles on small tabular datasets.
We have presented a theoretical and empirical study of offline sparse parity learning with neural networks; this is a provably hard problem which admits a multi-resource lower bound in the SQ model. We have shown that the lower bound can be surmounted using varied mixtures of these resources, which correspond to natural algorithmic choices and scaling axes in deep learning. By investigating how these choices influence the empirical "success frontier" for this hard synthetic problem, we have arrived at some promising improvements for MLP models of tabular data (namely, large width and sparse initialization). These preliminary experiments suggest that a more intensive, exhaustive study of algorithmic improvements for MLPs on tabular data has a chance of reaping significant rewards, perhaps even surpassing the performance of decision tree ensembles.
Broader impacts and limitations. The nature of this work is foundational; the aim of our theoretical and empirical investigations is to contribute to the fundamental understanding of neural feature learning, and the influences of scaling relevant resources. A key limitation is that our benchmarks on tabular data are only preliminary; it is a significant and perennial methodological challenge to devise fair and comprehensive comparisons between neural networks and tree-based learning paradigms.
We are grateful to Boaz Barak for helpful discussions and to Matthew Salganik for helpful comments on a draft version. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. Sham Kakade acknowledges funding from the Office of Naval Research under award N00014-22-1-2377. Ben Edelman acknowledges funding from the National Science Foundation Graduate Research Fellowship Program under award #DGE 2140743.
Tolstikhin, I. O., Houlsby, N., Kolesnikov, A., Beyer, L., Zhai, X., Unterthiner, T., Yung, J., Steiner, A., Keysers, D., Uszkoreit, J., et al. (2021). MLP-Mixer: An all-MLP architecture for vision. Advances in neural information processing systems, 34:24261-24272.
Wei, C., Lee, J. D., Liu, Q., and Ma, T. (2019). Regularization matters: Generalization and optimization of neural nets vs their induced kernel. Advances in Neural Information Processing Systems, 32.
Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. (2022a). Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466.
Yang, J., Lindenbaum, O., and Kluger, Y. (2022b). Locally sparse neural networks for tabular biomedical data. In International Conference on Machine Learning, pages 25123-25153. PMLR.
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. (2022). Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12104-12113.
Zhenmei, S., Wei, J., and Liang, Y. (2022). A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations.