ディフュージョントランスフォーマーにおける外れ値トークン問題の深掘り

画像生成タスクで優れた性能を発揮するDiffusion Transformers (DiTs)において、“外れ値トークン”と呼ばれる現象がモデルの安定性と生成品質に悪影響を及ぼすことが指摘されています。この外れ値トークンは、Vision Transformer (ViT)ベースのエンコーダとDiTのデノイザーの両方で発生し、特に中間層で顕著に現れる特性を持ちます。これらのトークンは、局所的なパッチ情報が限られているにもかかわらず、不均衡に高いアテンションを引きつけ、結果として不規則なアテンションパターンを生成します。

これまでの研究では、標準的なViTモデルにおいて外れ値トークンが主に最終層で発生し、レジスタートークンによって緩和できることが示されていました。しかし、DiTのような生成モデルでは、入力と出力が同じモダリティ（局所ピクセルまたは潜在表現）であるにもかかわらず、外れ値トークンは中間層で発生するという異なる挙動を示します。 興味深いことに、単純に高ノルムのトークンをマスクするだけでは性能改善が見られず、この問題が単に極端な値の存在によるものではなく、むしろ破損した局所パッチのセマンティクスに深く関連していることが示唆されています。 この発見は、DiTの内部動作において、トークン間の情報の流れと表現の堅牢性に関する根本的な課題を浮き彫りにしています。

Dual-Stage Registers (DSR)による革新的解決策

この外れ値トークン問題に対処するため、Appleの研究者らはDual-Stage Registers (DSR)という新しいレジスターベースの介入手法を導入しました。DSRは、DiTパイプラインを構成する主要な2つのコンポーネント、すなわち事前学習済みViTエンコーダとディフュージョンモデル（デノイザー）の両方に適用されます。

具体的には、エンコーダ側では推論時にのみ挿入される「テスト時レジスタートークン」が使用されます。一方、デノイザー側では、ディフュージョンモデルのトレーニング中に学習される「学習済みレジスタートークン」（本研究では36個）が導入されます。 これらのレジスタートークンの主要な機能は、外れ値の振る舞いを吸収し、画像パッチに対してよりクリーンなアテンションパターンを誘導することです。 トレーニング時には、エンコーダ側のレジスタートークン出力はディフュージョンモデルに供給される前に破棄されます。推論時には、両モジュールからのレジスタートークン出力が最終的な画像生成から破棄され、純粋な画像トークン出力のみが保持されます。 この二段階のアプローチにより、DiTのトレーニングと推論の両フェーズで外れ値トークンの悪影響を効果的に軽減することが可能になります。

生成品質向上への具体的な貢献とベンチマーク

DSRの導入により、Diffusion Transformersの生成品質は一貫して向上することが実証されました。この介入は外れ値アーティファクトを効果的に削減し、ImageNetや大規模なテキスト-to-画像生成タスクにおいて具体的な性能改善をもたらしています。

具体的なベンチマーク結果として、SigLIP2-Bエンコーダを搭載したRAE-DiTモデルでは、ImageNet-256データセットにおけるFIDスコアが5.89から4.58へと大幅に改善されました。 また、大規模なテキスト-to-画像生成タスクにおいても、GenEvalスコアが0.426から0.466に向上しました。 これらの結果は、DSRが単一のアーキテクチャに限定されるものではなく、SiT、JiT、RAEベースのデザインを含む様々なディフュージョンアーキテクチャに対して安定した効果を発揮することを示しています。 この研究は、外れ値トークンの制御が、より堅牢で高品質なDiTを構築するための重要な要素であることを強く示唆しています。

開発者・エンジニア視点での考察

1. DSR導入による安定したトレーニングと推論: 既存のDiffusion TransformerパイプラインにDual-Stage Registers (DSR)を統合することで、モデルのトレーニング安定性が向上し、特に画像生成品質のばらつきを低減できる可能性があります。実装は登録トークンの追加と出力のフィルタリングに焦点を当てるため、比較的手軽に組み込める改善策となり得ます。

2. 中間層におけるトークン挙動の監視の重要性: 本研究は、特にTransformerの中間層で外れ値トークンが顕著に発生することを示唆しています。DiTモデル開発者は、モデルの中間層におけるトークンのノルムやアテンションパターンを監視するデバッグツールを強化することで、潜在的な品質劣化の原因を早期に特定し、DSRのような対策を適用するタイミングを判断できます。

3. 既存ViTエンコーダの活用とDSRの適用範囲: 多くのDiTが事前学習済みViTエンコーダを使用していることを考慮すると、DSRのテスト時レジスタートークンは、既存のエンコーダに変更を加えることなく外れ値問題を緩和する柔軟な手段を提供します。これにより、既存の強力な視覚エンコーダを活用しつつ、生成パイプライン全体の堅牢性を高めることが可能になります。