PoM: Efficient Image and Video Generation with the Polynomial Mixer

基于多头注意力（MHA）的扩散模型已成为生成高质量图像和视频的普遍方法。然而，将图像或视频编码为补丁序列会导致高昂的注意力模式成本，因为内存和计算需求会随着序列长度的平方增长。为了解决这个问题，我们提出了一种可以替代MHA的即插即用模块，称为多项式混合器（PoM），它的好处是能够将整个序列编码成一个显式的状态。PoM 的复杂度与令牌数量呈线性关系。这种显式状态还使我们能够以顺序方式生成帧，从而最小化内存和计算需求，同时仍然能够并行训练。我们证明了多项式混合器是一种通用的序列到序列近似器，就像常规的 MHA 一样。我们将多项式混合器应用于多个扩散变压器（DiT），用于生成图像和视频，取代 MHA，并在使用较少计算资源的情况下获得了高质量的样本。代码可在 https://github.com/davidpicard/HoMM 获取。

该论文试图解决使用多头注意力机制（MHA）生成高质量图像和视频时，计算和内存需求随序列长度增加而呈二次增长的问题。这是一个在大规模图像和视频生成任务中普遍存在的问题。

论文提出了一种称为多项式混合器（Polynomial Mixer, PoM）的新方法，作为MHA的替代方案。PoM通过将整个序列编码到一个显式的状态中，实现了线性复杂度，从而显著降低了计算和内存的需求。这一方法不仅能够高效生成序列数据，还能支持并行训练。

1. PoM被证明是一个通用的序列到序列逼近器，与传统的MHA具有相同的理论保证。 2. 实验结果显示，使用PoM替换MHA的扩散变压器（DiT）能够在减少计算资源的同时生成高质量的图像和视频。 3. 论文提供了开源代码，便于其他研究者复现和进一步研究。 4. 未来的研究可以探索PoM在更多任务中的应用，如自然语言处理和时间序列预测。

1. "Attention Is All You Need" - Vaswani et al., 2017 2. "Vision Transformer" - Dosovitskiy et al., 2020 3. "Denoising Diffusion Probabilistic Models" - Ho et al., 2020 4. "Efficient Attention: Attention with Linear Complexities" - Shen et al., 2018 5. "Linear Transformers Are Secretly Fast Weight Memory Systems" - Katharopoulos et al., 2020