Scalable Diffusion Models with Transformers | DiT Explanation and Implementation

this video will be covering dit or diffusion models with Transformers for all our past diffusion videos we have been working with a unet architecture where we had a series of encoder blocks mid block and a series of up blocks with skip connections between encoder and decoder blocks of same spatial resolution the authors of dit replaced this unit architecture with a Transformer and similar to latent Fusion models work on latent images rather than image in pixel space for this the authors base their architecture on Vision Transformer or viit the way we'll go about with this video is first to do a quick review of vit if you aren't familiar with vit then I've already covered it in detail in previous videos so I would suggest to watch those first here we'll only refresh some aspects of it then we'll see the changes the authors make to get to dit and finally look at its implementation and results so in vit we had a sequence of Transformer encoder layers these layers work with sequence of visual tokens of the image and these visual tokens are basically representations of patches of height and width 16 along with these visual tokens we also add a learnable representation for CLS token for training vit on image classification task this CLS token's last layers representation was used to get us the classification scores using a fc layer to transform the input image into a sequence of patch tokens vit has a patch embedding block and the input image is actually fed to this block let's look inside the patch embedding block and a single Transformer layer of vit for patch embedding block the goal is to convert the image into a sequence of visual tokens for this we divide the image into patches of width PW and height pH and say for this image we happen to have these four patches each of these patches have number of channels times height time width pixel values which are then converted to a d dimensional representation which is the hidden Dimension at which the transformal layer is going to be working rather than using FC layer we could also do this with convolutions with patch height and width deciding the kernel sizes and strides and D being the output channels in order to fuse into this representation some knowledge about the order of patch tokens we use position embedding which for vit was representations for absolute positions of these patches so using a learnable embedding layer we got these position representations and they were added to the existing patch representations to get the tokens that will ultimately be fed to the Transformer layers for the classification task viit authors also append the CLS tokens representation with its position being represented as zero this is all about patch embedding block now onto the Transformer layers for the Transformer layers in vit we have a layer Norm followed by self attention with a residual connection and then another layer Norm followed by MLP which is two linear layers with activation and then again a residual connection then these trans former layers are stacked one after another to give us our viit model now we are going to start with changes from vit to get to dit for diffusion Transformers since we are interested in generating images say of 256 cross 256 resolution for now then we won't be needing these classification layers and the entire Transformer would be expected to take as input a noisy image and generate as output either the denoised image or noise prediction furthermore rather than working in pixel space the authors of dit work in the Laten Space by encoding the images using a pre-trained auto encoder specifically the same vae that was used by authors of latent diffusion models with a downscale factor of eight so the input images will actually be noisy latent images of 32 CR 32 with four channels as that is the configuration of va used and the output will also be a latent space image or noise prediction and after the last Transformer layer we would get the denoised image prediction and feed it to the pre-trained decoder of vae to get our generated image in pixel space now let's look at the changes in patch reading block if we would be doing the exact viit implementation then we would have something like this take the sequence of patches project them to Transformer Dimension via linear layer and then add learnable absolute position embeddings as position representation for dit for converting grid of patches to sequence of patch representation we would do the same projection via linear layer it's just that channels will now be whatever is the Z channels or latent dimension of our Auto encoder the only thing that is different from vit to dit is regarding position embeddings instead of using 1D learnable position embeddings the authors use 2D sinusoidal position embeddings so we consider the input as a grid of patches and now each patch instead of being represented by one position is represented by two positions position across X and Y AIS sinusoidal embeddings for both are used to get D by2 dimensional representation for positions across both these Dimensions so here for our example case the first patch will have D by2 dimensional representation for 0 and 0 for second it will be representations of 0 and 1 third will be 1 and Z and the last one will be 1 and one then both these representations are concatenated to finally get a d dimensional position representation for these grid of patches as a side note this is kind of something that the authors of vit also tried where instead of using 1D learnable position embeddings they Ed 2D learnable position embeddings but they did not find significant differences between 1D and 2D with regards to Performance for diffusion Transformers the authors don't provide any details on any kind of experiments they did regarding position embeddings other than the fact that sinusoidal embeddings are being used but when you look into their official implementation there it's 2D grid embeddings the second change is on the final layer side what we desire from the output of Transformer is ultimately an image but as of now it's going to return a sequence of D dimensional tokens the authors of dit add an unpacified block which has a layer norm and MLP and it takes in the D dimensional tokens as input and transforms each of them to a dimension of patch height time width times input channels which would be Z channels these can then be transformed to latent patches and then combined to get this latent image during training this is all we need but during inference we need to generate the image in pixel space so we would additionally call our pre-trained decoder of va to convert this unpacified latent image to image in pixel space now for diffusion we need to pass time step information as well to give the model some sense of the extent of noise present and also conditioning information like class or text if present in our unit models that we have been working with the way we have been handling it so far is that the time step information is converted to time embedding using sinusoidal embeddings and we would pass that to couple of FC layers to get the time step representation which is then passed to all resonant blocks where it's added between the two convolution layers for text conditioning we took the text representation from a text encoder and had all the blocks attend to this representation through cross attention layers and finally for class conditioning we got the class representation using an embedding layer and just added to the time step representation prior to passing it to the resonant blocks the authors of dit try out four different variants for passing the time step and conditioning information to the model but for conditioning the authors only experiment on class conditioning all the four variants use the same mechanism to convert the class and time step information to a d dimensional representation so the hidden size of Transformer so let's see that first the time step T is first converted to a representation using an embedding layer and similar to previous diffusion architectures that we have seen this is a sinusoidal embedding layer this time embedding then goes through two linear layers to give us a d dimensional time step representation so this particular conversion from time step to time step representation is exactly same as previous unit architectures except there we ultimately added these two resonant blocks to get a class representation the authors use a learnable embedding layer which gives us the D dimensional representation for each class now let's look at the four variants so this is our dit block and up till now it's exactly same as viit Transformer block which receives the sequence of tokens from previous layer or patch embedding block the first variant is called Inc context conditioning in this the authors add the class and time step representation as two additional tokens to the input sequence all the Transformer layers Now operate on this new sequence of tokens and obviously before the unpatch ify block after the last Transformer layers we would get rid of the first two tokens and work only with the patch tokens the second variant is cross attention where now we add a cross attention layer with a layer Norm prior to it between the self attention and MLP block the time step and class representations are now concatenated as a sequence of two tokens and this cross attention layer that we just added is actually going to attend to this two token sequence of class and text representation the third variant of dit block is adaptive layer Norm very similar in spirit to what we do in style Gan if you guys are familiar with it if not it's fine I think soon I'll be doing a video on that anyways so here the time step and class representation are added and this now goes through a MLP to give us 4D values these 4D values are used as Dimension wise aine parameters so scale and shift parameters after the layer Norm now layer Norm by default learns these parameters but in this variant we disable that and learn to predict these aine parameters from the conditioning input of time and class the fourth and final variant is very similar to the previous one and is called adaptive layer Norm zero Block in this from the MLP we now predict additional 2D values so a total of 6D values these are used to predict the scale by which we will multiply the output of both the MLP and attention block prior to the residual connection the weights for this MLP is initialized such that the scale predictions are zero initially which means the entire dit block kind of just acts as an identity function at the start similar to Identity mappings in residual Network to ensure better training of the entire network for their experiments the authors buil four different configurations of the model scaling from the smallest dit model which is a 12 layer Transformer with a hidden size of 384 to the largest dit model with 28 layers and a hidden dimension of 1152 apart from the increase of scale by increasing the Transformer depth or width the authors also explore additional dimensionality of scale that is the number of input tokens for that given the latent image size of 32 CR 32 the authors experiment with three different path sizes path size of 8 cross 8 giving us a sequence of 16 tokens patch size of 4 cross 4 so a sequence of 64 tokens and finally patch size of 2 cross2 giving us a sequence of 256 tokens to compare different conditioning design choices the authors train the dit XL model with path size of two with all these four block designs and to compare the performance they use FID 50k metric so what's this FID well this is a metric commonly used to compare generated image quality and was proposed in one of the earlier Gan papers the idea here is that we have a distribution of real world images present in our training as well as testing data set and a distribution of fake images generated by our generative model to compare the quality of different generative models we can use the distance between these two distributions and the model that generates a distribution of images that are closer to the distribution of real world images is a better one according to this metric the way we compute this metric is we feed a sample of real world images to a pre-trained model which is usually Inception V3 pretrained on image net we get the output of the model after feeding these images and here the output used is activations prior to the last classification FC layer of the model from these output features for all the real world images we compute the mean and variance of real world distribution of features similarly feed the distribution of generated images to the model and get the mean and variance of output activations for these generated images then we use this distance metric between two goian to get the FID so lower the FID the better as the distribution of generated images is closer to the distribution of real images and 50k just represents the number of image Samples used to compute this metric now back to the result of performance of different conditioning strategies the authors find that in context conditioning where we just added the text and condition as two additional tokens perform the worst cross attention block is very similar to Adaptive layer norm and the Adaptive layer norom zero block ends up achieving the best FID do keep in mind that this is specific to class conditioning the same result might not be necessarily true for something like text conditioning where even though you could do adaptive layer Norm with pooled text token embeddings but I'm guessing cross attention should perform better like here is an example of a later paper which trains dit for text to image generation and it also uses cross attention for text conditioning with the Adaptive lorm zero block design chosen as the dit block for class conditioning the authors then compare across different configurations to visualize the impact of both the dimensions of scale Transformer size and Patch size not surprisingly as you increase layers and the hidden Dimension size of Transformer the performance becomes better and better across all path sizes of 8 4 as well as 2 keeping the Transformer configuration fixed the smaller patch sizes which have the largest number of input tokens perform the best and as we increase the patch sizes keeping the remaining configurations of Transformer fixed we see worse fids and this is true across all the sized models looking deeper comparing different path size variations for the same sized model sayit XL if you calculate you will see that the parameter counts of these models aren't much different from each other in fact the parameters for path size 2 variant is actually slightly lesser than that of size 4 which is lesser than that of size 8 this is because of the patch embedding block and the unpatch IFI block parameters as they increase as one increases the patch size but still the FID of path size 2 is better than the others now while the parameters of path size 2 variant is not larger than that of four or 8 but when we look at flops the flops are considerably increasing as we go from XEL 8 to Xcel 4 because of Xcel 4 working with larger number of sequence tokens and same when we move from Xcel 4 to XEL 2 this indicates that it's scaling flops that act matter for performance the authors then compare model performance for all configurations with model G flops and find strong correlation between model flops and performance as measured by FID where as you keep increasing flops FID gets better and better suggesting that adding more and more flops or compute is a key component of getting improved performance with dit the largest model that the authors train in terms of flops is did XL with a path size of Two And it generates images with really high quality and in fact ends up getting an FID that's even lower than latent diffusion models and this is while having an architecture complexity as measured by G flops very close to ldm we will now go through an implement mentation of dit but before that let's just summarize all aspects we covered to get to dit from vit and take a look at all the different parts of dit model with adaptive layer NM zero block for dit we start with viit architecture and get rid of classification layers work in latent space rather than pixel space add time step and class condition to get conditioning rep presentation convert each viit block to Adaptive layer Norm zero block and finally add unpacified block at the end of Transformer layers like our previous diffusion model implementations we'll be training dit on CB HQ with similar kind of config that we have been working with just that now model params will actually be dit configuration parameters so this is the config and we are training dit B model with path size 2 here the bunch of parameters that we are concerned with are path size number of layers in our model its hidden Dimension the attention parameters and time embedding dimension now since this config is for training on latent images we also have autoencoder parameters here to start with understanding implementation let's get into patch embedding block which we saw is just going to have a linear layer to convert the input image from a grid of patches to sequence of tokens and also be responsible for baking in positional information into these tokens all Pat embedding block module has is one linear layer this linear layer is going to have input Dimension as channels in input image so this would be Z channels which is four times patch height time patch width and the output Dimension will be hidden size of Transformer which is 768 for dit B this initialization I have is just copied from the official repo in papers the author mentions spefic initialization for the Adaptive layer Norms MLP and the final MLP which we'll look at later but for all the other linear layers in the code this is the initialization that they use in the forward method we first transform the input image into a grid of patches and then from the grid of patches to a sequence of patch tokens each token of Dimension patch height times width times channels this rearrange function from enops takes care of that transformation for us converting batch cross channels cross H cross W where H is nothing but number of patches along height times patch height and same for width to batch size cross number of patches cross patch Dimension then using the linear layer we go from patch Dimension to Transformer dimension for each of these tokens now on to the second responsibility of patch embedding block which is adding positional information this method takes care of getting the 2D sinusoidal position embeddings for the patch tokens this gets as input the grid size after patch ifying so assuming input of 32 cross 32 and a hypothetical path size of 16 grid sizes along both height and width will be 32x 16 so 2 which means this method will return the 2D position embeddings for the 2 cross to Grid in that case we first get the indexes of positions along both dimensions and then create a grid of these indexes giving us the X and Y index positions for each of those patches then from the flattened X and Y positions for all patches we get a d by2 dimensional sinusoidal representation for both X and Y positions concatenating them gives us a position representation of Dimension D for each of the patch tokens and finally patch embedding block adds this positional information to the patch tokens and returns these tokens now let's move to the individual Transformer layer we'll only work with the Adaptive LOM zero block variant as that is the one that's ultimately used in dit and for this we first have the layer Norm for attention block and for all layer nors because this is the Adaptive layer Norm variant we'll disable the layer Norm from using its Define parameters this argument does exactly that without this pyo would actually go ahead and also include a weight and bias parameter but instead we want to regress these from the time step and conditioning input hence we disable these default learnable parameters for the LOM module then the attention block and this is exactly same as vit attention block which I have discussed discussed in detail in my viit videos but let's just do a quick pass over it this is going to have a linear layer to get the qkv representation for each of the tokens and another linear layer for output projection post the scale dot product attention in the forward method we get the qkv representations for all the tokens in the sequence compute attention weights and then after using these weights to get the value representations we use the output projection layer to get back to the Transformer hidden size dimension just like attention we also have a layer Norm for MLP and then the MLP block with the Default hidden dimension of this MLP block we four times the Transformer Dimension D now for each block or layer we use the conditioning input to predict six Dimension wise scale and shift parameters these would be the scale and shift parameters for the layer Norm of attention scale and shift for layer Norm of MLP and scale for output of both of these prior to the residual connection this linear layer computes all these six sets of values from the conditioning input that each block receives dit authors initialize this layer such that the alphas are all zero at initialization making the entire dit block as an identity block so we initialize this MLP weight and bias to be all zero for that note that this would would also zero out the layer Norm scale parameters at initialization in the forward method whatever is fed as conditioning input to the layer we first use that and the Adaptive Norm MLP to get the six scale and shift Dimension wise values then we first have the layer Norm of attention block and now instead of multiplying with scale and adding the corresponding shift we multiply it with 1 plus scale and because the scale output will be zero at initialization doing this effectively makes the layer Norm scale at initialization to be one so at the start this layer Norm scale and shift parameters won't be modifying the layer Norm output at all then we get the attention output and multiply it with post attention scale parameter and add it to the output and as we saw the post attention scale parameter at initialization will be zero making this entire operation equal to Identity at start same thing we do for MLP block layer Norm scale and shift then MLP and then post MLP scale and finally return output after the residual connection and again because both these post attention and post MLP scale parameters which were alpha 1 and Alpha 2 are initially zero this makes the entire dit block just an identity block at start now let's look at our main dit model class which is going to make use of these Transformer layers and Patch embedding block during initialization after getting all the parameters from config we instantiate our patch embedding block the time step information needs to be converted into a representation with same size as hidden size of Transformer so these linear layers take care of exactly that and they are going to receive the sinusoidal time embeddings as input these are all our Transformer layers so instances of the class that we just saw these three are the final layers which which we can refer to as unpacified block so we have a layer norm and when the dit blocks are adaptive layer Norm variant the final layer Norm also has the same scale and shift parameter prediction from conditioning input which is why again we disable the default learnable aine parameters of layer Nom this is the linear layer for predicting the scale and shift parameters and here we just need two sets of per Dimension value which will just be the dimension wise scale and shift values lastly the output linear layer which is going to convert each patch token from hidden dimension of Transformer to an output dimension of patch height time patch width time channels which will be the Z channels of our Auto encoder and we have the same layer Norm initialization of zero scale and shift at start Additionally the authors initialize the weights of final linear layer to also be all zeros in the forward method we call the patch embedding block convert time step to its sinusoidal embedding representation using this function which we have also used in previous videos with unit architecture and then use the time step FC layers to get a Time step representation of same size as the hidden size of Transformer then we go through the Transformer layers one after the other then use the layer Norm MLP to get the scale and shift values for this layer Norm and do the scale and shift accordingly then call the final linear layer to get the output of shape channels times patch height time patch width for each patch tokens at last this rearrange is kind of doing the opposite of what we did inside the patch embedding block this will transform the sequence of patches each patch of Dimension channel Time patch height time patch width to a grid of patches and that's it as far as dit model code is concerned we'll be training dit to predict the noise that was added in the image similar to previous diffusion videos so this returned output would be that noise prediction the training and inference portions are exactly same as previous diffusion videos no change at all for training on latent images we'll additionally have VA to encode the image or load saved latence if present the exact training sampling and noise Schuler implementation is the one that we have already seen in ddpm and ldm videos where the model is trained to predict the original noise that was added and once it's trained we sample starting from a random noise and go through the reverse process till we cover all time steps and get a prediction of original image if you aren't familiar with the specific implementation of diffusion and noise Schuler I talk about it in detail in my ddpm implementation video so do check that if you want to have a better understanding of of it there is one difference between my current dit implementation and the official one for training the authors actually learn the variance of reverse process as well as proposed in this paper but as of now the code that I have only provides training using the same original ddpm strategy where the variance is fixed however other than learning variance part rest of the implementation is exactly same same as official dit implementation so let's see the results of training this on CB HQ I found dit to be better not just in terms of final output from The Limited experiments that I did but even time taken to start getting decent outputs like for CB HQ within a few epochs itself I started getting a face like generated image and these are some results after training the model for about 200 epochs all the code that we have seen as well as the training and sampling code I've pushed to the repo so do try diit on your data sets and see how it fars in comparison to ldm for your use case and let me know in case you run into any issues with the code that's about all for this dit video thank you guys for watching these videos and I shall see you in the next one