Training Krea 2: What matters in generative model training — Sangwu Lee, Krea.ai

AI Engineer · 21 min · 156 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:01[music]
00:13Okay.
00:13So, uh I'm Sa.
00:15I'm from Korea.
00:16I'm going to be talking a little bit about uh how we recently trained uh our like image foundation model uh Korea 2 as well as and we also recently like open sourced
00:26uh the medium version of our model.
00:28So I'll be like talking about like how we like trained it mostly from research perspective.
00:33Uh 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.
00:43But today I'll mostly be talking about like research also public speaking is not one of my best abilities.
00:51So take that in mind but anyways I'll go ahead and start.
00:56So uh we rec as I said we recently open source our uh create 2 uh creative to medium uh variant of our like closed source model.
01:05Now it's like open source and I think people are enjoying it quite a bit.
01:09One thing we focused quite a bit was like stylistic uh diversity.
01:15So how we are going to 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
01:27train the model and then I'll like talk a little bit about like what were like the most effective levers for actually improving model performance and then the last thing would be you know like what I think would be are like kind of exciting directions for the next
01:42uh generation image model.
01:44So I'll go ahead and then start talking about that.
01:48So one of the trend that we've been like seeing recently is that you know existing models like chipt uh chach nanopar pro like the like the production grade uh models that are like from the big labs.
02:00One thing is that they kind of like focus on like slower generation but like very reliable output.
02:06So I don't know chachi BT2 or like nanobanara pro I think they take up to like a minute or like two to like give an output.
02:13typically is like something very acceptable.
02:15There's like hardly barely any like flaws.
02:17But one problem is that in order to get like good consistency, they have like kind of like significantly mode collapse their models.
02:25Because 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
02:33that exists and then like put it in a center frame.
02:35Uh and like this for example like this is one of the examples that I like to use.
02:40So like if you type like burning skull ch2 like it's 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
02:56also 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
03:06like a birthday card in that case like chatb2 or nano bar pro excellent solution, but sometimes like when you're like a creative studio, you don't quite know what you want yet and
03:17want to like slightly explore what kind of like visuals you want to make.
03:20That's something that we wanted to uh focus on a little bit.
03:24So, you know, with that in mind, like how do you actually like train a diffusion model?
03:30Um I hope most of the people in the audience like know how diffusion models work but just to give a small recap.
03:37How 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
03:51remove this noise to get like a valid image.
03:54So 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
04:05remove the noise uh noise to generate like an image.
04:08So at a high level like this is how a diffusion models are like trained.
04:12But right now most of the production uh production diffusion models at least the open source ones that are very competitive including ours use older encoder.
04:22So this is typically called a latent diffusion model where you use a autoenccoder 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
04:35compressed uh latent space so that it's a little bit more efficient to train.
04:39One of the main reason is that most people use diffusion transformers uh transformers or like some variant of that to train diffusion models.
04:48And as you know transformers at least the non-sparse ones uh they tend to have like n squared uh or n squared uh time complexity.
04:59So if you try to like model every single pixel that's very expensive.
05:03So 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.
05:16So one thing that I'll 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
05:26a lot of work just goes into like just feeding the model like what it wants.
05:29Like I mean this sounds stupid but I cannot iterate this more.
05:33That'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
05:46like making sure that the data is good and in our case like we wanted to focus on like stylistic diversity.
05:53So 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
06:11a 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 c like over sample
06:24uh typically what are considered like conventionally good images was also something that we had to like take into account.
06:31So with that in mind I'm going to describe some of the things that I think we did uh quite cool when we came to like data and kind of like things that we like kind of consider bad data are like duplicated samples over represented concepts.
06:46So these are typically taken uh care of by like dduplication 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.
07:00They sometimes consistently fail to like capture important aspect of the image which leads to certain biases.
07:06Uh which is also the third point.
07:08And 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 2000
07:18like characters on a 256 by 256 pixels because that's just going to be like too hard for the model to like learn at least at like low resolution stages.
07:28So that's one uh example.
07:30And 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
07:44synthetic 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 going to get like very like chatbt
07:55or like nanob aesthetic and at least for me like I can tell when like model has been very heavily trained or like distilled on like CHPD2 or like nanob pro and in the long run like that's just not how you want to like do like you don't I
08:09and you know as a researcher it also slightly hurts my ego if all I'm doing is distillation.
08:15So yeah and then you know captions are very important.
08:20So 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.
08:28Uh 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.
08:36If we know this is like a picture of a famous person, we make sure that that's included.
08:41And then we do a second pass with a visual language model to generate a very detailed caption for this image.
08:47And 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 what like JSON prompts or like other formats
08:57you would like to like uh feed the model.
09:00So like this our like captioning pipeline and this is actually one good example of what I consider like bad data.
09:09It doesn't look that bad.
09:10But 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
09:22blah blah but it would not mention the fact that it's framed on the wall with like on and have a white background and I know sometimes the captioner would like consistently like fail to mention this fact.
09:34So when you try to generate a painting of whatever it will be always hanged on a wall on a white wall which is probably not what the user wants.
09:41So 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
09:56design filters and then throw this kind of like data out or at least under sample it and you know dduplication we mostly use like hashbased solutions so we when we like train these models we need to use like like I don't know anywhere from two to like 10 billion images.
10:12That's a lot of images to run like filters on.
10:15So like first thing we do is like just tip calculate like P hash or like MD5 hash to do basic dduplication 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 dduplication method SSCD
10:29like cichlip to do uh se semantic duplication or like near dup uh remove near duplicates and like typically how we like also like design filter is that we use like large
10:42uh large vision language model I know for instance is 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 from from a large uh vision language model
10:59one 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 cichlip classifier and then you can bas you have like a very like cheap classifier
11:12that 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 cycllip sized uh for instance.
11:22So this is one of and also like this is if you guys are into like LM literature like this is one of the approaches that uh that essential web data is used.
11:31They also used the big LM big LM 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 distilled that down to like I don't know like 500
11:43million 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
11:55and another thing which I'm a little bit proud of uh is that we actually use like sparse order encoders uh for some of our filtering.
12:02Uh I don't know how many of you guys are still like into like sparse auto encoders but I we did some work or at least I did some work in doing
12:12uh sparse auto encoder research on like clip or these kind of like image models vision models and one thing that you can actually get out of SAPE is unsupervised
12:23tagging system.
12:23So one thing you can do is that once you train a SAPE 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
12:36let's say you feed this image to a sparse odo vision sparse odo encoder and it will give you like it will get activated for like features like horse black and white
12:46blur uh blurry like image blurry image so once so you can kind of use this as like offtheshelf of like unsupervised tagging system.
12:56And if one of these like have like something that you want to like filter on or over sample on like easy things are like signatures or like watermarks or like some kind of like border artifacts that I was like talking about.
13:07So if you have like a feature for that in your sparse auto encoder like this is one nice thing to remove uh to like remove like data that's kind of like undesirable
13:17uh in your data set.
13:19And then one thing that I think like people have quite liked about our open source create to model is like world knowledge.
13:26Frankly, 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
13:39the 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 90%
13:52tile of like page rank on Wikipedia is probably something quite important.
13:56for the model to know.
13:57And 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.
14:11And during this little bit of exploration like one thing that we did find was that there's a Barack Obama the horse.
14:19I 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 bra Obama the horse, but it's it's a nice horse.
14:34Anyways, uh yeah and and 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
14:43custom in-house classifier like different heristics 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.
15:01So again, can't emphasize of how important data is.
15:06And then like once your data is set up like we kind of go through this like very LLM uh inspired training pipeline where we do uh we progress we do like
15:17uh 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 to a long prompt.
15:31So 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
15:42like 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.
15:58So we've trained from 256 to 1k resolution.
16:02And then once you have your like pre-trained model, you kind of have this very malleable base to like train on.
16:08And similar to like LMS like if you pre-train LM is just basically autocomplete.
16:14But typically at least for LM you wanted to do like chat or like agent tech stuff tool calling.
16:20So you actually need to like mold this into like things that are useful for you.
16:24So 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.
16:33So we curate some like large scale mid-train data and then SFT data to kind of mold your distribution.
16:40And once this little bit of like molding is done, then we do like preference optimization, which is, you know, if you do CHP, they'll ask you like do you like this over this?
16:51So we collect bunch of these pairs where we use this for doing a little bit of like preference optimization to just polish out the model a little bit.
16:59This is where we get a little bit more opinionated need opinionated about like the kind of model we want to train.
17:06And then you know reinforcement learning is also like kind of is now extremely standard in LMS.
17:12That's what we do in diffusion land too.
17:14So we have like a pretty much like a gpo 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 don't know like how to like improve text rendering like
17:31have like better anatomy structure these kind of things and then the last step this is kind of become a almost essential step for production grade diffusion models you actually need to train a small LM that takes in like user prompt and then outputs is like a very long detailed prompt.
17:46Uh because typically longer detail prompt they're more in distribution with your models like training data that tends to like make better images and you know the kind of next step if you're into like uh LM literature is something that we are doing is uh doing like multi-expert
18:02on policy distillation that we are like currently working on.
18:05So 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.
18:17So 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.
18:30And 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 messes can like change every time code is something you can change very easily but you know data is like
18:44eternal like you can if you have good data that's going to be valuable no matter what the hot new training paradigm is.
18:51Uh simplicity and scalability I pref we very much prefer like methods that have like low number of hyperparameters that we need to tune.
18:58efficiency.
18:59Again, this goes ties back to just fast iteration like speed.
19:03And then like you know, our thing that I like to do is steal a lot from LM research so that I can just reuse their kernels and like research and like literature.
19:11So all of these are things that I kind of like found very useful to iterate.
19:16I 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.
19:25Now we are going the opposite way because we have a prom expander which is all the regressive decoder and then diffusion which is the encoder.
19:32So it's kind of getting reversed and with that in mind it's also starting to look a little bit like deli too where you know you have like a model
19:40that would like generate conditioning using like the regressive model or like diffusion model and then we feed that to our diffusion model.
19:47So it's kind of this prom expansion pipeline makes reminds me a little bit of like Dolly 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.
20:00I really like to you know simplify the stack so that we can get rid of VAE and then text encoders and then just train a single clean transformer.
20:09Uh so this something I always look forward to working on.
20:12And then you know vision language models has gotten like more powerful like this is from the original LDM paper but you know like bounding boxes before it was expensive to generate but now you can perfectly generate good bounding box for every image and you can condition the image model in this
20:29uh this is something that ideoggram rev have also like been working on and then you know this is from actually 2017 from uh FA Lee uh Lee's lab like you could do like scene graph to like image generation
20:42So again like image generation I think is really like a proxy for like BLM like progress.
20:47So kind of things I'm like excited for is like what are other what are interesting like textual ways to like describe on image now that we have more like more like powerful visual
20:58language models.
20:59So this is something that I'm quite excited about like bounding boxes you know scene graphs.
21:03I mean there's different things that could be like quite useful here.
21:06And then yes and this uh my shameless recording slide.
21:10So, if you enjoyed this talk, send me an email here, take a picture, and yeah, I think that's mostly it.
21:16I don't know if there's time for Q& A, but sorry that I went slightly over, but hopefully you enjoyed it.
21:22I'll stick around for under 30 minutes if anybody wants to talk to me.
21:25But yeah,