μPを用いた拡散変換器の効率的なスケーリング

拡散トランスフォーマーは、ビジョン生成モデルの基礎として登場しましたが、大規模でのハイパーパラメータ（HP）調整の高コストにより、スケーラビリティに制限があります。最近、vanillaトランスフォーマー向けに最大更新パラメトリゼーション（muP）が提案されました。これにより、小規模から大規模言語モデルへの安定したHP転送が可能となり、調整コストが大幅に削減されます。しかし、アーキテクチャと目的が異なる拡散トランスフォーマーにvanillaトランスフォーマーのmuPが適用できるかどうかは依然として不明です。本研究では、標準的なmuPを拡散トランスフォーマーに一般化し、大規模実験を通じてその有効性を検証します。まず、DiT、U-ViT、PixArt-alpha、MMDiTなどの主要な拡散トランスフォーマーにおけるmuPがvanillaトランスフォーマーと一致することを厳密に証明します。この結果を利用して、DiT-muPの堅牢なHP転送可能性を系統的に示します。特に、学習率を転送したDiT-XL-2-muPは元のDiT-XL-2よりも2.9倍速い収束を達成しました。最後に、PixArt-alphaを0.04Bから0.61Bへ、MMDiTを0.18Bから18Bへスケールアップすることで、テキストから画像生成におけるmuPの有効性を検証しました。両ケースにおいて、muPを使用したモデルはそれぞれのベースラインを超える性能を発揮し、僅かな調整コストで済みました。PixArt-alphaでは1つのトレーニング実行の5.5%のみが必要であり、MMDiT-18Bでは人間エキスパートによる消費量の3%でした。これらの結果は、muPが拡散トランスフォーマーのスケーリングに適した原理に基づいた効率的なフレームワークであることを確立しています。