Improving Visual Representation Alignment Generation with GRPO

2026-05-30

AI Summary

This work addresses the inefficiency in training diffusion Transformers caused by weak alignment between generative and discriminative representations, a limitation exacerbated by existing static alignment losses that lack adaptability. To overcome this, the authors propose Generative Representation Policy Optimization (VRPO), which introduces reinforcement learning into representation alignment for the first time. VRPO employs a dynamic reward mechanism to jointly optimize generation fidelity, perceptual quality, and semantic consistency, replacing fixed similarity constraints. Integrating a pretrained visual encoder, a diffusion Transformer architecture (e.g., DiT), and multidimensional reward signals, the method achieves a 1.8-point improvement in FID over REPA on ImageNet-256×256 and accelerates training by 2.3×, substantially enhancing both convergence behavior and generation quality.

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.

Problem

• representation alignment
• diffusion transformers
• static alignment loss
• adaptive optimization
• generative modeling

Innovation

• VRPO
• representation alignment
• diffusion transformers
• reinforcement learning
• generative modeling