Teaching LLMs to Self-Reflect with Reinforcement Learning with Maohao Shen - #726

0:00Here's the thing. This isn't how humans really naturally solve the hard problems. As humans, we are capable of extended thinking. And more importantly, we can self-reflect. If we realize we have made a mistake, we can backtrack, rethink, and correct the mistake. And sometimes we also propose better solutions. And the motivation is like how we can train the language model to do the similar things.

0:41All right, everyone, welcome to another episode of the TwiML AI podcast. I am your host, Sam Charrington. Today, I'm joined by Maohao Shen. Maohao is a PhD student at MIT. Before we get going, be sure to hit that subscribe button wherever you're listening to today's show. Maohao, welcome to the podcast. Yeah, thanks for having me for this podcast. We are going to be talking about your paper, Satori, Reinforcement Learning with Chain of Action Thought Enhances LLM Reasoning via Autoregressive Search. That is quite a mouthful, but looking forward to digging into it. Maybe we can start by having you share a little bit about your background and research interests.

1:23So I have a broad interest in various AI and machine learning problems. In general, my goal is to make AI system more intelligent and also more reliable. And a large part of my research focuses on quantifying the uncertainty of AI models. Because uncertainty can serve as a signal for users to decide whether they should trust an AI model's predictions. And this is especially important in those high-stake applications such as autonomous driving, health care, where those unreliable or misleading predictions could have serious consequences. And in recent years, I have also been working on multiple language model problems.

2:21And one of the most exciting challenges is moving beyond the chatbot and how to develop the AI system that can think more like humans and tackle those complex problems. And this is my ultimate goal. And Satari is my latest exploration along this direction. And I believe that making AI models with both smarter and more reliable is the key to achieve AGI in the future. What are some of the things you've worked on in the past? So in the past, I worked on several problems relevant to trustworthy machine learning. How to, as I said, how to enhance the uncertainty quantification performance of a machine learning system.

3:18Usually using some auxiliary models building on top of the neural network to predict the uncertainty or using some other algorithms such as Bayesian algorithms to estimate the uncertainty. And this project, Satori, is more like a new direction to be explored. And to talk a little bit about how this particular project fits into some of the current trends in AI research. which we're seeing a lot of development around reasoning models, around test time compute. Recently, DeepSeq R1 explored the use of reinforcement learning to aid in model training. There's so much going on. How does your project fit into that milieu?

4:07So our work is definitely in line with the broader research on improving reasoning in language model. but I would say it takes a different approach from the traditional methods. One of the most well-known techniques is chain-of-south reasoning, where you prompt the model to break down a problem step-by-step. And this kind of structural reasoning is really helpful for handling those complex problems. And previously, people tried improving language model by training the model on a lot of train of salt data. And those data are coming from either human annotation or distillation of the teacher model, for example.

4:57But as you can imagine, that's really expensive in terms of both data collection and training compute. And more recently, instead of just scaling up the training compute, people also started to think about scaling the test time compute. So basically, test time compute, scaling test time compute means they don't require training the model itself. They just push, they try to push the limit of an existing model by letting it generate more response. And a super simple example is majority voting. You ask the model the same question multiple times, get different answers, different response, and pick the one that is most frequent, basically.

5:54And it turns out this is quite effective. It's more accurate than just taking a single answer. And more advanced test and search methods, they usually rely on a reward model to guide the selection process. So instead of just counting answers, you can sample multiple solutions, having a reward model to score the different solutions and pick the optimal solution. And the intuition behind this kind of approach is there's even one correct answer among the samples. Then the reward model can help to find it. So some other approach, even more advanced, they try to score each reasoning step along the way and doing a tree search to find the best reasoning path.

7:06But I mean, what's the downside of this kind of approach? So these methods, they require running multiple models at once. They need to deploy two models at least, which gets expensive. Meaning the model that's scoring as well as the model that's... Right, the reward model and the generator. So two models. This is a two-player system. And they also throw away a lot of useful information you can tell. For example, if you sample 100 times but only keep the best result, you're essentially discarding everything else. This is quite inefficient. And that's why reinforced learning comes in. And this is a, I would say a quite different approach.

7:58So instead of just relying on huge amount of labeled supervised fine tuning data or complex test time search methods, reinforced learning allows the model to learn from its own generated response. And this is also the magic of reinforced learning. Once you define a reward signal, the model can essentially teach itself to get better by using its own knowledge. And this kind of knowledge is obtained from the pre-training stage of Lanxiu model. And we already know this is quite successful for some models like OpenAI's O1 model and DeepSeq R1. And this model already proved that reinforced learning could be used to train language model in an efficient way without requiring too much human supervision.

9:03So, right, so that's pretty much the entire picture about the reasoning model. And so with Satori, we are also inspired by this kind of new direction using reinforced learning. Especially since the OpenAI released a one model, that model showed there was more like there was a new path to be explored. And at the same time, we also realized reinforced learning for reasoning, language model reasoning is quite underexplored in academia. And back to September, we put together a team and started to tackle this tricky problems. And as you know, after a lot of trial and error, over the three months, we finally got Satori working.

10:03Yeah, that's a story about Satori. Right. And you also asked about DeepSeek R1. And I would say that's more like a concurrent work. It's an industry project. And they have done a fantastic job, I would say. And when their report came out, I guess it's late January. And we were actually really surprised to see how many of the high-level ideas overlap with Ours. and at the same time, it was also super exciting because it's like it confirmed that we were on the right track in the working toward more powerful reasoning models. So it's kind of a mixture of two feelings. It's surprising and it's also exciting.

11:02Can you dig into a little bit more the idea of applying RL to evaluating these chain of thought traces? Like how do you set that up from a research perspective? So when I was thinking about this research problem, I always started by thinking about how humans would solve a tricky task step by step. And then I tried to figure out how we can adapt this kind of human thinking process to AI systems. So as I mentioned earlier, a lot of researchers focus on the test and search methods, but I will say that may not be how human really solved the problem in practice. So let me give an analogy. So imagine a student working on a tricky, tricky mass province with a teacher standing nearby.

12:12For example, I'm the student, you are the teacher, and the teacher, the student might try different approaches, and you will pick the most promising solution. or maybe the students write out several reasoning steps and the teacher selects the best one and tell the students to build on that reasoning step and so on and so forth and this kind of method is definitely better than the student just working along with no feedback but here's the thing this isn't how human really naturally solve the hard problems. We don't always need any external guidance. As humans, we are capable of extended thinking, and more importantly, we can self-reflect.

13:09If we realize we have made a mistake, we can backtrack, rethink, and correct the mistake. And sometimes we also propose better solutions. And the motivation is like how we can train the language model to do the similar things, how we can make the language model solving a problem just like the human did. I also mentioned that I did a lot of work about unsending quantification. so actually in the language model community there is also a big challenge it's called hallucination which means the model will generate some response that sounds really competent but actually they are full of factual errors and for example take mass problem as an example so a typical language model might generate a step-by-step solution, output a bunch of equations, derivations.

14:23And even though there are a lot of mistakes, the language model cannot backtrack and catch the errors, even though those errors are actually quite obvious. So, right. So I've been working on a lot of approach to address this. One approach is certainly using unsenting quantification. As I mentioned, we can add some additional modules or algorithms to estimate how confident the model prediction. And if the unsenting is high, we know the answer is unreliable. and if uncertainty is low, then we trust the model predictions. But even though this kind of system is useful, it doesn't resolve the fundamental issue, right?

15:23It just tells us whether we should trust the model or not.

15:32So, but later it's more like OpenAI's AI's O1 model came out in September, and it was a game changer. So they already proved that the reflection and self-correction is not something theoretical. It can actually work in practice. So it's more like I want to know how to really uncover the technical details behind it. And we know the secret sauce of our model is using reinforced learning. Different from the classical training pipeline that teach the model how to reason using a bunch of supervised fine-tuning data, reinforced learning is quite different. It just assigns some reward to the model and led the model to interact with the environment.

16:33And secretly, the model can learn some reasoning patterns and the general abilities for reasoning. Right. So I think that's pretty much the main objective of our paper. We just want to uncover the secret sauce of how to use reinforced learning to really improve the reasoning capabilities of a language model. In the title of the paper, you refer to reasoning via autoregressive search. Expand on the idea of search in this context. What are you getting at there? So we know most language models use transformer architecture, right? most of language model is building on top of transformer. And transformer are autoregressive models.

17:27It sequentially generate a bunch of tokens based on previous tokens. And the search just means a long reasoning process with trials and errors. So combining these two words together, it's like autoregressive search means a single language model performs a human-like reasoning process with self-reflection and self-exploration of new strategies. When it encounters a potential mistake, it can self-identify the arrows and point out the mistake. It can also start a new solution if it realizes the current solution is actually suboptimal. and different from the classical test time search approach that relies on external guidance.

18:24I think the main objective of autoregressive search is teach a single language model to perform a test time search without any external guidance. And is the idea with the test time search that is it different than training the model to change the way it's generating the new token? Like how much evaluation is involved and reflection is involved? And is that all captured in the next token generation process or are there other aspects of the model? Like, you know, we've established that there's not external models, but are there internal modules or other things besides from the token generation pipeline that are coming into play?

19:18Because we know the classical approach, they typically rely on external guidance. There's a reward model to guide this process, like when the model wants to trigger self-reflection or the current solution is suboptimal and the reward model will come in and provide some feedback and select the optimal reasoning stat, for example. But if we want to achieve the similar thing, similar behavior using autoregressive predictions by generating a token in an autoregressive manner, there's a challenge, right? So it's like the traditional chain of thought approach doesn't really achieve this because it just produces a linear sequence of reasoning steps without distinguishing the types of reasoning actions.

20:19So we can treat self-reflection or proposing new strategy as different actions. The COT is more like a robot that can only move straight without the ability to change directions or or stop or rethink the approach. And, but we think about, but let's think about how a human solved the problem. You might go through a few reason steps, pause and start self-reflection. Maybe backtracking, verifying correctness, correcting mistake, or even coming up with new strategies to tackle this problem, right? So it's more like a robot navigating an environment. If it hits a war, it stops and it decides the next movement and explores a different path, different trajectory.

21:18So I think the core idea of our paper is this chain of action salt in this long title. So it explicitly defines different actions the model can take. and the unique technique we introduce is directly leveraging the special token to guide the language model's generation. Because we know autoregressing model can only produce a sequence of tokens, nothing else. But we can still use special tokens as some signal to tell the model how to take the next action. So we introduced a bunch of special tokens and each special token corresponds to one specific actions. And we basically define three types of actions, continue, reflect and explore alternative solutions.

22:23And continue special token encouraged the language model to build upon its current reasoning trajectory by generating the next intermediate step. And reflection tells the model to stop reasoning and verify the correctness of its previous reasoning steps. And exploration signals the model to identify critical errors in reasoning and explore a new solution, basically. So we recycle these three actions. The model can iteratively refine its reasoning and their reasoning process is more like how humans solve the problem. And so then the key challenge of the project is training a model to generate these new special tokens at the right time.

23:19Is that right? I would say the first stage of post-training has a challenge. we need use a format tuning stage as we propose in our paper to get the model familiar with the reasoning format of the chain of action salt reasoning. Or in other words, we should teach the model to get familiar with the special action tokens. When it encounters the special tokens, it needs to make the corresponding response and take the actions. Can you walk us through the training process and how you accomplish these goals? If we want to teach the language model to perform the COAT reasoning, there are actually two challenges.

24:07The first challenge is, as I said, it's called unawareness of the meta-action tokens. So this means that without any training and without any post training or fine tuning, the language model doesn't recognize that encountering these special meta tokens will require reflection or proposing alternative solutions. So we can think about the robot example. So without the training, the robot doesn't know how to properly take actions based on different situations, like how to grasp an object from the table. And the straightforward way to teach the robot, of course, is using some demonstrations provided by humans.

24:59And this technique is exactly called imitation learning. We need to construct some demonstration trajectories with this kind of chain of action salt reasoning format and let the language model emulate the behavior of the demonstration trajectories. And then the problem boils down how to construct this kind of trajectories. Human annotation might be an effective approach, but it's of course really expensive. So in our paper, we instead propose a multi-agent data synthesis framework. So there are two models, basically. It's also a two-player system. There is a generator and there is a critic model.

25:51So given an input problem, the generator will propose solutions and the critic model will step in and provide feedback. And these two models collaborate with each other to solve the tricky problem. And the generator might make mistakes. The critical model aims to identify this mistake and suggest an alternative solution and so on and so forth. So let me make an example. This is more like the scenario that the two students collaborate to solve a problem. or if we are solving a problem by ourselves, it's more like the situation that there are two voices in someone's head and they may or may not agree with each other, right?

26:47That's usually the case. But eventually we hope these two voices will end up one single solution and successfully solve the problem.

27:01you know, through a long debate process, right? The two boys compete with each other, debate with each other, and finally they successfully solve the problem. And we basically want to construct this kind of reasoning data. And then this kind of reasoning data is used to fine-tune a language model, a pre-trained language model. And we call this entire process format tuning. which aims to teach the language model follow the chain of action salt reasoning format. And after this kind of training, the model knows to tackle different problems using different actions triggered by corresponding special tokens.

27:53Is it learning that because that is the way that the critic is communicating to the generator when it needs to take a different action using those special tokens? The critic is more like providing the self-reflection part. It provides the feedback, the suggestions. And this kind of information can be viewed as a self-reflection process. Right. And this is triggered by this reflection special token. And also in practice, we observed that using only 10 ,000 this kind of demonstration trajectories, the model can already learn this chain of action salt format. So we basically have two stages of training.

28:44One is format tuning and the other is reinforced learning. So the main motivation for this reinforced learning stage is the model, although the model learned the chain of action thought format, but it may struggle to generalize its reasoning performance to unseen problems. I think this is expected because the model only ambulates the expert demonstrations. And the quality and the diversity of this demonstration usually, you know, constrain the performance of the model. Right, so in order to mitigate this issue, we need leverage reinforced learning to truly incentivize the model's reasoning capabilities, especially the capability to take different actions with respect to different situations.

29:48So we encourage the model to reason like an agent. It learns from the trials and errors and gradually build a policy that allows the model to get higher reward. But there is a second big challenge, the problem of sparse rewards. So let me explain So language model reasoning Is more like a long horizon Decision making process It needs to output Multiple intermediate steps Takes a lot of actions And finally It gives you a final answer And receives the rewards From the environment But we realize that The reward is only Received at the end And this causes a problem because the language model... It's attribution types of problems.

30:53Right. So if we encounter the problem that the language model cannot successfully solve the problem, there's no any reward. and as long as the model fails, it needs to start from the initial province, the initial state. Meaning so you're not teaching the model to, you're not giving it any partial credit. It's all or nothing. Yeah, so the reward is only at the end because for the mass province, the most effective reward is just checking whether the final answer is correct or incorrect and such reward is at the end, of course. Right. And so to address this challenge, we propose a key technique in our paper, which is called Restart and Exploration.

31:55Restart and Explore, right. So the idea is actually quite simple. Instead of starting from the problem statement or the initial state, we also let the model to start from the intermediate steps, the intermediate states. And it will increase the chance of the model that arrive at the correct answer. And it also encourage the model to trigger self-reflection and self-correcting its mistake. And this is similar to how a robot explore the environment instead of always putting the, placing the robot at the same starting point, we can always place the robot at some intermediate state that the model already, the robot already explored and we let the robot to restart from their intermediate state.

32:54And is a state in this case a sequence of tokens that have been generated? Right. In the language model case, that's the case. It's more like some intermediate reasoning steps. The model already generated in its previous response. Right.

33:16So consider the scenario that this intermediate state is actually a failure state. For example, if we let the robot to move an object from point A to point B, and the robot somehow fails at a certain point of point C at the middle, and we can actually let the robot to restart from an intermediate state, which is this point C, and makes a second attempt of exploration. and in practice we observed that this kind of restart and explore is actually quite effective for the language model to learn how to self-reflect and how to self-correct the mistake right and i think that's that's the that's the main contribution of the reinforced learning part Can you tie together the format tuning and the self-improvement?

34:22Are they totally independent steps or in what ways do they both need to be in place in order for you to achieve the goal? So this is more like you need a pre-training stage for the language model, right? So the pre-training stage aims to inject all the necessary knowledge and information into the language model parameter. and the format tuning as the first stage of all training is more like you teach the model to reason according to a certain reasoning format. And this format is defined using these special tokens and this is the so-called chain of action salt reasoning format. And this training stage is usually small scale because in practice, we only use 10 ,000 of training data.

35:18And the reinforced learning is more like how you can leverage this kind of special reasoning format to really solve the problem, like take different actions and finally solve the problem and get a reward from the reinforced learning. And this stage, this reinforced learning stage is usually large scale. Right, so in practice... How would you characterize the scale of that? In practice, it's more like the format tuning stage, we only use 10 ,000 examples. And reinforced learning, we use more than 300 ,000 examples. But that's actually quite efficient. Why? Because for reinforced learning, we only require the final answer as the target.