210M Parameter DiT Text-to-Image Model Trained on Single GPU

The Unconventional Path to Training a Diffusion Transformer

Training a 210 million parameter text-to-image diffusion transformer (DiT) from scratch on a single GPU is not the typical route. Most large-scale DiT models leverage massive clusters of high-end accelerators. However, one researcher set out to understand the end-to-end process, undertaking this ambitious project on a single NVIDIA RTX PRO 6000. The goal wasn't just to generate images, but to meticulously measure the process and uncover practical insights often lost in large-scale distributed training narratives. This endeavor yielded three key observations that deviate from commonly held assumptions about training such complex models.

The training run itself took 3.5 days, processing 4.2 million images at a resolution of 256x256 pixels. This duration, while substantial for a single GPU, is remarkably efficient for a model of this size and complexity, hinting at optimized training methodologies or perhaps a fortunate convergence of hyperparameters. The decision to focus on measurements rather than sample quality underscores the project's academic and engineering rigor. It’s an approach that prioritizes understanding the 'recipe' over showcasing the final 'dish'.

Learned Null Attention Slots as Data Sinks

One of the most striking findings is the behavior of learned null attention slots. Following the 'register-token' concept, the model incorporates 16 register tokens within the image stream. Additionally, two learned key/value slots were appended to the sequence. The researcher observed that these appended slots effectively become a sink for learned information. Instead of acting as generalized control mechanisms or auxiliary features, they absorb a significant portion of the model's learned representations. This suggests that the model, when given the opportunity, prioritizes consolidating information into these dedicated, learned slots rather than distributing it more broadly across attention mechanisms. This phenomenon is akin to a meticulously organized librarian who, instead of scattering notes across many desks, consolidates all new arrivals into a single, highly efficient processing bin. This concentration of learning in specific parameters could have implications for model interpretability and the efficiency of information flow during inference.

The implications for model architecture are significant. If learned slots act as data sinks, their design and placement become critical. Are they truly necessary, or could this capacity be better utilized elsewhere? The research implies a need for careful consideration of how auxiliary tokens or learned embeddings interact with the core attention mechanisms in diffusion models. It challenges the notion that more learned parameters are always better, suggesting that their placement and function are paramount.

The Trade-off: Latent Dimension vs. Training Speed

A second crucial measurement relates to the trade-off between the latent dimension size and training speed. For diffusion models, the latent dimension (often denoted as 'C' in architectural descriptions) significantly impacts computational requirements. A larger latent dimension allows the model to represent more intricate details and nuances in the data, potentially leading to higher quality generations. However, it also dramatically increases the computational cost per step. The researcher quantified this relationship, finding that doubling the latent dimension from 64 to 128 resulted in a substantial increase in training time, even on a single GPU. Specifically, the increase was not linear but closer to quadratic, meaning that a modest increase in model capacity leads to a disproportionately larger increase in training duration and resource demands.

This observation is critical for anyone looking to train or fine-tune diffusion models on limited hardware. It provides a concrete data point: if your goal is to achieve faster training cycles or to fit within the memory constraints of a single GPU, opting for a smaller latent dimension is a necessary compromise. The choice between a 64-dimensional latent space and a 128-dimensional one is not merely a hyperparameter tweak; it fundamentally alters the training feasibility on constrained hardware. It’s the difference between baking a complex cake in a standard oven versus attempting it in a toaster oven – the latter might work, but the process and the final outcome will be drastically different.