NormGuard: Reward-Preserving Norm Constraints in Flow-Matching Reinforcement Learning

2026-06-26

AI Summary

This work addresses the common degradation in perceptual image quality observed during reinforcement learning (RL) post-training for reward alignment, a phenomenon inadequately captured by existing reward proxies. The study identifies, for the first time, anomalous expansion of the velocity field norm as a key structural signature of this quality deterioration. To mitigate this issue, the authors propose NormGuard—a hinge-based penalty mechanism that activates only when the norm exceeds a predefined threshold, dynamically constraining its growth during training. Evaluated across two base models, three RL algorithms, and two reward functions, NormGuard consistently enhances image quality and forensic-level realism as assessed by multimodal large language models, with particularly pronounced gains in few-step inference settings. These improvements are not attributable to early stopping and are achieved without compromising reward performance.

Abstract

Reinforcement learning (RL) post-training improves the reward alignment of flow-based generators, but often degrades perceptual quality in ways that are not captured by the reward proxy. We identify a simple structural signature of this drift: across three post-training methods (NFT, AWM, DPO), RL fine-tuning inflates the per-step velocity norm $\|v_θ\|$ by $5\%$ to $15\%$ relative to the reference. A form of norm inflation has been studied in classifier-free guidance (CFG), where rescaling the velocity back to a reference norm at inference time can mitigate the resulting artifacts. However, this inference-time correction does not transfer cleanly to RL: rescaling $v_θ$ to match $\|v_{\text{ref}}\|$ at inference time neither improves reward nor fixes the quality degradation, because the inflation is co-adapted into the model weights. Furthermore, an adjoint sensitivity analysis shows that velocity magnitude rescaling carries no coherent first-order reward signal at the batch level, indicating that suppressing norm inflation is unlikely to remove a consistently reward-carrying component. Since inference-time renormalization fails while norm suppression carries no reward cost, training-time intervention is the appropriate strategy. Together, these findings motivate \methodname, a hinge penalty that activates only when $\|v_θ\|$ exceeds $\|v_{\text{ref}}\|$ and composes additively with any velocity-local base loss. Across two base models, three post-training methods, and two reward proxies, \methodname consistently improves MLLM-judged image quality and forensic realism while preserving reward, with gains that amplify under few-step inference and are not explained by early stopping.

Problem

• norm inflation
• flow-matching
• reinforcement learning
• perceptual quality
• reward alignment

Innovation

• flow-matching
• norm constraint
• reinforcement learning post-training
• velocity norm inflation
• hinge penalty