Building Generative Image & Video Models at Scale

Talk transcript

1. 0:00

   [upbeat music] Cool.

2. 0:15

   Um, I'm gonna get started. The sound's all right? Okay, yeah. I can, I can hear myself. Okay, great. Uh, thanks for coming, everyone. Uh, I'm, uh, Sander. I'm a research scientist at Google DeepMind, have been there for over a decade now.

3. 0:30

   Uh, I'm on the generative media team, so we work on models like, uh, Veo and Nano Banana. And I wanna-- this talk is kind of a bit of a behind-the-scenes perspective on, like, uh, everything that goes into training these, uh, models at scale.

4. 0:45

   Um, and so it's gonna be focused on, uh, diffusion models because that's sort of the, the mo- the modeling paradigm that at least nowadays people tend to use for generating audiovisual data, which is a bit different from, uh, the language modeling situation where we primarily rely on autoregression, although of course, uh, not exclusively these days.

5. 1:03

   Um, so yeah, I have, um... it's kind of a whirlwind t- tour of everything that goes into it. So I have eight sections, which is quite a lot to get through.

6. 1:10

   Um, they're not all equally long. They-- I have more to say about some things than others, so I have more to say about the modeling and the, and the sampling side of things, for example.

7. 1:19

   But I'll just try to, uh, cover a little bit of everything and then hopefully leave, uh, enough time for questions as well at, at the end, uh, about, about any of these things.

8. 1:27

   So the, the topics I wanna cover are, uh, data curation, uh, representation also because you don't necessarily just shove pixels into a transformer, right? You might use a more, uh, advanced representation, uh, to sort of simplify things.

9. 1:41

   Uh, modeling is gonna be sort of about the core mechanism behind diffusion models. I'd just like to explain that, uh, in an intuitive way, and I think that's, that's kind of useful to understand why these models are able to work as well as they do for, for audiovisual, uh, data generation.

10. 1:54

    Uh, we'll talk a little bit about architecture as well, like the, the core neural network architecture that's, that's behind all this. Uh, and then, uh, very briefly, uh, some things about training at scale.

11. 2:03

    I'll talk a bit more about sampling as well because, uh, diffusion models have sort of a very flexible, um, uh, sampling regime. Like there's, there's, there's sort of more stuff you can do with diffusion models than you can do with autoregressive models, so I wanna talk a little bit about that.

12. 2:17

    Uh, we'll talk about distillation, which in the context of diffusion is, uh, not so much about making the model smaller but more about reducing the number of steps that you need to get a good sample.

13. 2:27

    And then finally talk about control, like the various control signals that we wanna use to actually make the models do our bidding and m- make, make them do what we want them to do.

14. 2:35

    All right. Uh, before I start with that, I have a blog. Um, you don't have to remember all these individual URLs. Just go to sander.ai. You can find all my blog posts there.

15. 2:44

    They're mostly about diffusion these days, but in general I like to talk about sort of the, the intuition behind how generative, uh, models, uh, actually work.

16. 2:53

    Cool. So let's talk a little bit about data first. Um, I think, uh, it's really important to emphasize just how important data curation is in training these, these large-scale models.

17. 3:04

    Um, it's really, really essential for high-quality results. And like, for me, coming from a research background, it's actually something, like when I was doing my PhD, it's not something that's incentivized.

18. 3:13

    Like you're not incentivized to look at your data because you're actually incentivized to use a predefined data set that everyone else uses so that you can compare your results to the state-of-the-art.

19. 3:23

    Um, and so that's sort of something that we are sort of collectively having to unlearn in this era 'cause you really just have to look at your data to the point where I think time spent on improving the data is sometimes a better investment of that time than actually sort of trying to tweak the model and trying

20. 3:37

    to make the optimizer better or things like that. Uh, I think this is still sort of an un-- even, even today it's still sort of an underrated aspect of, of, uh, training these models.

21. 3:46

    Um, but yeah, I'm, I'm, I'm, I'm not gonna say too much more about that because this, it's a topic that you won't find a lot of publications on, and it's also a topic that, um, I won't be able to share a lot of details on because obviously it's, it's also part of the secret sauce of, uh, what

22. 3:59

    makes the models good. Um, so let's talk about representation next. So, um, the sort of typical digital representation of audiovisual data, or sorry, of visual data I should say, is, uh, pixels, right?

23. 4:10

    Grids of pixels. You get a two, 2D grid for images. You get a 3D pixel grid for, for video. Um, and you could actually just shove that into a diffusion model as is, and that's what we did at the start, you know, when diffusion sort of, um, started, uh, being a thing.

24. 4:25

    That's what people were doing. We were sort of directly training on pixels, and it worked remarkably well. Uh, but that's not actually what people do nowadays because once you sort of start scaling up the, the size of the objects that you're modeling, like the resolution or the duration in the case of video, these, uh, pixel tensors in

25. 4:41

    memory get very, very large, right? You, you have like... Okay, let's, let's say you have, um, a 1080p video, right? Th- thirty seconds of 1080p video at thirty fps.

26. 4:52

    We're talking many, like several gigabytes of, of data, right? To store that in memory. That's, that's one training example. So that's just not feasible. So what we end up doing is we actually use compressed representations, right?

27. 5:04

    But we're not just gonna use, uh, sort of pre-existing compressed representations. You know, there's, there's many sort of like video codecs out there. Uh, you know, there's like the JPEG standard for, for images.

28. 5:13

    We tend not to use those representations because those are really focused on making things small, and in the process, they actually obscure the structure of the data to a degree that, uh, generative modeling becomes really hard.

29. 5:25

    So what typically happens in this context is we actually learn our own compressor. We-- and, and we do that based on autoencoders. So this is like the, the top, uh, the top row on this slide is basically a neural network that consists of an encoder and a decoder.

30. 5:37

    Uh, and the task of this neural network is just like reproduce your input, right? You put an image in, you want the, the same image out, uh, but you're, you're sort of forcing it through a bottleneck in the middle, and that bottleneck is gonna learn what, what's typically called a latent representation.

31. 5:50

    So it's like a more compact, uh, representation, uh, of the data. And that is then the representation that we're actually gonna use as a basis for, uh, diffu- like for training our diffusion model.

32. 6:01

    So in the second stage, we take these latent representations, we extract them from our, uh, inputs. I, I usually use images as examples, uh, because obviously that's easier to show on a slide, but you can imagine this is a video as well.

33. 6:11

    It's, it's basically the, uh, the exact same procedure. Um- Uh, so we extract the latent representations, and then we train on top of that a generative model, which could actually be an autoregressive model or a diffusion model.

34. 6:22

    Um, sort of doesn't really matter, but people have found in general for audio-visual data, um, uh, diffusion sort of has the edge in terms of, uh, being able to get good results with, uh, certain parameter budget.

35. 6:34

    So that's what we mostly tend to use nowadays. Uh, and then once you've trained your, your generative model in latent space, of course, once-- if you sample from that model, you're gonna get a sample in latent space.

36. 6:44

    Uh, so it's a compressed sample, so you still wanna, uh, use the decoder that you trained before to bring that back to pixel space, um, as well.

37. 6:53

    And so what does that actually look like in practice? So I have this sort of concrete example here, and this is actually-- I think these dimensions sort of correspond to what the autoencoder in the original Stable Diffusion, uh, model does.

38. 7:04

    So Stable Diffusion is also a latent diffusion model. Uh, let's say you have, like, a, an RGB image, uh, two fifty-six by two fifty-six pixels. If you represent that as a ten-- as a three D tensor, uh, on the left, so these are the, uh, dimensions that you get.

39. 7:16

    And then when we extract our latent representation with our encoder, so one thing that's worth noticing is it still has the same spatial structure, it still has the same topology, but the resolution is greatly reduced, right?

40. 7:26

    So now we have thirty-two by thirty-two latent grid, and we sort of compensate for that reduction by adding a few extra channels so that we can encode some of the information that would get lost if you were to just, uh, resize the image, right?

41. 7:37

    If you just take the original image and resize it to thirty-two by thirty-two, you lose a lot of high-frequency detail. And these, uh, latent representations sort of try to hold on to that a little bit by adding-- by using these extra channels, uh, to encode that information.

42. 7:49

    Of course, it's still, uh, lossy compression because if you look at the total sort of tensor size of the latents versus the original image, it's still much, much smaller, right?

43. 7:58

    And that's the key idea. That's the point. We want to make this smaller, and especially for video also where you have this extra time dimension, you have even more redundancy, right?

44. 8:05

    You're, you're often able to reduce the size of these tensors that you have to work with by, by two orders of magnitude. And that's the difference between, like, being able to fit your data in memory and, and just not being able to do anything at all, right?

45. 8:19

    Um, but so one thing to, to note about these compression mechanisms is they're actually sort of slightly more primitive than sort of the, the s-- you know, like the H.265 or like the other sort of codecs that people tend to use these days.

46. 8:30

    And that's actually by design 'cause, um, you kind of wanna preserve a lot of the topological structure of the original, uh, pixel representation because, uh, our neural network architectures that we use to build these models, they kind of rely on the structure being present, right?

47. 8:44

    They have these strong inductive biases. And so we, we end up with these latent representations that are still-- that still have the same grid structure as the original, uh, pixels, but sort of at a, at a more, um, coarse grain scale, right?

48. 8:55

    And, and the image that I'm showing here is actually from a paper called EQ-VAE. It's a paper about, like, improving the training of these autoencoders. And they, uh, did a really nice thing.

49. 9:03

    They sort of tried to visualize the latent space by taking, uh, the principal components of the, of the l-- across the latent channels and then, uh, mapping that onto RGB so that you can actually visualize what these are.

50. 9:14

    And you can kind of see, looking at these, uh, latent visualizations, you can still tell what animal is in the image, right? So the latents are not really making abstraction of any semantic content of the image.

51. 9:24

    They're ma-- they're basically just sort of, um, abstracting the local texture and very fine grain structure, right? That's sort of the information that's sort of compressed and that's removed to some degree.

52. 9:33

    And everything else is just preserved as in, uh, the original, uh, pixel space. And that's, that's kind of by design.

53. 9:39

    All right. So once we have our latent representation, we can start doing some generative modeling. Uh, I've already mentioned you could do that either with autoregression or with diffusion.

54. 9:47

    With autoregression, you sort of have to turn everything into a sequence, and then you predict the sequence step by step. This is a very natural thing to do for language, which is already kind of sequence structured.

55. 9:55

    Slightly less na- less natural to do for an image, right? 'Cause an image is sort of a two D pixel grid, so you have to sort of decide an arbitrary sequence order in which to, uh, generate, uh, the pixels.

56. 10:06

    Um, diffusion is kind of a different approach to achieving a very similar thing, which is sort of this iterative refinement, right? You generate something not in one go but step by step.

57. 10:15

    And diffusion does that, uh, in a slightly different way, which is by first defining a corruption process, uh, sort of destroying information in the image or in the video by adding noise gradually, so adding more and more noise.

58. 10:27

    And then you learn a denoiser model that tries to remove that noise. Um, and then you can use that denoiser as, as we'll show, you can use that denoiser to sort of generate images or videos step by step.

59. 10:37

    So, uh, yeah, just, just to visualize what that looks like. So I have an example image here on the left. If you gradually add more and more Gaussian noise to that image, uh, one thing to note is that eventually, when you add a lot of noise, you can't see the image anymore, right?

60. 10:50

    So this sort of destroyed all the structure. But sort of in the intermediate stages, you can see that if I add a little bit of noise, it's sort of removing the detail, right?

61. 10:58

    'Cause for example, you might not be able to make out the whiskers anymore of, of this bunny, but you can still see the silhouette, so you can still tell that it's a bunny, right?

62. 11:07

    So you're sort of losing fine grain detail, uh, initially, but not the global structure. And only when you add more noise, that global structure also, uh, starts to fade.

63. 11:15

    So that's the corruption process. And then I'm gonna try and explain to you how we can use a denoiser that tries to reverse this process to do generative modeling, right?

64. 11:24

    So sort of an, an intuitive view of how diffusion models actually do this. So I, I like to do this with a, with a two D diagram. So images are-- and video are actually high-dimensional objects, right?

65. 11:34

    They have-- If you've sort of flattened an image into a single vector, it's very high-dimensional. Um, I'm gonna reduce that down to two dimensions so that I can actually plot it on a slide.

66. 11:43

    It's a slightly risky thing to do, like generalizing from low dimensions to high dimensions, but in this case, it's, uh, it's actually quite instructive. So I'm gonna try to explain how this backward process in, in diffusion actually works.

67. 11:53

    So we start here with, uh, a clean image from our dataset. So that's what I'm gonna call X nought. Uh, put that in the, uh, bottom left corner on the slide there.

68. 12:01

    And, um, and then I'm gonna add some amount of noise to it. So I'm actually gonna sort of simu-simulate this corruption process and add some amount of noise to it, and then we end up in an arbitrary, uh, uh, step of this corruption process.

69. 12:14

    We're gonna call that Xt. T is like a time step that's, that's indexing the, the corruption process. Uh, and so that's a noisy version of our image in the top right corner.

70. 12:23

    So what's gonna happen in diffusion sampling is we're gonna start from noise, right? But, uh, I'm just gonna, uh, sort of zoom in on a particular step in the process where we've sort of already partially denoised the image.

71. 12:33

    So we might end up at something like Xt, and then I'm gonna tell you, like, how does the, how does the process continue from there. So, uh, the first thing that our diffusion model is gonna do, as I said, it's a denoiser, right?

72. 12:44

    So it's going to try to predict a clean version of the image given the noisy version. Like, it observes the noisy version, that's all it gets, and then we ask the model, "Where could this have come from?"

73. 12:54

    Right? What image could have given rise to this particular noisy observation? And that's actually an ill-posed problem 'cause, as I just said, the noise actually removes information, right? So there's actually many different images that could have given rise to this particular noisy observation.

74. 13:10

    And so what ends up happening is the model just do-doesn't know which one to predict, and it has to make a single prediction, so it's just gonna predict the average of all of them, and that's, uh, why you get a blurry prediction out of a diffusion model.

75. 13:22

    Like, if you just ask a diffusion model to do a one-step prediction, you're gonna get a very blurry image that looks sort of like this. And the way to think about this is that it's not predicting a particular image per se, but rather it's sort of predicting a region of image space, w- the direction that we roughly

76. 13:36

    wanna go in. And so that's what I'm trying to visualize here with these, uh, with these, uh, circles, is that there's sort of a region of images that all po-potentially could have given rise to this noisy observation, and that's sort of roughly the direction that we wanna go in.

77. 13:49

    And so how does sampling work then? Well, we, we predict that direction, and then we take a small step in that direction. We don't make a big jump all the way to, uh, our predicted X nought because obviously then we would just end up with a blurry image.

78. 14:00

    That's not what we want. So we're only gonna take a small step and then ask the model again, basically, right? You can compare this to how, uh, optimization of neural networks works, right?

79. 14:09

    When you're doing optimization, uh, of course, that's happening in parameter space. Here we're ha- we're working in, in pixel space or in, in latent space. But when you're doing optimization in parameter space, you're doing something very similar.

80. 14:19

    Right? The optimizer gives you an update direction. You don't wanna take a step that's too large because actually this update direction is only valid locally, and it's, it's the exact same thing here.

81. 14:29

    Um, one thing that's a bit different from a typical optimization algorithm is that after we've done this small denoising step, uh, in many diffusion sampling algorithms, you actually add a little bit of the noise back, and this is, this is new noise.

82. 14:40

    This is not sort of like the noise that you just removed, but you just add a little bit of new noise. Obviously less than the amount that you just removed, right?

83. 14:46

    'Cause otherwise you wouldn't make any progress. And the reason you do that is often to sort of, um, avoid accumulation of errors along the process. Because obviously this denoiser that we've trained is a neural network, so its predictions are not perfect.

84. 14:59

    So it's gonna make mistakes. And if you keep feeding the network its own mistakes back, um, you're gonna go off the rails, right? And so what-- It turns out that this is actually a useful trick to, uh, to greatly reduce the risk of that.

85. 15:11

    Like, you add a bit, a little bit of new noise, and it sort of obscures the, the mistakes that the denoiser is making. Not all diffusion sampling algorithms use this trick, but many do.

86. 15:22

    Um, and then, yeah, as I said, the process just repeats. So we're now in a new iterate. We're in, in Xt minus one, which is basically a slightly less, slightly less, uh, noisy version of the image 'cause we've partially denoised it.

87. 15:32

    And we're just gonna, uh, ask our denoiser to make a new prediction, and that's gonna be different this time, right? It's still trying to predict where we came from, but now it has a little bit more information because there's less noise in the image.

88. 15:43

    So that corresponds to a new prediction, and that new prediction corresponds to a smaller region, uh, of input space, right? And so that region is just gonna sort of continually shrink as we go along until it sort of collapses to a single point.

89. 15:54

    And at, at that point, we've sampled, uh, an image, uh, from the distribution that we're trying to, uh, generate. Um, yeah, so this kind of continues, a-as I said, like, sort of exactly the same procedure.

90. 16:06

    Denoise, add a little bit of noise back, um, uh, and so on, and so on. And that's how we sample from a diffusion model. Um, so I wanna now take a slightly different perspective to try and explain why doing this particular thing works so well for images and video.

91. 16:22

    Uh, and for that, we're gonna, uh, bring out Fourier analysis. We're gonna do, um, frequency analysis of what's actually happening in a diffusion model. So I have here, um, four images from the ImageNet dataset.

92. 16:35

    Uh, and what I've done is I've taken their Fourier transforms. And a Fourier transform is a, of an image, is a 2D object. Right? An image is a 2D object, so you also get a 2D Fourier transform, and it's a complex valued thing.

93. 16:46

    Uh, so I'm gonna pull that apart into the magnitude and the phase. So the magnitude is the middle row, and then the phase is the bottom row. The phase is kind of hard to work with.

94. 16:54

    We're gonna ignore that for now. We're just gonna look at the magnitude spectrum. And then, uh, to make this even more convenient, we're gonna summarize the magnitude spectrum into, like, a 1D, uh, plot.

95. 17:05

    And we're gonna do that by radial averaging. So we're just gonna sort of take these radial slices and, and sort of, um, uh, see what they are and then just average them all together.

96. 17:15

    And if you do that and you plot that on a log-log plot, uh, here's what you get for these images. So what's very interesting about this, I think, is they look like straight lines, right?

97. 17:25

    And this is a log-log plot. So if you're familiar with scaling laws in language models, this says power law, right? There is a power law relationship, um, in the spectra.

98. 17:35

    And this is sort of a natural phenomenon with, with images and also with video. Like, if you, if you take a photo, you calculate the spectrum, you'll typically get this, this power law relationship between the, uh, the frequencies and then the corresponding, uh, power, uh, or energy of those frequencies.

99. 17:52

    So what does that mean for diffusion? Well, um, if you add noise to an image, here's what that looks like on the spectrum. So, uh, the spectrum of Gaussian noise is basically such that it contains all frequencies in, in equal measure, in expectation, right?

100. 18:10

     So it's sort of flat in the frequency domain. So you can see that here, like, this is, this is a spectrum of actual Gaussian noise, uh, this, this blue line, right?

101. 18:17

     So you get sort of this, this horizontal line. And then for an image, you're gonna get this power-law spectrum, so you get this downward sloping line. Now, if I add both of them together, if I add the noise to the image, I can again calculate the spectrum, and what you get then is the green line.

102. 18:30

     So you can kind of see that it's, uh, following the image spectrum, uh, up to the point where the noise starts to drown it out, right? And then it just follows the noise spectrum.

103. 18:39

     So what does that mean concretely for our corruption process? It means that the more noise you add, the more of the high-frequency components you're starting to obscure, right? So add a little bit of noise, you only obscure the highest frequency components.

104. 18:51

     If you add a bit more noise, then you're actually obscuring more of the lower and lower frequency components until eventually, if you add a lot of noise, you're obscuring all the frequency components in the signal and, and they're, they're completely indistinguishable.

105. 19:03

     Um, and so that observation, I've sort of summarized that myself as diffusion is basically spectral autoregression, right? Because it's essentially allowing you to generate images from coarse to fine, right?

106. 19:16

     You start with the low frequencies, and then you gradually add higher and higher frequencies, and that corresponds to coarse-grained features and fine-grained features. Um,

107. 19:25

     but of course, yeah, this, it's, it's, it's an approximation, right? It's not the same as hard autoregression, but it is sort of... In, in principle, it's, it's doing a very similar thing, just in a, in a representation space that makes more sense for, for image generation, where you can sort of go coarse to fine, which typically means

108. 19:39

     that you can sort of sketch out the semantics of your image before you add all the details, which is a very natural way, uh, to do, uh, image generation.

109. 19:47

     And it also means that when we train these models, we can actually put more weight on the frequencies, on the, on the, on the scales that matter more, uh, perceptually, and that's a, that's a very important aspect of this as well.

110. 19:59

     All right. So, uh, let me now talk a little bit about, uh, network architectures, uh, for these denoisers. Like, it's a, it's a-- The network itself is very simple, right?

111. 20:07

     It's just do-doing denoising. You feed in a noisy image, and you ask it, like, what is the, what is the clean version of this image? Uh, and so people initially started using UNets for this.

112. 20:15

     So UNets are convolutional neural networks, uh, that are sort of designed-- Uh, they were origi-originally designed for, for things like image segmentation, uh, but you can also use them for any sort of image, uh, restoration task or whatever.

113. 20:27

     Anything where the, where the output dimension is the same as the input dimension, basically. Um, and yeah, they worked, they worked really well for, for this particular use case as well, so people sort of adopted them.

114. 20:36

     Early stable diffusion models are, are UNet-based as well. Uh, but then people figured out, actually, can we use a transformer for this? Um, and the answer is yes, of course you can.

115. 20:45

     Um, there's, uh... You know, you, you, you obviously don't use the sort of exact recipe that you would use for a large language model. Uh, you don't use a causal mask, for example, because you're not gonna, uh, you know...

116. 20:55

     It-it's okay for attention to be fully bi-directional in this context, so it's actually slightly more expressive than, than a typical LLM transformer. Uh, but yeah, it works, it works just fine with transformers.

117. 21:04

     And because we've learned so much about how to scale transformers from the LLM side, it just makes practical sense to also use this, uh, for diffusion models 'cause we can reuse a lot of that knowledge.

118. 21:14

     Um, one more architectural aspect that I wanna talk about in the context of video generation is sort of the tension between autoregression and, uh, diffusion. So you can train a fully autoregressive video model, but it would mean that you have to take this sort of height times width times time cube and then flatten it out into a

119. 21:32

     sequence, right, and then generate it token by token. You can also treat this whole, uh, 3D volume of the video, this 3D pixel volume, as a thing that you just jointly noise and denoise, and that's what most modern, uh, video generation models do.

120. 21:47

     So they're sort of adding noise across the whole, uh, time dimension as well as the spatial dimensions. Um, that works pretty well, but there's actually, uh... It's, it's not a dichotomy.

121. 21:56

     Like, it's not a binary choice. There's actually an intermediate, uh, sort of, uh, setup where you do autoregression in time, so you generate frame by frame, but you use, uh, diffusion to generate each of the frames.

122. 22:09

     And that's actually a really nice sort of hybrid setup for, for many applications, um, especially when you w-wanna do things like real-time video generation, for example. You kind of have to do autoregression in time, right?

123. 22:19

     So, like, uh, Genie is a good example of sort of this, this, uh, this compromise.

124. 22:25

     Um, cool. All right. I'll, I'll talk very briefly about sort of what goes into actually training a model. Like, it's, you know, we, we, we scale these things up, right?

125. 22:32

     More parameters is better. Um, I would say they're probably not at the scale where large language models are today. They're just-- They tend to still be a, a little bit smaller.

126. 22:42

     Uh, I'll also, uh, explain a little bit why that might be. Uh, but, you know, we still scale them to, to, to reasonable sizes, and so it's, it's very important to figure out how to do parallelism and, and sharding across many chips.

127. 22:54

     Um, so up to a certain scale, you can rely on data parallelism, right? Just split up your batch across many chips, uh, for training. But at some point, you just need to go into model parallelism as well, so you actually need to start sort of spreading your model across, uh, different chips.

128. 23:07

     Um, we tend to use JAX to build these models. JAX has some really good tooling to-- that sort of does that for you, right? So you, you can, uh, you can use JAX or PJIT to basically ask it to, you know, um, to shard the model in a way that's gonna be, uh, ideally close to optimal.

129. 23:23

     And it sort of automatically minimizes communication between the chips for you, so you don't have to do that manually.

130. 23:29

     Cool. Uh, yeah, I think that's, that's all I'll say about that for now. Uh, move on to sampling 'cause I wanna talk a little bit more about sort of this contrast between stochastic and deterministic sampling.

131. 23:38

     So, um, I already mentioned that some of these algorithms add a little bit of noise back after every step. Other algorithms do not. This has different trade-offs. Um, uh, we'll talk about distillation a bit later.

132. 23:50

     In that context, it's actually really useful to have a deterministic sampling algorithm 'cause you sort of have a one-to-one mapping between initial noise and, and, uh, samples from your, uh, data distribution.

133. 24:01

     Um, but st- as I said, stochastic algorithms can be a bit more robust, uh, to accumulation of errors. So there's, like, interesting trade-offs there. Uh, the main thing I wanna talk about in context of sampling is this idea of guidance, right?

134. 24:12

     This is something, um, it's commonly associated with diffusion models, but you can actually also do it for autoregressional-- autoregressive models. It just seems to work exceptionally well for, uh, diffusion models specifically.

135. 24:23

     So what guidance enables you to do is sort of trade off sample quality for diversity. So if you crank up what's called the guidance scale, if you crank up this hyperparameter, um, diversity of your samples is gonna greatly reduce.

136. 24:34

     They're gonna start looking more alike. Uh, but the quality is gonna, is gonna massively improve to the point where these models are punching well above their weight. And this is what I, what I was saying earlier, why these models...

137. 24:44

     Why I think these diffusion models tend to still be a bit smaller than their sort of LLM counterpart, is that they have this very powerful trick that really allows them to, to punch above their weight.

138. 24:54

     I'll show you some examples in a bit. But first I wanna explain what guidance actually does. So I'm gonna go back to this diagram that I had before, and I'm gonna show you how guidance actually modifies the, the sampling procedure.

139. 25:05

     So we're gonna start the exact same way. We're just gonna, uh, try and predict where we came from. It's gonna be a blurry prediction. But now we're gonna make a second prediction as well, because we're gonna have trained our model with a conditioning signal.

140. 25:18

     So this could be a text prompt, for example, and we can basically sample from the model by giving it a text prompt, and also without a text prompt. So if this is the prediction, is the prediction without the text prompt, then I can make another prediction with the text prompt.

141. 25:31

     And typically, what's gonna happen is that's gonna be a slightly less blurry prediction, because obviously the number of images that this noisy observation could have come from is greatly restricted by what's described in the prompt, right?

142. 25:42

     If the, if the prompt says, "This is an image of a, of a, of a rabbit," and it gives you some details and stuff, then obviously you, you know a bit more about what the original image might have been.

143. 25:50

     So there's a difference between these two predictions, and it's precisely that difference that we're interested in, so we're gonna call this, um, delta. And, and then we're gonna basically amplify this.

144. 26:00

     We're gonna say, "Okay, if this is the delta between not having the prompt and having the prompt, what if we just, uh, scale this up?" And then we, we say, "This is our new direction that we're going to actually denoise in."

145. 26:12

     It turns out this is super powerful. It seems, like, it seems almost childish, like it seems so simple, but actually there's sort of like a Bayesian reasoning type thing behind this that makes this work really well.

146. 26:22

     Um, and then, um, yeah. So this is ac- basically the only change that we make to the sampling algorithm. So everything else proceeds as normal. Again, we might add a little bit of noise back, and then we just take, uh, the steps as normal.

147. 26:33

     So you do need to do two model evaluations per step, but you get a, a, a huge increase in sort of, um, uh, sample quality from that. In fact, let me show you.

148. 26:42

     So, um, so this is a comparison of some samples from a model, uh, without guidance and with guidance. Um, this is from a 2021 paper, so positively ancient, uh, in AI, right?

149. 26:54

     'Cause things move pretty fast. Uh, but this is one of the last sort of papers where they train diffusion models at scale, where they actually show the difference between not using guidance and using guidance.

150. 27:04

     After this paper, everyone just always uses guidance. Like, the- nobody ever turns it off anymore 'cause it's so essential to getting good samples from these models. And the differences you can see here are sort of, uh, you can see that reduction in dive- diversity, right?

151. 27:17

     Like, from the, from the left to the right, you can see the samples from, uh, uh, from on the right are actually much less diverse. But also, individually, they're much higher quality with respect to the prompt.

152. 27:27

     And so you get this trade-off. And if you only want one sample, why would you even care about diversity, right? So you can sort of use this, or you can even, uh, you know, reintroduce diversity in different ways if you wanted to.

153. 27:38

     Uh, but yeah, basically, uh, guidance is sort of a, a no-brainer these days. So everyone just always leaves it on, and I think a lot of people would be surprised at how bad today's models are if you take guidance away.

154. 27:49

     Like, just if you use them, um, as they are. Another example here from the same paper. So this is a GLIDE paper from OpenAI, uh, one of the first sort of, uh, pixel space diffusion models at scale.

155. 28:01

     Cool. Right. Let me talk briefly about distillation. Uh, so sort of to speed up, uh, the sampling procedure, what we're gonna try is we're not gonna try and make the, the model smaller, which is usually what distillation refers to in the context of large language models.

156. 28:14

     But here, what we're actually gonna try and do is, uh, reduce the number of steps that we require. So, uh, sampling from a diffusion model, you could say, is sort of tracing out this nonlinear path through input space, right?

157. 28:25

     I've showed you a few steps before. Now, here I'm sort of trying to interpolate that path with this dashed red line. I'm sort of trying to show, like, this is the trajectory that we might follow through input space during sampling, and that's a nonlinear trajectory.

158. 28:36

     Uh, and what a diffusion model does, it sort of gives you-- At each point on the trajectory, it gives you the direction that you should move in next, right?

159. 28:42

     So it, it predicts a tangent to this path, basically. Um, and then a question that arises is, okay, once we know what this path is, especially if, if we have, like, a deterministic sampling algorithm, there's, like, one path between a particular noise sample and a particular data point.

160. 28:56

     Once we know that path, um, why are we predicting the tangent to this path? Why can't we just predict where we're gonna end up? Right? Just, like, where we are.

161. 29:04

     Let's just try to predict where this path is gonna end up when all the noise is gone, and then we can just sample in one step, right? And this is exactly what consistency models do.

162. 29:12

     So you can take a diffusion model and then distill it into a consistency model, and this is just one of the many ways that people have developed, uh, to, uh, reduce the number of steps that are required to sample from diffusion models.

163. 29:22

     Now, in practice, when you do try to do this, uh, and try to sample in one step from a consistency model, it usually won't work that well, and it's ma- it's mostly because, like, you're, you're asking, uh, the neural network to do in a single pass what it did before in, like, fifty passes or something like that,

164. 29:37

     right? So, like, that's a tall order. So usually that's not gonna work super well out of the box. Um, but so one, one, uh, uh, trade-off that you can sort of make is you can basically say, "I'm gonna do consistency modeling, but only for, like, certain intervals of my sampling path so that I can still sample in,

165. 29:51

     like, maybe three steps, and maybe then I'm gonna get, uh, better results." This is just one of the, one of the many ways, uh, uh, to achieve this, but I just wanted to kind of show this, um, uh, intuitively what this is actually doing.

166. 30:02

     All right. Let's talk briefly about, uh, controls. This is the last thing I wanna talk about. So these models only are as useful as, um, our ability to actually influence, uh, at sort of a semantically, semantically high level, um, what they actually produce, right?

167. 30:18

     And so the canonical way that we've been doing this with these models is to give them a text prompt. You give them a text prompt, and then that describes what the image should look like, what the video should look like.

168. 30:27

     That's great. That works pretty well. But more and more, we're starting to see that people want more than that, right? They don't just wanna try to, uh... Ooh, lost my slide.

169. 30:35

     They don't just wanna try and, um, uh... Is this on my side or is this-

170. 30:40

     Oh, it's back. Awesome, yeah. Um, they don't just wanna try and describe everything in, in, in great detail in, in language. They-- You wanna be able to do reference-based generation, for example.

171. 30:50

     You wanna, you know, you wanna be able to generate a video and put yourself in it, right? And you're not gonna do that by describing what you look like in intricate detail.

172. 30:57

     You're gonna do that by taking a photo or a short video of yourself and then conditioning the model on that. So we want these other types of conditioning signals as well.

173. 31:04

     Uh, in video generation specifically, you might wanna sort of explicitly control camera motion, the speed of events happening, like the specific timing of events. That's not something that you should sort of try to shoehorn into a textual representation, right?

174. 31:15

     So you need sort of more advanced conditioning signals. And then a big question that arises is, when should we introduce these conditioning signals to the model? 'Cause often you don't have this kind of conditioning information for most of your pre-training dataset, so it can actually be quite useful to do that in post-training.

175. 31:30

     Like, pre-train a model with, with just text prompts and then maybe add some of these extra conditioning signals, uh, in a post-training phase. And of course, other things we can do in post-training are, like, preference tuning, uh, based on, based on human evaluations, like, uh, uh, using reinforcement learning or, uh, direct preference optimization.

176. 31:47

     All right. That was, uh, sort of a whirlwind overview of all eight of these things. Um, I'm going to stop there so we have some time for questions. Um, all right.

177. 31:59

     Um, thank you very much for listening. [audience applauding]

178. 32:15

     Yeah. So you, you were talking about how you increase guidance with increased quality. Is that kind of why lots of these models have a certain style? Like, if you look at a Nano Banana image and then look at one-

179. 32:25

     Uh, that's part of it. I mean, I think part of it is the, probably the post-training recipe that tends to sort of make these models a lot more opinionated.

180. 32:33

     I like to think of post-training as sort of taking... Like, pre-training sort of models the distribution of images, let's say, and then post-training is more about, like, which sliver of that distribution am I actually interested in, and that's a very opinionated type of thing.

181. 32:45

     So that's playing a big role in that. I think there's specific effects though that you can observe that are easily attributed to guidance. Like, if you get images that are, that look sort of very saturated, uh, that's usually a sign that the guidance scale is too high, right?

182. 32:57

     It's, it's like, uh, because the guidance scale is sort of amplifying this, um, this difference between the unconditional and the conditional prediction, but it's doing that at every noise level.

183. 33:07

     And so, um, if you do that at s- Depending on which noise level you're sort of manipulating this, this direction at, you sort of get these different effects. So what people have often s-started doing is, like, varying the, the s- the scale of the guidance across the sampling procedure.

184. 33:20

     Like, it often-- There's a, there's a paper that says it's actually best to not do guidance at the very start and at the very end, but you should sort of ramp it up in the middle, and that tends to give you the best sort of trade-off.

185. 33:29

     Um, but yeah, it's definitely playing a role, but I think probably a bigger role in that is played by the post-training. Yeah.

186. 33:36

     Yes.

187. 33:36

     How does, uh, the guidance or the conditional work for text diffusion as well?

188. 33:42

     Sorry, could you just repeat that question please, that one? Thank you very much.

189. 33:45

     Hi. Um, how would you do this guidance or, like, a s- conditional sampling for text diffusion, let's say?

190. 33:53

     Um, it-- You could use exactly the same, uh, procedure, um, provided that you have, you know, you have a sort of, uh, prompt that you, like you, that you're giving the model, right?

191. 34:03

     You can, you can u- treat that prompt as a conditioning signal, and then you could, uh, apply guidance to that. Now, how well that's going to work is, is sort of an open question.

192. 34:11

     Like, it, it clearly is applicable to this setting as well. But I think, uh, typically what we've found is that guidance tends to work best when there is sort of, um, a bit of a semantic gap between the thing that you're guiding with and the thing that you're generating.

193. 34:27

     And so that's why it works so exceptionally well with this sort of prompting setup. Like, if you guide on a text prompt, text prompt tends to capture sort of the high-level semantics of what's in an image.

194. 34:37

     And, and that seems to be, like, a really good way to sort of steer the model in a sort of semantic space, so to speak. But if you're, if you're guiding...

195. 34:45

     Like, this is-- You have the same problem in image generation. Like, if you do guidance on-- Like, let's say you provide a segmentation mask and you want the model to fill in the s- the, the, the different parts of the segmentation mask, right?

196. 34:56

     Actually, guidance will work a lot less well for that because it's sort of a low-level conditioning signal that's sort of at the same level of abstraction as the, the pixels that you're trying to generate.

197. 35:04

     It seems like guidance just works better when you have that gap. I think that's also why people don't typically use it as much for language models, um, even though it's perfectly applicable to, uh, autoregressive models as well.

198. 35:17

     Yeah. Um, over there.

199. 35:21

     Uh, a question back here.

200. 35:27

     Hi. Uh, you mentioned, uh, JAX is better suited for, uh, model sharding and parallelism. Uh, is there any particular reason for that, especially compared to PyTorch? Um, mic.

201. 35:36

     Um, okay. I have to caveat my answer with the fact that I've been at Google for over a decade. So, you know, PyTorch didn't exist when I joined Google, right?

202. 35:45

     So I've never had a lot of hands-on experience with, uh, with PyTorch. I think JAX, JAX is just kind of designed with that in mind from, from the get-go because, uh, you know, we've, we've had TPUs for a while at Google, and it was really just, um, I think, conceived to try and make, make that as easy

203. 36:03

     as possible, right? To, to, to make using TPUs as easy as possible. And with TPUs, like, you very rarely use just one chip, right? [chuckles] You kind of wanna take advantage of the fact that they have th- this really fast interconnect and sort of scale up, uh, pretty quickly.

204. 36:15

     Um, so yeah, I think it was sort of designed with this in mind from the get-go. But, you know, PyTorch can do these things as well, right? Uh, it's not, it's not unique to JAX.

205. 36:29

     Thank you, uh, for your presentation.

206. 36:30

     Sure.

207. 36:31

     If you slow down the process of denoising and noising, do you get better images or-

208. 36:36

     Slow down as in take more steps?

209. 36:38

     Yes, take, take more steps

210. 36:40

     ... up to a point.

211. 36:41

     Yeah.

212. 36:41

     Yeah, 'cause it- as I said, you're sort of, uh, you're approximating this nonlinear path through input space, right? So you're d- and you're approximating that by taking finite steps, so that means you're actually approximating that with sort of a, a, a piecewise linear path-

213. 36:55

     Mm-hmm

214. 36:55

     ... through input space, right? And so the more that path, that piecewise linear path starts to deviate from your true nonlinear pathway, the worse your samples are gonna get.

215. 37:04

     But there is obviously a limit and, uh, and there's also been a lot of research on, like, how can we train diffusion models or, like, models like this that actually produce straighter paths from the get-go so that you can actually get a more accurate approximation with fewer steps.

216. 37:18

     This is something called, uh, rectified flow or Reflow, uh, that, that sort of tries to dig into this. So yeah, and so up to a point is, is the answer.

217. 37:24

     Yeah.

218. 37:25

     Thanks. I, um, I was wondering if you could maybe talk a little bit more about how you could get these other types of control signals in for, like, a special camera movement.

219. 37:37

     Mm-hmm.

220. 37:38

     Is this something that you would do, like, with a simulator or something where you can, like, actually control all the parameters? Or, like, how-

221. 37:43

     That's a very good question. Um, I... Well, obviously I can't con- comment on what we actually do. [laughs] I can sort of speculate. I- that's actually... Like, I don't actually know the details about this because the team at this point is quite large, and, like, there's, like, people working specifically on what we call capabilities.

222. 37:59

     You know, like, uh, endowing the model with this ability to, for example, uh, be conditioned on, on, on a camera control signal and things like that. Uh, I'm not actually actively involved, uh, in that part of the work.

223. 38:11

     But yeah, it's like, it's, it's very, um... As I said, it's sort of like, um, it's quite important to have these conditioning signals that are sort of semantically abstract, so you kind of have to find the right representation that makes the most sense, uh, to condition a model on.

224. 38:23

     I think also what's very important is, um, you are sort of, you know, if, if, if, if your model is a transformer, you actually have quite a few ways to insert conditioning information.

225. 38:34

     Like, you could add a bunch of extra tokens. That's one way to do it. Uh, sometimes you might rather sort of broadcast the conditioning information to all the tokens.

226. 38:43

     Like, for example, if you, if we, g- the, the, the noise level input that we give to the diffusion model, like, to tell it how much noise there is in it, uh, we're going to typically broadcast that to all the tokens.

227. 38:53

     We're not gonna just add a, a single token for that. You know, and that's, that's, that's... And then there's, um, uh, I think there's like a third way that people have started doing conditioning.

228. 39:02

     But yeah, this, these are, these are all things that you have to think about when you sort of try to figure out how to introduce a new conditioning signal.

229. 39:08

     Um, yeah. I think we have time for maybe one more.

230. 39:16

     Well, I might pass that down to Thanks for the really, really nice talk. Thank you. Uh, just to understand, uh, in, in the denoising part, if you didn't add that error in each step, would it be fully deterministic- From a- ...

231. 39:36

     yeah. Yes. Okay. Yeah. So there's, there's basically, when, when you, once you have a denoiser, there are sort of two different types of generative models that you can construct.

232. 39:45

     One of them is stochastic, and the other is fully deterministic, right? And it, it, it's kind of a magical, uh, feat. Like, I still don't really believe that this works, but it really does.

233. 39:54

     Uh, so you can kind of, yeah, you can have a deterministic model. And as I said, this can be useful in many applications, like especially if you want to sort of go both ways.

234. 40:01

     Like, if you wanna go from your data distribution to your noise distribution and back, it's really useful. And then also for many distillation techniques, it's a requirement for the teacher to be deterministic.

235. 40:11

     So that's the case with these consistency models, for example. Um, you need, you need a, a, a deterministic sampling algorithm to start from. Yeah. Thanks. Cheers. All right. Thank you, everyone.

236. 40:21

     Um, I have a longer version of this talk on YouTube. Like, if you want a bit more detail, uh, for some of these sections, uh, you can check that out as well.

237. 40:28

     Cool. Thank you. [audience applauds] [upbeat electronic music]