GriDiT: Factorized Grid-Based Diffusion for Efficient Long Image Sequence Generation

TL;DR

GriDiT reframes long image-sequence generation by representing sequences as grid images of subsampled frames and factorizing synthesis into a low-resolution coarse sequence (Stage 1) followed by high-resolution per-frame refinement (Stage 2). The approach leverages a Diffusion Transformer with 3D positional embeddings and a Grid-based Autoregressive Sampling scheme to enable arbitrary-length generation with strong long-range coherence. Empirical results across SkyTimelapse, Taichi, Minecraft, and 3D CT datasets show superior quality, improved temporal consistency, and at least 2× faster inference than state-of-the-art methods, while requiring less training data in data-constrained domains. The method generalizes across domains without domain-specific priors and offers a simpler, scalable data representation for image-sequence generation, with notable potential impact in medical imaging and other long-form video applications. $L = 12N - 4$ captures the relation between iteration count and resulting sequence length in the AR sampling procedure, underscoring the efficiency of the sampling strategy, and 3D positional embeddings mitigate looping artifacts observed with purely 2D embeddings.

Abstract

Modern deep learning methods typically treat image sequences as large tensors of sequentially stacked frames. However, is this straightforward representation ideal given the current state-of-the-art (SoTA)? In this work, we address this question in the context of generative models and aim to devise a more effective way of modeling image sequence data. Observing the inefficiencies and bottlenecks of current SoTA image sequence generation methods, we showcase that rather than working with large tensors, we can improve the generation process by factorizing it into first generating the coarse sequence at low resolution and then refining the individual frames at high resolution. We train a generative model solely on grid images comprising subsampled frames. Yet, we learn to generate image sequences, using the strong self-attention mechanism of the Diffusion Transformer (DiT) to capture correlations between frames. In effect, our formulation extends a 2D image generator to operate as a low-resolution 3D image-sequence generator without introducing any architectural modifications. Subsequently, we super-resolve each frame individually to add the sequence-independent high-resolution details. This approach offers several advantages and can overcome key limitations of the SoTA in this domain. Compared to existing image sequence generation models, our method achieves superior synthesis quality and improved coherence across sequences. It also delivers high-fidelity generation of arbitrary-length sequences and increased efficiency in inference time and training data usage. Furthermore, our straightforward formulation enables our method to generalize effectively across diverse data domains, which typically require additional priors and supervision to model in a generative context. Our method consistently outperforms SoTA in quality and inference speed (at least twice-as-fast) across datasets.

Figures (18)

• Figure 1: (a) SoTA image sequence generation models treat image sequences as large tensors of ordered frames. In contrast, (b) our method factorizes image sequence generation into two stages. First, we learn to model the dynamics of the sequence at low resolution, treating the frames as subsampled image grids. Second, we learn to super-resolve individual frames at high resolution. Using the DiT's self-attention mechanism to model dynamics across frames, and paired with our sampling strategy, our method yields superior synthesis quality for sequences of arbitrary length while significantly reducing sampling time and training data requirements. (Notation used is the same as defined in section \ref{['para: Notation']}. $K=4$.)
• Figure 2: An overview of our method's training pipeline. We leverage DiT's self-attention for efficient, high-quality, and arbitrary-length image sequence generation using a two-stage process. (a) Stage 1 ( ): We learn to generate the coarse image sequence at low resolution. We organize the image sequence as grid images, comprising subsampled frames arranged in their sequential order. An unconditional latent DiT is trained to generate them. (b) Stage 2 ( ): We learn to refine individual frames in the generated coarse sequence via faithful generative $\times K$ super-resolution. We pose the problem as one of learning a conditional DiT model to restore the degradation caused by the lossy subsampling of images from our training dataset. ( : trainable. : frozen. $\downarrow K \,\,\, \& \,\,\, \uparrow K$: bilinear (lossy) downsampling and upsampling, respectively. The "Noise" function is further elaborated upon in section \ref{['subsec: SR-method']}.)
• Figure 3: Inferring a single grid image's sequence elements from our model entails: (1) synthesizing grid images using Stage 1 ( ), (2) splitting the grid into coarse frames, (3) adding fine information and super-resolving the coarse frames into individual output frames via Stage 2 ( ), and (4) stacking the ordered frames to form the sequence. ( : frozen.)
• Figure 4: We illustrate our Grid-based Autoregressive Sampling Algorithm used to sample arbitrary-length image sequences. The algorithm entails two steps. We start with step 1 wherein the first iteration starts with vanilla Stage 1 ( ) sampling. Every subsequent iteration uses an appropriately noised control signal from the previous iteration's output at each reverse diffusion timestep to generate four new grid elements, which are in spatiotemporal agreement with all previous grid elements. Upon transitioning to step 2, we interpolate eight new frames between each consecutive pair of 4 new frames generated in different sampling iterations of step 1 for superior temporal resolution. Finally, all new coarse grid-elements are super-resolved via Stage 2 ( ) and stacked in their sequential order for superior spatial resolution and refinement. Consequently, we obtain a long, high-quality image sequence. $N$ such iterations lead to an $L= 12N-4$ length image sequence, inducing a substantial gain in efficiency and quality over the SoTA. ( : frozen).
• Figure 5: Qualitative comparisons with the SoTA on the CT-RATE SkyTimeLapse, and Minecraft datasets. Images are arranged from left to right in their sequential order, i.e., frames 1 through 16. Synthetic CT volumes are generated at 512 $\times$ 512 resolution using $\times$4 SR whereas SkyTimelapse videos are generated at 256 $\times$ 256 resolution using $\times$2 SR. We use the standard $4\times4$ grid setting in Grid-based Autoregressive sampling (step-1) for both cases. Whereas we use a $\{8\times8$ grid, four-row control signal$\}$ setting for step 1 of sampling in experiments on the Minecraft dataset to intentionally allow room for $\times2$ SR, ensuring a fair comparison while yielding the desired 128 $\times$ 128 resolution. Our method yields superior performance in terms of spatiotemporal coherence and quality. (GT: ground truth, GenCT: GenerateCT, SG-V: StyleGAN-V)