[論文レビュー] Temporal Dynamic Quantization for Diffusion Models

Junhyuk So, Jungwon Lee+3 more|arXiv (Cornell University)|Jun 4, 2023

Stochastic Gradient Optimization Techniques被引用数 4

ひとこと要約

本稿では、時間ステップ情報に基づいて量子化インターバルを動的に調整する、拡散モデル向けの新規量子化手法であるTemporal Dynamic Quantization (TDQ) を提案する。この手法は推論オーバーヘッドを伴わず、生成品質を顕著に向上させる。従来の静的または一般的な動的量子化とは異なり、TDQは推論段階に計算コストを追加せず、post-training quantization (PTQ) および quantization-aware training (QAT) の両方を向上させ、低ビット幅（例：CIFAR-10 における W4A4 FID 4.48）で最先端の性能を達成する。

ABSTRACT

The diffusion model has gained popularity in vision applications due to its remarkable generative performance and versatility. However, high storage and computation demands, resulting from the model size and iterative generation, hinder its use on mobile devices. Existing quantization techniques struggle to maintain performance even in 8-bit precision due to the diffusion model's unique property of temporal variation in activation. We introduce a novel quantization method that dynamically adjusts the quantization interval based on time step information, significantly improving output quality. Unlike conventional dynamic quantization techniques, our approach has no computational overhead during inference and is compatible with both post-training quantization (PTQ) and quantization-aware training (QAT). Our extensive experiments demonstrate substantial improvements in output quality with the quantized diffusion model across various datasets.

研究の動機と目的

• 反復的ノイズ除去ステップにおける時間に依存する活性化分布の変化が、拡散モデルの低ビット量子化において性能劣化を引き起こす要因であることを解消する。
• 活性化統計の時間的シフトに適応できない静的および一般的な動的量子化の限界を克服する。
• 推論時に計算コストを追加せずに、時間依存の最適な量子化パラメータを生成する量子化手法を開発する。
• モバイルおよびエッジデバイスへの広範な展開を可能にするために、post-training quantization (PTQ) および quantization-aware training (QAT) と両立する。
• 4ビットの重みと活性化（例：4ビット精度）で高精度な画像生成を実現するとともに、フル精度ベースラインに近い性能を維持する。

提案手法

• 現在のノイズ除去ステップに応じて、量子化パラメータ（スケールおよびゼロポイント）を予測する学習可能なニューラルネットワークであるTDQモジュールを導入する。
• 推論計算グラフを変更せずに、既存のPTQおよびQATパイプラインにシームレスに統合可能なように設計する。
• すべての時間ステップにわたる活性化量子化誤差を最小化するために、微分可能量子化スキーム（例：LSQ）を用いてTDQモジュールをエンドツーエンドで訓練する。
• 性能と複雑さの最適なトレードオフを実現するために、実験的に4層のマルチレイヤーパーセプトロン（MLP）をTDQのバックボーンとして選定する。
• 逆ノイズ除去プロセス中に変化する活性化分布に適応する、時間依存の量子化インターバルを生成する。
• 前方伝搬時に各時間ステップごとに一度だけ量子化パラメータを計算することで、実行時計算を回避し、推論効率を確保する。

Figure 1: The forward and reverse processes of diffusion models.

実験結果

リサーチクエスチョン

• RQ1 時間ステップに特化した活性化分布に適応する動的量子化が、拡散モデルにおける低ビット量子化性能を顕著に向上させることができるか？
• RQ2 時間依存の量子化戦略は、静的および入力依存の動的量子化を上回り、低ビット幅でも生成品質を維持できるか？
• RQ3 提案されたTDQモジュールは、推論オーバーヘッドを追加せずに、PTQおよびQATフレームワークの両方と効果的に統合可能か？
• RQ4 TDQモジュールの層数は、最終的なモデル性能および収束安定性にどのように影響するか？
• RQ5 サンプリングスケジュールが変化しても（例：ステップ数が減少しても）、TDQはどの程度の性能を維持できるか？
• RQ6 量子化インターバルは時間ステップに応じてどのように変化するか？また、活性化ダイナミクスと相関関係があるか？

主な発見

• CIFAR-10でDDIMを用いたW4A4量子化において、TDQはFID 4.48を達成し、LSQ（7.30）や他のベースラインを顕著に上回る。
• LDM-Churchesデータセットでは、TDQはW4A4 FID 4.64を達成し、静的量子化（5.06）および動的量子化（5.17）を上回る。
• W3A3量子化において、TDQはCIFAR-10でFID 6.48を達成し、LSQ（7.63）や他の動的スキームよりも顕著に優れている。
• アブレーションスタディの結果、4層TDQが性能と複雑さのバランスが最良であり、FIDは4.48（2層では5.18、8層では4.88）に低下する。
• サンプリングステップ数を100から10に削減しても、TDQは一貫した性能を維持するが、LSQは著しく劣化する。
• 可視化により、TDQのインターバルが時間ステップに応じて活性化の変動に適応していることが確認され、約70％の層で強い時間的相関が観察され、特にアテンションおよびリーマンドブロックで顕著である。