TWed: Abraham Sanders on "Training Large Language Models to Reason in a Continuous Latent Space"

Key Takeaways

welcome everybody to our first um formal informal paper discussion TW and we're we're very thankful for Abraham Sanders to lead us off um as we've been socializing this term uh we're integrating uh toix group's paper discussion format with our Tweeds so we're kind of interlacing we'll have uh our next uh scheduled paper talk is Daniel Villa leading us in a uh that decomposing that paper about deep seek deep seek R1 paper in a couple weeks um and and we've got some some other goodies and there's still some slots available this term if you want to get some uh opportunities to present um and without further Ado take away thank you very much thanks um think there's anyone online right now just double check that they can hear me if there is onine I don't but that's fine I started it up there okay welcome to um our first uh combined um reading group our language and cognition lab combined with the twed series so this would be continuing our reading group this would be the 94th NLP related presentation quite a long way since 2021 started it um this is the first combined presentation of twed so that's exciting we have a campaign on the right and a party today we'll be talking about uh a paper which uh came out this this past year called training large language models to reason in a continuous Laten space um and this paper is by researchers at um at Facebook or fair at meta and also in collaboration with the UC and what is fair I know what meta is but I don't know what fair is Facebook AI research oh okay got it so it was called Fair before Facebook became meta so they just sort of okay carried the name forward right so first we'll go through our introduction we'll get into some background on Chain of Thought reasoning and talk about some attempts to do this the latent space of the llm and then we'll talk about the main content of this paper and its so just to introduce I think everyone around this table probably knows about um and you know we all know um and if and if you're if you're not familiar with it then this is just sort of a good uh very quick it llm is a Transformer model which is going to take in token by token and it's going to predict Auto aggressively the next token so in if you take the laurum ipsum you might feed in um to get a laser pointer here left oh there we go now you can see where my house is so you feed in a um the token you feed in the token low and it's going to predict REM you feed in rem it's going to predict so all of these mod uh all of these tokens are fed in sequentially the Transformer computes its attention to create um contextual representations of each token and finally a soft Max layer is going to give you a probability distribution over the next token so the important thing is that that is that this is these models are really just next token predictors and as such the question answering tasks that these models are usually used for uh are presented as such you feed in a context such as how do you roast a ham and the response is going to be well first you preheat the oven and then you do this and then you do that and it's going to to generate this textually one token at a time however more complex questions um such as Tommy ate three ham sandwiches daily for six months and Davey ate half as many sandwiches for day for 12 did Tommy eat more ham than Davey this type of question benefits from a reasoning chain so the reasoning chain is basically just a space where the llm can break the question down into easier Parts before answering and now I know you know folks who are who are you know involved in in actual logical reasoning might come to disagree that that this is reasoning at all but you know putting the definitions aside I call it reasoning here because in the llm literature' been talking about the word that's that's associated with it in that literature is reasoning so we'll put the the philosophical discussion on definitions to to discussion so we have here the question we have here a space where the model can generate a Chain of Thought to reason about or break down the question and finally can can make an answer and in most implementations this Chain of Thought just more text right so text could say something like assuming 30 days per month Tommy ate 30 * 3al 907 which is per month for 6 months which is 90 * 6 equal 540 half of three is 1.5 so daav 83 * 30 * 1.5 * 12 = 45 * 12al 540 sandwiches thus they both ate the same number of ham sandwiches so the answer would be no so it it's much easier to for the llm to say no after that going through those those you know steps as opposed to after just seeing the question in the cont so the question that this paper comes along to ask is is this optimal that we're generating all of this intermediate steps as text and the question can be further refined to well reasoning be done in the latent space of the model so instead of actually generating the text we're going to generate some some arbitrary uh sequence of maybe not even tokens which is going to um internally transform the state of the model in a similar way as that much longer so comparing the two paradigms we have text reasoning chains of course human readable while biton reasoning chains are not text reasoning chains are restricted to only exploring one possible reasoning path at a time whereas a latent reasoning process would be able to explore multiple parall paths simultaneously um text reasoning chain is inefficient because it requires many tokens um and whereas a latent reasoning chain can acire very few tokens maybe even less than 10 are they defining lat is that or is that a term is that something that it's um yeah it's it's explored in in Prior work we'll get to that um so most of the tokens that are in this textual reasoning chain are only there for fluency sort of to glue the L the the syntactic structure of the text together and don't hold critical information that's necessary for the llm to achieve its answer um whereas a latent reasoning chain would not be Bound By by linguistic fluency and finally uh The Final Answer um may not adhere to the generated reasoning chain it could go through these steps and then finally decide well we're just going to go and and output something else anyway and this has been explored by multiple works as with a latent reasoning chain that might still happen except it it it really doesn't matter since you're not generating anything that would that that should arguably be consistent with the answer it's just some some additional scrap space for the the model to to do Transformations and this is arguably and this this last point is sort of my own extrapolation here it's arguably closer to what humans do in that we think with mental abstractions and imagery with before speaking where with rather than actually verbalize our own internal thoughts out loud every time we plan before we speak so we're going to just quickly go through some background on the literature or Chain of Thought Reon so the term Chain of Thought or coot which I think I'll use for the rest of this this presentation just for prity um is originally originating in the work of uh Jason way at all and paper was called Chain of Thought prompting elicits reasoning and large language mod however Co is now widely used in the literature just to refer to any reasoning chain of any structure not necessarily the one explored by Jason way in 2022 um so Guided by this wonderful survey we're just going to go and look at a number of the most popular ones you've probably you've you've looked at llm engineering guides you've probably seen some of these so the first is uh is the concept of a scratch Pad this was explored in um 2021 where given some you know mathematical problem rather than directly trying to generate the answer the llm would would give generate a scratch Pad where it could then go and do intermediate steps kind of like what we were just showing before this is this particular work they specifically trained or fine-tuned meaning gradient updates to the weights the model to perform this task which is not typically What's Done anymore so this is why literature doesn't usually refer to this particular work as Chain of Thought but it's there sort of as a as as the origins of the of the idea that given more space think L do better jobs so coming into Ways work I don't know if you can really see this okay um so you have the model input in this case there's a a One-Shot example so there's a question and answer followed by another question and the model's got to give the answer but the the question is a mathematical reasoning problem the answer U is elicited immediately after the question and the uh the answer and of course the model is going to get wrong but now be now um if the model is prompted in its in its example to rather than just give the answer but to break the answer down into steps then it will follow that example and it will break the test answer down into steps and come up with the correct answer and this is what was coined as chain prompting in by way in colleagues in 2022 almost immediately following that up was zero thought zero shot coot um there we go that's much easier so um zero shot coot basically Compares okay well we can do a chain of we can do first of all few shot prompting which we know doesn't work to directly output the answer oh my God I there we go um zero shot prompting so just ask the question and get some answer out um we can do few shot Chain of Thought which is what we just saw previously and now we can do um zero shot chain of thought so all you have to do is add let's think step by step after your prompt and Magic the llm will will break the problem down and get a much better answer and they they showed dramatically incre increased um zero shot results something on the order of I don't want to put the results in the slides but something on the order of the difference between 18% and you know 80% on on a a test set of math problems just by adding this this uh incantation into the prompt the paper was about this line that's the thing I said yes it was a publication that was the publication damn we usually use um the examples in the paper that led to us putting chain of thought into chap BS was explain step by step tag it the same time yeah yeah I mean nowadays the models are smart enough to know what you what you want but back then they were experimenting with the original gpt3 which was uh at the time stones and N at the time gpt3 and one of Google's largest monster models were the only two capable of doing this and that was just three years ago okay so then we have um another followup paper called self ask um so here they're they you know again they're comparing well if you just pose the question and try to get a direct answer then you don't get a good result out but if you you could do Chain of Thought you still might not get the correct answer out um in this case they're asking who is president of the US when super conductivity was discovered um so I guess it requires discovering a who is president or what year was super conductivity discovered and then B who was president during that year so it's uh first does it first does that first step and then it comes up with um a hallucination of what the Second Step would be so the an the the the updated approach here was then to prompt the model explicitly to come up with a list of follow-up questions in a specific uh sort of Q&A format that's why they call it self ask and the model is able to do better with this more structured Chain of Thought that's tailored to to Breaking the qu the question down into into specific sub questions um then we have self-consistency self-consistency as usual Compares itself against standard Chain of Thought So you have prime it does the Chain of Thought it gets the answer self-consistency will repeat this multiple times sampling a range of different reasoning chains and then marginalize out those reasoning chains to see which of the final answers is most common across the the different reasoning chain so if you get to the answer is $18 twice and the answer is $26 once then $18 wins because it's more probable so so it's it's it's asking the model to uh to do a lot of reasoning and then saying okay what is the most consistent result then you have how a quantum computer works do it a thousand times so now now we have program aided language so as usual comparing against standard Chain of Thought um but they found that if you prompt the model to break down as code so for each piece of the of the answer you're going to write a line of code that's going to Define a variable and set it to a value that corresp to what's in the answer then the model will eventually WR a line of code at the end that will that will produce the answer so it doesn't actually the model doesn't actually have to do the math it just has to come up with a way to tell the computer how to do the math based on the the the question and then you can run in a python terminal and get the answer that's pretty a pretty clever approach and uh this is I've seen others like this but this two years um even just going further uh we have least to most prompting so least to most prompting is about problem decomposition least to most prompting will take a problem here is it takes Amy four minutes to climb to the top of a slide it takes her one minute to slide down the water slide closes in 15 minutes how many times can she Slide the slide before it closes so the language model is prompted in stages first first stage stages is to come up with questions that we need to answer before we can answer the main question so to self how many times do she slide before closes we first need to solve how long dop take um so then they they prompt the model to with that that secondary question it then out uh is able to do chain standard Chain of Thought to come up with the answer and then the second stage third stage here is to take that intermediate answer append to uh the the context we had before and then try to approach this the the main question with that intermediate answer already having been uh given how did it know they even pulled that out that origin question there the trip you mean how is how does it how did it know they even do that I guess based on the way they prompted it so so what we're not seeing here is sort of the setup the extra context to to construct where does it get the word trip this is all all well it's going to be prompted to to do this task it's basically to look at a complex question and come up with a list of of prerequisite questions um what we're not seeing here is if they had few shot examples demonstrating the model how to do that then they're not showing those few shot examples here on the slide now um we have the uh a comparison of of Chain of Thought self-consistency and tree of thoughts tree of life right so we have um well first standard input output prompting we have input followed by a Chain of Thought that's Chain of Thought prompting now we have self-consistency you sample a bunch of chains and pick the the majority vote now we have tree of thoughts which sort of breaks this down as as you you have a bunch of diverse re possible reasoning chains but we're going to try to do a research to Traverse which of those the steps in those chains is more probable than others to to kind of try to come up with the shortest path between the input and the output so that's yet another uh prompt engineering technique that is that that that attempts to do this type of uh you know reasoning chain to improve llms um now these reasoning chains have been integrated into larger Frameworks such as agentic Frameworks Frameworks where you put reasoning in a loop and basically get the llm to iteratively improve itself or to attack larger problems self- refine is one of those self- refine puts the the model in a loop where it first generates an answer um gets then prompts itself to get feedback on that answer and then refines its own question and then does that several times before it's it's it it's satisfied um for example um in a dialogue context the user says I'm interested in playing table tennis the response is I'm sure it's a great way to socialize stay active the model will ask itself for feedback on that and whether or not the response was engaging enough or whether the user would understand it so it says oh well you didn't provide information on table tennis or how to play it and the user will lack understanding of the uh or or and or and and the response shows lacking of understanding of the the user's needs and state of mind so then it will refine it and rewrite the response given the feedback as this is great to hear it's a fun sport requiring quick reflexes good high and high coordination Etc so Chain of Thought is being now put into a loop to create systems so on Chain of Thought I could I can make I could probably have prepared a three-hour presentation just on that but it's that's not what we came here to talk about so latent Co tries to take this approach of take of generating textual Le reasoning chains and compress it down into the Laten space of the model for the benefits that we just enumerated before you know it's it's it's more efficient you can explore multiple reasoning chains simultaneously Etc so the first attempt to do this and the the ones who who defined this problem of of latent Co was Den at all in their two Publications in 2023 and 2024 they introduced what they call implicit coot or iot and uh the approach is to try to um rewrite the latent representation of the question so that it reflects what you would get at the end of The Chain of Thought so here we have standard no Co right you see you have 12 * 3 equals and then the model will go through Transformer layers and then it will get the answer 36 and they can you know you can look at this in terms of random variables you feed in X and you get out Y no random well there's no there's no latent random variable there um here explicit Co is going to do with sequence generation it's going to generate 12 * 3 = 6 + 30 = 36 so the 6 + 30 equals intermediate piece here um intermediate piece here is representing a latent variable Z which is going to contain you know the the information needed to to uh prompt or to condition the model generated the right answer so implicit coot takes the or attempts to to compress this latent variable Z into the Transformations that happens as information moves up the layers of the Transformer cutting out the autoregressive the need to do auto regressive um steps to get there the goal is to keep the accuracy boost without explicitly generating the reasoning chain I included this just in case somebody has the slide deck later and wants just to read it um Y which model did they use it on this one the previous one this one there's a small open source model I forget exactly which one it was either like either the original gpt2 or or maybe one of the L models so what they did is there was a three-step uh training strategy first is to is to take a teacher which is a a llm that has an explicit text toal train of thought take the intermediate hidden States at each layer that correspond to each of these tokens um each one each subsequent token you take it from the next layer so that it has the ability to attend to the previous layer and then extract them and substitute them as the you know for the latent latent states of the final token of the answer that you get in the student model that does not use explicit textual Co so here 12 * 3 6 + 30 something here this is just 12 * 3 equals something and we want the latent representation of equals to match the lat representation of 6 plus 30 um there an additional objective is then to use the latent representation that um that's computed from 12 * 3 to be able to predict the internal latent that you get from your teacher model um they call that an emulator I'm honestly I'm not sure exactly um why they needed to do that but that's what they did um and then finally they put the two together so they have the emulator model generate the um the hidden States from equals and then they if you recall from the first step the student model is trained to use those hidden states to predict the answer the second stage trains it to to predict the uh hidden States from the latence of equals and then the third step sort of puts it together so it's going to implicitly compute the latent states of the the reasoning chain and then convert that back to a prediction of of the answer um so it's it's sort of a sort of an onerous training process because you have to go through those three stages and you need a student and a teacher model um however this the authors realized that they actually didn't need to do all of that instead what all they had to do was split the training up into stages where each stage they incrementally remove part of the re the textual reasoning chain until at the last stage there is no nothing left so if you take um let's say 21 * 43 = 84 + 63 = 804 and then at the second stage you remove the eight then you remove the four then the plus then the Z Z then the six then the three basically what this does is at each stage it incrementally updates the representation of this equal sign to be able to predict stages of the reasoning process that are further and further along until it can eventually predict the answer that would have been output from the reasoning process and they found that this works even better it outperforms the three-step knowledge distillation approach that they originally used um this is the the results that they had um I guess just the the important piece is if is to compare these two iot KD the knowledge distillation three-stage approach with icsi which is um the approach where they incrementally remove tokens and you can see that in all of these different tasks the incremental removal of tokens outperforms the um the knowledge distillation approach and all of these also are outperforming baselines that are either using no or explicit coot even using gpt2 small comparing against GPT 4 which is pretty crazy yeah that's pretty crazy not surprising though because they're they're deliberately training this on this task where GPT just left to its General devices so it's kind of a baseline in a way yeah um this was the results on the eighth grade or sorry uh grade school math problems they found that um explicit coot actually does better than the all of these methods except the implicit coot methods do a lot better than not using Co at all um a second approach was pause tokens pause tokens was this from goall at all in 2024 pause tokens basically just introduce some scratch space middle of the here so 5^ 2 + 4 is 25 + 4 is right so here it would say 5 + S Plus 4 equals and then pause pause pause these are learnable tokens which means the llm can update the parameters of those tokens as it trains and then um it will just continue on where it left off the idea being that additional computation and Transformations can happen in the attention mechanism in the space which would presumably encapsulate parts of the reasoning and the goal is to provide space for the LM to think without needing output text tokens this is kind of analogus if you're a computer vision person it's a analogous to the um to the registers that are used by the dino V2 model Dino V2 contains these disposable registers here whereas pause token model contains these disposable pause tokens here and they're ignored in the output in both cases so just a nice visual representation of of the effect that this has in in a vision Transformer is that each one of those disposable registers is attending to different parts of the image so if you have the coffee cup with sugar one of those registers is attending to the liquid one is attending to the spoon one is attending to the sugar one is attending to the outline of the cup and the classification token which Aggregates all of those is able to take a summary of all of those different individual pieces that the register tokens are able to attend to so it has the opportunity to sort of summarize what's in the image before going ahead and and outputting uh to the classification head and the same concept would apply here pause tokens would allow the model to you know do some summarization over the context before it continues to to you know move forward into the answering process so this brings us to continuous chain of continuous thought or coconut here coconut is again comparing against standard Chain of Thought where you Auto regressively go and say Here's my question generate the reasoning chain step or token by token and then finally generate the answer no instead we'll get the question put in a begin thought token and then basically take the output latent of the llm which is the last hidden State before the soft Max and feed it directly into the the the token embedding of the next position so think about that there's no tokens no tokens used in between these spaces it's literally just taking the output of the the hidden state of the of the llm Transformer block the top one and feeding it directly in as if it was a token so this is truly a continuous thought because it never actually leaves the the vector space that the that the llm uses internally it's not UNT it's not um it's not deten and converted back into the discret vocabulary space so our patent ipsum example basically saying well what you know what would happen if we tried to use this technique to model the the laurum ipom you might pass in your laurum here and then it'll do some internal thinking generate this uh latent representation feed it back in here generate another one feed it back in here and then it might generate the last token of the L nium and if you try to interrogate what one one of those tokens in the middle actually would output from the language modeling head you might get something that's unrecognizable although there is a there is a method that can be used to to probe the latent reasoning chain as we as we'll see soon so the way that this model is trained is in stage just like I so the original dat langu langu linguistic Chain of Thought data includes question step one step two step three step n followed by answer the first stage is just to insert the begin and end tokens here so it doesn't really do anything it just inserts these tokens and continues to to generate step one step two step three step and answer in text um the loss is the cross entropy loss it's computed standard language modeling loss it's computed on everything that comes after the eot token so all of these reasoning steps followed by the answer but not applied to the question or the reasoning steps the first stage is then to take the first reasoning step and replace it with a continuous thought so begin of begin thought one lat one continuous thought and thought then step two step three step in iner the next stage pulls out another text reasoning step and introduces another continuous thought and then it does that until stage n and by stage n we've replaced all of the remaining textual reasoning steps and we have the question begin thought continuous continuous continuous end thought answer so we've completely eliminated the textual uh reason chain yes um so in how do you generate training data for this kind type of models I mean it can be just web right it has be free composed libraries the training data would come from a data set where the where the Chain of Thought for one of these question answer pairs has already been written out yeah exactly so how how big are these uh databases um do they need to be handwritten no they already exist I mean each one only has um somewhere between between several thousands to tens of thousands of examples that's not much yeah because they're they're fine-tuning the existing pre-trained model to do this so are they training or fine tuning or fine tuning fine tun okay because you're talking about training right so yeah well I guess it it they say training procedure but what they're doing is they're starting with a a pre-train model okay that makes a lot of yes yeah so there's also a question in the chat by the way look at the chat um how will the model perform in a complex environment where different Co strategies may be needed for different problems um so they they U evaluate this on the grade school math data set and they also evaluate it on a tree based um logical reasoning problem so so I guess it it it would just have to be fine-tuned on the target domain just like we usually would except instead of using discret CH uh chains of thought we would be able to use these uh these these latent chains of or continuous chains of thought for more efficient um reasoning at inference time but ultimately you would still need data in their target domain and textual chains of thought in your data set in order to to go through this train process so when I ask the question not sure if you can actually hear yeah I can hear you oh I was thinking the original very first version chain of thoughts was actually zero shots which allows the model to sort of be given different strategies for different problems now once it's nicely trained it's like someone learned a bad habit if that habit happen to be bad um now if we wanted to learn a different habit we have to you have to sort of like was out um the over learn content especially for the tokens the last algorithm where you add in additional tokens and directory take the output logit to put them back as input token I was just like it must be hard for them to learn something new after this yeah the authors actually discussed this a little bit they so they so as you said the original Chain of Thought zero shot Chain of Thought was emergent from pre-training right um but here they are here they're they're focusing on fine-tuning on specific data sets that already have chains of thought written out for them but they talk about the potential to do um to incorporate these continuous thoughts in the pre-training process which was beyond the scope of this paper but if that it would be done that way at a sufficient scale then we would expect uh some type of zero shot Chain of Thought capability using continuous thoughts analogous to what we see with Co so this is sort of like you know put off to to future work but recognized by the authors all right um so at inference time um what happens is right just like we discussed we put in the question uh we put in the begin thought token we generate Auto regressively is continuous thoughts no tokens remember we're just taking latent output and feeding it back into into the next position's input um and then we well they'll generate a fixed number of these like maybe three of them or four of them and then input the end of uh end of thought token to condition the answer so the main results um were found by evaluating on three data sets one was the gsmk data set grade school math word problems kind of like the ones we were looking atto QA which is a Lo a logical reasoning task on a tree structured ontology of f of fictional Concepts so we'll look at what that means in a minute and then Pro QA which is uh a data set that these authors also contributed it's a harder version of pronto QA which requires more complex search to find the correct the correct path through that that um tree structured ontology so GSM this is an example of what you find in here you know your standard grade school math word problem my my uh eight-year-old daughter is probably working on one of these as we speak um Pronto QA so here you have a question which introduces a knowledge graph of fictional Concepts linguistically so you have something like bumpuses are not luminous chuses are amenable each jpus is a lorus gorus are shus each zpus is a grimus guses are rompus and dumus are not floral etc etc etc etc finally Stella is a bmus and Stella is a zpus true or false Stella is not floral so it's going to have to use these these G to have to use these rules in its Chain of Thought in order to decide whether or not Stella is floral so the the answer would be Stella is a zpus zpus is our gorus Stella is a gorus gorus is are Romus Stella is a Romus Romus is are Yus so Stella is a Yus each ypus is a lpus so Stella is a lpus and thus or every lpus is floral thus Stella is yeah Pros QA they deliberately reshape the structure of the questions and the and the of course the assertions so here the every every assertion here is that is basically a is a b type of assertion so every uh shus is a rmus every shus is a Yus every tempus every tpus is a Gus every gpus is a bmus Alex is a rmus etc etc etc and the question is is Tom a Limpus or a scampus so it has to decide you know which of these paths are we going to explore in order to figure out whether or not Tom is Olympus or a Scrump given the knowledge somewhere in here that Tom is a tpus so it's a three-step chain Tom is a turpis every tpus is a bmus and every bmus is Olympus thus Tom is Olympus do they go into being able to figure out um contradictions because on the previous one you had not so you could have gotten you're floral and you're not floral I don't in Pros QA I don't think they have any contradictions or any any negations at all but but what they say is that there are many there are many different paths longer paths that you can go down before you hit a den end than there were in the original um Pronto QA data set well what could be a contradiction because you could have a Scrump as as orange and a Scrump as is the the previous one something is not floral then something is floral this one didn't have any negations but this one had a negation yeah there every rompus or whatever that dumass are not floral and then someplace else something is floral so you know you could get to a path but I I guess the thing I was so Stella is a bmus Stella is a zpus depending upon how that categorization of Stellas is go is a bmus and a zus from the same class are they exclusive well they don't state that but but I know that you include that could you say could you have uh a bmus can't be a zpus you could and then then you wonder how it would work um we've been sitting here for a long time I mean these are the kinds of things that mess with reasoners but but this is I guess what I was gonna say is a a different point is two things one if if you had a a knowledge gra that had that information in it you know you could Define in your ontology that okay but the question relevant to this is how would you how would you get that sort of knowledge how would you teach it to reason in that that same way so it could detect that that exclusivity you can't have a a grumpus and a rumpus you know same classification yeah so in in this they they sort of skate around that problem by not actually using negations and uh the types of examples disjointness yeah right rather they they they focus on the depth of the planning or the depth of TR of graph traversal that's needed in order to get to the answer with the motivation that traditional Chain of Thought you can only explore one of those paths at at once continuous Chain of Thought you can explore sort of a bre bread search and you can explore multiple paths simultaneously but I guess the big point and you're kind of getting at it is the kind of reason we're talking about is not usually found in Chain of Thought reasoning right it's it's just what is the problem with Chain of Thought reasoning it doesn't right it's basically forward training but the fact that that one example had not maybe oh this could be cool they don't talk about negations at all in the paper uh but I think that's a great point and um when we when we look at the results I suppose if we were to do an analysis on the where it got wrong I bet I bet you would find stuff like that um okay so um really quick question if you don't mind to they test any questions that involve like asking if like two entities are from so or do they only check do they only ever test questions where they ask where if only one entity is one of those things or not um I didn't look at the data so I I don't know but the examples in the paper show a question about one entity with the the graph being populated with information about multiple okay but it is possible I just didn't look so so anyway this is an example graph that you can create so we have Alex um Alex is sitting here and Alex is a grimus and we want to know whether Alex is a gorus or a bmus so the correct path through here would be Alex is a grimus and a grimus is a gpus and thus and aorus is a bmus thus Alex is a bmus so because we can't get from Alex to gpus in this graph there's no Edge there's no D directed Edge that will let you take that path um so there's many opportunities for the reasoning chain to hit a dead end and if you're doing standard chain of thought you might have to backt track and it would be very inefficient um there's 23 nodes in each of these graphs 36 edges um the length of the shortest path is on average 3.8 and the shortest number of shortest paths on average 1.6 so here's the the results that they came up with and the baselines are all ones we've talked about we talked about I and we talked about pause tokens um so they found that um for eighth grade math problems grade school math problems coconut outperforms a coot and it outperforms no coot outperforms pause tokens does not outperform standard coot because grade school math does not require a lot of branching and planning it's just more you know numeric comprehension type stuff Pronto QA um coconut coconut um coot and pretty much all of the baselines except for pause token do equally well and the authors figure that you know this is the Pronto QA was too easy of a data set just because it didn't give enough an opportunity to get lost in that graph but in PR QA they found that no nonot and St and textual coot achieved pretty much the same performance whereas iot um and coconut both achieved almost 100% which so this is a clear evidence that when you have a a case where where where you have to do a lot of upfront uncertain planning and traversal of a graph of possible paths the the uh latent or continuous pain of thought is superior to standard text Chain of Thought So the conclusions that they have is um like I just said in planning intensive tasks we have much better results with with with the latent um they found that the llm will still need guidance from text coot to learn how to do this in the latent space um they had an ablation where they basically did not do the multi-stage curriculum and just replace the chain outright with the continuous thoughts and they got poor results on that um they found that the reasoning chain is able to do parallel reasoning paths in the latent space so what they did is they probed the they Ed the language modeling head to probe one of the latent um one of the latent thoughts and they found that the most likely tokens coming out of that probe were 9 and 180 when the problem was 3 * 3 * 60 equals what so it's either 9 * 60 = 540 or 3 * 180 = 5 40 so 9 and 180 represent intermediate variables of two alternate reasoning paths to get to the answer and latent uh the latent thought here contained information that pointed to both of those paths and also uh chaining these these continuous thoughts will enhance reasoning as they showed by showing what happens if you gradually increase the number of thoughts that you use per per textual step during training they get higher accuracy so we we're almost out of time here so okay we got five minutes we started late too so we're in we're in the last section we get through on so the authors for their analysis they want to understand well what's actually going on here when when this model is learning this method to do continuous Chain of Thought in the latent space what's actually what is actually happening here what is the model learning so they trained a special version of coconut where they at each stage of training they're going to mix the current stage with some with 30% of examples from the previous stage so the effect of that is it doesn't forget how to use smaller numbers of continuous thoughts it continually learns how to use more and more continuous thoughts but then can also be is also capable of being prompted with the eot token to shift back to the text ual reasoning modality at any time and what they do they use this case study and they tried a bunch of different things on it first they tried standard Chain of Thought which had this error here Alex is a it when Alex is a Limpus Olympus is a Scrump I don't have my laser where's my laser there's my laser Alex is a Limpus Limpus is a a rompus a uh a ypus is a rampus so it sort of jumped on a non-existing edge somewhere else on the graph I don't know where yasses are but it um it it sort of jumped here it hallucinated an edge that doesn't exist and then and then came to the conclusion that Alex is a gpus which which uh is clearly incorrect coconut with only one latent thought also went that route every lpus is a Scrump every Scrump is a bmus and now it hit a dead end so it just output Alex is a bmus because we're at a dead end which is a they classify as a wrong target type of error if they add a second uh continuous thought then it's able to go Alex is a uh oh sorry each Rus is a bmus so it's sort of Skips the latent thoughts carry it over here into this area and then it just completes every warus is a bmus thus Alex is a bmus so what they found is that as they increase the number of thoughts that they're using from zero to six the number of instances of those errors of wrong targets longer than the shortest optimal paths um hallucinated edges and incorrect labels in the case where you have no uh textual steps left and go directly to the answer the more um latent thoughts you use the more it gets the the the correct label immediately and does not exhibit those types of Errors more frequently and what the way it's actually doing that is the following so if we an if we break this down further and remember you have these two cases where you either have only one or two Laten steps so what coconut is actually learning under the hood is it's implementing a latent bread search first search or BFS search on the tree of possible reasoning paths each thought is representing a level of depth in the BFS search so the first thought represents um all of these options here in the first um step away from Alex sorry I got a five o'clock the uh the second PL then expands this first to consider all of the nodes that are neighbors of level one nodes here and it assigns an implicit value function basically a learned estimate of How likely the next node is from the current level to the correct answer and the numbers that you see here these numbers are are are estimates of the implicit value function that is assigned to each of those nodes in the latent reasoning and the way they're able to estimate that is by taking the latent reasoning process at this stage and forcing it to decode the rest of the reasoning process in text and looking at the probabilities given by the language modeling head for all of these possibilities so if you generate one thought and then say what what is the probability that the next token is going to be uh stus versus lumpus versus zorus versus grimus you find that there's an equally likely probability maybe slightly higher for lpus than grimp this but if you then go and generate a second thought and evaluate the probabilities of all of these neighbors that could either come from moving to stus or moving to lpus or to zorus or grimus you find that only those that that come from moving to grimus yield a high probability the rest yield probabilities are very close to zero so what's happening is at k equals 1 the thought at uh the K at the thought where equals only one choosing the Limpus path is more likely than choosing the grimus path however at k equals 2 the continuations from that would come from choosing lpus become improbable they're sort of pruned out and what happened is coconut sort of deferred the hard choice of picking between lpus and grimus at k equals 1 carrying information about both of those choices forward in continuous thought but then at step two at k equals 2 it will it will be able to figure out since these nodes are now closer to the the edge of the tree that they these up here are are irrelevant while these down here are likely to yield the correct answer this is General pattern there's a parallel reasoning chain happening multiple different paths are considered and then as we go forward the subsequent thoughts allow information to be pruned out of the representation of the continuous thought while the um correct path is Amplified which they show here they basically they're picking at the point of the first thought nodes that are higher up in the tree are going to um maintain higher probabilities but at at the second thought those nodes um are going to become less and less probable as the the models able ble to disambiguate um what the correct path is in traditional Co that would require um expensive generation and backtracking one path at a time so that brings us to the end um just two discussion questions that I came up with one is that training um with these types of continuous thoughts is inefficient it would require on the-fly autor regressive steps with each training example because unlike token embeddings these values here are conditioned on past outputs so so they can't be input in parallel to the model like a single forward pass the way you typically train in llm um and you can't pregenerate them cach them on dis and then feed them in in parallel because distribution shifts after every uh distribution shifts will happen in the outputs here after every gradient step so it sort of forces you to generate while training uh which is very inefficient um my second thought here is this could potentially be applicable to other modalities like for example World modeling in video if you want to predict what will happen later on in a video without rolling out and playing out every frame you could potentially use something like this um similarly spoken dialogue if you have a model that's generating spoken dialogue then maybe um you could use latent tokens or latent thoughts to represent the planning of what you want to say without actually putting a full text translation of the speech in the context those are my thoughts about you know challenges and potential applications of this approach um and that brings us to the end so we can open it up the discussion nice first just I I think that potentially if this could be applicable to something like logistical planning of sequence of events or moves uh under some condition ofed which obviously previous move somehow conditions what the next is and each is and then each time you make next decision so for there's a program from Air Force that ask for something like that so of course this is language but this we talking about moving equipment or troups or planes or whatever um in some sequence in order to get something delivered to particular place or or so so I wonder if that could be used for planning most efficient half you know given whatever you have at your disposal to reach some particular objective I suppose so I mean the the main benefit of this approach over train of thought if you're using an LM that is for planning and a lot of times you could argue that if you have the right problem set up for a traditional plan that would give you deterministic results you even want to use an llm but if you're using an llm for planning this approach has the primary benefit of allowing you to plan more efficiently rather than you know having to recurse through uh that first search of of text textual be um textual Chain of Thought So I don't know if it will do it better but it will definitely do it more efficiently that's that's one of the I think parameters is if you can do this more efficiently and you know yeah uh hin yeah so um have two questions the first one is uh do they mention uh how to when to stop thinking because three is the correct number for problem but what about other problems that might need more steps do they have any idea of how when to stop they don't have a general formula when to stop um and of course they because they're fine-tuning the model rather than pre-training the model with continuous thoughts they only know how it performs on these problems like like uh May's question earlier but they did talk about exploring two possible ways of determining when to stop for these tasks the first was to just use a fix mied number of thoughts like like they ultimately went with and the second was to use a to add a a um a logistic regression style head top of the model that would predict zero or one basically continue or stop and um so that that way they could have a dynamic number of thoughts and they found that on these tasks it didn't matter the performance was about the same either way but I don't think that would that would generalize to all TKS I think they would have to try more than just two types of tasks to really get a sense for what would be a good General stopping okay I actually have yet another question three in total okay so the chat so okay okay uh do you want to go for them first sure um go ahead hello can you hear me y okay I had a quick question regarding um the relationship between the steps that are in coot and the number of like thought steps so if we could go back to the figure um you is this one yeah I guess like maybe like a couple of steps couple of slides earlier I think you had like a table with with a green ones oh yeah the results table yes yeah I guess K here is the number of thoughts and then uh one slide before uh one slide before I think it was result table the table okay yeah um that's quite a few slides back here we are okay my apologies my memory is bad okay so here are the number of tokens corresponds to the number of steps or I guess what what are the number of tokens here number of tokens here would be the number of of uh Auto regressive steps needed in the case of um the case of coconut they're they're not really tokens but they they basically them as if they were tokens because they they have to perform a forward pass and and Auto with an autor regressive step in order to uh to do it so nine so nine tokens would be you've done nine forward passes whereas in coot they did 92 forward passes because they needed to generate 9 tokens got it okay that makes sense okay so here they try to compress I guess each Co step into like one forward pass is essentially what they're doing and yeah one one or two it's variable they they try different one for different test oh I see interesting okay but they've only tried it the same way implicit coot did right where they take out like one step and replace it with a thought step in sequential order so implicit coot did it one token at a time rather than one full full reasoning step at a time I guess what they would consider a reasoning step wasn't exactly defined I I assume it meant like sentence but um but they they did it with more of like a block of of text where like a full sentence of the reasoning chain would be removed at a time whereas Ico did it token by token oh I see okay um I guess perhaps step is defined given the data sets that they're using um maybe yeah did they try to just like do it in reverse order so if they take step one and replace it with a thought forward pass but keep step two three four until however many and then you get the answer but did they try to like just do like one step or one and two thought steps and then just skip straight up to the uh the final reasoning step and then the answer I guess the question that I'm implying to ask here is is the model learning to forward to model like the coot steps and it's SL in space as a byproduct of this training setup um or is it like is this the most efficient way like why can't why can't the model perform multiple steps or even all the reasoning steps in just one inference uh one inference uh forward pass um like it's learning to mimic human reasoning or human reasoning steps but it's a machine it can perhaps work differently yeah so two thoughts here so the F the first thing is um is whether or not they whether or not they they could directly directly skip to the end so they would go thought one thought two two and then skip like to reasoning Step n um no as far as I know they didn't try that so um yeah no I can't remember that them talking about that the the different abls that they did is here one was without curriculum where they don't use the text chain at all they just say replace it direct with all of the thoughts all at once and train it as a single stage without thought they essentially do the same thing as icot they go directly from from um the question to the answer and incrementally remove the steps the the the the reasoning steps as they go in each stage and pause as thought is they're using the pause token instead of using the continuous output those are the only ablations that they ran um and then the second thing was well why can't you just do it in one step so I think there that it depends on how it's trained so if they're training each thought or group of thoughts to correspond to a single reasoning step then you have still have that grid search problem where you you're where after each after each particular continuous thought the model has only traversed as far as it would have with part with a partial reasoning chain um if the training objective would always only output a single continuous thought and that single continuous thought is responsible for taking the place of the entire reasoning You could argue that maybe it would learn to do that but I would expect that you would need a much larger you know a much larger model in terms of the parameters or or how you know the the the even the just the the raw Vector size that you're using for your hidden state to just to cram all of that information in so um the the the model they used actually for their experiments is is gpc2 yeah that's right what's that and F 3.8 five no they didn't use five they just used gpt2 which is surprised that people are even still experimenting with that yeah okay I would be interested in seeing if P does better but um yeah I guess a followup a follow-up remark here is that like using the Lattin space to reason is definitely like an interesting idea idea um it obviously per like it's more efficient because you don't have to do as many forward passes as if you're do if as if you're doing traditional Co um it's like 9 to 92 right for Bronto QA um but I wonder like um yeah I wonder like if the models the bigger they get the more the less thinking they need to do so the forward passes they need to do um okay that was my last remark yeah I think you'd find in general yes because the larger the models are in the parameter space we know that they can memorize a lot of stuff really well so if you're just using benchmarks like this like gsmk then if the model can squeeze more into its memory then yeah it's going to do better on in in all of the settings so I think it would be hard to ultimately disentangle better performance due to just sheer memorization due to Perimeter capacity versus actual benefit of using um implicit coot but if you're if all the experiments are done using the same model size then you've control for that variable it's a so it's a possible a hypothesis why they went with uh the dust steel gpt2 could that be because also um because the more recent models they've been mucking around with them to you know there's they're messy models when it comes to simulating reasoning whereas gpt2 is a cleaner model with from that perspective it could be yeah I mean because there's no way it's like and and don't do your own reasoning modern models have also seen a boatload of of uh data in the in pre-training including many data sets that have been created on and are out there on the internet that may already condition it to to do well on coot style tasks so yeah I in some experiments that we did in Kristen's class like a year ago or so um there was evidence that there was you know they had trained on certain certain things that we were trying to so we got unreasonable results when we we shouldn't have got them yeah surprised my question would have been if no entering a new era of deep seek research where everybody jumps on this mod sence open source but actually it's a good counterargument against it um I wasn't actually referring to to deep seek but but but deep seek is a category of supposedly better at reasoning exactly okay but the same is true of uh uh gp4 one G4 whatever many you know these these these are things that have and we don't know exactly what they're doing there but this is the crucial difference is deep seek is open source whereas all of the open AI models are close on this why I was referring to deeps as a new standard because deeps allows uh really a transpar more transparent uh setup of experiments in theory a theory yes in theory I haven't released any data yet so yeah um there was a comment on the chat um from May so neither this or the original coot can deal with decisions that have branching factors this makes me wonder how far the approach can go so I guess I I feel like the authors are are claiming here that that the the thing that that that that this approach is best at is decisions with branching factors so interesting that you're saying that it can uh do you want to elaborate on that so then maybe I must Miss some of the critical aspect of this how does it deal with branching I mean branching in the steps sort of say like in Chain of sorts you can say if yeah let me get the the part where they talk about BFS so here this thing um I think this is branching on the content but not branching in the procedure so sort of like if this is a mass problem let's do this procedure if this is a I don't know second grade Mass problem that's still the other procedure oh I see sort of say that was what I had in mind not criticize the to this paper was just my general reflection I see so like like in terms of branching out to completely different problem solving strategies like starting maybe with mathematical reasoning and moving to logical reasoning within the same even even just slightly different strategies in processing in leading to the answer what if I'm just trying to imagine what if I'm getting a much more complex logical problem that I need to decide how to like how how that work um I remember Brown like he used to show me the the super long expression then you can sort of decide on which part you work on first I may not be able to elaborate myself well I was just trying to create a scenario where even if it's a simple problem you may need to you may need to brunch so what I'll say is that we've seen with de with um deep seek and with also with O3 mini which just came out this week um these models are using some kind of coot I don't know exactly what they're using behind the scenes but they do generate a coot in in their final results and they show it to you and that Co has been shown to be very very diverse like transitioning seamlessly between code generation to planning High higher level steps to to doing um to to doing you know more like like social and effective style reasoning like if you ask them to build a game that has elements of a story inside they might they'll be able to do that so they're able to within the same chain and and of course we don't know exactly whether it's one chain or whether they're chaining together multiple chains we we we don't know but these models are now starting to show such such capability of of branching off into different I saw that too I read a lot of uh how at least with GPT how it does it but we never know whether it's one model or multiple model they could have generated a chain of s through another separate model not necessarily even newer model and then follow that with their sort of like mean trained model yeah deep seek um is is a is a mixture of experts too so in you know by definition they actually do have multiple models and each different ones are being invoked for each token so it's uh in that at least in that case they they are that that that particular assumption definitely applies because there different different sets of parameters could be invoked that that are very good at specific things but we with the uh as early as gp4 gp4 0613 um we were seeing um if we asked it to do explain step by step give it a the particular question you give it you do Nal 10 nals 20 whatever and it's coming up with different plausible chain of thoughts uh chains each time it was it was not so even there there was subtle differences um a lot of them were the same strategy but slightly different but sometimes it would go off and and do something different but it was um de side you know was it they were different we should we're like 20 after right now so we should we should kind of wrap this up yeah I think um you had your fifth question okay 2 first okay um how do you do they talk about how to query the thoughts themselves like oh you got gave me an answer but how did you get to the answer yeah they they did um they did probe The Thoughts by directly passing the so since it's coming out of the final hidden state of the Transformer you can pass the continuous thought into the language modeling head except since there was no loss applied during training on those positions you might not get anything interpretable but they did actually prob it and they found that the most probable tokens if you wanted to predict what would come next in the text space did correspond to intermediate variables of their math problems so so it's not going to give you a nice human readable explanation but it it'll give you a sense of like what concepts are sitting there in that Vector um and that's up to interpretation really just it's only by the model itself right yeah provided by the model itself let's let's thank Abraham thank you forward toing