Krea 2: クリエイティブな探索のための基礎モデルとインフラストラクチャ

イントロダクション

近年、画像生成分野は顕著な進展を遂げてきました。拡散モデルやフローマッチングモデルを用いることで、以下のような能力が可能になりました。

• 高解像度の画像生成
• 鮮明でフォトリアルな質感の再現
• 安定した構造化と高密度テキストのレンダリング
• 広範な世界知識の暗黙的獲得
• ユーザーのプロンプトに対する精密な従順性

これらの進歩は、スケーラブルなトランスフォーマーアーキテクチャ、改良されたキャプション化およびテキストエンコーダー、向上した潜在表現、並びにパイプライニングされた事後学習技術など、複数の相互作用する要因によって支えられてきました。

しかしながら、この分野が信頼性の最適化を優先した結果、多くのシステムは限られた数のデフォルト美学的特徴に収斂してしまいました。これは効果的な製品ツールとしては優れていますが、クリエイティブな探索のためのエンジンとしての有効性を低下させています。ユーザーが単一の磨き上げられたデフォルトを受け取るのではなく、スタイル、気分、構成、そしてビジュアルの方向性などを探求し検索する必要がある分野において特にそうです。

これらの制限に対処するために、私たちは**クリエイティブな探索に焦点を当てた基礎モデル「Krea 2」**を発表します。Krea 2 のモデルは、以下の信念に基づいて構築されています。

• 「画像生成は表現力があり多様な美学的特徴をカバーできる一方で、クリエイターがそれらを自由にナビゲートできるよう制御可能な媒体であるべき」

開発アプローチ

• データインフラストラクチャとフレームワークのゼロから構築: 広範な世界知識とスタイルの覆盖を持つ包括的な事前学習データをカuration（選定）しました。
• 多段階パイプライネルによるトレーニング:
  • 事前学習 → 中間学習 → 教師ありファインチューニング (SFT) → 好ましさ最適化 (Preference Optimization) → 強化学習 (RL)
  • 各段階は、モデルの出力分布を順次洗練させることを目的として設計されています。

アーキテクチャと技術的特徴

• シンプルな高性能拡散トランスフォーマー (DiT) アーキテクチャの開発（徹底的なアブレーションの実施を通じて）。
• 収束加速コンポーネントの統合:
  • iREPA（初期学習加速）
  • 改良された VAE（変分オートエンコーダ：Qwen Image VAE & FLUX 2 AE）
  • Qwen3-VL（テキストエンコーダ）
• トレーニングの安定性と効率性を向上させるアーキテクチャ的改良:
  • グループ化クエリ注意機構 (GQA)
  • シグモイドゲート付き注意機構
  • 軽量タイムステップモジュレーション
  • テキストエンコーダー特徴のための多层階特徴集約

プロンプトとスタイルの制御システム

堅固なベースモデルであっても、ユーザーがその分布の関心のある部分を確実にアクセスできる限りのみ有用です。テキスト入力は短く曖昧であり、画像生成モデル学習時に使われる豊かなキャプションとの間にギャップが生じます。このギャップを縮小するため、以下の 2 つのシステムを構築しました。

1. プロンプトエクスパンダー

   • ユーザーの意図を書き換えずに、単純または未指定のプロンプトをより豊かなビジュアル方向性にマッピングします。
   • オープンソースの LLM をベースとし、二段階の SFT および RL パイプライネルでトレーニングされました。
   • 目的：画像品質向上だけでなく、クリエイティブな多様性を促進し、制御可能な探索を奨励すること。

2. スタイル参照システム

   • 言語が不十分な場合に、画像を通じてビジュアルな意図を表現することを可能にします。
   • コリー（内容漏洩）を防ぎながら、参照画像のスタイルや雰囲気を注入します。
   • スタイル強度に対する細やかな制御および加重されたスタイルミキシングを提供します。

成果: Krea 2 は Artificial Analysis ランキングでテキスト対画像タスクにおいて上位 10 モデルに入っており、独立系ラボからのモデルの中では2 位という成績を収めています。

---

データ

データカレーションの原則

「良いデータミックス」は単に高品質な画像のみから構成されるわけではありません。多様性と広範なドメインカバーが不可欠です。従来のモデルベースのフィルタリング（審美点数や IQA モデル）は暗黙的なバイアスを導入する可能性があるため避けています。

事前学習データセットへの適用条件:

• 重複したサンプルと過度に代表的な概念の排除。
• VLM が画像の重要な側面を一貫して捉えられないサンプルの排除。
• 望ましくないバイアスやアーティファクトを引き起こすサンプルの排除。
• 視覚的に複雑すぎてローレスOLUTION で信頼性を持ってモデル化するのが難しいサンプルの排除。
• AI 生成画像のサンプルの排除（合成データはモデルの出力分布にバイアスを招くため）。

キャプション生成プロセス:

1. OCR の実行: 各ターゲット画像から可視テキストを抽出します。
2. 増強キャプションの生成: OCR 結果とメタデータをカプチュンモデルに提供し、世界知識を取り入れたキャプションを生成させます。
3. 長形式への変換: 安価な LLM を使用して、広範なプロンプトスタイルへの曝露を確保します（高密度な監督と Fast 収束を実現）。

解像度による学習戦略

• 低解像度事前学習 (256px): テキスト対画像の整列性と構造の学習。安価な CPU ベースのフィルタで低品質画像を除去（ラプラスフィルタ、破損ファイルなど）。
  • 例：均一色背景や境界アーティファクトの生成を防ぐため、社内分類器や VLM によるフィルタリングを実装。
• 中間学習 (512px): トップダウンのカリキュレーション。特定のスタイルカバレッジと高品質画像ソースを明示的に選択。セマンティッククラスタリング（FAISS 使用）を行い、ロングテールの視覚概念を保持。
• 高解像度事前学習 (1024px): b16 ベーストレーニングへ移行し、より細粒度なスケーリングを使用。

SAE (スパースオートエンコーダー) の活用:

• SigLIP-2 エンベッディングに基づいて SAE をトレーニング。
• VLM による注釈付き特徴を用いた非監修タグシステムで視覚アーティファクトをフィルタリング。

エンティティカバレッジの確保:

• Danke を使用して英語版ウィキペディア上で PageRank を実行し、上位 90% の記事を保留。
• Wikidata メタデータに基づく全記事のフィルタリングと全文検索による評価。
• 稀なコンセプトを持つ画像を優先的にサンプリング。

SFT データ

個別のビジュアルドメインに焦点を当てた小規模な手動でキュレーションされたデータセットを使用します。品質がボリュームよりも重要です。

---

アーキテクチャ

アーキテクチャアブレーションは以下の 4 つのカテゴリに分類されます：

1. 安定性: ロスを減らし勾配スパイクを減少させるか？
2. パフォーマンス: モデルの収束を早めるか？高解像度でのトレンド維持か？
3. 効率性: パラメータ数、FLOPs、メモリ使用量を減らすか？
4. シンプルさ: モデルを単純化できるか？

コンポーネント	ベースライン	アブレーション	最終コンポーネント
注意機構	Multi head attention	GQA, MLA, ゲート付きシグモイド注意	GQA + ゲート付きシグモイド注意
MLP	GeLU MLP	SwiGLU	SwiGLU
リジューダブル接続	スタンダード	値リジューダブル、Laurel	スタンダードリジューダブル接続
テキストエンコーダー	T5-XXL	T5, Gemma, Qwen 2.5 VL, Qwen 3 VL	Qwen 3 VL
モジュレーション	ブロックごとの MLP モジュレーション	バイアス付き軽量モジュレーション	バイアス付き軽量モジュレーション
オートエンコーダー	FLUX AE	Qwen Image VAE, DC-AE, FLUX 2 AE	Qwen Image VAE & FLUX 2 AE
ブロック設計	シングルストリーム	ハイブリッド、並列シングルストリーム	シングルストリームトランスフォーマーブロック
ノーマライゼーション	LayerNorm	RMSNorm, Zero center RMSNorm	Zero center RMSNorm, QKNorm
位置エンコーディング	3D Axial RoPE	Golden Gate, MRoPE, Partial RoPE	3D Axial RoPE

主要な設計決定

トランスフォーマーブロック

• MLP の置換: GeLU MLP を SwiGLU に置換。一貫したパフォーマンス向上のため採用。
• 注意機構の選定: GQA は計算効率を向上させ、ゲート付きシグモイド注意はトレーニングダイナミクスを安定化させるため、両方を組み合わせました。
• ストリーム設計: シングルストリーム、デュアルストリーム、ハイブリッドストリームの比較を行いました。パフォーマンスの差は小さく、シンプルさのためシングルストリームを採用しました。

タイムステップ条件付け

• ブロックごとの MLP を、ブロックごとのチューナブルバイアスに置換。パラメータ効率を向上。
• 代替案として「タイムステップ情報の完全削除」や「トークンによるインコンテキスト条件付け」を検討しましたが、高解像度での性能劣化を防ぐためAdaLN ベースの維持を選択しました。

位置エンコーディング

• 3D Axial RoPEを使用。テキストトークンに対して RoPE インデックスをゼロに設定。
• Partial RoPE は低解像度では効果的でしたが、高解像度トレーニングでは最終的にパフォーマンスが劣化するため採用されませんでした。

オートエンコーダー

• DC-AE は復元誤差のために細部解決能力の上限を引き起こすため除外。
• Qwen Image VAEとFLUX 2 VAEは高速収束と優れた復元品質を提供するため採用。

リジューダブル設計とノーマライゼーション

• スタンダードなリジューダブル接続を採用（表現力向上のため）。
• **RMSNorm (Zero Center)**をデフォルトに採用。LayerNorm からの切り替えは品質劣化が極めて少ないため。

テキストエンコーダー

• Qwen 3 VLを採用。VLM が豊かな入力空間と強力な多言語一般化を提供するため。
• 浅い注意層を導入し、自律的 LLM の最後のレイヤーの特徴（画像生成に不適切）を集約せず、画像からの特徴も活用できるように設計。

---

最適化

• 主要アルゴリズム: パイライン全体でAdamWを使用。
• Muon の検証: 初期ステップでの高速収束は見受けられましたが、長期軸ではパフォーマンスが下回る傾向があり、安定性にも課題があったため、最終的には AdamW で維持しました（一部改良版の結果は将来のサイクルに活用予定）。

---

トレーニング

事前学習

• 解像度を 256px → 512px → 1024px に順次スケールアップします。
• iREPAを使用し、256px 段階で初期収束を加速。その後除去。
• PMA (Post-training Merge Adaptation) を適用し、メモリオーバーヘッドを増やさずに EMA に匹敵するパフォーマンス達成。

中間学習

• SFT ステージ前のモデル分布のウォームアップとして機能。
• 高忠実度生成と強力なドメインカバレッジの確保。

教師ありファインチューニング (SFT)

• 高審美的画像の専用データセットを使用。
• モデルを美的に望ましい方向にバイアス付け。
• チェックポイントマージを行い、一般化された SFT チェックポイントを生成。

好ましさ最適化 (PO)

• ポリシー分岐の緩和: 従来の DPO では勝者と敗者の両方が低品質になり得る問題があり、これに対処するために STPO のバリアントを開発。

強化学習 (RL)

• 多報酬 GRPO スタイル:
  1. 一般的な審美モデル（PO データでファインチューニング）
  2. プロンプト追従報酬
  3. テキストレンダリング報酬
  4. アーティファクトと構造報酬
• プロンプト固有の rubric 報酬: ユーザー意図の変動に柔軟に対応するため、各プロンプトを検証可能な要件に分解して評価。
• アーティファクト報酬モデル: 構造アーティファクト（余分な指など）を検出し、視覚的品質を犠牲にしたベンチマーク信号の改善を防ぐ。
• 分類器フリー (CFG) なしトレーニング: トレーニング初期に迅速に適応するため、ロールアウトおよびトレーニングは CFG なしで行い、推論時に必要な場合は有効化。

タイムステップ蒸留

• TDM (Trajectory Distribution Matching) を採用。シンプルで最小限のハイパーパラメータを持ちつつ、柔軟なマルチステップ生徒を育成可能。

---

プロンプト展開

高密度プロンプトはより良い画像を生み出しますが、ユーザーは短いプロンプトを使用します。これを分布マッピング問題として設定：

• システム: ユーザー意図を理解し、入力プロンプトを豊かなモデルフレンドリーキャプションに変換。
• 手法: LLM 上での SFT と RL を組み合わせ。
  • 合成データで「短いプロンプト → 拡張されたキャプション」ペアを生成。
  • RL を用いて直接エクスパンダーを最適化（イメージ報酬とプロンプトレベル検証可能な報酬の両方を使用）。
• 安全性: エクスパンダーが未プロンプトなコンテンツを導入しないように安全制約チェックを実装。
• 多様性崩壊の防止: DINOv3 エンベッディング多様性スコアを追加し、グループ内の視覚的多様性を報酬。

---

スタイル参照システム

ベースモデルの上に構築され、以下の機能をサポート：

1. 複数スタイルの滑らかなセマンティックミキシング
2. 各スタイル参照に対する強度の連続制御
3. 複雑なスタイルへの最先端適合

トレーニングには新規自己教師なし技術を開発し、好ましさ最適化ステップで出力を整列させました。

---

トレーニングインフラストラクチャ

分散トレーニングフレームワーク

• ゼロから PyTorch に基づいて構築。DTensor と torchtitan を使用。
• FSDP2と Megatron-LM スタイルのテンソル並列性を併用。
• torch.compileを主要最適化戦略として採用（cuDNN デフォルト、FlexAttention/FlashAttention 3 オプション）。

データ読み込み

• Parquetフォーマットを使用。
• データを事前シャッフルし、グローバルシャッフルを保証。
• ノード内接続 (NVLinkSharp)、ノード間接続 (InfiniBand) を使用。

信頼性とフェイルオーバー

• Weka ファイルシステム: Ceph から移行し、高いパフォーマンスと安定性を実現。チェックポイント時間は約 30 秒。
• 故障ノード管理: 故障ノードリストからソフト/ハードアンチ親和性を適用し、健全なノードでのみトレーニングポッドをスケジュール。
• 可観測性: DCGM メトリクス、InfiniBand エラー、PCIe リプレイカウンターなど細粒度の監視を行い、早期に故障を検知・対策。

システムインフラストラクチャ (Kubernetes)

• Kueueによるワークロードスケジューリング（優先度管理）。
• VK (Virtual Kubelet) を用いたクラスター外推論スケーリングの実装。
• 訓練と推論の共有プールを管理し、トレーニング時の容量調整を自動化。

---

データインフラストラクチャ (Krablet)

Krea 2 のデータ取り込み・カレーションには、PostgreSQL サーバー群を回りの**「Krablet」**システムを採用しました。

• 構成:
  • データの一片を保持する Postgres インスタンス。
  • 「Funnel」サーバー（非同期的なバッチ/キュー処理でロック競合を最小化）。
  • RPC サーバー群（PgBouncer 代替による接続管理）。
• 特徴:
  • メタデータのみで 208 TB スケール。
  • 高速な UPSERT トランザクション処理能力。
  • データの単一真理源 (Single Source of Truth) としての役割。

DAG ベースの作業フロー:

1. OCR ワーカー: 未処理画像を検出し、テキスト抽出を実行。
2. Embed ワーカー: OCR が完了した画像に対し、埋め込みを生成。

   • FOR UPDATE SKIP LOCKED

     セマンティクスを用いて、待機キューを自動管理。

利点:

• 再試行: 失敗した行のみをキュー末端で自動的に再試行。
• フェイルトランレンシ性: ワーカーのクラッシュが全体ジョブに影響しない。
• 動的ワーカー数: 需要に応じて自動スケーリング。
• 即時可視性: リアルタイムでの進捗確認。

---

議論と将来の仕事

スケール

次のサイクルでは、以下の方向でアーキテクチャと最適化を進める予定です。

• MoE の導入によるスパース注意の使用。
• ネイティブ 2K〜4K 解像度へのスケール。
• NVFP4 を用いた事前学習。

マルチティーチャーオンポリシー蒸留 (MOPD)

• 異なる専門家を単一生徒に蒸留する手法。組織的な視点から非常にスケール可能（ドメイン固有能力の分離と回帰）。

アーキテクチャ簡素化

• 現在、多数のコンポーネント（AE, DiT, Encoder など）が必要ですが、将来はこれらを単一モデル下に統合し、研究チームの調整を容易にします。

新しい機能

• 画像参照とネイティブ 2K/4K 生成のサポート。
• より幅広いプロンプトスタイル（JSON、バウンディングボックス、Markdown など）への対応。

---

結論

Krea 2—クリエイティブな探索のために構築された我々の最初の基礎モデルシリーズ—を発表しました。この技術報告では、モデルファミリーの背後にあるインフラストラクチャ、データシステム、トレーニングパイライン、および研究選択を共有しました。基礎画像モデルはまだ初期段階にあるため、さらに多くの研究を共有することを期待しています。