In short
Episode Notes: Is It Time to Rethink LLM Pre-Training? with Aditi Raghunathan - #747
Podcast Information
- Podcast Title: The TWIML AI Podcast
- Host: Sam Charrington
- Guest: Aditi Raghunathan, Assistant Professor at Carnegie Mellon University
- Episode Number: 747
- Episode Date: Not specified in the transcript
- Episode Link: [TWIML AI Podcast Episode 747](https://twimlai.com/go/747)
Overview In this episode, Sam Charrington interviews Aditi Raghunathan to discuss the limitations of large language models (LLMs) and how we can build models that are more adaptable and creative. They delve into Aditi’s award-winning paper, "Roll the Dice & Look Before You Leap," which critiques the creative constraints of current next-token prediction models. The discussion also touches on the concepts of catastrophic overtraining, memorization sinks, and the challenge of model reliability.
Key Topics
Background of Aditi Raghunathan
- Academic Journey:
- Background in complexity theory and deep learning.
- Influenced by adversarial examples and the unexpected failures of capable models.
Limitations of Current LLMs
- Benchmarking vs. Real-World Performance:
- Current benchmarks often do not reflect real-world usage scenarios.
- Significant gap between benchmark performance and user experience.
- Catastrophic Overtraining:
- Training on more data often leads to worse model adaptability during fine-tuning.
- Example: More data can adversely affect model performance after quantization.
Proposed Solutions
- Roll the Dice & Look Before You Leap:
- Roll the Dice: Introduces randomness at the start of generation for better exploration.
- Look Before You Leap: Encourages models to take "leaps of thought" and explore diverse outputs.
- Memorization Sinks:
- A proposed architectural change that isolates specific information to enable targeted unlearning.
- Allows models to retain shared knowledge while enabling specific unlearning of harmful or private information.
Creativity in AI Models
- Understanding Creativity:
- Current models struggle with open-ended tasks and generating truly novel ideas.
- Creativity can be decomposed into two categories:
- Combinational Creativity: Finding novel connections (e.g., humor).
- Exploratory Creativity: Discovering new concepts or solutions from existing structures.
- Training for Creativity:
- The need for new training objectives that promote creative thinking.
- Exploration of multi-token prediction and diffusion models as more effective training strategies.
Future Directions
- Adaptability and Update Mechanisms:
- Rethinking how models are updated over time without retraining from scratch.
- Exploring structured exploration in reinforcement learning settings.
- Call for Research:
- Emphasis on understanding LLM behavior through controlled experiments and formal hypotheses.
- Encouragement for the research community to focus on these underlying challenges to unlock new opportunities.
Conclusion Aditi Raghunathan's insights highlight the pressing need to rethink how LLMs are trained and utilized in real-world applications. The conversation emphasizes understanding the fundamental properties of models and fostering creativity, adaptability, and reliability in AI systems.
Key Takeaways
- The limitations of current LLM architectures need to be addressed to improve adaptability and creativity.
- Randomness and structured exploration can enhance the creative capabilities of models.
- Insights from Aditi's research suggest promising new approaches for training more effective models.
- Understanding and controlling model behavior through precise experiments is essential for future advancements.
Additional Resources
- Paper: "Roll the Dice & Look Before You Leap: Going Beyond the Creative Limits of Next-Token Prediction"
- Related Work: Research on memorization sinks and catastrophic overtraining.
Feel free to explore the complete show notes and additional resources linked in the episode above.
Written by AI. May contain mistakes. Listen to the episode to check what was said.
Transcript
Automatic transcript. May contain errors.0:00I'd like to thank our friends at Capital One for sponsoring today's episode. Capital One's tech team isn't just talking about multi-agenticentic AI. they already deployed one. It's called Chat Concierge and it's simplifying car shopping. Using self-reflection and layered reasoning with live API checks, it doesn't just help buyers find a car they love. It helps schedule a test drive, get pre-approved for financing and estimate trade in value. Advanced, intuitive and deployed. That's how they stack. That's technology at Capital One. You know, we measure performance on a benchmark and if that's all we care about, it seems like we can do really well because if you collect data that looks like the data you want to do well on, you can just, you know, throw a lot of compute at it.
0:45But like, does that actually solve the task if we, you know, just test it in a slightly different way that is also meaningful from a deployment perspective? But like, when does that break the models and why does that happen? And how does the, you know, the training dynamics or the data curation, like what aspects of these actually influence this behavior and what's the right way to intervene and make these problems, you know, go away?
1:17All right, everyone, welcome to another episode of the Twin Mill AI podcast. I am your host, Sam Charrington. Today, I'm joined by Aditi Raghunathan. Aditi is an assistant professor of computer science at Carnegie Mellon University. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Welcome to the podcast, Aditi. Yeah, thanks for having me here. Excited for the conversation. I'm excited for the conversation as well. We are going to be digging into one of your recent papers, which won the outstanding or an outstanding paper award at ICML 2025.
1:54That paper is called Roll the Dice and Look Before You Leap, going beyond the creative limits of next token prediction. But more broadly, your lab has been really digging into some of the limitations of current LLM architectures and opportunities and, you know, what we need to understand to better make use of AI models. And I'm excited to talk through those with you. To get us started, I'd love to have you share a little bit about your background and, you know, what got you excited about AI and machine learning. So I guess in my undergrad, I was always excited by complexity theory. And I really liked the elegance of thinking about what's possible and what's not.
2:38But at the same time, maybe like several others, I also had the age to work on something that had immediate impact, was practical. And when I started doing research was around 2015, 2016, when deep learning was really taking off. And this was post-ImageNet and so on. And so naturally, I got excited about it. But really, the moments that shaped my thinking around this was when I was at Stanford, I was attending a couple of talks, and one of them was on adversarial examples, which showed me how these really capable models can also fail in spectacularly kind of surprising and maybe seemingly dumb ways.
3:14and they also had some immediate, you know, practical questions around security, reliability in the wild and so on. So that's sort of what got me thinking about like, how do we think about, you know, these really capable systems that also have these failures and it also tied into my itch on, you know, complexity and think about abstractions and making precise statements about these things. And it kind of felt like a nice combination because a lot of these aspects cannot be captured by just specific numbers and a benchmark and really need to go one step beyond to think about how do we, you know, what do these models learn?
3:45When do they work? When do they fail? And so on. So that's sort of what got me into all of these questions. And of course, the field's really been changing rapidly. And so I've kind of gone along for the ride. And what's kind of interesting is that progress has really, you know, models have become very capable in a lot of ways, but some of these fundamental failures still remain. And so we can keep asking the same questions about various different models. And maybe it's also sobering that it's not that we've actually pushed the reliability as we have pushed the capability of these models. And so in some ways, the same questions have remained over time, but also a lot of the context of these questions has been changing.
4:21It's pretty spectacular what we can do with the models, and yet we're still talking about not really understanding how they work and how it's a little bit of magic. Yes, absolutely. And I think the part that's also concerning is when this lack of understanding becomes a real issue, when people think about getting models to be safe in some way, to not utter toxic content or to release dangerous information to people or manipulate people in some way that's dangerous. And a lot of our guardrails around these things are very brittle. And we currently don't have a way to do better because we don't really understand these systems.
4:57And so I want to really use our understanding in a way to shape actually making these models more reliable in all of these contexts. So talk a little bit about with that context in mind, how you've kind of crafted a research agenda? What are your core focus areas? Yeah, so I've always been interested in thinking about, like, this is maybe a catch-all phrase of distribution shifts, but it does capture this idea that, you know, we measure performance on a benchmark, and if that's all we care about, it seems like we can do really well because if you collect data that morally looks like the data you want to do well on, you can just, you know, throw a lot of compute at it, and it seems like that just works really well.
5:35But at the same time, like how does that, does that tell us anything about like how the model works in any situation that's slightly different, but you still expect the model to work well? So that's sort of the angle that you've taken in a variety of different questions that we ask is like we can minimize some loss or like try to target a certain benchmark, which is what a lot of these models are trying to do. But like, does that actually solve the task if we, you know, just test it in a slightly different way that is also meaningful from a, you know, from a deployment perspective? But like, does, like, when does that break the models and why does that happen?
6:05And how does the training dynamics or the data curation, like what aspects of these actually influence this behavior and what's the right way to intervene and make these problems go away? This growing gap between benchmark performance and the experience of using these models, I think is kind of a growing concern and one that we're hearing a lot of this particular moment of time, I think, because of the recent release of GPT-5, of which, according to the benchmarks, is the smartest model around, but the user experience for many has been lacking in a lot of ways. And I think that kind of exemplifies this idea of the benchmarking as we think of it today being somewhat inadequate.
6:52Yes. And maybe one concrete way of thinking about this that I feel like is very important, but people haven't thought too much about is many times we want to use this as a starting point and do some kind of fine tuning. And so this could just be for safety alignment, that kind of stuff that we can say, oh, maybe one company just takes care of that. But also a lot of companies have their own propriety data that they want to fine tune on, or like you want to personalize it in some way to your context, or you want it to improve over time and things like that, or the world is changing. So we don't have a good way of measuring this adaptability of these models, which I also think is actually a very fundamental question.
7:29And as an example, I was starting with a colleague, Graham Newbig, who kind of asked me this question several months ago on like, which model should we start off if we had some data that we wanted to fine-tune on? Should we take the model that like as such on our benchmark or like on this data works really well, like the zero short performance of the model as is? Does that automatically mean it's better after fine-tuning? And yes, and we realized the answer is like, really no. And for a lot of things, actually model performance on a benchmark kind of tracks performance, including reliability, actually, or robustness to a lot of distribution ships.
8:07What we found, like, you know, what the community has found is, in general, the more data you train on, like all of these numbers go up. But this one aspect of how easy is it to adapt these models after you train them on data, we actually found the reverse thing happens at some point. So if you take a small model and keep training on a lot of data, we see that eventually the model that has been trained on more data that you've thrown more compute on is worse as a starting point for fine tuning than an earlier checkpoint that you had. So this was actually a really striking result because it's one of the first realistic cases where more compute in a very non-contrived setting, like on high quality data, more compute is actually kind of making a model worse, not just that it saturates, but it's worse.
8:53And another context in which this happens is, you know, when people are trying to serve quantized models to improve efficiency, there also we find this kind of trend where at some point, like as you throw more data at these models and then you quantize them, the model trained with more data still is better. But then you see this U-curve where at some point showing more data actually means that this model is worse after quantizing, which is again a case where your downstream model is strictly worse, even though you've spent more compute at the best of intentions and shown good data. So I think this was like a really interesting kind of finding that I think people should think more about is like one important aspect of using these models is actually fine tuning and adapting.
9:32And our current, you know, push towards just doing what we're doing now by improving this pre-training or like the first step of this process, but really optimizing that doesn't mean it's going to be good after doing your fine tuning or post-training. And this is not just a theoretical limit. We actually see that, you know, a lot of people talk about how LAMA 3, so this is before the LAMA 4 release. So people are just talking about LAMA 3 versus LAMA 2. And a lot of people found that LAMA 3 was much harder to fine tune, you know, in like academics regularly fine tune models. And so, you know, this was sort of Graham's question too.
10:05And it kind of ties into this idea that it's because we have trade on so much data that at some point LAMA 3 is still really good at benchmarks, but it's just worse if you want to kind of use it for your task. And similarly, we ran experiments on the Olmo checkpoints because, you know, great work. They're releasing all these checkpoints so we can do analysis on them. And we found that the model, the smallest size model, 1B, I believe that was trained on 3 trillion tokens is actually worse than the model that was trained on lesser tokens, fewer tokens after we do this kind of fine tuning or post training on very realistic benchmarks that people care about.
10:38So this was kind of an exciting result that shows that we might need to sort of one axis of rethinking pre-training is like, how do we get a good starting point for fine-tuning? Is it specifically the ratio between the number of tokens that the model is trained on and the number of parameters that you found to be inversely proportional to performance and tuning? Yes. So for the precision result, that's exactly what we found. And, you know, we try to give different exponents, but it turns out that in our experiments, the exponents turned out to be the same. So it's literally the ratio. But in our fine tuning, we kind of didn't do precise like curves because a lot of that depends on the exact distribution of interest.
11:17So there's no like a clean mathematical result that would hold or the trend won't be the same for all data sets. But in general, we find that it does like a larger model can take in more tokens before it shows the sort of like inverse effect compared to a small model. So in some sense, more like a larger model can absorb more tokens efficiently, but it's not always exactly the ratio. It could be a different exponent. In some ways, that strikes me as like an intuitive result in the sense that, you know, when that ratio is larger, the model is potentially more overfit and it would be harder to unlearn and learn things as you're trying to fine tune.
11:58I feel like another way to think about this is usually when a model, when we train models and, And I've looked at learning dynamics and things like that for a while now. And the converging result in that space is models learn simple things first, and then they learn increasingly complex things. And so my colleague had this kind of visual image of trying to build something with cards. Usually you start off with something really solid, and then you kind of add more and more complicated things. But then that also means that that structure gets less stable. And so if you actually try to adapt these models in some way, then everything collapses.
12:33in some sense. And so I think we see something similar that like you're forcing the model to learn more and more complex things, which is good because it's fitting the data. But then that also means that the model is brittle in some way that like you try to push the model to be a little, you know, in some direction by minimizing some gradient steps in a direction, but that introduces so much noise or like kind of breaks the model and causes a lot of forgetting. So that's sort of what, and it's same thing for precision too, like kind of we add, you can think of that as like adding some kind of noise by changing the weights.
13:01And so the models become like less, like what, like they cannot absorb that noise and they just break. It strikes me that looking at this ratio is a fairly coarse-grained metric, but that it would be interesting to go even further. Like if you could, for example, understand the relative distribution of the training data on the model relative to the direction you want to fine-tune it to. Like you might have a model that is, you know, benchmark, benchmarks worse, you know, is smaller or other reasons why you might think it wouldn't do as well. But because of some like distribution overlap or something might do better.
13:39Is that a reasonable direction, you think? Yeah. So there are different ways to think about this. So like one is purely kind of is the model more stable in that it can like move in more, can move in different directions without like losing too much. right and so that's this token to parameter kind of you know could like roughly correspond to that and then the other questions like what you said is like how much do we actually need to move and we have some experiments in our you know this catastrophic overtraining icml paper where one proxy for how much to move is just the learning rate um you know so if if the model gets good performance even with a small learning rate then that roughly means the model hasn't moved too much and it kind of works.
14:18The distribution is closer in some sense. So that is that we do find some interesting trends over there as well. So if we start measuring or fine tuning on things that are very close, then we actually don't see this effect. The models are able to take in more tokens and still show good performance because we don't really end up updating the model much. And in the limit, if the fine tuning matches a pre-training, then we just get back the usual, like things go down as we show more data. But then as we have larger changes, then surprise, like these kind of trends start happening where like at some point the model is actually getting worse.
14:52And one challenge is that it's not very easy to say which distributions are close or hard because we might have some intuition for what this is, but like how the model stores information might be different, which is sort of why we look at learning rate itself as like a proxy for how much the model changes as a way of saying like how different the distributions are. And then when I hear catastrophic overtraining, I think of not just that there's like this inverse relationship, but also that there's maybe like a cliff, like you reach a point and then your fine-tuned ability kind of falls off a cliff and it's much worse.
15:30Did you kind of characterize what that point is? Yeah. So I guess like the way we, so what actually happens is like at some point, I mean, like overall in the, there is a regime where showing more data does help, like no matter what the model size, because the model is just learning stuff. And then at some point the model gets so brittle that whatever it learns from the more data kind of just like is overcome by its brittleness. And so like we start seeing this jump. So there is a point where more data starts hurting you. So it's like a U-shaped kind of situation. And so, and the point at which this U is really depends on kind of like we discussed, like it depends on the data that you're fine tuning on.
16:08And yeah, so we are able to characterize these precisely in the quantization setting because there's just fewer factors to count for because you just add noise. But in this fine tuning settings, a little bit more tricky because we can't exactly say what the distribution of interest is. But if we fix a certain distribution or a certain learning rate that we care about, then again, this point actually becomes predictable. So we didn't go into modeling this functional form too carefully in the fine tuning setting in our paper. And I think that would be very interesting future work for once we decide what kind of distributions you want to fine tune on and you want to make decisions on how to allocate compute and do the scaling laws and so on.
16:45But we do exceed trends where it looks fairly predictable with a nice mathematical form. So having done this research and identified these results, if you were to approach a future task and need to fine tune a model, how would it change the way you think about base model selection? You know, beyond the generalities, like would you, it's good to know because it helps you understand, but you would still like test everything and see how it performs on your data? Or is there just a whole set of models that you would no longer look at, for example? Yeah, that's a great question. I think if you are using a really, I guess what we found is the 1 billion model size, which again, you might want to use for efficiency reasons.
17:29So if it's trained beyond 2.5 trillion tokens or maybe even true trillion tokens because we didn't have too much granularity in the public checkpoints, I would say like, okay, maybe that's a point where you don't want to use that model unless you're sure that you don't want to change the model too much. And in terms of like, so that's sort of like one thing that's clear. In all the gray areas in between, I think we do have to do some preliminary experiments where we take a bunch of different checkpoints and see like what the fine tuning performance actually ends up looking like. And that would, you know, give us some sense of like, are we, which part of the U curve are we on?
18:04And so can we actually, so I would say like that's the understanding it gives us is that we kind of expect this shape and we can try to probe to see like which regime current models are. And so that would guide us like where the optimal could be. And this is the paper, overtrained language models are harder to fine tune. We'll link to all the papers that we discuss in the show notes for folks. and so you're just kind of contextualizing this and we'll get to creativity in a second but you published a really interesting blog post that kind of surveyed all of your labs papers at ICML and it was kind of broken up into limitations and opportunities and one of those limitations is overtraining and that's kind of what we just talked about.
18:47The next one was unlearning and the challenges associated with unlearning, which is related to this idea of fine-tuning. Talk a little bit more about what you've seen with unlearning. Yes, that's like, yeah, it's actually very related. And so like when you think about what is the use case of fine-tuning, like one is, you know, just to specialize and like push a few benchmarks, push a few numbers on your use case rather than the benchmark. But actually another important use case of all of these like post-training or fine-tuning methodologies is safety in some way. And so a lot of the alignment work is actually trying to teach the model post hoc what is good and what's bad.
19:24And similarly, we might try to unlearn harmful knowledge or unlearn private information and so on. So that's like a very safety specific use case of fine tuning. And we tried to look at why is it so hard to do unlearned English. There are so many papers that are published. And it very much reminded me of adversarial examples from my PhD where people had all these ideas for defenses. But then, you know, Carlini would break all of them. And so, and so kind of, yeah, so that, so yeah, so it was basically, it kind of felt like that. And if you look at some of the assumptions or kind of how the unlearning field has progressed, it's, people have tried to take the base model as it is, like, or like trick the starting point as it is, and then try to assume certain things that might be happening in these models and use that to get algorithms.
20:10And to maybe, maybe interject and to be more concrete about unlearning, you mentioned safety. But, you know, example might be you may have in the training data how to create a chemical weapon. And there's a whole, you know, line of work around like building guardrails to like detect that and suppress the model from talking about that. But another direction is to try to just extract that information, erase it from the model. And that's unlearning. Yeah. And it could be harmful information. It could also be private information that they shouldn't have trained on or someone wants to remove that information.
20:47And so it's something that exists in the model and you kind of want to remove that. And so I guess maybe the privacy angle also tells sort of why the guardrails feel sort of, you know, you kind of really just want to remove it from the model and want it to be like it wasn't ever trained on this data. So, you know, that's the other used case for all of these things. And so what people have been trying to do is, so you can start the first, like the first attack would be, or like, you know, first way to try to address this would be to fine-tune so that the model has high loss on all of these things that you want to forget.
21:18And then people realize that you can't actually do too well because then you kind of, you have to like really change everything a lot and that destroys a lot of information in the model. And then there's another sense that maybe that this information could be localized or maybe we can find specific neurons or specific subspaces and try to just erase those parts. And even that has actually limited success for two reasons. One is we have to first find a way to figure out where those neurons are, which store this information. And second is even, I mean, and even after we find that it's not clear, like what's the right way to erase that without destroying everything else.
21:53But the assumption here is that there exists such neurons in the first place. And what we show is that that actually is not true. That there's no reason in how we frame these models that encourages that should allow this information to be disentangled in this nice way. And it seems like maybe we see such separation, but because the separation is not very good, that's sort of why our unlearning methods don't work very well. And so we instead say, instead of constraining ourselves to work with the starting point that is not very good and that it's not really disentangled all of these things, what if we give ourselves a flexibility of like what if we could train our models in a way that enables this kind of downstream unlearning because we know that that is a use case which we might care about.
22:36So that inspired this work on memorization sinks, which is kind of taking the same idea or assumption that people implicitly try to make about models that maybe information is isolated to neurons, but instead of waiting for it to happen by magic, we are like, let's try to actually enforce that by design. And we find in our paper, both for experiments and analysis, that normal training does not actually lead to this assumption to be true because like, yeah, we have analysis in the paper, but kind of the main idea is like, we don't really encourage or force this sort of disentanglement. And it really depends on the bias of the training algorithm.
23:12And the current training algorithms don't have a bias to actually enable this kind of separation. But we can encourage this separation. The assumption that we're talking about is that knowledge is localized, essentially. Exactly, yes, yes. That knowledge is localized, yes. So there's no, like, if you look at the training objective, you're just passing gradients to all the parameters. There's nothing to do with it. So there's nothing about it. There is some separation that seems to emerge, but it's not perfect because it doesn't have to be. It's not trained to be. But what if we instead trained to encourage this sort of separation?
23:46And that's exactly the idea behind memorization sinks, where we're trying to say for every document, let's say that these are some specific neurons that are only updated on that document with the hope that all the information that is specific to this document, like goes to that neuron, and those neurons are not touched or updated on other documents. And so that's the main idea here. And we, of course, want to have shared neurons because we also want the models to actually learn from all of this information. And we don't want to train completely decentralized models because those won't be as capable.
24:17So we're trying, when you're doing pre-training, we want to model to learn something that is shared. And this is where, again, the beauty of the training or the inductive bias comes in, which is with this architecture, the model actually is incentivized to keep the shared information that is shared across all the neurons, sorry, across all the documents in the neurons that are updated on all the documents. That's just a strictly better solution. And the stuff that's very specific to a particular document is in these memorization neurons. And since those are not updated on any other documents, that information is sort of preserved, disentangled, kept aside.
Read the full transcript
24:49And we find through our experiments on a somewhat small scale, but we're scaling that up now, is that this architecture actually enables this kind of nice separation between what is special or like what is unique to your particular documents. That's all kept in a specific neurons, whereas what is shared is allowed to be learned in this shared neurons. And at test time, you can just drop out these memorization neurons and then you're good to go. A couple of questions. So are you, at train time, are you identifying the information that you will later want to pull out of the model? Yeah, that's a good question.
25:26So the way we're setting it up right now is that we are just having the units that we might want to remove. So if you feel like this document has information about a unit who might want to remove this information, then we kind of associate this entire document to a specific set of neurons. So it could also be like a topic. Like let's say we say that we don't want to be at the document level, but we want to be at the topic level. Then you want to have neurons for all the documents for a particular topic. we want to selectively kind of activate only those neurons. So we need to know kind of that abstraction that we might want to remove later.
26:01It seems like the number of these memorization sync neurons would be a hyperparameter and like the ratio of those to the total number of parameters is kind of an interesting mix. To what degree did you experiment with all of that? Yeah, we experiment with all of this. So we do need to, and like, so the models have to be a little bit bigger in this way because we do want more neurons. But one other kind of nice trick is we don't need them to have completely separate neurons. We just need to make sure that the neurons are somewhat orthogonal. And so we can just pick random high dimensional directions and they are almost orthogonal.
26:36So that kind of means we can have a, so we can, instead of operating on individual neurons, we take these neurons, but activate different subspaces that are fairly orthogonal. And that also works as well. So that's a trick to prevent having a really large model size. if you were to otherwise have a specific neuron for like every document or something like that. So that's one way in which we're able to like make sure that the model size, you know, doesn't go up way more than what it is. So we have ablations of paper where we find that for moderate increases in model size, we're actually able to kind of encourage this behavior.
27:09Are these neurons like localized in the architecture like to a particular layer or are they distributed? Are they, is kind of the topology totally learned or is it a priori like set up where the, these neurons are? Yeah. So what we do in our paper is we pick a random, for every document, we have a hash or like some kind of random neurons that that particular combination of neurons is activated. And so this effectively models like some one dimension in like one direction in a high dimensional space that's like like left to this new document uh one could consider smarter or like there might be improvements to this on like kind of which layers should we look at so we we only looked at the mlp layers and um we only introduced these things within the mlp layers sort of building on the intuition people had that facts or factual information is generally stored in the mlp layers uh but like i could imagine like there's a lot of research that one could do to like you know be more intelligent about this but we kind of learned this and right now it's just random.
28:15We could also try to kind of have, say, related documents to maybe share similar subspaces in some way or something. So we could even have a softer version of this where instead of everything being orthogonal, we can maybe put things that are slightly close together to be closer together or so on that might further push up performance. I think there's a lot of things that one could do here. And I'm very excited about this because I think it really kind of tells us ways in which we can get more control over these models by design. It also makes me think about kind of the overlap with the anthropic circuit tracing work and if there is some way to combine these techniques to better localize concepts.
28:57Yeah, that's a great point. So one thing that could be interesting is, so maybe natural training already has a propensity to have certain kinds of structures, but like it's not perfect. So maybe we can start with like, okay, these are structures that seem easy to enforce. So let's try to like actually hard enforce them. And so maybe that's a way to like take these fuzzy kind of interpretability stuff and like actually train a model to make those more hard and more concrete to enable better control. Yeah. And to a first order, we kind of do something like that, which is people have kind of thought that knowledge is isolated.
29:27And so we were like, okay, let's try to really isolate that because it seems plausible. So we could imagine doing other things like that. That paper is called Memorization Sinks, Isolating Memorization During LLM Training. It also brings to mind for me, you know, maybe it's just this word overlap of memorization, but a lot of conversation now is talking about the role of memory with LLMs and, you know, in particular how, you know, memory is not a very robust feature of like attention based LLMs and how memory architectures are a promising way to increase performance. Like, does that relate to this in any way?
30:08I mean, I can think of a philosophical way in which I think they're somewhat related in that I think in some sense, we want to like have, like another way that memorization actually could be useful is not just for the privacy, but it tells us like we can disentangle things that should be constant or kept in the model and things that should be updated. So for example, facts change, but we want to still preserve the ability to reason or the linguistic capabilities of the model. And so if we had ways to kind of disentangle those, then that helps with some of these things. And so you can think of memory in the context that you were saying as stuff that we kind of want the model to actually remember.
30:44And so the more we can kind of disentangle stuff that should be kept around versus stuff that's independent of that, I think in that sense, they are sort of related. And maybe there are some ideas that we can cross share on like how the exact architectures are developed and so on. Yeah, yeah, yeah. And I don't know what the kind of neuroscience implications would be or what we've learned from neuroscience, but it strikes me that like in humans, like fact memory and concept memory are different. And so we should have different things on the AI side. Exactly. And I think architectures that are scalable, but still kind of try to enforce something like this, I think are very promising.
31:26And so lastly, on the list of limitations you profiled is creativity. And that's where your big paper comes in. You know, talk a little bit about the motivation to explore creativity. So I guess like every AI researcher, I try to also use LLMs as much as I can to automate and make my life easy. And, you know, sometimes like may or may not have happened in real is that you prompt the model to generate homework problems. Right. Like especially because you want the model to generate something new or like, you know, I'm like, oh, can the model come up with something really clever, like a bit like to test students that I couldn't come up with or something.
32:04or new research ideas to, you know, write some new grant proposals or something. And I've almost always found that, or maybe always, like the models have never been able to give me something that's truly like, aha, I hadn't thought of that. So they're great at almost every other task that I try to use them for. Like they're great at summarization. They're great at like, you know, looking at what's common across a bunch of different things and drawing some stuff there. But they're not really good at these open-ended tasks that I give these models. And so that was sort of like the kind of motivation or like just, you know, there's always something that was lingering as like, yeah, I just don't feel like, or like, how do we think about that?
32:41And there's also work that people have been trying to do in the community about analyzing research ideas, like actually running human studies to see can models generate ideas or not. And there's been a lot of back and forth. So it's something that like is very fuzzy and like seems like things people are thinking about. And this was in collaboration with Vaisnav, who's a researcher at Google. And so we were having this conversation and he was also an author on this next limitation of next token prediction paper. And so I wanted to do some kind of a similar research of like, like, let's try to identify like what are the core principles that we might need for creativity?
33:16And can we test like, can that actually emerge from these training, like training objectives that we have from these models? And that's like a way that we can get an answer to this because having actual benchmarks and testing creativity and who knows what's in the training data, all of that makes it harder to answer this at school. scale. So we tried to take a different approach, be like, okay, let's just like start from first principles here and think about creativity, try to devise simple tasks and see in these concrete tasks, like what is the right objective? Like what does next token prediction do?
33:46How do we extract creative solutions? So that was sort of how that got in, how we started thinking about this. At the same time, I also had another project in my lab that looks at this from, you know, more, slightly more realistic settings of problem solving, like reasoning. And there again, we found that models, especially after you train them more, they started kind of collapsing in their solutions. So for people who are familiar with inference time settings, like you can either give the model one shot to answer or you can query it multiple times. And what we find is that when we try to train models, in general, their performance at this one shot setting does go up, but they also get worse at this, like, give it multiple times.
34:25So what it kind of means is that when models are trained too much, it seems like they start giving the same solution and try the the same incorrect thing rather than actually trying diverse solutions. So I was kind of convinced that this is actually a problem and that models are not very good at trying out creative diverse solutions, even on realistic tasks. And we kind of wanted to study this from a first principle. So that's sort of really what motivated us to think about creativity. And then once we had some really simple tasks to work with, that actually allowed us to make formal statements or give intuitive statements about what do different objectives do?
34:59like how do they perform? Like, can we do something better? And we found some really nice, interesting alternatives to the current paradigms that people have to improve creativity. And we are now trying to run these at scale and put some of these ideas and actually training larger scale models. And we're kind of excited to see where we can take in terms of making models more creative by changing the way we train them. So talk a little bit about what an objective for creativity means. Like that seems very difficult to capture in an objective. Yeah, so we drew inspiration from a lot of work in cognitive science, like particularly Bowdoin's work that tries to formalize some notions of creativity.
35:41So we are in no ways kind of really not thinking about like what's the right definition of creativity, but more like let's lean in on the definitions that cognitive scientists have taken about creativity and like use that. So as a very concrete example, there's this notion called a combinational creativity. And I think a nice example of that is like wordplay. So if we think of a joke, for example, let me, why did the scare, let me try to pull up exactly, I missed a punchline here. Yeah, I feel like I've seen this so many times that I'm like, it's not funny to me. So Yes. Yeah. So I guess the one example is why did the scarecrow win an award?
36:25And the punchline is because he was outstanding in his field. So if we can think about why this is creative or why this is funny is because scarecrow and award are two kind of seemingly different words, like unrelated. But there is this actual connection, like outstanding, that is somewhat novel or something that you hadn't thought about that like actually links these two words together. So that's an example of this combinational creativity where we're trying to see whether the model can find unexpected connections through a graph or find two words that have a common parent, two different interpretations of outstanding.
37:05And so can the model actually give new ways of doing this? And so the way we abstract this is, now let's say we have a graph and we teach the model all the edges in the graph. Can the model actually discover new siblings? So like two nodes that actually have a parent. So the two nodes don't look connected, but they have a parent. So can the model actually see some of these and then generate more of these? Like, can it find new connections in the graph? So that's one example of creativity that's abstracted really following, you know, Bowden's work on combinational creativity. And they have a lot of examples on like, you know, how a lot of things that we do that we think are creative are actually kind of finding these connections, unexpected connections between things.
37:45The other kind of creativity we consider is just exploratory, which is we just want to freeform find interesting things in the world. But, you know, of course, they still have some structure to them, which is, but it's just sort of a latent structure, but we just want to find new things from that structure. And like Borden explains, a lot of the work that artists do, for example, is this kind of exploratory creativity. And we try to write this down in like math and symbols that we can train models on as let's say we train on a bunch of like can the model like generate circles for example it's same in the following the graph kind of setting can the model find new circles or new triangles in the graph that the model has wasn't trained on so these are the ways we think about creativity like they're not perfect by any means but i think they capture like some of these core principles and they already show why certain training objectives might actually be good or bad for these notions of creativity?
38:38And how do we kind of, and we can, you know, dig more into that, like what's the right way to train these models or what could our current training be missing? And how should we be thinking about, you know, maybe new ways to pre-train these models to encourage this kind of creativity? Taking a step back and maybe getting a little bit philosophical, how do you distinguish the kind of creativity that you are thinking about broadly, meaning not necessarily, you know these particular constructs but you know more broadly the direction that you'd like to see LLMs go with you know I asked LLM to you know tell me a joke with a funny accent and like it can do these kind of creativity-ish kinds of things like how do you for someone who says oh I use JetGPT all the time and it's very creative like how do you distinguish?
39:28Yeah I think I feel like one notion of creativity is just, is it generating something that's not in the training data, right? And this was all it and something meaningful. And like, this was already for a long time, like not easy. And so it's like, it's very impressive that models can do this and can generate poems in the style of blah, blah, blah, right? Like we can do all of that stuff. I think the difference becomes what I feel is like open-ended versus like closed-ended. So in these cases, we're still kind of saying like roughly where we want to go. Like if I say write X in the style of why I kind of I'm telling the model what to do and then the model can fill out the path but what I'm thinking about like creativity is like take me to things that I've just never seen or thought of before and I think that's the part where I feel like uh like all our constructs as well are trying to say and even from our own use of these models like I feel like that's maybe where we want to think more um I yeah so I think for as another example like when I look into like these benchmarks or people trying to use models to for example kernel bench to come up with like solving problems in new ways, it feels like we still have to carefully kind of prompt the model on the strategies or the things it has already done and things it should explore and so on.
40:36So I think that part is still coming from our explicit instruction, even though it's generating new things. And so I kind of want to see, can they actually do this part on their own too? And I think that's what it would mean to meaningfully kind of go beyond what we are able to do and find new things. And if we are able to do that and we're able to do that at scale, then like we're actually going to really discover a lot of interesting new things and find new connections. And then how about the idea that, hey, I've already got a creativity dial, it's temperature and I can crank it all the way up.
41:06And wow, it's doing things that I totally didn't expect or didn't prompt it to do. Before we jump into that, I want to also say like when I was trying to prompt these models for this iClear, we were working on the talk and we tried to just get, can the model generate jokes that I can use, like the scarecrow kind of thing, like new jokes. And it was actually really bad. And I think we were joking that it's because models were trained with, you know, this kind of objectives that we are just showing in our paper are not very good. And so that was actually a case where I was trying to like get a model to like say something that was a new joke and like, but still have the structure and it wasn't able to, but maybe, I mean, and to the best extent that I could prompt engineer it.
41:48Meaning that, you know, maybe if you're thinking about it casually, like in just asking the model to tell you a joke, maybe you, you know, it'll come across something that you haven't heard before because there are a lot of jokes out there. But like, you know, if you really dig into it, you know, they're probably not novel and, you know, probably not funny. Yeah. And we couldn't get one that obeys the structure of like connecting. like for example the prompt would be like tell me a joke that's funny because it connects to unexpected entities right and like that was the kind of wordplay that we were trying to get at like so it's finding a new connection that like i hadn't thought about but like it was kind of bad at doing that so you could think of many different structure of jokes but like this particular aspect of like kind of this hidden punchline that the model discovered uh it didn't and i think there's a very deep reason like we explained in the paper is that a lot of the training kind of like the training data that doesn't really have supervision on sort of this, like it doesn't say first the punchline and then the joke that the model can do that.
42:50But instead, the punchline is sort of latent. And so that actually means that the models are not able to kind of learn the right structure, and they end up learning some things local that are like sort of memorization. And they don't learn this true process that we want, where think of something that's search over all possible things and find something new and then generate them. And so that's sort of the limitation that we try to show in the training of these models. Two things jump out at me. One is that I'm not sure that a human would do very well with that kind of instruction, like generate this joke on the fly, given, you know, this criteria funny.
43:23But also that, you know, there's something maybe orthogonal to the idea of, you know, structure and creativity or structure and like impulse. and like trying to convey both of those maybe is confusing to the LLM, like rules and like create. Yes, actually, that's a good point. So like a lot of the creativity is not just simply saying something crazy, but we want it to be still structured and meaningful. And if you think about like use cases like molecular biology or like, you know, drug discovery, it's not like we want to generate like random new things. We actually want the model to infer sort of the right latent structure and all of these and generate new things according to that structure.
44:02So that's really the kind of creativity that I'm thinking about and that we try to model in these tasks. And that's a great point that how do we get a model to try new things while still obeying some structure? And I think you were just asking me a question about the temperature. And I think that really leads into that where to crank up the temperature, you can start getting crazy things, but then the model is also going to be less structured. So we really want this sort of structured exploration from these models. And similarly, when I try to get a model to generate a homework problem, I don't want it to put together something like that's like meaningless.
44:37So that's, you know, creative. But I really want I want it to follow some certain logic or some structure there. That's an interesting idea that the most valuable examples of creativity are, you know, a balance of structure and impulse. Yes, absolutely. Yes, yes, yes. Yes. That paper is roll the dice and look before you leap going beyond the creative limits of next open prediction. Where does the roll the dice and look before you leap parts come into that? Okay. So first I'll talk about the leap because I think that ties in closer to what we just chatted about, which is sort of like this wordplay or like this connection, we are looking at tasks where there is a leap of thought that has to be made, which is not often spelt out in the training data, which is why models struggle to actually infer that thought.
45:23And so what we're saying is that models should actually be trained to kind of learn how to take those leaps or like take the structured leaps especially and rather than just showing the outcome of the leaps which is kind of what our current training data is because those just have they don't have the thought process that goes behind these. Which sounds a little bit like training on thought traces as opposed to training on answers like do you think there's a path there? Yes I think that that's certainly one way to do that but I think another part that we show in the paper is actually different training objectives like teacher-less training.
45:54So that could be things like multi-token prediction, where we are trying to discourage the model from getting the right answer by just looking locally. But if it has to actually generate the entire pattern or the entire set of things, then that actually encourages the model to have this global understanding. And people are excited about diffusion models lately. We find that diffusion models also have a similar thing where they don't show all the tokens and they mask out different things. And so that actually encourages the model to have some sort of planning or the ability to do more global things.
46:26And so we find these two objectives actually work a lot better in our experiments. One is multi-token prediction, where we just force a model to produce multiple tokens at a time rather than just one token, see the correct answer, and then the next token. And so in multi-token, the model does not get that local supervision, has to get everything right before it gets reward. and diffusion models can be thought of as having these different masks or different orderings. And so that also encourages the model to have more global understanding. So both of those are alternatives that might be worth pursuing more seriously if we care about getting these kind of diverse generations from the models.
47:06So that's the part about look before you leap. The part of the roll of the dice is also really interesting. I think one, like there's an underexplored aspect of like how to get diverse things from the model. Right now, there are two ways. Like one is we think about the diverse stuff ourselves. And like, you know, when people prompt the models, they're like, you give more specific instructions. And then you're like, actually, I want to do something else. And then we tell the model, go do this instead. So that's one way. The other way that, you know, machine learning people may be thinking about is like, just increase the temperature.
47:38And that like gives you more, you know, tokens. And like I said earlier, that also tends to really destroy the structure. And so we try to think about like, okay, how do we think about creativity? Like what's the right way to kind of elicit this randomness from the model that can be structured? And one way that seemed more natural than temperature sampling is first generate like a random idea and then follow that idea and generate all your tokens. So the right place to introduce randomness is probably at the beginning where the model commits to exploring something. and then once it decides to go down that path, if you keep increasing the temperature, then it's just going to go all over the place.
48:15Instead, we pick something and then we stick to it. And so at test time, if you want more diverse generations, we would sample new prefixes or new starting points and then let the model kind of do its thing. So that's where we're trying to say we should introduce randomness, which is like the roll the die. So we roll the die first and then once we are based on the outcome of the die, we kind of pick that thought and go with it. And so this, of course, it requires training the model to be able to do that. And we have a very simple way to do this in the paper, which is we just in the training data, we just have random prefixes that we prepend.
48:49So the model conditions on a random prefix that you can think of that as like some proxy for a random idea. And then it generates things. And at test time, we can get new ideas or new generations from the model by changing this prefix. Sorry? Random nonsensical prefix or a random sensical prefix or? Yes, great question. So in our paper, because our experiments are so simple, a random nonsensical prefix actually just works, which is surprising. We don't really understand why it works. That's very interesting future work too. But as a proof of concept, it seems like, okay, like there's maybe some meat here.
49:21And what we are looking at now in my group is having more meaningful prefixes, which actually capture like what exactly is the idea, like some semantics of the idea that could, and then condition on that to generate things. And if this paradigm of pre-training ends up working, this would also give us more controlled, diverse generations. Instead of prompting the model with explicit instructions, we can actually just change this prefix that we condition on in some nice way. And then we can get diversity in that sense as well. And so that's what we are trying to say that maybe we should think about is rather than doing temperature sampling, maybe we can tell the model to condition on new ideas that can then be randomized or can be made more diverse.
50:04And that way of training as well as inference might actually be a better way to get structured diversity from these models. And we mentioned GPT-5 and also this idea of inference times compute, inference time scaling. And one of the things that I think we're learning about this model is that at least at the higher tiers, like it does a bunch of parallel inferences and kind of integrates those together. And And in some sense, you might think that, okay, assuming this stuff in parallel, like you're going to get this diversity of thought. But also there seems to be this averaging effect that kind of like dries out the expression of that thought.
50:50And it strikes me that it's connected to kind of your work and maybe they need more random seeds in their parallel threads or something. Yeah, it also depends on how these like parallel thoughts are actually coming. I mean, we don't like, I guess, I don't know how exactly those thoughts are being generated, you know, behind the scenes there. So two comments. So one is like, if they're doing RL on sort of a starting point, the question is like, can we, how do we leverage the starting point appropriately to get the right diversity? Like, can you actually go beyond what the base model can generate?
51:21Like, are some things more likely than the others? In which case, like, can we correct that in some way? Like, what if we want to generate more of one kind than the other? So all of these kinds of things are still going to be quite challenging with like kind of the current paradigm. So hopefully if we do pre-training in a way that allows this kind of, you know, diverse control sampling, then there's more that we can get. But yeah, I think the thing that like is like basically unclear is like what, like how do we get that diversity and like how do we actually span things? Like if we just let it all be end to end from the model, then again, there might be some collapse and like the model might not be generating all the diverse things that we need.
51:59And yeah, and there's some of that that you can simulate by carefully collecting training data. But then, of course, like if you want to do stuff that's like very different from your training data, like how do all of these things work? Like, yeah, that's hard to say again with the parallel traces kind of thing. But I think morally it is sort of what we are saying is that we want the model to, you know, try a new idea and then stick to it rather than like, you know, introducing randomness at every step. And the random prefixes, the random prefixes in your case are at inference time. So it's like you're prepending that you're prepending this random prefix to the prompt.
52:38Yeah. So in our toy settings, there's no like notion of a prompt. I guess there's no instruction, but yeah, you can think of that as like prepending. Yeah, you're prepending. So you're giving your question and then you're adding the prompt and then letting, sorry, adding the random string and then letting the model generate. That's how it would work in practice. Yes, but you also have to train a model to be able to do this. So at training time too, we have to simulate this where the model has like random strings in the beginning that it uses in its, you know, while training. Meaning, so don't expect that you can just add some randomness to the beginning of a prompt with chat GPT.
53:12It's going to give you better answers. Yes, yes, yes, yes. And so you kind of talked about RL. Do you think that the approach is compatible with RL training? Yeah, this is fascinating. I think overall, any improvements that we can do to the base model in terms of how we sample, that directly translates to RL because a lot of the RL is like you try a bunch of things from the model and like you up like you know you make the model do more of what it's doing well and so if you just have better starting points like then then you just like will automatically improve all of these things and the way i think a lot of this could be worked useful is in ideas like structured exploration and so on and you know people keep bringing up i guess like exploration is kind of one of the biggest challenges to um you know to the next frontier and so if we want to uh do some kind of structured exploration over many different spaces like the more control you have over the kind of diversity, like the better, you know, like, yeah.
54:09So all of that, I feel like that's very compatible and would like actually be really useful for further ADL training. So I've referred back to this blog post a few times. And so there was these limitations and opportunities and we've actually covered the opportunities. It's this idea of memorization syncs, seed conditioning, which is like prepending the random text and multi-token learning. um yeah you've hinted at some future directions for your research but talk a little bit more broadly about like the direction that you see these various efforts going yeah i think a lot of like we really want to scale up a lot of these things and to see and to both you know for ourselves but also convince people that these are the interventions that are worth investing for training their you know future models uh and especially this adaptability part like i you I often think like, how do people keep their models fresh?
55:03Like right now, we are at a point where everyone is training new models every year. And so like, maybe we haven't thought about this issue, but at some point we're like, kind of not gonna be able to do that. And I think the straw man answer here is like just do retrieval, like that gets you the latest facts and everything's okay. And we had this other paper that was an oral at the previous conference where we kind of show how that models are not gonna be that good at using the context and overriding its parametric information. So this is not really a magic solution to keeping the models sort of updated.
55:33So I think we're trying to like, I feel like the aspect that's like kind of really interesting to think about is like, how do we actually make models that are easy to update? Like, what's the right way to decompose or disentangle these aspects of what should be preserved, like what should be updated? So I think I'm really excited about kind of that aspect. And I'm also, you know, really excited about scaling up these ideas for improving diversity, creative generations. and we're all like, we all have like, you know, we're looking at things like theory improving and, you know, trying to build systems that like help mathematicians.
56:04And I feel like every time I talk to all of these people who want to use LLMs, I keep coming back, like how do we make sure models can actually search and find creative things? Like for example, finding counter examples, like we can use LLMs, they would be great, but how do we make sure they find kind of that kind of structured exploration, like finding creative counter examples in some way. So I would be very excited about using these ideas in like all of these different applications that I think are top of mind for people right now in terms of the interesting applications that can really push the frontiers here across various notions of science.
56:38So those are kind of the things that we're thinking about. And of course, also all of this, I think we spoke a little bit about the fact of trying to understand these models. And I think a lot of our experiments are also trying to less be like, here is a state of the art method that gets high numbers on this, but really hopefully give some objectives that stand the test of time in terms of like showing you new understanding about the fundamental properties that go beyond specific data sets or specific, you know, training decisions that you might make. And that would help in guiding kind of the next generation that we try to train.
57:13So we are continuing our efforts along that axis as well. Well, Ajitji, thank you so much for sharing a bit about your research. It's very interesting stuff. Yeah, thank you so much. And I kind of wanted to especially sort of plug a little bit that I want more people thinking about trying to understand these models because we really need a lot of work there. And I think, you know, it's like a scientific question where these are complex systems. And we really have to think a bit like a scientist setting up the right controlled experiments, make formal hypotheses, test them out. And there's like a huge sort of opportunity here.
57:47And I think that will really unlock like how we further push models, especially as people are starting to see maybe diminishing gains from, you know, just pushing our current paradigm. And so like kind of I think there's a lot of opportunity there. And so hopefully we can have a bigger community looking at these aspects. Awesome. Very good. Very good. Well, thank you very much. Thank you.
58:17Thank you.
From the publisher
Today, we're joined by Aditi Raghunathan, assistant professor at Carnegie Mellon University, to discuss the limitations of LLMs and how we can build more adaptable and creative models. We dig into her ICML 2025 Outstanding Paper Award winner, “Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction,” which examines why LLMs struggle with generating truly novel ideas. We dig into the "Roll the dice" approach, which encourages structured exploration by injecting randomness at the start of generation, and the "Look before you leap" concept, which trains models to take "leaps of thought" using alternative objectives to create more diverse and structured outputs. We also discuss Aditi’s papers exploring the counterintuitive phenomenon of "catastrophic overtraining," where training models on more data improves benchmark performance but degrades their ability to be fine-tuned for new tasks, and dig into her lab's work on creating more controllable and reliable models, including the concept of "memorization sinks," an architectural approach to isolate and enable the targeted unlearning of specific information.
The complete show notes for this episode can be found at https://twimlai.com/go/747.




