Diffusion for Text: Why Mercury Could Make LLMs 10x Faster

0:00We were matching the perplexity, but we were able to be like 10 times fast.

0:03Stefano Ermon:That was super exciting to me, and I really wanted to see what happens if you train something bigger than a GPT-2 model, possible to build something commercially viable. And that's why I started the company to scale things up. The arithmetic intensity of inference workloads that we have today with an autoregressive model is very bad. The utilization is very low, and that's why people are building massive data centers or or even building custom chips, AI inference chips that are better suited for that kind of work. Basically, if you can generate more tokens per second, what this means is that for the same amount of hardware for the same number of GPUs, you can produce more tokens.

0:38Stefano Ermon:And so the cost per token is going to go down. And that's why we're able to serve our models much more cheaply than what you could get because we make better use of the existing hardware. So now the Mercury models that we have in production are significantly larger. They've been trained on more data. that's going to enable Mercury models to be even smarter. It's going to have much better planning and kind of like reasoning capabilities. And so that's going to enable a lot of agentic use cases that people really care about. They're going to make them really, really fast. Welcome, humans, to the Neuron AI podcast.

1:10I'm your host, Corey Knowles, and I'm joined as always by the man who can turn a GPU benchmark into a bedtime story, Grant Harvey. How's it going today, man? It's going great. don't put me on the spot to do that right this moment, though. I'd have to think of some mechanics there. Oh, well, here in just a few, we're going to be joined by Stefano Armand, Stanford University computer science professor and the founder of Inception Labs that created the mercury diffusion large language models. But first, Grant's going to share us a little context before we get in there. Yeah, so image diffusion models work in an entirely different way than the next token predicting GPT models.

1:53So we've invited Stefano today because he's taken that same technology and applied it to LLMs. And it has the potential to transform how AI is used in all types of settings from agents to complex enterprise workflows. Excellent. Well, before we bring him on, I want to take just a quick second to show you Mercury in action, because I think seeing it really matters and will keep you really interested. you'll understand why you need to be watching this video. So what you see here on my screen, this is the Inception Lab site. And if you go up top, you can go to Mercury Chat. And down here, I'm just going to grab one of their suggested prompts.

2:27And I love this one. Simulate a roundtable discussion between Einstein, Ada Lovelace, and Alan Turing. Now, you want to make sure you click this diffusion button because the diffusion button gives you the visual of how cool this is. So watch this. All right.

2:47watch how this works whoa and if you go there's so much cool into the typewriter effect of the ai is that not insane right yeah yeah yeah that's awesome it's really cool too when you do this like when you build a game in html5 like how quickly it can make you something like pong or or you know any kind of like 2d game it's it's amazing yeah it is it is well uh before we get to this interview, please take a quick second to like and subscribe to the channel so we can keep bringing you the most interesting people in tech and AI. And with that, welcome to the Neuron, Stefano. It's great to have you.

3:24Stefano Ermon:Thank you. Pleasure to be here. Good to see you again. Excellent. Well, we're so excited to have you on. As I mentioned before we started, we chatted in Vegas, did a short interview, and I've really been looking forward to this because I had so many questions when I walked away still that I was like, oh, we've got to get him on. We've got to get him on. So I guess to start, would you mind kind of explaining Diffusion in a fairly simple way for viewers who maybe aren't familiar? So Diffusion is a type of generative AI model. It's the kind of model that is commonly used to generate images, video, music.

4:03And you're probably familiar with the, you know, chat GPTs or Gemini's or Claude, where

4:08Stefano Ermon:kind of like see the models generate text kind of like left to right one token at a time. A diffusion model works very differently in the sense that it generates the full object from the beginning and then it refines it by kind of like fixing mistakes making it sharper making it look better and better and it's a very different kind of like solution that it's more parallel in the sense that the neural network is able to modify many components of the image or the text at the same time and that's why diffusion models tend to be a lot faster than traditional autoregressive models that kind of like work left to right one token at a time okay right right and how how is it that they're they're actually like reasoning over the original version that they they create like how like how do they know that the first version isn't good yeah so it's a great question yeah that's a great question and it really it really stems from the way the models are trained.

5:04Stefano Ermon:A traditional autoregressive model like a GPT model is trained to, there is a neural network and it's trained to predict the next token, the next word. And that's how you use it at inference time. You give it a question and then it will try to predict the answer left to right, one token at a time. The fusion language model, it's trained to remove mistakes, fix mistakes. So you kind of like start with clean text or clean code. You artificially add mistakes and then you train the model to fix those mistakes. And that's how the model is also used at inference time. You start with kind of like a full answer and then you refine it.

5:41Stefano Ermon:And so it's a very different way of training the models. It's a very different way of using the models at the inference time. And it happens at the speed of lightning. It's very fast. It's very fast because the key thing is in an autoregressive model, you get one token at a time. You have to process a massive neural network with hundreds of billions or trillions of parameters. And at the end, you only get a single token. Yeah, very inefficient if you think of it that way. You still need big neural networks, but each forward pass, you need to still evaluate the whole thing. But then at the end, you are able to modify more than one thing.

6:20Stefano Ermon:And so as long as you don't need too many denoising steps, too many diffusion steps, this can be really, really fast. Wow. Stefano, what did you... I know you've done a lot of research on this topic along the way. What did you see that convinced you diffusion for text was commercially viable now? Yeah, so it started... I mean, I've always been passionate about diffusion models. The kind of like original idea came out from my lab at Stanford back in 2019. back then pretty much all the image generative models were based on GANs generative adversarial networks which were very difficult to train, very unstable there is this kind of game between a generator and a discriminator it's a pretty complex kind of model to train and there is all kinds of issues in scaling that approach up and we came up with this alternative approach of training the model to remove noise and and then generating in a course-defined way.

7:25Stefano Ermon:And we showed that it was working much better. And eventually that took off and everybody kind of switched to diffusion models for image, video generation, mid-journey, SORA, stable diffusion. They were all based on those original ideas from my lab at Stanford. And since back then, I kind of tried to see, can we get this to work also on text and code generation? and it took a few years to figure out how to do it properly. But then in 2024, we kind of had a breakthrough. We figured out how to adapt the math and the underlying sort of like algorithms from continuous spaces to discrete spaces like text and code.

8:07Stefano Ermon:And we had some really promising results at the GPT-2 scale. So at Stanford, in university labs, you don't have access to a lot of GPU, a lot of compute. And so the largest model we were able to train was a GPT-2 sized model. But basically what we did is we took a GPT-2 sized model and we train it as a diffusion model and as an autoregressive model on the same data. And so what we found was that the quality was the same, so you were able to actually, if you think about perplexity, which is the way people usually use to figure out how good the model fits the data, what's the quality of the generations you get.

8:46Stefano Ermon:we were matching the perplexity, but we were able to be like 10 times faster. And so that was super exciting to me. And I really wanted to see what happens if you train something bigger than a GPT-2 model. Is it possible to build something commercially viable? And that's why I started the company to scale things up. And you've been scaling up from GPT-2 caliber since then, right? Since then, yes. Now the Mercury models that we have in production are significantly larger. They've been trained on more data. There is a lot of engineering work that went into post-training the models and making sure that they would be useful for tasks that people care about, like commercial kind of use cases of LLMs.

9:27Are they kind of similar in size? Should we be thinking about them the same way? Oh, like parameters and such? How do we compare them to a traditional language model?

9:40Stefano Ermon:Yeah, so the models are still fairly large. in terms of the number of parameters. We're still using actually similar architectures. So under the hood, it's still a transformer. Okay. So what we found is that that kind of architecture actually works pretty well. Also for, you know, as a backbone for a diffusion language model. And so, you know, we've kind of used what we knew worked well and is well supported by existing, you know, frameworks and open source code. And so, yeah, the neural networks are not that different. It's just like they're trained in a different way and they're used in a different way at the inference time.

10:23That's so interesting. I didn't even realize that, that essentially you're still looking at transformer technology, just not wrapped in an autoregressive approach, right?

10:33Stefano Ermon:That's right, yes. And I think perhaps that's actually suboptimal. I mean, people have kind of converged on transformers as being a really good architecture for autoregressive models. I mean, these days people also use them for diffusion models, like people use diffusion transformers. So it's kind of like an architecture that it's widely used across different modalities, across different kinds of generative models. but it's possible that there might be better architectures that shine even better once you change the generative model, it's no longer autoregressive so I think the design space is different, I think there is a lot more room for doing R &D and coming up with further improvements just by kind of matching the neural network architecture to the training objective and to the inference kind of computations that we do right that makes sense very cool when we talk about parallel generation what is parallelized are we talking about tokens spans edits is something i maybe don't know about yeah basically what's parallelized is that the the network is able to essentially modify multiple tokens at the same time wow and so that's kind of like what you were seeing i think if you tried our website and you kind of like see the animation of how the diffusion model works, you're going to see that it constantly changes the answer and it's not one token at a time.

12:06Stefano Ermon:Many things get changed at the same time and that's what makes it more parallel. It makes it much more suitable to GPUs. GPUs are built to process many things in parallel. We're going to like apply the same computation across across different data points effectively. And it's the kind of computation that we do when you sample from an autoregressive model does not map well at all to a GPU. It's a very memory bound kind of computation where you're going to spend most of your time moving around weights from slow memory to fast memory where you can actually do the computation. So the arithmetic intensity of the inference workloads that we have today with an autoregressive model is very bad.

12:54Stefano Ermon:The utilization is very low. And that's why people are building massive data centers because it does not, or even building custom chips, AI inference chips that are better suited for that kind of workload. We just saw a model drop a couple of days ago where in order to get it two and a half times faster, the cost is 6x. Right. Because it's a sequential computation. You cannot generate the third token until you've generated the first and the second. And so it's just a structural bottleneck. There is no way to parallelize it because there is sequential dependencies across the computation. And so you can't process something into the future until you've generated everything before it.

13:38Stefano Ermon:And so there's just no way to parallelize that. It's kind of funny. Basically, what you're talking about is like instead of scaling the amount of GPUs widely, it's like let's scale the amount of tokens that we're actually dealing with at a certain time. And then you can do less on a GPU. That's awesome. Exactly. Exactly. Like it's shifting from a memory bound regime to a compute bound regime where you basically are bounded by the number of flops that you have available on the GPU, which is a much easier quantity to scale. Like if you were a chief manufacturer, it's a lot easier to add flops that do increase memory bandwidth.

14:13So why wouldn't all the big labs like immediately switch to this? Like if it's that much more efficient, because I guess they have to train it. They don't have the same skills that you do. What's your thought?

14:22Stefano Ermon:Yeah, some of the labs are very entrenched in a certain stack. And so there's a big cost if they were to switch to something different. There is quite a bit of secret sauce involved in terms of like, what is the right way to train these models? what is the right way to even just sample from them. It's not as obvious as, okay, generate one token, append it, generate the next one, append it. There is not much you can do on the inference side if you have a traditional autoregressive model. But on a diffusion model, the design space for inference algorithms is much broader. And then there is also the issue of kind of like the MLCs level.

15:04Stefano Ermon:If you think about actually serving production workloads, there is a decent amount of open source and closed source, of course, solutions for ultra-aggressive models. Things like VLLM, SGLang, TensorRT, there are pretty mature serving stacks for ultra-aggressive models. For diffusion models, it's much earlier. We have our own stack, but it takes a significant amount of work if you want to figure out how to actually make things efficient in practice on real-world GPUs. And there's all kinds of optimizations that you can do at the system as well. Well, I've got to say, at a dollar per million output tokens, you seem to be doing okay with that.

15:53Stefano Ermon:For sure, for sure. That's the thing. I keep thinking when we talk like GPUs and what it takes to do the autoregressive approach, you know, this kind of, in a lot of ways, could be a smart way to sort of side skirt things like the current memory supply issue, the need to go acquire brink trucks of money and back them up at Jensen Wong's patio door. you know I really think that this is an interesting approach at a prime time for that yeah basically if you can generate more tokens per second what this means is that you know for the same amount of hardware for the same number of GPUs you can produce more tokens and so the cost per token is going to go down and that's why we're able to serve our models much more cheaply than what you would get if you were to use traditional ultra-aggressive models because we make better use of the existing hardware and so the costs are actually significantly lower.

16:58That makes sense. That makes sense. So how does this behave with long context? Are we looking at it getting specifically any more expensive or is parallelism playing more of a role as they grow?

17:18Stefano Ermon:Yeah, so that really depends more on the architecture than whether it's a diffusion model or an autoregressive model. Right now, as I mentioned, we're still using self-attention, which unfortunately scales pretty poorly with the context length. So I would say there is no difference. It's not better, it's not worse than an autoregressive model as you think about longer context. Our models are supporting roughly 100k tokens of context length. we could potentially scale it up more. Again, it's not something that is very different. If you think about an autoregressive model versus a diffusion model, it's more a function of the underlying architecture.

18:02Stefano Ermon:And in fact, we can actually use alternative architectures that scale better with respect to the context line, like state-space models or other attention variants that are more efficient. We have some preliminary results, So everything is compatible with different kind of backbones, but not in the production process at the moment. That's cool. Yeah, I was wondering, like, I've always wanted to ask a researcher this. Have you seen anything over the past year or like six months even that has lit you up in terms of like, oh, this could be a good alternative for that memory context problem? Or are you like, we still haven't seen anything that's even close?

18:45Nothing particularly.

18:47Stefano Ermon:I think it's just like a fundamental problem for which it's going to be hard to get a real breakthrough. There is just inherent trade-offs. I think of them in terms of sufficient statistics. What do you store about your past and how do you keep track of... You want to remember the things that are useful. You want to discard the things that are not useful. And that's just fundamentally a hard problem. Like there is no, there's always, you know, there is some kind of no free lunch involved where ahead of time, you don't know what you should remember and what you should discard. And some things are going to be useful for something and they're going to be not useful for something else.

19:25Stefano Ermon:And so I think it's a fundamentally very difficult problem where you have to make trade-offs. Although I will say with that equation, being able to work in parallel, I think makes you maybe perhaps make those trade-offs or you can make some of those calculations more efficiently. Would you agree with that, if that makes sense? Yeah, I think it changes a little bit the design space in terms of how many flops you have access to and how memory bound you are, but not fundamentally. You still need to process and you still need to be able to look at all the context, all the past information to be able to generate good quality answers.

20:12Stefano Ermon:whether you do them one token at a time or you do them in parallel you kind of have to look at the past and so there is something pretty fundamental there you know i would say probably outside of most you know maybe enterprise and software applications but you're dealing with what say the average worker uses 100k context is is plenty for most things you can really do a lot in that range. Yeah, I was kind of wondering how coherence plays together by not going left to right. And I guess that's me thinking of it through how my human brain works. Yeah. So essentially, there is an element of error correction.

20:55Stefano Ermon:And so the models are trained to fix mistakes. And then they constantly revise the answer. And so initially, the answers are not coherent and then they get increasingly better as you throw essentially more compute they also think of it as a as another dimension over which you can scale compute at test time so test time inference test time compute to trade off quality for speed or cost and so that's kind of like the fundamental trade-off that is exposed by a diffusion language model it just provides you an axis to control quality versus speed. And there is a fundamental trade-off. Like if you want to be, you know, if you don't want to do too many denoising steps, too many passes over the output, the quality is not going to be as good as what you would get if you refine it many, many, many times.

21:50Same as with if I'm running, you know, stable diffusion and CompUI or something and you're doing your denoising runs.

22:00Stefano Ermon:That's exactly the same. Exactly, exactly. So, you know, even in the context of image and video generation, you can usually control the number of denoising steps. The more denoising steps you take, the higher the quality. But of course, the more expensive it becomes, the more time. How does this work in a Gentic context then? Like, how would it be different? Like, let's say, like in something like Cloud Code, like where, you know, it's going out, it's doing a bunch of tool calls, it's reasoning. How does it play out in that scenario if different? So it's related in the sense that, you know, it needs to output something.

22:38Stefano Ermon:And then, you know, if there is a tool call involved, then it kind of you need to essentially wait until the result of the tool call comes back. But it's not too different. Like essentially, we're able to serve the models in an open AI compatible way. We support tool calls. And so people are already using our models in, you know, a variety of authentic frameworks, including some of the open source ones. I think people have figured out ways to also use it in cloud code. I think there are some wrappers that allow you to use other models. I've been eyeballing how to connect it to my OpenClaw. That works.

23:19Oh my gosh. Have you done that actually?

23:22Stefano Ermon:Have you played with OpenClaw? I haven't, but some of my team members have and they've used Mercury. So I know for sure it works. Excellent. That's cool. Excellent. I haven't tried to connect it yet. I've still been stumbling my way through, but it's very much on my list because I keep thinking of like, I have such a variety of tasks. I'm having this thing do for me now. And I think some of these, that super speed would be so huge on. And so much more affordable, frankly, than running plenty of the other models I have access to as well. so it's definitely on the radar and I feel like in many agentic flows there's a lot of value for this approach yeah completely agree I think at that point speed of interaction with the environment becomes the key bottleneck you really want to be able to use the tools that you have access to and collect the information needed to solve the task and the faster you interact the faster you can take actions, collect feedback reason, decide what to do next the better the experience the better the model is going to work and so speed I think we're seeing it also from other labs people are pushing more and more okay let's make the models faster, faster, faster because that really improves the user experience Yeah because you kind of have to scale that at the same time you're scaling intelligence otherwise I feel like what you're going to wind up with is a 48 minute wait for everything you ask for like you kind of have to lift them together Yes, yes um you mentioned earlier so like one of the the um unique things and benefits of being the person who basically made this um uh is that you know all of the tricks of how to you know run it and deploy it um and you know there's a lot of frameworks that support regular llms but maybe not for as many frameworks or you know you have an in-house framework for dealing with diffusion would you ever make like an open version of that like do you want everyone to go through you like What's your business plan there, I guess?