AI Engineer World's Fair 2026
Training Krea 2 - What matters in generative model training.
Slide chapters
6 slide and demo groups · Timestamps are approximate. 28 additional groups were omitted from this pass.
Read the talk
Training Krea 2 — What Matters in Generative Model Training
Stylistic diversity depends on more than sampling: Krea’s training pipeline connects data curation, caption quality, scalable filtering and staged optimization to the images a model can imagine.
From a talk by Sangwu Lee
Before you start: Basic familiarity with neural-network training and transformers is helpful; the article introduces the diffusion and data-curation mechanisms it uses.
What if you do not know which image you want?
What should an image generator optimize for when the user is still discovering what they want? Krea 2 makes stylistic diversity a central objective. Sangwu Lee introduces its released Medium variant through the research decisions behind data curation and training; a separate colleague’s talk covers training infrastructure. The release includes available weights and inference code, with distinct licenses: the repository identifies Apache-2.0 code and community-licensed weights.
Lee contrasts that objective with the reliable outputs of large production systems, including Nano Banana Pro and the model labeled ChatGPT 2 in his slides. Lee estimates that these systems can take one or two minutes to return an image, without specifying version, resolution or service conditions. Reliability can also favor familiar compositions: his example is an average-looking person placed in the center of the frame. He interprets this narrowing of outputs as mode collapse.
The prompt burning skull makes the distinction visible. The slide shows four varied Krea 2 images above four similarly composed orange skulls labeled ChatGPT 2. The latter outputs are individually acceptable; the missing property is variety across attempts. Krea’s goal is faster generation that lets someone iterate through different visual ideas.
That changes what counts as a useful product. A user requesting a specific poster or birthday card may value a dependable finished result. A creative studio exploring a visual direction may instead need variation, quick iteration and controls for steering the search. The target is not simply a good image, but a useful range of possible images.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Learn to denoise a compressed image
Diffusion training starts with an image, adds noise and asks a model to learn how to remove it. Once trained, the model can start from random noise and progressively turn it into an image. Training supplies examples of the denoising task; generation repeatedly applies the learned operation.
Krea performs this work in a spatially compressed representation. An autoencoder encodes pixels into a latent space, the diffusion model generates within that space, and a decoder converts the result back into pixels. This is the central computational motivation for latent diffusion: a diffusion transformer processes a shorter sequence instead of directly modeling every pixel.
Dense attention has a quadratic dependence on sequence length. For n image tokens, its pairwise attention work scales as:
Spatial compression reduces the sequence that incurs this cost. The autoencoder adds another component to the system, but makes the generative model’s training problem more economical.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Do not filter away the aesthetic
Once the architecture is fixed, much of the research work becomes deciding what to feed it. For a model intended to preserve stylistic diversity, conventional image-quality filtering can remove exactly the examples it needs. A low-resolution CRT image may score poorly as a clean photograph while being an excellent example of a desired aesthetic. Oversampling conventionally attractive images changes the model’s distribution, even if every individual filtering decision looks reasonable.
Undesirable data therefore needs a more specific definition:
- Redundancy: duplicate samples and overrepresented concepts need deduplication and clustering-based rebalancing.
- Missing supervision: a captioner may repeatedly omit an important visual property, leaving the model to absorb it as an unexplained default.
- Resolution mismatch: a densely text-filled image may be useful eventually but too difficult to learn from at a 256 × 256 training stage.
These are different problems and require different filters. A single aesthetic score cannot resolve them.
Krea also tries hard to exclude AI-generated images from pretraining. Lee describes synthetic aesthetics as unusually sticky: distillation can quickly produce an attractive model, but it can also import the recognizable appearance of its teacher. That shortcut conflicts with stylistic independence. The exclusion is specific to pretraining; it does not imply that synthetic examples have no role in later preference training. Lee also admits a less technical motivation: simply distilling another model is less satisfying research.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
A good image can still be bad supervision
Caption construction is a pipeline rather than a single request to describe an image:
- Extract visible text with OCR. Accurate text supervision matters for learning to render text.
- Add available metadata. If the image depicts a known person, include that identity.
- Generate a detailed caption with a VLM. Use the image and the gathered information to capture relevant visual content.
- Rewrite the caption into the training format. Once the description is sufficiently complete, convert it into JSON prompts or another desired representation.
The order matters: reformatting a caption cannot recover information that never entered it.
The painting example exposes that failure. A captioner describes the painting’s subject but repeatedly fails to mention that it is framed, hanging on a wall and surrounded by a white background. The screenshot shows a colorful painting of a pink flowering tree against a pale wall. During training, the wall remains in the image even though it is absent from the description. A later request for a painting can consequently produce a painting hung on a white wall, rather than the artwork alone.
The image itself is usable. The problem is the pairing between pixels and language: important context is consistently unlabelled. After unsuccessful attempts to make the captioner capture it, Krea filters out or undersamples this kind of data. Data quality includes whether the supervision can explain the image, not just whether the image looks good.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Spend expensive judgments where they matter
Lee describes image-training collections ranging from two to ten billion images. At that scale, even a reasonable per-image operation becomes consequential. Krea begins deduplication with inexpensive hashes, including pHash and MD5, before applying embedding methods to a reduced collection.
The later stage uses SSCD and SigLIP representations to address near duplicates and semantic redundancy. These representations serve different purposes: SSCD is designed for image-copy detection, while SigLIP supplies language-image representations useful for semantic comparisons.
| Stage | Representation | Role |
|---|---|---|
| Initial pass | MD5, pHash | Cheap duplicate reduction |
| Reduced corpus | SSCD descriptors | Copy and near-copy detection |
| Reduced corpus | SigLIP embeddings | Semantic comparison |
The sequence reserves more expensive comparisons for the images that survive the inexpensive pass.
Classifier construction follows a related pattern. First, use a large vision-language model to establish a useful judgment—for example, whether an image appears AI-generated. Refine that behavior through a system prompt or fine-tuning, then distill its decisions into a small SigLIP classifier. The large model helps define the filter; the small model applies it across the corpus. Lee describes the resulting classifier as somewhat reliable, with an inference cost suitable for billions of images.
Lee connects this approach to Essential-Web v1.0: a large language model supplies taxonomy or annotation judgments that a compact model can reproduce at pretraining-corpus scale. He cites an approximately 500-million-parameter annotation model. The point is to distill the curation decision, avoiding the cost of running the largest model over every document or image.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Find unwanted features—and missing concepts
Sparse autoencoders provide another route to dataset tags. Lee’s earlier work on CLIP and other vision representations suggests a practical use: train an SAE on visual features, feed it an image and inspect the sparse features that activate. His example activates features corresponding to a horse, black-and-white appearance and blur. These become an unsupervised tagging system without requiring a manually specified label for every training image.
Once a useful feature exists, it can guide filtering or oversampling. Signatures, watermarks and border artifacts are examples of unwanted properties that an SAE feature might capture. The condition is essential: the method can filter on a property only if the learned representation exposes a suitable feature. It can also preserve or increase a desired style instead of automatically treating unusual visual properties as defects.
Coverage requires the opposite question: what should be present? Krea uses Wikipedia concepts to check its dataset’s world knowledge, although Lee says he does not know how much this procedure helped. Articles are ranked by PageRank, and selected concepts are checked against the image collection using plain-text or embedding search. The accompanying technical account clarifies the selection as the top 90% of English Wikipedia articles by PageRank, followed by representability filtering. Lee relates this to CLIP’s broad concept coverage; CLIP’s documented query construction used Wikipedia frequency and article search volume rather than this PageRank procedure.
This search also uncovered Barack Obama the horse. Lee does not establish that this particular horse entered the training set and thinks it probably did not. The example nevertheless shows why checking concept coverage can reveal more than familiar names. Across the overall curation process, Lee reports approximately 30–40 in-house classifiers, heuristics and filters. Together they determine which visual distribution the otherwise fixed architecture can learn.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Learn semantics first, then shape the distribution
With the data prepared, Krea follows a sequence familiar from language-model development: pretraining, mid-training, supervised fine-tuning, preference optimization and reinforcement learning. A separately trained prompt expander accompanies the image model. Resolution also follows a curriculum: learning what a horse looks like does not initially require the pixels needed to render its finest details.
Krea trains progressively from 256 to 1K resolution. Lower-resolution training establishes text-to-image semantics; higher-resolution stages develop structure and detail. The resulting pretrained model is a malleable base, much as an autocomplete language model is a base for later chat or tool-use training.
Mid-training and supervised fine-tuning then shape that base toward intended applications. Krea curates illustration, graphic design, photography and cinematic imagery, using larger mid-training collections and more focused SFT data to mold the output distribution. Pairwise preference judgments subsequently polish the result. These comparisons make the process more opinionated: choosing one image over another teaches the model which outputs the team wants it to favor.
Reinforcement learning adds feedback on generated images. Krea uses a GRPO-inspired method: the model generates images, sends them to reward servers and learns from the returned feedback. Text rendering and anatomical structure are examples of capabilities targeted by this loop. Lee describes the flow and its objectives, without supplying a reward formula or measured improvement.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Translate user intent into training-distribution language
A production system must also bridge the gap between a short user request and the detailed descriptions used during training. Krea trains a small LLM to turn the user’s prompt into a longer, more detailed prompt. Lee describes this as an almost essential production step: descriptions closer to the image model’s training distribution tend to produce better images.
For example, a prompt-expansion training pair can preserve a simple subject while adding visual detail. The following Python record illustrates the input/output contract with a horse prompt:
python
import json
training_pair = {
"user_prompt": "a horse in a field",
"expanded_prompt": (
"A photograph of a horse standing in a grassy field, "
"with its full body visible, natural daylight, "
"and an open landscape in the background."
),
}
print(json.dumps(training_pair, ensure_ascii=False, indent=2))
The learned component produces the expanded description; the image model consumes that description as conditioning. This additional specificity helps explain the distribution match, but also makes prompt expansion a place where the system chooses details the user left open.
Beyond prompt expansion, Krea is working on multi-expert on-policy distillation. Separate experts specialize in capabilities such as photography or text rendering, then transfer those capabilities into a single student. Matching each specialist in its area is the goal of this ongoing work, rather than a completed result reported in the talk.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Make the next experiment cheap
The useful research levers are not limited to a particular optimization method. Infrastructure that permits fast iteration makes it easier to discover which data and training choices matter. Curated data has a longer useful life than a particular implementation: code and fashionable training methods change, while a good collection can remain valuable across those changes.
Lee favors simple, scalable methods with few hyperparameters to tune. Efficiency matters because it shortens the experimental cycle. Borrowing from LLM research also provides more than conceptual inspiration: it creates opportunities to reuse kernels, implementation work and an extensive body of literature. These choices reduce the cost of trying the next idea.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Simpler models, richer descriptions
Prompt expansion creates an architectural echo. In Lee’s comparison, an autoregressive decoder first expands the text, followed by an encoder-like diffusion transformer—a reversal of the familiar encoder–decoder ordering. It also resembles DALL-E 2 at the level of a model producing conditioning for another generative model. The analogy has a precise limit: DALL-E 2’s prior generates a CLIP image embedding, not an expanded text prompt, and its diffusion image generator is called a decoder.
That growing stack motivates a different aspiration: remove the VAE and separate text encoders, and train one clean transformer. Lee presents this as a direction he wants to pursue, not the architecture already used for Krea 2.
At the same time, stronger vision-language models can make the conditioning richer. The original latent-diffusion work already included bounding-box conditioning; Lee argues that generating useful boxes for training images is now much more practical. He names Ideogram and Reve as pursuing related directions. Boxes give the generator explicit spatial information that a conventional caption may leave implicit.
Scene graphs extend the idea from object locations to relationships among objects. Lee points to earlier work from Fei-Fei Li’s lab, while the displayed figure contrasts sentence and scene-graph conditioning leading to generated images. His broader proposition is that image-generation progress tracks progress in VLMs: better visual understanding can produce better descriptions to train and condition the generator.
The remaining question is which representations should carry that understanding. Bounding boxes and scene graphs are promising because they describe aspects of an image that prose may omit or express ambiguously. The technical ending returns to the same issue exposed by the painting on the wall: what the training system can describe determines what the image model can learn to control.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Resources
From the talk
Official inference instructions and links to RAW and Turbo checkpoints, with guidance for sampling and fine-tuning.
A web-corpus curation system that uses a compact model to annotate documents with a structured taxonomy.
Language-supervised visual representation learning, including the construction of a broad concept-query list for collecting image-text pairs.
Pretrained image-copy descriptors and instructions for embedding extraction, similarity measurement and evaluation.
The original SigLIP method for training image-text representations with a pairwise sigmoid objective.
The foundational latent-diffusion paper, including compressed image representations and text or bounding-box conditioning.
A two-stage image generator combining a text-conditioned image-embedding prior with a diffusion decoder.
Further reading
- Krea 2 Technical ReportArticle
The team's detailed account of data curation, architecture, staged training, prompt expansion and infrastructure.
Johnson, Gupta and Li Fei-Fei's method for generating images from object relationships through predicted layouts.
Read the complete timestamped transcript
6 slide groups. Extracted from the recording; expand builds to follow changes. Timestamps are approximate. 28 additional groups were omitted from this pass.
- 0:00
[upbeat music] Okay. So, uh, I'm Sangwu.
- 0:15
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.
- 0:28
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.
- 0:43
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.
- 0:56
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.
- 1:09
One thing we focused quite a bit was, like, stylistic, uh, diversity.
- 1:15
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.
- 1:28
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.
- 1:43
So I'll go ahead and then start talking about that.
- 1:48
So one of the trend that we've been, like, seeing recently is that, you know, existing models like ChatGPT, uh, ChatGPT, Nano Banana Pro, like the, like the production grade, uh, models that are, like, from the big labs.
- 2:00
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 Nano Banana Pro, I think they take up to like a minute or like two to, like, give an output.
- 2:13
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
- 2:31
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.
- 2:40
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
- 2:56
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.
- 3:07
In that case, like ChatGPT 2 or Nano Banana 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.
- 3:20
That's something that we wanted to, uh, focus on a little bit. So, you know,
- 3:27
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.
- 3:35
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?
- 3:48
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.
- 3:57
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.
- 4:08
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.
- 4:22
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.
- 4:33
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.
- 4:48
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.
- 5:03
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.
- 5:16
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...
- 5:26
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 reiterate this more. That's why I put like again, like really data is quite like everything.
- 5:37
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.
- 5:49
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-
- 6:00
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.
- 6:14
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.
- 6:31
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.
- 6:46
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
- 7:08
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
- 7:25
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.
- 7:43
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, Nano Banana aesthetic.
- 7:57
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, Nano Banana Pro. And in the long run, like, that's just not how you want to, like, do.
- 8:07
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.
Two stage captioning pipeline
OCR extracts visible text, and optional metadata helps a vision-language model produce a detailed caption. The caption is then rewritten into prompts in the desired format.
Text in image
Two stage captioning pipeline Accurate text and metadata information IMAGES OCR OPTIONAL METADATA DETAILED CAPTION REWRITE PROMPTS
- 8:17
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.
Two stage captioning pipeline
OCR extracts visible text, and optional metadata helps a vision-language model produce a detailed caption. The caption is then rewritten into prompts in the desired format.
Text in image
Two stage captioning pipeline Accurate text and metadata information IMAGES OCR OPTIONAL METADATA DETAILED CAPTION REWRITE PROMPTS
- 8:28
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.
- 8:41
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,
- 8:58
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
- 9:19
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.
- 9:29
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.
- 9:41
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
- 9:57
this kind of, like, data out or at least undersample it.
- 10:01
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.
- 10:12
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.
- 10:22
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.
In-house classifiers
Krea distills decisions from a fine-tuned or prompted large vision-language model into a small SigLIP classifier, making it cheaper to filter images at billion-image scale.
Text in image
In house classifiers Distill to < 1B models Large VLM Finetune / System Prompt SigLIP Distill to small classifier
- 10:37
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?
- 10:52
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,
- 11:04
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
- 11:21
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 dataset used.
In-house classifiers
Krea distills decisions from a fine-tuned or prompted large vision-language model into a small SigLIP classifier, making it cheaper to filter images at billion-image scale.
Text in image
In house classifiers Distill to < 1B models Large VLM Finetune / System Prompt SigLIP Distill to small classifier
- 11:31
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,
- 11:46
like, a pre-training level corpus. Otherwise, it'll be somewhat expensive and frankly inefficient use of your GPUs to do so.
SAE filters
A sparse autoencoder trained on a vision model turns the horse image into activated features such as horse, black and white, and blur. Those features can help filter or oversample images.
Text in image
SAE filters Unsupervised tagging system SAE Horse Black and white Blur
- 11:55
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
- 12:13
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.
- 12:23
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.
SAE filters
A sparse autoencoder trained on a vision model turns the horse image into activated features such as horse, black and white, and blur. Those features can help filter or oversample images.
Text in image
SAE filters Unsupervised tagging system SAE Horse Black and white Blur
- 12:35
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.
- 12:49
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
- 13:01
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
- 13:18
data set. And then one thing that I think, like, people have quite liked about our, uh, open source Krea 2 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,
Knowledge coverage
Krea selects Wikipedia concepts in the top 90% by PageRank, then searches its dataset by text or embeddings to check coverage. The speaker is unsure how much this improved world knowledge.
Text in image
Knowledge coverage PageRank based wikipedia coverage top 90% PageRank concepts Search for wikipedia concepts
- 13:36
is that you can actually take the Wikipedia, the entire Wikipedia, and for each article or concept, you can compute the PageRank 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, PageRank on Wikipedia, it's probably something quite important for the model
- 13:56
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.
- 14:11
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.
Knowledge coverage: Wikipedia example
While checking Wikipedia concepts for dataset coverage, the team found an article about a horse named Barack Obama. The speaker does not think this horse was in the dataset.
Text in image
Knowledge coverage PageRank based wikipedia coverage Barack Obama (horse) From Wikipedia, the free encyclopedia Barack Obama (1998[1] – 14 September 2018) was an Anglo-Arabian horse from New Zealand that competed in international endurance events beginning in 2009, the same year that Barack Obama became president of the United States. Racing career [edit] In 2014, Jenny Champion began to ride the horse, and the team had fifteen starts and six wins.[2][3] In 2018, Champion began to fundraise to be able to participate in the 2018 FEI World Equestrian Games, with a goal of NZD$80,000. She trained with Barack Obama constantly, making sure that he was well-groomed after each ride.[4] Together they represented New Zealand at the games in Tryon, North Carolina. The 120-kilometre race that Barack Obama participated in was abandoned due to poor weather conditions, and he was taken into the Tryon International Equestrian Center to be treated for kidney issues. Shortly thereafter, Mark Round, Barack Obama's owner, elected to euthanize the horse.[2][3][5][6] Following the death of Barack Obama, samples were taken from his carcass.[2][3][7] Death [edit] According to Sarah Dalziell-Clout, the Equestrian Sport NZ High Performance Director, Barack Obama's death was unexpected. She said that the horse was in good health prior to the race, and that the team veterinarian, Nick Page, was confident in Barack Obama's readiness.[8] FEI President Ingmar De Vos, as well as the FEI Veterinary Director and President of the Endurance Veterinary Commission, held a press conference stating that the race was abandoned over concerns for the health of the horses, with high heat and humidity causing a large number of the horses present to show signs of distress.[9] Of the 95 starters in the race, 50 saw veterinary treatment, including Barack Obama.[8] A veterinary report revealed that the horses treated were severely dehydrated from excessive sweating, and required fluid therapy.[10] Barack Obama Jenny Champion riding Barack Obama Species: Horse Breed: Anglo-Arabian Sex: Gelding Born: 1998, New Zealand Died: 14 September 2018 (aged 20), Tryon, North Carolina, US Cause of death: Animal euthanasia Occupation: Racehorse Years active: 2009–2018 Owner: Mark Round Interface text: 1 language Article Talk Read Edit View history Tools
- 14:27
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 end of the day, like, we have, like, around, like, 30 to...
Filters, heuristics, and classifiers
Krea used 30 to 40+ filters, heuristics and classifiers. The speaker says that, once the architecture is chosen, data determines model quality.
Text in image
30 ~ 40+ filters, heuristics, and classifiers Data is the engine that shapes the model distribution.
- 14:38
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.
- 14:47
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:01
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.
- 15:24
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.
- 15:34
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.
- 15:45
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:58
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.
- 16:08
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.
- 16:20
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.
- 16:33
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?"
- 16:51
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.
- 17:06
And then, you know, reinforcement learning is also, like, kind of, is now extremely standard in LLMs. That's what we do in diffusion land too. So we have, like, a pretty much, like, a GRPO inspired, uh, method where,
- 17:20
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.
- 17:35
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.
- 17:46
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
- 18:05
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.
- 18:17
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.
- 18:30
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.
- 18:40
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.
- 18:51
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.
- 19:03
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.
- 19:11
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.
- 19:25
Now we are going the opposite way because we have a prompt expander which is autoregressive decoder, and then diffusion which is the encoder, so it's kind of getting reversed.
- 19:34
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, autoregressive model or, like, diffusion model.
- 19:45
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.
- 19:55
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.
- 20:09
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 LDM 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
- 20:27
on this. Uh, this something that Ideogram, Reve 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.
- 20:38
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...
- 20:51
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.
- 21:03
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.
- 21:14
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.
- 21:25
But yeah. [audience applauding] [upbeat music]