MMOE: Modernizing Diffusion Transformers with Efficient Expert Design
Abstractの概要
本論文は、大規模言語モデルから効率重視のMixture-of-Experts (MoE) 設計を導入し、SiT型の拡散Transformerを近代化したMMOEを提案している。このアーキテクチャでは、ルーティングされたエキスパート、共有および軽量なエキスパート（コピー/ゼロ/定数ルートなど）、ゲート残差ルーティング、アテンション残差集約を組み合わせており、ノイズ除去目的関数、条件付けインターフェース、およびサンプラーは変更せずに維持している。実験は主にクラス条件付きImageNet生成について、8基のH100 GPU、バッチサイズ256、40万学習ステップという制御された単一ノードの計算予算下で行われた。本研究では、各コンポーネントが収束、メモリ、学習コスト、およびルーティングの挙動にどのように影響するかを分析するための条件を揃えた内部比較を重視している。
新規性
際立った貢献は、拡散Transformerに単にMoE層を追加することではなく、最新のLLM-MoEの複数の効率化メカニズムを、統一された拡散Transformerのバックボーンへ体系的に適応させたことである。また本論文では、計算量ゼロの軽量エキスパートやアテンション残差の再利用を含むこれらのコンポーネントを統制された近代化プロセスで調査し、それらが品質とコストのトレードオフに与える複合的な影響を評価している。
成果
ImageNet-256での条件を揃えた学習およびサンプリングプロトコルにおいて、MMOEは記録された各チェックポイントで本論文の密（dense）および中間的な疎（sparse）ベースラインの中で最も低いFIDを達成し、メインの50ステップODE・5万サンプルの設定において、密なSiTが40万ステップで5.20であったのに対し3.75に達した。また、より小規模なモデルスケールやImageNet-512でも密なSiTから改善が見られ、軽量なエキスパート設計によりAMOEと比較してブロックあたりのアクティベーションメモリが削減され、AMOEの測定学習時間は120時間から67時間に短縮された。ルーティング分析では、ブロックごとの安定した特化、初期ブロックでの軽量ルートの大幅な利用、およびノイズ除去中の隣接ステップ間での低いルーティング切り替え率が示されている。
論文の注目点
- MMOEは、外部の学習インターフェースやサンプリングインターフェースを変更することなく、ルーティングエキスパート、軽量なコピー/ゼロ/定数エキスパート、ゲート残差ルーティング、アテンション残差集約をSiT型の拡散Transformerに統合している。
- 本論文の統制された内部比較において、アテンション残差集約が単一コンポーネントとして最大の品質向上をもたらし、完全なMMOEモデルはテストされた疎なバリエーションの中で報告された最高の品質とコストのバランスを示している。
- 著者らは自身の主張を単一マシンの環境下で条件を揃えた内部ベースラインとの比較に明確に限定しており、これまでに発表された拡散MoEシステムとの比較は、学習予算やプロトコルが異なるため直接的ではないと述べている。