Neural Networks Are Elastic Origami! [Prof. Randall Balestriero]

At a very early stage during training, the training accuracy—or whatever the training metric is—grows and then starts to plateau. You might think, “Okay, training is done,” but the test metric is still not at its high point. If you keep training much longer than you normally would, although the training metric continues to plateau, the test one will suddenly start increasing. Your network will start assembling its weights so that it can extrapolate to test samples.

The high-level summary is that grokking was observed in some very specific settings: a specific task, a specific model, and a specific weight initialization. What we tried to look at were more general settings and whether the same delayed property emergence was happening. It turns out that it is not happening in terms of training and test accuracy. Those two metrics evolve at roughly the same rate and plateau at roughly the same time, but when you start looking at test accuracy under adversarial noise, that is when you see the emergence of robustness to adversarial noise on test images, much after you reach the plateauing area of clean training and clean test accuracy.

When you have a neural network with any piecewise-linear nonlinearity—it can be ReLU, leaky ReLU, sawtooth, or anything that is piecewise linear—what it basically does is cut up the input space into linear regions and then linearly map each region to the output. The theory says that you have this spline, where the neural network boundaries—the spline boundaries—are defined by the network weights, and the mapping is also defined by the network weights.

It is a lot like origami, except that the origami is not just folded; there is a little bit of stretching as well. I would say it is elastic origami. The neural network takes an input space, or an input domain, and plays origami with it, turning it into some warped space. Then the output, because it is linear, is going to be a hyperplane in the embedding space, so it cuts this origami up with a straight line such that the decision boundaries separate the classes.

We have seen in experiments that, for some samples, there is more partitioning around them. That would mean there is more nonlinearity needed there, or that the function is more complex. There are also parts of the space where the cutting up is sparser. That is a harder region, whereas where there is more cutting up, the model is more prone to adversarial attacks, because if you move a little bit, you are crossing more nonlinearities. More of the activations will be changing in the network.

You can take a couple of different 2D slices to estimate what the partition statistics look like. The more important thing is estimating the partition statistics around a locality. These statistics can tell you how the model behaves differently for different samples. For example, if you have a couple of samples from a demographic to which the neural network assigns less expressivity, you can say that the neural network is algorithmically biased toward those samples because it is not as expressive in those regions. For other samples, or for a particular class, it might assign more regions.

Based on these statistics, you can infer how the model behaves differently for different sets of samples, instead of coming up with a general notion of what a neural network generally does. This allows us to think about what a neural network does in different parts of a region on the manifold, or off the manifold. These notions could help us converge toward a more general understanding of neural networks than just a one-line statement that this is what a neural network does.

One of the basic results in spline function approximation, which has mostly been studied for 1D-to-1D regression, is that when you design a spline, you have to choose what your partition of the space is and the degree of the polynomial that you use in each region. You might say, “Okay, piecewise-affine functions are not really the best. You just have an affine mapping per region, which is not rich enough to represent a lot of variety within each region.” But it turns out that if you have to choose between positioning your regions well—finding a good partition—and increasing the degree of your polynomial, it is much better to fine-tune or position your partition according to the data, even if you remain piecewise affine.

You can get extremely good approximation power, although each region contains a really simple function. This is much better than keeping the partition the same and just increasing the degree of the polynomial. The takeaway is that not having smoothness, and being merely piecewise affine, is okay. In fact, it is optimal if you can position those regions well.

That is what you do with networks. The partition adapts when you train the parameters of your model because the partition and the affine mappings are tied together. By learning one, you learn the other. In practice, we see that most of the regions get concentrated around the training points and are also extrapolated based on the rule given by the architecture. If you go to other parts of the space that are neither near the data distribution nor near the extrapolation rule, you have much less regional concentration. This is important because the smaller the regions, the more precise your approximation is.

Professor Randall Balestriero, welcome back to MLST, and congratulations on your new role at Brown.

Thanks very much. I am happy to be back, especially to speak about the latest research and splines. I am very happy to be here.

You invented, or perhaps co-invented, the spline theory of neural networks, which revolutionized my understanding of deep learning. I think we should have a quick refresher. What do we mean by the spline theory of neural networks?

First, just to put some things in a better context and avoid some legal issues later, splines have a very rich theory. They have probably been used since the 1980s, or maybe even before, for function approximation. Most of the research was done for 1D, 2D, or maybe 3D input spaces, because that is where most of the function approximation was needed, perhaps for partial differential equations and things you could observe.

What we did around 2018 was try to understand current deep networks. At that time, it was mostly convolutional networks, ResNets, or MLPs with ReLU activation, max pooling, and these types of nonlinearities. It turns out that when you have this type of operation—an affine operation such as a dense mapping or convolution, followed by ReLU or max pooling—and you keep interleaving and composing those layers, the entire input-output mapping is itself a continuous, piecewise-affine spline.

What that means is that the input space of the network—perhaps the space of images, if you do MNIST or CIFAR classification—is a huge, high-dimensional space. It is cut up into polytopal convex regions, and within each of those regions, your network is just an affine mapping. Overall, it is continuous, but within a region it is simply an affine mapping.

You can characterize the geometry of those regions: where you have more regions than in other places, what your network is actually learning, why it is able to extrapolate, and what the impact of the architecture, regularization, and other factors is.

There is some quite technical language here, so I really want to drive home to the audience that it is a bit like when we train a neural network, we train it to become a honeycomb.

Yes.

There is this structure, a bit like a lattice, and inside the holes of the honeycomb are the decisions that the neural network makes.

Exactly. Another parallel would be K-means, or some clustering method. There, you learn a partition of your space depending on which cluster your point is assigned to. You have this sort of structure with regions. With K-means, the mapping is piecewise constant: within the region, or cluster, you are assigned to the same cluster, and when you go to the next region, you are assigned to another cluster.

This is the type of partition geometry you get with deep nets. It is very similar: you have those regions as well, but within each region, instead of being constant, you have an affine mapping. You also have an extra constraint: the regions are not independent from one another. They can learn even when you are very far away, in an extrapolative regime.

The type of geometry you get is really akin to what you have with K-means or K-nearest neighbors. The key thing with deep nets, and one of the key benefits, is that the way you learn your partition—how you learn those regions—is not restricted to where you have data. It is defined everywhere in the space. That is why you get much better extrapolation performance, even in high dimensions, which you do not get with K-nearest neighbors, for example.

One of the reasons I like spline theory is that a lot of people say neural networks do this magical emergent reasoning. When you understand a neural network as computing these spline partition boundaries, it is a little bit like a locality-sensitive hashing table, which does not seem conducive to reasoning or doing any different kind of computation.

I think there are 2 points to make. First, as you said, you have this sort of locality-sensitive hashing or template matching. You just try to fit locally what the region geometry and the affine mapping are. You could think of this as really ad hoc and brutal.

But where you actually have some sort of reasoning or intelligence emerging is in the fact that the way you learn in one part of the space impacts how you learn in another part of the space, even if you do not have data there. That is where things become more complicated than K-nearest neighbors or K-means. It may even be more human-like, because from one example you learn something that you can reuse on another example that you did not see during training and that is very far away in your space.

It is not ad hoc in that sense. It is able to extrapolate hidden rules in a very efficient way.

Now we are going to get to the most exciting part today, which is that you have a paper. You wrote this with Imaad, and we have some great content with Imaad coming out as well. It is called “Deep Networks Always Grok and Here Is Why.” Give us the elevator pitch.

I want to thank Imaad because he has been pushing a lot, and that is how this paper got to this stage. Unfortunately, he could not be here for visa reasons, but huge kudos to him.

The high-level summary is that grokking was observed in some very specific settings as a delayed emergence of generalization. Test-set accuracy starts growing much after the training accuracy has already plateaued. But this was in very specific settings: a specific task, a specific model, and a specific weight initialization.

What we tried to look at were more general settings, such as CIFAR, computer-vision tasks, convolutional networks, and ResNets, to see whether the same delayed property emergence was happening. It turns out that it is not happening in terms of training and test accuracy. Those 2 metrics evolve at roughly the same rate and plateau at roughly the same time.

But when you start looking at test accuracy under adversarial noise, that is when you see the emergence of robustness to adversarial noise on test images, much after you reach the plateauing area of clean training and clean test accuracy. We found adversarial grokking happening in most of the settings we examined. This happens even though you do not perform any adversarial training. It is simply a consequence of very long training and the emergence of sparse solutions, where new geometric properties of your network naturally emerge.

Maybe we should linger on a couple of things here. First, let us talk about adversarial robustness. What happens when a neural network is not robust?

What we do is a very standard way to attack a network. You take an image, or whatever the input is, feed it through your network, and then, based on gradient information, determine the best direction of noise to add to the original input so that you can fool the network into predicting an incorrect class.

This is a white-box attack: you use the gradient information of your network to fool it. You can have a training set on which the network achieves, say, 100% accuracy on clean images, and then, just because of really small perturbations that you cannot see by eye, you can reduce its accuracy to random guessing.

It is a very efficient attack, and many people try to make networks robust to it through adversarial training. During training, they sample those attacks and try to make the network robust to them. Here, we show that you actually become robust to them naturally through very long training episodes.

I have always looked at spline theory as a kind of mental intuition pump for the adversarial-examples problem. We will show a graphic on the screen now, but essentially, when you look at a spline partition, there are crisscrossing, overlapping splines. It is a complete mess. It is very chaotic, and this is the reason why the networks are so brittle: it is easy to push a test example over the boundary so that the network behaves differently.

Exactly. When you see all the partition regions in this figure, what you have to think about is that when your input moves a little bit and goes from one region to another, your mapping undergoes a change in nonlinearity. That produces a change in the prediction.

The more of those regions you have, the more nonlinear kinks, or points, you have, and therefore it is much easier to perturb your network. But when you reach this very long training stage, and therefore this sparse solution, you no longer have so many regions everywhere in the space and around your training points. Instead, the regions start migrating away from the training points and test points and concentrate near the decision boundary.

Because you now have much wider regions around those points, your network is just affine within a much bigger region of the space. Therefore, it is easier to control its sensitivity to noise.

You are saying—and perhaps we should introduce grokking first and then come back to this—that grokking is this phenomenon of delayed generalization. I know that is an oversimplification, because whenever I speak to people who study grokking and say that, they tell me it is an oversimplification.

Basically, you train a network for ridiculously longer than you normally would, and then things happen later in training. Can you explain that?

On specific settings, most of which involve constrained or simple tasks, you can show that at a very early stage during training, the training accuracy—or some other training metric—grows and then starts to plateau. You might think, “Okay, training is done. I can stop training there.”

But the test metric, which might also be accuracy or something else, is still near random or only slightly higher. If you keep training much longer than you normally would, although the training metric continues to plateau, the test metric suddenly starts increasing and converges on its own later.

During training, there is still gradient information that makes your weights change. Those changes do not have an impact on the training metric, but eventually your network starts assembling its weights so that it can extrapolate to test samples.

There is an interesting broader story here, because I know you are very interested in the training dynamics. There are different stages of training. At the beginning, the network learns relatively simple features, and as we progress through the training stages, it learns increasingly complex features.

This is a little strange, though, because what you are showing is a local decomplexification. I would imagine that things would get more complex as it learns high-frequency information, but the network is learning a local decomplexification to stretch out those boundaries. Why does that happen, and what learning signal is it using to do that?

As you said, there are 2 stages of training, a bit like a double-descent dynamic. During the first stage, you start from random initialization, and your network focuses aggressively on the points. It uses a lot of regions around them, so it is a form of memorization. It is still able to extrapolate, but it is really focusing on the actual points, with lots of regions near them, a very smooth mapping, and no simplification anywhere in the space.

Once you reach this stage and keep training much longer than you usually would, there is still gradient information conveyed through your loss. The network will decomplexify around those points. The regions start to migrate and move away, and instead they focus on where the decision boundary is.

All the allocated parameters are trying to fit the decision boundary very precisely instead of remembering where the points are. Because of this, the radius of the regions increases around those points. You get adversarial robustness, and the network goes from something really smooth—almost uniformly smooth—to being piecewise constant, which is what it should theoretically converge to in the limit.

You could argue either way about whether it is complexification or decomplexification. The remarkable thing, as you say, is that it moves from looking at the training examples to looking at the regions.

There is a great figure that I am going to show on the screen again. After this grokking phase has happened, you get these partitions emerging. It looks a little like a topological map or a contour map, and the partitions look very much like a Voronoi diagram. It looks as though the partitions are equidistant, with the boundaries equidistant between the points. It looks like a mountain structure, not a single decision; it is many decisions that have been squashed together.

Exactly. When you see those huge concentrations of partition regions, all of them are squashed together around the decision boundary to bring representation capacity to the transition from one class to the other.

That is what you want. When you are near a point and in its neighborhood, you just want a simple affine mapping, or even a constant mapping. You do not need to put a lot of parameters or regions there. You want to put everything at the decision boundary, because that is where you need to go from one class to the other. You need actual curvature in your mapping.

This transition goes from uniform smoothness in the space to a piecewise-constant mapping, and that is what brings adversarial robustness. Again, this emerges only after very long training, because it is not the first thing that the network learns.

It is a latent solution that emerges after very long training, perhaps because the gradient norm needed to reach it is very small, or perhaps because it is fighting the implicit bias of your architecture. As we also showed, as a function of the strength of regularization, you can control the rate at which this emergence happens.

The more you regularize with things like weight decay, the more you fight this sparse solution. Therefore, it may not happen at all, or it may happen very late during training.

This is related to sparsity, right? We know that when we do iterative magnitude pruning, we train a dense network and then pull out all the low-magnitude weights. The sparse network is strangely more robust, even though we have taken away most of the weights.

Isn’t it interesting that after this grokking phenomenon, the network we get resembles a sparse network?

It is exactly related. In another paper, we showed that you can prune in a refined way as a means of simplifying the partition, making it focus on the decision boundary rather than the points.

You can prove a relationship between pruning, its impact on the geometry of the partition, and other methods, such as collapse in the rank of your parameters and different types of regularization. There is a one-to-one correspondence between all those things.

The nice thing is that the partition gives you a single geometric object in which you can visualize all of them and understand where and why they are beneficial. It is related to the Lottery Ticket Hypothesis and iterative pruning. When you do iterative pruning, you probably switch from a complicated solution with a uniform distribution of regions to one in which you aggressively remove most of the regions that are useless for your task—probably the ones near your points.

You bring the model closer to this later stage, but through active pruning of your model parameters or units.

How is the phenomenon related to double descent?

This is also related to neural collapse, which we mentioned before. There is a dynamic in which, at first, there is a kind of memorization, and then you start to learn to extrapolate. Here, we do not look at it in terms of capacity as the number of parameters; we look at the geometry of the partition, but you have the same phenomenon.

If you look at the local-complexity measure that we derive, you can see the dynamics, or migration, of the regions. At the beginning of training, when you start to see an increase in training and test accuracy, a lot of regions become concentrated near the points. This is like the first ascent that you get after the first descent, and this is where you get no robustness at all.

People will stop training at this point because their training and test accuracy look good, but they do not get any robustness. If you keep training, the regions suddenly migrate away. You see a second descent, and this is when robustness eventually emerges. The regions have moved away, so the ones near the training and test points have a bigger radius, and therefore you get robustness.

This gives you a new way to look at double descent or training dynamics through the geometry of the partition. That is really nice because, until now, most of what was studied involved looking at the loss function or training and test accuracy. Those metrics may be very task-specific and black-box, in the sense that you just look at a network, F, and do not really dive into it.

Now you have a new way to look at what is happening within your network in terms of this partition. One major future question is whether we can rederive most of those results in terms of the geometric properties of the partition.

We use regularizers because, certainly in the old-school view, by deliberately constraining networks to be simple, they train better. What we are talking about now is that we actually want a particular type of complexity.

I think you said to me earlier that if we are not careful about our regularizers, we might not get this effect in the first place.

There are many ways to regularize a network, including implicit and explicit regularization. Even methods such as batch normalization, which we might think are simply ways to make training easier because they normalize the activations, act as regularizers.

In another paper, we show that if you use batch normalization, for example, you actively try to concentrate the regions near your training points. When you employ techniques like this, you fight the sparse solution, and therefore you cannot get there at all, or you can get there only through even longer training episodes.

Weight decay also has a strong bias. It says that a solution near zero is preferable and brings your model to a nice solution. Usually, these regularizers try to enforce smoothness in your mapping in the L2 sense. Because of this, you do not try to reach a piecewise-constant solution, which is the one that gives you adversarial robustness and the one that you eventually reach during the adversarial-grokking phase.

That is why there are also many implicit biases introduced through standard regularization that we may need to rethink if we want to speed up the emergence of adversarial grokking.

One of the themes we will talk about today is that neural networks do incredibly interesting things later in the training dynamics. That is great if you are Meta or Google, because you have these big GPU clusters and can train neural networks beyond what is within reach of a normal person.

We want to capture that behavior and make neural networks do it earlier. I know you have already done some work with a geometrically inspired regularizer that made the boundaries orthogonal to each other. Could we design a regularizer that encourages this grokking behavior sooner?

That is a very good question. As you mentioned, right now we reach this adversarial robustness through very long training. Because we do not perform adversarial training, the robustness we get seems to be very good across different types of adversarial attacks. That is a beneficial property that people want.

The problem is that you need to train for so long, so this is not a solution that everyone can access. A major research direction is therefore how to speed up the emergence of adversarial grokking so that everyone can benefit from it.

There are many ways to build regularizers based on the geometric understanding of the partition. For example, something that is easy to compute, even for very large networks, is the distance from a point to the nearest boundary of the region in which it lives. This quantity is differentiable with respect to the model parameters and is very fast to obtain, which means you can use it as a regularizer during training.

That is just one example. There are many ways to derive a differentiable regularizer that can enforce a constraint on your partition. You could also avoid using a regularizer and instead build the architecture with those constraints. One typical example is a network with no biases. In that case, the type of partition you get is constrained to be central: all the regions are cones extending from zero to infinity.

That is an architectural constraint. If you do not have biases, you constrain your partition to always look like this, and therefore you do not need an extra regularizer. There are many ways to build hard constraints into your architecture or the parameterization of the weights, or to build a different regularizer that you can use during training.

After speaking with Sarah last week, about the EU AI Act and the executive order, we discussed how they had a hard, absolutist limit on the number of compute FLOPs—a tally of the amount of computational operations we perform.

We are in a regime where people think there is a commensurate relationship between capabilities and compute, so they just throw more compute at a problem. What we need to be doing is smart FLOPs. More broadly, we need people like you who actually have a theory of neural networks, because without the spline theory of neural networks, it would not be possible to have the mental model needed to design new regularizers in a principled way. It is so important.

One of the key benefits of using splines is that they give you not only theoretical guarantees and solutions, but also visualizations that are really easy to interpret. Even if you have never used splines before, it is easy to look at figures of the partition and the regions and see why one partition is better than another for a specific task or in terms of robustness.

The good thing is not just having a theoretical understanding of deep networks, but having something that any nonexpert can visualize, act on, and use to train state-of-the-art models better. You want it to be tractable, interpretable, and easy for everyone to understand.

As you said, this can provide strong insights. For example, in this paper we see that if you have an allocated number of FLOPs, the common approach might be to say, “Let me fit the biggest model I can on my GPU, and then use whatever FLOPs remain for training time.”

What we show is that, depending on the application, if you want adversarial robustness, you may want to use a much smaller model and allocate more FLOPs to training time instead. With this different FLOP allocation—this smart FLOP allocation—you get a model that is adversarially robust after a long training episode.

That is why we should return to questions such as where to allocate your FLOPs and what properties you want your network to have.

Talking about your results a little bit, how does it change depending on the type of problem, the dataset, and so on? What have you seen?

We saw a few trends. For example, as the size of the dataset increases or decreases, and as a function of the noise in your labels, the emergence of grokking can be delayed or not. Those are things to keep in mind, and they could later be used to speed up grokking.

For example, you might use a smart curriculum for training or a teacher-student training strategy. Those are the types of approaches you could try in order to speed up the emergence of grokking.

More and more papers are trying to look at this because all these things are intertwined. If you think again about the partition, you have to think of it as adapting to the dataset you have. That is one of the beauties of the splines that deep networks use, as opposed to standard splines: the partition adapts through gradient descent of the weights to your data distribution and the loss you have.

If you change the way your points are distributed, the number of points, the dimension in which they live, or the parameterization of your network, it will affect the partition you learn and therefore the geometric properties that are beneficial or not for you.

All those things are intertwined, and the spline interpretation of deep nets allows you to understand this relationship precisely. You can get provable guarantees about why you get one partition rather than another. Now you can make informed decisions about how to parameterize your model to reach a state that is good for your downstream task.

This might be a good time to talk about your local-complexity measure. The reason I bring this up is that, once we have a measure of complexification and grokking, we can change the hyperparameters, the prediction architecture, and so on to optimize this effect.

Downstream, perhaps we could have a principled way of designing the architecture to optimize grokking. How did you design that complexification measure?

We are trying to look at how many regions there are, or how complicated the partition is, near a specific point. A good proxy is simply to count the number of regions that are nearby.

You can think of this as roughly equivalent to how many bits of information your network has in that neighborhood. We use a proxy to obtain that number of regions very quickly, even when using a very large model. It is based on counting, layer by layer, the number of vertices within the epsilon ball.

We perform ablations to show that this is a good proxy for the number of regions. That is what we measure. In the local-complexity graphs, you see that this number has different dynamics during training.

The interesting thing is that it peaks when training and test accuracy reach their plateauing phase, but then it starts decreasing much before adversarial grokking happens. This means the metric is more sensitive to training dynamics and to what is happening in terms of the geometry of your deep network.

The metric starts showing the complexification around the point much before adversarial grokking occurs. This is useful because you have a sensitive metric that tells you how the geometry of your model evolves. It opens a new door to early stopping and to understanding when your network is done training.

Even if you are not interested in adversarial grokking, suppose you simply want to know when the first stage of training has happened. You can look at this metric and ask when it is plateauing. When it reaches this first plateau, you know you can stop training, although you never computed training or test accuracy.

A proxy metric that characterizes the geometry of your partition may be all you need to understand the stage your network is in and determine whether you need to stop training or change some hyperparameters or architectural choices.

Some people have used this idea before. There is a paper by Atlas Wang from Austin in which they show that you can perform neural architecture search by looking at statistics about the partition. Those statistics are informative enough to do neural architecture search and obtain a good model to train.

Why do neural networks learn low-complexity, low-frequency features first?

There is a lot of work on implicit bias, or simplicity bias, where the network learns the simplest rule first, such as spurious correlations. That is something people are trying to address because you might think you have learned the right solution, but you have actually learned a shortcut solution that could put you in a bad situation once you deploy the model.

There are some cases in which we can explain why this happens. For example, if you learn by reconstruction, there is a bias that comes from the dataset. The simplicity bias is transferred through the eigenspectrum: which frequencies of your image have the most gradient information? It turns out that they are the low-frequency components, and you can prove why this is the case when you learn representations by reconstruction.

In the most general setting, though, I think this remains an open question. It involves the implicit bias of your architecture, the way you train, and perhaps other factors such as batch normalization and data augmentation.

You implied that high-complexity features are less likely to be shortcut features. What is the intuition for that?

At least for perception tasks, you can show that when you look at high-frequency features visually, you have fewer spurious correlations. For example, there are fewer correlations between the background and the actual object in the image, because that information is not present in the high-frequency part of the image.

By removing those features, you remove the opportunity for the network to learn those shortcut solutions. Instead, it has to focus on the actual shape of the object it is trying to classify. Rather than asking whether the background is grass or a beach and using that to identify the object, it has to focus on the object itself.

The term “emergence” comes up quite a lot. It is a woolly word, and people talk about it in relation to grokking. Sometimes it is overestimated because people use logarithmic plots, and the grokking is not quite as transient as they think it is.

Does it make sense to use the word “emergence”?

It does to me because it is a phenomenon that happens on its own. After you reach the first stage of training, you do not actively change the learning rate or regularizer to make it happen. You simply keep training, and suddenly this new property emerges on its own.

To me, the term “emergence” fits as long as you do not actively do something to make it happen.

I am looking at this diagram again, which shows the honeycomb, or topological map. It is a good example because it shows a clean partitioning between the examples. Are there more complex examples, and might there eventually be a kind of metapartitioning scheme that partitions the space even further?

From this understanding, we know what type of partition we are trying to reach. The next question is how to enforce that in a network. Is there a way to impose it through the parameterization of the model, through regularization, or through pruning?

Little by little, once we understand the geometric property we need and how it translates into the parameters of the weights, we can derive methods to reach that stage actively and earlier during training.

Again, the spline partition makes this much more intuitive than treating the network, F, as a black box. You look at the images and say, “We want to increase or reduce the radius of the regions.” You can identify which units are responsible for each partition boundary and determine which ones you need to prune to make the region bigger or smaller.

All of these things are tied together. Because we understand them, it becomes much easier to act on the network and reach that solution faster.

For people at home, I cannot emphasize enough how big a result this is. It is absolutely brilliant. I have had many adversarial-robustness researchers on the show, including Nicholas Carlini and Andrew Ilyas. For years, people have said this is an intractable problem: anything you do to fix robustness does not really fix it, and you simply reduce the headline accuracy.

You have shown that you can do it in an optimal way. What do you think this means for the whole space?

There are several questions in there. First, in our setting, we do not use strong regularization. Many previous results may have been obtained under strong regularization, and perhaps that makes universal robustness impossible, or at least much harder to achieve.

We show a new way in which people could approach the problem. The other nice thing is that we do not perform adversarial training. Because of this, we do not overfit the robustness to a specific type of attack.

If you train with adversarial examples, you may overfit your robustness to that specific adversarial attack. Because of this overfitting, you may become more sensitive to another type of attack. We do not perform adversarial training, so this robustness emerges naturally.

There is no reason to think that it is overfitted to a specific attack, because that attack was not used during training. Since it is an implicit emergence that happens on its own, it should be much more universal than the active methods people have used to obtain adversarial robustness.

I think this opens a new door for revisiting these results and seeing whether there is a new method, or a new compromise, that can be obtained through this implicit emergence.

Congratulations on this amazing work. My love goes out to Imaad. You have both been driving this work for so long, and I am so happy that you are finally showing the world how important spline theory is.

The nice thing is that these results are not specific to vision or to a particular architecture. That is one of the beauties of splines: your network is a spline regardless of the data modality or the input dimensionality.

Whatever insight you get about a spline partition transfers across many applications. When you derive something new or write a new paper, you do not solve just one problem but a whole family of problems. That is the power of this theoretical understanding: you can do one proof, and it will be useful for many applications.

It might be worth saying that people could think the spline theory is only about MLPs. I want to emphasize that every neural network is technically an MLP, right? Whether it is a Transformer, where every self-attention layer has an MLP, a graph-convolutional network, or something else, it is all an MLP.

Exactly. You can think of a convolution as an MLP with a circulant matrix, so it is simply a constrained parameterization. You still have the same structure: the whole network is an interleaving of affine mappings and nonlinearities.

This is true for all current architectures. In another paper, the Polarity paper, where we control the quality and diversity of generated samples, we show that you can use the spline formulation and spline results to improve state-of-the-art results on large architectures.

This is not only for toy examples or small-dimensional settings. It can give you answers and actionable solutions for state-of-the-art models across modalities.

Very cool. We are going to move on to your next paper. Your second paper is “Learning by Reconstruction Produces Uninformative Features for Perception.” This is about the difference between reconstruction, such as an autoencoder where you reconstruct an image and calculate the mean-squared error between the original and reconstructed images, and contrastive or noncontrastive models, where you look at the difference in the latent space.

In this paper, we tried to give some answers and explanations for phenomena that had been observed empirically. The first is that if you learn a representation by reconstruction, the representation is a good baseline but not state of the art. You need some fine-tuning to bring its quality up for the specific downstream task you are trying to solve.

The second observation is that the quality of your representation for solving a task does not align well with how good the reconstructed samples are. Even if the reconstructed sample looks good by eye, you may still need to keep training for a very long time before the representation becomes useful for a downstream perception task.

Those 2 observations have been known for a while. The question is whether we can explain why they happen and perhaps derive better methods from that explanation.

The main takeaway is that reconstruction methods in input space—in pixel space, for images—use mean-squared error, and most of the gradient information comes from the low-frequency components of the images. Those features are not the ones that are useful for perception tasks.

There is a nice example in which we look at what information is encoded in the low-frequency features and in the high-frequency features. You can easily see by eye that the low-frequency features are not enough to tell you which class an image belongs to, but the high-frequency features are enough. The bias that comes from the dataset and its spectrum is copied into the autoencoder, but because that bias does not align with the downstream task, you get a suboptimal representation.

I am going to show the figure on the first page now. It shows the eigenspectrum, and on the right-hand side it shows the features that the neural network learns first. They have more mass, so they dominate, and then later the higher-frequency features are learned.

There are also a couple of examples. The image on the left is the high-frequency component, which is very recognizable. The one on the right is the low-frequency component, and it is just blurry.

In the graph showing the distribution of eigenvalues, the components corresponding to the low frequencies have the highest eigenvalues. That is where you have most of the energy in your image.

As we show in the paper, this is what dominates the gradient information. That is what you naturally learn first because it gives you the biggest reduction in mean-squared error. In a sense, you could ask which frequency gives you the biggest reduction in MSE, and the answer is the low-frequency component.

Because of this, and because we use gradient descent, we learn it first. We are copying a bias that is already present in the dataset. We learn those components first, and then, if you train long enough and have enough capacity in your autoencoder, you start learning the high-frequency details.

Those details have much less amplitude and therefore provide much less gradient information. Only then do you learn features that become useful for perception tasks, because those are the features that contain the information needed to tell you which class an image belongs to.

You said there is a new method now. People are starting to use reconstruction—something like an autoencoder—because it is easy and accessible. But it picks up this dataset bias, because many datasets are dominated by low-frequency attributes.

That is dataset-specific, though, and you can actually add noise to fix it.

It is dataset-specific because you are trying to mimic the bias in the eigenspectrum of your dataset. That eigenspectrum will differ depending on whether you have a background, different translations in your images, different numbers of classes or objects, and so on.

The bias, or the misalignment between reconstruction and perception features, is dataset-specific. The simpler the dataset—MNIST or SVHN, for example—the more aligned the two tasks are. In a sense, every bit of information in the image becomes useful for both reconstruction and perception because there is no background, nuisance variable, or noise. If you learn to reconstruct, you learn to recognize.

That is not the case for realistic images with higher resolution, color, backgrounds, and many varieties of objects, such as ImageNet. There, the alignment becomes very poor.

People now use masked autoencoders, which are a different version of denoising autoencoders with a different noise strategy. In a denoising autoencoder, you take the original image, add some noise or perturbation, and then try to reconstruct the original image. You are trying to denoise the distribution you used.

Denoising autoencoders traditionally used additive isotropic Gaussian noise, but masked autoencoders mask large blocks of the image. You therefore have a different type of noise strategy.

What we show in the paper is that you can manipulate the noise strategy to counteract the bias in the dataset. You are telling the network, “I know you are trying to mimic this bias, but let me make your life harder for this part of the spectrum or for these types of features.” Instead, the network focuses on the other side, which is probably better for your downstream task.

Through careful tuning of the noise distribution, you can realign learning by reconstruction with learning perception features. But it remains a very active area of research. How can you design a noise distribution automatically if you do not have access to labels or a prior?

You were saying earlier that, for example, pink noise has a logarithmic power spectrum. You can design the noise spectrum to preferentially focus on the low-frequency features, or on some other part of the spectrum.

If you tell me what type of frequencies your downstream task should focus on and which ones are useless, then, from that information alone, you can reverse the process and determine the noise strategy. When you learn a representation by reconstruction, you can avoid encoding the useless features and focus only on the ones useful for your downstream task.

But this requires some expert knowledge, and that is not always easy. If you have another downstream task, such as depth estimation or counting the number of trees in an image, it is not clear that the right distinction is simply high versus low frequency.

Depending on your downstream task, it may be difficult to define the noising strategy that makes the most sense. You also need to implement that noising strategy. If it is too complicated or slows training down too much, it may not be practical.

The punchline is that if you use a reconstruction loss, you inherit a set of dataset biases that can cause problems downstream. Therefore, we should use contrastive methods. Some members of the audience may need a refresher on what that means. How does it work?

In most reconstruction-free self-supervised-learning methods, you observe different views of the original image. This can be because you apply different data augmentations, extract adjacent frames from a video, or have different viewpoints of the same building.

You pass those inputs through your network and compare their representations in the embedding space. You want all those different views to have the same representation. This is a form of comparison, but it happens in the embedding space rather than by reconstructing the original image and comparing the reconstruction with the original input.

Is the intuition that, after a network has reached maturity, it focuses on high-complexity information? When you compare things in the latent space, you get that for free. You are focusing on the type of representation you want.

Because you work in the embedding space, you can disregard many details about the input image that you do not need. Of course, this depends on how you define the data augmentation and how you sample positive pairs, but it is much easier to disregard things because you are not comparing against the original pixel-space image using mean-squared error.

You can therefore control much more easily which features your network focuses on. With reconstruction, the only way to say, “I do not want to focus on the leaves of the tree,” would be to devise a new loss that is invariant to them. That is a huge research program, and there may be no easy or tractable solution.

Working in the embedding space is a useful and efficient proxy. You can continue using MSE, but in this new space, where it is easier to disregard information about the input.

The third paper is “Characterizing Large Language Model Geometry Helps Solve Toxicity Detection and Generation.” Give us the elevator pitch.

There are 2 key components in this paper. One goes back to splines. If you look at a single layer of an LLM, you can decompose it into 2 large blocks: the multi-head attention block and the following MLP block.

This is true for each layer of most current LLM architectures. If you look at the MLP block alone, you can interpret it as a spline. Whether it uses a ReLU activation or a SwiGLU activation, it is the same under this spline viewpoint.

You can try to understand the region in which a given prompt falls. Again, this is a purely geometric characterization: is the region large or small, and what geometric characteristics does it have?

We derive 7 very simple features that characterize this. There are 7 features per MLP block, so the number grows linearly with the number of layers, but the total number of features is small. Even for a 70B model, you have about 500 features that fully characterize a given input prompt.

We then ask whether these features are informative about the prompt. Even if you perform a simple t-SNE visualization—a very basic unsupervised dimensionality reduction into 2D—you can see that the features cluster based on the modality of the prompt. You can tell whether the prompt comes from mathematics, law, medical data, or another dataset.

You can also see that prompts cluster according to whether they contain toxicity. This is interesting because it shows that, by characterizing the geometry of the partition and the region in which a prompt falls, you can obtain a strong characterization of what the prompt is about.

This can be applied to any pretrained LLM. It does not require expert knowledge, and the features are easy to extract. You can use them for different downstream tasks. You can feed them into a linear layer for toxicity detection, or if you want to do something else, you can train a model to predict whatever quantity you want from those features.

This is another great example of people doing unsupervised representation learning by looking at the space, rather than looking at the geometry—the partition boundaries.

You have created a set of features that statistically describe this geometry. One feature might be the average distance from a point to the boundary, for example. These features alone work well: you perform t-SNE on them and see clear structure, but you can also build a linear probe or a random forest, and they are significantly more informative than other methods.

We compared this with the most-downloaded models on Hugging Face for toxicity detection. We used those models as a baseline and compared them with this approach, where we extract the geometric features and train a linear head for toxicity detection.

Not only can we make predictions with latency as low as those models, but we also obtain a better detection rate. It is a competitive solution. It is useful because you can reuse a pretrained model and control how you use the features for your downstream task.

Here, the downstream task was toxicity detection, but you can fit the method to anything you want. You can also control the tradeoff between latency and accuracy because you extract features layer by layer. In this case, you could use only the features from the first 3 layers to get low latency and good accuracy, or use more layers to improve accuracy at the cost of latency.

This differs from current solutions, where you might retrain a new LLM and treat toxicity detection as a new task for the language model. Here, you do not have to use the model as a black-box detector.

Another benefit is that it is easy to derive the features. Suppose you know nothing about your model. Each layer has a representation that might be 10,000 dimensions, and you have one such representation for each layer and each token. You do not know how to make sense of that.

You cannot simply concatenate all the representations, because you might end up with a million-dimensional representation of your prompt. If you want to learn a linear probe on that, you need to experiment with regularization and feature selection, which becomes cumbersome.

With the spline intuition, we know how to derive informative features that characterize your prompt. We need only a few hundred of them, so you can learn a linear probe even if you have only a few samples for the downstream task.

That is incredible. I am looking at your results now. The most famous model for toxicity detection, Martin Hart's model, has been downloaded 1.2 million times in the last month. On the OmniToxic dataset, that model has an area under the curve of around 73.5%. Yours, using only a linear probe on Llama 2 7B, achieves 99.18%, and your latency is comparable to the best model.

The nice thing about this method is that, because you extract the features per layer, you can control the latency-accuracy tradeoff. You can use only the first layer or the first 2 layers to reduce latency further, perhaps losing some accuracy, or use more layers to improve the accuracy.

You can also do much more than toxicity detection. As they did in the study, you could use the features for data filtering, to determine which prompts should or should not be used for training, or to compare models.

You could try to create a new LLM that is orthogonal to the current one in terms of these features. There are many things you can do because the features are differentiable, which means you can use them during training. They open the door to using them as additional regularizers or training objectives.

That is important because we were talking earlier about designing a new regularizer. These features are differentiable, and you can use them in many different ways: interpretability, robustness, sparsity, and other parts of the training dynamics.

You can use them as regularizers, but you could also use them to derive adversarial attacks. Going back to the first topic, because you can differentiate through the features, you could manipulate a prompt to make it look more toxic or less toxic.

You can do many things because you have differentiability. The features can be computed very quickly on the fly, which means you can use them from the beginning at every training step as an additional regularizer or training objective.

The second part of the paper looks at the intrinsic subspace of the prompt. Can you tell us about that?

The first part was about understanding the MLP block and seeing what we could do from that geometric understanding. The other part of an LLM layer is the multi-head attention block. Here, we try to understand what geometrically characterizes a given input prompt.

We derive a scalar proxy: the intrinsic dimension of the space in which the prompt lives. In short, you can derive this as a function of the sparsity of the attention that you get.

It is easy to see that current training prompts have a particular intrinsic-dimension distribution. We can therefore try to create new prompts with increased or decreased intrinsic dimension.

You might initially ask what this is useful for. One application we found in the paper is that if you artificially increase the intrinsic dimension so that the prompt resembles a point far from the training data, you may move into a space where your LLM was not well aligned. You can bypass the RLHF mechanism and make the LLM generate toxic content.

This is natural because the way you prevent toxic generation is through training. You tell the model, “Do not say this here, and do not say that there.” But whether the model extrapolates that rule is an open question.

We show that once you explore a new part of the space that was not used during training, through this manipulation of intrinsic dimension, you can make an RLHF-trained LLM generate toxic answers. With normal prompts, it will say, “I cannot say this because I am not allowed to.”

There is a real theme here: a spectrum of complexity in neural-network training dynamics and representations. You have a figure showing a relationship between context length and intrinsic dimension.

A better way of explaining it might be in terms of the complexity of the representations. RLHF currently has a complexity limit, which means it is capable of addressing only low-complexity representations or things with a small intrinsic subspace.

When you use a very long context length, you are reminded of the interview I did with the University of Toronto students recently. They have a self-attention controllability theorem, and they say that as context length increases, controllability increases. That means you can make a language model say anything.

That is related to increasing the intrinsic complexity of the representation.

The higher-dimensional the space you work in, the harder it is to control what your LLM will do. The number of samples you need to control it grows exponentially, unless you have a smart parameterization that allows you to learn to extrapolate from a few samples.

That is not the current way fine-tuning works. As you increase the context length, and as you increase the relatedness between tokens so that the self-attention mask is less sparse, you move into a part of the space that was never seen before. It is a very high-dimensional part of the space, so it is less likely that the people who created the LLM were able to control what happens there.

We are simply in a very high-dimensional space where no one can guarantee anything at this point.

This is a real problem for steerability, alignment, and interpretability. I was speaking with Nora Belrose the other day about concept scrubbing. You can scrub concepts out of neural networks, and it works well early in training, but as the network complexifies, it adapts and counteracts the concept scrubbing.

We have a real problem: as neural networks complexify, we can no longer control what is happening.

It is also difficult to find a solution by acting only on the data. Based on our vision and the way we learn, we might think that removing information from the data is all we need to do. But in high-dimensional spaces, many things happen between different input dimensions.

The concept may still be embedded in another part of the data, and the network will find it if it is a useful shortcut solution. That is why we need to be careful about relying on visual inspection and 2D reasoning. Those intuitions may not scale to high-dimensional spaces because many other things are happening, and the network will pick up on them.

There is a huge need for provable solutions. We need to understand what the network is trying to learn and control it geometrically to prevent it from learning shortcut solutions. We also need guarantees that we have reached a safe model, or whatever property we want.

It cannot be achieved only by acting on the dataset or through empirical methods. As we show here, you can always find a new part of the space where nothing has been seen before, simply because these spaces are gigantic.

You need a better parameterization or better control of your network if you want provable guarantees for RLHF extrapolation.

To bring it home for people, the example of this jailbreak is that you add a huge context length, and suddenly you have made the prompt impervious to RLHF. RLHF has a complexity ceiling, and when you complexify the prompt, you move outside the controllable space of RLHF.

That seems like a major problem. What can we do to RLHF to fix this?

This jailbreaking is not specific to one LLM architecture or one LLM setting. In the paper, we showed examples with Llama 2, but you also experimented with ChatGPT and saw the same thing.

It does not seem specific to the architecture or to the way RLHF is implemented. There is a fundamental problem: controlling the behavior of a deep network in a very high-dimensional space, where you cannot visit all the places in that space.

One way to approach it is to find a better parameterization of the network, so that learning from a few examples lets you generalize or extrapolate to many other parts of the space. This is about finding the right parameterization of the model or the right way to perform RLHF.

It is a huge problem involving extrapolation in general: what to do with high-dimensional data and how to control your model's behavior everywhere in the space from only a few training samples.

One way to attack RLHF is to increase the context length. One way to increase the success rate is not just to add extra context made up of random tokens, but to add related concepts. Because the added tokens contain related concepts, the attention is less sparse, and you have a much higher chance of jailbreaking the RLHF mechanism.

These attacks are also not easy to detect. If you simply repeat a phrase many times, it may be easy to safeguard the model against it. But if you add natural English sentences containing concepts related to the toxic prompt, it is difficult to detect because the sentences are normal English.

Even so, this can jailbreak RLHF at a higher rate than random tokens. There is a lot of work to do in understanding the relationship between adding related context and adding random context, and how all those things interact.

You have been working closely with Yann LeCun for many years at Meta, and you are now at Brown. What is your research plan for the next year?

My plan is to increase the number of provable guarantees we have in current learning solutions, whether we are working with text, computer vision, or multimodal datasets.

We need to dive deeper into what is happening behind the curtains: the training dynamics, the biases, and the regularization that we learn from the data. We need to rethink basic things that we have been doing forever, such as regularization and training, so that we can control these effects and give provable answers to users and practitioners.

There are many things to do because we do not question anything nowadays. It turns out that we need to question most of the methods we are using if we want to make progress.

This can involve splines, but there are many other tools we can use. The goal is that whenever a method is not working, you should be able to give a precise answer rather than simply saying, “Try another hyperparameter and come back to me in 2 days.”

Everything we can do to develop useful theory and provable guarantees that are tractable for industry-scale problems is what we are going to work on over the next few years.

We were briefly talking earlier about the deep-learning theory book with Shai Shalev-Shwartz and Dan Roberts. Do you have any broad views on other theories of deep learning?

There are many theories, and many of them provide practical insights. For example, there is a paper by Greg Yang showing that, from a theoretical characterization of training dynamics and what happens in a current network, you can perform cross-validation on a small network. Given the hyperparameters you find, you have a rule for extrapolating them so that they are also the best hyperparameters when you use a larger model.

There are many practical results like that which come from theoretical studies of deep networks. What we are missing is something easily accessible to everyone. A lot of current theoretical studies of deep networks require a significant mathematical background, so they are not easy to access for people who did not do a PhD, or at least a bachelor's and master's degree, in mathematics.

That is why I like the spline formulation. You can make progress even through visual inspection of the model. I always try not only to do the theory, but to make it accessible so that anyone who reads the paper can learn something from it and train better models the next day.

There are many different theoretical viewpoints on deep learning, and each has tried to develop new solutions. Now I think we need to assemble them into one framework that is independent of the modality, independent of the architecture, and accessible to everyone.

Are you going to publish a paper with Ellie Pavlick in the next year?

That is the hope.

I am a huge fan of Ellie. I had her on the show, and I think she is working a lot on negation in LLMs at the moment.

It is very interesting to talk with her because of the way she thinks about language, how we learn, and how LLMs learn. There are many things we can try to learn from that, and it is very complementary to using splines to explain LLMs.

I think there will be some fruitful collaborations with her over the next year.

That is what is good about machine learning, and it is what fascinated me about the papers we discussed today. You have a huge background in self-supervised learning, vision models, and so on.

The theme we have been discussing is this spectrum of complexity in representations and training dynamics. It is interesting how you transferred that into RLHF. It seems like a slightly different domain, but science is about reusing knowledge from different domains and cross-pollinating ideas.

The intuition you get in one modality can easily transfer to other modalities as long as you do not overfit too much to a specific architecture or viewpoint. That is the beauty of splines: they occur regardless of the modality, so whatever insight you get about a spline partition transfers across datasets, whether they involve images, text, or something else.

That is what happened in this paper. You came from the LLM expertise side, and I came from the spline perspective. Together, it was easy to develop the solution because everything becomes transferable once you have the right intuition.

People should not be afraid to explore new dimensions, new data modalities, or new architectures. That is how you get the best insights, which become complementary in the end.

If people want to learn more about spline theory, where would you point them?

The best approach is probably to look at the last few papers we have written with Rich Baraniuk about splines. There are several different ones: some focus on generative models and uncertainty quantification, and this one focuses on LLMs.

People should find the paper closest to their current expertise so that they do not have too many things to learn at once. Start by learning the spline partition and the basic ideas, and then reach out to us. Feel free to send us emails or messages on Twitter.

Splines can be cryptic if you look at papers from the 1980s and 1990s, because the way people thought about them was very different. Do not begin with a 30-year-old paper on spline approximation and approximation theory, because it may confuse you more than anything. Start with figures and current papers, and reach out if you have questions.

There is a bonus question. Recently, Kolmogorov-Arnold Networks came into the limelight. They are also a kind of spline approach, aren’t they?

They present an alternative to current MLPs in which they hard-code some spline activation functions into the architecture. That is a great example of using prior knowledge to define an architecture that works well for specific problems.

Most of the problems they looked at were small-scale and low-dimensional. That is where you need expert knowledge to design what the partition should look like and what type of spline to use in order to get the best results from a small training set.

Splines offer a nice way to do this because the geometry is visually interpretable. You can see what you need, and if you hard-code most of those properties, you implicitly need less training time and fewer training samples to learn something meaningful.

That is another area where splines are useful. If you are an expert in your domain, your data, and your downstream task, you can transfer that expertise in terms of geometric properties. Splines give you a nice way to do this in practice, create new architectures, and obtain new models.

Randall Balestriero, it has been an absolute honor and pleasure to have you back on. Thank you so much.

Likewise. Thank you very much for the invitation.