Personalization for Text-to-Image Generative AI with Nataniel Ruiz - #648

0:09All right, everyone, welcome to another episode of the TwiML AI podcast. I am your host, Sam Charrington, and today I'm joined by Nataniel Ruiz. Nataniel is a research scientist at Google. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Nataniel, welcome back to the podcast. Thank you so much. Yeah, it's been a while. It's been like three years, maybe. It has been about three years, maybe a little bit more, and I'm really looking forward to chatting with you about what you're working on nowadays, which is personalization for generative AI models.

0:41Before we jump into that, I'd love to have you share a little bit about your background with our audience. Absolutely. So yeah, a lot has happened since then. A lot of my research topics have changed, but I think for very interesting directions. My background is I was born in Bolivia. When I was 18, I left for France and I did my undergrad there in computer science in Paris at Equipity Technique. And then I came to the US, I did a master's at Georgia Tech, and then I did my PhD at Boston University, which I just finished in March this year. And then I joined Google right after. So I was doing a long internship at Google while in my last year where we did a lot of the work that kind of set me on this path right now.

1:17And then, yeah, I joined Google full time as a research scientist. So I'm working there right now. Thank you. Last time we spoke, you were we were talking about deep fakes and adversarial attacks. And that was part of your graduate work. Are you still working on those topics? Yeah. Funnily enough, it didn't even make it into my thesis. My thesis ended up being about simulating images and video in order to train and test machine learning models in efficient ways. That was kind of like the topic. So a lot of I tried to like make it very kind of related, all the papers that went into that. So the deepfake and adversarial attack against deepfakes work didn't really kind of relate as much into that topic.

1:56But I did some cool work there. I think we were a little bit early with that idea, which the kind of like core idea is to use an adversarial attack on images in order for people that are trying to like kind of generate deepfakes out of those images. to basically thwart them from generating deepfakes. And so it's basically an adversarial attack. Like the most basic idea is to do an adversarial attack against a conditional image translation network in order for the network to completely fail when used on this specific input. And in that way, kind of, we were a little bit early because the methods that were going around then and there didn't work very well with very few images.

2:32The outputs were in that realistic. And I think that's kind of had a bit of a revival now with all of the incredible generative models for images and video that have been popping up. So this is actually kind of taking a little bit of a renaissance, that idea. And now there's a bunch of work. For us, it was a bit hard to get this published or understood at the time. And now it's very, I think, easy for people to understand it when they see the quality of the generations now. So they understand the value of needing to protect identity and from being replicated, et cetera. So there's a lot of work from Alexander Madri's lab at MIT that has been kind of reviving this issue, which is really cool and really important, I think.

3:10And so your focus today is on, broadly speaking, as I mentioned, personalization for generative AI models. Talk a little bit about just the problems you're exploring and trying to solve there. Yeah, so kind of chronologically how this happened, how I got interested in, I guess, is I had work at Apple that's called MorphGan, which was basically, this is kind of what set me on the deepfake disruption route, with adversarial attacks, because we realized that if you train a model on, like this was again at the time, it was like 2019, I think, you train a model on a lot of images, like paired images of faces, then you were able to kind of like manipulate these faces and generate like one shot kind of deepfakes of people.

3:53So it was pretty cool in order to like, you know, have your avatar and be able to like talk, but obviously has some repercussions with respect to privacy that set me into that route. So that's where I had some experience in generative models, and then kind kind of got recruited here at Google for an internship on personalizing generative models in order to generate objects. For example, you want to generate this specific can in different scenarios, different contexts, and you only have a limited amount of images of that can or this dog or something. So that's when we set out to do that. But this was right at the moment where Dali2 had come out and Imagine had come out from Google.

4:26So we got early access to Imagine internally, and we were able to kind of set out that route of like personalization for generative models at the time, for diffusion models that were generating such good quality images. And surprisingly, you know, simple techniques work incredibly well with diffusion models. And that's kind of how we came up with Dreambooth, which is basically personalizing a diffusion model or a generative model for a specific subject in order to generalize it in different contexts and situations. And then we discovered they could do different styles, et cetera. And so that kind of set out my current direction, which is personalization of generative models.

5:00And is Dreambooth, it's the beginning of the title of the paper of that name, but is it also a product? Is it a model? Is it a system? What all is Dreambooth? Yeah, I would say Dreambooth is kind of like a method or an algorithm. That's what I would say. And it's like, the idea is personalization of a generative model with few input images, the personalization, yeah, of a generative model using few images of a subject for the subject driven generation. That's something I think the term we came up with, subject driven generation, which is basically you want to grab the model and you want to query it in different ways.

5:35And then you want to generate novel pictures of that subject. So it could be like a cat or, you know, like some type of - So I upload a few pictures of my cat or my dog and say, okay, now show me Sparky in Paris at night with a orange sky or something like that. And the idea is that this algorithm will generate images based on that prompt with that subject. Exactly. So that, I mean, I think you explained it very well. So that's like, it was kind of like a new problem because methods at the time weren't able to solve this problem before. And it was kind of a problem because the capabilities of current models weren't there yet.

6:13And then when the models had come out, like the diffusion models like DALI 2 and Imagine, And then it was kind of our contribution there is to create a method that is able to kind of tackle this problem. So like formalize this problem, tackle this problem. And also in the paper, we have a lot of like scientific contributions, such as like how do you measure success in this scenario? And we created the first data set that is large for this problem. So we kind of set out a new direction for research in this area. Are there things that, for folks that are tangentially familiar with diffusion models, but not familiar with their details, are there elements of the details, the workings of diffusion models that are important to fully understand what you've done and the way that you've used them?

7:04So diffusion models kind of like, I think I guess already existed like about seven years ago, I think, in 2016, if I'm correct. I'm 100 % sure. But then like in 2019, I think there were a bunch of papers that really made them work very well. And then 2019, 2020, and then, you know, more and more like training on larger data sets, getting some of the techniques right. And then they were kind of, now they've superseded like GANs for the best generative models. The core idea is you're basically learning to denoise images. So I guess a very natural application scenario was to do like denoising or super resolution as applications with this, which basically, you know, the core idea of the diffusion model is you grab an image that's clean, you noise it intermediately with a specific amount of like noise, and then you train a model that tries to denoise that image in one step.

7:58And then in inference, you go from, you know, if you want to generate new images, like people realize that you could generate new images after you train this model, if you just grab full a full noise, also full noisy image, and then you can go fully into the clean image regime, iteratively denoising that image. And that actually works really well, which is kind of like, it's an interesting kind of concept, because it's an asymmetric thing where your training is different from your inference, but it works extremely well. One thing that really makes it work is like the new techniques that have been developed like inside of Google and OpenAI and Berkeley in the last like three to four years.

8:32And there's a bunch of papers on that. And then one of the things that really makes them amazing is that they're much more stable than GANs. They're easier to train basically because, you know, GANs, you have a lot of stability issues. You can have mode collapse. You have to kind of be an expert in order to get your GAN to train well. And a lot of people are trying to do that. But I think with diffusion models, you have much more flexibility in what you're doing with the architecture, what you're doing with certain things, and it's kind of easier to get them to work. And that can change everything.

9:01I also think it's like objectively a really good idea. So that's why it works. Is it fair to think about the way you set up the problem as, you know, the traditional way of doing image generation is you are kind of conditioning the image generation on some texts? Do you think of the subject-driven generation as additional conditioning or a constraint or something else? Like, does that question make sense? Yeah, it makes a lot of sense. I think like you can think of conditioning the whole procedure of generating new images for this kind of cat or something as, yeah, condition on the data set. But in practice, what we do is we don't kind of create a new conditioning kind of like pipeline in the network.

9:45What we just do is we just do fine tuning on the network, on these images, basically. So that's how it like technically this is how it happens. So I wouldn't call it like, yeah, conditioning per se, but you know, more recent models are trying to do faster Dreambooth and they do take this conditioning route. So conditioning on a small data set of images. What we did, which is the contribution that we have for Dreambooth is like, okay, you have the set of images of the cat. Then what happens if you fine tune them? Like, okay, so we figured out a way to fine tune these diffusion models on the set of images of this cat by just adding like a personalized or a personal token that is rare that kind of identifies this cat.

10:22So it can be like a string of characters that is not, that doesn't have a strong prior in the model. And we condition like these. So we train the model using this kind of small data set of cat images with prompts A, V, where the V is like the rare identifier and then the name, or so the class name cat. So the model can localize, you know, very well actually and if you do early stopping with this very simple technique of fine-tuning the model on the small data set with these prompts that we came up with then it works really really well and you can start generating the cat it doesn't overfit because you do some early stopping and it was able to like memorize this subject so then it can like regenerate in different even sometimes different poses and you know with different lighting and different circumstances and also like different styles so a lot of very surprising things that we get out of these diffusion models.

11:15You make it sound very simple. So yeah, I guess, honestly, I love that about it. I think it's something that I think we were kind of first to see that this was possible with diffusion models. And because of several things, probably like we don't have like a clear answer on why this simple algorithm works. But, you know, possibly it's because these models are like huge, you know like much larger than previous models they're also trained on much more data it's trained on pairs of text and images so a lot of those reasons and also like some properties of diffusion models can like help with avoiding memorization because if you try to train again on very few samples then you get mode collapse very easily so you start generating the exact same samples very quickly with diffusion model this happens more slowly and it can have something to do with kind of maybe how the training is constructed, where you kind of train the fusion model to denoise partially noised images.

12:14And maybe that makes it, you know, overfit less fast, you know. So those are all hypotheses on why it works. But we kind of were the ones that figured out, you know, oh, this works. And this super simple idea works. We don't even have to fine tune a subset of the parameters, find which parameters to fine tune. We can just fine tune the whole model on this small data set with this very simple caption and it works. So I really love that the idea is simple. And it has like, I think that's where it really exploded because people could use it pretty easily or implement it pretty easily. And it just worked.

12:44You know, what I would have imagined to happen is, you know, less about overfitting and more about the generated images being kind of reminiscent of the subject, but not really fully capturing the subject. But the samples that you show, like the subject is there and it doesn't seem like it's distorted. Yeah. It does a really good job of recreating the subject in the scenario. and granted that you just said that you know we don't understand fully why this works all that kind of stuff any any thoughts on why that part works why the the subject is recreated so you know carefully yet without the overfitting yeah like the subject details are very well preserved i think for the one the model that we were using imagine at the time it was a great model trained on a very large data set so i think what it can do is like it already has like a concept of For example, for dogs, it was pretty easy to train it on dogs because it already has concept of these, you know, specific dogs.

13:45So it can mix its prior with the details that you're trying to teach it, which are specific kind of like details in the fur or colors or shapes, et cetera, of that dog. So it's kind of leverages is very strong prior in order to learn these things. And for like pixel level diffusion models, which is basically you don't have any kind of like down sampling into latent space at all. I think there is no real issue because you could overfit to specific images very easily. You just like train it all a long time on an image. So it will learn how to generate almost anything you give it, especially if it has this like already pre-trained, like it's a large pre-trained model.

14:22So it will be able to memorize things much faster than a random model, I think. And for the latent case, for the, for example, like stable diffusion, which is a latent diffusion model, then the only kind of obstacle that I see it from like actually reproducing the subject with good fidelity would be the encoder, the autoencoder, right? And the autoencoder is able to kind of reproduce most images with pretty high fidelity. It's not perfect, but it's able to do it. And then with SDXL, actually, it is like much better autoencoder. So it can reproduce many, you know, most images almost perfectly. So what's SDXL?

14:56Oh, Stable Diffusion XL, which I think is has been released like a couple months ago. So yeah, I think I would have expected it to maybe be able to very easily reconstruct the subject, but then it can't reconstruct it in other poses or in other contexts that easily, or it would have had trouble. I think the amazing thing about it is that it does learn the subject and it's able to preserve these capabilities. That was kind of the biggest surprise. And then also like do it in different styles and different accessories. Like we just kept discovering new and new things that the model could do when it was personalized.

15:31And can you talk about the training process for diffusion models generally and then how the training process for Dreambooth? is Dreambooth using an off the shelf diffusion model? You said Imogen, imagine. Yeah, that's the first model we used. And then we have also experiments on stable diffusion, I think 1.4 when that came out. And then since then, Dreambooth has been kind of re-implemented open source on all of the open source models, which are not many. So like stable diffusion, all the stable diffusion versions and stable diffusion XL. So, okay. So you're taking the off the shelf model and fine tuning it with the subject images.

16:13Yeah. And it's like, I took, so we took a bit of inspiration. We took inspiration from computer vision a little bit, but one thing where this was kind of already working for it was language models. So you had like a very large language model and then you fine tune it for a specific task. And then it did really well for that task. And one thing that I was noticing is that when you did this, it wasn't just that it was learning new things and it had to learn a ton of new things. It was that the prior of the large language model that had already been pre-trained was so strong that you were bringing out kind of knowledge onto the surface.

16:47And I think that's a pattern that you see nowadays is like these models have been trained on so much data and they're so large and they're basically so good almost everywhere, but they have to be good on average. So then once you fine tune in a specific scenario, a specific subject or a specific kind of for large language models for a topic or something, then you really get, you dig out that prior. You maybe teach it new things, right? But you dig out that prior and you kind of specialize it a little bit. It's basically the model that was already good. You give it a little bit more knowledge to memorize, and then you also bring out that prior, that part of the prior that you're really interested in for your specific task, and then it does really well.

17:23So that's one of the things that inspired me was the large language models. Yeah, and our team in general, I think. And so does the, you know, when you do this pre-training, you mentioned you've got the, you kind of come up with this unique identifier that you use as part of the prompting. To what degree is the model able to retain the distinction between my cat that we've fine-tuned on and cats? Like if you, you know, said a group of six cats, you know, with my cat in the center, like, is it just going to generate a bunch of copies of my cat? Or does it actually, you know, is it actually able to retain the, you know, the ability to generalize cats?

18:02Yeah, exactly. So this is one of the questions we grapple with at the beginning is that if you fine tune it even for a few iterations on a specific subject with this like prompt AV cat, then it'll kind of tie the word cat to your cat too. And then it won't be able to generate other cats. So it kind of becomes, so we call it that language drift, where it forgot the meaning of the word cat. So we have some techniques in the paper. So like one specific technique that's called the prior preservation loss that avoids it from losing its prior. And it was an interesting idea where, again, like Dreambooth is the core idea of like personalizing generative model for a subject, right?

18:37And the simple algorithm of fine tuning for that subject. And this is kind of an extension on that. It has some advantages, some disadvantages. but definitely one thing that it does do is it avoids language drift. So what we did is you just generate a bunch of images of other cats and you also mix those in with a different prompt, so a cat without the identifier, and you train the model on both kind of data sets of your cat and other generated cats. So this is one of the first situations where you could generate your own data with the model and feed it back into the model to achieve something.

19:09So I think that was kind of cool. We call that autogenous because it's like generating its own, like training samples. And I think that idea actually influenced some future ideas. For example, the Suti paper, this subject driven text image generation via apprenticeship learning is basically like, that's like a fast dream thing that I participated in that was at Google, like really good work by some colleagues. And the idea was to generate a ton of like dream data sets. So you have certain subjects of like one, you know, this cat and then this bag or And then you don't have enough images of them in different contexts.

19:43So like generate a bunch of data sets of them in different contexts. So dream booth data sets, and then retrain the whole model with a conditioning kind of conditioning pipeline. So like condition the model on small data sets and train this in a supervised way to generate new images. And basically you get fast dream booth for like a wide array of subjects and it could be specialized, but also, you know, in that paper, we show that it works for like a wide array of subjects. And the more data you can generate, the better it can become. And when you say a wide array of subjects, meaning the subjects are predetermined?

20:18No. As opposed to... I mean, like you can insert a new subject, like you can use it on a new subject and it has a very high chance of working, especially if the number of subjects and variety that you've trained it on are large, which they are in that paper. Okay. Yeah. So one thing that's interesting there is like exactly that thing. So we were talking about that prior preservation loss. It's just like, wow, we reuse the images of the model to train it. That's the marking of like, wow, those image generation models were already really good that we could use their own data to train them, train like for some things on.

20:47And then that has been reused in other work, I think, to some measure of success. That method sounds, you know, super hacky, but kind of amazing that it works. Yeah, yeah. I mean, you could use real images, but it just takes too long to look for real images. So that's why we use images of the same model so that it could just fit in a training loop basically. But there are some caveats, by the way. If you keep training a model on a lot of generations by the same model, you can get artifacts that kind of start reoccurring. So you can get a feedback loop. It's like basically if you train LLMs with, we've been hearing all about this, train LLMs with LLM outputs, then they become more and more limited and more and more biased and the outputs will start to become really bad at some point.

21:34But for one round of training, it works. Is the prior generation an automated pipeline that looks at the image that you upload and figures out what images to generate? Or was that all done manually? You could do it manually, but you could also do it in an automated way. I think the only thing you need is just the class of the subject. And I think more and more... That's what I was kind of speculating at. Yeah, and you can have that as user input. But I think more and more people also have realized that for some things, you don't need the class of the subject. So I don't know, some people also have stopped using that in some situations.

22:09We kind of came out with the first version of Dreambooth, I like to think of it. And then a lot of people online have improved it more and more. So for example, we train all of the parameters in the model. And then there's been like kind of Dreambooth versions that train like only the cross attention and self attention layers, which is kind of just realizing that you can do Dreambooth, but with less amounts of layers. You don't have to train the convolutional layers of the network. But it's still the same basic algorithm. And then there's also doing Dreambooth, but with lower rank adaptation. So instead of training the full matrix for a linear layer in the cross-attention self-attention layer, then you have a low-rank approximation of that, which is the LoRa paper.

22:51And then that has become the main way of doing Dreambooth because it's so parameter efficient, but it's still using the core idea of Dreambooth, which is basically the subject-driven generation and then prompting. And then also another extension of Dreambooth that was kind of like found out by artists very quickly after we released, which was super cool. Like one of the coolest things was following the art community, honestly, and their excitement about Dreambooth and other algorithms. So they just figured out that you could do it on style with enough images and with careful enough prompting, you could actually do Dreambooth on style.

23:23So you could learn a style And then you can have new images of other things in that style. But then there's also, we have like, you know, more recent work that does focus on style. Meaning subject-driven image, like how's it, how's what you're describing different from style transfer? Yeah, it's very different. So we called it subject-driven because the only thing we have in the Dream with paper are subjects that are kind of in different contexts with different accessories, in different styles, but it's still like you're learning the subject. But what people figured out very quickly is that instead of subject-driven, you can just change it to concept driven and your concept could be a style.

23:56So you learn a style, you personalize the model for style, and then you generate new things using that specific style. So let's say you have, you're an artist and you have a very distinctive style, like it's like watercolor and like certain colors that you use a lot. And then certain, I don't know, like lines that you, you know, you usually highlight certain things more than others. And you have like a portfolio. So you have a hundred images and people have realized that they can fine tune on those images using Dreambooth style. So like very similar to Dreambooth, but just for style. And then now they can generate new things, like let's say a dog in my style.

24:29And then they were able to do that. And this is kind of the genesis also of like future work that inside of Google that I also participated in that is able to like do personalization for style. So this is different than style transfer in the sense that you're not grabbing an image and saying, I'm going to use this style and transfer it to this image. So like make this image in that style. The content is not preserved, like the content is fully generated by the model, the diffusion model or the generative model. And the style is the thing that you had learned and that you're reapplying into this image, basically.

24:58Just as it knows, as the model knows how to generate something in the style of Van Gogh, then you just added a new one, you know, Nathaniel style, even though I'm not an artist and my style is probably not very good. And if you were comparing two pipelines, one that does generic generation and then apply some style transfer method to this dream booth style or i forget the what is that one called style booth is does it have a name a distinct name or is it just the yeah so so dream with people started using dream with their style very early on but then we have a paper at google called style drop which uses a different yeah so that uses a different architecture that's not a diffusion model so it's a transformer based image generation model Muse.