Inside s1: An o1-Style Reasoning Model That Cost Under $50 to Train with Niklas Muennighoff - #721

0:00At a high level, both S1 and R1 are attempts to replicate O1 from OpenAI. And I think the key difference is that for R1, they try to go the full way, replicate the entire pipeline as much as they could. Whereas for S1, we focused on what is the minimal approach, the simplest way to get those two exciting things about R1, which is the strong reasoning performance and the test time scaling. What is the minimal recipe to get that?

0:46All right, everyone, welcome to another episode of the TwiML AI podcast. I'm your host, Sam Charrington, and today I'm joined by Nicholas Munikhoff. Nicholas 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. Nicholas, welcome to the podcast. Thanks for having me. You are welcome to be on. I'm looking forward to digging into the conversation. We are going to be talking about your very recent paper on S1, Simple Test Time Scaling. I'd love to have you share a little bit about your background and kind of how you got started in the field.

1:26You're a first year student there at Stanford, right? I originally did my bachelor's at Peking University and I was focused on business and finance. But then midway through, I realized that artificial intelligence was going to be probably the most important paradigm of our time. So I decided to completely switch my focus and do AI research instead. And then after graduating, I spent some time in industry at Hugging Face and others and eventually landed here in the Stanford PhD program as a first year student. Let's talk a little bit about S1 and how this project came about. you saw what was happening with the O1 models, presumably.

2:05And was that the impetus? Yeah, it was. Interestingly, it was right at the start of the PhD that they announced O1. So it felt like the perfect thing to jump on and do research on this topic. And I think there were two things that we were really excited about with the O1 release. The first one was the really strong reasoning performance that they showed. So gains on math tasks, but also science questions. And the second really exciting one was that those gains seem to get better with more compute at test time or test time scaling. Can you talk about other approaches that folks have taken to dig into test time scaling and try to replicate what OpenAI did with O1?

2:51At a high level, there's two different approaches to test time scaling. One is you scale in parallel, and the other is you scale sequentially. In parallel means that you're running multiple computations independently. So for example, you take one model and run the question through it multiple times, and then you aggregate those answers somehow. The most common way to aggregate them is majority voting, where you just take the majority answer. And similar to classical assembling techniques in machine learning, this will slightly improve the performance of the model because you're not bound to one-off errors or mishaps of the model.

3:28And you can do more fancy things there like have a reward model decide which output to pick and so on. And on the sequential side, you have the model generate a single reasoning trace or reason in a single chain and then iteratively refine what it's doing and improve on its previous current answer. And intuitively, this one should scale better because you're not making potentially the same mistake multiple times in parallel. But the model is reasoning in one continuous chain. So it can build on intermediate steps and try to arrive at a better answer. But the two of them are not exclusive. You can combine them and there is some synergy.

4:10And in the OpenAI 01 blog post, even though the 01 model with its long chains of thoughts that it generates, it is sequential test time scaling. They do show that they can add test time parallel scaling on top of it with majority vote and further improve the performance. So they generate multiple chains in parallel and then aggregate them with majority voting. Is there work being done that looks at, rather than aggregating via voting, aggregating via an additional model processing stage? So maybe, you know, summarization or something like that, or training a model to look at the output of these individual parallel runs and kind of synthesize them.

4:55So there are approaches that leverage reward models in this process. So I think best of N or even some tree-based approaches like rebase, which we also explore in the S1 paper, where there is a second reward model that assigns a reward to the current intermediate or the final output of the model. And then based on that, you just take the one with the highest score or the few ones with the highest score or something like that. And that way, hopefully arrive at a better answer. Yeah, I think I was thinking less of like using the reward model to assign scores and more like using a model to synthesize the text that the other models spits out.

5:36because if the models are, you mentioned how if the models are just generating text in parallel, then you, a disadvantage relative to sequential is that the, you know, other models aren't necessarily learning from the output of the, you know, the other N-1 models. But if you've got a model sitting in front of that, that can potentially learn from all of the outputs, maybe it can generate a higher quality. output oh that seems like a an interesting research idea so maybe we should maybe explore that so it could be i guess it could either be at the very output level or even as you hinted at during the generation of each of these models in parallel maybe there's some synchronization step where all of the models say okay this is where i'm currently at those are the mistakes i've made and then sort of they try to improve based off of that i think that's very exciting um yeah trying to either put another model at the end and improve based on all of these parallel tries.

6:41I haven't seen any work on that yet, but I think it's very exciting. Your work came out right around the time of the DeepSeq R1, you know, excitement. Well, it came out around the time of DeepSeq R1 and the subsequent excitement, I should say can you contextualize s1 with r1 and like how those projects relate to one another i think at a high level both s1 and r1 are attempts to replicate o1 from openai and and i think the key difference is that for r1 they try to go the full way replicate the entire pipeline as much as they could Whereas for S1, we focused on what is the minimal approach, the simplest way to get those two exciting things about O1, which is the strong reasoning performance and the test time scaling.

7:41What is the minimal recipe to get that? And in terms of the results that you saw, I think the headline was that S1 was trained with about$50 worth of compute and matched or exceeded the O1 preview scores on some benchmarks. Can you elaborate on what you saw? So on the cost front, our final training run took 26 minutes on 16 H100s. And if you ran 16 H100s on one online platform, so I looked at Prime Intellect, for example, it costs around, I think,$40 per hour. So that would be 20 for 26 minutes. So very cheap for that final run. And I think that's where... That's the cost for 16? That's a cost for two nodes, 16, yeah, on Prime Intellect.

8:31It's a great platform. And we, of course, also use resources to generate the data and create other parts. But for the final training run, that's all you need. And since we published all of our data sets, anyone should be able to replicate our results with those 26 minutes of 16 H100s. Yeah. I think the performance was better than 01 preview on some math benchmarks, like AME and Math 500. But on GPQA, which is like common PhD science questions, it was worse. And also O1 Preview is an older version. So I think at this point in time, they're already at O3. So of course, need to contextualize the results a little bit.

9:11But nonetheless, I think given the very cheap recipe, it is quite impressive that we could get this strong performance. Awesome, awesome. And we'll talk a little bit later on, I think, about how you might go about evolving your approach to keep up with advances. But I did want to dig into the training recipe in particular. And a couple of the major steps there, one of them was the data set curation. And then this interesting bit that you did with that you called budget forcing. But let's start with the data set. How did you pull that together? Initially, we collected some of the hardest questions we could find online.

9:50So we took some really hard Olympiad questions from covering subjects like math, physics, and others, but also brain teasers from quant interviews and these kinds of difficult problems where you'd think that a lot of reasoning is necessary. And we assembled 59 ,000 of those. And then we went through a process to filter them down to just 1 ,000. And the key criteria in this process were focusing on quality, diversity, and that are very difficult. And the way we quantified these three metrics is that for difficulty, we looked whether previous models could solve them. And if they can't, then we thought these would be very challenging.

10:37So ideally, like a previous poor model cannot solve them, but we still know that they are solvable either by a really strong model being able to solve them or some indication that there's a solution or something like that. And then for diversity, we made sure to sample from very different fields. So not just math, but also astronomy, biology, and various others. And finally, for quality, we had some very simple filters, like ensuring that there is no duplication of questions, that the questions are all properly formatted and so on. And ultimately, this led to our final 1 ,000 example selection.

11:16And that number 1 ,000, did that come from your intuition? You know, let's see what we can do with 1 ,000, or we should be able to get this down to 1 ,000? Like, where did that come from? There is a great prior paper called Lima, Less is More for Alignment, And in that paper, they showed that with just 1 ,000 samples, they were able to train a model that was a pretty good chatbot. So based off of the original Lama models, they used these 1 ,000 samples to fine-tune it. And they framed this as the superficial alignment hypothesis, which is that to align a model, you actually don't need that much data.

11:55You can make it much simpler. And we wanted to revive that, but for reasoning. And that's where the 1 ,000 number came from. And it's also somewhat catchy if we can do 1 ,000. So you generated this data set of 1 ,000 questions or 59 ,000 questions, but then you filtered that down to 1 ,000 questions. You mentioned, and then the next step, I think, was doing distillation. And you did that against Gemini. Can you talk a little bit about distillation and the inspiration for that and, you know, how that process worked? Yeah, of course. Initially, we used the Google Gemini Flash Thinking API, and we just ran it over even the entire 59 ,000 questions.

12:48And then we took from the API the thinking trace and the solution. So at that point in time, they already had a reasoning model, and we directly took the traces and answers from that model. But later, after R1 was released, we switched to the R1 reasoning traces and found that actually they are much better. And we were able to train a slightly better model with that that we just released recently called S1.1. You use the Gemini thinking experimental and I am curious if you've gone back to look at that. And I'm only asking this because in my own experiences with that model, it's changed a bit over time.

13:35And when you talk to Google, it's like it's experimental. We're gonna change it without notice. It's not version, that kind of thing. But when I've used it for coding, I noticed like some pretty dramatic changes in the way it responds to questions over the course of the past, like say five weeks or four weeks? Did you ever go back and look at it again? Yeah. Or did you notice a shift in the course of the project that kind of changed the way you, changed the kind of results you were seeing out of that model? Yeah, the biggest change was, it seems like after our paper got released, they removed the option to get the thinking trace from their API.

14:19and besides that they've also followed up with newer and presumably better thinking models so maybe that also motivated that change and yeah it's changing can you get the thinking trace from those models no i think they removed that option for now okay but that's also partly the reason why r1 is a much better option if you want to generate thinking traces besides the reason that it leads to better performance as we showed with our follow-up. Because you can, because they make the thinking traces available. Exactly, yeah. In talking about the dataset curation, you mentioned difficulty and the criteria was roughly, you know, models have a hard time with these questions, but then the distillation process is let's ask a model these questions and then capture its answer and its reasoning trace.

15:10You know, presumably those are kind of at odds with one another, Like the models are going to generate wrong answers and think about them in wrong ways, but then you're using that information for training. Talk a little bit about that paradox. I think one hypothesis is that you're bottlenecked by the capability of your teacher. So in our case, for Gemini, on one of the benchmarks, Amy, it got 60%. And after distillation, we got 57 with our model. so we decided to stop there at that point in time because it felt like there's not that much room to go further if we keep distilling from Gemini because we're already very close to its performance we're just using its thinking traces and its answers so how much further are we able to go but then with R1 which had much better performance than Gemini we went a lot further and performance improved so I think that is a an important thing to think about is how good is your teacher ideally of course, it's as good as possible.

16:11But not to say that you can't be better. Maybe it is possible to get better than the teacher you're distilling from. But with 1000 samples, that's especially difficult. Yeah, yeah. So the veracity of the answers that the teacher produces is a challenge for this training. Yeah. And it does get many of these questions wrong. So in our case, we actually trained on many wrong questions from the teacher because we didn't find there to be a big difference if the questions are actually correct by the teacher or not. It's much more important to show the model how to do the reasoning and arrive at some answer, even if some of them are wrong.

16:51That's usually fine. Oh, interesting. And did you kind of capture and track a metric, you know, correctness of the training data? And like, what was that percentage? Yeah, we didn't capture that. But I recall a recent paper called Stream of Search, where they also show that with, if I remember correctly, around 50 % of answers being incorrect, they were able to train a better model than one that was only trained on 100 % correct reasoning traces and answers. and part of the reason there was also that i think the incorrect ones they had a lot more reasoning and backtracking because presumably the model knows it's not quite the right time yeah exactly and so it like goes a lot further and that's potentially why it leads to better performance so the intuition is like you're fun what you're trying to do with uh um with the supervision is to teach the model how to think and even if it doesn't get the right answer if you're seeing it kind of go through thinking cycles, that's going to provide some signal to the model.

17:59Yeah. And especially the backtracking in the thinking is really important. So moments where the model goes, wait, is this correct? Let me go through it once more. And these are, of course, more likely with really hard questions, maybe even where the model ends up with the wrong answer. Talk a little bit about the budget forcing aspect of the training process. Yeah. So after training on those 1000 samples, we got a good reasoning model with strong performance on math and science questions similar to 01 and 01 preview. But we didn't yet have the test time scaling. So the second goal, how do we get better performance if we give the model more test time compute?

18:40And we can go into details why this is desirable. But the way we got there is... Well, yeah, let's do that and just start it by elaborating on what you mean. So you have a model that you're training, you're distilling the kind of curated data set and like you're fine tuning this model. When you say you didn't have test time scaling, what does that mean specifically? Yeah, it means that we run the model on those math benchmarks, but it always uses the same compute. So say we ask it, what is 10 plus 10? No matter what we do in the simple setup, it just always does a forward pass through the model and then provides an answer.

19:32there was no clear way at the time for us to change the test time compute so how much time does it spend to think about this answer and and the reason through the question and of course for very hard questions intuitively we'd like the model to think for longer because they are harder it should take more time and for simple questions we don't want it to spend that much test time compute so it should give us an answer very quickly and so can we produce like a plot that It shows that if we provide the model with more of this test time compute, its performance improves continuously. Got it. So relating back to the whole thinking fast and slow idea that kind of kicked off some of this test time scaling work.

20:19Awesome. Awesome. And so how did you implement that? Yeah. So the way we went about it is, again, trying to find a very simple approach. And the one we put forth is called budget forcing. And the idea is we have some test time compute budget. And we usually measure that one in tokens, which is how many text units the model produces. And given this budget, we just force the model to generate that long of a reasoning trace. So either we cut it off. So after it has generated, say, a thousand tokens, we just cut off the thinking. and then the model will be forced to provide its current best guess, current best answer.

21:02And that works well for cutting it short. But the key question for us was how do we extrapolate it? So how do we make the model think longer when it wants to stop and wants to provide its current answer? And for that, we force the model to continue thinking by injecting weight into its current reasoning trace. And when we inject weight into the model's reasoning trace... And to be clear, sorry for interrupting, we're talking about the word or like the token weight, not weights as in parameter weights or anything like that. Yeah, of course. We're talking about the word weight and actually the string, which is then transformed into a token, just right when the model is currently generating its answer to some question we give it.

21:48And so the model might be something like, okay, given what I've reasoned above, the final answer should be three. And then it wants to finish and stop the generation. And what we do is we disallow that and just put wait, and then let the model continue reasoning through it and generate further. And of course, what the model will do is it will be like, wait, is this correct? Let me go through it once more. Because the way we train these models is through next token prediction, where they learn to predict future words. And so if we provided the word wait, then the future word after that is most likely something like, wait, is this correct?

22:28Wait, let me do this once more. And we actually have some of these wait examples, even in the 1000 samples. So some of the reasoning traces, they are naturally something like, wait, let me recheck this. And so the model learns to generate these tokens after wait. And then of course, when it rechecks, it goes through it once more. And if the previous answer was incorrect, then it is likely that it will change it and improve on its answer to get a correct answer. Versus if it is already at a correct answer, you might think that, okay, might this impair the model and lead it to an incorrect answer?

23:06And while this can happen, it's generally less likely. It might do something like, wait, let me double check this. And then it'll be like, it seems my original answer was correct. And the reason we think that it's less likely to change to an incorrect answer is because verification is easier than solving. So when the model verifies its current answer, then it's easier for the model to know whether or not it's already correct. Did you try other words besides wait? Yeah, we tried alternatively. And also no word at all, like just letting the model continue. And wait seemed to provide the best performance, but we definitely haven't exhausted the list of possible words there.

23:47and yeah, I'm excited to see what other people try. And can you elaborate a little bit on how exactly that word is injected in? Is that in kind of the inference generation loop itself or someplace else? Yeah, exactly. It's in the inference loop generation. Specifically, we set the end of thinking token. So the way our generation works is that we put in a question and then after the question, there's a delimiter that denotes the beginning of the reasoning process or thinking process and then there's a end token for this one or end of thinking token after which the model will generate its final answer based on its thinking process and whenever the model tries to generate this end of thinking token delimiter we disallow that so we have it as a stop token in our generation pipeline.

24:41And then instead inject wait right into the model's current generation, and then it'll just continue generating off from there. And of course, after wait, it's very unlikely that the model will just try to put the end of thinking token again, because now a new sentence started, so it'll do some more thinking and reasoning and hopefully get a better answer. And now the insertion of wait doesn't have anything to do with the correctness of the answer, how far the model had gotten or any of that. It's just random. We do it whenever the model wants to finish. So presumably it already has an answer.

25:17And the question is just, is the answer correct or not? And for that, we didn't do any further optimization. But I think you could do something like check with a reward model first or have the model check itself. Is it worth continuing further? Or should I just end the thinking here and provide my final answer? I think there's definitely optimizations there to be done. Okay. So are you always providing some constant number of weights for each generation? Or is the number of weights variable or random? Yeah, we always pick the same number and found in our experiments that up to four can lead to performance gains.

26:01but beyond four times eventually the model will go crazy because it's like getting weighted all the time and then it'll just go off into an infinite loop and bad things happen but up to four times it's still relatively natural and the model will properly use them to improve its performance so yeah we got pretty consistent math gains up to that degree when i'm hearing this i'm hearing like okay for this training run you're always saying wait four times that is training the model to always think, you know, four lengths worth. But I guess what is not true there is that the length of thinking tokens is constant.

26:40The length of thinking tokens is just whatever the model generated based on next token prediction. And so that's going to be variable even if it's always four weights. Yeah, that's a great point. So we can combine it with a maximum, or we could also set a minimum token budget for the thinking trace. But generally we find just adding weight is enough to get us some additional test time compute with better performance and varying that a lot doesn't bring out further gains. And maybe to tie back to the budget forcing and cutting off the generation. So if we put these together, then what we arrive at is like a plot where we have different test time compute budgets.

27:21So either say we cut off at 1 ,000 tokens, we cut off at 2 ,000 tokens, We cut off at 3 ,000 tokens, and then we force wait one time. We force wait two times, three times, four times. And each of those dots has better performance, and we get a very clean test time scaling trend. Meaning each of those is a separate independent test. And so if you're waiting four times, you're not using the cutoff. That's a separate thing? Exactly. We run separate evaluations. So the idea is that ahead of time, you have an idea of how expensive your question is to the model or how much you want the model to think about it.

28:04And based on that, you set your test time compute budget constraints for the model and then let it go with that. You could make it variable as well. So having the model itself decide. But that's kind of already happening in the original thinking trace. So naturally, the model will use different amounts of tokens for different questions. So we already have that. It's just if you additionally want to set a budget constraint to enforce something, you might say the model should answer in one minute or in 10 minutes, then you should use that. Is it minutes or is it tokens? It is tokens. Time-based?

28:39Yeah, it is tokens, but at the end of the day... But you've got a constant tokens per second. Exactly. So depending on how fast your GPUs are or how you're deploying it, you can convert that. I assume in a UI interface, you would turn it into minutes. Or as OpenAI is doing, in their API, they have three parameters for a reasoning effort key. They're low, medium, and high. And so people can choose what they want there. Not exact time control, but also some amount of control. Going back to the dataset, there was a mention in the paper about decontaminating samples. Talk a little bit about that and what that means.

29:18Sure. One big problem in large language model training is that we're training these models in huge amounts of data. And when we evaluate them, there's a chance that some of that evaluation data was in the training data. And then the evaluation isn't really proper anymore because the model has seen these things already. It might just be reciting its training data and not generalizing or really improving performance. And so that's why we employ decontamination. In our case, what we do is we take all of the evaluation samples and check for any matches in our training data. And if there's any match, then we discard the sample from our training data.