Tokenless

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

The talk by Sangwha Lee discusses Krea 2's open-source medium variant, emphasizing stylistic diversity and rapid iteration over the consistency-focused approach of larger models. A significant portion details their robust data curation pipeline, including unique methods for deduplication, filtering out AI-generated images, utilizing sparse autoencoders for unsupervised tagging, and ensuring world knowledge coverage. He outlines an LLM-inspired multi-stage training process, culminating in a prompt expander, and shares insights on fast iteration and future directions for image generation, highlighting the increasing integration of VLM advancements.

6 min read
Training Krea 2: What matters in generative model training — Sangwu Lee, Krea.ai
source: https://www.youtube.com/watch?v=-tviRdpmHvs

Sangwha Lee introduced Krea 2’s recently open-sourced medium variant, highlighting its core philosophy of prioritizing stylistic diversity and rapid generation over the “reliable, but slow and uniform” outputs of major production models like ChatGPT-2 or Nano Banana Pro. He illustrates this with the “boring average person” analogy: production models achieve consistency by generating the most generic, centrally framed output, while Krea 2 aims to support creative exploration, especially for studios that are still discovering their visual direction.

Diffusion Model Fundamentals

Lee provided a brief recap of diffusion models, explaining that they learn to progressively remove noise from an image. Most competitive production diffusion models, including Krea 2, employ a latent diffusion model architecture. This involves using an autoencoder to perform generation in a spatially compressed latent space, significantly improving training efficiency due to the O(N^2) time complexity of transformers on raw pixel data.

The Primacy of Data

A central theme reiterated throughout the talk is that data is “quite everything” once the model architecture is locked. Lee emphasized that the vast majority of effort shifts to data curation, especially when aiming for specific goals like stylistic diversity. This involves consciously avoiding overly aggressive filtering based on conventional aesthetic scores, which might inadvertently remove valuable, diverse data (e.g., low-resolution CRT videos that some users might find aesthetically pleasing).

Rigorous Data Curation and Filtering Pipeline

Krea 2 employs a sophisticated multi-stage data pipeline to ensure high-quality and diverse training data:

  • Rejection of AI-generated Images: Krea 2 actively refuses to train on AI-generated images. Lee stressed that while synthetic data might offer a shortcut, its “sticky aesthetic” quickly leads to models adopting a “ChatGPT or Nano banner aesthetic,” losing original diversity and individuality. He expressed a personal aversion to models whose core strength is merely distillation.
  • Addressing VLM Captioning Failures: The team discovered instances where Vision Language Models (VLMs) used for captioning consistently failed to capture crucial visual context. A notable example is a painting photographed on a wall: VLMs would describe the painting but omit the frame and white wall, leading models trained on such data to always render paintings as if they were hanging on a white wall. Such biases necessitated the design of custom filters to either remove or undersample problematic data.
  • Captioning Pipeline: Their internal captioning process involves:
    1. Running OCR for text rendering.
    2. Adding optional metadata (e.g., famous people).
    3. A second pass with their VLM for detailed caption generation.
    4. Rewriting captions into suitable formats (e.g., JSON prompts).
  • Large-Scale Deduplication: For datasets ranging from 2 to 10 billion images:
    1. Initial hash-based deduplication (pHash, MD5) handles exact duplicates at scale.
    2. Subsequent embedding-based methods (SSCD, SigLip) are used for near-duplicate removal on smaller, pre-filtered sets.
  • Distilling VLM Judgments into Cheap Classifiers: To efficiently filter vast datasets, Krea 2 distills decisions from large VLMs (e.g., whether an image is AI-generated) into smaller, cheaper classifiers (e.g., SigLip-sized models). This allows for reliable filtering across billions of images, a strategy borrowed from LLM literature for text data quality control.
  • Sparse Autoencoders (SAE) for Unsupervised Tagging: Lee highlighted the innovative use of Sparse Autoencoders (SAEs) as an unsupervised tagging system. Once trained on a vision model, an SAE can take an image and activate sparse features (e.g., “horse,” “black and white,” “blurry”). This capability is leveraged to identify and filter undesirable elements like watermarks, signatures, or border artifacts in the dataset.
  • World Knowledge Coverage: Inspired by the original CLIP paper, Krea 2 ensures adequate representation of “world knowledge” by analyzing Wikipedia concepts. Articles with high PageRank (top 90th percentile) are identified, and the corresponding concepts are verified to be present in the dataset through text or embedding search.
  • Overall, Krea 2 employs 30 to 40 custom in-house classifiers and heuristics for data filtering.

LLM-Inspired Training Pipeline

Krea 2’s training pipeline is heavily influenced by strategies developed for large language models:

  • Progressive Resolution Training: Starts with low-resolution pre-training (256x256) to teach the model fundamental text-to-image capabilities and semantics. This is progressively scaled up to 1K resolution to learn structure and fine details.
  • Mid-training and Supervised Fine-tuning (SFT): After pre-training, the model undergoes mid-training and SFT using curated data tailored for specific downstream use cases (e.g., illustration, graphic design, photography, cinematics), shaping its output distribution.
  • Preference Optimization: Analogous to ChatGPT’s feedback mechanism, this stage collects user preferences (“like this over this?”) to further polish the model and instill more opinionated stylistic choices.
  • Reinforcement Learning (RL): A GRPO-inspired method is used, where the model generates images, receives feedback from reward servers, and learns to improve aspects like text rendering quality or anatomical correctness.
  • Prompt Expander: An “almost essential” component for production-grade models, this small LLM takes a concise user prompt and expands it into a long, detailed prompt. This is crucial because longer, more descriptive prompts are typically more in-distribution with training data and lead to better image generation.

Key Learnings and Future Directions

Lee shared several insights critical for Krea 2’s development:

  • Fast Iteration: Emphasizing robust infrastructure, simplicity (low hyperparameters), and efficiency as paramount for rapid development cycles.
  • Leveraging LLM Research: Actively borrowing kernels, research, and literature from the LLM space to accelerate progress.
  • The Inverting Stack and DALL-E 2 Resemblance: He noted an interesting trend where the stack is “inverting” back towards architectures like DALL-E 2. The combination of an auto-regressive prompt expander (decoder) feeding into a diffusion model (encoder) mirrors DALL-E 2’s prior-decoder setup.
  • Simplifying Architectures: A personal goal is to simplify the diffusion model stack by ideally removing VAEs and text encoders to train a single, clean transformer.
  • Advanced Conditioning with VLMs: Lee expressed excitement about future directions involving more powerful Vision Language Models. These could provide richer conditioning signals for image generation, such as accurate bounding boxes or scene graphs, moving beyond simple text prompts. He views advancements in image generation as a proxy for VLM progress.
  • Multi-expert Positive Distillation: Krea 2 is exploring training specialized experts (e.g., for photography, text rendering) and then distilling their collective capabilities into a single “student” model that can effectively match the performance of each expert.

Share post

Stay In The Loop! Subscribe to Our Newsletter.

Get updates straight to your inbox. No spam, just useful content.

Recommended for You

Reinforce Adjoint Matching: Scaling Diffusion RL

Reinforce Adjoint Matching: Scaling Diffusion RL

Andreas Bergmeister presents Reinforce Adjoint Matching (RAM), a novel method for scalable reinforcement learning post-training of diffusion models. RAM sidesteps the complexity of existing RL approaches by deriving a consistency loss that directly corrects the pre-training target with a reward signal, eliminating the need for costly SDE rollouts or reward gradients. This approach achieves superior performance and up to 50x faster convergence on challenging image generation tasks like complex composability, text rendering, and human preference alignment.

Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal

Taking Reinforcement Learning Cross Datacenter — Nan Jiang, Modal

Nan Jiang introduces "Adam absorption" to revolutionize RL model synchronization. By exploiting finite precision serving and small Adam steps, less than 1% of served model weights actually change, allowing for 500MB patches instead of 500GB checkpoints. This enables a distributed "bulletin board" architecture, decoupling trainers from global rollout fleets and unlocking elastic, cross-region GPU capacity for RL.