U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers

要約

拡散トランスフォーマー(DiTs)は、潜空間画像生成のための拡散タスクにトランスフォーマーアーキテクチャを導入する。一連のトランスフォーマーブロックを連鎖させる等方性アーキテクチャにより、DiTsは競争力のある性能と優れたスケーラビリティを示すが、一方で、DiTsによるU-Netの放棄とそれに続く改良は再考する価値がある。この目的のために、U-NetアーキテクチャのDiTと等方的なDiTを比較する簡単なおもちゃの実験を行う。その結果、U-Netアーキテクチャは、U-Net誘導バイアスの中でわずかな優位性しか得られないことが判明し、U-NetスタイルのDiTの中に潜在的な冗長性があることが示された。U-Netのバックボーン特徴が低頻度優位であるという発見に触発され、我々はクエリ・キーバリュー・タプルに対してトークン・ダウンサンプリングを行い、計算量の大幅な削減にもかかわらず、さらなる改善をもたらす自己注意を行う。ダウンサンプリングされたトークンを用いた自己アテンションに基づき、本論文では一連のU字型DiT（U-DiT）を提案し、U-DiTモデルの並外れた性能を実証するために広範な実験を行う。提案するU-DiTは、DiT-XL/2の1/6の計算コストで、DiT-XL/2を上回ることができる。コードはhttps://github.com/YuchuanTian/U-DiT。

要約(オリジナル)

Diffusion Transformers (DiTs) introduce the transformer architecture to diffusion tasks for latent-space image generation. With an isotropic architecture that chains a series of transformer blocks, DiTs demonstrate competitive performance and good scalability; but meanwhile, the abandonment of U-Net by DiTs and their following improvements is worth rethinking. To this end, we conduct a simple toy experiment by comparing a U-Net architectured DiT with an isotropic one. It turns out that the U-Net architecture only gain a slight advantage amid the U-Net inductive bias, indicating potential redundancies within the U-Net-style DiT. Inspired by the discovery that U-Net backbone features are low-frequency-dominated, we perform token downsampling on the query-key-value tuple for self-attention that bring further improvements despite a considerable amount of reduction in computation. Based on self-attention with downsampled tokens, we propose a series of U-shaped DiTs (U-DiTs) in the paper and conduct extensive experiments to demonstrate the extraordinary performance of U-DiT models. The proposed U-DiT could outperform DiT-XL/2 with only 1/6 of its computation cost. Codes are available at https://github.com/YuchuanTian/U-DiT.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL