What is OpenAI’s Sora Diffusion Transformer (DiT)?

Introduction

OpenAI Sora is back and making waves with its first-ever commissioned music video, “Washed Out – The Hardest Part” This mind-blowing creation was stitched together from 55 individual clips, each generated by Sora itself.

Sounds almost too good to be true, right?

Well, believe it! Back in February 2024, OpenAI’s Sora took the world by storm, showcasing its incredible ability to craft high-definition videos from simple text prompts. This new technology is leading in Generative AI, powered by a strong architecture called the Diffusion Transformer (DiT). In this blog let’s dig deeper into this magical technology behind Sora – DiT.

Diffusion Transformer (DiT) = Diffusion + Transformers

At the core of Sora lies the Diffusion Transformer (DiT) architecture, a novel approach to generative modeling. DiT combines the strengths of diffusion models and transformers to achieve remarkable results in image generation. Let’s break down the key components of DiT:

Diffusion Models

Diffusion models are a class of generative models that learn to gradually denoise a noisy input signal to generate a clean output. In the context of image generation, diffusion models start with a noisy image and iteratively refine it by removing noise step by step until a clear and coherent image emerges. This process allows for the generation of highly detailed and realistic images.

Transformers

Transformers are a type of neural network architecture that has revolutionized natural language processing tasks. They excel at capturing long-range dependencies and understanding the context within a sequence of data. In Sora, transformers are employed to process and understand the textual descriptions provided as input, enabling the model to generate images that accurately reflect the given prompt.

Integration of Diffusion Models and Transformers

The Diffusion Transformer (DiT) architecture seamlessly integrates diffusion models and transformers to leverage their respective strengths. The transformer component processes the textual input and generates a latent representation that captures the semantic meaning of the description. This latent representation is then used to guide the diffusion process, ensuring that the generated image aligns with the provided text.

Sora has been trained on a vast dataset of image-text pairs, allowing it to learn the intricate relationships between visual and textual information. During training, the DiT model is trained to minimize the difference between the generated outputs and the ground truth. The diffusion process is applied to the hidden states, and the denoising network learns to estimate and remove the added noise. The model is trained using a combination of maximum likelihood estimation and adversarial training techniques.

At inference time, the model starts with random noise and iteratively denoises the hidden states using the trained denoising network. The denoised hidden states are then passed through the decoding layer to generate the final output tokens.

How does DiT work in Sora?

Suppose we’ve to generate a video using a text prompt and a series of diffusion steps. 

Here’s a simplified breakdown of what’s happening above:

1. Setting the Stage

• We have a video clip as input.
• We also have a prompt describing the video content, in this case, “sora is sky”.
• We’re at a specific diffusion step (t = 3) in the training process.

2. Preparing the Data

• The video is divided into small squares called patches (imagine a grid overlayed on the video). In this case, each patch covers 4 consecutive pixels across space and time (2 pixels horizontally, 2 pixels vertically, considering multiple frames).

3. Feature Extraction (Understanding the Video)

• Each patch is processed by a visual encoder (shown as a yellow box). Think of this as extracting key features from the image data.
• The encoder uses weights and biases (adjustable parameters) along with a ReLU activation function to transform the raw pixel values into a lower-dimensional representation called a latent vector (shown as a green box). This reduces complexity and allows for better noise handling.

4. Adding Noise (Training the Model)

• Noise is intentionally added to the latent features based on the current diffusion step (t). Higher steps have less noise.
• This is similar to how a language model might be trained by removing words from a sentence and asking it to predict the missing word. By adding noise, the model learns to remove it and recover the original information.

5-7. Conditioning the Noisy Data (Guiding the Model)

• Conditioning helps the model generate a video relevant to the prompt.
  • The prompt “sora is sky” is converted into a text embedding vector (a numerical representation).
  • The diffusion step (t) is also encoded as a binary vector.
  • These are combined into a single vector.
• This combined vector is used to estimate “scale” and “shift” values (also adjustable parameters).
• Finally, the estimated scale and shift are applied to the noisy latent features, creating a “conditioned” noise latent. This injects information from the prompt to guide the model towards generating a video containing “sora” in the “sky”.

8-10. Refining the Conditioned Noise (Focusing on Important Features)

• The conditioned noise latent is fed into a Transformer block, a powerful deep learning architecture.
  • The Transformer uses a technique called “self-attention” to identify the most important relationships within the data.
  • This attention is then used to emphasize relevant information in the conditioned noise latent.
  • Finally, a pointwise feed forward network further processes the data to extract additional features.

11. Training (Fine-tuning the Model)

• The model predicts what the original noise was, based on the conditioned noise latent.
• The difference between the predicted noise and the actual noise (ground truth) is calculated as a loss.
• This loss is used through backpropagation to adjust the weights and biases (red borders) in the model, making it better at predicting noise and ultimately generating realistic videos.
  • It’s important to note that the weights and biases of the visual encoder and decoder (blue borders) remain fixed during this training step. These are pre-trained for efficient feature extraction and generation.

12-14. Generating the Video (The Payoff)

• Once trained, the model can be used for generation.
  • The predicted noise is subtracted from the conditioned noise latent to obtain a noise-free latent.
• This latent representation goes through the visual decoder (another yellow box) which reverses the encoder’s operations.
  • The decoder outputs a sequence of patches.
• Finally, the patches are rearranged back into the original video format, giving us the final generated video content.

Benefits of DiT in Sora

The Diffusion Transformer architecture brings several benefits to OpenAI’s Sora language model:

1. Improved Expressiveness: By treating the hidden states as a continuous diffusion process, DiT allows Sora to learn a more expressive and flexible representation of the input data. This enables Sora to capture subtle nuances and generate more coherent and contextually relevant outputs.
   
2. Enhanced Generalization: The diffusion process helps Sora generalize better to unseen data. By learning to denoise the hidden states, Sora can handle noisy and incomplete inputs more effectively.
   
3. Increased Robustness: DiT’s denoising capability makes Sora more robust to perturbations and adversarial attacks. The model can generate stable and consistent outputs even in the presence of noise or adversarial examples.
   
4. Scalability: The DiT architecture is highly scalable and can be applied to large-scale language models like Sora. It allows for efficient training and inference on massive datasets.

Conclusion

DiT is a significant leap forward in AI-powered video generation. While the full details of Sora remain under wraps by OpenAI, the capabilities showcased suggest a bright future for this technology. DiT has the potential to revolutionize various fields, from filmmaking and animation to video game development and even education. As research progresses, we can expect even more impressive and nuanced video generation with the help of DiT.

Stay tuned to Analytics Vidhya Blogs to get latest updates on Sora!

Himanshi Singh

I’m a data lover who enjoys finding hidden patterns and turning them into useful insights. As the Manager - Content and Growth at Analytics Vidhya, I help data enthusiasts learn, share, and grow together. 

Thanks for stopping by my profile - hope you found something you liked :)