embedding similar to token embeddings from an LLM vocabulary. Now that we sorted out the interface, the remaining question is how does the model condition the image generation process on the time and class inputs? Well, the diffusion transformer paper looked at four different ways. They're shown in figure 3 here, but we'll build them progressively. Today, models use one of these options or a combination of them, so it's worth understanding all of them. Option one treats time and class as extra tokens and simply concatenates them to the input patches. We'll call them conditioning tokens from now on. This technically works. Self attention allows the conditioning tokens to influence all image tokens. But time and class are now competing with image tokens for attention bandwidth and might end up being ignored. This calls for a conditioning mechanism that gives time in class special treatments compared to the other tokens. Option two adds a new cross attention block. Notice that this is not self attention. It's cross attention between image tokens which are the queries and conditioning tokens which are the keys. In other words, each image patch attends to time and class. So the attention bandwidth is distributed only between the two labels. This setup removes the risk that conditioning tokens are ignored, but it's quite expensive computationally, the most expensive out of all four options. The additional attention adds 15% more gigaflops. Option three is more cost effective and also a lot more interesting. It takes inspiration from an older paper called film that was published back in 2017. At the time they were working on visual reasoning, automatically answering questions about an image. Back then, convolutional neural networks were the dominating paradigm for image processing. CNN internally operate on feature maps, sometimes called channels, which are transformations of an image focused on a particular aspect, be it highle contours, fine grain edges, or shapes. When answering questions, some feature maps are more relevant than others. For instance, when interested in the specific dog breed, it's helpful to look at the fine grain details in the fur. When figuring out the rough position of the dog in the scene, the highle contours are a less distracting signal. Film is a strategy to automatically amplify or suppress feature maps through a linear transformation. The scale gamma and shift beta are derived from the question embedding. Each feature map gets its own set of parameters which are uniformly applied to every feature in the channel. Generally this strategy is called feature modulation. The diffusion transformer borrows this idea in a creative way. Say each patch translates to a flat embedding of size 1024 equal to the transformer's hidden dimension. In the DIT, the equivalent of a feature map is a slice of features that spans across all patches. So the first embedding dimension corresponds to feature map one. The second embedding dimension corresponds to feature map 2 and so on. Basically we'll end up with 1,024 feature maps each requiring its own scale and shift parameters. For image generation, the conditioning signal is no longer a question, but rather the class and time inputs. And it's particularly intuitive why we would want to modulate features based on time. Early on at t= 1, when the noise is so high, the model needs to focus on highle structure. And as it gets closer to tals 0, it could focus on finer details. The implementation is quite straightforward. The time and class embeddings are mixed together. One option is addition followed by a multi-layer perceptron that produces the modulation parameters gamma and beta. The DIT authors are injecting modulation in two separate places before the multi head attention and the feed forward block. Note that one and two are now indices of the group. In both of these groups, gamma and beta are vectors of 1024 values. This option is also called adaptive layer norm or add ln. That's because of its similarity to layer normalization which also applies a scale gamma and a shift beta. The difference is these parameters are dynamically determined by the conditioning signal. And we've almost reconstructed the full diffusion transformer. Option four brings the final missing piece. It's something that is known to stabilize training not just for the DIT but also for older architectures like the ResNet or the diffusion unit. It has to do with the residual block, the one that combines the current activations with the ones from the previous layer. It was empirically observed that at the beginning of training when the model weights are random, it's helpful to ignore upper layers and focus on the lower ones. effectively reducing the depth of the network in early stages. It's quite intuitive that this simplifies the learning problem. To achieve such an effect, it's common to add a simple scaling layer right before the residual block. This simply multiplies the current activation by a parameter alpha. If we initialize the parameters in such a way that alpha is zero at the beginning of training, it will basically silence the current activation and implicitly reduce the depth of the network. The DT defines alpha as a learned function that takes in the conditioning signal. So the time and class now influence the strength of the current activation. Conditioning on time is particularly helpful. Time is correlated with the amount of noise in the image and therefore with the difficulty of the task. It allows us to softly control the depth of the network based on how noisy the input image is. And that is a diffusion transformer. Its biggest limitation is that it's class conditional instead of allowing prompting through free form text. Now there is subsequent work that extended the DIT to text sprints. In particular, the Pixart alpha model built on top of option two, the one with cross attention. They simply replaced the single class embedding with a sequence of token embeddings produced by an off-the-shelf language model named T5. While this technically works, it's a bit suboptimal. The text tokens are static. They're not influenced by the image. But in certain cases, like image editing tasks, this is very restrictive. A prompt like move the bat to the right means different things when referring to this image versus this image. The multimodal diffusion transformer or MMDIT proposes a more principled approach. It treats images and text inputs symmetrically. They both get their own separate encoder parameters. The two modalities communicate through a joined attention mechanism. Basically, the image patches and text tokens are concatenated into one big flat sequence. The sequence attends to itself. Image patches get to see all the other image patches as well as the prompt tokens and vice versa for the prompt in a symmetric way. This dual conditioning makes sure that neither modality is processed in isolation which removes ambiguity. Note that the joint attention module doesn't come with its own parameters. The queries, keys, and values are computed separately by the two modality specific encoders. The joint attention box simply computes and applies the attention scores. Finally, the joint attention output is bifurcated again. Image tokens flow to the left, text tokens to the right. They're each passed through their own linear output projections and continue to travel through the rest of their encoder. Even though that DIT and MMDIT were