Improving Visual Representation Alignment Generation with GRPO

Abstract: Recent diffusion transformers have demonstrated strong image synthesis capabilities but remain inefficient to train due to weak alignment between generative and discriminative representations. While representation alignment frameworks such as REPA improve convergence by aligning noisy denoising features with pretrained visual encoders, their externally supervised alignment loss is static and lacks adaptivity during training and inference. Existing methods rely on fixed cosine alignment or contrastive objectives, which cannot dynamically balance representation consistency and generation quality, resulting in limited discriminative benefit and failing to optimize alignment in a task-adaptive manner. To address this, we propose VRPO, a reinforcement-based optimization strategy that replaces REPA's static alignment loss with a generative representation policy optimization objective. Instead of enforcing a fixed similarity constraint, VRPO treats representation alignment as a reward-guided process: the model receives adaptive rewards based on generation fidelity, perceptual quality, and semantic coherence between the diffusion features and pretrained visual embeddings. This formulation enables the generator to continuously refine its internal representations toward semantically meaningful directions while improving image quality. Our VRPO-driven training seamlessly integrates into diffusion transformers, introducing negligible computation cost and preserving full compatibility with SiT and DiT architectures. Extensive experiments on ImageNet-256x256 demonstrate that our VRPO-Alignment substantially enhances both convergence and fidelity, achieving up to +1.8 FID improvement and 2.3x faster training compared to REPA under identical compute budgets.

• The paper introduces VRPO, a reinforcement learning approach that reformulates representation alignment for diffusion transformers by using adaptive reward signals.
• It decomposes rewards into fidelity, semantic, and stability components, significantly accelerating convergence and improving image quality on datasets like ImageNet.
• Experimental results confirm VRPO's superiority over static methods, showing enhanced FID, CLIP-scores, and robust transferability across diverse datasets.

VRPO: Reinforcement-Based Representation Alignment for Diffusion Transformers

Problem Formulation and Motivation

Diffusion transformers have achieved state-of-the-art results in high-resolution image synthesis. However, current approaches struggle with inefficient training dynamics resulting from weak alignment between generative features and discriminative representations. Existing methods such as REPA enforce alignment via static cosine or contrastive losses, leading to stationary supervision that cannot dynamically adjust as model behavior evolves. This causes over-constrained updates, limited semantic guidance, and slow convergence—particularly on large-scale datasets like ImageNet. The paper introduces Visual Representation Policy Optimization (VRPO), which redefines representation alignment as a reinforcement learning (RL) problem: the diffusion model acts as a stochastic policy, receiving rewards adaptively based on fidelity, semantic coherence, and stability between its features and those of a frozen vision encoder (e.g., DINOv2).

Figure 1: Overview of the VRPO framework, showing reward-guided optimization along semantic, fidelity, and stability dimensions within the diffusion denoising process.

VRPO Framework and Objective Design

VRPO decomposes the total reward into three axes:

• Fidelity Reward ($R_\text{fidelity}$): Quantifies perceptual and statistical realism of generated samples via proxies for FID and LPIPS.
• Semantic Reward ($R_\text{semantic}$): Measures cosine similarity between internal diffusion features and encoder embeddings (DINOv2).
• Stability Reward ($R_\text{stability}$): Penalizes abrupt latent feature transitions across denoising timesteps.

The composite reward $R$ (with learnable weights $\alpha, \beta, \gamma$) guides the policy gradient update, replacing the fixed alignment loss of REPA. Adaptive reward normalization and dynamic scaling ensure that VRPO allocates learning capacity to semantic alignment early in training, and shifts focus to visual fidelity and stability as representations converge. The gradient is injected selectively into early-to-mid transformer layers, preserving the specialization of deeper blocks for texture refinement.

Theoretical Guarantees

The paper provides formal proofs for:

• Convergence of variance-adaptive policy gradient updates under Robbins-Monro conditions and bounded reward variance.
• Pareto optimality of inverse-variance reward weighting, ensuring stable multi-objective optimization.
• Prevention of feature-mode collapse by regularizing temporal Lipschitz constants through the stability reward.
• Zero asymptotic bias in the generative score, mitigating the permanent offset introduced by REPA’s static loss.

Integration and Computational Efficiency

VRPO is fully compatible with SiT and DiT backbones, requiring no architectural modifications. The reward computation depends solely on model outputs and frozen encoders, introducing negligible training overhead (~3%), while achieving substantial gains in convergence speed (up to 2.3$\times$ faster).

Experimental Results

Extensive evaluation on ImageNet-256$\times$256 shows that VRPO outperforms DiT, SiT, REPA, and REG in both FID and Representation Alignment (RA), with strong improvements in CLIP-score and training efficiency. The adaptive reward mechanism yields robust semantic disentanglement and consistent global structure in generated samples, especially for complex object classes.

Figure 2: Qualitative comparison of generated samples on complex classes. VRPO maintains global structural coherence and sharp textures, outperforming DiT and REPA.

Ablations demonstrate that semantic reward is the dominant factor for alignment, while fidelity ensures perceptual realism, and stability prevents oscillatory training dynamics. The optimal balance is empirically achieved by favoring semantic weight moderately. Layer-wise analysis confirms that early-layer gradient injection maximizes performance and efficiency.

Figure 3: Visualization of feature activation maps. VRPO achieves precise semantic disentanglement, closely matching encoder attention heatmaps and foreground isolation.

Generalization and Transferability

VRPO-trained diffusion transformers retain $>$92\% alignment when transferred to out-of-domain datasets (CIFAR-10, CelebA-HQ) without fine-tuning, outperforming static alignment and regularized methods in both FID and semantic retention. This robustness underpins VRPO’s applicability as a backbone for synergistic generative and discriminative tasks.

Practical and Theoretical Implications

Practically, VRPO dramatically accelerates the training of high-fidelity diffusion transformers, reducing compute budgets and environmental impact. Theoretically, it bridges the gap between score-based generative modeling and discriminative representation learning, unifying them via reward-guided policy optimization. This allows diffusion models to internalize meaningful semantic priors while maintaining generative flexibility.

Future directions include extending VRPO to multimodal or out-of-distribution domains, and leveraging negative reward signals for alignment with safety or fairness constraints.

Conclusion

VRPO advances the field by recasting diffusion representation alignment as a reinforcement learning paradigm, balancing semantic structure, visual fidelity, and convergence stability through dynamic reward signals. It yields strong improvements in efficiency, fidelity, and transferability, setting a new standard for integrated discriminative-generative optimization in diffusion transformers (2606.00583).