生成模型研究
TinyDiT Releases a Single-GPU Training Recipe: Registers Absorb 90% of Cross-Attention in a 210M-Parameter Image Model
TinyDiT trains a 210M-parameter DiT from scratch on a single RTX PRO 6000 in 400,000 steps over 3.5 days, while releasing its code, weights, and stage-by-stage measurements. The results show that flow-matching loss is a poor proxy for image-generation progress, whereas registers, timestep shifting, and compiled training provide more actionable signals.
TinyDiT is not intended to challenge large commercial image-generation models. Instead, it reduces modern text-to-image training to an inspectable, single-GPU experiment. The authors trained a 210M-parameter diffusion transformer on 4.2 million 256² images using a single RTX PRO 6000. The 400,000-step run took 3.5 days. The FLUX.2 autoencoder and flan-t5-base text encoder remained frozen, while the original work covers the DiT, data pipeline, training recipe, and evaluation.
The architecture uses cross-attention, 2D RoPE, QK normalization, SwiGLU, and adaLN-single. Compared with the authors’ previous design, adaLN-single freed up 27% of the parameter budget, which was reallocated to a 16-layer model with a hidden width of 896. More notable are the registers: 16 learnable image tokens are inserted at layer 3, while each cross-attention layer also includes two null key/value slots. Measurements show that, in the middle layers at moderate noise levels, the two null slots absorb roughly 90% of cross-attention, while EOS falls to about 4%. The register vector norms reach 4–13 times those of the image tokens, supporting the interpretation that they act as global scratch space, though this is not yet causal evidence from an ablation study.
Training uses rectified flow. Because FLUX.2 latents have 32 channels, the timestep shift is set to 2.8, placing half of the samples in the region above 74% noise. Measured on actual training steps, `torch.compile` delivers a 2.4× speedup while using roughly half the memory, although each of the five aspect-ratio buckets requires its own static graph. Final FID improves from 33.7 to 27.0, FD-DINOv2 from 570 to 218, and object-detection accuracy from 65% to 90%. Over the same period, flow loss declines only from 0.805 to 0.754, indicating that it is better suited to monitoring training health than assessing visual quality.
The limitations are also clear: the model still struggles with legible text, close-up faces, crowds, counts above three, and complex geometry. The metrics come from the authors’ own fixed-prompt evaluation pipeline and have not yet been independently reproduced. The weights also depend on frozen external components, so engineers should verify the full dependency chain and licensing before interpreting “trainable on a single GPU” as complete end-to-end independence.