PoM:基于多项式混合器的高效图像与视频生成
计算机视觉与模式识别
2024-11-20 v1 人工智能
机器学习
摘要
基于多头注意力(MHA)的扩散模型在生成高质量图像和视频方面已变得无处不在。然而,将图像或视频编码为补丁序列会导致代价高昂的注意力模式,因为内存和计算需求均呈二次方增长。为缓解这一问题,我们提出了一种替代MHA的直接替换模块,称为多项式混合器(PoM),其优势在于能将整个序列编码为显式状态。PoM相对于序列长度具有线性复杂度。这种显式状态还允许我们以顺序方式生成帧,在最小化内存和计算需求的同时,仍能进行并行训练。我们证明了多项式混合器与常规MHA一样,是通用的序列到序列近似器。我们将几个Diffusion Transformers(DiT)适配为使用PoM替代MHA来生成图像和视频,并在使用更少计算资源的情况下获得了高质量样本。代码可在https://github.com/davidpicard/HoMM获取。
引用
@article{arxiv.2411.12663,
title = {PoM: Efficient Image and Video Generation with the Polynomial Mixer},
author = {David Picard and Nicolas Dufour},
journal= {arXiv preprint arXiv:2411.12663},
year = {2024}
}