How Feature Learning Can Improve Neural Scaling Laws

Abstract: We develop a solvable model of neural scaling laws beyond the kernel limit. Theoretical analysis of this model shows how performance scales with model size, training time, and the total amount of available data. We identify three scaling regimes corresponding to varying task difficulties: hard, easy, and super easy tasks. For easy and super-easy target functions, which lie in the reproducing kernel Hilbert space (RKHS) defined by the initial infinite-width Neural Tangent Kernel (NTK), the scaling exponents remain unchanged between feature learning and kernel regime models. For hard tasks, defined as those outside the RKHS of the initial NTK, we demonstrate both analytically and empirically that feature learning can improve scaling with training time and compute, nearly doubling the exponent for hard tasks. This leads to a different compute optimal strategy to scale parameters and training time in the feature learning regime. We support our finding that feature learning improves the scaling law for hard tasks but not for easy and super-easy tasks with experiments of nonlinear MLPs fitting functions with power-law Fourier spectra on the circle and CNNs learning vision tasks.

• The paper demonstrates that feature learning nearly doubles scaling exponents for hard tasks beyond the kernel regime.
• It employs a solvable DNN model to differentiate scaling behaviors across easy, hard, and super-easy tasks using RKHS norms.
• The study outlines compute-optimal strategies, informing resource allocation for efficient training of modern neural networks.

Insights into Feature Learning and Neural Scaling Laws

The paper "How Feature Learning Can Improve Neural Scaling Laws" by Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan presents a contributory exploration into the phenomenon of neural scaling laws, particularly how feature learning affects these laws. The authors propose a solvable model for understanding scaling behavior in neural networks beyond the traditional kernel limit, which is a regime where feature learning plays a significant role.

The central thesis of this work is the identification of distinct scaling regimes based on task difficulty—hard, easy, and super easy tasks—and how feature learning affects these tasks differently. The model explores these phenomena through the lens of deep neural networks (DNNs) by dissecting model performance in relation to the size of the network, the training time, and the amount of data available.

Key Findings

1. Task Difficulty and Scaling Laws: The study delineates tasks into categories based on difficulty. Easy and super-easy tasks are defined with respect to the reproducing kernel Hilbert space (RKHS) norms. For these tasks, which are within the RKHS of the initial neural tangent kernel (NTK), the scaling exponents remain consistent regardless of whether the network is in the feature learning regime or the kernel regime. However, for hard tasks that lie outside of this RKHS, the paper demonstrates that feature learning notably enhances scaling behavior.
2. Doubling of Exponents for Hard Tasks: A pivotal claim in the research is that feature learning nearly doubles the scaling exponent for hard tasks, offering a different optimal computational strategy for scaling parameters and training periods in the feature learning regime. This insight was buttressed by rigorous analytical and empirical evidence within the study, involving nonlinear MLPs on tasks with power-law Fourier spectra and CNNs tasked with vision learning.
3. Compute Optimal Strategies: The paper details strategies for allocating computational resources most effectively under the constraints of feature learning. These strategies are informed by the derived scaling laws and differ according to the difficulty level of tasks. The model posits specific exponents for the optimal compute-efficient scaling law across various data and learning scenarios.

Theoretical and Practical Implications

The implications of this research extend both theoretically and practically. Theoretically, it challenges existing paradigms by interposing feature learning dynamics into the discussion of scaling laws, traditionally dominated by kernel-based analyses. This model aligns scalable training with real-world high-dimensional settings by enhancing understanding of how neural networks inherently prioritize certain tasks due to feature re-weighting.

Practically, the insights gleaned from this research could inform the design and tuning of neural networks—particularly LLMs and vision models—by indicating when feature learning can be advantageous. In practice, adopting these insights can lead to smarter initialization, parameterization, and optimization techniques that leverage feature learning for task-specific training efficiency. Moreover, these strategies could refine curricula and data-sampling methodologies, potentially reducing redundancy and maximizing resource utility in computing environments.

Speculation on Future AI Developments

The study opens several prospective avenues for advancing AI, particularly in understanding the hidden dynamics of feature learning in deep networks. Future research might explore more complex architectures or diverse datasets to see how broadly these insights can be applied. Furthermore, examining how other elements such as hyperparameter choices, batch sizes, and specific network designs might further impact these scaling relations could yield additional economization and efficiency in AI deployment at scale.

The extension of this model's application from controlled scenarios to more varied datasets will arguably test the robustness of these claims, especially in intricate real-world tasks that involve multi-task learning environments where feature sharing is possible.

In summary, this work stands as a significant contribution to the understanding of neural scaling laws, particularly how feature learning reshapes scaling for hard tasks. The nuanced approach to deciphering these laws promises deeper optimization strategies for training large neural networks efficiently, which marks an important step in the trajectory of advanced AI research and application.