Pareto Frontiers in Neural Feature Learning: Data, Compute, Width, and Luck

TL;DR

This paper explores how neural network architecture choices, like width and initialization, influence resource tradeoffs in feature learning, demonstrating that wider, sparsely-initialized models improve sample efficiency and can outperform traditional methods on benchmarks.

Contribution

It introduces a theoretical and experimental framework linking network width and initialization to resource tradeoffs and sample efficiency in feature learning, especially in sparse parity tasks.

Findings

Wider networks improve sample efficiency in sparse parity learning.

Sparse initialization enhances the probability of finding lottery ticket neurons.

Wide, sparsely-initialized models can outperform tuned random forests on benchmarks.

Abstract

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.