Pushing the Boundaries of State Space Models for Image and Video Generation

TL;DR

This paper advances state-space models for image and video generation by building a large-scale diffusion SSM-Transformer hybrid, demonstrating high-quality, temporally consistent visual outputs aligned with complex prompts.

Contribution

It introduces the largest-scale diffusion SSM-Transformer hybrid model, exploring its capacity for long sequence visual generation and demonstrating high-quality results.

Findings

Generated high-resolution images up to 2K.

Produced temporally consistent 8-second videos.

Model aligns well with complex text prompts.

Abstract

While Transformers have become the dominant architecture for visual generation, linear attention models, such as the state-space models (SSM), are increasingly recognized for their efficiency in processing long visual sequences. However, the essential efficiency of these models comes from formulating a limited recurrent state, enforcing causality among tokens that are prone to inconsistent modeling of N-dimensional visual data, leaving questions on their capacity to generate long non-causal sequences. In this paper, we explore the boundary of SSM on image and video generation by building the largest-scale diffusion SSM-Transformer hybrid model to date (5B parameters) based on the sub-quadratic bi-directional Hydra and self-attention, and generate up to 2K images and 360p 8 seconds (16 FPS) videos. Our results demonstrate that the model can produce faithful results aligned with complex text prompts and temporal consistent videos with high dynamics, suggesting the great potential of using SSMs for visual generation tasks.