Neural Tangent Kernel Perspective on Parameter-Space Symmetries

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Neural Tangent Kernel Perspective on Parameter-Space Symmetries".

Jane: The paper was written by Ori Shem-Ur, Khen Cohen, Aviv Orly and Yaron Oz from Tel Aviv University.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Jane: We also have Lu with us today — senior AI researcher at Tsinghua.

Tom: We also have Meng with us today — lead engineer at a mysterious AI startup.

Jane: We also have Lalam with us today — the in-house Large Language Model.

Tom: Alright, let's get started.

Title and Authors: Tom: Welcome back to the show, everyone. Today we’re digging into a paper that’s been making the rounds, titled “Neural Tangent Kernel Perspective on Parameter-Space Symmetries.” Jane, I’ve got to say, that title alone had me leaning forward.

Jane: It did for me too, Tom. And the author list is Tel Aviv University all the way — Ori Shem-Ur, Khen Cohen, Aviv Orly, and Yaron Oz. That’s a physics department, not a computer science one, which already tells you something about the angle they’re taking.

Tom: Right, and that’s exactly what got me excited. They’re treating neural networks like physical systems, with degrees of freedom, and asking why they simplify when you make them wide. That’s a physicist’s move.

Jane: Exactly. And the title hints at something deeper than just “wide networks behave linearly.” It’s about symmetries in the parameter space. The authors are saying that the reason these networks linearize isn’t a mathematical accident — it’s tied to how the derivatives of the network don’t talk to each other at initialization.

Tom: So when you have a huge network, the first derivative and the second derivative, the Hessian, they’re essentially uncorrelated. That’s the “weak correlations” they keep hammering on. And that’s the symmetry — the system is blind to its own higher-order structure.

Jane: And that’s the part that gets me. They’re not just saying it happens. They’re saying it’s equivalent. If you have linearization, you have weak correlations, and if you have weak correlations, you get linearization. That’s a two-way street.

Tom: That’s a strong claim, and it’s the kind of thing that could reframe how we think about the neural tangent kernel entirely. We’re not just talking about a limit anymore; we’re talking about a cause.

Jane: And the cause is this beautiful, almost physical idea — that the system has no built-in bias toward any direction in parameter space. It’s symmetric. And that symmetry is what makes it behave like a linear model.

Tom: So when we look at the authors, coming from high-energy physics and string theory backgrounds, it makes sense they’d see it this way. They’re used to thinking about symmetries as the fundamental organizing principle.

Jane: Absolutely. And that’s the lens we’re going to keep as we go deeper. Next up, we’re going to get into the actual summary of what they proved, and I promise it’s going to be worth the wait.

Paper Summary: Tom: Alright, we’re back, and we’re still on “Neural Tangent Kernel Perspective on Parameter-Space Symmetries.” Jane, you’ve had a minute to sit with the actual results. What’s the core claim?

Jane: The core claim is that linearization in wide neural networks isn’t a happy accident. It’s a direct consequence of what they call “weak derivative correlations.” Basically, at initialization, the first derivative of the network — the gradient — and all the higher derivatives, like the Hessian, are essentially independent. They don’t correlate.

Tom: And that’s not just a nice observation. They proved it’s equivalent. If you have one, you have the other. That’s Theorem three point one and three point two in the paper. It’s a biconditional.

Jane: Right. And the way they prove it is pretty elegant. They use a Taylor expansion of the network’s update rule. When you take a gradient descent step, the change in the function is a sum of terms, each involving a derivative correlation. If those correlations are weak, all the higher-order terms vanish, and you’re left with just the linear term.

Tom: So it’s like if you’re pushing a cart and the wheels are perfectly aligned — you only need to think about the forward force. But if the wheels are misaligned, you get all these sideways forces, and the motion gets complicated. Weak correlations mean the wheels are aligned.

Jane: That’s the intuition. And they go further. They show that the rate of linearization is governed by a single sequence, m(n), which for typical networks is the square root of the width. So the wider you go, the weaker the correlations get, and the more linear the network becomes.

Tom: And they also tackle stochastic gradient descent, which is what everyone actually uses. They show that even with random batches, the deviation from linearity stays bounded over time. That’s Corollary four point one, and it’s a big deal because most prior work only handled deterministic gradient descent.

Jane: That’s the practical win. You don’t need to assume a fixed dataset. You can have randomness in the training process, and the linearization still holds. That’s much closer to how real models are trained.

Tom: And they back it up with experiments. They ran networks on MNIST, CIFAR-ten and Fashion-MNIST, with different activations and depths, and the correlation decay matches their theory.

Jane: So the summary is: weak correlations cause linearization, linearization is equivalent to weak correlations, and this holds even in stochastic settings. That’s a complete package.

Tom: And it’s a package that’s going to have real consequences. Next, we’re going to talk about what this means for actually building and training networks — the improvements and the practical advice.

Improvements Suggested: Tom: We’re still on “Neural Tangent Kernel Perspective on Parameter-Space Symmetries,” and now I want to get into the “so what.” Jane, what does this paper actually suggest we do differently?

Jane: The big one is that it gives us a diagnostic tool. You can measure the derivative correlations of your network at initialization, and that tells you immediately how linear your training is going to be. You don’t have to run the whole training to find out.

Tom: That’s huge. You can check, before you even start, whether your network is going to behave like a kernel method or like a feature learner. And the paper even shows how to compute those correlations efficiently — you don’t need to build the full Hessian, just a few directional derivatives.

Jane: Right, they use a clever chain rule trick to avoid the O(N2) cost. You can get the second-order correlation with just O(N) work. That makes it a practical tool, not just a theoretical one.

Tom: And then there’s the learning rate. The paper shows that if you rescale the learning rate, you can push the system toward or away from linearization. That’s Theorem three point two. So you have a knob to turn.

Jane: And that connects directly to the “lazy training” regime that Chizat and others talked about. But here, it’s not just about a scale factor — it’s about how that scale factor interacts with the correlations. You can see exactly why a smaller learning rate makes the network more linear.

Tom: There’s also a point about activation functions. They show that the growth of the derivatives of your activation function controls the linearization rate. If your activation has exploding higher derivatives, you’ll need a wider network to get the same linear behavior.

Jane: That’s a design guideline. Pick an activation whose derivatives are bounded, and you’ll get linearization sooner. That’s the kind of concrete advice engineers can use.

Tom: And for the architecture folks, they generalize this to any network that fits the tensor programs framework, which covers a huge range of architectures — CNNs, RNNs, attention models. So the advice isn’t just for fully connected networks.

Jane: The one caveat, and I think this is important, is that they’re not saying linearization is always good. In fact, they spend a whole section on why the NTK limit might underperform — because it lacks bias. So the improvement here is knowing when you want linearization and when you don’t.

Tom: So it’s a tool for understanding, not a mandate. That’s a healthy way to think about it. Next up, we’re going to wrap this up and talk about the big picture.

Conclusion: Tom: And we’re back for the final stretch on “Neural Tangent Kernel Perspective on Parameter-Space Symmetries.” Jane, give us the send-off.

Jane: The paper gives us a unified explanation for why wide neural networks linearize: weak derivative correlations at initialization. It proves this is both necessary and sufficient, and it extends the result to stochastic gradient descent, which is what we all use in practice.

Tom: And it gives us practical tools — a way to measure correlations cheaply, a way to control linearization through the learning rate, and guidance on activation functions.

Jane: It also raises a really interesting question about bias. The authors suggest that the NTK limit might underperform precisely because it’s too unbiased. Real networks keep a little bit of correlation, and that acts like a prior that helps with real data.

Tom: That’s a provocative idea. It flips the usual narrative that kernel methods are biased and neural networks are unbiased. Maybe it’s the other way around.

Jane: Exactly. And that’s the kind of thinking that could lead to new architectures that deliberately introduce beneficial correlations, rather than just trying to kill them all.

Tom: So we’re saying goodbye to this paper, but the ideas are going to stick with us. It’s a physics-flavored view of deep learning that gives us both clarity and new questions.

Jane: And that’s the best kind of paper. We’ll be thinking about this one for a while. Thanks for listening, everyone — we’ll see you on the next one.