MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Abstractの概要
本論文は、トークンレベルの視点から拡散トランスフォーマーの表現アライメントを研究し、すべての拡散トークンをクリーンな画像のエンコーダ特徴量にアライメントさせることは、拡散入力がノイズを含み、情報を持つコンテンツがタイムステップごとに変化するため、ミスマッチを生じさせると主張している。フルトークンのアライメントでは、アライメントの勾配が空間トークン全体に均等に分布せず、むしろ安定した空間的偏りを示し、特定のトークンパターンに依存していることを著者らは示している。これに対処するため、訓練時にクラストークンを常に保持しつつ、ランダムにサンプリングされたトークンのサブセットに対して表現アライメントを行うMaskAlignを提案している。また、この手法は、マスキング前の情報損失を減らすために、軽量な事前マスク・トークンミキシングブロックを追加している。SiTバックボーンを用いたImageNet 256x256での実験により、収束速度、アライメントの安定性、生成品質、およびステップごとの訓練コストを評価している。
新規性
特徴的な貢献は、拡散訓練における表現アライメントを、完全なトークンマッチングの目的関数ではなく、トークンのサブセット問題として再構築したことである。MaskAlignは、アライメント勾配の集中とトークンサブセットの摂動下での不安定さを示すトークンレベルの分析を動機として、ランダムなサブセットアライメントと軽量な事前マスクトークンミキサーを組み合わせている。
成果
ImageNet 256x256において、MaskAlignはREGなどのアライメントのベースラインと比較して収束と最終的なサンプル品質を向上させ、REPAに追加した場合にもそれを強化する。SiT-XL/2のでは、訓練イテレーションの観点から、通常のSiT-XL/2の約77倍の速さでFID 8.3に、SiT-XL/2 + REPAの約30倍の速さでFID 5.9に到達し、REGと比較して40万イテレーションでFIDを3.4から2.8へ、100万イテレーションで2.7から2.4へと改善した。さらに、REGに対してステップあたりの訓練時間を11.6%削減し、マスキング下でのアライメント損失のギャップがはるかに小さいことを示し、より安定したアライメント挙動を示している。
論文の注目点
- トークンレベルの分析により、フルトークンの表現アライメントは、パッチトークン全体に均等なアライメントをもたらすのではなく、空間的に安定した高勾配トークンの集中を生み出すことが示された。
- MaskAlignは、訓練中にランダムにサンプリングされたトークンのサブセットをアライメントし、パッチトークンをドロップする前の情報を保持するために事前マスクトークンミキシングを使用する。
- ImageNet 256x256での実験では、複数の訓練予算においてREGよりも優れたFID、同等のFIDレベルへのより速い収束、低いステップごとのコスト、およびトークンサブセットの摂動下での安定性の向上が報告されている。