Generative Representation Learning for Scalable Latent Tokenization

TL;DR

RePR (Baade et al., Stanford) proposes a self-supervised representation-learning objective that factorizes the pointwise mutual information (PMI) between two views of the data and estimates that PMI with a jointly-trained masked generative model, then supervises the encoder’s dot-product similarity to equal the PMI. The payoff is a representation that is simultaneously (i) principled — no negatives, no batch statistics, no non-masking augmentations, and no feature collapse under watermarked adversarial data where DINO and LeJEPA collapse — and (ii) scalable: compute spent on the tokenizer encoder yields predictable, monotone improvement in downstream generation. Combined tokenizer-plus-JiT diffusion is compute-optimal versus raw pixel-space diffusion at matched budget, hitting FID 2.18 with JiT-L on ImageNet 256² and extrapolating to 7–8× fewer FLOPs than DINOv2 / WebSSL baselines on ImageNet and the GPIC text-to-image dataset.

Key claims

Contrastive learners implicitly factorize the PMI of the data, and a jointly-trained generative model can directly estimate PMI as the information difference between its own conditional predictions — supervising the encoder’s dot product to match that PMI gives a representation loss with no negatives, batch statistics, or non-masking augmentations [§Thread tweet 4, §Fig. “Model diagram”].
On watermarked data where every patch contains a seed identifying the batch index — a worst-case shortcut for discriminative objectives — RePR continues to learn useful features while DINO and LeJEPA collapse to the shortcut [Thread tweet 7 / Fig. “watermark experiment”].
Scaling the tokenizer encoder alone (holding the downstream generator fixed) produces predictable, monotone improvement in generated-image FID; each point on the sweep is a Latent Forcing model trained with no GAN, no LPIPS, no auxiliary decoder, and no external models [Thread tweet 8 / Fig. “tokenizer encoder scaling laws”].
Allocating 20% or 35% of total (tokenizer + diffusion) FLOPs to the tokenizer outperforms JiT at matched compute, so tokenizer + diffusion is a strict Pareto improvement over raw pixel diffusion on ImageNet 256² [Thread tweet 9 / Fig. “compute optimality”; Abstract: FID 2.18 with JiT-L].
At the extrapolated end of the fits, RePR representations match the quality of DINOv2 at 7× fewer FLOPs on ImageNet and WebSSL at 8× fewer FLOPs on the text-to-image dataset GPIC [Abstract].

Method

The objective takes two views $x, y$ of the data, embeds them into vectors $R(x), R(y)$ via the encoder being trained, and supervises their dot product to equal the PMI of the two views. The PMI is computed by a jointly-trained masked generative model as the information difference between conditional predictions. This replaces the usual contrastive-learning hacks (negatives, batch statistics, strong augmentations) with a single generative-information target.

Two practical corrections are needed: the PMI is dominated by high-frequency bits (addressed via timestep weighting) and becomes intractable as input dimension grows (addressed via bootstrapping — the authors note that one can take an already-learned representation and learn a second encoder on top of it, going from a pixel encoder to a latent encoder for better downstream performance). The paper and project blog cover the specifics; the thread flags both as load-bearing.

Downstream, the representation is used as the latent for a diffusion generator (JiT-L on ImageNet 256²), and the combined-compute Pareto is measured against raw pixel diffusion.

Results

ImageNet 256²: FID-50K of 2.18 with JiT-L on the RePR latent, outperforming direct pixel-space diffusion at matched compute [Abstract].
Compute optimality: training the tokenizer for 20% and 35% of total tokenizer+diffusion FLOPs outperforms the no-tokenizer JiT baseline — a new Pareto frontier for pixel-space generation [Thread tweet 9].
Encoder scaling laws: holding the diffusion generator fixed and scaling only the encoder gives predictable improvement in generated-image FID with no sign of saturation across the swept range [Thread tweet 8].
Representation efficiency at extrapolation: 7× fewer FLOPs than DINOv2 on ImageNet, 8× fewer FLOPs than WebSSL on the GPIC text-to-image dataset at matched representation quality [Abstract].
Collapse robustness: on watermarked ImageNet, DINO and LeJEPA collapse to the batch-index shortcut; RePR continues to learn meaningful features [Thread tweet 7].

Why it’s interesting

RePR is a third design point in the Representation Autoencoders cluster, orthogonal to both the frozen-encoder lineage (RAE on DINOv2 Diffusion Transformers with Representation Autoencoders, FAE’s compression-side rework One Layer Is Enough: Adapting Pretrained Visual Encoders for Image Generation) and VTP’s from-scratch joint CLIP + DINO-SSL + reconstruction pre-training Towards Scalable Pre-training of Visual Tokenizers for Generation: here the encoder’s training signal is the generator’s own information difference, so there is no external DINO teacher to replace, bootstrap from, or distill. If the compute-optimality and encoder-scaling-law claims hold up, this is the first filed representation-learning recipe where “spend more encoder FLOPs” is a Pareto improvement over “spend them on the generator” without the collapse-avoidance hacks that have forced DINOv2/v3 and LeJEPA into multiple heuristic rewrites. The approach also contests Self-Flow’s “external-DINO supervision is a transitional crutch” position Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis from the opposite direction: Self-Flow replaces the external teacher with internal self-distillation; RePR replaces it with a jointly-trained PMI oracle. Finally, the “take a learned representation and encode it again” bootstrapping property gives a natural knob for the pixel-space-diffusion crowd (PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion, PixelDiT: Pixel Diffusion Transformers for Image Generation) to iterate latent spaces without ever leaving the self-supervised regime.