「Diffusion Transformers（DiTs）」とは。わかりやすく解説

「Diffusion Transformers（DiTs）」とは

Diffusion Transformers (DiTs) は、画像生成タスクにおける拡散モデル（Diffusion Models）にトランスフォーマー（Transformer）アーキテクチャを組み込んだ新しいモデルのことを指します。DiTsは、従来の拡散モデルで広く使われているU-Netバックボーンの代替として設計されています。

1. 背景とモチベーション

• 拡散モデル（Diffusion Models） は、画像生成タスクにおいて最近非常に注目されている手法です。特に、U-Netアーキテクチャをバックボーンに用いた拡散モデルは、GANsに代わって優れた生成性能を示しています。
• 一方で、トランスフォーマーは、自然言語処理や視覚タスクにおいて、特にスケーリング（モデルサイズやデータ量の増加）に対する優れた特性を持っています。これを拡散モデルに適用することで、より高いパフォーマンスが期待できるのではないか、というのがDiTsの出発点です。

2. DiTのアーキテクチャ

• トランスフォーマーブロック：DiTは、トランスフォーマーのSelf-Attentionメカニズムを用いて、画像の潜在表現を操作します。これにより、ViT（Vision Transformer）に近い形で画像パッチを処理します。
• パッチ化：画像を小さなパッチに分割し、それをトークンとしてトランスフォーマーに入力します。これにより、画像全体を一度に処理するのではなく、局所的な情報を効率的に学習します。
• 拡散過程の条件付け：トランスフォーマーブロックにおいて、タイムステップやクラスラベルなどの追加情報を条件付けとして組み込むことができます。これにより、単なる画像生成だけでなく、条件付き生成も可能になります。

3. 主な特徴

• スケーラビリティ：トランスフォーマーの特徴であるスケーラビリティが、DiTsでも活かされています。モデルのサイズ（パラメータ数）や計算量（Gflops）を増やすことで、生成品質が向上する傾向があります。
• 計算効率：従来のピクセル空間での拡散モデルに比べ、DiTsは計算効率が高いことが示されています。特に、潜在空間での拡散を行うことで、Gflopsを抑えながらも高品質な生成が可能です。

4. 性能

• DiTsは、ImageNetのような高解像度画像生成タスクにおいて、従来のU-Netベースの拡散モデル（例えばADMやLDM）を上回る性能を示しています。特に、計算効率の面でも優れた結果を出しています。

5. 今後の展望

• DiTsは、現在の拡散モデル研究における一つのブレークスルーであり、さらなるスケーリングや改良の余地があります。また、テキストから画像を生成するモデル（例えばDALL·E 2やStable Diffusion）におけるバックボーンとしての適用も考えられます。

6. まとめ

• Diffusion Transformers (DiTs) は、トランスフォーマーを基盤とした拡散モデルで、従来のU-Netベースのモデルに代わるものとして設計されています。高いスケーラビリティと計算効率を特徴とし、画像生成タスクにおいて優れた性能を発揮しています。

拡散モデルについて

「Diffusion Transformers（DiTs）」という新しいモデルを紹介します。これは、画像を生成するための「拡散モデル」に使われるアーキテクチャです。基本的に、DiTsは「トランスフォーマー」という既存の技術をベースにしています。トランスフォーマーは元々、テキストや画像の処理で非常に効果的な技術ですが、このモデルはその特性をうまく活かしつつ、画像生成のタスクに最適化されています。

DiTsは、「Vision Transformer（ViT）」という画像を扱うトランスフォーマーのバージョンに基づいています。ViTは、画像を小さなパッチ（部分）に分割し、それらを順番に処理していきます。DiTsもこの方法を採用していて、画像を細かく分けて処理することで、より効率的に学習できるようになっています。

この新しいアーキテクチャは、従来の方法の良いところを引き継ぎつつ、より高度な画像生成ができるように設計されています。図3は、このDiTアーキテクチャ全体の仕組みを示しており、ここではDiTsがどのように画像を処理していくのかを説明しています。

潜在拡散モデル（LDMs）について

LDMs（潜在拡散モデル）は、画像を生成するための効率的な方法です。従来の方法では、画像そのものを扱ってモデルを学習させますが、高解像度の画像になると計算にかかる負荷がとても大きくなります。これを解決するために、LDMsは次の2つのステップを使っています。

1. 画像を圧縮する: まず、画像を「エンコーダ」と呼ばれる仕組みを使って、小さな情報（潜在表現）に圧縮します。これにより、元の大きな画像を扱う代わりに、より小さな情報だけを扱うことができるようになります。
2. 圧縮された情報を使って学習する: 次に、圧縮された情報（潜在表現）を使って、画像を生成するためのモデルを学習させます。この段階では、エンコーダは固定されていて変更されません。

新しい画像を生成するときは、まず学習済みのモデルから潜在表現を取得し、それを「デコーダ」と呼ばれる仕組みを使って元の画像に戻します。

この方法により、LDMsは、計算量が少ないにもかかわらず、良好な性能を発揮します。これは、元の大きな画像を直接扱う方法（例えば、ADMというモデル）と比較しても、計算効率が高いことを示しています。研究では、LDMsをトランスフォーマーという技術に適用し、さらなる性能向上を目指しています。

要するに、LDMsは、画像を一旦圧縮してから効率的に扱うことで、計算負荷を減らしつつ高性能な画像生成を実現する方法です。

DiTの入力仕様

図4. DiTの入力仕様についてわかりやすく説明します。

この図では、DiT（Diffusion Transformers）がどのようにデータを処理するかを説明しています。

1. パッチサイズとは？
   • 画像を小さなブロック（パッチ）に分割する際のサイズです。例えば、パッチサイズが p×pp \times pp×p なら、画像を p×pp \times pp×p の小さなブロックに分割します。
2. 空間表現の意味
   • ここで「空間表現」とは、元の画像が変換されたデータのことです。具体的には、VAEという仕組みでノイズを加えられたデータです。このデータは、形が I×I×CI \times I \times CI×I×C となっており、これは画像の幅と高さが III、色チャンネルが CCC であることを意味します。
3. パッチ化とは？
   • 画像をパッチに分割することです。分割後、各パッチを順番に処理するために、長さ TTT のシーケンス（順番に並べたデータ）に変換します。この長さ TTT は、画像のサイズ III とパッチサイズ ppp によって決まります。具体的には、T=(I/p)2T = (I / p)^2T=(I/p)2 という式で計算されます。
4. パッチサイズと計算量の関係
   • パッチサイズ ppp が小さいほど、画像をより細かく分割することになります。その結果、処理しなければならないデータの量（シーケンスの長さ）が増えます。これは計算量（Gflops）が増加することを意味します。

簡単に言うと、画像を小さなブロックに分けて処理しますが、そのブロックが小さいほど処理する量が増えて、計算が複雑になります、ということです。

Patchifyについて、わかりやすく説明します。

1. Patchifyとは？

• Patchifyは、画像を小さな部分（パッチ）に分割して、それぞれの部分を処理する方法です。DiT（Diffusion Transformers）では、最初にこの「パッチ化」というステップを使います。

2. 具体的な例

• 例えば、元の画像が256 × 256ピクセルで、色のチャンネルが3（RGB）だったとします。この画像は、まず32 × 32 × 4の形に変換されます（これは「空間表現」と呼ばれます）。
• 次に、この空間表現をさらに小さなパッチに分けます。このパッチのサイズはパラメータで決められます。

3. パッチ化のプロセス

• 各パッチを「トークン」と呼ばれるデータの単位に変換します。これらのトークンは順番に処理され、全体として画像を再構築します。

4. パッチサイズが変わるとどうなる？

• パッチサイズ ppp を小さくすると、画像がより細かく分割されます。その結果、トークンの数 TTT が増え、DiTが処理しなければならない計算量（Gflops）も増えます。例えば、パッチサイズを半分にすると、トークンの数は4倍になりますが、これにより計算量も大幅に増加します。

5. パラメータ数への影響

• パッチサイズを変えると計算量は増えますが、それによってモデルのパラメータ数（学習される値の数）には大きな影響はありません。

6. DiTの設計空間

• DiTの設計では、パッチサイズ ppp を2、4、8などに設定して、それぞれ異なるバリエーションを試すことができます。

簡単に言うと、Patchifyは画像を小さく分けて順番に処理する方法で、パッチサイズを小さくすると処理が増えるけど、学習する内容にはそれほど影響しないということです。

DiTブロック設計について、わかりやすく説明します。

1. パッチ化後の処理

• まず、画像が小さなパッチに分けられます（これを「パッチ化」と呼びます）。その後、これらのパッチは「トークン」というデータに変換されます。

2. トークンの処理

• 変換されたトークンは、次に「トランスフォーマーブロック」と呼ばれる処理ユニットを通過します。これは、トークンの内容を分析して、次のステップに必要な情報を取り出すためのものです。

3. 追加の条件情報

• 画像に加えて、DiTは他の情報も処理することができます。例えば、ノイズが加えられた画像のタイミング情報（タイムステップ ttt）、画像が属するクラス（クラスラベル ccc）、さらに自然言語などの追加の情報が含まれます。

4. 4つのトランスフォーマーブロックのバリエーション

• これらの追加情報をうまく処理するために、4つの異なるトランスフォーマーブロックの設計を試しています。それぞれの設計は、標準的なトランスフォーマーブロック（ViT）に少しだけ、でも重要な変更を加えたものです。

5. 図3に示される設計

• 図3では、これらのブロックの設計が具体的にどのように行われているかが示されています。

簡単に言うと、DiTブロック設計では、画像を処理するためのブロックに、追加の情報をうまく処理できるようにした4つのバリエーションがあり、それぞれが画像生成において重要な役割を果たしているということです。

インコンテキスト・コンディショニングについて、わかりやすく説明します。

1. 何をする手法？

• この手法では、画像に加えて「タイムステップ ttt」や「クラスラベル ccc」といった追加情報も一緒に処理します。

2. 具体的な方法

• まず、タイムステップ ttt とクラスラベル ccc を、それぞれ「トークン」というデータの単位に変換します。次に、これらのトークンを、画像トークンの列の一部として追加します。

3. 特別な処理はしない

• 追加されたトークンは、特別扱いせず、他の画像トークンと同じように処理します。つまり、これらのトークンも画像の一部として扱われ、特別な設定や変更は不要です。

4. ViTの仕組みを活用

• この手法は、ViT（Vision Transformer）で使われる「clsトークン」と似ています。clsトークンは、ViTで画像全体の特徴をまとめるために使われるトークンです。インコンテキスト・コンディショニングでは、特別な修正を加えずに、標準的なViTの仕組みをそのまま使うことができます。

5. 効率性

• この方法は、追加の計算負荷（Gflops）をほとんど増やさないので、効率的です。

簡単に言うと、インコンテキスト・コンディショニングは、タイムステップやクラスラベルといった追加情報を画像データに混ぜて処理する方法です。この方法は簡単で、特別な修正を加えずに既存の仕組みをそのまま使えるという利点があります。

図5. 異なるコンディショニング戦略の比較について、わかりやすく説明します。

1. 何を比較しているのか？

• 図5では、画像生成モデルがどの方法（コンディショニング戦略）で一番良い結果を出せるかを比較しています。

2. コンディショニング戦略とは？

• これは、モデルに追加の情報（例えば、画像のタイミングやクラス）をどうやって取り込むかという方法のことです。

3. 3つの戦略

• adaLNZero（適応型レイヤーノルムゼロ）: モデルが最初から安定して働くように、特別な方法で初期化する戦略です。
• クロスアテンション: 画像と追加情報を別々に処理し、後でそれらを結びつける方法です。
• インコンテキストコンディショニング: 追加情報を画像データに混ぜて一緒に処理する方法です。

4. 結果

• 図5によると、訓練の全段階を通じて、adaLNZeroが他の2つの方法（クロスアテンションやインコンテキストコンディショニング）よりも一貫して良い結果を出しています。これは、モデルが安定して効果的に学習できることを示しています。

簡単に言うと、adaLNZeroという方法が、他の方法よりも優れた性能を発揮して、画像生成において最も良い結果を出した、ということです。

クロスアテンションブロックについて、わかりやすく説明します。

1. クロスアテンションブロックの目的

• 画像生成モデルにおいて、追加の情報（タイムステップ ttt やクラスラベル ccc）を画像処理にうまく反映させるための仕組みです。

2. 具体的な方法

• タイムステップ ttt とクラスラベル ccc を、画像のデータとは別に、長さ2のシーケンスとして一緒にします（これが「クロスアテンション」の「クロス（交差）」の部分です）。
• この情報を使って、トランスフォーマーブロックを修正します。具体的には、画像トークンを処理する「自己注意ブロック」の後に、「クロスアテンション層」を追加します。

3. クロスアテンションの役割

• クロスアテンション層は、画像の情報とタイムステップやクラスラベルの情報をうまく結びつける役割を持っています。これにより、モデルはより正確に画像生成の条件を理解することができます。

4. オリジナルデザインとの関連性

• この方法は、Vaswaniらが提案した元々のトランスフォーマーモデルのデザインや、LDM（Latent Diffusion Model）でクラスラベルを処理するために使われる方法に似ています。

5. 計算コスト

• クロスアテンションを使うことで、モデルの計算量（Gflops）が増加します。具体的には、約15％の追加の計算コストがかかります。

簡単に言うと、クロスアテンションブロックは、画像に加えてタイムステップやクラスラベルといった情報をモデルにうまく取り込むための仕組みです。ただし、この方法を使うと計算が少し重くなります。

クロスアテンションブロックについて、わかりやすく説明します。

1. クロスアテンションブロックの目的

• 画像生成モデルに、追加の情報（タイムステップ ttt やクラスラベル ccc）をうまく反映させるための仕組みです。

2. 具体的な処理の流れ

• タイムステップ ttt と クラスラベル ccc を、画像データとは別に扱います。これらを一緒にして、長さ2のシーケンス（連続したデータの列）として連結します。
• その後、この情報を画像トークン（画像の各部分を表すデータの単位）と一緒に処理します。

3. トランスフォーマーブロックの変更

• トランスフォーマーモデルには「自己注意ブロック」という処理がありますが、クロスアテンションブロックでは、これに加えて「クロスアテンション層」を追加します。このクロスアテンション層が、画像の情報と追加情報（タイムステップやクラスラベル）を結びつける役割を果たします。

4. 参考となるデザイン

• この方法は、トランスフォーマーの元々の設計（Vaswaniらのデザイン）や、LDMというモデルで使われる方法に似ています。

5. 計算コストの増加

• クロスアテンションを使うと、計算量が増えます。具体的には、計算量（Gflops）が約15％増加します。

簡単に言うと、クロスアテンションブロックは、画像生成において追加の情報をうまく取り入れるための仕組みです。しかし、その分計算が重くなるというコストがかかります。

適応型レイヤーノルム（adaLN）ブロックについて、わかりやすく説明します

1. 何をする仕組み？

• adaLNブロックは、モデルがデータをうまく処理できるように調整するための仕組みです。この調整は「正規化」と呼ばれる方法を使って行われます。

2. 従来の方法との違い

• 通常の正規化（レイヤーノルム）では、データを一律に処理しますが、**適応型レイヤーノルム（adaLN）**は、状況に応じてデータの処理方法を変えることができます。

3. 具体的な処理

• adaLNでは、データの調整に使うパラメータ（スケールのパラメータ γγγ と シフトのパラメータ βββ）を、直接学習するのではなく、追加の情報（タイムステップ ttt やクラスラベル ccc）から計算します。これにより、データの処理がより柔軟になります。

4. 計算効率の良さ

• adaLNは、他の方法と比べて、計算にかかる負荷（Gflops）が少ないため、非常に効率的です。

5. 同じ関数を適用

• adaLNは、すべてのデータ（トークン）に対して同じ処理を適用します。これは、他の方法にはない特徴です。

簡単に言うと、adaLNブロックは、データを賢く処理するための仕組みで、特に計算が少なく効率的に動くように設計されています。また、すべてのデータに同じ処理を行うことで、一貫性を保ちながらも柔軟に対応できるのが特徴です。

adaLN-Zeroブロックについて、わかりやすく説明します。

1. adaLN-Zeroブロックの背景

• モデルの学習を始める際に、モデルの一部を特定の方法で初期化（設定）すると、学習がスムーズに進むことがわかっています。これは、モデルが安定して動作するための工夫です。

2. 初期化の重要性

• 例えば、ResNetというモデルでは、各ブロックを最初から「恒等関数」という特別な形に初期化するのが有益であることが研究で示されています。これにより、モデルの学習が早く、安定して進みます。同様の方法が、他のモデルでも使われています。

3. adaLN-Zeroの仕組み

• adaLN-Zeroブロックでは、モデルの中で使われるパラメータ（γγγ や βββ）をゼロで初期化することで、最初から安定した学習ができるようにします。
• また、さらにもう一つのパラメータ ααα を導入して、モデルが残差接続（学習がしやすくなる工夫）を行う前に、これを適用します。

4. 具体的な変更点

• adaLN-Zeroは、元々のadaLNブロックに対して、初期化の方法に工夫を加えています。具体的には、γγγ と βββ だけでなく、新たに ααα というパラメータも追加し、それらをゼロから学習するようにしています。

簡単に言うと、adaLN-Zeroブロックは、モデルが安定して学習を始められるように、特別な初期設定を行う方法です。この方法を使うと、モデルが早く、確実に学習を進められるようになります。

表1. DiTモデルの詳細について、わかりやすく説明します。

1. DiTモデルのバリエーション

• DiTモデルには、サイズに応じて4つのバリエーションがあります。
  • Small (S): 小型モデル
  • Base (B): 標準モデル
  • Large (L): 大型モデル
  • XLarge (XL): 最大モデル
• これらは、ViT（Vision Transformer）という元々のモデルの構成に基づいています。

2. 初期化方法

• MLP（多層パーセプトロン）という部分を初期化するとき、特定のパラメータ ααα をゼロに設定します。これにより、DiTブロック全体が「恒等関数」という、何もしない状態でスタートするように設定されます。これによって、モデルが安定して学習を始められるようになります。

3. 計算量の効率性

• この特別な初期化方法（adaLNZero）は、計算の負荷（Gflops）をほとんど増やさないという利点があります。つまり、効率よく処理ができるようになります。

4. DiTの設計空間

• DiTの設計には、次の4つの要素が含まれています。
  • インコンテキスト: 追加情報をトークンとして画像に混ぜて処理する方法。
  • クロスアテンション: 画像と追加情報を別々に処理してから結びつける方法。
  • 適応型レイヤーノルム（adaLN）: 状況に応じてデータを調整する方法。
  • adaLN-Zero: 特別な初期化で安定して学習を始める方法。

簡単に言うと、DiTモデルには4つのサイズがあり、特別な初期化方法を使って効率よく学習できるように設計されています。また、モデルには異なる処理方法を取り入れて、より良い性能を引き出せるように工夫されています。

モデルサイズについて、わかりやすく説明します。

1. DiTモデルの構造

• DiTモデルは、「ブロック」と呼ばれる小さな処理ユニットをたくさん積み重ねて動作します。具体的には、隠れ次元サイズ ddd で動作する NNN 個のブロックを順番に適用していきます。

2. トランスフォーマー構成のスケーリング

• モデルのサイズや性能を調整するために、ブロックの数 NNN、隠れ次元サイズ ddd、そして「アテンションヘッド」という処理単位を一緒にスケーリング（大きくしたり小さくしたりすること）します。この手法は、ViT（Vision Transformer）という技術に基づいています。

3. 4つのモデルサイズ

• DiTには、4つの異なるサイズがあります。
  • DiT-S: Small（小型）モデル
  • DiT-B: Base（標準）モデル
  • DiT-L: Large（大型）モデル
  • DiT-XL: XLarge（最大）モデル
• これらのサイズは、それぞれ異なる計算量（Gflops）を持っていて、小さいモデルは0.3 Gflops、大きいモデルは118.6 Gflopsに及びます。

4. スケーリングパフォーマンスの評価

• モデルサイズが異なることで、どれだけパフォーマンスが変わるかを評価することができます。表1には、各モデルの具体的な構成の詳細が示されています。

5. 設計空間に追加された構成

• DiTの設計には、これらの4つのモデルサイズ（B、S、L、XL）がすべて含まれていて、それぞれ異なる用途やパフォーマンスの要求に応じて使い分けることができます。

簡単に言うと、DiTモデルには4つのサイズがあり、それぞれのサイズでどれくらいの計算が必要か、どの程度のパフォーマンスが出せるかが異なります。この仕組みを使って、さまざまな用途に応じたモデルを設計できるようになっています。

トランスフォーマーデコーダーについて、わかりやすく説明します。

1. 何をする部分か？

• トランスフォーマーデコーダーは、DiT（Diffusion Transformers）の最後のステップで、画像のデータを最終的な出力に変換する部分です。

2. 具体的な処理

• DiTが処理した後、画像データ（トークンのシーケンス）を、予測されたノイズや共分散と呼ばれる出力に変換する必要があります。これらの出力は、元の画像と同じ形にする必要があります。

3. 線形デコーダーの役割

• この変換を行うために、「線形デコーダー」という仕組みを使います。線形デコーダーは、トークンを元の画像の形に戻す役割を果たします。

4. 最後の処理

• デコードされたトークンを、元の画像のようなレイアウト（配置）に再配置します。これにより、モデルが予測したノイズや共分散の情報を、元の画像と同じ形に整理することができます。

5. 設計空間の要素

• DiTの設計には、パッチサイズ、トランスフォーマーブロックの構造、モデルサイズといった要素が含まれており、これらを調整してモデルの性能を最適化します。

簡単に言うと、トランスフォーマーデコーダーは、モデルが処理したデータを最終的な画像の形に戻す部分です。このプロセスでは、トークンを元の画像の形にデコードして、予測結果を正しく配置します。