Training Data Locality and Chain-of-Thought Reasoning in LLMs with Ben Prystawski - #673

0:03All right, everyone. Welcome to another episode of the TwiML AI podcast. I am your host Sam Charrington. Today, I'm joined by Ben Prostovsky. Ben is a PhD student at Stanford University. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Ben, welcome to the podcast. Thank you so much for having me. I'm excited to dig into our conversation. We're going to be talking about your work on large language models and reasoning, and in particular, the Why Think Step-by-Step paper, Reasoning Emerges from the Locality of Experience, which you'll be presenting at NURPS 2023.

0:43Before we do that, I'd love to have you share a little bit about your background and how you came to work in AI and cognitive science. Yeah. So I guess my interest in cognitive science started out with an interest in language I've had since I was a teenager. So when I was like 14, I started learning languages as a hobby and got really into that. But then eventually I realized that what I actually liked in learning a language was reading the grammar books. I was one of like a very small minority of people in the language learning world who'd like really liked the grammar. And I got more interested in sort of language itself and the structure of how language works.

1:17At the same time, I was, I really liked computers and I liked to program computers. And so by the time I got to university and had to sort of decide what to study, I was like, hmm, can we get computers to understand language? That seems like a reasonable way to combine these interests. And so I got interested in computational linguistics and natural language processing. And at some point during my degree, I took a, well, my second year, I took a cognitive science class. And when I took that intro to cognitive science class, I realized that a lot of what interested me about language was really insofar as language told us about cognition and how the mind works more generally.

1:54So that really led me to get interested more in cognitive psychology and computational cognitive science. So I started reading about Bayesian models of cognition and doing some research at a few labs at the University of Toronto, where I did my undergrad. And at the same time, it was like the late 2010s or so. And so the neural networks boom was happening and was especially big, I think, at the University of Toronto. And so that's kind of the background that led me to what I'm doing now in grad school, which is trying to understand the mind, trying to understand machine learning models, and trying to understand the relationships between both of those.

2:33How do you craft a research program or an agenda around trying to understand the mind, trying to understand machine learning? What are your specific areas of focus and interest recently? There are two areas that I'm really interested in. The first is, why is it useful for us to work through some reflective reasoning process, right? It's easy, I think, to take for granted that we can reason about problems and it can lead us to answer questions better. But if you actually think about it, we don't get any new data from the world when we reason. And yet mathematical proofs can be surprising, right?

3:10We can often learn things or form strong beliefs about the world just by working through mental processes. Why should we even expect that to be useful in the first place, I think is one question that I'm really interested in in the reasoning space. At the same time, another bit of my work focuses on cultural learning. So how is it that humans can build up big bodies of knowledge over many generations? There's this idea that what's really distinctive about humans is not so much our individual intelligence, that we do have a fair bit of that compared to other animals, but it's our ability to communicate with each other and build bodies of knowledge and share knowledge with others, really to build culture.

3:50It's the idea that Yuval Harari popularized in Sapiens. Yeah. There's also Michael Tomasello is one of the first people to put forward this idea. He has this book, The Cultural Origins of Human Cognition from the 90s. I don't think he's the first person to come up with this hypothesis, but that's one sort of early popular book in this area. Yeah. I'd love to have you riff on reasoning and how you think about it in the context of machine learning and large language models. This is one of the things that I've been really interested in exploring. I've come to articulate it as LLMs can display reasoning, but they don't necessarily mechanically reason.

4:30Some people push back on that. Others agree with it. I'd love to have you react to that before, broadly talk to the way you think about reasoning in this context? Yeah, I think about reasoning, I think in a fairly general way in that I more or less run with the definition of reasoning as intermediate computation. So when you have a question, you can produce an answer directly sometimes, right? You can ask me, Ben, what did you have for breakfast this morning? I can say, I ate an apple. That's straightforward. I don't really need to walk through some complicated mental process to come up with an answer to that.

5:07But if you ask me, say, the answer to a math question, I have to do a lot of work to go from the question to the answer, or at least if I want to have any hope of getting the answer right. I could guess a number, but that's not likely to be very helpful. And so I walk through some internal mental process before giving an answer. And there's a similar pattern in large language models in that we can ask them to map directly from questions to answers, right? We can give them a prompt that is a math question. We might give them an example of a question and an answer, and then another question and get them to produce another answer.

5:43And they can do that. They will give answers. But we could also either give them examples of reasoning toward an answer beforehand, or we could even just tell them something simple like, let's think step by step, or you should reason toward answering this question. And that often induces them to produce some intermediate output, generate some tokens before giving an answer to the question you gave. And as a number of papers have shown, in particular one by Jason Wei in 2022, another one by Maxwell Nye, showed that often doing intermediate computation is helpful to producing better answers, not just different answers.

6:20So the argument there then is that the intermediate output or any output before the ultimate output is the intermediate computation and therefore qualifies as the reasoning by your definition. That's right. Yeah. As opposed to like an offline private reasoning or out-of-band reasoning before producing the answer. Is the distinction interesting or meaningful? I think it is possibly meaningful in that a language model, when it does reasoning, it's trained only on the reasoning that we put in language, right? Humans reason both mentally without saying anything, and sometimes we think out loud. But it's only the thinking out loud that makes it into the training set.

7:10So while language models are trained on human reasoning data, they're not trained on all of the human reasoning that exists. So I think there is an important distinction there between the reasoning we verbalize and the reasoning we don't verbalize. But I do think it's worth calling both of those reasoning. And so saying that generating a chain of thought before answering a question is a kind of reasoning. Got it. Got it. Awesome. I love that. Tell us a little bit about the framing for the Why Think Step-by-Step paper. Yeah. So in this paper, we're interested in essentially, what is it about the data that large language models are trained on that leads chain of thought reasoning to be useful?

7:52So what is it about the experience we need to have in order for it to be useful to do some intermediate computation when answering questions? And so we give models this probabilistic inference task where essentially they're trained on samples from a base net that we generate randomly. And given some trained model, we can estimate conditional probabilities, so probability of B given A, using one of three different estimators. One of them we call direct prediction. So we can prompt the model with the name and the value of one variable, and then the name of another variable, the one we want to infer.

8:27And then we can just get the probability of one or zero going in that next token. And so we can just get the probability the model assigns to one or zero and normalize. And that's our probability estimate. Another way we can prompt the model is through what we call free generation. So we can give it the name and the value of a variable we want to condition on. And then we can essentially run the model forward, generating names and values of whatever variables we sample from it until eventually it generates the name of that target variable, the variable we want to infer the value of. And then after it's done that, we can do that process again where we look at the probability of one and zero coming next.

9:07And then we can get an estimate that way. In practice, when we run this estimator, we resample those intermediate variables 10 times. And that helps reduce the variance in estimation. If you're introducing a lot of intermediate variables and you're sampling values from the language model, the variance of that estimator is a lot higher. And so we do have to average over a number of different samples of intermediate variables to get a reliable estimate. This is similar to a prompting method in the reasoning literature called self-consistency, where you're resampling a chain of thought multiple times and looking at what answer is most common across different chains of thought.

9:43And so we can think of that as a variance reduction technique, essentially. The third estimator we might consider, which represents like ideal reasoning, if you knew sort of the best set of steps to work through, we call scaffolded generation. So in this, we're going to compute a scaffold, which is the smallest set of variables that collectively de-separate the variable we've observed from the variable we want to infer. And the idea being if we know that the values of that de-separating set, then all of the influence that that observed variable has on the target variable had to go through there.

10:14And so if we know all those values, then the observed variable is now independent of the target variable. Is the idea with that estimator to identify an optimal, minimal number of steps? Yeah, we're essentially looking at a positive control condition. We want to say, what's the best you could do if you knew exactly the best set of steps to instantiate? We can then ask how free generation compares to scaffolded generation. So if we just run the model forward, how close is it to generating the optimal set of reasoning steps kind of automatically? In the training data here, we decided to hold out certain pairs of variables that have high mutual information but aren't adjacent to each other.

10:54So we want to have a set of pairs of variables where we've learned their relationships with their surrounding variables, but we haven't learned their relationships directly. And so they weren't in the training set, and we have to sort of generalize out of distribution using the learned conditional probabilities in the rest of the training set. Got it. What are some concrete examples of these pairs of variables that you've used? So in this setting, we have only variables in randomly generated base nets that have names like x51, and they have values. If you want an example of how this might reflect a real-world scenario, you could imagine training a language model on data from Wikipedia.

11:31And then if you want to ask that model, what is the climate of France's capital? It's likely not going to be great at giving the right one word answer. Because if you look at the Wikipedia page for France, and most of the Wikipedia pages where France is mentioned, the capital has an oceanic climate never really shows up. But we do know that the capital of France is Paris. And we do know that in the Wikipedia article for Paris, it mentions that Paris has an oceanic climate. And so by thinking through this intermediate variable, by saying the capital of France is Paris, now we can connect what we've learned about the relationship between France and Paris with what we've learned about the relationship between Paris and its climate.

12:11You mentioned these estimators that you present in the paper. What is the overall goal that you're driving towards with the paper? Our goal is to develop an account of how the training data in language models drives the effectiveness of chain of thought reasoning. So what are the properties in training data that make chain of thought reasoning effective? To answer that question, we train the models on different kinds of data. So we have a few different training conditions. The most important, I think, are a locally structured training condition, where in every sample that we train the model on, written out as a string, we don't show all the variables at once.

12:48We only show a local neighborhood of variables that consists of one random central variable, and then all of the variables within some distance k, where we sample k from like an exponential distribution of that central variable. And so there we only really see some central thing and then everything that's closely connected to it. We compare this against what we call a fully observed training condition, where a model is trained on all of the variables in the base net, except for like half of those variables in the held out pairs, so they don't see them together. And we can compare A, is reasoning helpful in both of these training conditions.

13:25And we can also ask whether free generation differs from scaffolded generation. And we find that in the locally structured training condition, free and scaffolded generation perform a lot better than direct prediction in terms of the mean squared error in estimating those conditional probabilities for held out pairs. But in the fully observed condition, there's a little bit of a benefit to scaffolded generation, but free generation really doesn't help very much. One idea as to why might be that it tends to generate a lot more intermediate variables because locality is helpful in that if you're trained on local data, you tend to instantiate reasoning intermediate steps that are kind of close to the variable you're conditioning on, because that's what you've seen in training.

14:05And so just kind of throwing whatever happened to co-occur with the thing you're conditioning on does pretty well. It does essentially as well as scaffolded generation. And it's probably unique to our simplified setting where we have like 100 variable Bayes Nets. But it's interesting that you get a little bit of reasoning ability for free just by observing locally clustered things. And so is the training data in this case, is it generated by the Bayes Net? That's right. So we'll draw a million samples from a Bayes Net, and then we'll format it as a string, just mentioning the name of the variable, an equal sign, and then the value.

14:42And that's a sample. And we just concatenate a million of those together and that's our training corpus. You're training the LLM from scratch on this data as opposed to like fine-tuned so it doesn't know anything other than you know this type of data that you've exposed it to. That's right we train them from scratch we use a smaller version of the GPT-2 architecture so this is you know a lot smaller than your GPT-3s. How do you know that these results generalize to LLM is trained on kind of natural language and the types of architectures that have evolved. I hear that you're using something similar to GP2, but the data sets used to train LLMs are typically very different.

15:23There are a few things that we think will generalize and then a few that we think are open questions as to how closely they map onto the natural language setting. So a document tends to be about a few topics at most. It tends not to be about the entire world at the same time. And so the fact that documents have topic structure that they're not about everything. They're about a small subset of the world and that there are some subsets of the world that you talk about more than others. If it's only trained on adjacent pairs of variables, but then you want to generalize out of distribution to infer things about more distant pairs, doing scaffolded generation where you generate the intermediate variables first is always going to have lower bias than doing direct prediction.

16:05And so we should expect some similar version of this finding in simpler settings to generalize to different architectures, to different kinds of sequence models. That said, understanding that we think these aspects of our findings will generalize, natural language is really different than 100-valued BayesNet, right? For a lot of different reasons, right? There are a lot more things in the world. The fact that you can just throw random intermediate variables into your context window, and typically it does well if you've trained on local data, isn't something I think you can do in natural language.

16:42I don't think you can just kind of free associate and then usually find your way to the right answer. There's probably a lot more planning is a bit of a loaded term, but there's a lot more, I think, deliberate choosing of useful steps in both human reasoning and the reasoning we often see in larger language models trained on natural language. There's also the fact that in natural language, you can talk about things that you haven't directly observed. We have a lot of abstractions in language, like society, like the day of the week. And the fact that we can think about these concepts probably lets us bridge otherwise disparate topics a lot more.

17:23And I think we'll probably see something similar in language models where we can make analogies or we can reason about unobserved variables, which doesn't show up in our setting. You mentioned earlier that one of your key interests This is learning about the human reasoning and how human reasoning works. What did this paper illuminate about human reasoning for you? It illuminates that the way that we experience the world probably drives the usefulness of reasoning. So just as there's this local structure in documents, I think there's a local structure in human experience in that we see the world through a first-person perspective, right?

18:02We can only ever see a little tiny subset of all of the things that we might possibly think about at any given time. And so because there's this local structure in human experience, it could be that our reasoning processes let us bridge different things that we've seen locally together. I think that gives a little more insight into that question I mentioned in the beginning of like, why should we expect reasoning to be useful at all for humans? One possible answer is just like, we learn some relationships in the world really well because we see them together a lot. But there are other relationships that are a little bit more distant in our direct experience.

18:39And so we need to chain together reasoning steps in order to make good inferences in those domains. That said, my advisor Noah Goodman and I are also interested in running some human experiments, trying to test hypotheses generated here. So you could ask, if we can experimentally control the data that humans see, can we expect reasoning to work more in some settings than others? There's a paradigm used a lot in psychology where you can give people a task and in one condition, you can give them like five seconds to answer. And then in another, you let them take as long as they want to think. And we could maybe think of these as analogs of direct prediction and free generation, right?

19:23In one case, you don't have time to work through any mental steps and another you can. And so we can ask, what tasks do people perform differently on when you give them time to think versus not? And could that have to do with the statistical structure of the information that they're seeing? Are there lessons here for folks that are using LLMs and how they construct prompts? I'm not so sure about constructing prompts. Our setting only studies what people call zero-shot chain of thought. So asking a model a question and telling it lets things step by step and letting it generate sort of whatever intermediate reasoning traces there are.

20:04You could also ask questions about few-shot chain of thought, where you give it a few examples of questions and then reasoning steps and then answers. We could ask in a similar setting to what we have now, how do the in-context examples of reasoning traces influence the reasoning that the model produces for that last question you give it? We haven't done that in this setting, but it would be an interesting thing to do. There's a, yeah, it reminds me of this DeepMind paper that actually was a big inspiration in designing this work. Data distributional properties drive in-context learning in Transformers by Stephanie Chan and others.

20:39You could probably do something kind of similar for in-context learning reasoning. We haven't. We do think, though, there might be some insights in data curation for language models. One thing we find is that the language models learn a lot more slowly, both to do direct prediction and to do free generation, when you give them the fully observed data, because there are a lot of pairs of variables that just don't really influence each other. And giving only the locally structured subsets of samples lets a model learn a lot faster. It needs to see way fewer tokens before it can get pretty good at answering questions.

21:12So it might be that curating data sets that might be smaller, but that consist of good representations of the relationships between concepts could lead to more data efficient training of models. In the way language models are typically trained, do documents preserve the document structure of the internet? Does that preserve some of that locality or do chunking and randomization and those kinds of things break that? And what you're suggesting is that maybe incorporating or allowing more of that to pass through might help the language model? My understanding is that if you're training a model on a lot of text documents, you present the documents either sort of in the order the web crawler found them or in a random order.

22:00You don't present a Wikipedia article on one concept and then the Wikipedia articles that it links to. Locality is broken at the document to next document boundary. As humans, we never see a clear break, right? We see and hear and feel some inputs, some stimuli, and then we move on in this sort of continuous experience to see related things over and over again. So there's a stronger version of locality that governs human experience than that governs the training of language models. It could be that training models with documents that are close together, closely related, maybe web articles that link to each other or that we've curated to be on similar topics, might let us learn those local dependencies even more strongly.

22:47I think the idea that LLMs trained on text could process training data sets and come up with denser, richer, downstream training sets that are somehow tuned for reasoning is kind of one potentially interesting idea that comes out of this. It kind of suggests that that might be a thing to explore. Yeah, so that is largely what my lab mate Eric did in the STAR paper, where their reasoning is in natural language and you're fine-tuning on your own outputs of successful reasoning. The question of how do you learn qualitatively new ways of reasoning in natural language seems tricky, but I've seen examples of training on your own outputs or also the idea of curating data of really good reasoning traces.

23:39Can we take lots of worked explanations for math problems and train a model on that and produce a model that's better at working step-by-step through math problems? There are a lot of, I think, practical versions of this idea of training on successful reasoning that you can do. There's also the question, I guess, which is, why does that help? There are still questions about developing a fundamental understanding of what's going on there that I think is often easier to study in these simplified settings. but there's still work to be done in understanding how those insights translate to the full setting of natural language.

24:13Well, Ben, thanks so much for taking some time to share with us what you're working on and kind of talk through some of your thoughts on LLM reasoning. Thank you so much for having me. It's been great.