Calibri: Diffusion Transformer のパラメータ効率型較正による生成品質の向上

本論文は、Diffusion Transformer (DiT) の生成タスクにおける性能を大幅に向上させるための、パラメータ効率の高い手法「Calibri」を提案しています。従来のモデル全体を微調整（Fine-tuning）するのではなく、DiT ブロックの出力に適用する学習済みスケーリング係数（スカラー）のみを最適化することで、生成品質の向上と推論ステップ数の削減を同時に実現しています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

1. 背景と問題定義

近年、Stable Diffusion 3 や FLUX に代表されるように、画像生成の分野では UNet アーキテクチャから Diffusion Transformer (DiT) への移行が進んでいます。DiT はスケーラビリティに優れていますが、その内部構造には以下の課題が潜んでいることが示唆されました。

• 層ごとの寄与の不均一性: DiT は一見均一なブロックの連続ですが、各ブロックが最終生成に寄与する度合いは均等ではありません。
• 既存研究の示唆: 一部の「重要な層（vital layers）」を除外すると出力が劇的に変化することや、逆に特定の層を無効化すると品質が向上するケースがあることが報告されています。
• 仮説: 標準的な DiT アーキテクチャの重み付けは最適化されておらず、ブロックごとの出力を単純にスケーリング（再重み付け）するだけで、モデルのパフォーマンスを大幅に改善できる可能性があります。

従来のアライメント手法（RLHF, DPO, DDPO など）は、モデル全体を微調整する必要があり、計算コストが非常に高いという問題を抱えていました。

2. 提案手法：Calibri

Calibri は、DiT の構成要素への寄与を較正（Calibration）し、生成品質を向上させるパラメータ効率型のアプローチです。

2.1 基本的な考え方

各 DiT ブロックの出力に、学習された単一のスカラー係数 $s$ を乗算します。

• 最適化問題: 報酬モデル（Reward Model）による評価値を最大化するよう、係数ベクトル $c$ を最適化するブラックボックス最適化問題として定式化されます。
  $$c^* = \arg \max_{c} R(c)$$
• パラメータ数: 最適化するパラメータは非常に少なく、モデルサイズによらず約 $10^2$ 個（数百個）の係数のみです。

2.2 較正の粒度（Granularity）

内部層の較正には、以下の 3 つの粒度が提案されています。

1. Block Scaling: 1 つのブロック内の Attention と MLP レイヤー全体に共通のスカラーを適用（粗い調整）。
2. Layer Scaling: ブロック内の各レイヤー（Attention, MLP など）に個別のスカラーを適用（より細かい調整）。
3. Gate Scaling: マルチモーダル DiT（MM-DiT）において、視覚トークンとテキストトークンのゲートごとに個別のスカラーを適用。

2.3 最適化アルゴリズム

勾配を使用しない進化戦略である CMA-ES (Covariance Matrix Adaptation Evolution Strategy) を採用しています。

• 報酬モデル（HPSv3, Q-Align, ImageReward など）を評価関数として使用。
• 候補となる係数セットを生成し、評価して分布を更新するプロセスを反復します。
• 勾配計算が不要なため、大規模モデルに対しても効率的に適用可能です。

2.4 Calibri Ensemble

複数の較正済みモデルを組み合わせる手法です。

• 異なる係数セットで較正された $N$ 個のモデルの出力を重み付けして統合します。
• Classifier-Free Guidance (CFG) の枠組みと統合可能であり、条件付きモデルと無条件モデルの 2 つを較正して組み合わせることで、多様性と精度を両立させます。

3. 実験結果

FLUX.1-dev, Stable Diffusion 3.5 Medium (SD-3.5M), Qwen-Image などの最先端モデルを用いて評価を行いました。

3.1 定量的評価

• 生成品質の向上: どのベースラインモデルにおいても、HPSv3（人間評価に近いスコア）、ImageReward、Q-Align などの指標で顕著な向上が見られました。
  • 例：FLUX において、HPSv3 が 11.41 から 13.48 へ向上。
• 推論ステップ数の削減: Calibri を適用することで、高品質な画像生成に必要な推論ステップ数が大幅に減少しました。
  • FLUX: 30 ステップ → 15 ステップ
  • SD-3.5M: 80 ステップ → 30 ステップ
  • Qwen-Image: 100 ステップ → 30 ステップ
• Ensemble の効果: 較正済みモデルを 2 つ組み合わせる（N=2）ことで、さらに性能が向上し、最適な推論ステップ数が 10〜15 ステップまで低下しました。

3.2 定性的評価と人間評価

• 人間評価: 200 人のユーザーによる評価において、Calibri 適用モデルはベースラインに対して「全体の好み（Overall Preference）」および「テキストとの整合性（Text Alignment）」で有意に高い勝率（約 52-54%）を示しました。
• 可視化: 生成された画像は、より鮮明で、指示に従った内容が正確に反映されていることが確認されました。

3.3 他のアライメント手法との比較

• Flow-GRPO との比較: Flow-GRPO（大規模な微調整手法）と比較して、Calibri は更新パラメータ数が 10 万倍近く少ない（1878 万パラメータ vs 216 パラメータ）にもかかわらず、同等以上の性能を達成しました。
• 組み合わせの可能性: Flow-GRPO などで既にアライメント済みのモデルに対して Calibri を適用することで、さらに性能を底上げできることが示されました。

3.4 コスト

• 較正コスト: 1 回のオフライン較正に要する GPU 時間（H100）は、モデルによって 32〜356 時間程度です。これは一度きりのコストであり、推論時の計算負荷は増えません。
• 効率性: 較正により推論ステップ数が半分以下になるため、実用上の推論速度は 2〜3 倍向上します。

4. 主要な貢献

1. パラメータ効率の高いアライメント手法: 数百パラメータのみの最適化で、大規模 DiT モデルの生成品質を劇的に改善する手法「Calibri」を提案。
2. ブラックボックス最適化の適用: 勾配不要な CMA-ES を用いることで、報酬モデルを直接最適化対象とし、微調整なしでモデルを「較正」する枠組みを確立。
3. 推論効率の飛躍的向上: 生成品質を維持・向上させながら、必要な推論ステップ数を大幅に削減（例：30→15 ステップ）。
4. Ensemble によるさらなる性能向上: 複数の較正済みモデルを組み合わせることで、単一モデル以上の性能とロバスト性を獲得する手法を提案。

5. 意義と結論

Calibri は、大規模生成モデルの微調整が抱える「計算コストの高さ」と「パラメータ膨張」という課題に対する画期的な解決策です。

• 実用性: 計算リソースが限られた環境や、リアルタイム性が求められるアプリケーションにおいて、高品質な画像生成を低コストで実現可能にします。
• 汎用性: 異なるアーキテクチャ（FLUX, SD3, Qwen など）や、異なるアライメント目標（PickScore, GenEval など）に対して広く適用可能です。
• 新たな視点: DiT の内部構造において、重み付けの最適化が生成品質に決定的な影響を与えることを実証し、モデル解釈性や制御の新たな研究方向性を示唆しています。

本手法は、生成 AI の実用化における「品質」と「効率」のトレードオフを打破する重要な技術として位置づけられます。