拡散モデルから一致性モデルへの進化

拡散モデル（Diffusion Model）は高品質な画像生成が可能ですが、数十甚至は数百回の反復ステップが必要という課題があります。この問題を解決するために 등장したのが**一致性モデル（Consistency Model）**です。拡散過程を加速し、少数ステップ（理想的には1ステップ）で画像生成を実現することを目指します。

常微分方程式（ODE）による解釈

従来の拡散モデルでは、前向き過程（加噪過程）と逆向き過程（復調過程）の両方においてランダム性が存在していました。Consistency Model関連の論文では、このランダム過程（確率微分方程式、SDE）を**決定論的な過程（常微分方程式、ODE）**に変換することを提案しています。この変換により、生成過程が制御可能になり、一致性モデルの数学的基盤が確立されます。

一致性モデルの核心概念

Consistency Modelの基本思想は以下の制約として表現されます：

加噪軌道上の任意の点から、その軌道の起点（クリーン画像）への写像が一貫している

数学的には、関数 (f: (x_t, t) \rightarrow x_\epsilon) が存在し、すべての時間ステップで (f(x_t, t) = f(x_{t'}, t')) が成立します。つまり、拡散プロセスを経た画像がどの時間 (t) にいても、同一の関数 (f) で初期状態 (x_0) に直接変換できるべきです。

モデルのアーキテクチャ

この制約を満たすため、モデルは次のような構造を持ちます：

[ f_\theta(x,t) = c_{skip}(t) \cdot x + c_{out}(t) \cdot F_\theta(x,t) ]

ここで、境界条件として (c_{skip}(\epsilon) = 1)、(c_{out}(\epsilon) = 0) を満たす必要があります。具体的な係数は：

[ c_{skip}(t) = \frac{\sigma_{data}^2}{(t - \epsilon)^2 + \sigma_{data}^2}, \quad c_{out}(t) = \frac{\sigma_{data}(t - \epsilon)}{\sqrt{\sigma_{data}^2 + t^2}} ]

復調過程が進行するにつれ（(T \rightarrow \epsilon)）、(c_{skip}) の値が増加し、現在の復調画像が結果に占める割合が大きくなります。(F_\theta) はUnetなどのニューラルネットワークを表します。

学習アプローチ

蒸留による最適化

事前学習済みの拡散モデルを使用した蒸留ベースの学習では、教师モデルと学生モデルの間の整合性を保つ損失関数を設計します：

[ \mathcal{L}{CD}^N(\boldsymbol{\theta}, \boldsymbol{\theta}^-; \phi) = \mathbb{E}[\lambda(t_n) \cdot d(\boldsymbol{f}{\boldsymbol{\theta}}(\mathbf{x}{t{n+1}}, t_{n+1}), \boldsymbol{f}{\boldsymbol{\theta}^-}(\hat{\mathbf{x}}{t_n}^{\boldsymbol{\phi}}, t_n))] ]

この損失関数の意味を理解するために、実際のプロセスを見てみましょう。サンプルに対して加噪を行い (x_{t_{n+1}}) を取得した後、事前学習済み拡散モデルを一回だけ適用して (x_{t_n}) を予測します。次に、両方の出力を損失関数で制約し、モデルが異なる時間ステップ間で一貫した予測を行うことを強制します。

予測処理はODEソルバーを使用して実行されます：

[ \hat{x}{t_n}^\phi = x{t_{n+1}} - (t_n - t_{n+1}) \cdot t_{n+1} \cdot \nabla_{x_{t_{n+1}}}\log p_{t_{n+1}}(x_{t_{n+1}}) ]

DDIMやDPM++などはこのODEソルバーの実装例です。

欧拉法（Euler Method）の基本式：(y_{n+1} = y_n + h \cdot f(t_n, y_n))。深層学習では大部分の関数をニューラルネットワークで「推定」するため、(\nabla_{x_{t_{n+1}}}\log p_{t_{n+1}}(x_{t_{n+1}}) \approx s_\theta(x_{t_{n+1}}, t_{n+1})) となり、(s_\theta) は学習済み復調ネットワークです。

直接学習による最適化

蒸留手法に加え、直接的なモデル学習アプローチも存在します。この方法では教師モデルを使用せず、学生モデルのみを訓練します。

LCM（Latent Consistency Model）とLCM-LoRA

潜在空間（Lantent Space）における一致性モデルであるLCMは、Stable DiffusionがVAEで画像を潜在空間に圧縮し、DFモデルで学習、最後にVAEデコーダで出力するアーキテクチャを活用しています。LCMでは主に2つの改良が行われています：

1. Skipping-Step（スキップステップ）

元のConsistency Modelでは、連続する時間ステップ間の損失計算において、ステップ間の距離が近すぎるために勾配が微小化する問題があります。スキップステップはこの問題を解決し、損失計算が (d(f(x_{t_{n+k}}, t_{n+k}), f(x_{t_n}, t_n))) のように変更されます。

2. Classifier-Free Guidance（CFG）の導入

損失関数は以下の形式になります：

[ d(f(x_{t_{n+k}}, w + c, t_{n+k}), f(x_{t_n}, w + c, t_n)) ]

ここで (c) はテキスト条件を表します。CFGは改良されたODEソルバーとして機能し、生成品質を向上させます。

実験結果の分析

スキップ幅 (k) の影響： DPM-Solver++およびDPM-Solverでは、約2000ステップの反復でLCMは4ステップのサンプリングのみでFIDがほぼ収束することが確認されています。

Guidance Scale の影響： 異なるLCM反復回数とGuidance Scaleの比較実験により、(w) の増加はCLIP Scoreを向上させます一方、FID指標（多様性）は低下することが明らかになっています。興味深いことに、2、4、8ステップのLCM反復ではCLIP ScoreとFIDの差が小さく、LCMの蒸留性能が非常に優れていることを示しています。

LCM蒸留トレーニングの実装詳細

実装を通じてLCMのトレーニング過程を理解しましょう。

学生モデルの処理

事前学習済みUnetモデル（unet = UNet2DConditionModel.from_pretrained）を関数 (f_\theta) として使用します。Consistency ModelのODEベースのアプローチでは、「経路」の整合性を保証し、モデルに(z_{t_{n+k}}) から (z_{t_n}) への予測を学習させます。

実装例：

batch_size = latents.shape[0]
step_size = noise_scheduler.config.num_train_timesteps // args.num_ddim_timesteps
# DDPM の topk 値を使用（例如：1000 // 50 = 20）
selected_indices = torch.randint(0, args.num_ddim_timesteps, (batch_size,), device=latents.device).long()
# t_{n+k} を取得
future_timesteps = solver.ddim_timesteps[selected_indices]
# t_n を取得
current_timesteps = future_timesteps - step_size
current_timesteps = torch.where(current_timesteps < 0, torch.zeros_like(current_timesteps), current_timesteps)

# 境界条件のスケール係数を計算
skip_start, out_start = compute_scaling_coefficients(future_timesteps, ...)
skip_current, out_current = compute_scaling_coefficients(current_timesteps, ...)

# ノイズを付加
noisy_input = noise_scheduler.add_noise(latents, noise, future_timesteps)

ノイズ取得後、モデルへの入力として予測ノイズを計算（noise_pred = unet(noisy_input, ...).sample）し、原点サンプルを逆推定します（pred_x0 = recover_original_sample()）。前述の公式 (f_\theta(x,t) = c_{skip}(t) \cdot x + c_{out}(t) \cdot F_\theta(x,t)) を適用：

noise_pred = unet(noisy_input, ...).sample
predicted_original = recover_original_sample(noise_pred, future_timesteps, noisy_input, noise_scheduler.config.prediction_type...)
student_output = skip_start * noisy_input + out_start * predicted_original

教师モデルの処理

学生モデルの処理と類似していますが、CFGを適用するため処理がやや複雑になります。 핵심は次の通りです：

noisy_input = noise_scheduler.add_noise(latents, noise, future_timesteps)

# LoRA を無効化して教師モデルとして動作
accelerator.unwrap_model(unet).disable_adapters()

with torch.no_grad():
    # 条件付き予測（条件埋め込み c を使用）
    cond_output = unet(noisy_input, future_timesteps, ...).sample
    cond_x0 = recover_original_sample(cond_output, future_timesteps, noisy_input, ...)
    cond_noise = recover_noise(cond_output, future_timesteps, noisy_input, ...)

    # 無条件予測（条件埋め込み = 0 を使用）
    null_embeds = torch.zeros_like(prompt_embeds)
    null_pooled = torch.zeros_like(encoded_text["text_embeds"])
    null_conditions = copy.deepcopy(encoded_text)
    null_conditions["text_embeds"] = null_pooled
    
    uncond_output = unet(noisy_input, future_timesteps,
        encoder_hidden_states=null_embeds.to(weight_dtype),
        added_cond_kwargs={k: v.to(weight_dtype) for k, v in null_conditions.items()},
    ).sample
    uncond_x0 = recover_original_sample(uncond_output, future_timesteps, noisy_input, ...)
    uncond_noise = recover_noise(uncond_output, future_timesteps, noisy_input, ...)

    # CFG による x0 と eps0 の推定
    final_x0 = cond_x0 + guidance_weight * (cond_x0 - uncond_x0)
    final_noise = cond_noise + guidance_weight * (cond_noise - uncond_noise)

    # ODE ソルバーで次の点を推定（時間を遡って解決）
    previous_state = solver.ddim_step(final_x0, final_noise, selected_indices).to(unet.dtype)

この処理では、教師モデルが学生モデルに対して確定的な復調パスを提供し（ODE計算 통해 直接計算）、学生モデルがノイズから高品質サンプルを生成する方法を学習することを可能にします。

損失計算

future_timesteps = solver.ddim_timesteps[selected_indices]
current_timesteps = future_timesteps - step_size

# LoRA を有効化して学生モデルとして動作
accelerator.unwrap_model(unet).enable_adapters()

with torch.no_grad():
    teacher_noise_pred = unet(previous_state, current_timesteps, ...).sample
    teacher_x0 = recover_original_sample(teacher_noise_pred, current_timesteps, previous_state)
    reference_output = skip_current * previous_state + out_current * teacher_x0

if args.loss_type == "l2":
    loss = F.mse_loss(student_output.float(), reference_output.float(), reduction="mean")

以上の処理を通じて、学生モデルの予測結果（student_output）と教師モデルの提供する確定パス（previous_state）を使用して、スキップステップを含む損失 (\mathcal{L} = |f_\phi(x_{t_s}, t_s) - \text{sg}[f_\theta(x_{t_e}, t_e)]|^2) を計算します。

まとめと今後の展開

一致性モデルは拡散モデルの生成過程を根本から変革する可能性を持っています。ランダムな生成過程を決定論的な過程に変換することで、(T \rightarrow t_0) のすべての点が「一本の線"上にあることを保証し、(f(x_t, t) = f(x_{t'}, t')) が成立するため、モデルが (t+1) の生成内容に依拠して (t) 刻の内容推断する必要がなくなります。

LCM、LCM-LoRA、さらにTCDなどの後続研究は、Consistency Modelの原理に基づいて実用性と効率性をさらに向上させています。