Generative Representation Learning for Scalable Latent Tokenization

Alan Baade; Kyle Sargent; Vincent Sitzmann; Eric Ryan Chan; Justin Johnson; Li Fei-Fei; Gordon Wetzstein

cs.LGOct 14, 2026

Generative Representation Learning for Scalable Latent Tokenization

Authors: Alan Baade, Kyle Sargent, Vincent Sitzmann, Eric Ryan Chan, Justin Johnson, Li Fei-Fei, Gordon Wetzstein

Organizations: Stanford University · Massachusetts Institute of Technology · University of Michigan

Abstract

We introduce an approach for self-supervised representation learning designed to be used as a tokenizer for generative modeling. Our objective supervises a representation loss by factorizing the pointwise mutual information of the data, which we estimate by taking information differences between conditional predictions from a learned generative model. We justify our algorithm theoretically and empirically, demonstrating that our generative approach to representation learning avoids dimensional collapse and feature suppression in cases where commonly used methods like DINO and LeJEPA fail. Moreover, we show our representation learning technique is scalable; expending more compute leads to predictably better latent tokenization. For the first time, we perform a careful analysis of the combined flop efficiency of tokenizer and generative model and demonstrate that our approach is a compute-optimal improvement over direct diffusion for pixel-space generation, obtaining an FID-50K of 2.18 with JiT-L. Applied to ImageNet and the text-to-image dataset GPIC, we show tokenizer encoder scaling laws for diffusability, where extrapolating our approach achieves representations in 7x fewer FLOPs than DINOv2 on ImageNet and 8x fewer FLOPs than WebSSL on GPIC.

Figures & tables

Figure 1: Architecture. Left: 1) Our Prediction Model (bottom) estimates weighted data likelihoods for a target patch

p

using an MAR architecture, which combines an MAE Encoder

E_{\theta}

and Decoder

D_{\theta}

with a diffusion prediction MLP head

M_{\theta}

. 2) The Representation Model

R_{\theta}

(top) outputs vector representations supervised to factorize a weighted PMI. Right: We enforce spatial equivariance in

R_{\theta}

by hiding absolute position information but outputting a feature for every patch.

Figure 2: Discriminative representation learning can suffer from feature suppression when a small subset of features, such as a watermark, is sufficient to distinguish examples. This degrades representation quality, as reflected by collapsing ImageNet class-probing performance. RePR remains robust to feature suppression.

Figure 3: Tokenizer Encoder Scaling Laws: RePR predictably scales as a tokenizer encoder. When holding size and training duration of a generative model constant, greater RePR size and FLOPs improve generation performance. Iterative feature prediction (Sec. 5.4 ) yields further improvement.

Figure 4: Total (Encoder + Diffusion) training compute versus FID-10K for JiT and JiT with Latent Forcing using RePR, showing that representation learning is compute-optimal for generation.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Monte Carlo PMI convergence versus number of Monte Carlo Integration Steps, measured as the MSE between the PMI when estimated with

N

steps and the PMI when estimated with 4096 steps. Pixel space prediction, default configuration RePR-173M. Furthermore, sharing the noise among different likelihood calculations strongly reduces variance. We empirically choose

N=32

steps at which diffusion integration uses

1.3\%

of total flops during training.

Model	rFID	Reconstruction PSNR
DINOv2-B ( Oquab et al., 2024 )	0.49	18.09 dB
RePR-B Iter. 3 600ep (Ours)	0.47	21.72 dB

Appendix

Table 5: Reconstruction performance for RePR-B Iter 3, 600ep, using the decoder configuration from RAE ( Zheng et al., 2026 ) .

config	ImageNet	ImageNet	GPIC
decoder	JiT-B ( Li and He, 2026 )	JiT-L ( Li and He, 2026 )	MiniT2I ( Wang et al., 2026 )
hidden dim	768	1024	768
backbone depth	12	20	17
output experts Baade et al. (2026)	2 $\times$ 1024 dim	2 $\times$ 1024 dim	-
text encoder	-	-	FLAN-T5-large
max text tokens	-	-	128

Appendix

Table 7: Latent Forcing ( Baade et al., 2026 ) decoder training at 256

\times

256 resolution with 16

\times

16 image patches. JiT-B used for scaling laws, JiT-L used for baseline comparisons.

Figure 6: Uncurated class-conditional Imagenet Generations, from our compute optimal JiT-L.

Figure 7: Feature cosine similarity between patch features and a center patch feature shown in red, with and without spatial equivariance (Sec 5.2 ).

Figure 8: Example Clean and Watermarked images from Imagenet, used in Sec. 6.1 . We follow the watermarking strategy from Enigma (2026) exactly, except for that we shrink watermarks from

32\times 32

16\times 16

pixels, so that every individual patch provides sufficient signal for suppression.

	Iteration 1: Pixel		Iteration 2: Latent
Encoder (width)	Epochs	FID	Epochs	FID
20M (256)	110	41.61	$210+110$	36.24
	210	40.19	$210+210$	34.89
	410	37.01	–	–
	610	37.21	–	–
44M (384)	110	40.93	$210+110$	31.64

Appendix

Table 8: ImageNet scaling result tables, evaluated using FID-10K after 80 epochs of generative model training. Iteration 2 epoch counts include 210 epochs of Iteration 1 training (200 epochs + 10 epochs of learning rate decay). The DINOv2 baseline obtains 17.95 FID at

1.1685129\times 10^{22}

encoder-training FLOPs.

Encoder (width)	ImageNet Epochs-Worth	FD	ImageNet Epochs-Worth	FD
	Iteration 1: Pixel		Iteration 2: Latent
44M (384)	110	514	$210+110$	416
	210	479	$210+210$	393
	410	456	–	–
78M (512)	110	498	$210+110$	384
	210	463	$210+210$	365

Appendix

Table 9: GPIC scaling result tables, evaluated using FD-DINOv2-10K after 1 GPIC epoch of generative model training. For consistency between approaches, epochs are measured as effective epochs on ImageNet. For reference, 80 ImageNet epochs is roughly equal to 1 GPIC epoch. The WebSSL-DINO-1B reference obtains an FD-10K of 277.1 at

9.78\times 10^{21}

encoder-training FLOPs. MiniT2I obtains an FD-10K of 546.3.

Representation Learning

Representation learning develops data transformations that expose useful structure for prediction, transfer, generation, and control.

See all 804 papers

FeaTok: Feature-Loss Tokenization for Image Generation
Oct 15, 2026Zhang Xiaofeng, Samuel Lavoie, Nicolas Beltran-Velez +6Image Tokenization Image Generation
Mila - Québec AI Institute · Université de Montréal · FAIR at Meta +4
Balancing Image Compression and Generation with Bootstrapped Tokenization
Jun 4, 2026Haozhe Chi, Jinghan Li, Hao Jiang +4Learned Image Compression Image Tokenization
Peking University · Central Media Technology Institute, Huawei
Representation by Design in Generation: Cross-View Class-Token Alignment in Diffusion Transformers
Sep 28, 2026Xiaoyu Wu, Yifei Wang, Chen WeiFlow Matching Diffusion Transformer
Carnegie Mellon University · Rice University

Oct 15, 2026J/K move · Enter open · S save

FeaTok: Feature-Loss Tokenization for Image Generation

Zhang Xiaofeng, Samuel Lavoie, Nicolas Beltran-Velez, Jianhao Yuan+5

Mila - Québec AI Institute · Université de Montréal · FAIR at Meta+4

Image Tokenization Image Generation T2I Generation Conditional Image Generation

Open paper AskPDF arXiv

Discrete image generators learn the distribution of tokens produced by an image tokenizer. These token representations should be easy for generators to model, not merely enable faithful image reconstruction. Yet these tokenizers are commonly trained to reconstruct pixels, which prioritizes low-level fidelity and may disrupt representation structure useful for generation. Motivated by feature-predictive self-supervised learning, we introduce feature-loss tokenization (FeaTok), which trains discrete tokens by reconstructing pretrained features rather than pixels. Ablations across SigLIP, SigLIP2, DINOv2, and DINOv3 show better generative performance with FeaTok. In class-conditional generation, FeaTok advances the quality-compute Pareto frontier, reducing gFDr6 at 256x256 by up to 29 percent over the strongest discrete baselines and outperforming the strongest unguided continuous baseline (RAEv2) with 38 times less sampling compute. FeaTok also achieves higher human-preference reward-model scores than prior art without reward-based training and surpasses continuous and discrete text-conditional baselines on three benchmarks with only 500 fine-tuning updates after pre-training.

Abstract

Representation Learning

Explore similar work

FeaTok: Feature-Loss Tokenization for Image Generation

Balancing Image Compression and Generation with Bootstrapped Tokenization

Representation by Design in Generation: Cross-View Class-Token Alignment in Diffusion Transformers