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

AI summary

Overview

Research area: Computer vision, generative modeling, and self-supervised learning (SSL) — specifically the design of compact latent spaces for image synthesis.

Technical level: Advanced. Familiarity with diffusion/flow matching models, variational autoencoders, and vision transformers (especially self-supervised ones like DINOv2) is assumed.

Scope: The paper introduces RepTok, a framework that turns a single fine-tuned token from a self-supervised vision transformer into a complete latent representation for image generation, achieving competitive quality at a fraction of the usual training and inference cost.

What This Paper Is About

Modern image generators (diffusion and flow matching models) are expensive because they operate either on raw pixels or on 2D grid-shaped latent spaces produced by VAEs. Those 2D grids carry a lot of spatial redundancy that is wasteful to model. This paper asks whether a much more aggressive compression is possible: could an entire image be represented by a single continuous token drawn from a self-supervised vision transformer, while still supporting faithful reconstruction and high-quality generation? RepTok answers yes, with light fine-tuning of only the [cls] token and a jointly trained generative decoder.

Key Contributions

1. SSL encoders as direct generative latents. The authors show that self-supervised vision transformers can serve as the latent space itself — not merely as a training guide (as in REPA) — provided the [cls] token is lightly adapted to carry reconstruction-relevant detail.
2. Single-token continuous latent space. RepTok represents an image with one continuous vector (~768 dimensions), eliminating the spatial redundancy of 2D grid latents used by VAE-based diffusion models.
3. Attention-free, highly efficient generative pipeline. Because only one token exists, token-to-token attention becomes unnecessary; the authors use a pure MLP-Mixer for ImageNet generation, cutting training FLOPs by over 90%.
4. Natural extension to text-to-image. By adding cross-attention to a frozen language backbone, RepTok reaches competitive zero-shot MS-COCO performance in under 20 hours of training on four A100 GPUs.

Main Findings

• A single token is sufficient for faithful reconstruction. With just one continuous token, RepTok achieves an rFID of 1.85 and PSNR of 14.94 on ImageNet 256×256 — competitive with or better than tokenizers using 32 to 128 discrete tokens such as TiTok and FlexTok.
• Strong generation at a fraction of the cost. RepTok-L reaches gFID 2.06 (1.88 with CFG) on class-conditional ImageNet, using roughly 1.7% of the training FLOPs consumed by SiT-XL/2.
• Attention-free generation works in the single-token regime. An MLP-Mixer alone — no self-attention — produces competitive generations, saturating at ~40 inference steps versus 64 for SiT, and running about 9× faster end-to-end per sample.
• Cosine-similarity regularization is essential. λ in the cosine loss balances pixel-wise reconstruction (PSNR) against generative quality (gFID): higher regularization yields a smoother, more generate-able latent at modest reconstruction cost; lower regularization gives sharper reconstructions but poor generation.
• The SSL prior matters more than the architecture. Replacing DINOv2 with a randomly initialized encoder produces acceptable pixel-wise reconstructions but a totally unstructured latent space and gFID of 128.5. CLIP, MAE, and DINOv2 all work well; DINOv2 is best.
• The latent space retains smooth interpolation. Linear interpolation between two tokens produces gradual, semantically and spatially coherent transitions (object shape, orientation, presence), unlike unCLIP which only varies high-level semantics.
• Text-to-image scales with the language backbone. Under a fixed 200k-iteration training budget, swapping from CLIP to InternVL to Gemma-2B consistently improves zero-shot COCO metrics, since the language model is frozen and can be scaled independently without increasing the generator's training cost.
• Register tokens carry more local detail than [cls]. DINOv2's [reg] tokens give better pixel-wise metrics while [cls] gives better semantic alignment — an interesting diagnostic about what each token type encodes.

Methodology in Plain English

The pipeline has three steps.

Stage 1 — Adapt the encoder. Start with a frozen, pretrained SSL vision transformer (DINOv2 by default). Unfreeze only the [cls] token embedding; keep every other weight frozen. Train a flow-matching decoder jointly with this token so that the token learns to absorb the low-level visual details needed for reconstruction. To prevent the token from drifting into a chaotic, hard-to-generate space, add a cosine-similarity loss that pulls the adapted token toward the original frozen token. This single loss term replaces the usual KL or vector-quantization regularization used in VAEs and TiTok.

Stage 2 — Train the latent generator. Freeze the adapted encoder. Encode every training image into a single vector. Then train a small flow-matching model — an MLP-Mixer for class-conditional ImageNet, a cross-attention transformer for text-to-image — to sample new vectors from that distribution.

Stage 3 — Inference. Sample a latent token from the generator, then decode it through the Stage 1 decoder back into pixels. Because there is only one conditioning token, the generator is cheap and the decoder's strong conditioning means only ~25–40 function evaluations are needed.

Why This Matters

Research impact. The paper challenges a widely held assumption that generative latent spaces must be 2D grids of many tokens, and shows that pretrained SSL features can play the role of a VAE's latent space with only minimal adaptation. This connects two previously separate lines — self-supervised representation learning and latent generative modeling — in a particularly direct way. It also suggests that a good frozen feature space is itself a form of regularization.

Real-world applications.

• Cheap text-to-image and class-conditional synthesis, useful for teams with limited compute who cannot afford DiT-scale pretraining.
• Low-cost image compression and reconstruction, since a single token suffices to reproduce an image with competitive FID and PSNR.
• Image editing and interpolation, exploiting the smooth, semantically structured latent space for attribute and pose morphing.
• Latent-space research infrastructure, providing a compact representation that can be plugged into reinforcement learning, planning, or retrieval pipelines.

Industry relevance. The >90% reduction in training FLOPs and roughly 9× inference speedup are directly relevant to any company training or serving generative models. The ability to swap in a larger frozen language backbone for free (from the generator's perspective) makes it attractive for products that already own large text encoders.

Future Directions

• Scaling to higher resolutions and modalities. The authors show preliminary 512px fine-tuning; extending the single-token idea to video, audio, or 3D would test how far the compactness argument goes.
• Exploiting further structure inside SSL tokens. DINOv2 provides both [cls] and register tokens — the paper's ablation suggests there is unused signal that could be combined or mixed.
• Better regularization than a cosine loss. The λ trade-off between reconstruction and generation is a first-order knob, not a solution; alternatives might yield a latent space that is simultaneously faithful and generation-friendly.
• Interaction with newer / stronger SSL backbones. Since RepTok's quality tracks the SSL prior, advances in self-supervised vision should transfer directly — an open question is whether the adaptation step (only [cls]) remains sufficient as backbones grow.

Target Audience

Generative modeling researchers and graduate students working on diffusion, flow matching, or latent-space design; practitioners looking for compute-efficient alternatives to DiT/VAE pipelines; and SSL researchers curious about downstream uses of pretrained representations beyond classification and segmentation. Readers should already be comfortable with transformers, latent variable models, and the basics of flow matching to fully appreciate the technical framing.

Authors’ abstract

We introduce Representation Tokenizer (RepTok), a generative modeling framework that represents an image using a single continuous latent token obtained from self-supervised vision transformers. Building on a pre-trained SSL encoder, we fine-tune only the semantic token embedding and pair it with a generative decoder trained jointly using a standard flow matching objective. This adaptation enriches the token with low-level, reconstruction-relevant details, enabling faithful image reconstruction. To preserve the favorable geometry of the original SSL space, we add a cosine-similarity loss that regularizes the adapted token, ensuring the latent space remains smooth and suitable for generation. Our single-token formulation resolves spatial redundancies of 2D latent spaces and significantly reduces training costs. Despite its simplicity and efficiency, RepTok achieves competitive results on class-conditional ImageNet generation and naturally extends to text-to-image synthesis, reaching competitive zero-shot performance on MS-COCO under extremely limited training budgets. Our findings highlight the potential of fine-tuned SSL representations as compact and effective latent spaces for efficient generative modeling.