Source-linked AI summary

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

Ming Gui, Johannes Schusterbauer, Timy Phan, Felix Krause, Josh Susskind, Miguel Angel Bautista, Björn Ommer

arXiv:2510.14630v2cs.CV

TL;DR

Latent diffusion reduces generation costs through compression, but its 2D grid still contains spatial redundancy. RepTok adapts a pooled SSL token into a single continuous latent and trains a generative decoder around it. The paper reports competitive class-conditional ImageNet generation and competitive zero-shot MS-COCO results under a limited training budget.

  • Problem

    Latent diffusion reduces computation by compressing images, but its 2D grid structure still fails to exploit natural images’ spatial redundancies.

    Source: §1 · p. 1
  • Method

    RepTok lightly fine-tunes an SSL encoder’s class token and jointly trains a flow-matching decoder, using cosine similarity to preserve the latent’s smooth geometry.

    Source: §1 · p. 2, §3.2 · p. 6
  • Results

    RepTok achieves competitive class-conditional generation at a fraction of transformer-based baselines’ computational cost and competitive zero-shot MS-COCO results with under 20 hours of training on four A100 GPUs.

    Source: §1 · p. 2, §5 · p. 10
  • Takeaways & Limitations

    A fine-tuned pooled SSL representation can serve as a compact latent for faithful reconstruction and efficient image synthesis.

    Source: §1 · p. 2, §5 · p. 10

    A single token may limit fine-grained detail in complex or high-resolution scenes, and precise control over object location and scene composition remains an open challenge.

    Source: C LIMITATIONS · p. 17

Abstract

from arXiv · show

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.

1 INTRODUCTION

Latent diffusion reduces image-generation costs by moving generation into a compressed space, but its 2D grid still contains spatial redundancy. RepTok adapts a pooled SSL token into a compact, reconstruction-capable latent and uses it for efficient generation.

1 INTRODUCTION

  • Latent diffusion lowers computational costs by compressing images, but its 2D grid structure does not exploit natural images’ spatial redundancies.The compressed space abstracts away imperceptible details so generation can focus on semantic content.Source: §1 · p. 1
  • TiTok compresses images into short 1D discrete sequences, while REPA uses cosine similarity to align diffusion features with pretrained SSL representations.These prior directions motivate RepTok’s direct use of an adapted pooled SSL token.Source: §1 · pp. 1–2
  • RepTok uses the pooled [cls] token as a compact latent that retains semantic information and enough spatial detail for high-fidelity reconstruction.This single-token representation eliminates spatial redundancies in 2D grid-based latents.Source: §1 · p. 2
  • RepTok lightly fine-tunes a pretrained SSL encoder’s semantic token and jointly trains a decoder with flow matching and cosine-similarity regularization.The regularizer keeps the adapted latent close to the original smooth, semantically structured space; the decoder needs no perceptual or adversarial auxiliary loss.Source: §1 · p. 2
  • RepTok achieves competitive ImageNet generation at a fraction of transformer-based baselines’ compute and extends to competitive zero-shot MS-COCO synthesis with under 20 hours of training on four A100 GPUs.The compact latent and lightweight architecture reduce training cost by over 90%.Source: §1 · p. 2, Figure 1 · p. 2

2 RELATED WORK

Prior work uses compact latent spaces and pretrained representations to improve generative modeling, while tokenization methods explore shorter or continuous representations. RepTok builds on these ideas with a single continuous pooled SSL token.

2 RELATED WORK

  • Latent-space generation moved from compact codebooks to diffusion in learned spatial latents, which preserve semantics while abstracting perceptual detail.The passage describes latent-space methods as improving sampling quality and efficiency.Source: §2 · pp. 2–3
  • The pipeline adapts the SSL encoder’s [cls] token with a reconstruction decoder, trains a generator on frozen encoder outputs, then decodes generated latents into pixels.The latent is z = E(x).Source: Figure 2 · p. 3
  • REPA accelerates diffusion training by aligning diffusion features with DINO embeddings, whereas RepTok adds low-level visual information for faithful reconstruction and generation.Other SSL-feature approaches use full spatial grids; RepTok instead uses only a pooled semantic token.Source: §2 · p. 3
  • TiTok, ElasticTok, and FlexTok explore compact discrete token sequences, while RepTok uses a continuous latent suited to differentiable diffusion training.Diffusion Autoencoders require an additional noise-space subcode for reconstruction, whereas RepTok reconstructs from one latent alone.Source: §2 · p. 3
  • RepTok augments pretrained SSL representations so an image can be encoded as one continuous token for high-fidelity reconstruction and synthesis.The figure’s third row indicates the token count used for reconstruction.Source: Figure 3 · p. 4

3 METHOD

RepTok adapts a pretrained SSL encoder’s class token into a single continuous image representation and jointly trains a generative decoder. Flow matching supports both reconstruction and generation, while cosine alignment preserves the pretrained latent space’s structure.

3.1 PRELIMINARIES

  • Flow matching learns a vector field between a prior distribution and the target data distribution, and sampling integrates the learned ODE.The section describes a linear interpolant between the prior and target, with Euler integration as one sampling approach.Source: §3.1 · p. 4

3.2 REPTOK: REPRESENTING IMAGES AS A SINGLE TOKEN

  • RepTok uses a single continuous token from a pretrained encoder with a generative decoder, extending beyond multi-token encodings such as TiTok.The decoder is jointly trained with the encoder using flow matching; a pretrained 2D VAE keeps decoding in a learned latent space.Source: §3.2 · pp. 4–5
  • Fine-tuning only the [cls] token injects reconstruction-relevant detail while the remaining encoder weights stay fixed.Freezing the token gives poor reconstructions, whereas unconstrained encoder training pulls it away from the pretrained space; cosine alignment limits that drift.Source: §3.2 · pp. 5–6
  • Figure 5 compares ground truth, frozen-token, and fine-tuned-token reconstructions; Table 1 compares tokenizers for reconstruction and class-conditional ImageNet generation.The supplied table caption identifies the evaluation scope but does not provide metric values.Source: Figure 5 · p. 5, §3.2 · p. 5, Table 1 · p. 6

3.3 SINGLE TOKEN GENERATION FOR IMAGE SYNTHESIS

  • A separate flow-matching model generates samples in RepTok’s continuous latent space, whose frozen SSL encoder provides structure without explicit KL or vector-quantization regularization.The latent is typically represented by one token in R1×768.Source: §3.3 · p. 6
  • Single-token inputs allow ImageNet generation with an attention-free MLP, while text-to-image uses attention for text conditioning at low cost.The paper attributes the reduced need for token interactions to the aggressive compression of the input representation.Source: §3.3 · p. 6

4 EXPERIMENTS

The experiments evaluate RepTok on ImageNet class-conditional generation and text-to-image synthesis, as well as its computational efficiency and design choices. They report competitive generation with substantially lower training compute, and examine generalization, scaling, and reconstruction–generation trade-offs.

4 EXPERIMENTS

  • Experiments assess reconstruction with rFID, PSNR, SSIM, and LPIPS, and generation with gFID; models operate at 2562 resolution.Implementation and training details are provided in the Appendix.Source: §4 · pp. 6–7

4.1 CLASS-CONDITIONAL GENERATION

  • 1.7% of SiT’s FLOPs are required by RepTok, which remains competitive on ImageNet generation while using a single continuous token and an MLP-only generator.The comparison accounts for first-stage encoder inference cost; Table 2 reports FID, parameter counts, and training FLOPs.Source: Table 2 · p. 7, §4.1 · pp. 7–8
  • RepTok achieves significantly higher PSNR and lower FID than RCG, with results indicating stronger performance across pixel-level and perceptual metrics.Table 3 covers reconstruction performance on ImageNet 2562.Source: §4.1 · p. 7, Table 3 · p. 7

4.2 ENABLING REPTOK FOR T2I GENERATION

  • RepTok achieves substantially lower training cost than prior text-to-image models while maintaining competitive zero-shot FID on COCO after 200k training iterations.Larger frozen language backbones consistently improve performance across all reported metrics.Source: Figure 7 · p. 8, §4.2 · p. 9
  • After fewer than 20 hours on 4× A100 GPUs, RepTok produces realistic text-to-image generations that capture fine details and closely follow prompts.These results follow 200k training iterations.Source: §4.2 · p. 9, Figure 8 · p. 9

4.3 ABLATIONS

  • The method generalizes to MAE and CLIP encoders, while a randomly initialized encoder yields high pixel-wise reconstruction but an unstructured latent space with high generation FID.The authors attribute smooth, low-dimensional manifolds that promote stable generations to a semantic prior.Source: Table 4 · p. 8, §4.3 · p. 9
  • Stronger cosine-similarity regularization improves gFID at the cost of PSNR, while mild regularization improves generation with minimal reconstruction degradation.The loss weight λ balances preservation of high-level semantics against reconstruction of low-level visual details.Source: §4.3 · pp. 9–10, Figure 9 · p. 10
  • The latent-space MLP-Mixer scales with parameter size and FLOPs, and larger variants with continued training yield improved performance.The largest evaluated XL model contains 516M parameters.Source: §4.3 · p. 10, Figure 10 · p. 10

5 CONCLUSION

The conclusion presents RepTok as a compact latent-space framework that adapts self-supervised representations for generative modeling. It reports competitive generation and argues for the potential of SSL representations as lightweight generative-model latents.

5 CONCLUSION

  • RepTok adapts self-supervised representations into a compact latent space for generative modeling by fine-tuning the class token and regularizing it with cosine similarity.A flow-matching decoder uses the resulting single continuous token, which retains smooth geometry while adding reconstruction-relevant information.Source: §5 · p. 10
  • RepTok enables faithful reconstruction and efficient synthesis without costly attention mechanisms or auxiliary losses.The conclusion attributes this to coupling the adapted token with a generative decoder trained via flow matching.Source: §5 · p. 10
  • The single-token formulation achieves competitive class-conditional generation and extends to more complex text-to-image settings.The conclusion frames fine-tuned SSL representations as potential compact and effective latents for efficient generative modeling.Source: §5 · p. 10

A IMPLEMENTATION DETAILS

The implementation uses a DiT-XL/2 flow-matching decoder and a standard MLP-Mixer generator. It specifies the decoder’s training setup and the Mixer’s architecture and conditioning inputs.

A IMPLEMENTATION DETAILS

  • The generative decoder is DiT-XL/2, trained for one million steps with AdamW, a 10^-4 learning rate, 2000-step linear warm-up, and global batch size 512 on 8 H100 GPUs.The implementation uses RoPE, RMSNorm, and SwiGLU, which the authors report improve decoder stability and performance.Source: A IMPLEMENTATION DETAILS · p. 15
  • The decoder concatenates the SSL embedding with patch tokens and applies full self-attention over all tokens.Its implementation also uses RoPE, RMSNorm, and SwiGLU.Source: A IMPLEMENTATION DETAILS · p. 15
  • The MLP-Mixer has hidden dimension 1280, depth 28, and expansion factors 4 for the channel MLP and 2 for the token MLP.CLIP text embeddings or class tokens are concatenated with the noisy image token as conditioning information.Source: A IMPLEMENTATION DETAILS · p. 15

B ADDITIONAL RESULTS

Additional results compare token types, interpolation behavior, and decoder inference settings, and examine extensions to higher resolution and additional tokens. They also include qualitative samples and comparisons with unCLIP and alternative SSL priors.

B ADDITIONAL RESULTS

  • Register tokens contain more appearance, location, and orientation information than the frozen [cls] token, but neither token type yields proper pixel-wise reconstructions.The table comparison reports that register tokens improve pixel-wise metrics, indicating stronger local information.Source: B ADDITIONAL RESULTS · p. 15, Table 6 · p. 15, Figure S4 · p. 15
  • Reconstruction FID improves as function evaluations increase, then saturates around 20; the authors hypothesize decoder conditioning reduces the need for further refinement.The comparison uses the ImageNet validation dataset.Source: B ADDITIONAL RESULTS · p. 15
  • Single-token interpolations show gradual changes in semantics and low-level details, including object shape, size, emergence, and rotation.The appendix reports smooth changes in object spatial configuration, especially rotation.Source: Figure S1 · p. 16, B ADDITIONAL RESULTS · p. 16
  • Compared with unCLIP, RepTok preserves semantic and structural continuity in interpolations, whereas unCLIP representations primarily capture semantics and lack low-level detail.The qualitative comparison uses a pretrained Stable Diffusion 2.1 unCLIP checkpoint.Source: Figure S2 · p. 16
  • Fine-tuning the encoder–decoder for 100K iterations adapts the model to 5122 resolution from weights pretrained at 2562.Additional text-conditional and uncurated class-conditional generations are also shown.Source: B ADDITIONAL RESULTS · p. 16
  • Additional global tokens can improve reconstruction but worsen generation, so RepTok restricts its method to [cls] tokens and leaves multi-token extensions for future work.Such extensions require SSL encoders with additional non-spatial tokens because spatial tokens have no straightforward global repurposing.Source: B ADDITIONAL RESULTS · p. 16, Figure S5 · p. 17

C LIMITATIONS

The limitations concern reconstruction detail and control: a single token may underspecify complex or high-resolution scenes, while cosine regularization trades pixel fidelity against latent smoothness. ImageNet-only text-to-image generation also struggles with compositional multi-object prompts.

C LIMITATIONS

  • A single token may limit fine-detail expressiveness in complex or high-resolution scenes, and fine-grained object placement and scene composition remain open challenges.The experiments show preservation of some low-level spatial structures, but do not resolve fine-grained control.Source: C LIMITATIONS · p. 17
  • Cosine-similarity regularization can suppress low-level detail, limiting applicability where very high visual reconstruction fidelity is critical.Stronger regularization improves latent smoothness and structure for stable generation but can degrade pixel-wise reconstruction.Source: C LIMITATIONS · p. 17
  • ImageNet-only text-to-image generation struggles with compositional prompts requiring multiple objects, while fine-tuning on more diverse data alleviates this issue.The passage attributes the difficulty to ImageNet’s object-centric bias and limited exposure to multi-object scenes.Source: C LIMITATIONS · p. 18
Loading 2510.14630v2…