PoM: Efficient Image and Video Generation with the Polynomial Mixer

要約

マルチヘッド アテンション (MHA) に基づく拡散モデルは、高品質の画像やビデオを生成するために広く普及しています。
ただし、画像またはビデオをパッチのシーケンスとしてエンコードすると、メモリとコンピューティングの両方の要件が二次関数的に増大するため、コストのかかる注意パターンが発生します。
この問題を軽減するために、シーケンス全体を明示的な状態にエンコードする利点がある、多項式ミキサー (PoM) と呼ばれる MHA のドロップイン代替手段を提案します。
PoM は、トークンの数に関して直線的な複雑さを持ちます。
この明示的な状態により、フレームを逐次的に生成し、メモリと計算要件を最小限に抑えながら、並行してトレーニングすることもできます。
多項式ミキサーが、通常の MHA と同様に、汎用のシーケンス間近似器であることを示します。
MHA の代わりに PoM を使用して画像とビデオを生成するために複数の拡散変換器 (DiT) を採用し、より少ない計算リソースを使用しながら高品質のサンプルを取得します。
コードは https://github.com/davidpicard/HoMM で入手できます。

要約(オリジナル)

Diffusion models based on Multi-Head Attention (MHA) have become ubiquitous to generate high quality images and videos. However, encoding an image or a video as a sequence of patches results in costly attention patterns, as the requirements both in terms of memory and compute grow quadratically. To alleviate this problem, we propose a drop-in replacement for MHA called the Polynomial Mixer (PoM) that has the benefit of encoding the entire sequence into an explicit state. PoM has a linear complexity with respect to the number of tokens. This explicit state also allows us to generate frames in a sequential fashion, minimizing memory and compute requirement, while still being able to train in parallel. We show the Polynomial Mixer is a universal sequence-to-sequence approximator, just like regular MHA. We adapt several Diffusion Transformers (DiT) for generating images and videos with PoM replacing MHA, and we obtain high quality samples while using less computational resources. The code is available at https://github.com/davidpicard/HoMM.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google