Learning in Large Neural Networks

Abstract

In this thesis, I will summarize my recent works on theoretical frameworks for learning, generalization, scaling limits, and scaling laws of large neural networks. In the first part of the thesis, we will examine the kinds of limits attained by training randomly initialized networks. The two limits of primary interest are the infinite width and infinite depth \textbf{feature-learning} limits. The large width limit will take the form of a dynamical mean field theory (DMFT), where neurons asymptotically decouple and the macroscopic dynamics of the network are governed by population averages over the neurons in each hidden layer. This theory computes the dynamics of the learned representations of data in each hidden layer of the network throughout training as well as the dynamics of the network output predictions. Asymptotic corrections to this mean field limit will be computed from fluctuations around the DMFT saddle point. The mean field limit will then be stressed tested on several realistic networks such as convolutional networks and transformers on realistic computer vision and language modeling datasets. We next investigate the infinite depth limit of residual neural networks, where each hidden layer is a trainable perturbation of the identity map. When the residual branches are scaled correctly, these models admit infinite width and depth limits which are computable with a DMFT where intermediate layers approach a continuum limit described by the solution to a set of stochastic integral equations. We empirically show that scaling the depth and width correctly enable hyperparameter transfer, where optimal hyperparameters (learning rates, batch sizes, momentum values, etc) in small width and small depth models are the same in large width and large depth models, reducing the need for hyperparamter tuning. We will extend these mathematical techniques to analyze several distinct infinite-parameter limits for transformer models. Next, I will describe simplified models of learning which enable a statistical analysis of generalization in a data-limited or parameter-limited regime. As in the previous section, we present both static and dynamic versions of these results. For data distributions and architectures that generate power law spectra for their limiting kernels, these theories provide predictions for the scaling laws with respect to the key computational and statistical resources: training time, model parameters, and total available data. I show that this setting provides a toy model of compute optimal scaling laws where model size and training time are traded off optimally. Lastly, we use similar mathematical techniques to analyze a toy model of hyperparameter transfer in randomly initialized deep linear networks. Lastly, we move beyond randomly initialized deep networks and attempt to address how \textbf{structured neural representations}, such as the cortical representations of external stimuli in real brains, encode an implicit learning bias. We start with a simple neural circuit trained with the delta rule, finding that the spectral decomposition of the population code controls which learning tasks can be learned in a sample-efficient manner. Next, we examine how different learning rules alter feature learning dynamics and inductive bias in multilayer neural networks, extending the DMFT for multilayer networks to other biologically plausible learning rules. Lastly, we illustrate how the geometry of neural codes which respect symmetries in the data generating process controls the manifold capacity of linear readouts.