This is a post explaining my paper with Kaarel Hänni on complexity of infinite-width networks. I will explain the result, why it matters, and how the mathematical idealizations can interact with real structure in neural nets. This leads to some threads I am excited to pull on more in the future, via some new speculations on where interpretable structure can live.
Introduction
Our paper to some extent (and up to some important details) concludes an analysis of generalization complexity in Bayesian infinite-width neural nets. It resolves the question of when a function can be learned in polynomial sample complexity (in complexity theory, polynomial size is a fundamental asymptotic size class that means not too large. Informally, things outside the polynomial class are understood to be effectively unreachable/ uncomputable at realistic scales). More precisely, it asks whether (or when) generalization properties in infinite with neural nets reduces to generalization in finite width neural nets.
Here a crucial observation is that there are different infinite-width limits even within the same architecture, and they lead to different generalization behavior. The difference between these limits in Bayesian learning is in the choice of the learning prior (analogous to initialization or, more directly, to weight regularization in SGD learning) as a function of width. While seemingly a technical choice, this turns out to be extremely important for the infinite-width phenomenology. The relationship between initialization and generalization is familiar in usual SGD learning: for instance, the amusingly named omnigrok paper (Ziming Liu et al.) shows that modular addition generalizes better under small initialization (or strong weight regularization), and similar phenomena are observed in Bayesian learning where stronger regularization (corresponding to smaller initialization) tends to generalize better. For example the paper that explains Bayesian generalization in modular addition implies that changing the regularization is one of the levers that interpolates between the memorizing and the generalizing regimes.
In the infinite-width limit, we are inherently looking at a family of models that depend on a hidden-width dimension (often called in theory papers), which is taken to infinity. There are two natural scalings for the regularization strength, related to different choices of natural matrix norm (see also Greg Yang et al.), where a regularization parameter scales either weakly, as or strongly as . The weakly regularized limit is called the standard limit and the strongly regularized limit is called the mean-field limit. There is also a natural interpolation between the two depending on a general real parameter , where one must be careful about differences in scaling between the early layers and the last layer to avoid certain degeneracies -- I will not get into this here[1].
Details on complexity and setup (hidden by default)
A few details here to fix concepts for people who want to track the rigorous results.
When we talk about polynomial complexity, we mean that we ask for bounds on this parameter which are some fixed polynomial in terms of the input dimension, as is usual in complexity theory of circuits. Note that the infinite-width setup always assumes that while the hidden layers are taken to infinity as the large number H, the input dimension stays fixed, as does the target or label function We will later consider another parameter which is the dataset size D; this is assumed much larger, and in the case of all boolean inputs the total datasize is exponential, . The fact that the data dimension d is understood as polynomial while the total possible data is exponential is crucial for learning theory.
We assume that our activation function is bounded. This assumption is relatively harmless: e.g. for any given Relu neural net, we can replace the Relu by a clipped Relu so long as we upper-bound at a sufficiently large value, and empirically using bounded activations (such as sigmoids) does not significantly affect performance.
Most importantly, we view the depth as a constant parameter (small compared to ) and our bounds are allowed to depend on nonpolynomially; more on this below.
The result
Once the technical dust settles, we have at least two different settings to theoretically investigate generalization behavior of overparameterized neural nets in the limit: the mean-field setting and the standard setting. Our paper implies, essentially, that a very wide overparameterized neural network generalizes poorly in the standard limit and generalizes "as well as possible" in the mean-field limit.
The first part of this statement is known (and we cite it instead of re-proving it). In the infinite limit of standard scaling, learning reduces to linear regression in a certain nonlinearly re-embedded data manifold associated with a so-called Gaussian process kernel. Crucially, the re-embedding depends on the training data but is independent of the labels , so informally all the label learning must occur in a single linear-regression layer. If the labels depend on some deep sequential structure where representations are transformed layer-to-layer (such as what happens in a Boolean circuit), then you're out of luck. More precisely the network can learn any function eventually with enough datapoints (by essentially solving linear regression in a suitably parameterized linear space of all functions), but this will take exponential data size D.
But the mean-field limit is known to develop deep, layerwise structure that depends on . This happens via a phenomenon called "feature learning" (this is a silly name, that refers to the fact that hidden layers are learning at all as opposed to the standard scaling where in the limit they are "frozen" to their prior configuration. It neither assumes nor implies any natural notion of features, other than the tautological features associated to the activations in hidden layers, or rather their relative configuration up to rotation matrices which is encoded by the activation dot product matrix).
The question our paper resolves is how much of the deep structure that can develop in a narrow network can survive in the infinite-width deep limit. And our answer is, informally, "all of it". Here I am of course being pithy. We don't really have a good notion of "structure", but if we did there would certainly be structures (mechanisms, circuits, features and so on) that behave differently at finite (polynomial) and infinite width. The thing we actually show is that if sufficient structure can be found in a large dataset to generalize well at some finite (more precisely polynomial) width in an underparameterized setting, then the same good generalization behavior survives to infinite width.
More precisely, one of our main results is the following. Fix a boolean[2] target (or label) function . Assume that when we train on this target at some polynomial width , the function generalizes after polynomially many data points (we will unpack the relevant notion of generalization below). Then we deduce that the infinite-width limiting model (and indeed, any sufficiently wide model) also generalizes after polynomially many data points.
Conversely, we show that if the infinite-width model generalizes after polynomially many data points then there is some polynomial width such that the width- model learns the target (also in polynomial data).[3]
Here our notion of generalization is worth unpacking. Given a Bayesian model (determined by the weight scaling, the width , and a choice of tempered Bayesian posterior), we define the generalization loss at data size as the expectation over tuples of boolean datapoints of the expected test (i.e. population) MSE loss of the trained model (i.e. a model sampled from the Bayesian posterior):
This value has a sensible extension to width We say that a model generalizes relative to a small error parameter at width and data points if Our main result is that if for some polynomial then we also have for some polynomial . And conversely we show that implies that there is some pair of polynomially bounded such that
In other words, we show that "learning polynomial data gives small error" is true at infinite width if and only if it is true at polynomial width.
Why think about infinite width
Before moving on to details and caveats (though note that there is a particularly big caveat associated to the asymptotics w.r.t. depth), let's step back and see why this is interesting. After all, no one wants to train infinite-width networks. It's also not true that one should deduce immediate insights here about empirical neural nets from this, for example that "more regularization is good for generalization always". First, empirical models are not "close" to the infinite-width limit, and even when infinite-width analyses are stretched down to realistic width, many new phenomena emerge (some kinds of emergent regularization, certain instanton phenomena) and intuitions don't directly carry over. Second, empirical models are trained dynamically (via SGD or Adam or similar) and Bayesian learning results must be understood with a grain of salt.
Instead, I think that there are two main reasons why our result is interesting:
They imply a width-independent notion of heuristic bias, which may be more elegant and natural to study, and which might help give cleaner models of the relationship between data, parameter size, and generalization.
Our results suggest that the theoretically simpler circuits and mechanisms found at infinite width are sufficient to capture any regularities that are captured in finite-width neural nets. Various "natural parsimony" heuristics (like natural latents) suggest that two different predictive models of the same data are likely to find similar patterns and circuits, and thus interpreting the infinite-width model may capture some of the same mathematical structures used by a finite-width learner.
I find the promise of item 2 particularly exciting here. Infinite-width limits are physical objects: they can be described by a certain set of integral equations on fields, which are associated to training data points (rather than weights). While the resulting object has many parameters (called observables), it gives a new perspective that naturally incorporates tools from statistical field theory of many-particle systems. In particular, the ability to control sample complexity opens up a new perspective related to scale and renormalization. The key upshot of our results is that a simple training objective leads to a simple statistical description, in a way that is independent of width. We discuss this idea and a few speculative implications in the following sections.
Small tasks have small complexity
Throughout this discussion, we are carrying around two numbers which can be relatively big. The first is the training data size and the second is the width . Generalization on a task implies that, informally, the data size never needs to be taken to infinity: there is some finite amount of data (only dependent on the task function ) which is sufficient to fully learn it at least to some finite desired level of population accuracy (a.k.a. "test" accuracy). A priori this "characteristic" data count D is a function of the hidden width, . Our result implies that in rich scaling there is some more fundamental data size such that the target is already learned after seeing random datapoints at any width (here and elsewhere I am sweeping under the rug some discussion of coarseness, related to the parameter and the temperature of tempered Bayesian learning).
The upshot is that (modulo some discussions of coarseness and error bounds), the function learned after seeing any large data fraction is equivalent to the function learned after seeing data. If the target is simple and the width is large, we have , and this puts us solidly in the mean field regime. Heuristically, this means that if we were to train a giant multi-billion parameter model[4] on some simple task like induction learning, and the training is in a Bayesian regime with mean-field regularization, we expect the model to be well-described by a mean-field prediction. This means that a single finite physics-flavored description simultaneously predicts the behavior at any width.
Parallel (or superposed) tasks have emergent regularization
Now it's likely that for a very complex task like language prediction, the effective data complexity is very large if we want perfect performance, or at least very large compared to the parameter count of modern llms. But if we only care about explaining a very small amount of performance, or a simple component of the model, its effective data requirement may be small. Famously it is often reasonable to think of neural nets as not so much a single complex algorithm as a set of interlocking but somewhat independent algorithms: informally, a bag of heuristics. Some components or subproblems may be easier while others may be harder, and there are likely various ladders of complexity associated to a precision parameter or some other notion of scale. For example in some coarse decomposition, an LLM has pieces that are (at least early in training) just following bigram patterns or induction circuits, or offsetting days of the week via modular addition.
A highly simplified model here is to view the subproblems as existing in parallel, or in superposition. If a single component is simple enough and we learn just that component in a mean field-regularized fashion, the resulting physics should fully predict the weight distribution as explained in the previous section. However we cannot deduce that the collection of parallel circuits are given by some simple combination of the resulting mean fields. When a neural network performs multiple tasks in parallel or in superposition, these tasks are forced to share resources such as weight norm and entropy. But if we're lucky we can predict the tasks as effectively executing some set of superposed infinite-width limits once we incorporate the shared resources via emergent regularization. My notion of a shared resource is intentionally left vague, but an example of a problem where this can be fully worked out is the sparse denoising model here. A particular shared resource that must always be tracked is weight norm: since different parallel circuits must share weight norm, we expect a tuple of many circuits to mutually regularize, leading to an emergent regularization phenomenon where the effective regularization felt by any particular small circuit is much larger than what it would experience if it were the only thing being learned. This leads to a situation where, for example, under standard scaling a superposition of many circuits behaves as if each circuit were undergoing mean field scaling by itself.
Extensions and limitations
It should be clear that the above discussion leads to a particular orientation on neural net interpretability. One can conjecture that a large neural net decomposes into small simple circuits, each of which has low complexity (low effective width-independent complexity ) and sees the empirical width as effectively infinite. They interact and regularize each other via shared resources and hierarchical relationships, but are otherwise independent and simple. After accounting for emergent regularization, some circuits behave as if they are in the mean field regime while others behave as if they are in the standard (i.e. Gaussian process) regime. I really like this as a cartoon and a thread to pull on. This picture is related to my excitement about variants of the renormalization agenda in the org I work in (Principles of Intelligence).
At the same time I don't expect this to be the full story, for several reasons.
First, the separation of phenomena into mean field and Gaussian is a false dichotomy, and it is possible to come up with nice examples where both mean-field and Gaussian process phenomena interact in a significant way and must be understood together to describe the system. A complete analysis of the sparse denoising theory (not yet published) should be such a case. A general study of such phenomena, given by a 2-particle mean field theory, is related to past and upcoming work by my colleague Nischal Mainali. A related phenomenon is explained in the "Proof idea" section below.
Second, neural nets learn dynamically (SGD and so on), whereas the mean field limits are cleanest in Bayesian learning. This is not as bad as it might seem, since much of the mean-field analysis of neural nets applies also to SGD (and is believed to apply to other gradient-based optimizers like Adam, etc.). In fact, Kaarel and I are also discussing a version of the paper that would apply for dynamical learning, with additional caveats and details.
Third, mean field-style analyses may break down when studying underparameterized (or more generally "not-very-wide") networks via a neuron specialization degeneracy. Here instead of interpolating between different simple circuits in a superposition-like way, neurons may decide to specialize to a single circuit, or more precisely a circuit may depend on a single neuron or a small set of neurons in a robust mechanistic way without redundancy; this would break a fundamental assumption necessary for the standard infinite-width techniques. This phenomenon and how it fits into the various infinite-width limits is discussed in particular in the paper "Mitigating the Curse of Detail" by Noa Rubin et al. In fact, there is much evidence that this kind of monosemantic behavior does occur in some circuits (a bunch of classical examples in vision models are provided by Chris Olah's interactive account here). Phenomenology of this kind is genuinely new relative to the thermodynamic phenomena that occur in infinite-width limits, and is more or less directly analogous to instanton symmetry-breaking effects. Exact predictions for such systems would likely require a separate analysis similar to work of Saxe on symmetry breaking (see for example Saxe-McClelland-Ganguli), or the analysis of symmetry-broken instanton phenomena in physics.
For this reason and others, I don't expect a complete theory of interpretability of finite networks which is reducible to "mean field theory plus inductive regularization phenomena". However I am very hopeful that some of the simpler and more polysemantic mechanisms in large neural nets (such as induction heads) can in fact be explained in this way, or at least use the same mathematical mechanisms as mean field learners. Optimistic signals are that this is true for modular addition (where both narrow SGD and infinite-width mean-field models learn the same "clock" Fourier-basis mechanism). We also see a deep-network example of this (mechanistic agreement between narrow and mean field solutions) in an example of an SGD-learnable deep boolean XOR-like task in upcoming work joint with Sergey Alexeev (whose MARS project I am mentoring).
MFT and neuron independence
A second aside before moving on: often when I talk to people about mean field theory, they assume that this is a simplistic model that postulates independence between neurons, and hence cannot capture phenomena that require collaboration between many neurons. Under such a model it seems strange to them, and perhaps even un-believable, that a mean field formalism can capture the full complexity of algorithms learned by a neural net.
I think that this misconception makes sense, and is important enough to devote a section to.
The simple answer is that while the mean field limit does imply a certain kind of neuron independence, the independence is conditional, on a relatively complex self-organizing "macroscopic" structure which (as we show in our paper) is powerful enough to encode algorithms of arbitrary complexity.[5]
Informally, a first-order way to think of mean field theory is that neurons have jobs. More precisely in a mean field cartoon, each of ~infinitely many neuron indices at a given hidden layer is labeled by a "job" latent from some finite set. The weight between a pair of neurons at adjacent layers is then drawn as a (statistical) variable that depends only on the two jobs, and not the associated indices. Thus each neuron's preactivation is a weighted linear combination, associated to its job, of activations from job postactivations of neurons at the previous layer. The key point (and the key reason why mean field techniques are interesting for interpretability) is that the number of jobs can be large, but is in some sense finite, depending on the number of training data points; while the number of neurons is assumed to be infinite and indistinguishable. Thus the onus of interpretability is moved away from neurons and into fields, which I called jobs. Since the number of possible fields scales with the number of training data points, this formalism is able to represent very complex mechanisms where neurons at different layers collaborate. For example if a neural net is learning a finite boolean circuit, its jobs at any given layer can be associated to the finitely many hidden gates of the circuit (and this is heuristically what we observe for XOR), with each neuron type reading the two neuron types associated with its parents in the circuit. The full formalism of mean field has its interacting fields indexed not in terms of informal "job labels" but in terms of activation kernels, and a rough translation between these two points of view is a central intuition pump for our paper.
Crucially, the dimensionality of the field theory (the "maximal number of jobs") is related to the training data dimension. As I explained above, our results imply that structure saturates at some polynomial data, independent of width, and in the large-width regime it is this value (rather than the parameter count) that encodes the complexity (the dimension of "job space"). This finite set of parameters defines the landscape of possible neuron jobs and their interactions. Neurons sampling this space of jobs are now conditionally independent, but interact in structured ways mediated by this parameter. This implies the existence of a certain platonic, width-independent mechanism for implementing any task that can be learned by a neural network. Of course the platonicity here, as in most cases, is elusive: different choices of activation function, re-weightings of layers, and of course questions of dynamics/SGD can lead to a much more hyperparameter-sensitive picture. However the upshot is that nevertheless the model reduces to a more streamlined abstraction controlled by physics - I think this is exciting.
Details and caveats
I'll explain an intuition pump for our proof in the section below, but first here I will point out the assumptions we make on the model, some of which are cosmetic and others which are significant.
The main cosmetic assumptions we make (which we think can be replaced by weaker assumptions) are:
that the activation function is bounded (e.g. a sigmoid or a clipped Relu)
that the data distribution (on training and test data ) is the uniform distribution on the boolean cube (i.e. coordinates drawn uniformly from ).
The bounded-activation assumption is there to avoid allowing neuron specialization in the infinite-width limit, and the finite data distribution on the cube is useful (more fundamentally) to keep the data bounded and, less fundamentally, to work with a finite total data distribution of size --- large enough that it's "huge" (larger than any polynomial) in the dimension parameter but still smaller than the width parameters approaching infinite-width limits, which are assumed to be on an ever larger super-exponential scale.
The most significant issue in our paper is that the various bounds in our paper scale very poorly in depth (we use the letter to denote depth, as is standard). While at a fixed depth the bounds are polynomial (so generalization complexity at some finite width implies generalization complexity at infinite width, with a fixed polynomial independent of the target ), the dependence on of the bounds is exponential; so the more precise statements have bounds that look like
where is a fixed degree- polynomial that depends only on the activation function. This can be practical for small (like 2 or 3), and in fact already gives quasiexponential improvements over standard scaling for boolean circuits (like XOR) at depth (for the data dimension/ number of bits). But for deep networks, our result as it stands leads to horrible bounds with only theoretical interest. While we can console ourselves that this is a theory paper, we think this is not the best one can do.
The state of our rigorous mathematical knowledge currently has this ugly factor of in the exponent that makes our current bounds extremely sensitive to depth. However, we believe that under an additional Lipschitzness condition (related to this work by Linsefors and Bushnaq) we can make the polynomial dependences depth-independent and we are actively working on an extension in this direction. This argument requires datapoints to be distinct (like the boolean cube we work with, or text strings) and effectively forces the deep neural nets that are "allowed" in our distribution to perform some denoising/ error correction in data-space basins around the datapoints. Interestingly, this kind of behavior is in fact observed in practice, most remarkably and carefully in a series of posts by Stefan Heimersheim and his students Francisco Ferreira da Silva and Matthew Shinkle (see for example this post). Thus on the level of heuristics and ideation, I expect our results to be robust, at least to some extent and under some idealization, to the large-width limits of empirical deep networks. These empirical stability phenomena make me hopeful that some of the mean field-inspired ideas sketched out in the previous sections apply to deep neural nets, despite the fact that our formal Bayesian guarantees have a horrible depth instability that give unusable empirical complexity bounds for deep networks.
Proof idea
Our result works, to first order, by essentially formalizing the idea of "jobs" in the section on neuron independence. This is easiest to frame in a certain richer version of mean-field (where we move an exponential parameter that interpolates between standard scaling and mean-field scaling even further away from standard, corresponding to even stronger regularization, and consider the limit ). Here the infinite limit corresponds to finding a minimal L2 norm neural net solving the task (where the total L2 norm is an appropriate weighted sum over layers). Given a network with small norm at finite width, we can "blow it up" to a network with the same norm, at infinite width by a process we call cloning. This works differently at middle layers vs. embedding/unembedding layers. The middle layer story is representative (and more interesting). It works by taking a narrow middle-layer matrix and simply writing out a block matrix consisting of identical "clones" of in an square, with a prefactor of :
Taking gives us an infinite-width "limiting structure". You can observe that the large matrix has the same norm (sometimes called "Frobenius norm") as the original matrix. You can also see that this limiting structure leads to the same input-output behavior as the original neural net. If the activations at layer for some datapoint are given by the vertical vector then we see by layer-to-layer induction that the activations at layer are again a repeated version: repeated times. We now see that the next-layer preactivation is again the concatenation , repeated times. Since the activation function is applied recursively, this completes the induction. Conversely, given a large neural net with hidden dimension we can stochastically reduce it to a smaller net of size by a process we call subsampling: simply remove at random a collection of neurons, and any weight coefficients that abut them, leaving a neural net of hidden width . Here (carefully eliminating outlier neurons) we can guarantee, using statistical concentration bounds, that the resulting net has norm and input-output behavior that are close to the original network, with error bounded by a width-independent (but target-dependent) constant times . Thus by paying some finite precision cost (corresponding to the difference between the accuracy assumption and the accuracy infinite-width guarantee above), we can show that the resulting narrow subsampled net solves the task with admissible error. These arguments imply that if a small-norm neural net that performs well on a task exists at some polynomial-size width, then such a net also exists at infinite width, and conversely good performance at infinite width implies the existence of a performant net at finite widths. A priori these guarantees of existence do not imply anything about generalization, and the generalization result requires some additional learning-theoretic (PAC-Bayes flavored) arguments (the main idea is that if a function has good expected performance on a random dataset that is much larger than its description length, then by simple entropy/ counting arguments it must generalize well; though deducing infinite-width generalization results involves checking that these arguments are in some strong sense robust to cloning and subsampling).
While the above sketch can be converted to a proof in the rich limit (where Bayesian learning reduces to minimizing the weight norm of a successful learner), some additional arguments are needed in the case (whereas infinite-width learning at collapses to Gaussian process learning, equivalent to standard scaling).
For the mean-field regime we need to track entropy in addition to minimizing weight norm. The hardest case turns out to be exact mean field , where the entropy contribution to free energy scales (as ) at exactly the same rate as the weight norm's contribution. The resulting math is much more involved, and the behavior that we end up seeing is an interesting mix of weight norm-minimizing phenomena in big or "relevant" data directions vs. macroscopic isotropic noise phenomena (analogous to standard scaling/ kernel learning) in small or "irrelevant" data directions, with the cutoff set by the accuracy scale . This is an interesting new structure somewhat related to Nischal's 2-particle picture above.
A somewhat technical footnote: the linked paper works out the most interesting, "critical" mean field regime called the Mei-Montanari regime, where the proof becomes significantly more involved. Here we sketch the results and the proof focusing on a simpler case of "ultra-rich" mean field scaling , which is worked out in an upcoming companion paper. The results in the two cases are roughly the same, with a couple of differences we point out along the way in footnotes. The takeaway for readers without extensive context is that as we vary between different mean field regimes associated to different choices of a scaling parameter, one choice requires more work than others but all ultimately have the same complexity properties.
A technical footnote: here we are giving the weaker result that applies both at critical and at rich mean field scaling (see previous footnote about the scaling exponent ). In the case of rich scaling our result applies to any initial width , not necessarily polynomial, and moreover we deduce polynomial data generalization at all widths above some polynomial width . We expect the same stronger result to hold in the critical regime, but it needs a technical strengthening of a physics-flavored lemma which we currently do not know how to prove.
Here one should use bounded activations and sufficiently large error bounds (i.e., high temperature), and be careful with depth instability; this is explained more below.
Here I develop a simplified economics-inspired picture of such structure similar to the one in this post with Lauren Greenspan. The formal picture with concentrating Gram-matrix fields is explained in various sources, for example these lecture notes. I hope to give a simpler sketch of the math in future posts with Lauren (and Nischal is also working on some excellent notes related to this).
This is a post explaining my paper with Kaarel Hänni on complexity of infinite-width networks. I will explain the result, why it matters, and how the mathematical idealizations can interact with real structure in neural nets. This leads to some threads I am excited to pull on more in the future, via some new speculations on where interpretable structure can live.
Introduction
Our paper to some extent (and up to some important details) concludes an analysis of generalization complexity in Bayesian infinite-width neural nets. It resolves the question of when a function can be learned in polynomial sample complexity (in complexity theory, polynomial size is a fundamental asymptotic size class that means not too large. Informally, things outside the polynomial class are understood to be effectively unreachable/ uncomputable at realistic scales). More precisely, it asks whether (or when) generalization properties in infinite with neural nets reduces to generalization in finite width neural nets.
Here a crucial observation is that there are different infinite-width limits even within the same architecture, and they lead to different generalization behavior. The difference between these limits in Bayesian learning is in the choice of the learning prior (analogous to initialization or, more directly, to weight regularization in SGD learning) as a function of width. While seemingly a technical choice, this turns out to be extremely important for the infinite-width phenomenology. The relationship between initialization and generalization is familiar in usual SGD learning: for instance, the amusingly named omnigrok paper (Ziming Liu et al.) shows that modular addition generalizes better under small initialization (or strong weight regularization), and similar phenomena are observed in Bayesian learning where stronger regularization (corresponding to smaller initialization) tends to generalize better. For example the paper that explains Bayesian generalization in modular addition implies that changing the regularization is one of the levers that interpolates between the memorizing and the generalizing regimes.
In the infinite-width limit, we are inherently looking at a family of models that depend on a hidden-width dimension
Details on complexity and setup (hidden by default)
A few details here to fix concepts for people who want to track the rigorous results.
When we talk about polynomial complexity, we mean that we ask for bounds on this parameter which are some fixed polynomial in terms of the input dimension
We assume that our activation function is bounded. This assumption is relatively harmless: e.g. for any given Relu neural net, we can replace the Relu by a clipped Relu so long as we upper-bound at a sufficiently large value, and empirically using bounded activations (such as sigmoids) does not significantly affect performance.
Most importantly, we view the depth
The result
Once the technical dust settles, we have at least two different settings to theoretically investigate generalization behavior of overparameterized neural nets in the limit: the mean-field setting and the standard setting. Our paper implies, essentially, that a very wide overparameterized neural network generalizes poorly in the standard limit and generalizes "as well as possible" in the mean-field limit.
The first part of this statement is known (and we cite it instead of re-proving it). In the infinite limit of standard scaling, learning reduces to linear regression in a certain nonlinearly re-embedded data manifold associated with a so-called Gaussian process kernel. Crucially, the re-embedding depends on the training data
But the mean-field limit is known to develop deep, layerwise structure that depends on
The question our paper resolves is how much of the deep structure that can develop in a narrow network can survive in the infinite-width deep limit. And our answer is, informally, "all of it". Here I am of course being pithy. We don't really have a good notion of "structure", but if we did there would certainly be structures (mechanisms, circuits, features and so on) that behave differently at finite (polynomial) and infinite width. The thing we actually show is that if sufficient structure can be found in a large dataset to generalize well at some finite (more precisely polynomial) width in an underparameterized setting, then the same good generalization behavior survives to infinite width.
More precisely, one of our main results is the following. Fix a boolean[2] target (or label) function
Conversely, we show that if the infinite-width model generalizes after polynomially many data points then there is some polynomial width
Here our notion of generalization is worth unpacking. Given a Bayesian model (determined by the weight scaling, the width
This value has a sensible extension to width
In other words, we show that "learning polynomial data gives small error" is true at infinite width if and only if it is true at polynomial width.
Why think about infinite width
Before moving on to details and caveats (though note that there is a particularly big caveat associated to the asymptotics w.r.t. depth), let's step back and see why this is interesting. After all, no one wants to train infinite-width networks. It's also not true that one should deduce immediate insights here about empirical neural nets from this, for example that "more regularization is good for generalization always". First, empirical models are not "close" to the infinite-width limit, and even when infinite-width analyses are stretched down to realistic width, many new phenomena emerge (some kinds of emergent regularization, certain instanton phenomena) and intuitions don't directly carry over. Second, empirical models are trained dynamically (via SGD or Adam or similar) and Bayesian learning results must be understood with a grain of salt.
Instead, I think that there are two main reasons why our result is interesting:
I find the promise of item 2 particularly exciting here. Infinite-width limits are physical objects: they can be described by a certain set of integral equations on fields, which are associated to training data points (rather than weights). While the resulting object has many parameters (called observables), it gives a new perspective that naturally incorporates tools from statistical field theory of many-particle systems. In particular, the ability to control sample complexity opens up a new perspective related to scale and renormalization. The key upshot of our results is that a simple training objective leads to a simple statistical description, in a way that is independent of width. We discuss this idea and a few speculative implications in the following sections.
Small tasks have small complexity
Throughout this discussion, we are carrying around two numbers which can be relatively big. The first is the training data size
The upshot is that (modulo some discussions of coarseness and error bounds), the function learned after seeing any large data fraction is equivalent to the function learned after seeing
Parallel (or superposed) tasks have emergent regularization
Now it's likely that for a very complex task like language prediction, the effective data complexity
A highly simplified model here is to view the subproblems as existing in parallel, or in superposition. If a single component is simple enough and we learn just that component in a mean field-regularized fashion, the resulting physics should fully predict the weight distribution as explained in the previous section. However we cannot deduce that the collection of parallel circuits are given by some simple combination of the resulting mean fields. When a neural network performs multiple tasks in parallel or in superposition, these tasks are forced to share resources such as weight norm and entropy. But if we're lucky we can predict the tasks as effectively executing some set of superposed infinite-width limits once we incorporate the shared resources via emergent regularization. My notion of a shared resource is intentionally left vague, but an example of a problem where this can be fully worked out is the sparse denoising model here. A particular shared resource that must always be tracked is weight norm: since different parallel circuits must share weight norm, we expect a tuple of many circuits to mutually regularize, leading to an emergent regularization phenomenon where the effective regularization felt by any particular small circuit is much larger than what it would experience if it were the only thing being learned. This leads to a situation where, for example, under standard scaling a superposition of many circuits behaves as if each circuit were undergoing mean field scaling by itself.
Extensions and limitations
It should be clear that the above discussion leads to a particular orientation on neural net interpretability. One can conjecture that a large neural net decomposes into small simple circuits, each of which has low complexity (low effective width-independent complexity
At the same time I don't expect this to be the full story, for several reasons.
First, the separation of phenomena into mean field and Gaussian is a false dichotomy, and it is possible to come up with nice examples where both mean-field and Gaussian process phenomena interact in a significant way and must be understood together to describe the system. A complete analysis of the sparse denoising theory (not yet published) should be such a case. A general study of such phenomena, given by a 2-particle mean field theory, is related to past and upcoming work by my colleague Nischal Mainali. A related phenomenon is explained in the "Proof idea" section below.
Second, neural nets learn dynamically (SGD and so on), whereas the mean field limits are cleanest in Bayesian learning. This is not as bad as it might seem, since much of the mean-field analysis of neural nets applies also to SGD (and is believed to apply to other gradient-based optimizers like Adam, etc.). In fact, Kaarel and I are also discussing a version of the paper that would apply for dynamical learning, with additional caveats and details.
Third, mean field-style analyses may break down when studying underparameterized (or more generally "not-very-wide") networks via a neuron specialization degeneracy. Here instead of interpolating between different simple circuits in a superposition-like way, neurons may decide to specialize to a single circuit, or more precisely a circuit may depend on a single neuron or a small set of neurons in a robust mechanistic way without redundancy; this would break a fundamental assumption necessary for the standard infinite-width techniques. This phenomenon and how it fits into the various infinite-width limits is discussed in particular in the paper "Mitigating the Curse of Detail" by Noa Rubin et al. In fact, there is much evidence that this kind of monosemantic behavior does occur in some circuits (a bunch of classical examples in vision models are provided by Chris Olah's interactive account here). Phenomenology of this kind is genuinely new relative to the thermodynamic phenomena that occur in infinite-width limits, and is more or less directly analogous to instanton symmetry-breaking effects. Exact predictions for such systems would likely require a separate analysis similar to work of Saxe on symmetry breaking (see for example Saxe-McClelland-Ganguli), or the analysis of symmetry-broken instanton phenomena in physics.
For this reason and others, I don't expect a complete theory of interpretability of finite networks which is reducible to "mean field theory plus inductive regularization phenomena". However I am very hopeful that some of the simpler and more polysemantic mechanisms in large neural nets (such as induction heads) can in fact be explained in this way, or at least use the same mathematical mechanisms as mean field learners. Optimistic signals are that this is true for modular addition (where both narrow SGD and infinite-width mean-field models learn the same "clock" Fourier-basis mechanism). We also see a deep-network example of this (mechanistic agreement between narrow and mean field solutions) in an example of an SGD-learnable deep boolean XOR-like task in upcoming work joint with Sergey Alexeev (whose MARS project I am mentoring).
MFT and neuron independence
A second aside before moving on: often when I talk to people about mean field theory, they assume that this is a simplistic model that postulates independence between neurons, and hence cannot capture phenomena that require collaboration between many neurons. Under such a model it seems strange to them, and perhaps even un-believable, that a mean field formalism can capture the full complexity of algorithms learned by a neural net.
I think that this misconception makes sense, and is important enough to devote a section to.
The simple answer is that while the mean field limit does imply a certain kind of neuron independence, the independence is conditional, on a relatively complex self-organizing "macroscopic" structure which (as we show in our paper) is powerful enough to encode algorithms of arbitrary complexity.[5]
Informally, a first-order way to think of mean field theory is that neurons have jobs. More precisely in a mean field cartoon, each of ~infinitely many neuron indices at a given hidden layer is labeled by a "job" latent from some finite set. The weight between a pair of neurons at adjacent layers is then drawn as a (statistical) variable that depends only on the two jobs, and not the associated indices. Thus each neuron's preactivation is a weighted linear combination, associated to its job, of activations from job postactivations of neurons at the previous layer. The key point (and the key reason why mean field techniques are interesting for interpretability) is that the number of jobs can be large, but is in some sense finite, depending on the number of training data points; while the number of neurons is assumed to be infinite and indistinguishable. Thus the onus of interpretability is moved away from neurons and into fields, which I called jobs. Since the number of possible fields scales with the number of training data points, this formalism is able to represent very complex mechanisms where neurons at different layers collaborate. For example if a neural net is learning a finite boolean circuit, its jobs at any given layer can be associated to the finitely many hidden gates of the circuit (and this is heuristically what we observe for XOR), with each neuron type reading the two neuron types associated with its parents in the circuit. The full formalism of mean field has its interacting fields indexed not in terms of informal "job labels" but in terms of activation kernels, and a rough translation between these two points of view is a central intuition pump for our paper.
Crucially, the dimensionality of the field theory (the "maximal number of jobs") is related to the training data dimension. As I explained above, our results imply that structure saturates at some polynomial
Details and caveats
I'll explain an intuition pump for our proof in the section below, but first here I will point out the assumptions we make on the model, some of which are cosmetic and others which are significant.
The main cosmetic assumptions we make (which we think can be replaced by weaker assumptions) are:
The bounded-activation assumption is there to avoid allowing neuron specialization in the infinite-width limit, and the finite data distribution on the cube is useful (more fundamentally) to keep the data bounded and, less fundamentally, to work with a finite total data distribution of size
The most significant issue in our paper is that the various bounds in our paper scale very poorly in depth (we use the letter
where
The state of our rigorous mathematical knowledge currently has this ugly factor of
Proof idea
Our result works, to first order, by essentially formalizing the idea of "jobs" in the section on neuron independence. This is easiest to frame in a certain richer version of mean-field (where we move an exponential parameter
\[W^\ell_{\text{cloned }\times N} = \frac{1}{N}\begin{pmatrix}
W^\ell&\dots &W^\ell\\vdots&\ddots & \vdots \W^\ell&\ldots &W^\ell
\end{pmatrix}.\]
Taking
While the above sketch can be converted to a proof in the rich limit
For the mean-field regime
A somewhat technical footnote: the linked paper works out the most interesting, "critical" mean field regime
The Booleanness here is a cosmetic assumption made for technical ease, and can be weakened.
A technical footnote: here we are giving the weaker result that applies both at critical and at rich mean field scaling (see previous footnote about the scaling exponent
Here one should use bounded activations and sufficiently large error bounds (i.e., high temperature), and be careful with depth instability; this is explained more below.
Here I develop a simplified economics-inspired picture of such structure similar to the one in this post with Lauren Greenspan. The formal picture with concentrating Gram-matrix fields is explained in various sources, for example these lecture notes. I hope to give a simpler sketch of the math in future posts with Lauren (and Nischal is also working on some excellent notes related to this).