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

TL;DR

RepTok proposes representing an image with a single continuous latent token derived from a pre-trained self-supervised vision transformer, by fine-tuning only the [cls] embedding and training a flow-matching decoder. A cosine-similarity loss keeps the adapted token aligned with the frozen SSL space, preserving its smooth, semantically structured geometry for stable generation. The approach uses an attention-free MLP-Mixer as the latent-space generator and a 2D VAE-based decoding pathway to enable efficient training on ImageNet-scale data, achieving competitive class-conditional generation with drastically reduced compute. It naturally extends to text-to-image synthesis by incorporating cross-attention for prompts, delivering competitive zero-shot MS-COCO results under limited training budgets. Overall, RepTok demonstrates that SSL representations can serve as compact, effective latent spaces for efficient generative modeling when appropriately fine-tuned and regularized.

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.

Figures (19)

• Figure 1: Comparison of our single-token MLP-Mixer generator against transformer-based baselines (DiT, SiT), as well as representation-aligned models like REPA. RepTok attains competitive generative performance while reducing training cost by over 90% owing to its compact latent space and lightweight architecture. All results reported without CFG. For fair comparison, we use our encoder and decoder that have been trained on general data.
• Figure 2: Overview of our pipeline. (a) Joint fine-tuning of the [cls] token of SSL encoder $\mathcal{E}$ and training of the generative decoder $\mathcal{D}$ for image reconstruction. (b) Training of the generation model $\mathcal{G}$ to synthesize frozen encoder outputs, which constitute the latent space $z = \mathcal{E}(x)$. (c) Inference pipeline, where the latent space $z$ is first generated and subsequently decoded into the pixel space.
• Figure 3: We introduce RepTok, a compact visual tokenizer that builds upon pre-trained SSL representations. Our approach augments these representations with additional necessary information to enable images to be faithfully encoded as a single continuous token, which allows for both high-fidelity image reconstruction and synthesis. The third row indicates the number of tokens for reconstruction.
• Figure 4: Latent space interpolation. We observe smooth transitions not only in semantic content but also in spatial configuration. This indicates that our method successfully integrates low-level spatial information while preserving the properties of the pretrained encoder's latent space, and facilitates generation within the learned representation. We provide more samples in the Appendix.
• Figure 5: