Dory - NVIDIA's Looped Transformer

0:00 Okay, so this is three models running the same prompt. But actually, it's not three models. It's the same model, but with different amounts of looping going on in the actual layers of the model. So

0:14 this is a 4 billion parameter model from Nvidia. And you can see that going through this, we've basically got different token speeds for each of the different versions of the model that's

0:25 actually running here. And each of these versions is looping a number of times. So this one is the basic one loop version. Then we've got a two-loop version and then we've got a fourloop

0:35 version going on here. And Nvidia's own composite score puts this model ahead of the Gemma 4 26B. But that's not actually the interesting part here. For me, the interesting part here is really about

0:46 how it does it. The model is the same. The only thing that changes here is how many times the model runs one block of its layers for each token. once for one loop, twice for two loops, four times

1:00 for four loops. And if you look at each of these that's running the same prompt, we get very different final answers out. And we get different amounts of tokens that actually got us here. So sure

1:11 enough, the fastest one here was the one loop version. And it looks like it did get to the right answer there. The twoop version just gave us a five and didn't even give us the breakdown in the

1:22 answer. But the for loop version went into a detailed analysis of breaking down what it was. You can clearly see it's got nice latte there and it's coming to the right answer and it's

1:34 giving us a full explanation of how it got to the right answer. Okay, so this is Dory, a looped transformer, also called a latent recurrent transformer. So this is a 4B model that comes from

1:46 Nvidia's Neotron Labs. It came out a couple of weeks ago and while this is mostly a research release, it really does show where loop transformers could actually be going. So, I've been wanting

1:56 to do a video on loop transformers for a while. And this is a good excuse because it's a loop model that you can actually download and run on your own hardware. Unfortunately, most of what people are

2:06 putting out there about this is just blog posts or walkthroughs and not actually showing you the result of the model. So, unfortunately, there's no blog post for this. The model card is

2:16 basically all you get. So I decided to actually install this, have a play with it, and build a UI for actually showing you sort of what this actually does. So the basics of this, it's a 3.97 billion

2:26 parameter model. It's got up to 256K of context. And Nvidia's actually released this on their open MDW license, so you can use it commercially if you want to. They've also released a beef bloat 16

2:39 version, an MXFP8 version, and an NVFP4 version along with other GGUFS etc. So the model itself was pre-trained on 6 trillion tokens. And interestingly, the first 4.5 trillion of those tokens were

2:53 done natively in NVFP4. They then switched to 16bit, gave it 300 billion tokens of context extension, and finally did SFT and reinforcement learning with GPO. All right. So, how does the loop

3:06 actually work? So, a normal transformer runs every layer once from the input right through to the output. What Dory does is it splits its layers into three separate blocks. There's an input block

3:18 at the start and an output block at the end. And each of those only runs once. The middle bit is where we've got the recurrent block. And that's the part that loops. If we sort of look at how

3:29 those layers are actually spread out, you've got 16 for the input block, got 16 for the output, and then you've got the recurrent block, which is actually 40 layers. So, if you look at the

3:39 animation here, what you're seeing here is a token going round that recurrent block. And we can set it to be once round, twice round, four times round, eight times round. And you can think of

3:51 it as those representations as it goes through the recurrent part are just refining those representations to get better and better now. So it doesn't go totally off on its loops. You can see

4:05 that it's actually sort of readding what came out of the input block to it each time. So in many ways it's very much like the ResNet models where you've got the original input being added to a

4:18 modification of it. The difference here is just that that modification rather than being different sets of weights each time for each block, it's basically the same set of weights that are just

4:29 getting reinforced here. So you can kind of think of this as every lap that it goes through, the model gets reminded of what the token is actually working on. So in some ways, it's kind of like

4:38 rereading the question or rereading the input basically before you have another shot at getting the best output. Now, if we were to unroll it like this, you can see that we're getting a much deeper

4:51 model here. Even though the weights for that repeated block are the same each time, compute-wise, we've got a lot more depth going on there. The other thing that's interesting here, and that makes

5:01 it perhaps more practical, is that the output block is trained to read the state after any of the loops. So, one loop, two loops, four loops are all valid models for this. You don't have to

5:14 retrain it for any particular number of loops. You can basically just set that at inference time and you're making use of eight times through these shared weights. So this whole idea is what

5:25 Nvidia calls this elastic inference compute. And as I turn up these loops here, the compute per token is going up, but the weights in memory don't change at all. In this case, it's still only

5:37 3.9 billion parameters sitting on the GPU. So already you should start to see that one of the huge advantages of this is that if we can just loop through the same weights multiple times, we don't

5:51 need to have GPUs with huge amounts of memory to be able to run a model like this. And you can see here at one loop, we've basically got 72 layers of effective depth. At two loops, we're up

6:03 to 112. If we jump to four loops, we're already at 192 layers there. But still, it's the same amount of GPU RAM being used for this. And you can see looking at the for loop version here, DOI is

6:17 actually doing roughly the compute of a 9B dense model. Even though the memory footprint of the weights in here is the size of a 4B model. Now, it's not a total win for your memory, etc.

6:29 Obviously, one place where the memory will grow is the KV cache. In per loop mode, every loop actually keeps its own cache. So more loops here means that you will need a bigger KV cache and memory

6:41 to actually support that in here. Now you can actually run it in a sort of last loop mode where each recurring attention layer keeps just one cache and then overwrites the old ones. And you

6:52 can see if we go from one loop to four loops there that last loop doesn't change at all. So that is something that you can actually try out to see how it goes. Now, Nvidia says that the accuracy

7:04 and output length stay close to per loop. So, to give you a sense of it, at a 32 context window and two loops, you're looking at roughly 2 gig per sequence versus just one loop or the

7:17 last loop out and that it's only going to require 1.32. So you could imagine that this is where you set it up so that if you're serving long context with lots of concurrent users, then this is a

7:29 setting you certainly would look at trying out to just use the last loop as opposed to the perl loop KV cache. Okay, so you can see when I actually run this, I built a little demo in here that

7:40 allows us to run these at once. I'm actually running multiple versions of the model. So, just to be sort of fully open with you here is one of the challenges with this at the moment is

7:52 that when you start VLM, you have to tell it how many loops you want it to do. So, I'm actually running three versions of the actual same model so that I can do this. Anyway, you can see

8:05 that the prompt gets sent to all of these. Now, it's not a hard and fast rule that one of them will have more output tokens than the other. Certainly the speed will be faster on the one loop

8:17 versus the two loop versus four loops etc. I think depending on the actual problem you'll see the quality of the output change as well. Often I see that the for loop one does seem to get better

8:31 quality outputs than the one loop one for example but certainly not all the time and it's not something that I would just sort of say it's the same as sort of turning up reasoning or something

8:41 like that. You'll also find for a lot of prompts that they actually get the same quality of output going through this. In this case here, you can see that basically the fourth one actually

8:54 generated the most amount of tokens just by a fraction over the one loop one. And obviously the wall time for that was actually a lot longer. Obviously, one loop is going to decode faster, etc. We

9:05 can also come over to a chat and see what the actual chat is like when we're looking at different outputs. So there does seem to be some difference in the reasoning, but I wouldn't say it's a

9:18 huge difference in the sense that you're going to get much better all the time from four as per three. If I run the same thing with four loops, you will notice you get different kinds of

9:30 responses out. Looking at this one, actually it got to an answer much quicker. And we didn't get the numbered thinking in the actual reasoning here. We just got sort of more general

9:40 thinking as opposed to the two loops which gave us numbered thinking etc. Okay. So this is the part that for me I was most curious about was I wanted to see what the model's actually doing on

9:52 each of the laps of this. So to run this I basically ported this to sort of plain pietorch and just use the bflat 16 weights in here and then I can read out every internal state of the loop. So

10:04 this is what I'm calling sort of like a loop lens in here. It then allows us to match VLM's log probabilities and to see how these outputs actually change. So you can see that if this is the sort of

10:16 rough outputs for one loop, we can actually see that there are certain things that it paid attention to more at six loops versus one loop. And you can see for example that the model gets more

10:26 confident about predicting certain tokens after more loops. So at one loop here it's predicting this bracket only at 15%. After two it's at 68 374 75 etc. Right up to sort of five and six

10:43 loops here seems to be the same. So it is going to be something to see as these models get built for production and not just for research. How do we actually know for what kind of questions do we

10:54 want to actually do this? I talked about that at the moment with VLM I need to run a different instance of VLM at each point and it is kind of cool that you can imagine putting a router in front of

11:04 that so that if you do have harder questions it routes to the six loop version as opposed to the one loop version. We can also imagine in the future that VLM just does that itself.

11:16 All right. So going back into their benchmarks against Gemma 426B Dory at one loop wins one out of 10 of these benchmarks. You can see at two loops that's already now jumped up to four out

11:30 of 10 here. And sure enough, if we go to four loops, it's now winning five out of 10 here. Now, interestingly, where this seems to win is at aentic work. So, instruction following and long context,

11:43 things like Bench verified, the ones where it tends to lose are things that you would kind of expect that it's not just about reasoning, it's about knowledge. So knowledge obviously

11:55 requires a lot more weights to store the knowledge. And sure enough on things like GPQA, humanities last exam, etc. This doesn't win. It just doesn't win at all. It actually does quite a lot worse.

12:07 And it kind of makes sense that that's probably because the model just doesn't have enough weights compared to the 26B version. So remember the 26B version still only got a small amount active,

12:20 but it's got all those extra experts that are storing those weights that can be turned on when it needs to basically get the results out there. Just to remind you that Gemma 26B is actually a

12:31 mixture of experts model with 4 billion active parameters per token. So the compute is roughly sort of the same for these. It's just how they're doing that compute and what weights are actually

12:42 being used is actually quite different. And you can see if we come down here and look at the amount of loops versus the quality that we're getting here. Certainly more loops seems to score

12:54 higher with the same amount of weights. And then don't forget we're just looking at a 4B model here. You could imagine this just doubling or tripling in size and then being able to compete with

13:05 mixture of experts models that are sort of 50 billion or greater. So anyway, just to finish up, I think this is a really interesting direction that people are going in. Loop transformers is

13:16 definitely a thing. We've had many rumors that that's what GPT6 is using. Supposedly, sort of one versus two versus three loops depending on the different sizes, etc. I would say if

13:26 you're going to deploy this today, probably you want to run two loops as the default. And I think that's what Nvidia is recommending anyway. And then maybe you have another version of it

13:36 that's running actually at sort of four loops where you send harder agentic tasks or things that do need that extra thinking in there. I think the really big idea here is that this is a

13:46 different way to scale. Instead of the model just writing longer and longer chains of thought, it thinks harder per token rather than more tokens. And the cool thing with that is that you get to

13:58 choose how hard when you actually deploy it. So the model's certainly not perfect. I don't think this sort of research version is going to replace anything that people are doing in

14:08 production that much. But I do think as we sort of look at how to sort of fine-tune these kind of models for the different numbers of loops etc. you could imagine suddenly being able to do

14:19 a lot of different tasks with a similar kind of size model that doesn't need a huge memory footprint like a large etc. Anyway, I put the links to the model card and all the details about that in

14:32 the description. If people are interested in me doing a broader video about loop transformers, let me know. I'll sort of start looking at some of the other ones that are being proposed

14:41 and maybe even look at some ideas of how to train a small one, etc. And let me know in the comments where you think this is going. If it is true that GPT6 is one of these loop transformers, we've

14:52 got to expect that the open models are going to catch up pretty quickly to something like that. Anyway, as always, if you found the video useful, please click like and subscribe, and I will

15:02 talk to you in the next video.