Alex Damian (Harvard), Jerry Liu and Yasa Baig (Stanford) | Stanford Hidden Layer Podcast #105

Transformers and Adam thing has always been very funny to me because there's just this folklore that Transformers are really hard to train by SGD. So if Adam just never came about, I don't know where we would be. But I think I've given this hot take before, but I do think that within a few years, we'll be looking back at the way that we were training models now and just have a lot of thoughts. I think we're doing things very wrong right now. I think Adam in particular, it's incredibly useful, but I do think it has huge drawbacks. And I do think that it's going to limit the models that you can train. I think there are a lot of things about the modern transformer that do just seem really strange. A lot of architectural decisions that do just seem really strange. And I think we're more driven by systems and the need to scale than by some thinking about, oh, what is the right way to represent this thing? So I think one of the things that's always very surprising to me is when you're scaling up a transformer, the, like you scale up the embedding dimension with the number of heads in a way that keeps like the head dimension completely constant. And actually, like if you look at like the FlashAttention implementation, like it's stuck at 256 max. Like if you want to do like a larger, like head dimension or a different head dimension, then you're kind of stuck. I mean, you can like hand roll something, but this does feel like it makes sense why this like came about because like in the forward pass, you want to keep most of your matrices for the MatMols like square, right? So like you do like embedding dimension to embedding dimension and you kind of chop it up. But at the same time, like from an actual like representation learning perspective, this does seem a little strange. It seems a little artificial. It's like for deep learning theory, like what is, what is your goal here? It's like, if you want to start from basic principles and understand a very complicated machine like ChatGPT, it's like trying to understand a dishwasher from first principles. Like you start with Schrödinger equation and then you build up to like, why does the dishwasher work? Can you prove the dishwasher works? And it's like, There are a lot of moving parts. Maybe we understand some individual parts. Maybe we have some rough sense of how they fit together. But I don't think that, like, the whole— the expectation should not just be to, like, zero-shot from first principles, like an LLM. I think this is way too much to expect. My motivation is almost entirely trying to understand the system that actually works. Like, every day I wake up and just, like, the fact that deep learning works blows my mind. And so my goal is just to understand that. I'm generally skeptical of the idea that like the solution to deep learning theory is just going to be deeper math. There's just going to be like some new tool that's going to unlock everything and suddenly we're like, ah, now I understand. I think a lot of the work to be done is just like clarifying thoughts and observations that people already have. I don't think that super deep math is necessary. I think you just have to look at things the right way. I think the right way evolves a lot and it just relies on asking good questions and like the answer to these questions, I don't think always has to be super deep math.

Yeah, well, hi folks. Welcome to another epoch of the Hidden Layer Podcast. I'm a new guest here, I guess. I'm Jerry Liu. I'm a 4th-year PhD here at ASU.

Yalsa, I'm a 2nd-year PhD student here.

Yeah, we're really excited to have Alice Samian. So currently a Kepner Fellow at Harvard, postdoc, and soon-to-be MIT professor. Do you want to do a proper introduction?

Yeah, sure. So I'm Alex. I just finished my PhD at Princeton under Jason Lee. I'm currently working at Harvard as a Kempner Fellow. And yeah, I'll be starting, like Jerry Liu said, next year at MIT as a professor in math and EECS. And I broadly work in deep learning theory, optimization, representation learning, and related questions.

Yeah, awesome. So, okay, very different type of guest from what we've seen before, but I think you're doing a ton of really exciting work and I'm excited to See if the folks who watch the podcast will appreciate it. We'll do our best. Yeah, okay. To start off though, I think this is actually kind of funny. So we were all at Duke in undergrad together. You guys have known each other for longer.

Yeah, I think basically 8 years at this point, something like that. Yeah, something like that. I don't know if you know this, but you're actually the first person to ever buy me a boba in my entire life.

Wow, I had no idea.

Yeah, I have a very critical memory of this. Well, I was like— My college experience. When I was 19, I was like, I got a random text from a strange messenger. Actually, this has been wrote with two. Um, and it was just like, what's your favorite boba? Like, I've never had boba before. And then we went to get a ca— we went to a cafe like several weeks later. It's like, uh, crazy. Anyway, sorry.

Very sorry you've received that thing.

You know what I mean?

So that was the first time I met Alex. And then, um, I like, then I, I heard about him from every single per— I didn't actually ever tell you this story, but, um, When I was an undergrad, I was trying to figure out where I wanted to do research, and I would talk to a lot of physics and math professors, and they would all just know you. I never brought you up, and you were, I think, a junior at the time, or maybe even a sophomore. And I was talking to probably the biggest biophysicist at Duke, who is Harry Greenside, and he was kind of like, oh yeah, you're like Alex Damien. I'm like, how do you know Alex? He's like, we all know Alex.

Yeah. So on my part, you were teaching the Putnam competition thing in, I think it was my sophomore year. I remember you just read a textbook on this stuff for the summer or something. I just remember you, you brought up some, there was some like crazy matrix identity things that you get by like setting things up as block diagonal systems or something. Anyway, all sorts of details. And then the year after, I think we were both in Cynthia Rudin's lab and then you actually handed off this like very cool product you were doing at the time. It was like GANs and solving inverse problems with StyleGAN2 and all sorts of crazy stuff. But yeah, now you're in a very different world. You're working on ML theory stuff, working on optimization mostly.

Yep.

Okay. It's very exciting to hear about this. So, okay. I want to first take some time and let you talk about your most recent work, which we're very excited about, on central flows and adjustability.

Yeah.

Okay. And these are terms that maybe are not super familiar outside the ML theory world. So do you want to give us a brief discussion of this?

Yeah.

I mean, central flows we just kind of invented, but adjustability has been around for a few years now. Actually, a lot of this work started or was at least inspired by Some of my experiences at Duke and working with Cynthia were like, I was doing more empirical work, like more computer vision work. And I think one of the things that we very consistently found that was very frustrating is I didn't feel like I could really— it felt like you would get so many performance gains just by tuning the optimizer a little better or doing weird things that didn't really make a lot of sense to me. And so going into my PhD, I just kind of wanted to understand what was going on with optimization and deep learning. And so I had a couple of papers on SGD and some dynamics, but there was this really cool paper by a collaborator of mine, Jeremy Cohen, in 2021 called, I think the title was like, Gradient Descent in Neural Network Typically Occurs at the Edge of Stability. And essentially what he discovered in this paper was that if you can easily categorize gradient descent into two regimes, stable and unstable, depending on how large the Hessian of your loss function is. And somehow the idea is that if your loss function is very flat, it's very easy to optimize and you can use a relatively large learning rate. Whereas if it's very sharp, then it's very hard to optimize and you have to use a much smaller learning rate. And so if you use too large of a learning rate, then you'll blow up and you call this unstable training. And so long as you're not blowing up, you call it stable training. And this seems kind of nice. And then like, if you like read an optimization theory book, it'll tell you how to set your step size in terms of the Hessian in order to be in the stable regime where optimization can actually happen. And what this paper discovered is that this actually isn't really indicative of the way deep learning operates in practice. And the reason is that the hessian just tends to grow during training so that whatever learning rate you pick, it'll eventually become so large that the hessian becomes unstable for whatever learning rate you're using. And then you do kind of blow up a little bit, like you are in this unstable regime, but then somehow training still works. Like the loss keeps coming down. Like you're looking at the loss curve. It seems fine. Like it's a little spiky and there's this like huge loss spikes, but like Progress seems to be happening. And it doesn't seem so catastrophic, but then like from a theory perspective, it kind of violates almost everything that we knew about or why gradient descent should converge.

So this is like a very like kind of like simple to understand case where like neural nets do something crazy that my traditional theory like can't really explain. Yeah, exactly.

Okay, that's awesome.

And then I actually had some prior work with an amazing collaborator, Ishan Nachani from Princeton, on trying to understand what was actually going on. And we've discovered this like hidden feedback loop in gradient descent that would kind of prevent it from blowing up. And so it's true that if you're in a region that's too spiky, that gradient descent can't really operate properly and you are supposed to blow up. But then after you begin to blow up, there is this like hidden feedback mechanism in gradient descent that restabilizes training. And then that's exactly what you see with these like spikes in the loss where when you begin to blow up, you'll begin to see the loss spike. And then this feedback loop kind of kicks in, restabilizes training and the spike comes back down. And so we kind of understood this, but The theory wasn't complete and there were a lot of like side cases that we couldn't, like edge cases that we couldn't understand. And it also didn't extend to more interesting optimizers like Adam. Like this was really a gradient descent argument. And so in this more recent work with Jeremy Cohen called Central Flows, we, it's less rigorous. We kind of adopted the role of physicists where we're just like doing vibes math and like looking at equations and just deciding which terms to keep and which ones not to. But we ended up getting these like really nice heuristic explanations predictions for which trajectories these optimizers take during training and what's happening during this edge-of-stability phase where they're kind of blowing up, but kind of not blowing up.

Awesome. Yeah.

Yeah.

So, okay. I want to get back to the actual approach you guys did and kind of like how you, you do your ML theory work and kind of what perspectives you take there. But yeah, maybe first, could you give a bit of a summary of like your main takeaways? Like what exactly is a central flow and yeah, kind of the high level of the analysis there. Yeah, sure.

Here, the idea was just that during these explosions and these instabilities, a lot of these optimizers are actually super oscillatory. And you kind of expect this for stochastic gradient descent because the noise can kind of introduce these oscillations. But it turns out that even gradient descent is a super oscillatory optimizer. And if you actually plot it in weight space, you'll see these jagged edges. And so the sum for flow is just some smooth process that tries to smooth out all of these edges. And give like a more holistic view of the trajectory that gradient descent is taking. And so it tells you like the overall path that's taken. You can read off some properties of the optimizer. You can try to reason about convergence rates. And essentially it's all just an attempt to make a version of optimization theory that actually applies to the regimes that people train it. Right.

Okay. So maybe, okay, digging into this edge stability and the connection here first, I think this is kind of, well, very unintuitive. Maybe if you haven't dived into it, but can you give us a sense of like what exactly is, well, first of all, what exactly is going wrong? Let's say if we observe that neural networks are always at the edge of stability, like if I'm just a practitioner, should I be worried about this? Does this mean that like our training's slow or it's not finding the best answer or something?

Yeah, I mean, this is a totally fair question. I think the default answer, and there's a few exceptions that I'll get to, is like you should not be worried about this at all. You shouldn't try and shrink your learning rate. Larger learning rates still do train faster. This is something that empirical researchers have known for forever, but it turns out that you can set these learning rates way larger than theorists imagined. And essentially what the trade-off is that you'll converge faster, but you'll get a much spikier loss curve. And whether or not this is worth it is going to depend on how stable the training is. If the loss goes down, but there's spikes in the loss, you don't really care. But if those spikes are catastrophic, this could actually be bad. And so there's some numerical stability issues, but then on top of that, it is possible that at some point, if you use too large of a learning rate, you won't find a good solution. You should shrink it near the end. But I personally have not seen that many situations where this is a huge issue. At least for the full batch case or large batch setting, use the largest learning rate you can. Again, it's something that practitioners have just known for ages, but it's really hard for theorists to reconcile this with the theory that we're familiar with.

Right, right.

And so, okay, so one piece of your theory as well is you guys have claims about like, kind of, if you give me a neural network, like, kind of what is the largest learning rate that like is, or like, you know, is this learning rate going to blow up, for example, right? If I'm like at a certain point in the loss landscape, whatever, am I going to have these catastrophic spikes or are these things actually kind of manageable? Is that, is that correct?

So that actually, I think it's fairly hard to predict. And unfortunately it depends a lot on like quantization issues and like whether or not you'll actually end up blowing up. It turns out to be really, really hard to predict. Uh, I think this is actually one of the failure points of the theory is that like, In trying to develop a prescriptive theory, I still don't have a good answer for how you should set your learning rate other than the standard trick of you just train for some amount of time, grid search some learning rates and see what works the best.

Gotcha.

It's more like trying to understand the mechanism for why you don't blow up when you use large learning rates.

I see. I see. Okay. So it's kind of extending classical theory in a new setting that hasn't been analyzed.

Yeah, exactly.

But I think it's still a very incomplete theory.

Yeah. Okay. Interesting. Okay. Yeah. I wonder if it's worth maybe getting into a bit of the details of like kind of the method, right? So you mentioned you were like thinking as a physicist as opposed to like your traditional mathematician's hat or something.

It's a very strange feeling for me.

Yeah. Can you, yeah. Can you talk a bit more about that? Like what does that look procedurally?

Yeah. Yeah. So, I mean, I think this goes back to this paper with Yisheng where there we're thinking much more like, you know, mathematicians and learning rigorous math. But we had to begin by just understanding the correct mechanism for like why you don't blow up. And there the procedure is very fun. It's actually a procedure, we worked on this together and it's a procedure I've kept ever since. So what we did is like, you can like take a training run and then go until the first loss spike and you just like save a checkpoint right before the loss spike. You like load the weights and then you step through the next like 5 steps of training as carefully as you can. And we just Taylor expanded everything. So what we actually did is like, I mean, we Taylor expand all the things, right? So at first we were like, okay, we know what happens if you Taylor expand to second order, because then your loss function is a quadratic. And we understand gradient descent on quadratics really well. We know exactly when you're going to blow up. And the theory says you should blow up. So of course we tested this and we Taylor expanded to second order and we ran the dynamics and we said, yes, if the loss function were quadratic from this point, you would blow up. Unsurprising. But the real loss doesn't blow up. So let's go on. So then we did like the dumbest thing possible, which is just Taylor expand to one further order. So now, you know, now we have a cubic and like, this is like actually fairly straightforward to do in JAX where it's just like, if I give you a loss function, I can easily construct the cubic approximation of this loss locally and just run gradient descent on the cubic approximation to the loss. And then we saw exactly what you see in practice where like the loss spikes up and then comes back down and like the curves matched up really nicely. So we're like, oh my gosh, okay, it's all cubics. So like, this is what people were missing. They were thinking too much about gradient descent on quadratics and you really like need to think about gradient descent on cubics.

Okay.

Great. So then like that still doesn't really give an answer. It says that you only need to think about cubics. And so then we're like, okay, what are the relevant directions? So you start like PCAing the iterates and then you find what the relevant directions are and you're like, oh, this is actually a low dimensional system. And then you project everything onto these low dimensional systems. You get a bunch of terms in this cubic, like, like if you have a cubic in a few variables, like it can have quite a few terms and you just check each one of them one by one. Like, am I allowed to delete this one? Can I delete that one? You delete all the terms that don't seem to matter. And eventually you're left with like a very compact representation. And so like afterwards, once we understood the mechanism in that paper, we then tried to prove everything rigorously and we got pretty far. But then in this follow-up paper with Jeremy, we just use this original intuition from Taylor expanding all the things, which is like, we know the terms that matter and we know the terms that don't. So let's just like do a heuristic derivation that only keeps the terms that matter without justifying why they're the terms that matter.

Right.

Okay.

And so you, so you kind of, at the end of the day, you have this very like clean picture and intuitive picture of what's actually going on. And kind of the formulas look very clean, at least at a high level. And in some sense though, you kind of just guess the answer and then you, you just check that it kind of works.

Yeah, that's exactly right.

Yeah.

It's a very different way of thinking, but I think that I have appreciated that once the system becomes complicated enough, it was actually very hard to track things manually. Right. Yeah. And I think one of the huge issues in deep learning optimization theory and deep learning theory in general is that if you want to prove something like, yes, this approximation is good, it's not clear what your building blocks are. Like, what are you allowed to assume about the loss function? What are you allowed to assume about smoothness properties? And yeah, there are a lot of challenges in making this type of theory rigorous, but all right, interesting.

Yeah, I mean, it's interesting to hear you say this. So I actually am a physicist by training, and it's like, what you're describing is like bread and butter. It's just sort of like, oh, we need our field theory to work out, but then fortunately it blows up in the second term. So we say, add the third term so it stabilizes to make the science work out, right? Like, this is like this most standard trick ever. I mean, to the point on just guess the solution, I mean, we like to say it's like a clatterly chosen ensemble motivated by a certain hypothesis. But I mean, my diatribe aside though, I mean, I am curious to sort of take the conversation a bit more general. So I mean, of course your work is fantastic. And as a physicist, I actually really quite enjoy these kinds of phenomenological approaches to deep learning theory. So it's like, high level is like, you know, we can't necessarily like reason about everything from first principles, but we can like take some empirical observations at a high level It's like, well, okay, this seems to be true axiomatically, take this to be true and then like see how far we can get. And I mean, as you know, central law stuff, so you can't necessarily do always a rigorous job, but you can do a very useful job. And maybe this is like, you can't really necessarily write down all the error bounds, but for a physicist, this is really, really exciting. We really enjoy these kinds of phenomenological theories. But I would like to, I'd be curious to hear your opinions about like what is just generally like the state of deep learning theory, perhaps just for the audience broadly who maybe is like a lot more familiar with sort of the world going, like the world happening in engineering, all these crazy API models or crazy models API competing, right? Like what is theory right now? Or like, what is your impression of the field? Like, what do you think is exciting? And like, what do you think are like different components to how people log to each other?

I mean, I think deep learning theory is like a, it's a massive collection of like many communities and a lot of people are interested in different parts of it. But I think that, I think one strong divide that I feel, that I'm thinking about when I'm thinking about deep learning theory is that deep learning theories that try and touch data and representations tend to just be much, much harder than deep learning theory that manages to abstract that away. I think that's been the main bottleneck to developing really interesting theories for deep learning. Optimization, for example, I think is— I think optimization theory for deep learning is relatively successful because you don't have to think about the data itself. You don't have to think about where it comes from. I don't have to know that I'm working on images. Whether I'm working on text, like you can kind of black box all of this into like some notion of a loss function and then just ask about properties of this loss function. And I think that this was like incredibly powerful and it actually makes sense because some of the lessons in deep learning really do seem to be universal. Like the fact that you can train an MLP somehow feels like the same phenomenon, like as the fact that you can train like a ResNet or like a transformer, like all of them are magical, but it's not clear that these should be treated as three separate phenomena. So it makes sense to work at this like higher level of abstraction. Whereas all deep learning theory that tries to actually say something deep about the representation learning process is kind of forced to grapple with the messiness of real data. And I think this is when things get very difficult. So here I'm lumping in questions about feature learning, questions about generalization are very, very hard. I think what it means for something to generalize in images versus text is just very different. Generalization in an ImageNet sense is very different from generalization in an LLM sense where you're asking about length generalization. Or like generalization to like skill composition. So I think these types of things are like kind of just by nature more task-specific, more data-specific. And there I think tools are much more limited. I think most of the successful approaches for trying to study properties of data and the way to interact with neural networks tend to be forced to make really strong modeling assumptions. And so you kind of say, what is the property of the data that I want to extract? Maybe it has some low-dimensional structure, maybe it has some hierarchical structure. And you come up with like a synthetic toy task where you kind of bake all of this in and then you ask, how does a neural net learn to reason about it? These approaches, I think, have been fairly successful, but it is like, I think it's a fair critique that it doesn't say something directly about like what happens when you pre-train on the Atari internet or even ImageNet, right? It's very hard to get theory that specific. And I don't even entirely know what such a theory would look like.

Right.

Yeah. I mean, I guess my impression of the field broadly is like coming from physics, I think people often make this analogy of the current state of deep learning theory is often like very comparable. And I'm in this conversation several years ago, it's sort of like what quantum mechanics was pre-1926. We had all of these crazy empirical results, some Keplerian phenomenological laws. We have all these scaling laws, for example, in the ML community we talk about all the time, but we don't have Newton. We don't have the unification for all these disparate empirical results, or Dirac in the case of quantum mechanics. How true do you think this is? I, somehow, not being a specialist at all, or really even a member in this community, I always feel like it's a lot of partial formalisms. There's folks that work on NTK limits, there's folks that work on mean field neural networks, there's folks that work on all these simplified regimes of trying to understand loss and optimization. And there's work like yours, for example, which is even, say, an IRReP structure, you would agree to make these very big assumptions on the actual model architecture itself. Do you think there's some hope for unity here? Or is there a unity here that I'm just not observing?

I don't think I would quite agree that we're at this point where there's a lot of empirical observations and we don't have a mathematical formalism because I think quantum mechanics, like the equivalent of like the Schrödinger equation is just going to be like the gradient descent update. We have it, like we know the rules of the game, we can simulate it. It's just a question of like when you have trillions of particles, trillions of like neurons interacting and being applied to real data, like what happens as a result of this system? And I think that's a very different question from like what the laws are.

Well, I think that the analogy here perhaps is like maybe a bit closer to statistical physics, right? Like, I mean, to your point, like, yeah, we have like the, you know, we have the fundamental equation of motion in the system. And this is also true in statistical mechanics. At the end of the day, your particle system obeys Schrödinger's equation. But in practice, no statistical physicist is— or I shouldn't say no statistical physicist is right, but a large agenda of stat mech as a field in physics is to say, well, we have this equation of motion. In practice, I don't need to solve Schrödinger's equation for a billion particles to simulate the motion of a basketball or understand how a gas becomes a liquid, right? But instead, emergent separate phenomenological laws that are entirely self-contained and often things that you can at least derive in simplified cases can come out of the system, right? And I guess my question is, do you think that— do you know some equivalent that might exist? Is there an equation, a fundamental equation of motion at a higher order of abstraction, perhaps away from GDE, like an emergent law essentially? And why is it central flows?

Not semiflow, but like, yeah, I mean, I do think this is a huge challenge. And like, I think the answer is like, I really hope so. And I think there should be, but as of right now, like, I think this is the main challenge that all people in deep learning theory are kind of grappling with, which is like, you have the fundamental law of like, if I take the data, I can write down what it means to run gradient descent on the data. But like, you want to know what happens if you actually like run this system. And then I think everyone is looking for some levels of abstraction. So I think a lot of these phenomena that people come up with that are, people are like observe and try and document and study are exactly looking for this type of thing. Like you look for a phenomenon and they test how robust it is and is it robust enough to be able to build things on top of? I think one of the real challenges is that often the phenomena are not that robust and they might be true for images, but not for text or something. And then like it becomes much less clear how to, how to build stuff on top of it. And so I would agree that this is like the huge challenge of deep learning theory, which is like finding the right abstractions.

It's kind of the historical inverse of what happened in a lot of particle physics though, in the sense that you guys are kind of like starting with like your Schrödinger equation, you're searching for thermodynamics laws, whereas in the history of classical mechanics, or sorry, I should say classical thermodynamics, we had all these macroscopic equations. That's how energy flows. If you didn't really fully understand what energy was, this is what entropy is. We didn't even have a microscopic equation for what entropy was. We just kind of knew it was this quantity we could measure and obey these really reliable properties. And then half a century later, we would have a theoretical insight into it. But here we're starting off on the part of it like, oh gosh, it's like We need to define stuff like order parameter that's meaningful to us or some sort of abstract thing. What is the right thing to even study? I feel like it's not clear. Maybe that gets at your— Yeah, right.

Sorry. But you're jumping in though. There are all these equivalent of these higher order things that people do study empirically and theoretically, like neural scaling laws and in-context learning about these things. Maybe one take here, and my understanding is there's a lot of theoretical work in each of these individual things, maybe one take could be that it's too, maybe too much to expect to like have an all-encompassing thing, but in some sense we are making progress, I guess, right? Like what's, I don't know, what was your take? Should we hope for more here?

Like, no, I think we definitely should hope for more. But I think like one of the other like interesting lessons in deep learning theory is that, and this applies to the different phenomena that you mentioned, like many phenomena that people either discover or rediscover in deep learning theory are not at all unique to deep learning. So like, scaling laws, for example, like you can see like different scaling laws in like linear regression settings. There's a bunch of papers on this. And I think like from the perspective of like, what does it mean to understand scaling laws? Like, okay, partially it's just scaling laws begets scaling laws in the sense that if your data has a scaling law, then you should expect your loss curve to also have a scaling law. But like, moreover, it's just like, I think you can distill different phenomena into much simpler settings and observe at least the same phenomena and trying to understand it there. That doesn't necessarily mean that you can take that observation and scale it back up to the full neural network. So in a scaling law setting, like in a linear regression setting, we can predict the scaling law exponents. If you tell me the data has the scaling law, I can predict that the loss curve will have some other scaling law that's related. But if you just directly take that and apply it to a transformer, it'll obviously give you the wrong answer. The scaling law must come from something else. And I think this is missing the whole feature learning aspect. It's missing a lot of the magic in a transformer. But like, yeah, I don't know how to really study this. Then there's an entirely different group of people just studying this feature learning process and how the feature learning process changes the representations that you learn. And like, yeah, I think one of the interesting things in deep learning theory is also just that there's so many different people and so many different groups working on the same problem from so many different perspectives. And it's actually very rare for two of these perspectives to intersect. There's very few situations where this has happened. Like, um, We're all in a blind room. It's like the blind room and the elephant. You all see part of it and you see something and it's real, but you just don't quite see how the picture comes together.

I guess the other part of Yaso's question is like, like toward the analogy as well, right? Like each group kind of has their own set of tools, let's say that they kind of prefer to play with. And if you have your set of tools, like you will kind of tend to get certain types of results or like even study different types of problems because those are more approachable, right? And I guess one question here is whether the set of tools that the ML theory community has broadly, let's say, is kind of sufficient to answer whatever larger questions we want.

Yeah, I mean, I think the way that I always think about this is, do we— I think, just to keep beating this analogy, right? In a very real sense, we were incredibly lucky in quantum mechanics that the correct language to use to describe quantum systems just so happened to be linear algebra, which was not standard pedagogical stuff at the time, I think, for physicists, but it was a very well-developed theory and we could just borrow so much intuition and insight from basically how linear algebra worked and all this excellent work mathematicians have been doing for almost centuries at that point. I would be amazed if this is true in deep learning theory. Do you think basically we need fundamentally novel mathematics at the end of the day to build our theory? In the same way Newton basically built calculus from first principles to describe motion. Or is it going to be like more like QM where like, I mean, certainly plenty of mathematical innovation has since come as a product of saying quantum mechanics and linear algebra fundamentally, right? But like a lot of that tools and like, like baseline work had been built for us, right? Even in general relativity, same thing, right? Like we differential geometry was like a very well or like relatively young, but like, you know, theoretically rich discipline that Einstein could pull from. Like, do you think that, do you think that math exists for us right now in deep learning theory?

I mean, okay. I mean, obviously very hard question to answer just because like, I don't really know what the answer is.

Sure.

To deep learning theory. And so I don't know what form it'll take, but I do think that I'm generally skeptical of the idea that like the solution to deep learning theory is just going to be deeper math. There's just going to be like some new tool that's going to unlock everything. And suddenly we're like, ah, now I understand. I think a lot of the work to be done is just like clarifying thoughts and observations that people already have. I don't think that super deep math is necessary. I think you just have to look at things the right way. I think the right way evolves a lot and just relies on asking good questions and like The answer to these questions, I don't think always has to be super deep math. I do think it's going to draw on a bunch of different parts of math. Like I think the central flows paper was entirely on perturbation theory. I think a lot of stuff on like initialization schemes and stuff like μp scaling, right? For this, this is a lot of like Gaussian process theory and this is also very nice. And so I do think it's going to draw on a lot of different parts of mathematics, but I don't think it's going to necessarily be like this new theory of deep learning is just some 10,000-page textbook that everyone has to spend 5 years studying and finally understand how everything fits together. I think that it has to be something higher level than that.

Interesting. I'd never heard that take before. I spoke to a biophysicist once and I asked them if they felt that way about theory in biology, which is a thing I worked on for several years. And his answer was, I don't think modern mathematical tools are even remotely tractable for any of these problems. And I don't know, maybe I was more optimistic than Eos Craps, maybe like naively, but it's good to hear that you seem to be optimistic about this.

I am, but I also think that part of it is like trying to clarify expectations of like what we can and cannot understand. I wasn't actually there for this, so I've just like heard this story and I'm now going to repeat it. But there was a panel at which someone asked a question like this to Andrea Montanari, and he responded like, it's like for deep learning theory, like what is, what is your goal here? It's like, if you want to start from basic principles and understand a very complicated machine like ChatGPT, it's like trying to understand a dishwasher from first principles. Like you start with Schrödinger equation and then you build up to like, why does the dishwasher work? Can you prove the dishwasher works? And it's like, there are a lot of moving parts. Maybe we understand some individual parts. Maybe we have some rough sense of how they fit together. But I don't think that like the whole, the expectation should not just be to like zero shot from first principles. Like, an LLM. I think this is way too much to expect.

And then like tomorrow someone will come up with a paper being like zero-shot solution to every problem if you're an LLM. That's how it goes.

I guess the other piece is like kind of all of the empirical work, right? I think there's been a lot of very fundamental understandings that have been coming out of like, say, MechInterp, for example, right? Like a completely different set of tools, very large empirical driven by ideas in neuroscience and whatnot. What do you think the interplay should be kind of between these like more empirical, you know, investigations and like the more kind of math-heavy, stats-heavy stuff?

Yeah, I mean, I think this is also going to change a lot person to person. My personal take is that like, my motivation is almost entirely trying to understand the system that actually works. Like every day I wake up and just like the fact that deep learning works blows my mind. And so my goal is just to understand that. And if you want to understand that, I think there is like a strong burden to like actually like follow empirical work and like run experiments on some of the stuff that you do. For MechInterp specifically, I think it's almost asking a different question in the sense that it's saying, once the system has worked, can we understand what it's doing? Sure. But it's not explaining how it's managed to extract these features from the data to be you, right? Like it's saying that you have some transformer and in the forward pass, you can extract circuits. You can look at neurons and linear combinations of neurons and kind of interpret what they're doing. But I think to me, that is the magic. Like, again, MechIntrep has done a really good job of isolating the magic and being like, look, this is what, what's so amazing. But then to me, it's not explaining where the magic comes from. Like I'm feeding in raw data, like raw tokens. And somehow this is automatically arising during training. This to me is the magic.

Brett, I mean, for the record, like this is also like, um, not necessarily a criticism, but like a point in AlphaFold, for example, in like my corner of the world in biophysics, right? Like, um, it's like a, or people talk about protein folding, protein structure problem being solved in some sense, right? It's like, yeah, you can go from like a sequence or in practice, like I said, a sequence processes and get a reasonable resolution for some particular set of proteins that we can crystallize really well. But also it's like, if you start off with just a big chain of molecules, it just folds into this crazy molecular machine. And it's great that we can see what that machine looks like, but how on earth did you get there? Right? It's like, you have this enormous loss function landscape, right? In the deep learning setting, it's sort of like, there is this solution minima, but gosh, how on earth did you even find it? I mean, I think if you did the back-of-an-envelope calculation, it would take, at least for the protein case, it would take, if just doing a naive search would take more times ever passed in the universe to stumble across this. So clearly there must be some clutter mechanism by which these things happen, right? But yeah, I mean, it's super exciting. I think, sort of pivot back a little bit though. So we had a conversation once, like 5 years ago, I guess at this point, on what you thought the future models would be. I don't know if you remember this, but it was sort of like, one day, I just want the future LLMs or future, I guess it was prior to AllenGPT existing, right? But future models to just throw out a lot of architectural Deep learning should be radically different. I still believe this. I would love to hear your opinions on generally, what do you think perhaps the future of models would be and how are we going to get there? To what extent do we pull insights from fundamental understanding-based work, theory-based work, trying to understand models and how they help to motivate us to make progress? And just one last asterisk on this is in the deep learning community, we often talk about the bitter lesson, which at this point has become kind of like a I mean, this is just my hot take, but it's like people sort of use it to mean like anything. It's like you can kind of find justification for any position from it. You say, well, you know, it's the bitter lesson. This is anecdotally how people use this in the community. Maybe, you know, Richard Sutton is watching this.

I'm sorry.

Right. But like, it feels like sometimes, like weirdly enough, like the AI community specifically is like just so confident the models will eventually be smarter than us. Like we've actually kind of like strangely become like rather quite like doomer on the notion that we could like understand and just build better ones. It's like, just like scaling is all you need, bro. Right. But anyway, all to say the original point was like, how do you think, what do you think the future models could be? And like, where do you think theory plays a role in getting us there?

Yeah. I mean, I think, okay. One important caveat is like when it comes to like architecture design and especially specifically like architecture design that scales, this is like well outside of my wheelhouse. So I'm happy to get some hot takes, but like, this is not—

That's the part of this podcast.

That's what we need. Yeah. I think for the conversation that we had a while ago, and this is still something I'm thinking about now, is like, I really don't believe in like many of the building blocks of modern architectures. Like, I think like we talked about this a while back and then you have this recent paper on MLPs where I don't feel like MLPs make that much sense as like if the evidence is that MLPs are just acting as like semantic lookup tables, there should be better ways of doing that. Either both like more parameter efficient, more compute efficient, like we shouldn't be having to do all of this compute just to do like a lookup. Right. Um, and so like replacing that with something that feels more targeted towards lookups just seems like a reasonable direction, but again, not something I'm working on. Um, same thing with self-attention, same thing with like tokenization feels very like ad hoc. Yeah. Um, I guess in reference to the bitter lesson, the tokenization does feel like it kind of goes against the bitter lesson where you should just let it figure it out. You shouldn't be manually trying to tokenize things and then feed them into the transformer. You should be feeding in as raw data as possible and let the model figure it out. Yeah.

And I mean, people in this lab specifically, and a lot of our research— Tolga and I have worked on this.

I know.

It's a surprisingly very, very hard problem, right? Again, we talk about these AI models as incredible. And like, well, you know, what is like one of the most basic units of human language? Just like just semantic bucketing. Like, can you figure out what a word or phoneme is? It's like the modern LLM during training just never solves this problem. It's like a very basic chunking procedure. Seemingly it's like intractable, right? Yeah.

But then as far as like architectures at scale, I think I honestly don't really have anything to say. Like, I think this is always shocking to me that I think there are a lot of things about the modern transformer that that do just seem really strange. A lot of architectural decisions that do just seem really strange. And I think we're more driven by systems and the need to scale than by some thinking about, oh, what is the right way to represent this thing? So I think one of the things that's always very surprising to me is when you're scaling up a transformer, you scale up the embedding dimension with the number of heads in a way that keeps the head dimension completely constant. And actually, if you look at the flash attention implementation, like it's stuck at 256 max. Like if you want to do like a larger, like head dimension or a different head dimension, then you're kind of stuck. I mean, you can hand, like hand roll something, but this does feel like it makes sense why this like came about because like in the forward pass, you want to keep most of your matrices for the matmuls like square, right? So like you do like embedding dimension to embedding dimension and you kind of chop it up. But at the same time, like from an actual like representation learning perspective, this does seem a little Strange. Seems a little artificial. I don't know. I don't, yeah. The entire field of architecture design just feels heavily constrained because like you're always going to be bottlenecked by like how efficient something is on like modern GPUs and modern like computer architectures. And this does end up like some weird co-evolution thing.

Yeah.

Yeah.

Oh, for sure. I mean, I mean, I think the, I think the caveat to this point though is that like, you know, folks at NVIDIA and like all these ASIC startups here around the Bay Area have like clearly shown that you can build an architecture that just works. Like the hardware will like adapt and like, I know we spent, like we should design for the hardware, the hardware will design for us. So yeah, I think this is actually like an extremely awesome part of the ecosystem.

Yeah.

But I guess the other like element of the co-design is the optimizer, right? Which is like a lot closer to your wheelhouse. And I mean, I guess one take I've heard is like, okay, the transformer is the thing that's dominant because Atom is the thing that's dominant. And then they just kind of co-evolve and you're stuck, right? Like maybe some other architecture with like some radically different optimizer actually is going to be better. It's like a lower well or something, but like the way to get there would have to be like a huge systematic change that we just like can't find as little researchers crawling along a little lost landscapes. Do you have any thoughts about progress on the optimizer side? Let's say both procedurally, how would we actually make progress in getting better optimizers? And also maybe on recent developments on the optimizer side of the world.

Yeah.

So I think this is definitely much closer to my wheelhouse now. There's a co-evolution of Transformers and Adam thing. It's always been very funny to me because there's just this folklore that Transformers are really hard to train by SGD. So like if Adam just never came about, like, I don't know where we would be. But, um, I think I've given this hot take before, but I do think that within a few years, we'll be looking back at the way that we were training models now and just like have a lot of thoughts. Like, I think we're doing things very wrong right now. I think Adam in particular, like, it's, it's incredibly useful, but I do think it has like huge drawbacks. And I do think that it's going to limit the like models that you can train. In terms of recent developments, I think that MuON is very interesting and it's the last in a longer line of work. I think probably originally going back to Shampoo. But I think this holds a lot of promise in the sense that it makes sense to— it's one of these very few optimizers that actually opens this black box of the neural network when we're trying to do optimization. Where Adam just treats the parameters of a neural network as a really long vector and you don't really care where these come from, but this is actually using some matrix structure. And so this makes sense. But I think more generally, just thinking about information flow through the network, trying to connect some of the ideas from signal prop, like the MUP-esque ideas to optimization, I think is just like a super interesting direction to work on. And I think there'll be a lot of successors to Mewon that will perform similarly better as well. But I think my own personal take for how to make progress on optimization is that we have a lot of optimizers that work surprisingly well, like shocking well. And I think Adam is one of them. And I really want to first take the time to understand what they got right. Because between like now and like 2015, when the Adam paper first came out, there are a lot of different optimizers that have been proposed and none of them managed to like definitively dethrone Adam. And so it clearly is doing something right. And it's not something that we fully understand. And so I just feel like as a general principle, It makes sense to try and figure out what it's doing right and why it's so dominant before we go trying to design new optimizers. Because there are definitely some lessons there to be taken.

Yeah. Which your recent work gets into.

My recent work gets into. Yeah, it's definitely just partial part.

So yeah, what types of questions or what are the types of statements you think might be true, I guess, about Adam? Is it kind of data distributional? It just happens to be really good for natural language because of whatever, or like low-dimensional, right?

Like, curvilinear properties or whatever.

Or architectural stuff.

Yeah. We need architectural. I kind of like, it's, I think especially in transformers, the different layers get balanced, like the, the, the, the size of the gradients in different layers can become radically different during training. I think even like within a layer, like different rows or columns can have radically different size during training. And so just having something that adapts in a reasonable way to this, I think is just going to be necessary. This is especially true, for example, like you have these like softmax layers in a transformer. So this like kind of gives like logistic regression vibes. And for something like logistic regression, you really don't want to just use like vanilla gradient descent. You have much better optimizers. And it's just because like, as you begin to saturate the softmax, the gradient shrinks to zero. And then you can just take larger step sizes and you should take larger step sizes. And so it's just like this potential is definitely there. But I think more generally, the principle I'm most excited about, and this actually does relate to the central flows work, is that I think both my work and a lot of other work has shown that a lot of simple optimizers, Adam SGD included, actually exhibit surprisingly sophisticated optimization strategies, but they're very implicit. And a lot of this is just from the fact that the optimizer is kind of bouncing around locally, either because of this edge of stability effect I described earlier, or just because of the noise in the process itself. And anytime you're comparing gradients at two nearby points, the difference in these two gradients is a Hessian vector product, which kind of gives these optimizers implicit access to Hessians. And I think they actually use this automatically to do good things. And so I think part of the success of Adam is that it was just like some update rule where when you work through all of this machinery, the like implicit way it uses the Hessian is just quite smart. And so in terms of like developing new optimization strategies, I think, I think we should be bolder in terms of what types of estimators in the Hessian we should be looking at. Because I think you actually can derive like very good estimates of the Hessian very cheaply using just gradient penalty loss.

Yeah.

Do you have any thoughts about kind of getting closer to like a full second order type approach where you actually spend more and more compute, like getting very, like even computing, like, you know, hashing of some little mini block or something? Like, do you think that's kind of a realistic, like a path forward to better optimizers?

Yeah, I think you once told me, I asked for you, it's like, do you think there'll be better optimizers out there? And your answer was like, I think everyone thinks this is true. Second order is one, should it— wait. In theory do better, but for whatever reason, all the off-the-shelf implementations just do not work.

So I mean, I think this comes down to like, well, okay, it also depends on what you're looking for out of the optimizer. So there was recently this paper out of Harvard where they were trying to compute, it's an approximate like full Gauss-Newton step, and they were trying to train an LLM using these Gauss-Newton steps. And so here it's closer to like a full Hessian inverse. And you could empirically train an LLM with like on the order of like 30 steps of gradient descent. I mean, these are very large batch steps. It's like a lot of samples getting put into it, but just as like a proof of concept of like where optimization could be in terms of iteration complexity, that's quite good. Like that's really a small number of steps.

Modern problems require truly ancient solutions. Like first optimizer we ever developed as a species. Pretty good.

Yeah.

Well, the first one would probably be gradient descent. All right. Like the second optimizer we as a species developed. It's in the name, right?

Yeah. No, I mean, I think there is just a lot on the table here. And I think a lot of it doesn't necessarily come from like me looking at optimization and being like, oh, they're all doing it wrong. This is what you should do better. Because obviously if that were true, I would just have released the optimizer. But more just like, it does seem like there's a lot of inefficiencies in the system. And the same way, like you can look at a transformer and look at the MLP blocks and just be like, there has to be something here. I think I can look at optimization right now and just be like, there has to be something here. Like, um, and I really like these like types of proof of concept experiments where it's just like, how much faster could we be doing if we really like could approximate the hash?

Right.

And like just as inspiration to actually try and come up with better— Yeah.

Yeah.

Good interactive algorithm, but for optimization. Sure.

Yeah.

Super interesting.

Yeah.

Okay. I think I have to ask the opposite question as well as the François step-in for today, because I think I'd be incredibly remiss if I did not ask you at zeroth-order optimization.

Okay.

And I know you were on this paper, Izzo? Mazzarella? Yeah, which is at Princeton, zeroth-order optimization, which does surprisingly well, I would say, for fine-tuning.

For fine-tuning.

Are you a believer in zeroth-order optimization?

I mean, I believe it has its place in very memory-constrained settings. I don't think it's ever going to be the de facto for pre-training, or at least if it is going to be a de facto for bigger-scale problems, I think it does just need to grapple with like, like I think the margins are just going to be very small because in the end of the day, like a gradient step is more expensive than a forward pass, but really only in memory. Otherwise it's like a constant overhead for the actual runtime. And like memory, you can then use gradient checkpointing to reduce further. And so like the baseline of like a well-tuned implementation of Backdrop is just like a really strong baseline to beat. And so then I think there's just a tall hill to climb for zeroth order. I see. Um, But on the other hand, like it is true that you can optimize architectures for inference really well.

So maybe. I mean, for the record, just to your point on galactic algorithms, like, you know, I would like to point out the one example of AGI we actually have, you know, like marginalizing out by all the technical details, it was discovered through essentially a zeroth order method. Sorry, what do you mean? That's true. That's true. I mean, evolution is a lot more, I mean, it's like an open, I mean, a lot of open research questions is a lot more clutter than just like naive random search, but like, I mean, you know, it worked once.

That's true, but I also wouldn't call evolution naive random search, right?

Oh yeah, of course. At least from what literature I have read, and I've talked to a few people, it doesn't see an evolution in access to gradients and fitness potentials.

And I think it's got implicit bias.

Oh yeah.

Oh no. I think it gets far beyond any knowledge. It's got a— or maybe like the E. coli, just on the Cassian second order outside. It's really good.

Yeah.

Yeah, well, okay, I think this may be a good opportunity to ask some broader questions, I guess, towards closing the interview. But okay, you've talked a lot about kind of the future and lots of promising directions you're excited about. Looking backwards, has there been a take in ML that you believed like a couple years ago that you believe is now fundamentally untrue?

Hmm. Yeah, I mean, I think the first thing that comes to mind is like, I do think there's a huge shift in the types of questions deep learning, that like deep learning theory was attempting to tackle during this transition from like the ImageNet era to like the LLM era. And I think one of the biggest things is just like a shift in generalization focus where there's this like massive body of literature on the connections between like flatter minima and generalization. It includes the, like, if you use, like, larger learning rates, smaller batch sizes, then SGD will find flatter minima, and then these flatter minima will generalize better. And then people were trying to develop optimizers that would do this even more. You're right, it's like Schwarzner's aware minimization. And there's just, like, this huge, huge story here where it is, like, why do these, like, models generalize? They generalize because of this, like, implicit bias of the optimizer. And I think that this general story is just not, maybe it's not as relevant for the LLM era, but also I just think is not necessarily true for the LLM era in a sense that I'm not convinced that you necessarily get like better generalization with like smaller batch sizes or larger learning rates. And more, I think the question, the types of questions that people are asking in generalization are just different. Like it's not really in-distribution generalization anymore.

Just so.

So then I think like, yeah, this like general Discussion around generalization, I think, was not necessarily the most fruitful. I think the other thing here is like, there's a lot of focus on getting generalization bounds, like asking like, why do neural nets generalize? And then trying to figure out some notion of capacity control, like trying to formalize Occam's razor, that like simpler models should generalize, but then what exactly is a simple model?

All that work in statistical learning theory, for example, right?

Yeah, yeah, exactly.

Right.

And so like, there's this huge focus on like, like in statistical learning theory for simple models, we have a very good idea of what constitutes capacity control, right? And we don't quite understand this for neural networks. And I think here, one of the really challenging things in retrospect is that it's very hard to know what success looks like because you can ask philosophically, why should a neural network generalize? And here you actually have answers. There are capacity— there are generalization bounds for neural networks, not necessarily that tight, but they do exist. And you can get parameter-independent generalization bounds so that As the network gets wider, the generalization problem is not getting worse. And so this at least tells you philosophically that if you had enough data points, the neural net should do a good job, right? Assuming you could also train to zero loss, right? But then it doesn't like, somehow the counterargument to this is like, oh, but like when you actually compute them, they're vacuous. And I think this was like a different, this is not a critique I necessarily agree with in the sense that it's not clear to me that it's a reasonable expectation to expect a generalization bound that to me is more of like a philosophical thing, like why shouldn't something generalize, to then predict exact generalization performance. And especially if your generalization performance is like, you get 70% test accuracy, for a generalization bound to actually measure that error correctly, the constants have to be really tiny. Your margin for error is really small. And now I feel like you're getting much more into high-dimensional asymptotics, the type of statistical learning theory that's actually trying to make it very precise numerical progressions.

And most theorists don't enjoy grinding out those constants, right? That's the worst part of the job.

That's the worst part of the job, right? Even if she somehow had a good notion of capacity control and can prove some inequality, it could still be vacuous, right? And so I feel like in general, I'm not quite sure what the conclusion of this line of work was.

Yeah. Interesting. Okay.

Yeah.

So I guess turning it the other way now, are there problems that you think are very promising and maybe kind of ML theory as a community doesn't, hasn't, you know, put enough focus on, let's say, you want more folks to work on?

Interesting. I think—

And why is it central flows? No, no, no, no, no.

I mean, okay, I do think optimization theory is like very fruitful, but I also think a lot of people are working on optimization theory. I don't think, I don't think I would say it's like underexplored by any means. But I think there was like when, after deep learning became popular and when deep learning theory started to become like a major field, one of the central questions was focusing on these hierarchical representations and hierarchical representation learning. And it's not that people are not working on this, but I do think that a lot of people worked on it for a while and they kind of moved on to other more general things. And a lot of the basic questions about how do you learn hierarchical representations Or in what sense is data, like real data hierarchical? I think a lot of these questions are still underexplored and not that many people are continuing to work on them. And so, yeah, can you train a deeper network and try and actually understand the development of hierarchical feature learning? Partially this problem is just very hard, but I also just feel like there have to be some breakthroughs coming in the near future. So I think maybe that.

I mean, I can tell you here at Casey, we're all very excited about hierarchical feature learning. Yeah, my personal bias working on just biological data, like, I mean, a lot of the priors in, say, like, you know, like a lot of biology is like, it's very hierarchically organized, right? You know, you have like molecules, these like make cells, cells form tissues, tissues form organs. I guess I'm just plugging my own work at this point. But, you know, it's very exciting. Like, I mean, yeah, like we can have a really more meaningful way to get our hierarchies reversed.

Yeah, yeah, of course.

I mean, it's huge.

Yeah.

Well, I feel like all of this, like both your answers to the past two questions have like centered kind of around this like question of having better data results in an L theory centering around data. And I think what's interesting too is, at least my understanding of some of the most recent frontier LLM releases too, has been a lot of the recent progress in quality has been because they've actually empirically looked at the data level. They're like, turns out you got to really dig into this, denoise, throw out garbage stuff, or synthetic data, all of that. Yeah. Do you have any thoughts here about, do you see it? Paths towards saying more interesting things about data in the theory world?

Yeah. So I think that, okay, as I mentioned at the very beginning, I think theory is like, or theory is struggling with this data question. I think it's just the hardest part of deep learning theory by far. And I think part of it is just that real data is very messy and we don't really have good mathematical models for what real data looks like. Yeah, yeah, yeah. The reason that we have transformers is that we don't have a model for natural language. And so like there's a chicken and egg problem.

I mean, I always make this comment to people. It's just sort of like, well, I'm using transformers. I don't have a theory. It's like, why are you running it? If you do sort of ML for physics, you're competing against some like ridiculously good models that are like interpretable and they're beautiful and they have like centuries of work on them. I mean, bio is a bit like, you know, way messier. We benefit a lot more from these models. But yeah, anyway, sorry.

No, but I think that the answer here really is just like, and this is not like underexplored. I think a lot of people are working on it, but I do think it's like one of the most interesting and exciting directions in deep learning theory, which is just coming up with like interesting models for data, which try and capture some aspect of real data that we're trying to understand. And it's never going to be like fully realistic, but if you want to capture like semantics, or if you want to capture like some sort of structure, this is something that you can try and encode in a more synthetic problem and then actually do real theory on. That's like tractable because you actually have a mathematical model to work with. Whereas if you're trying to just directly do things on real data, it's very hard.

Right.

Yeah.

I feel like, and the question of like, what is, a valid toy problem in some sense is just the hard bit, right? Yeah, exactly. There's no good answer.

There's no good answer, right?

But I think the best you can do is like, you try and capture some aspect of something that's real and then do as much like cross-checking to see if like you're actually capturing the phenomenon that you were originally trying to observe. So if you're like trying to train a model and you observe some like spike, maybe it's like a Grokking-esque thing on like modular arithmetic. They're like, this is a modular arithmetic is already like a toy enough setting that like Interesting things happen. You observe rapid skill acquisition, which you can also observe in large models. And then at the same time, it's not that complicated of a generating a procedure. You can actually try and study this. I think just coming up with more things like this. And then the attempts at trying to understand the training dynamics of just trying to fit massive amounts of information into the weights. I think this is where scaling law research is going and just trying to cram information into very small numbers of parameters. So like, I do think there, it just comes down to like, isolate the part of the data that you want to understand, come up with a nice model that encapsulates that, and then see how well you could study that. Yeah.

Yeah. You want a bunch of hydrogen atoms for different aspects of data and just like really study the essence of detail and like, you know, see how we can, see what we can learn.

Yeah.

Yeah.

No, easier said than done, of course, but yeah.

Yeah.

Yeah.

A ton of exciting stuff to do. Um, okay. Maybe we can close with We've been talking a lot about all the theory work that you find exciting, that you've been doing yourself. Maybe closing with the empirical stuff, is there lines of work there that you pay attention to the most, that maybe you think there's maybe interesting questions to move into the theory world? Or yeah, what do you pay attention to in those spaces?

I think recently, the attention sinks line of work, I found it very fun. Just because I think it also connected with things that I've been thinking about in weird unexpected ways, because I was thinking a lot about Optimization and loss spikes. These are also connected to loss spikes. I finally saw the— this is like a very, I think a very fun line of work. Let me think. I think a lot of empirical papers that are optimization focused, I find very interesting because I think these are like the easiest way of actually getting some sort of like ground truth for like basing a theory off of. And so there was recently like a very fun paper out of Google that was trying to, I mean, part of it was like they looked at like the Jackson kernels to try and compute like actual, like per example gradients and use this to get like a different estimator in the Adam denominator. The idea being that if Adam is really estimating the variance of your gradients, then like you can estimate the variance better. Like if you compute a mini-batch and then you really want like the expectation of like each individual entry in the mini-batch's gradient squared. And so to do this, they kind of had to be internals a little bit, and then they observed, hey, it performs worse. And it actually is consistent with a lot of other evidence that the gap between SGD and Adam is largest when the batch sizes are really large. And so first, I think this justifies studying large batch or full batch Adam as a good testing ground to see where this gap really lies. But I also just think it's nice because it challenges a lot of existing theory that Yeah. If you're, if you actually think of it as trying to estimate like a Fisher information matrix second moment type of thing, and you estimate that worse and performance gets better, it does challenge the original assumption that this is right. That's doing it.

Yeah.

Yeah. Yeah. Okay.

I mean, I'm glad we're closing this podcast with, I think the essential theme that we should all take away, which is everything's broken and there's still so much water to you.

I don't think everything's broken. I think things are working like surprisingly well. Sure.

Sure.

That's both can be true. I think that can be true, but awesome. Okay.

Well, yeah, actually we should. And well, first of all, echoing that point 100%. It's actually, it's actually a great time to do research. It's actually just fantastic. Every problem is, there's so many great open problems, but I think we should actually end with, so you're starting at MATH in the fall.

Yes.

Give me a message to put out into the world for prospective students or postdocs or.

Yeah, I mean, I'm just like generally excited to work on deep learning theory and just see where this adventure takes us. So I'm interested in like optimization, as I mentioned a lot in this podcast, but I'm also interested a lot in representation learning, in statistics. This is like not something that's like super deep learning focused. I'm really interested in like computational statistical gaps and these fundamental limits of learning. So if you're interested in this, please reach out. I'm happy to chat. And yeah, that's pretty much it.

We'll add links in the YouTube if you're on YouTube. Awesome. Well, Alex, so great to have you here. Thanks for sharing your work.

Thank you so much. Thanks for having me.