【DIFFUSSM】Attentionに依存しない拡散モデル

Image generation 2024年07月29日

3つの要点
✔️ 高い画像生成能力をもっている拡散モデルは高解像度での使用には計算上の課題があり、パッチ化などの現在の手法によって処理は高速化されるが、画質が犠牲になる
✔️ 主な計算リソースはAttentionに使われるので、Attentionへの依存がなくなると多く計算コストを削減できる
✔️ DIFFUSSMはAttentionの代わりにSSM機構を利用して、計算効率を向上させながら、生成画質を保持できる

Diffusion Models Without Attention
written by Jing Nathan Yan, Jiatao Gu, Alexander M. Rush
(Submitted on 30 Nov 2023)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

急速な画像生成の進歩は、ノイズ除去拡散確率モデル（DDPM）によって推進されています。DDPMは潜在変数を反復的にノイズ除去することで高忠実度のサンプルを生成しますが、高解像度へのスケーリングには計算上の課題があります。特にSelf-Attention機構がボトルネックとなり、計算コストを低減するために表現圧縮が使用されています。

高解像度アーキテクチャはパッチ化やマルチスケール解像度を採用しますが、これらは空間情報の低下やアーティファクトを引き起こします。DIFFUSSMはSelf-Attentionメカニズムを排除し、状態空間モデル（SSM）を使用して計算量を削減します。DIFFUSSMは画像の細かい表現を処理し、効率を向上させるために砂時計型アーキテクチャを採用しています。ImageNetでの実験では、既存の手法よりも少ないGflopsでFID、sFID、Inception Scoreが向上することが確認されました。

提案手法

本論文の目標は、パッチ化のような「長さの縮小」を必要とせずに、高解像度で長距離相互作用を学習する拡散アーキテクチャを設計することです。Transformerを利用した拡散モデル（DiT）と同様に、このアプローチは画像を平坦化し、シーケンスモデリングの問題として扱います。しかし、Transformersとは異なり、このアプローチではシーケンスの長さに対して亜二次の計算を使用します。

State Space Models (SSMs)

SSM（State Space Models）は、離散時間系列を処理するためのアーキテクチャの一種です。これらのモデルは、次の式で入力シーケンス $u_1$, . . . $u_L$ を処理し、出力 $y_1$, . . . $y_L$ を生成する線形再帰型ニューラルネットワーク（RNN）のような振る舞いをします。

ここで、

となります。このアプローチの主な利点は、Transformersや標準的なRNNなどの代替アーキテクチャと比較して、線形構造を用いて再帰ではなく長い畳み込みを実装できることです。具体的には、FFTを使用して $y_k$を$u_k$から計算することで、$O(Llog L)$の計算量を得ることができ、より長いシーケンスに適用できます。ベクトル入力を処理する時、異なるD個のSSMsをスタックし、D個のバッチFFTを適用することができます。

線形RNN単体では効果的なシーケンスモデルではありませんが、適切な連続時間の状態空間モデルからの離散時間の値を使うことで、安定かつ効果的なアプローチになります。連続時間のSSMパラメータ化と離散時間パラメータを学習し、近似によって同等の結果を得ることができる対角化されたSSMニューラルネットワーク、S4Dをバックボーンモデルとして使用します。

DIFFUSSM Block

DIFFUSSMの中心コンポーネントは、長いシーケンスの処理を最適化するためのゲート付き双方向SSMです。効率を向上させるために、MLPレイヤー内に砂時計型アーキテクチャを組み込んでいます。この設計では、Bidirectional SSMsの周囲でシーケンス長を拡張および縮小し、特にMLP内でシーケンス長を短縮するように交互になります。完全なモデルアーキテクチャは図1に示されています。

具体的には、各アワーグラス層は、短縮された平坦化された入力シーケンス$I ∈ R^{J×D}$を受け取ります。ここで、M = L/Jはダウンスケールとアップスケールの比率です。同時に、双方向SSMsを含むブロック全体は元の長さで計算され、グローバルなコンテキストを十分に活用します。ここで、$σ$は活性化関数を示します。$l ∈ {1 . . . L}$、j = $⌊l/M⌋$、$m = l mod M$、$D_m = 2D/M$として以下を計算します。

各層でこのゲート付きSSMブロックをスキップ接続で統合します。さらに、図1に示すように、各位置でクラスラベル$y ∈ R^{L×1}$とタイムステップ$t ∈ R^{L×1}$の組み合わせを統合します。

実験

クラス条件付き画像生成

本実験では、ImageNet 256x256と512x512データセットで、クラス条件付き画像生成タスクを通じて、提案手法の有効性を検証します。実結果は表１にまとめられています。先行研究、特に従来の拡散モデルと比較して、計算量（Gflops）を大幅に削減できたと確認できます。

ImageNet 256x256データセットでは、いくつかの指標でDiTを上回る性能を、ImageNet 512x512では少ない学習で競争的な結果を達成できました。

つまり、計算量を削減し、生成画質の保持という目標が達成できたと考えられます。図２は各データセットにおける生成の例です。

非クラス条件付き生成性能

DIFFUSSMがLDMと同等のFIDスコアを達成し、トレーニングコストも同等であることを示しています（差が-0.08と0.07）。この結果は、DIFFUSSMが異なるベンチマークや異なるタスクに適用可能であることを強調しています。LDMと同様に、このアプローチはLSUN-BedroomsにおいてADMに対して優れたパフォーマンスを発揮しませんでした。しかし、これは、ADMの総トレーニングコストの25％しか使用していないためです。

モデルのスケーラビリティと砂時計型アーキテクチャの効果

異なるサンプリング設定でモデルをトレーニングし、潜在空間の圧縮の影響を評価しました。結果は図３（右）に示されています。通常のモデル（M = 2）とパッチサイズ2を適用したモデル（P = 2）の比較では、通常のモデルの方が優れたFIDスコアを示し、トレーニングステップが増えるにつれてその差が広がりました。これは、情報の圧縮が高品質な画像生成に悪影響を与える可能性を示唆している。

また、3つの異なるサイズのDIFFUSSMをトレーニングし、スケールアップによるパフォーマンスを評価しました。最初の400kステップのFID-50kを計算した結果、大きなモデルはFLOPsをより効率的に利用し、スケーリングによりトレーニングの各段階でFIDが向上することが確認されました。結果は図３（左）に示されています。

図3．アブレーション
左図：異なる隠れ次元サイズ（D）を持つDIFFUSSM
右図：異なるパッチサイズ（P = 2）およびダウンサンプル比（M = 1）を持つDIFFUSSMのFIDスコア

まとめ

今回の記事では、Attentionに依存しない拡散モデルであるDIFFUSSMを紹介しました。このアプローチは、表現の圧縮を必要とせずに長距離の隠れ状態を処理できます。結果として、256x256の解像度ではDiTモデルよりも少ないGflopsで優れたパフォーマンスを達成し、高解像度でも少ないトレーニングで競争力のある結果を示します。

ただし、いくつかの制限が残っています。まず、無条件の画像生成に焦点を当てており、完全なテキストから画像へのアプローチには対応していません。また、マスク付き画像トレーニングなどの最近のアプローチがモデルを改善する可能性があります。

それでも、DIFFUSSMは、大規模な拡散モデルを学習するための代替アプローチを提供します。Attentionのボトルネックを取り除くことで、高忠実度のオーディオ、ビデオ、3Dモデリングなど、長距離拡散を必要とする他の分野への応用の可能性が広がると考えています。