PoM: Efficient Image and Video Generation with the Polynomial Mixer

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 https://github.com/davidpicard/HoMM.

• The paper introduces PoM as a novel alternative to MHA, reducing computational complexity from quadratic to linear for high-resolution generation.
• It proves that PoM acts as a universal sequence-to-sequence approximator, establishing its theoretical viability for various generative applications.
• Experimental integration into Diffusion Transformers shows PoM yields comparable quality to MHA while significantly reducing resource consumption.

Polynomial Mixer: A New Approach to Efficient Image and Video Generation

The paper "PoM: Efficient Image and Video Generation with the Polynomial Mixer" introduces the Polynomial Mixer (PoM), a novel alternative to Multi-Head Attention (MHA) in Diffusion Models for image and video generation. The motivation behind PoM is to address the quadratic scaling limitations of MHA in handling high-resolution images and videos. PoM reduces complexity to linear with respect to the number of tokens, enabling more efficient generation workflows without compromising output quality.

Key Contributions

1. Introduction of PoM: The Polynomial Mixer is proposed as a replacement for MHA. It allows the encoding of the entire sequence into an explicit state, facilitating sequential frame generation with linear complexity, contrasting MHA’s quadratic complexity.
2. Universal Approximation: The authors prove that PoM can serve as a universal sequence-to-sequence approximator, similar to MHA, thus establishing its theoretical validity in diverse generative applications.
3. Implementation in Diffusion Transformers: PoM was incorporated into several Diffusion Transformer (DiT) architectures for image and video synthesis. The results were comparable in quality to MHA-based counterparts but achieved with reduced computational resources.
4. Theoretical and Empirical Analysis: The paper provides a theoretical underpinning of PoM’s ability to operate as a drop-in replacement for attention mechanisms in transformers and supports these claims with empirical evidence across image and video generation benchmarks.

Key Findings and Numerical Results

• PoM demonstrated significant improvements in computational efficiency, particularly at higher image resolutions. For instance, at 4096x resolution, PoM’s total computation time was substantially lower than that of MHA while maintaining comparable image quality.
• The models employing PoM achieved competitive Fréchet Inception Distance (FID) and Inception Scores (IS) on the ImageNet dataset, aligning closely with current state-of-the-art results achieved with MHA.
• The paper's evaluation also included a qualitative analysis of generated images, showing fidelity and diversity across multiple image resolutions, meeting the standards of modern generative models.

Implications

Practically, the adoption of PoM in generative models represents a stride toward more resource-efficient high-resolution content generation. This efficiency is crucial as demand for real-time and on-device generation grows in applications from entertainment to professional content creation. Theorem-backed confidence in PoM's flexibility within existing neural network architectures makes it an attractive area for future research in AI model efficiency optimization strategies.

Theoretically, PoM presents a substantive shift in how transformer architectures can handle sequence data. Its linear complexity opens avenues for research beyond image and video generation, potentially benefiting fields such as multimodal AI and large-scale LLMs.

Future Prospects

The Polynomial Mixer is poised to be an influential component in the ongoing development of efficient AI models. Future research could explore adaptive learning mechanisms with PoM, expanding its utility to other sequential tasks benefiting from linear complexity. Additionally, as video content continues to dominate digital media, PoM’s ability to handle temporal sequences efficiently could prove transformative in video-based applications.

In conclusion, the Polynomial Mixer offers a compelling alternative to Multi-Head Attention, marking a significant advancement in the efficiency of diffusion models. Its introduction is bound to influence both practical applications in generative art and theoretical exploration in computational efficiency paradigms.