[論文レビュー] Accelerating Diffusion Decoders via Multi-Scale Sampling and One-Step Distillation

Chuhan Wang, Hao Chen|arXiv (Cornell University)|Mar 20, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

diffusion ベースのイメージトークナイザーのための二段階加速フレームワークを提示： (1) 多段階の粗→細サンプリングによる O(log n) デコード、(2) スケールごとの蒸留を単一ステップのデノイザーへ適用し、 fidelity を保ちつつ大幅な速度向上を実現。

ABSTRACT

Image tokenization plays a central role in modern generative modeling by mapping visual inputs into compact representations that serve as an intermediate signal between pixels and generative models. Diffusion-based decoders have recently been adopted in image tokenization to reconstruct images from latent representations with high perceptual fidelity. In contrast to diffusion models used for downstream generation, these decoders are dedicated to faithful reconstruction rather than content generation. However, their iterative sampling process introduces significant latency, making them impractical for real-time or large-scale applications. In this work, we introduce a two-stage acceleration framework to address this inefficiency. First, we propose a multi-scale sampling strategy, where decoding begins at a coarse resolution and progressively refines the output by doubling the resolution at each stage, achieving a theoretical speedup of $\mathcal{O}(\log n)$ compared to standard full-resolution sampling. Second, we distill the diffusion decoder at each scale into a single-step denoising model, enabling fast and high-quality reconstructions in a single forward pass per scale. Together, these techniques yield an order-of-magnitude reduction in decoding time with little degradation in output quality. Our approach provides a practical pathway toward efficient yet expressive image tokenizers. We hope it serves as a foundation for future work in efficient visual tokenization and downstream generation.

研究の動機と目的

• イメージトークン化のデコーダのモチベーションと遅い推論の問題を指摘する。
• 計算量を削減するための粗→細の多段階デコード方式を提案する。
• 待機時間を削減するために各スケールの蒸留を単一ステップデノイザーへ導入する。
• ImageNet-1K で主要な速度アップと競争力のある再構成品質を示す。

提案手法

• 低から高解像度（32×32 から 256×256）へとデノイズする S スケールを用いた MMDiT ベースのエンコーダ–デコーダを持つ多段階拡散デコーダを使用する。
• 分類器フリーガイダンスに導かれた velociy-field denoising 目標を用い、各タイムステップごとに Euler アップデートをスケールごとに実行する。
• ステージ1 の訓練ではエンコーダとデコーダを共同学習；ステージ2 では各スケールのデコーダを同じ潜在コードを条件として1 ステップデノイザーへ蒸留する。
• 蒸留は固定教師モデル、学生デコーダ、識別器を用い、マルチスケール再構成、知覚的損失（LPIPS）、敵対的項を含む損失を用いる。
• ステージ2 の蒸留により全体のステップ数を 50–100 から各スケール1ステップの計4へ削減する。
• ImageNet-1K（256×256）で、rFID、PSNR、SSIM、スループットを用いて他のトークナイザーと比較評価を行う。

実験結果

リサーチクエスチョン

• RQ1 拡散トークナイザーは、知覚忠実度を損なうことなくリアルタイムまたはほぼリアルタイムの再構成を達成できるか。
• RQ2 粗→細の多段階デコードは、品質を維持しつつ計算コストを削減できるか。
• RQ3 各スケールの蒸留は、多段拡散を各スケールの1ステップデノイザーへ変換し、大きな品質損失なしに実現できるか。
• RQ4 再構成忠実度とデコード速度のスケール間のトレードオフはどうなるか。
• RQ5 提案手法は rFID、PSNR、SSIM、スループットの観点で既存の拡散トークナイザーおよび非拡散トークナイザーと比較してどうなるか。

主な発見

モデル	トークン数	rFID↓	PSNR↑	SSIM↑	スループット (img/s)↑
当方（第1段階後）	128	0.91	23.27	0.752	2.76
当方（第2段階後）	128	1.09	24.74	0.800	87.16
Diffusion FlowMo（FlowMo）	256	0.95	22.07	0.649	1.44
DiTo（DiTo）	256	0.78	24.10	0.706	0.19
表1：ImageNet-1K 256×256 解像度でのトークン化比較。
TiTok-S-128	128	1.71	17.52	0.437	7.31
LlamaGen-16	256	2.19	20.67	0.589	4.55
Cosmos DI-16x16	256	4.40	19.98	0.536	9.55

• 多スケールサンプリングは全解像度サンプリングに対して最大10×の速度アップを達成し、O(log n) デコード計算量を実現する。
• 各スケールの蒸留により総 denoising ステップを約4（各スケール1ステップ）へ削減し、教師モデルより30×超の高速デコードを実現する。
• 蒸留済みの多段階デコーダは ImageNet-1K の 256×256 で rFID 約1.09、PSNR 24.74、SSIM 0.80、スループット 87.16 img/s を達成する。
• 拡散トークナイザー DiTo および FlowMo と比較して、蒸留済み多段階モデルはそれぞれ最大459×および60×の速度アップを実現し、品質は競争力がある。
• Table 1 は最終段階の蒸留モデルが拡散トークナイザーの忠実度に近づきつつ、非拡散トークナイザーの中にはスループットで凌駕するものがあることを示している。
• Table 2 は教員モデルと蒸留された学生モデルの大幅なスループット向上を示し、rFID の増分は控えめである（例：四スケール蒸留モデルで 0.18 増加）。