The NEW Geometry Behind AI Reasoning (Princeton, Berkeley) transcript

Hello community, so great that you are back. Today we optimize here the reasoning architecture of a transformer. So, let's start. Now, just think about here what you have to pay for Claude Sonnet 5 at max. This is here much more than Claude Opus 4.8 max, and this is coming close here to Fable 5 here. And given here my last video where I showed you here the performance of Sonnet 5 max, I think, "Wow, this pricing is really something, eh?" Now, if we have a close look here and they give us here detail by artificial intelligence, artificial analysis, sorry.

And they analyzed here, let's say here this, I don't know, $2.75 cents here for the weighted average cost here per artificial analysis index tasks, what is here the main, what is happening inside the machine, and you see the main bulk that we you pay for and that I pay for is reasoning. The reasoning that is happening into a Fable 5 or an Opus or a Sonnet. So, if we can optimize the reasoning, we can get here much cheaper, much faster, and much better results.

So, therefore, let's talk about how to optimize AI reasoning, and let's make it a little bit more complex. Let's go for multi-hop reasoning, not a single-step reasoning, eh? So, imagine looking inside a transformer. Now, let's start at the beginning. At the first glance, everything is straightforward, no? The input tokens become the embeddings, the embeddings become the hidden states, the hidden states flow through all the dozens of layer, I don't know, 48 layers, 96 layers, whatever you have built your transformer architecture, and finally the logic produce an answer, a probability distribution.

Now, you understand that every layer here performs two jobs simultaneously, and those two jobs are a little bit contradictory, no? Because at first we have to compute and integrate new information, new knowledge, new data. But at the same time, we have to preserve here the previous useful information. So, careful how and to what extent we override something, yeah? So, every computation is a slightly perturbs here the representation that future computation might still consume.

So, this is now also a layer dependent problem. Now, let's say the representation of the word Bob, you know, Bob is beautiful, you know, immediately starts changing. Because think about it, in a transformer, every attention block, every MLP, every feed forward, every residual update moves the hidden state, the hidden vector representation a little bit, a little tiny bit. Because this is a continuous vector space, we do have noise.

So, eventually the representation still dig contains Bob, we can decode Bob out of the representation, but it is no longer the mathematical vector representation that a downstream computation were trying to consume. So, suddenly we have a new problem. We have now an information transport problem inside the transformer architecture. It is our GPT, for example. It is not anymore an information storage problem. This is now really how we transport your information between the layers.

And I think this is really significant. We have a brand new study. This is here Disco Loop, looping here discrete embeddings and continuous hidden states for multi-hop reasoning. And I think they were from, my goodness, I forgot the authors. I think it was Stanford and Berkeley. I will put it here in addition when you see this video. Now, many questions required that an AI model internalizes here the multi-step reasoning.

Let's say within a single feed forward pass before generating the answer, yeah? And the authors said, "Okay, let's study this through a two-hop reasoning." Simplest, a two-hop reasoning. A representative task where the model must compose multiple pieces of its parametric knowledge nesting in the harness nested in skill empty file nowhere yet the core at LLM itself within a single forward pass of the transformer. Now, you know we encounter immediately problems that we we know we are familiar with this.

A latent distribution shift now at a macro level. Standard non-recurring transformer sole reasoning using here adapt-wise circus arrangement right now. So, example a layer two retrieves a fact A a layer 10 retrieves here a more complex fact B that has some interlinks here between other facts. But, in general this is rigid non-generalizable allocation of the parameters. Now, you know also there's another problem that we call the reversal curse.

Reversal curse is really simple. Let's say a model is trained A is B, but it fails to automatically infer B is A. And this exposes here a fundamental structural limitation of how our standard auto-regressive transformer architecture can handle represent and compute parametric knowledge. This is just here A is B B is A. There are times when the this is true, but there are also some exceptions when it is not true. But, whatever when scaled up to a multi-hop reasoning things get really chaotic.

Because it turns out that this directional vulnerability compounds now, and this leads to severe out of distribution generalization failures. And we do want that our AI is able to do things that are not in the textbook that was not in the pre-training here of the internet. So, simple example. When the LLM learns Alice's son is Bob, you see A to B, it optimizes its multi-layer perceptron structure here to route here the token Alice closer to Bob.

However, the reverse mapping B to A is often never explicitly optimized here in the gradient descent landscape. So, this is a problem. This is something the model has never learned. And I know a lot of you might say of my users, "Hey, but come on, this is clear, no? This is artificial intelligence." Yeah, but if it's not in the training data, the AI is not able to do it. So, this is the reversal course where the AI normally fail.

Okay. So, what does it mean for us? This means in a parametric residual stream, and I have on this here for our LLM, the reverse course and the multi-hop failures manifest itself as a depth local storage problem. So, when facts are learned independently during the pre-training, the AI models does not align them at the exact layer depths that are required for the composition at the inference time. Because how should it know what crazy query I will have 2 years into the future here?

So, it does it doesn't know where the exact layer depths to be optimized here for my future queries. So, it builds a neural network and not a that is here that has this particular function. Let me explain this. Standard non-recurring transformer, well, our classical transformer that you use every day suffer from a depth local storage problem. Facts learned in the earlier layers are suddenly unavailable because of the noise that adds up over the layers where the second hop retrieval happens.

And suddenly, our term Bob is not anymore available. And then you might ask, "Why? How is this possible?" Let me explain. So, the authors of today's study said, "Let's do here this simple two-hop reasoning task." It's a symbolic two-hop reasoning task. And as the most fundamental version of implicit reasoning, no? So, we have an example. Alice's son is Bob, and Bob's wife is Carol. Now, do you note that a composite query now that I asked 2 years later is, "Who is the wife of Alice's son?" Now, we understand that Bob is here our bridge function between these two statements, and we have here logical deduction that we know that the answer here is Carol.

But, our normal recursive uh our normal standard transformer architecture fails in this. Why? Non-recurring transformer solves the two-hop reasoning through a depth-wise circuit by assigning different roles to layers of different depth. So, an example, early layers can recover the first-hop bridge entity, that is remember Bob, from Alice's son is Bob, where later layers must use this bridge to retrieve the second-hop answer, Carol, from Bob's wife is Carol.

So, now they're looking here for this bridge entity in the particular mathematical space where they're encoding them. They can happen, and you know what? They don't find it anymore. If that fact is observed during the training only as an atomic fact, Bob, and never in their compositional second-hop role, so if your pre-training data were not thousand and thousands of examples of exactly this complexity, then the training provides little direct pressure to make it available at the required depth.

This means the neural pathway in our neural uh network never learned this particular task. And I know you want to say, "But hey, this is an artificial intelligence. It knows it." No. If it has not been trained on it, it does not know it. And you know there's this idea of emergence of superintelligence, but this is for the marketing department. Every subscriber to my channel that is an absolute beautiful and highly intelligent people person will say, "But hey, I know a solution, of course, no?

Because this perspective here motivates that we use something we just had, I don't think a week back, looped transformers. Why? Because by applying the same transformer block recurrently, and now you know why I stumbled about recurrently so often in this video, a looped model reuses now the same learned parametric memory across all the reasoning steps. So, this means the first loop can recover the bridge Bob, and the sec- the next loop can query the same memory to resolve the second hop.

So, we do have Bob available. Congratulation, yes, celebration. Isn't this great? If you want to learn more about the loop transformer, we talked about it here before the new looped world models here 11 days ago, sorry. We you had here 10 minutes in 2 seconds, we talked about here exactly here on the power of loop transformer, and I showed you here this publication by Google and Toyota Technology Institute of Chicago, 2024.

So, we do a we know a lot of looped and non-loop transformer model and what is their ability. A normal transformer stores information across depth. Layer eight, let's say, compute Bob, layer 26 suddenly needs Bob, but as I told you, Bob, unfortunately, is not available anymore, it disappeared here in the noise. Looping now fixes this. A loop transformer here fixes this. So, instead of moving through different parameters, the model repeatedly applies identical parameters, memory becomes reusable, everything is great.

And you all have to say, we checked it, and you know what? It is a little bit better, but it does not work really like we hoped, and it is still failing to generalize here in the perfect sense. It is still imperfectly. Generalization is a little bit better, but not what we hoped for. At this point in the video, you might say, "What? What is happening? Even if a looped transformer is is not a solution to this problem." There is something much much deeper in the architecture of a transformer architecture.

So, now it gets interesting. Think about it. The orders now are designed to control symbolic two-hop reasoning task. We already did this here 5 minutes ago and conduct now a mechanistic interpretability analysis on vanilla looped transformers. And they say, "Listen, now we open the black box. Now I have enough of this system. Now I want to know exactly, step by step, what is happening." And they build a beautiful experiment that's a little bit too complicated here to explain here in this 5-minutes video.

Have a look at this. Now, what is the result? It turns out after the first loop, the bridge entity Bob is often decodable from the continuous hidden vector. Great. But, the representation of this vector, the mathematical representation in this synthetic mathematical space of some thousand dimensions, remains noisy and is geometrically misaligned with the clean discrete embedding of the same entity that the next loop would ideally consume.

So, you know, we have matrix multiplication, we have vector algebra. And if suddenly the representation of this knowledge as a vector representation has a geometrically misaligned feature, now you're in trouble. Why? The loop one consumes here at loop transformer, the loop one consumes clean token embeddings. No, absolutely, vocabulary, beautiful, gorgeous. Loop two now consumes, and we have here the first one, so there's already a little bit of noise here.

Some continuous hidden mixtures, and there's also a elements in. So, you can say that this in a very restricted way, the second loop is essentially operating out of distribution already. Mhm. Geometry. So, let's talk about geometry. They already said, "Okay." They looked at this, did all the experiment and say, "You know what? We have an idea." We have now the idea of adding a second channel into the transformer architecture.

We modify the transformer itself. We have a looped transformer architecture that keeps the memory reuse benefit of recurrence. Great. But now we add a decoded embedding channel as a second channel that now mixes. And it mixes in a very particular mathematical way and I want to explain this to you. Why? Because we want to reduce the loop-wise representation mismatch that is happening here on a vector basis. So, again, challenge.

Beautiful. Early layers. The later layers don't find it. And then you have a representation geometry where you would have you the actual output but the embedding here, what it should be is a different vector representation. This is here simple in a in a three-dimensional simplification, but I think you get the idea. The vector is not there how we thought it would be. It is not available for future vector operation because the vector drifted away with the noise.

And now we have a problem, yeah? So, this means that the looping with our transformers fixed your discharge question but not the representation. And the representation of the vector is the main problem why we have not a multi-hop reasoning improvement. Now, the authors decided to really test this out. This is here really the training the the results that they have here. Look at this. The probability of the bridge entity that it is Bob is here really the probability is one.

This is 100%. So, they know it is a Bob. But, if they look now here at the cosine similarity between here our hidden representation and here the embedding vector representation, they see it is not identical. It is here the cosine is 0.327. So, this is now an absolute indicator that the representation is not aligned. So, simple design principle between loops, the residual stream should carry not just a continuous hidden state computed by our function f theta, but also a clean discrete embedding read off through the LM head.

So, the authors said, let's propose here disco loop. A looping architecture based on here the loop transformers as we talked about it that augments now the residual stream with both the discrete embeddings, this is new, and the continuous hidden states across the loop recurrence. So, they say we had we feed now a second channel into the transformer mathematics that is now, if you want, a correction channel for this problem that we have with the vector representation.

So, they decided we inject now discrete token embeddings of a very particular kind into the residual stream to stop the continuous vectors from drifting away. And of course, they drift here into what we call error hallucination nonsense, whatever you call it. So, how to do this? A little bit of further information if you're really interested because loop two is is just exact same block of transformer layers being recycled.

It fundamentally expects its input to reside on a clean data manifold of the standard embeddings, no? But, it turns out that our hidden representation is significant off manifold. So, this means it already suffered an internal distribution shift. So, the core correction is not a complete substitution now that we say, "Okay, we just go and look at embedding we take here the clear embedding vector from the very first from loop one." If it would do that if we were to have no back propagation in the system.

So, what we do, we just need a correction term. So, the core correction is not a complete vector substitution and it is not a hard argmax projection. But, the authors of this code loop decided to have here a mixed channel residual architecture. So, it actively preserves the drifting continuous vectors because there's a lot of information here about the context in this vector. You just cannot exchange it here for a vocabulary vector embedding.

While dynamically, you have to make sure that we are back to the norm back on the data manifold while dynamically adding here a correction term, a softly reconstructed pre-signed vector on top of it. So, we have a little vector on top of the main vector and saying, "Hey, you know, nudge over in this direction." So, they maintain two parallel channels of information flow inside the residual stream now. So, they have the continuous channel keeps you the noisy drifting vector with all the context information intact because it holds you the complete memory of the original prompt and the reasoning plan and everything.

And we need this for multi-hop reasoning. But, now we have dum dum dum, firework, the discrete channel added now. And this reads here the continuous channel, figures out the top predicted token, constructs a clean geometric anchor or correction term and adds it back to the stream as a navigational correction vector and nudge. Of course, this has a beautiful mathematical representation here for you if you really want to enjoy this here.

Yeah, we do have here an alpha parameter. This is a token wise gate along the sequence. This is real simple. If we come here to a vector that has to be corrected, the gate opens. If we have acted that is just perfect, the gate does not have to open at all. What is here really the the if you want intellectual beauty here is phi. Phi is a soft decode and encode operator. This is now an operator that that they found how to construct and they applied independently at each position.

So, we really modifying the transformer architecture significantly. You have here all the information, but I think formula four and formula five, if you understand the core of the paper, this is it. So, by passing now the hidden states with the operator phi, the network executes a soft argmax over the vocabulary logic and constructs now a mathematical pristine vector O and W. And this discrete continuous recurrence eliminates now the directional decay.

You can see that they really went here for the loop one. They manually changed here the the architecture here of the transformer architecture. They manually injected this in their experiment here different embedding, different real presentation here and they found out, yeah, this is the way to do it. Because look at this. Here we have the OD accuracy here and we have here at the alpha parameters, let's say we we started zero, we have here about 8%.

So, this is the baseline that we have. For the two hop accuracy here, 8% is horrible a result for a two hop reasoning, no? But look at this. If we apply now this correction term continuously to the two hop reasoning experiment that the authors did, within when you open here the gate more and more, so the mechanism becomes now stronger and stronger, you see at about 0.5, I would say 0.6, the orange line comes close to 100%.

So, what you achieve is you started 8% for the two hopper accuracy reasoning and you come close to 100 or 98% within very short term. So, this proves that the vanilla recurrent layer fail solely because of the continuous drift and not from a lack of learn parameters. Because this was another argument that the authors said, "You know what? Maybe we need a bigger model. Maybe we need here scale up. Maybe we need here, I don't know, 1 trillion or 2 trillion parameters." Nonsense.

It is just the architecture. It is the complexity of the neural architecture that we have to optimize. This is where the real learning happens when we get where we can really have an advantage. So, because doing now a hard non-differentiable argmax replacement breaks the gradient descent during the training and therefore the authors had to come up with a solution. As I told you, this is here phi. This is now a differentiable decoded encode operator and at every loop the network executes simply this uh equation.

So, what does it do? It takes the noisy continuous vector h, converts it to a softmax probability distribution p, and uses those probability as weights to pull out a linear combination of the clean token embeddings from w. We have our correction term. Now, they do this here in the stage one train accuracy and they show us here, so if you want to see this here in the results, I have here in my own words, I think this is an explanation you can follow.

So, here the stage one train accuracy visualizes here the neural network training dynamics. So, very early we have a stage one accuracy hits 100%. This is your 100% and we have here this cool loop is here in orange. So, means the network was relying here if you want in a vanilla loop here on a shortcut. And the shortcut is simply that AI memorizes here composite queries inside a single feedforward pass. It is not really arguing.

It is not really I must say reasoning. It has not learned how to reason this. This is just a memorization effect. It's 100%. And then, let's say here at epoch 500 And then, at epoch 500 a massive grokking-like phase transition occurs, no? The stage one accuracy collapses here almost near zero where our two-hop accuracy here skyrockets here. So, this could show that the induction gate enforces non-inductive bias that strictly punishes a single pass memorization and forces now the weights to learn true compositional graph traversal across multiple loops.

So, you see another framing if you want it fixes this by applying internal logit lens auto-correction, no? It projects the noisy hidden state back onto the clean vocabulary manifold and interpolating the discrete token embedding back into the residual stream before the next loop begins. Anyway, how you want to frame it, I think I hope I made it clear what is happening. So, to achieve true out-of-distribution multi-hop reasoning inside the LLM's latent space, the network requires structural mechanisms, new added structural mechanisms to continue collapse, continuous drifting states back onto the discrete mathematical pristine geometric anchors.

We need those correction terms, otherwise the LLM would simply drift off. So, this is what those residual token embeddings are. They are the clean anchors, if you want, that save now the AI model from drowning in its own computational noise. Remember, we are here operating with statistical methodologies and we are not having a deterministic calculation. So, in short, in the end, I think you must engineer mixed channel architecture, or you should engineer mixed channel architecture, that periodically project the drifting continuous state against a discrete hardened manifold, our vocabulary, or in my next video I will show you a file system to realign the neural trajectory.

So, you see, the problem was never that the model was not big enough or not complex enough. Problem was a representation geometry that failed us and that we never noticed that this was the problem that we have such a high failing rate at a two-hop reasoning complexity. Absolutely fascinating. And yes, of course, I have an outlook because I found another publication published the very same day on the other side of our planet and they talk about a very similar problem that we will encounter exactly 10,000 steps later.

And guess what, it will not be an intrinsic LLM-based problem that we have to solve, but now, in my next video, maybe we go into the AI harness, into the external deterministic world, and there we have also a massive problem for the reasoning performance that we want to increase, but I think more about this in a later video. I hope you enjoyed it. Would be great to see you in my next video.