LaMamba-Diff: Linear-Time High-Fidelity Diffusion Models Based on Local Attention and Mamba

要約

最近の Transformer ベースの拡散モデルは、顕著なパフォーマンスを示しています。これは主に、入力トークン間の全ペアの相互作用を計算することによってグローバルとローカルの両方のコンテキストを正確に捕捉するセルフアテンション メカニズムの能力に起因します。
ただし、二次関数の複雑さにより、長いシーケンスの入力に対して重大な計算上の課題が生じます。
逆に、Mamba と呼ばれる最近の状態空間モデルは、フィルター処理されたグローバル コンテキストを隠れた状態に圧縮することで線形の複雑さを提供します。
圧縮はその効率性にもかかわらず、必然的にトークン間のきめの細かいローカル依存関係の情報損失につながります。これは効果的な視覚的生成モデリングにとって重要です。
これらの観察に動機付けられて、自己注意と Mamba の長所を組み合わせ、線形の複雑さでグローバル コンテキストとローカルの詳細の両方をキャプチャする Local Attendal Mamba (LaMamba) ブロックを導入します。
効率的な U-Net アーキテクチャを活用することで、当社のモデルは優れたスケーラビリティを示し、256×256 解像度の ImageNet 上のさまざまなモデル スケールにわたって DiT のパフォーマンスを上回ります。同時に、大幅に少ない GFLOP と同等の数のパラメータを利用します。
ImageNet 256×256 および 512×512 の最先端の拡散モデルと比較すると、当社の最大のモデルは、DiT-XL/2 と比較して GFLOP を最大 62% 削減するなど、顕著な利点を示し、同等以下のモデルで優れたパフォーマンスを実現します。
パラメータ。
私たちのコードは https://github.com/yunxiangfu2001/LaMamba-Diff で入手できます。

要約(オリジナル)

Recent Transformer-based diffusion models have shown remarkable performance, largely attributed to the ability of the self-attention mechanism to accurately capture both global and local contexts by computing all-pair interactions among input tokens. However, their quadratic complexity poses significant computational challenges for long-sequence inputs. Conversely, a recent state space model called Mamba offers linear complexity by compressing a filtered global context into a hidden state. Despite its efficiency, compression inevitably leads to information loss of fine-grained local dependencies among tokens, which are crucial for effective visual generative modeling. Motivated by these observations, we introduce Local Attentional Mamba (LaMamba) blocks that combine the strengths of self-attention and Mamba, capturing both global contexts and local details with linear complexity. Leveraging the efficient U-Net architecture, our model exhibits exceptional scalability and surpasses the performance of DiT across various model scales on ImageNet at 256×256 resolution, all while utilizing substantially fewer GFLOPs and a comparable number of parameters. Compared to state-of-the-art diffusion models on ImageNet 256×256 and 512×512, our largest model presents notable advantages, such as a reduction of up to 62% GFLOPs compared to DiT-XL/2, while achieving superior performance with comparable or fewer parameters. Our code is available at https://github.com/yunxiangfu2001/LaMamba-Diff.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google