USP: Unified Self-Supervised Pretraining for Image Generation and Understanding

TL;DR

USP is a masked-latent-modeling pretraining recipe that produces one ViT weight-set usable to initialize both a discriminative backbone (classification, segmentation) and a diffusion transformer (DiT / SiT) via careful weight adaptation. The masked reconstruction runs in a frozen VAE’s latent space (rather than pixel space, as in MAE), so the pretraining input distribution matches what the downstream diffusion model consumes. Reported speedups are 11.7× on DiT-XL/2 and 46.6× on SiT-XL/2 vs. training from scratch, while linear-probe classification and ADE20K segmentation match or exceed MAE at the same epoch budget. The framing is a Bitter-Lesson pitch: replace REPA’s external-DINO teacher with plain masked-latent pretraining and initialization.

Key claims

• USP-initialized DiT-XL/2 reaches FID 9.73 at 400K steps and 8.93 at 1.2M steps without CFG on ImageNet-256, versus DiT-XL/2 from-scratch at FID 9.62 after 7M steps — 11.7× convergence speedup at parity FID [Table 2, §3.2.1].
• USP-initialized SiT-XL/2 achieves 46.6× convergence speedup vs. the from-scratch SiT baseline; USP SiT-XL/2 400K reaches FID 7.38 vs. REPA (DINOv2-B) 7.9 at the same step count [Table 3, §3.2.2].
• Weight initialization is load-bearing and non-trivial: naively loading MAE weights into DiT-B/2 gives FID 43.55 — worse than from-scratch (42.62) — because the downsampling PatchConv is left random and the input distribution during pretraining doesn’t match the diffusion input [Table 7(a), §3.4].
• Pretraining in the VAE latent space (not pixel space, unlike MAE) is what makes the initialization transfer to diffusion; VAE normalization (mean/std = 0.5 rather than ImageNet stats) is used at both stages to keep the input distribution consistent [§2.2, §3.1].
• The last generation-critical layers of DiT can be randomly reinitialized with graceful degradation (last 2 layers: FID 29.63; last 6 layers: FID 37.58; full USP: 28.26), consistent with the observation that early layers carry recognition features and late layers specialize to generation [Table 7(g), §3.4].
• Total training cost (pretrain + finetune) to reach FID 9.6 is ~15× cheaper than DiT-from-scratch and ~5× cheaper than REPA, and the pretraining cost is amortized across classification/segmentation/generation heads [Table 4, §3.2.4].
• USP maintains ImageNet-1k linear-probe accuracy (66.9% ViT-B, 74.5% ViT-L) above MAE (65.1% / 73.7%) at 800 epochs, and matches MAE on SFT accuracy — the pretraining recipe does not sacrifice recognition to gain generation [Table 5, §3.3.1].
• SFT-initialized DiT-B/2 (from an ImageNet-classification-tuned encoder at 82.6% Top-1) gets FID 39.35 — better than random init (42.62) but much worse than USP self-supervised init (28.26); over-tuning the encoder for discrimination hurts generation transfer [Table 8, §3.5].
• Post-generation-training linear-probe analysis shows the optimal discriminative layer in USP-DiT-XL/2 shifts to layer 20 (of 28), versus REPA’s handcrafted alignment at layer 8, suggesting USP allows the network to self-select the recognition/generation split [Figure 4, §3.5].

Method

USP replaces MAE’s pixel-space reconstruction with masked latent modeling in a frozen VAE’s latent space. Given an image x, a frozen SD-VAE projects it to a 32×32×4 latent; a 2×2 PatchConv converts this to 16×16 patches, which are then masked (default ratio 0.75). A ViT encoder processes unmasked patches; an asymmetric decoder reconstructs the masked patches under MSE with per-patch normalization. The decoder is discarded after pretraining, leaving a ViT + PatchConv that can initialize downstream tasks.

For diffusion transfer, three careful adaptations are needed: (1) reintroduce trainable bias/scale in AdaLN-Zero so the pretrained scale/bias survive; (2) bicubic-interpolate positional embeddings from 224² to 256² for ImageNet generation; (3) drop the class token. The VAE, PatchConv, and encoder weights all transfer; no additional teacher forward is needed at generation-training time (unlike REPA), so USP incurs no extra GPU memory or FLOPs during diffusion training. Pretraining uses only weak augmentation (RandomHorizontalFlip) so VAE-latent features can be cached, giving pretraining throughput comparable to MAE despite the extra VAE encode.

Results

• ImageNet-256, no CFG: DiT-B/2 USP 400K reaches FID 28.26 vs. baseline 42.62; DiT-L/2 USP 400K reaches FID 15.04 vs. baseline 23.03; DiT-XL/2 USP 400K reaches FID 9.73 vs. baseline 19.94 (all at 800-epoch pretrain; 1600-epoch pretrain marginally further improves) [Table 2].
• USP DiT-XL/2 with CFG (marked ‡) reaches FID 2.33 at 1.2M steps, edging out DiT-LLaMA-XL/2 2.42 at 2.5M steps [Table 2].
• Against acceleration methods at DiT-XL/2, 400K steps: USP FID 9.73 vs. REPA (DINOv2-B) 12.3 vs. EQ-VAE 14.5 [Table 2].
• Against acceleration methods at SiT-XL/2, 400K steps: USP FID 7.38 vs. REPA (DINOv2-B) 7.9 vs. EQ-VAE 16.1 vs. ReaLS 14.24 [Table 3].
• ImageNet-1k linear probe: ViT-B-USP 66.9% vs. ViT-B-MAE 65.1%; ViT-L-USP 74.5% vs. ViT-L-MAE 73.7% at 800 epochs [Table 5]. ADE20K semantic segmentation improves ~0.5 mIoU over MAE at matched compute [§3.3.2].
• Compute accounting: reaching FID 9.6 costs 98 units total (36 pretrain + 62 finetune) for USP vs. 622 units for DiT-from-scratch (7M steps) vs. 1904 units for REPA (850K steps + heavy DINOv2 alignment overhead) — a 15× and 5× reduction respectively; pretraining cost amortizes across classification and segmentation heads [Table 4].

Why it’s interesting

USP is a direct architectural competitor to REPA and RAE within the Diffusion training efficiency concept: it lands in the “no external DINO teacher, no extra forward pass during diffusion training” quadrant that No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA) and Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis (Self-Flow) also target, but via a fundamentally different mechanism — initialization, not auxiliary loss. Where SRA and Self-Flow keep the diffusion training loop unchanged and add an internal self-distillation loss, USP moves the entire “learn discriminative features” work into a separate pretraining stage in VAE latent space and then hands off warm weights. The reported 11.7× DiT / 46.6× SiT speedups over from-scratch are in the same order as RAE-family results but achieved with no representation-alignment loss at all.

The result also complicates the Unified Multimodal Models question of whether generation and understanding share representations: USP’s Table 8 shows that a supervised-fine-tuned ImageNet classifier initialization is much worse for diffusion than a self-supervised one — over-specializing for discrimination erodes the transfer to generation. This is a mechanism-level datapoint for why understanding→generation transfer needs to preserve generative reconstructibility of the features (which self-supervised masked modeling does, and cross-entropy classification doesn’t), aligning with UniMRG’s finding (Generation Enhances Understanding in Unified Multimodal Models via Multi-Representation Generation) that the type of understanding signal matters. Notably, USP is a training-time cousin of Representation Autoencoders: both use a frozen VAE latent as the substrate, but RAE replaces the VAE encoder with a frozen DINOv2 for the diffusion input, while USP keeps SD-VAE and instead uses masked reconstruction of those latents as the pretraining objective.

See also

• Diffusion training efficiency — USP contributes a new “warm-start” recipe distinct from external-alignment (REPA) and internal-self-distillation (SRA / Self-Flow) axes.
• Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis — Self-Flow makes the “no external encoder needed” argument via internal EMA-teacher self-distillation; USP makes it via pretraining-then-initialization. Both challenge REPA’s DINO dependency, from different directions.
• No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) — SRA also removes the external teacher during diffusion training; USP removes it entirely (no auxiliary loss anywhere in diffusion training).
• REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training — HASTE identifies that late-stage REPA hurts; USP sidesteps the issue by never adding a REPA loss.
• Diffusion Transformers with Representation Autoencoders — RAE uses a frozen DINOv2 as the encoder; USP uses a frozen SD-VAE and pretrains the DiT itself. Complementary latent-substrate choices.
• Generation Enhances Understanding in Unified Multimodal Models via Multi-Representation Generation — UniMRG’s finding that intrinsic-representation generation aids understanding; USP’s Table 8 finding that classification-SFT initialization is worse for generation than self-supervised init is the same phenomenon read the other direction.