Text Diffusion — Brendan O’Donoghue, Google DeepMind

I've got a couple of people still filtering into the room, but it's mostly intro stuff for the first couple of slides so they won't miss anything. Okay, welcome everybody. My name is Brendan, I'm a research scientist at DeepMind.

I'm talking today about text diffusion, which is kind of a more forward-looking research area at DeepMind. So you're probably familiar with image and video diffusion, which is kind of state-of-the-art for these modalitiesright now where you, you know, you take ground truth, say, image, you add noise to it in training, and then you train a neural network to remove that noise gradually.

And then at inference time you just initialize the picture with pure noise, and then you iteratively refine out the noise to recover back to, you know, whatever image or video or audio or whatever you're looking for. And the principle is essentially the same for text, for text diffusion, where you start with a clean sequence of tokens, like a clean sentence or something like that, and then you gradually add noise, you corrupt it somehow.

There's lots of different ways to do that. You can do it in a continuous or discrete way, but let's just say discrete for now, which would in this case just mean adding random tokens or replacing tokens with other random tokens.

And you do that for a bunch of different noise levels, and you train the neural network to try to fill in, to try to correct the mistakes, basically, in the text. And then at inference time you initialize the sequence of tokens to just pure noise, like pure random discrete tokens from the vocabulary, and then you iteratively refine through that to fill in the information in the order that the neural network wants to do to recover back to, say, a clean sentence.

And then in practice, so we, you know, I saw, I showed you some GIFs here of what it looks like for images. You get very similar-looking outputs for text, where it kind of, you know, starts off all noisy and then gradually, like, fills in the text, and you get clean, relatively clean outputs at the end.

Okay, so the team I'm on, we had a research demo release one year ago now called Gemini Diffusion, which was a variant of a Gemini model which did text diffusion instead of autoregressive next token generation. And that was like a research preview that was open to about 100k people.

And, you know, we're still, you know, keep posted for new developments in that direction upcoming soon. And we did have some good numbers at the time, but again, it's a year ago, which is like, you know, prehistoric times in this field.

Our kind of main comparator model was Gemini 2.0 Flash light at the time, because that was the architecture we were branching from, the text diffusion model. And we basically had very similar quality across the board there, mostly, you know, a little bit of advantage in code, a little bit of disadvantage in some other areas, but kind of relatively similar performance at much better latencies.

But again, yeah, this is a year ago, so I wouldn't fixate too much on these numbers. Okay, so what's the difference between autoregressive generation and diffusion? So in the standard vanilla Gemini, Gemma, you know, GPT, whatever, generation of text, you do this, you know, you have some context that comes in and you want to generate some response to that, and the model does it one token at a time.

So it generates the first token, and then condition on that generates the next token, and so on. Whereas in diffusion, you have, you know, context will come in, whatever that is, and it'll initialize, like I mentioned, like a long sequence of tokens.

It could be hundreds, it could be thousands, it could be shorter, depends on the model, to be random noise. And then it iteratively refines that canvas to remove the noise over the course of a few denoising steps. So rather than one token at a time, it does the entire block together, but over a couple of iterations.

So it's not just one pass, it does multiple passes, but it gets to attend to the future tokens and so on. So it's kind of a different way of generating text. So that obviously has some pros and cons.

So the main pro that people really like, and probably is the biggest advantage that text diffusion models have, is that it's just faster inference. It just generates faster tokens per second, because it makes much better use of the hardware, the TPU and the GPU.

And I have some slides on that to explain why. But some other advantages are it can do bidirectional attention within this canvas of tokens. So, you know, autoregressive models can only attend to the past. They have causal attention within their, you know, their transformer, whereas a text diffusion model is not restricted to that.

It can attend to the future. And that has some interesting properties, like it can do self-corrected generation based on future tokens. So it could, like, you know, do some reasoning, see that it got the answer incorrect, and then go back and fix the reasoning and do it again.

I have a demo of that. Because this process is iterative, it does a number of steps to respond. That means the model can actually do adaptive computation. It turns out that you can train the model to spend more time on harder problems and less time on easier problems.

And like diffusion models in general, you can do things like in-place editing, where you say, like, fix the last tokens and give me the prefix that corresponds to those tokens and stuff. But the main disadvantage it has, and the reason why it's not kind of used everywhereright now, is lower throughput for large batches.

So

autoregressive models, they're slow, but you can have a big batch of queries together, and then if you push that through the neural network on the GPU, each individual user is slow, but you make good use of the TPU by doing that.

And so you can serve a lot of queries. And so you keep your costs down. You can serve a lot. Whereas since text diffusion does multiple forward passes on the same data, it, you know, hits a compute threshold basically earlier, and it's, even though it's lower latency for any one user, it tends to be, you know, lower throughput overall.

So higher costs to serve. Andright now, you know, if you've played with Claude recently, you'll know that they have some throughput concerns. So

people really care about throughputright now. And so no one's landing text diffusion into any of these big models, primarily because of that disadvantage. It's just too expensive to serve, even if it is much lower latency. Okay, so just leading into why does it have lower latency, in case you're not familiar with kind of the architecture of how GPUs and TPUs run today.

So in a GPU, there's like a Tensor Core, which does these, like, big matrix multipliers. It's very efficient. It has a lot of flops or hops or, you know, whatever. And then the memory that sits on the TPU GPU, this HBM, that's where the weights and the activations and everything are stored.

And that has to transfer over from the memory, all the weights and the activations and the KV cache, into the Tensor Core in order to do the computation. And so it has to flow through this bandwidth channel. And that bandwidth channel is very tight.

It turns out that both GPUs and TPUs have a lot of flops and not that much bandwidth. It's quite hard, it's expensive to put bandwidth onto these chips, and it's easy to put flops. So because of that ratio, you can, you know, you can, if you do, the more flops you do for each streaming amount of data you put through, the better.

So when we're serving an autoregressive model, these chips are memory bound. They're basically bottlenecked by this bandwidth. So when you do autoregressive next token generation, for each token, you're doing one token at a time, let's say batch size one, you have to stream over the entire neural network and all the KV cache and everything to get one token, and then you do it again for the next token and so on.

Whereas for text diffusion, you're generating, say, 256 tokens, you still stream over everything, but if you can do that less times than the number of tokens, this iterative refinement process, then you'll get a speed up. So if you can do, say, 24 passes to generate 256 tokens, you'll be doing 10 times fewer memory transfers than an autoregressive model.

And if you are truly memory bound, then you'll be 10 times faster, something like that. So that's the real reason, that's the hardware reason why text diffusion models have much lower latency than autoregressive models. Okay, so, you know, we had this Gemini diffusion demo, maybe some of you got access to it, last year, and, you know, that was able to hit, you know, something like 2,000 tokens a second pretty consistently, depending on the length of the query.

Obviously, it depends. The longer sequence it's generating, the less it's prefill-dominated, and so you can really lean into these very long sequences of very fast tokens. But if you're only generating one token, for instance, you'll be just dominated by the cost of the prefill.

And the tokens per second number that was reported on this webpage was incorporated prefill and everything like that. So this was 2,000 tokens a second as genuine raw tokens that you would receive in your web browser. Okay, so that was the kind of a whirlwind tour of text diffusion and its main advantage, which is latency.

But I want to dig in a little bit into some of the other advantages that text diffusion has, which are kind of a little bit less talked about in the literature, but I think are pretty cool, and this is why I'm excited about it.

So at I/O last year, Google I/O last year, they showed this demo for the text diffusion model, which is, you know, this is a really easy prompt, but, you know, lots of models actually make mistakes on it. So the prompt is, if you go to the next slide, "What is the square root of 81 times two-thirds squared plus, you know, blah, blah, blah?"

And I think the answer is 39 to this problem. And so, you know, you pass that into the model and you ask the Gemini diffusion to respond to that. And after one forward pass, these are the tokens it's generated.

So one forward pass through the model, it's starting to respond. Now, it's doing this iterative refinement process, so one forward pass is not all it's going to do. But after one forward pass, it has this. So it has answer equals, and then it says 60.

It's not correct, but that's what it's guessing for now. And then it starts to do the reasoning. So it has solution, calculate the square root of 81, and so on. After two forward passes, it's changed 60 to 49, and it's gotten a little bit further into the reasoning.

So it's gotten, like, you know, five steps into the reasoning. Two squared equals four. And some of the blue tokens are kind of still going to change. And then after three forward passes, it's actually gotten all the way through the reasoning.

So it gets the answer correct at the end, 36 plus 3 is equal to 39. And it's gone back and fixed the original response to, say, 39. So it had a mistake twice, 60 and then 49, but once it finished the reasoning, it was able to return back and fix the mistake that it made at the start.

Now, it's going to do a couple more forward passes in order to fix, you know, some of these tokens that aren't quiteright in the text. But overall, that's basically the structure of the output that it'll return. And this is, you know, this is a property that text diffusion models have, this ability to do bidirectional reasoning.

So to not only see the past, but also see the future that it's going to utter, it's going to respond, and also to use that information to do self-correction. So it made a mistake, but it was able to, you know, had another forward pass, it was able to go through and fix that mistake.

At the time, you know, much, much bigger models than the one we were serving made a mistake for this problem. So both ChatGPT-4o, which was new at the time, and Gemini 2.5 Flash, which was brand new at the time, both made an error on this exact problem.

So you give them the exact same prompt, and then they would say, you know, remember the answer is 39? GPT-4o said 40, because that's the best guess it can do at that one token. It went through the reasoning, and then it said 30.

It did manage to figure out it was 39. It said, "I made a mistake. I made a mistake. It's 39, not 40." Gemini 2.5 Flash also made a mistake, said 42, and then it actually just stuck to its guns and never changed it and said 36 plus 3 is 42.

So it, like, incorporated the error into its reasoning later. And these are way bigger models than the Gemini diffusion model. So it really is a property, a flaw of autoregressive models that the text diffusion models don't have. And, you know, you can fix this with modern reasoning thinking models, but, you know, then you're just kind of punting the problem into something else.

But anyway, okay, so that's one advantage, which is bidirectional reasoning, self-correction. Another one is what I hinted at before, which is dynamic computation. So you can give the model more time at inference, more forward passes. You give it a bigger budget, and it can just do better.

It's not exactly monotonic, but it is roughly monotonic that the quality across every eval basically just continues to go up. Because it gets, even if the solution is almost entirely clean and correct, it gets to look at it and see that it made a mistake and then fix it.

So you get this nice kind of curve where you always see, as the number of denoising steps, which is the forward passes, increases, you get, overall, the quality gets higher. And these are just six coding evals that we monitor internally.

On top of that, a slightly different concept is the model can do adaptive computation, which is that you can allow the model, you train it in a way, to determine itself when it is finished. And then for easy responses, it can use a little bit of compute, and for harder responses, it can take longer.

So here's just three examples from the Gemini diffusion model, which is what are the first 100 digits of π? This is actually 100 tokens. It looks like a short response. It's actually 100 tokens. And it only takes four steps to do that.

Because the model, it's an easy prompt, it's an easy response, because you've just memorized the 100 digits of π. You can just output it. Whereas, you know, for an autoregressive model, at the same time, it would have only done four tokens.

So that's a very easy one. Slightly more challenging is to write a little bit of code. So that takes, you know, 18 forward passes to generate FISBAS. And then something more complicated is explaining quantum mechanics in a single paragraph, and that took 31 denoising steps.

It just took its time, just for whatever reason, decided to spend longer on those ones. And so the model naturally gets to decide, I guess, to determine when it's going to finish and return the response. And typically, we see that harder evals take more time.

So this is a year ago now, so the evals are kind of old-school. But, you know, on the harder end was GPQA Diamond, which for the model size we were targeting, it was quite a hard eval, and that took a long time for it to respond to those ones.

Whereas on the other end are like MBPP, which are mostly basic Python programs. It was very easy to, you know, it took very little time for it to respond to these ones. And this was entirely determined by the model itself.

Just easier problems, easier prompts, it could respond to quickly, and harder ones, it decided itself to spend more time reasoning. Okay, so that's another property, which is the dynamic and adaptive computation. Lastly is the kind of fast in-place editing.

So diffusion models in general have this very nice property where you can, so take an image, like in this example, you know, cut something out of it or whatever, or give it a little prompt, and it'll fill it in, and you can use the context that you haven't cut out to fill in the piece you've cut out correctly.

And so you can use that for, like, clever image editing, things like that. And the reason it can do that is because, you know, it's not autoregressive. There are autoregressive image generators,right, which go left toright, top to bottom, like raster order, generating pixels.

You know, but diffusion doesn't work like that. It'll just see the entire image and then start to denoise it. And because of that, because it gets to see every pixel, every pixel gets to see every pixel, it can fill in the missing information and do it in a way that's kind of consistent, whatever prompt you're giving.

So we can do something similar. So I have a couple of demos here. Can you see that? Yeah, so this is just some code, and you say, "There's a bug in this code. Can you fix it?" And it'll just make the edit in the correct place.

Like, it won't, you can barely see that, but it's just, you do a little fix here of the indices. And you can say things like, "Can you add documentation?" It'll go in. It's not just one by one generating all the tokens.

It's doing a clever editing procedure to actually fill in the correct edits here. You can do that with, you know, more general text. Like, you can take a story and then say, "Add a middle paragraph." And because it can see the first and the third paragraph, it can fill in the paragraph in a way that's consistent with the rest of the story, for instance.

And this is just in-place editing, basically. Okay, so that's, those are some of the advantages. Don't have a lot of time. The biggest advantage, like I mentioned, is this low latency. And, you know, we really lean into that.

And I just want to show you a couple of demos of some of the things that people internally have built to kind of show what the advantage of low latency can give you. So it's not just the same thing faster.

It can really unlock some new, some really new applications. So in your own work, you're all AI engineers. It'd be interesting to see when the next diffusion model comes out from our team, what the low latency could unlock and kind of what new applications can be built.

So here are just some demos. So this is Wikipedia. Let me just pause it, actually. This is Wikipedia, where everything is generated on the fly, even the HTML. So this is actually being, whoops, this is actually being generated by the model on the fly.

So there's a webpage with the HTML and the text and everything being generated on the fly. So it looks like regular Wikipedia. And when you click on it, the latency is low enough that it can just fill in the page as if it was a real Wikipedia page.

So that's kind of Wikipedia generated on the fly by just a very low latency model. We have a similar thing where we did it for Reddit. So now all the responses to your posts will be by bots. They weren't already.

And it's generating fake comments. So the Gemini diffusion model was not an image-generating model. So this demo links in the, which was our state-of-the-art image generator model at the time, which I think was Juno, wasn't it? It was before Nano Banana.

So it's the two of these models working together to fill in the webpage. So the image generation model is a little slower, but you can see that it's like, you know, you can invent any Reddit you want, sharks, in this case, and it'll generate the page with the text.

The images follow a minute later, and then you can interact with this website as if it was a real website with, you know, real users and so on. Just being all the comments, all the images, all the HTML, everything is being generated entirely on the fly here.

This is being generated by the model. I love this one. This is my favorite one. This is an operating system also being entirely generated on the fly. So every click here is generating the next page of the operating system.

So it looks like a real operating system, but it's all being generated by the model on the fly, responding to every click. So every time you enter, like, the README, it generates the text, but it also generates, you know, the web pages you can do, you know, go back to desktop and so on.

I think this, yeah, okay. And then this is a demo from someone on Twitter who used

the Gemini diffusion API. Or sorry, not the API, the webpage to do some vibe coding with his voice. So I really like this one.

Create a to-do app.

Add 10 random to-dos.

Allow to-dos to have a completed state. Mark four random to-dos as completed.

Allow me to sort to-dos by name and by state.

Allright, let's see if this works. Sort by name, sort by state. Testing, enter. It was added to the bottom. We'll try deleting a few. Let's add one. Everything's working. Please convert this to dark mode.

And this was literally 15 seconds of work.

Okay, so that was, yeah, that was someone outside of our team, so he could say that. Yeah, vibe coding by voice. But just in general, I think that, you know, low latency models can really unlock some new experiences for users and new products.

And so we're excited to see what people will do when the next generation comes out. Okay, and on that note, thank you very much.

Questions? Yeah?

Yeah. Do you use the same corpus of text to pre-train? Diffusion models as for the autoregressive models. Also, how does it work for reinforcement learning for human feedback? So do you use the same data too?

Yeah, yeah. So, yeah, we use all the same data, yeah. I mean, the algorithms have to change a bit, but we use all the same data, yeah.

And also maybe, can you distill these models? Are there any techniques for doing that?

You can distill them, yeah. Yeah, I'm not sure if there are any published ones.

Can you share more about it?

I don't think so. Maybe there's some externally, but.

Will we have in the near future, like, general availability of these?

There's a, yeah, so we're going to release something soon, yeah.

Cool.

Yeah, yeah.

Yeah, so as you've scaled the model, like, more parameters for it, have you find, like, the adaptive computation increases, like, the number of steps increases with, like, bigger models for even, like, easy prompts or not?

Yeah, so there's a, the bigger models tend to require less steps for the same output. So they kind of, you know, even if the model is getting bigger and the flops per forward pass are getting bigger, they tend to reduce the forward passes they need.

So it's kind of a, you kind of have some sort of a diminishing cost of serving even the biggest models, yeah.

So the next question, do you, how do you price the, if you're going to price it? Just tokens doesn't matter for diffusion.

I don't know. We haven't got to that yet. I'm a research scientist.

Yeah.

How do you define the size of the answer? Because for images, you expect that I have this image and I have the output frame. How do you do that in coding or text?

So there's a few different ways. The easiest way is to just fix some window length and then just iterate on that. So it's like autoregressive, blockwise autoregressive. That's kind of the standard way to do it. But you can have, like, a head that will predict the length of the response and stuff like that if you want it, yeah.

But you have to fix the outcome in order to diffusion model to work.

No, it still can generate unlimited text, but it's just, if you fix a window length, then it just does that window length autoregressively if it needs to generate many, many windows of text. Yeah.

I'm not sure you mentioned this, but when you gave that example of the denoising steps being different ranges for different, you know, difficulties of problems.

Yeah.

Could you ahead of time set, like, a limit on the denoising steps that you require for a problem? So you can almost just understand what your latency is going to be ahead of time.

Yeah, yeah. These are all with a limit, but they just finish earlier than the limit.

Allright, okay.

Yeah.

If you had multiple windows, like you just said going through, can they then go back and attend to previous windows? Is that kind of, as the window goes, it's?

Yeah, I mean, you could potentially, but for us, we just set it in stone and continue.

Right, and you just move on.

Yeah, yeah. Lots of questions here.

So you mentioned it's can be slow on larger batches. Would it make sense to, like, a hybrid between autoregressive and diffusion, where it's autoregressive not like a token at a time, but like a chunk at a time?

Yeah, yeah, that's how it works, yeah.

That's how it works, okay.

Yeah.

So it starts to pre-fill with autoregressive style and then it goes into diffusion mode, or?

Oh, no, so it's pre-fill is the same. It's just you've got some context and you pre-fill. And then after that, the generation step is typically in blocks of some fixed size, like 512 or 1000 or 32 or whatever you want.

And then that's autoregressive.

Yeah.

I get that you're doing the denoising inside, like, a bidirectional embedding space. But how do you, what is the process of getting back to token IDs? Is it like latent diffusion instead of?

So you can do that. The easiest way to do it is to just have, like, logits at the top, just vanilla prediction head at the top.

It's after the denoising model, like that.

Well, so this is all, like, discrete diffusion,right? So it's always tokens in, tokens out. Like, if I show you the.

For every step.

For every step, yeah. So if I go back here, it's always like a discrete corruption process, and then you fill in a discrete token back in. So it's always, so you're always in a discrete space. But you can do it in latent spaces, but, and people have done that, but most of the text diffusion models and literature is discrete diffusion today.

Yeah.

Is there only a Gemini diffusion?

There might be, one day. Yeah.

Yeah, maybe, what's your outlook? Like, do you think one of the architectures will win in the future, or will they have different use cases?

I think for now they have different use cases. So, you know, if you think about what a low latency model provides, that's like, that's worse throughput. So what's that trade-off? Is on-device applications. So we are in a couple of on-device applications already within, like, the alphabet ecosystem.

So robotics, things like that, where you want to run a model on the device itself, so your phone or your robot or whatever. And then, but you want it to be low latency, and you're not batching with thousands of other queries, like Gemini being served in the server side.

So you want the lowest latency model. Quality isn't really any, they're the same quality, basically. So then you may as well pick the low latency one, because you don't have the throughput concerns.

What do you mean, like, in the future, like, can we get quality up to the bar with the current, like, proxy models?

Quality isn't the concern. It's the throughput for serving in a big batch setting. Yeah, yeah. Yeah.

So currently, you are getting more or less the same quality of runtime models in your research?

Yeah, yeah.

Okay. So what about reasoning processes? Because we couldn't really verify what we wanted to do in this.

Sorry?

What about reasoning tasks? Is it possible to verify?

Yeah, yeah. Yeah, you can do RL, yeah. You just need to change the algorithm, but you can still do it. Yeah.

Yeah.

Is there much scope for combining diffusion and autoregressive text models?

Yeah, you can do that. Yeah.

How's your team done that?

Yeah.

Yeah.

Can you, like, initialize with a smaller model so you don't have to start from noise? So you can have to do less forward pass?

The problem with that is, if you train it with noise, it expects noise. You'd have to train it with a small model as you, like, the output. That just adds complexity. So we don't usually do that.

Yeah.

Is there a paper, like, a certain version of that where I can actually go and implement and learn more about it?

Well, not from our team, but there's a bunch of literature out there, yeah.

Okay.

Yeah.

Which are respectable for both?

Yeah. You get the idea, I think.

Yeah.

Yeah.

Any other questions? There's a lot of questions, though. I've gone through them all. That's good. Oh, one more, okay.

Allright, this is a random one. What happens if you ask it to generate noise?

I don't know. I don't think we've tried to do that.

Probably would work. We'd probably do something. Yeah.

Cool, okay. Thanks, everybody.