Training Krea 2: What matters in generative model training — Sangwu Lee, Krea.ai
Summary
Sangwu Lee details the training methodology for Krea 2, an open-sourced image foundation model, emphasizing that data curation is the most critical factor after architecture is locked. The talk contrasts production models (like ChatGPT-2) which prioritize consistency and reliability but suffer from mode collapse, with Krea 2's focus on stylistic diversity and fast iteration. Key techniques covered include advanced deduplication using hash-based methods (pHash/MD5) followed by embedding-based approaches (SigLip), specialized captioning pipelines to capture contextual details (e.g., a painting being framed), and multi-stage training from low to high resolution.
Key takeaways
-
Data is Paramount for Model Quality
9:59
After locking in the architecture, data curation determines model performance. The goal was to maximize stylistic diversity while filtering out undesirable content (e.g., images where captions consistently fail to capture important context) and avoiding reliance on synthetic AI-generated training data due to its 'sticky' aesthetic.
-
Balancing Consistency vs. Diversity
2:20
Large production models often achieve high consistency by focusing on the 'average' subject (e.g., a boring average person), leading to mode collapse and limited stylistic range. Krea 2 was designed to optimize for fast generation and broad stylistic exploration, catering to creative studios that are still defining their visual goals.
-
Advanced Deduplication and Filtering
16:55
Deduplication runs first on basic hashes (pHash or MD5) across billions of images, followed by embedding-based methods like SigLip to remove near duplicates. Filters are also applied using large vision language models (VLMs), which are then distilled into cheap classifiers for scalable filtering over massive datasets.
-
Multi-Stage Training Pipeline
The training process progresses through several stages: low to high resolution pre-training (learning semantics first, then structure/detail), supervised fine-tuning (SFT), preference optimization (collecting user pairs for comparison), and Reinforcement Learning (RL). The final step involves training a prompt expander LLM to generate detailed prompts that align with the model's trained data distribution.
Technical details
-
Latent Diffusion Models (LDMs)
323s
Diffusion models are trained by adding noise to an image and teaching the model to progressively remove it. Most competitive open-source models use LDMs, which employ an autoencoder to spatially compress raw pixels into a latent space for efficient training, mitigating the high time complexity associated with modeling every single pixel.
-
Sparse Autoencoders (SAEs)
1156s
SAEs can be used as an unsupervised tagging system. By feeding an image through a trained SAE, the activated sparse features can identify specific attributes (e.g., 'horse,' 'black and white,' 'blur'), allowing users to filter or oversample data based on these learned signatures, useful for removing watermarks or border artifacts.
-
World Knowledge Integration
1200s
To improve world knowledge coverage, the process involves taking the entire Wikipedia corpus. Concepts are ranked by PageRank, and keywords from high-ranking articles are used to ensure their presence in the training dataset via standard text or embedding search.
-
Training Resolution Scaling
The model is trained progressively, starting at low resolution (e.g., 256) to learn basic semantics and then scaling up to high resolutions (e.g., 1K) to capture structure and fine detail.
Mentioned resources
Channel & topics
Watch on YouTube · Back to latest
This independent, AI-assisted summary is provided for commentary and informational purposes. It may contain errors or omit important context. Please watch the original video for the creator's complete presentation. Video, thumbnail, and related copyrights belong to their respective owners.