Representation Entanglement for Generation: Training Diffusion Transformers Is Much Easier Than You Think

TL;DR

REG entangles the noisy image-latent input of a SiT/DiT with a single noisy DINOv2 class token, then jointly denoises both. Unlike REPA, which only uses external DINO features as an intermediate-layer alignment target at training time and discards them at inference, REG keeps the class token in-loop end-to-end — at sampling time the model reconstructs both the image latent and the class token from pure noise. On class-conditional ImageNet 256², SiT-XL/2 + REG reports 63× faster convergence than vanilla SiT-XL/2 and 23× faster than SiT-XL/2 + REPA, while adding only ~0.4% FLOPs and ~0.5% latency for the extra token.

Key claims

A single extra token (a DINOv2 class token, noised on the same schedule as the latent and concatenated along the sequence) is enough to fundamentally change DiT training dynamics — SiT-XL/2 + REG hits FID 2.7 at 1M steps without CFG, versus SiT-XL/2 + REPA’s 6.4 at 1M steps and 5.9 at 4M steps [Table 1, §4.2].
Entanglement alone — concatenating a DINOv2 class token to the input with no REPA alignment loss — already reduces SiT-B/2 FID from 33.0 to 26.7 (~19% relative), demonstrating that the headline gain is not coming from the alignment loss but from the in-loop discriminative input [Table 6, §4.3].
Ablations of what to concatenate isolate the source of the gain: a learnable token gives marginal improvement (FID 24.4 → 23.3), an average of latent features gives almost nothing (24.4 → 24.1), an average of DINOv2 features helps (24.4 → 16.9), but the DINOv2 class token is best (24.4 → 15.2) — i.e., high-level discriminative semantics is what matters, not the act of adding tokens [Table 5, §4.3].
The encoder family matters but the result is robust across self-supervised choices: DINOv2-B 15.2 < DINOv2-L 17.4 < CLIP-L 21.3, all of which beat REPA’s 24.4 baseline at the same backbone and training budget [Table 3, §4.3].
With CFG, SiT-XL/2 + REG matches SiT-XL/2’s 1400-epoch FID in 80 epochs (17× fewer) and surpasses REPA’s 800-epoch FID at 480 epochs, reaching FID 1.40 on class-conditional ImageNet 256² [Table 2, §4.2].
REG is also a representation-learning improvement: CKNNA scores between REG’s intermediate features and frozen DINOv2-g features exceed REPA’s across all measured training steps, network layers, and timesteps — the model learns more discriminative features than REPA despite REPA explicitly aligning to those features [Fig. 3, §4.4].
Computational overhead is negligible: +0.30% parameters, +0.38% FLOPs, +0.49% latency over SiT-XL/2 + REPA, while reducing FID by 56% and raising IS by 50% at matched setup [Table 4, §4.2].

Method

REG modifies the SiT forward process to operate on a paired input: the patchified noised image latent z_t plus a single linearly-projected noised DINOv2 class token c_t. Both are noised on the same SiT interpolant schedule with their own independent Gaussian noise but the same time t, then concatenated along the sequence dimension and fed to the transformer. The training loss has three terms: (1) velocity prediction on the image latent, (2) velocity prediction on the class token (weighted by a small λ_c), and (3) a REPA-style cosine-similarity alignment between an MLP-projected intermediate-layer feature and the spatial DINOv2 features. At inference, both image latent and class token start from independent Gaussian noise and are denoised jointly via the standard SDE Euler–Maruyama solver — no auxiliary network is needed to provide the class token.

The architectural innovation is in what the network reads at every denoising step, not in the loss or the schedule. REPA’s alignment supervision is only active during training; at inference the network has no access to DINO features. REG’s class token is present at every step of the reverse SDE, so the global discriminative signal is in the rollout, not just in the gradients.

Results

ImageNet 256² no-CFG, SiT-XL/2: REG FID 3.4 at 400K steps vs. REPA 7.9 at 400K and SiT-XL/2 8.3 at 7M; REG reaches FID 2.2 at 2.4M steps [Table 1].
ImageNet 256² with CFG, SiT-XL/2: REG FID 1.40 at 480 epochs, sFID 4.24, IS 296.9, Precision 0.77, Recall 0.66 — matches or beats SiT-XL/2 (1.86 at 1400 epochs, 17× longer) and REPA (1.42 at 800 epochs) [Table 2].
Smaller-model transfer: SiT-L/2 + REG at 400K steps reaches FID 4.6, beating SiT-XL/2 + REPA at 4M steps (5.9) — 10× fewer steps, smaller model [Table 1].
Reported wall-clock convergence speedups are 63× over SiT-XL/2 and 23× over REPA at the same target FID [Fig. 2(e), §4.2].

Why it’s interesting

REG sharpens an existing debate on this wiki. Diffusion training efficiency already tracks three competing recipes for fast diffusion convergence: external alignment (Improved Baselines with Representation Autoencoders RAEv2, REPA), internal self-distillation (No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) SRA, Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis Self-Flow), and latent redesign (Unified Latents (UL): How to train your latents UL). REG sits in a fourth position: keep the external DINO prior, but pipe it through the input rather than an intermediate-layer alignment loss. The numbers are striking — 23× over REPA without the architectural cost of an EMA teacher — but the Slack-flagged caveat is real: REG is demonstrated on class-conditional ImageNet only, where a single class token cleanly summarizes the global condition. It is not obvious the same recipe transfers to text-to-image (where the class-token-equivalent would have to come from a frozen VLM and would not be a single token) or to T2V (where Self-Flow’s evidence is that external alignment with V-JEPA hurts). REG therefore reads as the strongest known external-supervision result in the regime where Self-Flow’s scaling critique doesn’t bite, which makes it the right control for any future internal-vs-external comparison at scale.