Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

TL;DR

RepTok represents an image as a single continuous token: the pooled [cls] embedding of a frozen DINOv2 encoder, fine-tuned only at the cls position and paired with a flow-matching decoder. A cosine-similarity loss to the original frozen [cls] keeps the geometry smooth enough to model. Because the latent is a single token, the downstream generator can be attention-free (MLP-Mixer); the combined system reaches ImageNet-256 gFID 3.22 at 4.1k PFlops total training compute, versus 12.1k–212.5k PFlops for DiT-XL / SiT-XL / SiT-XL+REPA at comparable FID. RepTok also scales to T2I, hitting competitive zero-shot MS-COCO numbers after ~20 hours on 4 A100s.

Key claims

Pooled [cls] tokens of a frozen DINOv2 encoder, with only the [cls] embedding fine-tuned, are sufficient as a single-token latent space for faithful image reconstruction (rFID 1.85) and competitive generation (gFID 3.22 with CFG) on ImageNet-256 [§3.2, Section 3.2 Table].
Fine-tuning the [cls] token alone is the load-bearing design choice: freezing it gives poor reconstruction, fully unfreezing the encoder destroys the well-regularized SSL latent geometry, and only the cls-only intervention strikes the balance [Fig. 5, §3.2 cosine-loss subsection].
A cosine-similarity alignment between the fine-tuned [cls] and the frozen-encoder [cls] (Equation 3) is required to keep the latent space generatable; ablating it produces a “w/o prior” failure where the generator cannot capture the latent distribution [§4.3 ablation, Fig. 9].
In the single-token regime, token-to-token interactions are unnecessary, so an attention-free MLP-Mixer trained with flow matching can serve as the latent generator without quality loss [§3.3, Fig. 1].
RepTok matches strong SiT-XL+REPA quality (FID 1.42 with CFG=1.5) only loosely (3.22 vs 1.42) but at a fraction of training compute — RepTok 4.1k PFlops vs SiT-XL+REPA 143.9k PFlops at 4M / 7M iterations respectively [Table 3].
Generation–reconstruction trade-off is governed by the cosine-loss weight γ: stronger regularization (higher γ) improves gFID at the cost of PSNR, and a “mild” γ improves generation substantially while only minimally hurting reconstruction [§4.3, Fig. 9].
RepTok generalizes across SSL backbones (DINOv2, MAE, CLIP) but breaks completely with a randomly-initialized encoder — a semantic prior is what gives the latent space its smoothness and makes it modellable [§4.3].
Classifier-free guidance gives only limited gains under RepTok — the same phenomenon RCG reports for purely-semantic codes; in Table 3 RepTok without CFG = 3.4 vs with CFG=1.5 = 3.22 [§4.1 quantitative comparison].
The single continuous-token formulation supports smooth latent-space interpolation between two images — gradual transitions in shape, size, rotation, and emergence — which spatial-grid latents do not afford [§4.1, Fig. 4].
RepTok extends to T2I with cross-attention to a frozen language backbone; with a 120M COYO subset and 200k steps on 4 A100s (~20 hours) it reaches competitive zero-shot MS-COCO FID, and larger frozen language backbones (CLIP → InternVL → Gemma-2B) monotonically improve all metrics without affecting generator training cost [§4.2, Fig. 7].
RepTok beats the closest prior 1-token method RCG on both perception and pixel-level metrics (FID@50K 1.85 vs 3.20; PSNR 14.94 vs 9.31) — the fine-grained-information injection into [cls] is what closes RCG’s reconstruction gap [Table 2].

Method

RepTok has two stages.

Stage 1 — single-token autoencoder. A frozen DINOv2 ViT (only the [cls] token embedding is trainable) encodes an image into a single continuous vector $z_{\text{cls}}$ . A flow-matching generative decoder (MMDiT-style: concatenate $z_{\text{cls}}$ with the noisy image tokens) is trained jointly with the cls-only fine-tune via the standard flow-matching loss (Eq. 2). A cosine-similarity alignment loss (Eq. 3) holds $z_{\text{cls}}$ close to the original frozen [cls], with weight γ controlling how far the token may drift. No KL, no VQ, no perceptual or adversarial loss is needed — the SSL prior is the regularizer.

Stage 2 — latent-space flow. With the encoder frozen after Stage 1, images are pre-encoded and a separate flow-matching generator is trained on the $z_{\text{cls}}$ marginal. For class-conditional ImageNet, the generator is an MLP-Mixer with class-embedding concatenation and CFG-style class-drop; no attention is used because there is only one token. For T2I, 4 learnable tokens are concatenated with the noisy [cls] and cross-attended against a frozen language backbone (CLIP / InternVL / Gemma-2B). The DiT-XL-class generator is 276M parameters at 23.0 GFLOPs/iter — roughly 1/5 the FLOPs/iter of DiT-XL/2 (118.6) or SiT-XL+REPA (140.5).

Results

Reconstruction (ImageNet-256): RepTok rFID 1.85 with 1 token, beating LDM 32×32 (0.90 rFID but 1024 dims), and competitive with TiTok-B/64 (rFID 1.70) and FlexTok d18-d18 (rFID 1.61) despite using 32–256× fewer tokens [Section 3.2 Table].
Class-conditional generation (ImageNet-256): RepTok gFID 3.4 (no CFG) / 3.22 (CFG=1.5) at 700k steps; 5.4 already at 100k steps. Compute: 4.1k PFlops total. Comparators: SiT-XL/2 reaches FID 2.06 (CFG=1.5) at 212.5k PFlops; SiT-XL/2+REPA reaches 1.42 (CFG=1.5) at 143.9k PFlops [Table 3].
T2I (zero-shot MS-COCO): competitive FID at substantially lower training cost than prior T2I models, after only 200k iterations and ~20 hours on 4 A100s; larger frozen LM backbones monotonically improve FID + CLIP + validation loss [§4.2, Fig. 7, Fig. 8].
Ablation — encoder choice: DINOv2 (best generation), MAE, CLIP all yield strong reconstruction and generation; randomly-initialized ViT yields good pixel reconstruction but a generator-uncapturable latent distribution, so gFID collapses [§4.3].
Ablation — cosine-similarity loss: tunable γ produces an explicit gFID-vs-PSNR Pareto curve; mild γ improves gFID while costing very little PSNR; full ablation in Fig. 9.
vs RCG (closest prior single-vector method): RepTok FID@50K 1.85 / PSNR 14.94 vs RCG 3.20 / 9.31, attributed to RCG using purely-semantic codes while RepTok injects reconstruction-grade detail into [cls] [Table 2].

Why it’s interesting

RepTok is the cleanest filed expression of “the SSL [cls] itself is the latent” — orthogonal to where Luma’s existing wiki entries put DINO supervision. Diffusion Transformers with Representation Autoencoders (RAE, Saining Xie) uses all 256 DINOv2 patch tokens as the latent and trains a wide DiTDH on them; Improved Baselines with Representation Autoencoders (RAEv2) adds multi-layer aggregation; both stay in the spatial-grid paradigm and inherit DiT compute. RepTok throws away the spatial grid entirely — one continuous token, attention-free generator — and reaches gFID 3.22 at ~3% of SiT-XL+REPA’s training PFlops. The comparison is not apples-to-apples (RAE/RAEv2 are still ahead on absolute FID at higher compute), but RepTok puts a concrete lower bound on how much latent dimensionality matters once the prior is strong enough.

The complementary direction is Unified Latents (UL): How to train your latents (UL), which argues the right fix is to redesign the encoder–prior coupling so bitrate is an explicit hyperparameter; RepTok agrees that latent design is the lever, but takes the opposite extreme — minimal capacity, maximal prior. The CFG-insensitivity Terrance flagged is notable: RCG reports the same, and the natural reading is that single-vector conditioning is too weak a signal for guidance to amplify. Whether this is a ceiling (Wurstchen-style: cheap to start, hard to scale) or a parameterization issue (the class embedding is concatenated, not adain’d or cross-attended) is exactly what a scaling sweep would settle, and the paper does not run one.