Diffusion Transformerにおける外れ値トークンの抑制

背景と課題

Vision Transformer（ViT）には、ごく少数のトークンが異常に大きなノルム（ベクトルの大きさ）を持ち、注意機構（Attention）の重みを不釣り合いに集める一方で、局所的な情報をほとんど持たないという「外れ値トークン」現象が知られている。これまでこの現象は主に識別モデルで研究されてきたが、画像生成モデル、特にDiffusion Transformer（DiT）における役割は十分に調べられていなかった。

本研究は、近年主流になりつつあるRepresentation Autoencoder（RAE）-DiTパイプラインに着目する。このパイプラインでは、事前学習済みViTエンコーダで画像を潜在表現に変換し、DiTがその空間で拡散モデルとして動作する。

提案手法

著者らはまず、外れ値トークンが以下の二箇所で発生することを示した。

• エンコーダ側: 事前学習済みViTエンコーダ自体が外れ値表現を生み出す
• デノイザ側: DiT自身も内部に外れ値トークンを発達させ、特に中間層で顕著

重要な観察として、単純に高ノルムなトークンをマスクするだけでは性能が改善しなかった。これは問題が「極端な値そのもの」ではなく、「局所パッチの意味情報が壊れていること」に起因することを示唆している。

そこで著者らは Dual-Stage Registers（DSR、二段階レジスタ） を提案する。レジスタとは、外れ値の役割を肩代わりさせる追加トークンのことで、本研究では用途に応じて以下を使い分ける。

• 訓練可能レジスタ: 学習が可能な場合に用いる
• 再帰的なテスト時レジスタ: エンコーダを再学習できない場合に、推論時のみで適用
• 拡散レジスタ（diffusion registers）: デノイザ側に適用する専用のレジスタ

この2段階構成により、エンコーダとデノイザ双方の外れ値問題に対処する。

結果と意義

ImageNetおよび大規模Text-to-Image生成の双方で、提案手法は外れ値由来のアーティファクトを一貫して削減し、生成品質を向上させた。

意義として、外れ値トークンの制御がDiTを強化するための重要な要素であることが示された点が挙げられる。これまでViT研究で蓄積されてきた外れ値の知見を、生成モデル設計に橋渡しする成果といえる。

実務での使いどころ

DiTベースの画像生成システムを開発・運用する場面で、生成画像の局所アーティファクト改善に役立つ。事前学習済みエンコーダ（再学習困難なもの）を再利用する場合でも、テスト時のレジスタ導入で改善が見込めるため、導入コストが比較的低い。Text-to-Image基盤モデルのチューニングや、ImageNet系の生成ベンチマークでの品質改善にも適用できる。

注意点・限界

本要約は論文Abstractに基づくため、レジスタの具体的な実装詳細、追加計算コスト、適用するモデル規模ごとの効果差については原論文を参照する必要がある。また、評価はImageNetと大規模Text-to-Imageで行われているが、それ以外のドメイン（動画生成、3D生成など）への一般化は本論文の対象範囲外である。RAE-DiT以外のパイプラインへの適用可能性についても明示的には述べられていない。

実務での使いどころ（要約）

DiTベースの画像生成プロダクト（Text-to-Image、ImageNet系生成）の品質向上に直接活用できる。特に、生成画像に局所的なノイズや不自然なパッチが現れる場合、外れ値トークン制御によって改善が期待される。事前学習済みViTエンコーダ（DINOv2など）を再利用するパイプラインで、追加の学習コストを抑えつつテスト時にレジスタを導入する手法が実装上の参考になる。生成基盤モデルを内製・改良するチームにとって、アーキテクチャ改善の一手段として検討価値がある。