Generative Representation Learning for Scalable Latent Tokenization

1 INTRODUCTION

• Prior tokenizer scaling work couples encoders with decoders, although the cited studies report that encoders can worsen with scale.The paper frames tokenizer encoder scaling laws as faster convergence of generative models with more tokenizer compute and data, with nonzero compute allocated to tokenization.Source: §1 · p. 1
• Existing representation objectives may suppress information needed for generation, while classic VAE and VQ-VAE tokenizers face gradient bottlenecks and fail to learn semantics.Generative modeling samples all latent variables, unlike probing that may need only a few semantic bits.Source: §1 · pp. 1–2
• RePR uses a generative model to estimate weighted PMI, which its representation learner explicitly factorizes without negative samples.The paper says this approach is robust to feature suppression and directly supervises tokenizer patch features.Source: §1 · p. 2
• The paper evaluates RePR on ImageNet and GPIC and reports generalization across unconditional, guided, and unguided generation settings.The reported FLOP-wise comparison uses Latent Forcing without a decoder or perceptual, adversarial, or external-model losses.Source: §1 · p. 2
• At equal FLOPs, RePR improves JiT-L FID-50K on ImageNet to 2.18 and extrapolates to equivalent tokenizer quality 7× faster than DINOv2 and 8× faster than WebSSL.The comparison covers ImageNet and GPIC, respectively.Source: §1 · p. 2

2 RELATED WORK

• Contrastive objectives can suppress scene variables beyond object class, especially for strongly conditioned views where hard negatives are difficult to find.Prior work also uses hand-crafted augmentations to highlight desired variables.Source: §2 · p. 2
• Predictive representation methods include pixel-target approaches such as Context Encoders and MAE and latent-target approaches such as BYOL, Data2Vec, JEPA, and MAGE.These examples are presented alongside contrastive learning as established representation-learning approaches.Source: §2 · pp. 2–3
• PMI factorization has been linked to Word2Vec and contrastive learning, but recent work reports that large-scale neural representations fail to match the hypothesized structure.This paper directly supervises representations to match PMI rather than relying on emergence at scale.Source: §2 · p. 3
• Pixel diffusion uses transformers such as JiT, while latent diffusion adds a tokenizer stage before generative modeling.RAE and REPA report strong ImageNet FID, but their DINOv2-b latent space is distilled from a 1.1-billion-parameter model trained on over 100 times ImageNet data, leaving encoding FLOP efficiency unclear.Source: §2 · p. 3

3 BACKGROUND

• PMI is the information gain on an outcome x when conditioned on a different outcome y.It captures how much more or less often two observations co-occur than expected under independence.Source: §3 · p. 3
• Contrastive embedding dot products converge to an estimate of PMI up to an unidentifiable scalar offset.The relationship connects embedding similarity to a low-rank, weighted factorization of the PMI matrix.Source: §3 · p. 3
• A generative model that supplies log probabilities can estimate mutual information between two data samples.The paper's method uses generative models to estimate PMI between views and directly supervise their feature embeddings.Source: §3 · p. 3

4 METHOD

• RePR embeds two masked views and directly supervises their embedding dot product to match PMI estimated by a conditional generative model.This replaces negative-sample InfoNCE, which the paper says is challenging in strongly conditional settings and does not scale.Source: §4 · pp. 3–4

4.1 PMI ESTIMATION FROM GENERATIVE MODELS

• For diffusion models, the paper estimates probabilities from the ELBO and obtains PMI by subtracting log-probability estimates; its approximation approaches true PMI as prediction improves.The VLB equals log probability when the denoiser is optimal and endpoint terms vanish.Source: §4.1 · p. 4

4.2 TIMESTEP WEIGHTING

• Timestep weighting biases representation similarity toward information that remains shared across selected noise scales, as stronger corruption progressively destroys fine-grained information.The weighting is motivated by the greater semantic or perceptual importance often carried at lower signal-to-noise ratios.Source: §4.2 · pp. 4–5

4.3 BOOTSTRAPPED PMI ESTIMATION

• Bootstrapping replaces a costly full-view PMI estimate with incremental PMI for successive target tokens, whose increments telescope to the full PMI.This makes the likelihood-ratio target easier to estimate than the whole-input PMI at once.Source: §4.3 · p. 5

4.4 REPRESENTATION OBJECTIVE

• The complete representation objective combines timestep weighting and bootstrapping for input views x and y and an incremental amount of information z.Source: §4.4 · p. 5

5 ARCHITECTURE

• RePR estimates weighted data likelihoods with a prediction model and distills their weighted PMI into patch representations.The prediction model uses an MAR architecture; the representation model outputs vectors for every patch.Source: §5 · p. 5, Figure 1 · p. 6

5.1 PREDICTION MODEL

• MAR combines an MAE encoder and decoder with a diffusion MLP head, sharing encoder and decoder cost while averaging likelihood estimates across noise levels.The lightweight head can be run repeatedly after the shared encoder and decoder computation.Source: §5.1 · p. 6

5.2 REPRESENTING PATCH FEATURES

• Directly supervised patch features and relative-position cross-attention impose spatial structure while withholding absolute position information from the encoder.The paper motivates direct supervision by instability in unregularized intermediate features.Source: §5.2 · p. 6

5.3 MASKING STRATEGY

• Randomized masking ratios across two potentially overlapping views supervise a range of context sizes for reconstructing the telescoping series.Reusing encoder and decoder computations across views and target patches reduces additional work.Source: §5.3 · pp. 6–7

5.4 ITERATIVE FEATURE PREDICTION

• Recursive training can yield additional gains by replacing incremental predicted patches with same-location features from a previous training cycle.The previous representation model embeds the full image, and feature calculation is shared across repeated views.Source: §5.4 · p. 7

5.5 JOINT OPTIMIZATION

• The prediction and representation models are jointly optimized, with the representation model trained once every four iterations to speed training.The paper reports that the representation model converges relatively easily.Source: §5.5 · p. 7

6 EXPERIMENTS

• Experiments evaluate RePR for image generation on ImageNet and the text-to-image dataset GPIC, including feature suppression, encoder scaling, and compute-optimal pixel diffusion.Source: §6 · p. 7

6.1 FEATURE SUPPRESSION

• On watermarked ImageNet, LeJEPA and DINOv3 collapse to representing the shared watermark, while RePR remains robust to feature suppression.The watermark uniquely identifies each image, making identification sufficient for contrastive losses.Source: §6.1 · p. 7, Figure 2 · p. 7

6.2 SCALING LAWS

• Extrapolation predicts DINOv2-level performance at roughly one-seventh the training FLOPs and WebSSL-level performance at roughly one-eighth, following power-law scaling across about two compute orders on ImageNet and GPIC.At fixed generative-model size and training duration, increasing RePR size and FLOPs improves generation; iterative feature prediction improves it further.Source: §6.2 · p. 8, Figure 3 · p. 8

6.3 COMPUTE OPTIMAL GENERATION

• At equal total FLOPs, RePR with Latent Forcing improves unguided FID-10K by 2.61 for conditional and 4.18 for unconditional generation versus pixel JiT-L; guided generation reaches FID-50K 2.18.The comparison accounts for encoder and generative-model training compute.Source: §6.3 · p. 8, Table 3 · p. 9, Figure 4 · p. 9

6.4 ABLATION STUDIES

• Second and third training iterations substantially improve generative performance, while the representation-loss weight is optimal at small values.The ablations use ImageNet Top-1 accuracy as a proxy because generative-model training is costly.Source: §6.4 · p. 9, Table 1 · p. 9

6.5 GAN+LPIPS DECODING

• At small scales, RePR is an empirically competitive RAE tokenizer, outperforming JEPA-H and Mocov3 in comparisons using a GAN+LPIPS decoder.The RAE decoder uses LPIPS and a pretrained DINOv1/8 GAN requiring 4.7 × 10^20 FLOPs, more than any model trained in the study.Source: §6.5 · p. 9, Table 2 · p. 9

• RePR shares contrastive learning’s theoretical guarantees without requiring negative samples.Source: §7 · p. 10
• RePR remains robust to feature suppression in settings where prior representation-learning methods collapse.Source: §7 · p. 10
• The paper establishes diffusability scaling laws and reports RePR-based generative models outperforming direct diffusion at matched training FLOPs.Source: §7 · p. 10

• Generative AI assisted with implementation, particularly coding, while the core RePR algorithm was human-written.Source: AI USE STATEMENT · p. 10
• LLM-generated code was extensively verified and primarily used for data visualization and boilerplate.Source: AI USE STATEMENT · p. 10
• Generative AI also assisted with proof verification, paper-readability feedback, and internet searches.Source: AI USE STATEMENT · p. 10
• The authors report reviewing all AI-assisted work and taking responsibility for the final content.Source: AI USE STATEMENT · p. 10

• The algorithm samples masks and a patch index, then encodes views A, B, AB, and B with the selected patch.Source: §A.1 · p. 16
• A masked image is noised and passed to the model for a weighted masked autoregressive reconstruction loss.Source: §A.1 · p. 16
• PMI is estimated by comparing squared prediction errors from conditional predictions with and without the AB representation, averaged over N integration steps.Source: §A.1 · p. 16

A.2 ADDITIONAL EXPERIMENTS

• Sharing noise across likelihood calculations reduces variance; the authors choose 32 Monte Carlo steps, using 1.3% of training FLOPs for diffusion integration.Source: Figure 5 · p. 16

A.3 IMPLEMENTATION DETAILS

• The 24-layer architecture uses 18 encoder layers and six-layer decoder and representation branches, with the representation branch split between self-attention and cross-attention.Source: §A.3 · p. 17
• RoPE supplies encoder position information, while center-relative queries extend representations beyond image borders using padded query positions.Source: §A.3 · p. 17
• Because PMI is asymmetric and not positive semidefinite, the representation head distinguishes incremental-information inputs to enable perfect SVD reconstruction at optimality.Source: §A.3 · p. 17
• The reported epoch counts normalize processed image views by ImageNet sample count, and the appendix specifies iteration schedules and FLOP accounting for comparisons.Source: §A.3 · p. 17

A.4 PROPERTIES OF PMI

• RePR-B Iter 3 reaches PSNR 21.72, 3.63 dB above DINOv2 with the RAE ViT-XL decoder.Source: Table 4 · p. 17, §A.4 · p. 18

A.5 QUALITATIVE RESULTS

• The qualitative ImageNet generations use the Latent Forcing CFG-Interval and AutoGuidance setup, with FLOP calculations including AutoGuidance.Source: §A.5 · p. 20, Figure 6 · p. 20

A.7 RAW NUMBERS FOR SCALING PLOTS

• ImageNet FID and GPIC FD scaling fits project DINOv2-level and WebSSL-level results at 1/7.77 and 1/7.96 of their respective encoder-training compute.Source: §A.7 · p. 22