Training Krea 2 - What matters in generative model training.

[upbeat music] Okay. So, uh, I'm Sangwu.

I'm from Krea. I'm gonna be talking a little bit about, uh, how we recently trained, uh, our, like, image foundation model, uh, Krea 2, as well as-- and we also recently, like, open sourced, uh, the medium version of our model.

So I'll be, like, talking about, like, how we, like, trained it mostly from research perspective. Uh, later in the day, uh, like, my colleague will also give a little bit of details on, like, infrastructure and training infrastructure and, like, everything that needed to be, like, set up for that.

But today I'll mostly be talking about, like, research. Also, public speaking is not one of my best abilities, so take that in mind. But anyways, I'll go ahead and start.

So, uh, we rec-- As I said, we recently open source our, uh, Krea 2, uh, Krea 2 medium, uh, variant of our, like, closed source model. Now it's, like, open source, and I think people are enjoying it quite a bit.

One thing we focused quite a bit was, like, stylistic, uh, diversity.

So how we are gonna, like, go about this is that I'm going to, like, describe a little bit of the training, like, pipeline, uh, that we have, like, used to, like, curate data, curate data, actually, like, train the model, train the model.

And then I'll, like, talk a little bit about, like, what were, like, the most effective levers were actually improving model performance. And then the last thing would be, you know, like what I think would be, or like kind of exciting directions for the next, uh, generation image model.

So I'll go ahead and then start talking about that.

So one of the trend that we've been, like, seeing recently is that, you know, existing models like ChatGPT, uh, ChatGPT, Nanite Pro, like the, like the production grade, uh, models that are, like, from the big labs.

One thing is that they kind of like focus on, like, slower generation, but, like, very reliable output. So I know ChatGPT 2 or like Nanite Pro, I think they take up to like a minute or like two to, like, give an output.

Typically, it's like something very acceptable. There's like hardly barely any like flaws. But one problem is that in order to get like good consistency, they have like kind of like significantly mode collapse their models because I don't know, like if you're trying to render a person, the easiest and most reliable way to, like, render a person is

render the most boring average person that exists, and then, like, put it in a center frame. Uh, and like this, for example, like this is one of the examples that I like to use.

So like if you type like burning skull, ChatGPT 2, like is very consistent. All the outputs are, like, fine, but you know, there's barely any diversity. On the other end, like kind of the thing that we kind of like focus on is like faster generation so that like people can like iterate with like different visual ideas and

then also have like little bit of knobs on kind of... Because, you know, when people like come in sometimes, like they want very-- something very specific. They want to generate like a poster or like a birthday card.

In that case, like ChatGPT 2 or Nanite Pro are like excellent solution. But sometimes like when you're like a creative studio, you don't quite know what you want yet and l-- want to like slightly explore what kind of like visuals you want to make.

That's something that we wanted to, uh, focus on a little bit. So, you know,

with that in mind, like how do you actually like train a diffusion model? Um, I hope most of the people in the audience like know how diffusion models work.

But just to give a small recap, how it works is that you have an image and you add a little bit of like noise, and then you ask the model to like, okay, like how here is an image with some of the noise added in?

Like how would you like remove this noise, uh, remove this noise to get like a valid image? So like at a high level, this is how diffusion models are like trained.

And once you train a model to do this, you can just like get a random noise and then tell the model to like progressively remove the noise, uh, noise to generate like an image.

So at a high level, like this is how a diffusion models are like trained. But right now most of the production, uh, production diffusion models, at least the open source ones that are very competitive, including ours, use, uh, autoencoder.

So this is typically called a latent diffusion model, where you use a autoencoder to like instead of taking the raw pixels, you first like spatially compress it and then decompress it, uh, later.

But the generation works in a spatially compressed, uh, latent space, so that is a little bit more efficient to train. One of the main reason is that most people use diffusion transformers, uh, transformers or like some variant of that to train diffusion models.

And as you know, transformers, at least the non-sparse ones, uh, they tend to have like O, like O n squared, uh, O n squared, uh, time complexity. So if you try to like model every single pixel, that's very expensive.

So now, uh, and that was the initial motivation for like why people kind of like started with like latent, uh, latent diffusion models because now they can actually, like model this a little bit in a more efficient space.

So one thing that I will say is that like really like data is like quite everything that goes into the model. Like typically you lock in your architecture and then you just...

A lot of work just goes into like just feeding the model, like what it wants. Like, I mean, this sounds stupid, but I cannot iterate this more. That's why I put like again, like really data is quite like everything.

Like again, like this, this just comes up again and again. And typically after you lock in your architecture, like most of the work actually goes into like data curation, like making sure that the data is good.

And in our case, like we wanted to focus on like stylistic diversity. So making sure that like we don't necessarily like remove like remove like remove like data-

Data unintentionally to, like, cut stylistic diversity. Like, that was also very important focus for us. For instance, like, some people, like, think, I know, like, low resolution CRT videos are, like, a bad image, but some people like that kind of, like, aesthetic.

So making sure that we have, like, good coverage and don't just rely on, like, very standard, like, aesthetic scores or, like, image quality scores to, like, cut, like, oversample, uh, typically what are considered, like, conventionally good images was also something that we had to, like, take into account.

So with that in mind, I'm gonna describe some of the things that I think we did, uh, quite poor when we came to, like, data and kind of, like, things that we, like, kind of consider bad data or, like, duplicated samples, over-represented concepts.

So these are typically taken, uh, care of by, like, deduplication and, and, uh, clustering, uh, based, uh, rebalancing. And then, you know, there are certain samples that, like, vision language models we use to generate the captions, they sometimes constantly fail to, like, capture important aspect of the image, which leads to certain biases, uh, which is also the

third point. And then, like, you know, sometimes when you train on, like, very low resolution, it doesn't make sense to, like, put train on an image that has, like, I don't know, like, twenty hundred, like, characters on a two hundred fifty-six by two hundred fifty-six pixels because that's just going to be, like, too hard for the model

to, like, learn, at least at, like, low resolution stages. So that's one, uh, example. And then obviously AI images. Uh, distillation is, like, a very big, like, topic, but we try very hard to, like, remove any AI images, like, at all because it does, like, provide you a shortcut to, to get you, like, a good model.

But synthetic data is, like, so sticky to the model that once you, like, start training on AI image data, sure, your model is good, but you kind of lose the point because then you're gonna get, like, very, like, ChatGPT or, like, Nanobanner aesthetic.

And at least for me, like, I can tell when, like, a model has been very heavily trained or, like, distilled on, like, ChatGPT 2 or, like, Nanobanner Pro. And in the long run, like, that's just not how you want to, like, do.

Like, you don't, I... And, you know, as a researcher, it also slightly hurts my ego if all I'm doing is distillation. So, yeah.

And then, you know, captions are very important. So, like, just to briefly describe our, like, captioning pipeline, how we, we take the images, and then we run OCR because text rendering is quite important.

Uh, so we make sure that we first extract out all the text that are, like, visible in the image, and then we also add, like, optional metadata. If we know this is, like, a picture of a famous person, we make sure that that's included.

And then we do a second pass with a vision language model to generate a better detailed caption for this image. And once we have, like, a caption that sufficiently captures all the things that are relevant about the image, then we can, like, rewrite it to, like, whatever, like, JSON prompts or, like, other formats you would like to,

like, uh, feed the model. So, like, this is our, like, captioning pipeline. And this is actually, uh, one good example of what I consider, like, bad data. It doesn't look that bad, but one of the issues that we had with this kind of images is that we would, like, try many things with the captioner, uh, and it

would say that, "Oh, this is a painting of blah, blah, blahdy, blahdy, blah," but it would not mention the fact that it's framed on a wall with, like, a, an, and have a white background.

And I don't know, s-sometimes the captioner would, like, consistently, like, fail to mention this fact. So when you try to generate a painting of whatever, it'll be always hanged on a wall, on a white wall, which is probably not what the user wants.

So this is an example of, like, you know, the image is fine. It-- You can train on it. But if it's a kind of image that, you know, LLMs or vision language models cannot, like, consistently, like, capture important aspect of the model, like, this is an example where we just had to, like, design filters and then throw

this kind of, like, data out or at least undersample it.

And, you know, deduplication, we mostly use, like, hash-based solutions. So we-- When we, like, train these models, we need to use, like, like, I don't know, anywhere from, like, two to, like, ten billion images.

That's a lot of images to run, like, filters on. So, like, first thing we do is, like, just to calculate, like, pHash or, like, MD5 hash to do basic deduplication.

And then once we get to, like, a smaller, like, size, that's when we bring in, you know, some of the embedding based, uh, deduplication method, SSCD, like, SigLIP to do, uh, se- semantic de-duplication or, like, near dup-- uh, remove near duplicates.

And, like, typically how we, like, also, like, design filters that we use, like, large, uh, large vision language model. I know, for instance, you design a prompt to, like, get a large language model to, like, learn, does this look like a AI image or not?

And then, like, we can-- After, like, we get a good, like, fine-tune or a system prompt from a large, uh, vision language model, one thing that, uh, we can do is we can, like, distill this data,

this kind of, like, decision and, like, knowledge to, like, a very small, like, SigLIP classifier, and then you can bas-- You have, like, a very, like, cheap classifier that is somewhat, like, reliable and, you know, like, typically if you run a classifier over, like, a billion images, you do need things to be, like, SigLIP sized, uh, for

instance. So this is one of... And also, like, this is, if you guys are into, like, LLM literature, like, this is one of the approaches that, uh, the essential web data is used.

They also use the big LLM, big LLM to, like, come up with, like, some kind of taxonomy or, like, classifiers to, like, judge whether this text data is, like, good or not in this, and then distill that down to, like, I don't know, like, five hundred, uh, million parameter model so that you can actually run this over,

like, a pre-training level corpus. Otherwise, it'll be somewhat expensive and frankly inefficient use of your GPUs to do so.

And another thing which I'm a little bit proud of, uh, is that we actually use, like, Sparse Autoencoders, uh, for some of our filtering. Uh, I don't know how many of you guys are still, like, into, like, Sparse autoencoders, but, uh- I, we did some work, or at least I did some work in doing s- uh, sparse

autoencoder research on, like, CLIP or these kind of, like, image models, uh, vision models. And one thing that you can actually get out of SAE is a unsupervised tagging system.

So one thing you can do is that once you train a SAE on a vision model, what you can do is that you can feed an image, and then it will give you, like, sparse features that get activated.

For instance, let's say you feed this image to a sparse auto- vision sparse autoencoder, and it will give you, like... It will get activated for, like, features like horse, black and white, blur, uh, blurry, like, image, blurry image.

So once... So you can kind of use this as, like, off-the-shelf, like, u- unsupervised tagging system. And if one of these, like, have, like, something that you want to, like, filter on or

oversample on, like, easy things are, like, signatures or, like, watermarks or, like, some kind of, like, border artifacts that I was, like, talking about. So if you have, like, a feature for that in your sparse autoencoder, like, this is one nice thing to remove, uh, to, like, remove, like, data that's kind of, like, undesirable, uh, in your

data set. And then one thing that I think, like, people have quite liked about our, uh, open source Krea2 model is, like, world knowledge. Frankly, I don't know how much this has helped, but this is one of the things we did, uh, is that, you know, this is also something that, like, CLIP, the original CLIP paper did,

is that you can actually take the Wikipedia, the entire Wikipedia, and for each article or concept, you can compute the page rank of each of the concepts. And then if it's, like, if it has, like, quite high, it's in the top ninety percent percentile of, like, page rank on Wikipedia, it's probably something quite important for the model

to know. And what we do is, uh, what we do is we kind of, like, take these keywords and then make sure that... And then just run, like, I know, like, standard, like, plain text search or, like, embedding search just to make sure that, like, these kind of concepts are in our data set.

And during this little bit of exploration, like, one thing that we did find was that there's a Barack Obama the horse. I actually don't know if this ended up in our data set, but you know, like, there are many interesting things in the Wikipedia article, so something fun to share.

But actually, I don't think we had this specific Barack Obama the horse, but it's, uh, it's a nice horse. Anyways, uh, yeah, and, and at the day, like, we have, like, around, like, 30 to...

We-- I counted, uh, and then I think we ended up having around, like, 30 to 40, like, custom in-house classifier, like, different heuristics and, like, filters that we've used.

And again, like, data is very important because once you lock in your... I think now we're... I mean, you're probably going to train, like, a transformer. You know what you're training, and then, like, data is, like, what really determines the quality of your model.

So again, can't emphasize on how important data is. And then, like, once your data is set up, like, we kind of, we kind of go through this, like, very LLM, uh, inspired training pipeline where we do, uh, we progress, we do, like, uh, low resolution to, like, high resolution pre-training, mid-training, supervised fine-tuning, preference optimization, uh, reinforcement learning.

And we also, like, train our, uh, prompt expander, which takes in user prompt and then expands that out to a long prompt. So this is pretty straightforward, you know.

At-- Typically, most people, like, train, uh, start training at low resolution because that's where the model actually learns, like, text-to-image capabilities. Like, you know, like, it needs to know, like, how a horse looks like.

You can do-- train that at low resolution, and then you can progressively, like, scale up your training resolution so that it first learns semantics and then learns, like, structure, detail, like these kind of things that can be learned at high resolution later.

So we've trained from two hundred and fifty-six to one K resolution. And then once you have your, like, pre-trained model, you kind of have this very malleable base to, like, train on.

And similar to, like, LLMs, like, when-- If you pre-train an LLM, it's just basically autocomplete. But typically, at least for LLMs, you want it to do, like, chat or, like, agentic stuff, tool calling.

So you actually need to, like, mold this into, like, things that are useful for you. So in our case, we curate, like, illustration, graphic design, photography, cinematics, you know, kind of like data you want, you have in mind for your downstream use case.

So we curate some, like, large scale mid-training data and then SFT data to kind of mold your distribution. And once this little bit of, like, molding is done, then we do, like, preference optimization, which is, you know, if you do ChatGPT, they will ask you, like, "Do you like this over this?"

So we collect bunch of these pairs where we use this for doing a little bit of, like, preference optimization to just polish up the model a little bit. This is where we get a little bit more opinionated, need opinionated about, like, the kind of model we want to train.

And then, you know, reinforcement learning is also, like, kind of, is now extremely standard in LLMs. That's what we do in DiffusionLen too. So we have, like, a pretty much, like, a GRPO inspired, uh, method where,

where we have model, like, generate images, and then we send that to the reward servers. And then based on this, like, feedback, we teach the model, like, I know, like, how to, like, improve text rendering, like, have, like, better anatomy structure, these kind of things.

And then the last step, this has kind of become a almost essential step for production grade diffusion models. You actually need to train a small LLM that takes in, like, user prompt and then outputs, like, a very long detailed prompt.

Uh, because typically longer detailed prompt that are more in distribution with your model's, like, training data, that tends to, like, make better images. And, you know, the kind of next step if you're into, like, uh, LLM literature is something that we are doing is, uh, doing, like, multi-expert, uh, on policy distillation that we are, like, currently working

on. So we would, like, train experts that are, like, specialized in, like, photography, text rendering, and, uh, different capabilities, and then kind of, like, merge all of these capabilities into a single student.

So w- So that, like, we can have, like, kind of all... We can have a student that can effectively match the capabilities of each of the expert in, like, whatever you want the expert to be, uh, good at.

And, like, in my experience, like, kind of like things that mattered for me was, like, infrastructure to iterate fast. And also, again, like, data is important. Like, methods can, like, change every time.

Code is something you can change very easily but, you know, data is, like, eternal. Like, you can-- If you have good data, that's gonna be valuable no matter what the hot new training paradigm is.

Uh, simplicity and scalability, I pref-- we very much prefer, like, methods that have, like, low number of hyperparameters that we need to tune. Efficiency, again, this goes, ties back to just fast iteration, like speed.

And then, like, you know, our thing that I like to do is steal a lot from LLM, uh, research so that I can just reuse their kernels and, like, research and, like, literature.

So all of these are things that I kind of, like, found very useful to iterate. I think I only have, like, one minute, but kind of things that I find funny is that, you know, we started with, like, encoder and, like, decoder, uh, transformers.

Now we are going the opposite way because we have a prompt expander which is all the regressive decoder, and then diffusion which is the encoder, so it's kind of getting reversed.

And with that in mind, it's also starting to look a little bit like DALL-E 2 where, you know, you have, like, a model that will, like, generate conditioning using, like, all the regressive model or, like, diffusion model.

And then we feed that to our diffusion model. So it's kind of this prompt expansion pipeline makes-- reminds me a little bit of, like, DALL-E 2, so that's kind of a funny observation.

And, you know, like, there's bunch of stuff that go into, like, diffusion model training. I really like to, you know, simplify the stack so that we can get rid of VAEs and then text encoders and then just train a single clean transformer.

Uh, so that's something I always look forward to working on. And then, you know, vision-language models has gotten, like, more powerful. Like, this from the original LDN paper. But, you know, like, bounding boxes, before it was expensive to generate, but now you can perfectly generate good bounding boxes for every image, and you can condition the image model

on this. Uh, this something that Ideogram Rev have, have also, like, been working on. And then, you know, this is from actually two thousand seventeen from, uh, Fei-Fei Li, uh, Li's lab.

Like, you could do, like, scene graph to, like, image generation. So again, like, image generation I think is really, like, a proxy for, like, VLM, like, progress. So kind of things I'm, like, excited for is, like, what are...

what are... what are interesting, like, textual ways to, like, describe an image now that we have more, like, more, like, powerful visual language models? So this something that I'm quite excited about, like bounding boxes, you know, scene graphs.

I mean, there's different things that could be, like, quite useful here. And then, yes, and this, uh, my shameless recruiting slide. So if you enjoyed this talk, s- send me an email here, take a picture.

And yeah, I think that's mostly it. I don't know if there's time for Q&A, but sorry that I went slightly over. But hopefully you enjoyed it. I'll stick around for another thirty minutes if anybody wants to talk to me.

But yeah. [audience applauding] [upbeat music]