Data as particles and mean field theory abound

Epistemic status: Much of this post is a reframing of known results using slightly different tooling. The framing of "data as particles" has been conceptually useful for me so it seemed useful to share. Some steps in the derivations are stated and handwaved rather than writing out in full where the steps are generally standard. The dynamical mean field perspective presented later is part of a broader research agenda, and a paper deriving more formal results about it will be out soon!

Author Note: The automated LLM written detection really wanted to consider this post as LLM written. So maybe putting this at the start will help. This is not written in any substantial way by an LLM. It might be slop, but it is human generated slop.

Very early in my machine learning career I worked on group equivariant neural networks. These for a while became the dominant kind of neural networks in things molecular property prediction with the early advent of things like Schnet. This, through a series of events pulled me into the world of the quantum many body problem, and then into a quantum chemistry lab, where I would do my PhD (though not technically in quantum chemistry). This corrupted me though, because suddenly everything started to look like a many body problem.

Such a perspective though has lead to some nice insights, and I think it offers a way of reframing some results in learning theory in a way that opens up some new tools or ways of thinking about learning. In particular, I have found it useful for studying learning under gradient based optimizers, and thought other people might find it useful as well. While more formal research papers building around what is in this blog post are coming, I think the nature of formal research papers tends to make the core takeaways and intuition somewhat obscure. This post attempts to alleviate some of this by being a bit more pedagogical. There are a few sections which can be summarized as follows:

We explain how one can think of datapoints as particles, and the model parameters as defining bonds between them.

Using this perspective, we derive a mean field theory of the data, and show that this leads naturally to the Bayesian free energy of singular learning theory.

We show how one gets a dynamical mean field theory of the data when considering the behaviour of SGD in the large data/small learning rate limit.

Taking this perspective, we show that the dynamical mean field theory is effectively second order, and relate this to the static mean field of singular learning theory.

This post does get a little mathematically dense in places. We do try and make it somewhat more legible for those unfamiliar with mean field theory/many body physics though!

In some sense the starting point for this is the simplest object in machine learning: the loss function. For a model with parameters and (for pedagogy) let's consider a finite dataset with the drawn i.i.d. from some data distribution , so the empirical loss becomes

with the population loss defined via the standard limit.

One thing about the loss that sometimes goes underappreciated is that it looks a lot like what we might call a "parametrized Hamiltonian". To see this note that we can read the empirical loss as a total energy written as a sum of energies of individual "particles" in the system. Each datapoint is a particle, and is its energy when the model has parameters .

Notice that every particle shares the same . To see why this is important imagine a black box with three lightbulbs on top who are wired in some unknown way, whose brightness reports the energies , , , and a panel of twenty unlabeled knobs. We know nothing about what the knobs do, but we want to understand how the lightbulbs are wired. Now imagine we start twisting the knobs and just seeing what they do. Knob 7 brightens A and B a little and dims C. Knob 12 brightens all three by different amounts. Knob 3 mostly does something to C. This is kind of messy, hard to interpret exactly what each knob does. So we start just kind of going nuts, twisting and turning the knobs in random ways and looking at the outputs. We sort of ignore exactly what knobs do what, and just kind of see how the lightbulbs change together. Looking at the bulb, one can see that maybe changes which change A are very likely to change B, maybe C always dims if both A and B brighten, etc. So we may not know exactly what each knob does, but we can kind of infer how strongly each bulb is wired to each other bulb

We can pretty simply put some math to this for the loss. Starting with the simple "two body" story, imagine we perturb and for simplicity just take . The bond between two particles is the covariance of their energies under this perturbation:

Averaging over random perturbations sums out the parameter index , leaving you with indices in the data space, so the parameters are exactly what gives the relationships between particles.

The above construction only captures two-body structure. It tells us whether two datapoints tend to respond together if we perturb parameter space. But there can be more complex relationships between several datapoints which cannot be reduced to simple pairwise bonds. For example, perhaps and are only weakly correlated, as are and and and , but whenever the loss of and both increase, the energy of almost always decreases. There is then some structure belonging to the triple as a whole which is invisible if we only inspect the pairwise bonds.

This is the basic idea of a higher-order correlation. A three-point correlation asks how the energies of three particles vary together, a four-point correlation asks the same for four particles, and so on. There is an important subtlety, however. A raw three-point correlation can contain structure that is already explained by lower-order relationships. What we really want is the part of the three-particle relationship that is not captured by any of the lower order junk. This is captured by the joint cumulant (which we will talk about more later)

and similarly at higher orders. Cumulants therefore give us what in physics are often called connected or irreducible -body correlations.

The word irreducible is reflects the fact that we only want to capture only the parts that cannot be expressed by lower order terms. For example, at second order,

so the bond is the irreducible two-body interaction. At third order, measures whatever remains of the joint behavior of the three particles after the lower-order structure has been accounted for. At fourth order we ask what genuinely four-particle structure remains, and so on. Each order therefore answers roughly the question: what is new here that I could not have inferred from the orders below?

In this sense, the data distribution define a whole hierarchy of interactions of different orders, all mediated through the same shared parameter vector . Two datapoints can be coupled because changing changes both of their losses together, but collections of three, four, or more datapoints can also have collective responses which cannot be decomposed into pairwise effects.

Next, we consider what happens if instead of considering how changes with respect to , we consider how responds to changes in the population loss .

Moving from looking at bonds between individual particles to bonds between particles and some bulk structure made of many other particles is known in physics as amean field approach. For those unfamiliar, the idea of a mean field approach is actually rather straightforward. Suppose we have a huge number of particles in a system. Taking a particle and looking at how it interacts with every single other particle can be difficult and unnecessary. What we then do is replace the infinite collection of individual particles by their average effect. If, say, all the particles are magnets, to know how behaves from the electromagnetic interactions with all the other , we don't need to model the individual magnetic interactions. Instead we just pretend that all the other are actually just a big field. This is why it is called a mean field, as we can replace complex interactions by creating a simple field to represent them.

Each is like a weight that allows us to adjust the influence of a given datapoint. Setting every gives the log partition function of the posterior, . So why is this called the cumulant generating function? Consider the first derivative, , which gives the mean energy of the particle. The matrix of the second derivatives gives the bond matrix, and defines the variance of the system. The third derivative defines the skewness, which is the three body interactions. So the th cumulant defines the body interactions. The general notation for the th cumulant tends to be .

The cumulant generating function we just wrote out seems made of individual data interactions. However, we can actually decompose this into a "bulk" part and an individual part. We show explicitly how this works in the next section.

The general trick in a mean field theory, as discussed, is to replace particle-particle interactions with particle-bulk interactions. In our case, the "bulk" is the population loss. How do we then move from the interactions between per-sample losses to the interactions between the bulk and the samples? To show this we are going to start with the lowest order data-data interactions. First, we split each particle's energy into a population part and its own deviation:

Tildes always mean centered over the population, so has mean zero under for every . The total energy becomes

We can now build out our mean field. To do this, substitute into both slots of and expand. Covariance is bilinear, so two sums of two terms give four pieces describing interactions:

The first three terms all pass through the bulk. That's the same move as in the magnet picture: particle doesn't interact with directly but with the field the bulk produces, and does the same. To write those three terms more compactly, expand the coupling to the bulk the same way:

So already contains two of our four pieces: the bulk with itself, and with the bulk. Likewise contains the bulk with itself and the bulk with . Adding therefore picks up all three bulk-mediated terms, but it counts twice. To avoid overcounting we subtract the overlap once and get:

The mean-field bond is built entirely from each particle's coupling to the bulk, so it needs only one number per particle rather than one per pair. The residual is whatever two particles share that does not go through the bulk, and it is what one would call the "two-body correction."

If we think about this, it tells us that two datapoints can look correlated for two very different reasons. There might be some general set of structures in the model (like induction heads) that broadly lower the population loss because a lot of samples get a slightly lower loss by exploiting this mechanism. This looks like bulk coupling. The other possibility is two samples might rely on much more specific, targeted structures like the knowledge of the birthdays of all the royal family. These look much less like bulk coupling as probably not much of the bulk actually needs this information to perform well.

The bond decomposition was the bulk-particle case of something that holds at every order, meaning that we can decompose interactions into bulk interactions with interacting particles/datapoints. The math here looks a bit messy, but it really is just the higher order version of the previous section.

Here the cumulants are taken under the prior . Now substitute into every slot. Joint cumulants are multilinear, so an th-order cumulant splits into pieces, one for each way of filling each slot with either or a . For , these are exactly the four routes from before: bulk with bulk, with bulk, bulk with , and with .

To organize the pieces, sort them by , the number of slots that carry a particle's own identity. If of the slots are individual, there are ways to choose which ones. Each of the remaining bulk slots contributes , since doesn't care which index it came with. Using , this gives

Every term in this sum is one kind of interaction, and the three pieces of the decomposition are just three ranges of . This is all pretty nasty but we can actually put some meaning to what these orders actually are.

This piece depends on the weights of individual samples only through their total . It cannot tell which individual particles are present, only how much weight the bulk carries.

The terms have no bulk slots. They are the direct correlations between private deviations at every order, and they also get their own cumulant expansion:

The terms form . These are interactions in which particles and the bulk both take part. The terms, , describe a single particle coupling to the bulk. They are the all-orders version of the piece inside from earlier. Terms with describe several particles interacting while the bulk also participates, so some genuine many-body structure lives here too. In physics we would say it is "dressed by the bulk". So now we can write the whole cumulant generating function like:

Now here is a cool trick. The cross terms can be absorbed exactly by letting the bulk set the background measure, and the remaining fluctuations of individuals are then measured on top of it:

This decomposition looks like it's doing something clever, but it's actually just a split of the reweighting. Since every particle's energy is the population loss plus its own deviation, , the total reweighting splits cleanly:

The first piece only knows how much total weight you've put on the data, and it reweights the prior the same way no matter which datapoints carry that weight. The second piece is where individual identities live: each datapoint's weird little quirks, weighted by its own . Because the exponential turns this sum into a product, we can do the reweighting in two stages where we first let the bulk reshape the prior into , and then average the individuals over that reshaped measure rather than the original prior.

Seen this way, the two terms answer two different questions. asks: if the data were a featureless bulk with total weight , where would the parameters end up under the Bayesian posterior? The second term asks: given that the bulk has already put the parameters there, how much do the specific datapoints you actually have, with their individual weirdness, still move things around? This actually shows up in singular learning theory.

The cross terms are precisely what converts connected cumulants measured under the prior into connected cumulants measured under the mean-field measure . Now taking , that measure is the population posterior normally studied in SLT. The free energy therefore splits exactly into a mean-field part and a fluctuation part:

where is the learning coefficient (RLCT) and its multiplicity. The classical result of SLT is asymptotically that

As . We can interpret this as saying that the bulk determines the leading asymptotics, and the identities of individual particles enter only at . So, SLT is something like a theory of the bulk.

So far the mean-field construction has been static. We replaced the collection of individual data particles by the average field generated by the population. Learning, however, is a path through parameter space. That is, we we have a path through parameter space then our mean field construction must evolve in time. This sounds like a dynamical mean field theory. So how do we go from the static mean field to the dynamical one?

Before giving some of the more mathy bits, the core construction admits a nice intuitive description. The first thing to do is imagine we have a fixed starting point , a data distribution , and some fixed number of steps we are going to run SGD for, and take the learning rate to be sufficiently small. To start off, suppose we run SGD with a batch size of 1 for steps, and we look at the distribution of all the paths it took during that time. We then do the same thing for batch size 2, then 3, and so on, taking . As this happens, the distribution of trajectories starts to asymptotically concentrate around the deterministic gradient descent path.

To see why this is a mean field theory, think of each datapoint's gradient as a force: a push on the parameters in the direction that would lower that datapoint's loss. Just as with the loss, we can split each push into the bulk's part, the population gradient, and the datapoint's individual part, its deviation from the bulk. Every SGD step is then the bulk's push plus the average of the private pushes of whichever datapoints happened to be drawn. The thing then is that the interactions between each datapoint on the loss level is evolving, and that evolution is actually governed by how their gradients interact. For example, if two different samples both give the same gradient direction, following the gradient of one basically must pull the other one along, meaning their evolution is coupled.

Below we formalize this and show that it is effectively what gives the standard stochastic differential equation description of SGD. For those familiar with DMFT we give the basic outline of the two-time force covariance and the response kernel.

which is the force exerted by a data particle on itself/others, and suppose datapoints are drawn from a population distribution . The population "force" is the gradient of the population loss

A particular sequence of batches therefore determines an entire trajectory through parameter space. With , each update is determined by a single particle, so there is a large cloud of possible trajectories. With , each step averages two forces and the cloud becomes narrower. With it narrows further. As becomes large, the stochastic gradient concentrates around the population gradient and the bundle of SGD trajectories collapses toward a single mean path.

The stochastic gradient has therefore been split into a population mean field and a fluctuation around it. This is actually the standard gradient noise description of SGD.

This tells us how larger batches suppress the fluctuations. Specifically, it says that averaging over datapoints suppresses the force fluctuations like , or equivalently suppresses their variance like .

Notice that the stochastic force is actually generated precisely by pairwise disagreement between data particles. If and are independent draws from , then

If almost every particle tries to move the parameters in the same direction, then replacing one sampled particle by another barely changes the update and the trajectory cloud is narrow. If different particles pull strongly in different directions, then which particles happen to enter a batch matters and the cloud becomes broad. At small batch sizes, any change can be very broad. So, for large enough batches the fluctuations around the average path are generated by two body data-data interactions. Next, we are going to write this as a DMFT.

So, what we would generally do here is move to the SGD SDE limit and work in that for formal reasons. But, for the sake of this blog post it is sufficient to just write it out as generic weakly fluctuating dynamics:

where for finite-batch fluctuations (where we drop the learning rate for simplicity) and just assume we have absorbed a very small learning rate into the gradient and noise. The choice of is because this is essentially like a semiclassical limit. The sign of is conventional and will not matter for its covariance. Moving to the SGD SDE is formally straightforward if one desires.

As , the path approaches . We then write the expansion around the path as

Here, we use a bit of a trick which is actually important for downstream theory, but whose formal impact we will leave for later work. Expanding the dynamics to first order gives

A force fluctuation applied at time is therefore propagated to time by . If

so the same data particle can produce correlated forces at widely separated points on the path. This is one of the reasons why the DMFT perspective is cool. It kind of gives the natural construction of influence functions over time.

For those familiar with DMFT, a more fully self-consistent DMFT-like theory appears at the next orders, when the fluctuations modify the statistics of the force and generate memory or self-energy kernels. But we have found that the leading order theory is actually sufficiently rich for our purposes. Also, there is a relationship with the NTK literature, and the neural tangent hierarchy, but, since this is meant to be a simple blog post we are going to leave that for some future work.

One way to interpret this is that under the usual finite-moment and independence conditions on the data is that the Gaussian noise-Hessian closure description is actually very good theory for certain parts of training. The (anistropic) Gaussian noise tends to hold across all of training with the exact distribution changing over time. Now, while this is not what this post is about, the local quadratic contribution of the dynamics can be thought of as describing the movement of training during a "phase". The is, the local dynamics admit such a description apart from isolated times during training which look like a phase transition. This is the subject of an upcoming paper, and is one of the nice things you get out of the DMFT perspective

It is important to note that the statement is not that the loss itself is always well approximated by a quadratic function of the parameters. It is that the collective force fluctuations around the mean trajectory become Gaussian, and hence quadratic at leading order, because they are produced by averaging many data particles, and that at standard points during training the leading order term of the deterministic dynamics dominates. This gives what one would call a second order effective theory of the dynamics, which can be explored and manipulated.

The above constructions present the basic framework underlying some of my personal research directions as of late. Whether or not they are useful for anyone else remains to be seen but if you have made it to this point, thanks for reading!