PTQ4DiT: Post-training Quantization for Diffusion Transformers

要約

最近導入された拡散トランスフォーマー (DiT) は、従来の U-Net から離れ、トランスフォーマーのスケーラブルな性質を取り入れた異なるバックボーン アーキテクチャを使用することで、画像生成における優れた機能を実証しました。
DiT の高度な機能にもかかわらず、特にリアルタイム アプリケーション向けの DiT の広範な導入は、現在、推論段階でのかなりの計算需要によって妨げられています。
ポストトレーニング量子化 (PTQ) は、低ビットの重み付けとアクティベーションを使用することで計算とメモリのフットプリントを大幅に削減できる、高速でデータ効率の高いソリューションとして登場しました。
ただし、DiT への適用性はまだ検討されておらず、DiT の独特な設計により、簡単ではない困難に直面しています。
この論文では、DiT 向けに特別に設計された PTQ 手法である PTQ4DiT を提案します。
我々は、DiT に固有の 2 つの主な量子化の課題、特に極端な大きさを持つ顕著なチャネルの存在と、複数のタイムステップにわたる顕著な活性化の分布の時間的変動を発見しました。
これらの課題に取り組むために、私たちはチャネルごとの顕著性バランス (CSB) と Spearmen の $\rho$ ガイドによる顕著性キャリブレーション (SSC) を提案します。
CSB はチャネルの大きさの相補性特性を利用して極値を再分配し、アクティベーションと重みの両方の量子化誤差を軽減します。
SSC は、バランスの取れた顕著性を動的に調整して活性化の時間的変動を捉えることにより、このアプローチを拡張します。
さらに、推論中に PTQ4DiT によって生じる余分な計算コストを排除するために、DiT のオフライン再パラメータ化戦略を設計します。
実験では、当社の PTQ4DiT が同等の生成能力を維持しながら DiT を 8 ビット精度 (W8A8) に量子化することに成功し、さらに 4 ビット重み精度 (W4A8) への効果的な量子化を初めて可能にしたことが実証されています。

要約(オリジナル)

The recent introduction of Diffusion Transformers (DiTs) has demonstrated exceptional capabilities in image generation by using a different backbone architecture, departing from traditional U-Nets and embracing the scalable nature of transformers. Despite their advanced capabilities, the wide deployment of DiTs, particularly for real-time applications, is currently hampered by considerable computational demands at the inference stage. Post-training Quantization (PTQ) has emerged as a fast and data-efficient solution that can significantly reduce computation and memory footprint by using low-bit weights and activations. However, its applicability to DiTs has not yet been explored and faces non-trivial difficulties due to the unique design of DiTs. In this paper, we propose PTQ4DiT, a specifically designed PTQ method for DiTs. We discover two primary quantization challenges inherent in DiTs, notably the presence of salient channels with extreme magnitudes and the temporal variability in distributions of salient activation over multiple timesteps. To tackle these challenges, we propose Channel-wise Salience Balancing (CSB) and Spearmen’s $\rho$-guided Salience Calibration (SSC). CSB leverages the complementarity property of channel magnitudes to redistribute the extremes, alleviating quantization errors for both activations and weights. SSC extends this approach by dynamically adjusting the balanced salience to capture the temporal variations in activation. Additionally, to eliminate extra computational costs caused by PTQ4DiT during inference, we design an offline re-parameterization strategy for DiTs. Experiments demonstrate that our PTQ4DiT successfully quantizes DiTs to 8-bit precision (W8A8) while preserving comparable generation ability and further enables effective quantization to 4-bit weight precision (W4A8) for the first time.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google