Skip to main content
Back to path

Welcome to a Grasp lesson

Create a free account to keep studying and create your own study path.

Lesson 38

Temporal Attention for Video Generation

Recommended resources

Hello! Welcome to your next lesson on video generation.

In our last session, we implemented temporal consistency using autoregressive methods, specifically latent variable propagation (LVP). This approach builds a video frame-by-frame, where each new frame is conditioned on the previously generated one. While effective, this sequential process can struggle to capture long-range or non-linear temporal dependencies, such as a character returning to a starting position in a loop.

Today, we shift from a sequential to a parallel paradigm. We will explore how to use temporal attention to allow every frame in a video clip to attend to every other frame simultaneously. This non-autoregressive approach provides a global temporal context, enabling the model to generate all frames of a clip in parallel and maintain a more holistic understanding of motion and consistency. Your learning outcome for this lesson is to apply temporal attention for coherent video generation.

The Challenge: The Computational Cost of Full Spatio-Temporal Attention

A video is a 3D data structure: two spatial dimensions (height, width) and one temporal dimension (frames). A naive approach to attention would be to treat the entire video as one long sequence of pixels or patches and apply standard self-attention.

However, for a video clip of F frames, each with P patches, the complexity of self-attention would be . This is computationally prohibitive for even short, low-resolution videos. The solution is not to abandon attention but to factorize it.

Factorizing Attention: Spatial and Temporal Separation

The standard approach in modern video diffusion models is to separate, or factorize, attention into two distinct, more manageable steps:

  1. Spatial Attention: This operates within each frame. Each patch attends to other patches in the same frame to model spatial relationships and object structure.
  2. Temporal Attention: This operates across frames. Each patch attends to the patch at the same spatial location in all other frames to model motion and temporal evolution.

This factorization dramatically reduces the computational cost. Instead of one massive attention calculation, we have many smaller ones.

Spatio-temporal Diffusion Framework with Co-attention Mechanism
This diagram illustrates the core idea of factorized attention. A spatio-temporal input is processed by separate Spatial Attention and Temporal Attention modules. This separation of concerns is the key to making attention practical for video.

The "magic" that makes this factorization possible lies in clever tensor reshaping. Let's assume our video data, after an initial embedding, has the shape (B, F, P, D), where B is batch size, F is frames, P is the number of patches, and D is the feature dimension.

  • For Spatial Attention: We treat each frame as an independent "image." We can reshape the tensor to (B * F, P, D). Now, a standard self-attention mechanism applied to the P dimension will perform attention within each frame, across a batch of B * F frames.
  • For Temporal Attention: We want to model the trajectory of each patch location over time. We reshape the tensor to (B * P, F, D). Now, self-attention applied to the F dimension will compute the relationships between the same patch location across all frames, for a "batch" of B * P patch trajectories.

Let's see this explained in the context of a Diffusion Transformer (DiT) adapted for video.

Video Generation with Diffusion Transformers | Generative AI

The video 'Video Generation with Diffusion Transformers' provides an excellent explanation of how to implement this spatial and temporal modeling.

Please watch the section from 08:43 to 12:25. Pay close attention to how the input tensor is rearranged and reshaped differently for spatial and temporal attention layers.

This concept is also central to state-of-the-art models like Stability AI's Stable Video Diffusion (SVD). Let's confirm this by looking at the SVD paper.

VIDEO DIFFUSION MODELS FOR KEYFRAME INTERPO- ...

The paper 'Generative Inbetweening: Adapting Image-to-Video Diffusion Models for Keyframe Interpolation' details the internals of Stable Video Diffusion (SVD). Its description of temporal self-attention will reinforce what you just saw.

Please read Section 3.2, 'Temporal Self-Attention'. It explicitly describes reshaping a spatio-temporal tensor X to X' to operate solely on the temporal axis. This is the same principle shown in the video.

Architectural Patterns for Combining Attention

Now that we have our spatial and temporal attention building blocks, how do we arrange them in a network? There are several effective patterns. The ExplainingAI video you just watched also covers the most common variants.

Video Generation with Diffusion Transformers | Generative AI

Let's continue with the same video to see how these blocks can be architecturally combined.

Please watch from 12:25 to 16:00 to understand the four main architectural variants. Then, watch the results discussion from 17:29 to 20:20, which compares their performance. This is key for making architectural decisions.

As you saw, the main architectural variants are:

  • Variant 1 (Alternating/Interleaved): A sequence of alternating spatial and temporal layers (or blocks). This is a common and high-performing pattern. [Spatial_Layer -> Temporal_Layer -> Spatial_Layer -> ...]
  • Variant 2 (Sequential): All spatial layers are processed first, followed by all temporal layers. [Spatial_Layers...] -> [Temporal_Layers...]
  • Variant 3 (Factored Block): Both spatial and temporal attention are performed inside a single Transformer block, one after the other, before the MLP. [Spatial_Attn -> Temporal_Attn -> MLP]
  • Variant 4 (Parallel Block): The attention heads within a single block are split, with half performing spatial attention and the other half performing temporal attention in parallel. Their outputs are then combined.

The video's experiments show that the alternating (Variant 1) and factored (Variant 3) approaches perform best. This makes intuitive sense, as they allow for the continuous exchange of information between the spatial and temporal domains throughout the network's depth. SVD also uses an interleaved approach, as mentioned in the resource What's AI by Louis-François Bouchard.

Diving into the Implementation

Your preference is for practical, implementation-heavy details, so let's examine how this is realized in code.

A great way to see this architecture in action is through a from-scratch implementation. The text2video-from-scratch repository provides a clear Unet3D implementation.

FareedKhan-dev/text2video-from-scratch

This GitHub repository implements a text-to-video model with a 3D U-Net. We'll focus on its U-Net architecture, which cleanly demonstrates the interleaved spatial and temporal attention pattern.

Please review the section 'UNet3D Block'. Focus on the Unet3D class definition. Notice the loops for self.downs and self.ups where spatial_attn and temporal_attn are called sequentially within each level of the U-Net. Also, pay special attention to the definition of temporal_attn at the top of the class.

The most elegant part of that implementation is the definition of the temporal attention module:
temporal_attn = lambda dim: EinopsToAndFrom('b c f h w', 'b (h w) f c', Attention(dim, ...))

This line is a concise, powerful implementation of the reshaping we discussed. It uses einops to:

  1. Take an input of shape b c f h w (batch, channels, frames, height, width).
  2. Rearrange it to b (h w) f c. Here, (h*w) becomes the new "batch" dimension for the attention module, f is the sequence length (the frames), and c is the embedding dimension.
  3. Pass this reshaped tensor to a standard Attention module, which now operates along the f dimension.
  4. Automatically reshape the output back to b c f h w.

Now, let's watch the code walkthrough from the ExplainingAI video, which reinforces these concepts in a DiT context.

Video Generation with Diffusion Transformers | Generative AI

This final video segment walks through the forward method of a video diffusion transformer, highlighting the practical differences from an image-only model.

Please watch from 36:23 to 44:21. Focus on how the code handles: The initial reshape for the uniform frame patch embedding. The creation and addition of temporal positional embeddings. The loop over Transformer layers, with distinct calls and reshapes for spatial and temporal blocks.

Test your understanding!

You have a video latent tensor x with shape (B, F, P, D), where B=2, F=16 frames, P=256 patches, and D=768 dimensions. You need to prepare it for a temporal attention layer.

What will be the shape of the tensor x_temporal that is fed into the core MultiHeadAttention module, and which dimension does the attention operate on? Describe the reshaping operation using either torch.permute/view or einops.

Show answer

The goal is to treat each patch's trajectory across time as a sequence.

Reshaping Operation:

  • Using torch functions: 
    # Original shape: (B, F, P, D) -> (2, 16, 256, 768)
# We want to group B and P together to form a new batch dimension.
x_temporal = x.permute(0, 2, 1, 3) # -> (B, P, F, D) -> (2, 256, 16, 768)
x_temporal = x_temporal.reshape(B * P, F, D) # -> (512, 16, 768)
  • Using einops: 
    from einops import rearrange
# Original shape: (B, F, P, D)
x_temporal = rearrange(x, 'b f p d -> (b p) f d')

Final Shape and Attention Dimension:

  • The shape of x_temporal will be (512, 16, 768).
  • The attention mechanism will operate on the sequence length dimension, which is dimension 1 (the one with size 16), corresponding to the frames.

Conclusion

In this lesson, we made the leap from sequential, autoregressive methods to a parallel, globally-aware approach for video generation using temporal attention.

Key Takeaways:

  • Factorization is Key: Full spatio-temporal attention is computationally infeasible. We factor it into separate spatial attention (within frames) and temporal attention (across frames).
  • Reshaping is the Mechanism: The practical implementation of this factorization relies on clever tensor reshaping to present the data to a standard attention module in the desired format (e.g., (B*P, F, D) for temporal).
  • Interleaved Architectures Perform Well: Arranging spatial and temporal attention blocks in an alternating or interleaved fashion throughout the network is a robust and high-performing design pattern.
  • Temporal Positional Embeddings are Necessary: Just as spatial positional embeddings inform the model about patch locations, temporal positional embeddings are crucial to give the model a sense of frame order and direction.

Next Lesson Preview:

We have now established a strong foundation for generating coherent video clips. In the next lesson, we will apply this knowledge to a more specific and practical task: video inpainting. We'll investigate how to adapt our spatio-temporal diffusion models to fill in masked or missing regions across multiple frames, a task that requires the model to understand and respect both spatial context and temporal motion.