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θ and Decoder Dθ with a diffusion prediction MLP head Mθ . 2) The Representation Model Rθ (top) outputs vector representations supervised to factorize a weighted PMI. Right: We enforce spatial equivariance in Rθ 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 × 1024 dim
2 × 1024 dim
-
text encoder
-
-
FLAN-T5-large
max text tokens
-
-
128
Appendix
Table 7: Latent Forcing ( Baade et al., 2026 ) decoder training at 256 × 256 resolution with 16 × 16 image patches. JiT-B used for scaling laws, JiT-L used for baseline comparisons.
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×32 to 16×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×1022 encoder-training FLOPs.
Iteration 1: Pixel
Iteration 2: Latent
Encoder (width)
ImageNet Epochs-Worth
FD
ImageNet Epochs-Worth
FD
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×1021 encoder-training FLOPs. MiniT2I obtains an FD-10K of 546.3.
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.
Zhang Xiaofeng, Samuel Lavoie, Nicolas Beltran-Velez +6
Mila - Québec AI Institute · Université de Montréal · FAIR at Meta +4
Despite progress in image tokenization, standard methods encode redundant information by mixing all granularities within each token, thus redundancy persists between tokens. The mix of information of different granularity also complicates the training of generators. This paper introduces SelfBootTok, a method that resolves this by cleanly decomposing information into global and local token groups. Through self-bootstrapped learning, the model predicts local details exclusively from global tokens, shifting the burden of visual details from the generator to the tokenizer. Consequently, our generator is far more efficient, requiring only global tokens and reducing computation by approximately 40%, while delivering superior reconstruction and generation. Moreover, this paradigm scales elegantly: by leveraging more data or parameters to self-supervise local representation learning, SelfBootTok achieves a new state-of-the-art gFID score of 1.56 using only 64 tokens.
Haozhe Chi, Jinghan Li, Hao Jiang +4
Peking University · Central Media Technology Institute, Huawei
Generative and representation learning remain asymmetrically connected: semantic representations are used to improve diffusion generation, whereas the models' own representations are often treated as a by-product of synthesis. We ask whether diffusion models can instead be trained to learn substantially stronger semantic representations without sacrificing generation quality. SelfFlow takes a step in this direction by introducing self-supervised patch alignment into flow matching, but its main gains remain in faster convergence and improved generation. Inspired by DINO and iBOT, we extend this framework with cross-view class-token alignment to further strengthen semantic representations. Specifically, we form two independently noised, dual-timestep observations of each image and align each student class-token representation with the stop-gradient EMA-teacher target from the other observation. This objective is optimized jointly with the inherited flow-matching and local patch objectives. Notably, although the additional objective acts only on the class token, it strengthens both class-token and patch representations. Compared with a matched two-view baseline, ImageNet linear-probing accuracy improves by 9.4% using the class token and 10.1% using mean-pooled patch tokens, while frozen-backbone VOC2012 segmentation improves by 3.6 mIoU. These representation gains are achieved while maintaining comparable ImageNet generation FID. In text-to-image training, the same objective also improves generation FID, reducing it from 2.52 to 2.37 at matched checkpoints. Our results show that representation need not remain a by-product of generation or merely a tool for improving it: it can be directly optimized as a first-class capability of diffusion pretraining alongside generation.