Abstract
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 this https URL.
Abstract (translated)
基于多头注意力（MHA）的扩散模型已成为生成高质量图像和视频的普遍选择。然而，将图像或视频编码为一系列补丁会导致昂贵的注意力模式，因为内存和计算需求呈二次增长。为了缓解这一问题，我们提出了一种名为多项式混合器（PoM）的多头注意力替代方案，它可以将整个序列编码到一个显式的状态中，从而带来好处。多项式混合器相对于标记数量具有线性复杂度。这种显式状态还允许我们以顺序方式生成帧，最小化内存和计算需求，同时仍能进行并行训练。我们展示了多项式混合器是一种通用的序列到序列近似器，就像普通的多头注意力一样。我们将多项式混合器替代扩散变压器（DiT）中的多头注意力，用于图像和视频生成，并在使用较少计算资源的同时获得高质量样本。代码可在以下网址获取：[此处为提供的链接]。
URL
https://arxiv.org/abs/2411.12663
