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

Resumen Técnico: NormGuard: Restricciones de Norma Preservadoras de Recompensa en el Aprendizaje por Refuerzo de Flujo-Matching

Declaración del Problema
El post-entrenamiento mediante aprendizaje por refuerzo (RL) se ha convertido en un enfoque estándar para alinear modelos generativos basados en flujo con las preferencias humanas. Sin embargo, este proceso conduce frecuentemente a la "sobreoptimización de la recompensa", donde los modelos explotan las imperfecciones en el modelo de recompensa (proxy), resultando en una degradación de la calidad perceptual (por ejemplo, sobre-afilado, cambios de color, iluminación antinatural) que el modelo de recompensa no logra penalizar. Las estrategias de mitigación existentes, como el paro temprano (early stopping) y la regularización KL, tratan la deriva del post-entrenamiento como una cantidad agregada única. Estas restringen la desviación total de la velocidad del modelo respecto a una referencia, pero no logran distinguir entre los componentes de la deriva que portan señales de recompensa y aquellos que impulsan artefactos visuales.

Observaciones Claves y Diagnósticos
Los autores identifican una firma estructural específica de esta deriva: el ajuste fino por RL infla consistentemente la norma de la velocidad por paso ($\|v_\theta\|$) entre un 5% y un 15% en relación con la referencia pre-entrenada ($v_{ref}$) a través de diferentes métodos (NFT, AWM, DPO) y modelos (SD3.5, FLUX).

El artículo investiga dos posibles soluciones antes de proponer una nueva:

1. Renormalización en el Tiempo de Inferencia: Inspirados en el Guía Libre de Clasificador (CFG), donde el reescalado de la velocidad en la inferencia mitiga artefactos, los autores probaron reescalar la velocidad ajustada por RL de vuelta a la norma de referencia durante la inferencia. Descubrieron que esto falla para los modelos entrenados por RL. A diferencia de CFG, donde la inflación es una combinación explícita de cabezales de velocidad en la inferencia, la inflación inducida por RL está co-adaptada en los pesos del modelo. El reescalado en la inferencia preserva la recompensa pero falla en corregir los artefactos visuales (sobre-afilado, iluminación antinatural).
2. Análisis de Sensibilidad Adjoint: Los autores analizaron si suprimir la magnitud de la velocidad durante el entrenamiento eliminaría una señal de recompensa coherente. Utilizando un análisis de sensibilidad adjoint en 6,400 muestras, encontraron que el cambio de recompensa de primer orden bajo el escalado multiplicativo de la velocidad exhibe relaciones señal-ruido muy altas (de 3× a 100×) y heterogeneidad de signo entre los prompts. La media a nivel de lote (batch) permanece cerca de cero, lo que sugiere que suprimir la inflación de la norma es improbable que elimine sistemáticamente las ganancias de recompensa.

Metodología: NormGuard
Basándose en el hallazgo de que las correcciones en el tiempo de inferencia son ineficaces y que la supresión de la norma no conlleva un costo de recompensa consistente, los autores proponen NormGuard, un regularizador en el tiempo de entrenamiento.

• Mecanismo: NormGuard es una penalización de tipo hinge (bisagra) unidireccional que se activa solo cuando la norma de la velocidad actual excede la norma del modelo de referencia. Se añade de forma aditiva a cualquier pérdida base local de la velocidad (como las utilizadas en NFT, AWM y DPO).
• Formulación: Para un estado muestreado $(x_t, t, c)$, el regularizador se define como:
  $$R_{norm}(\theta) = \lambda \cdot \max\left(0, \frac{\|v_\theta(x_t, t, c)\|^2 - \|v_{ref}(x_t, t, c)\|^2}{\|v_{ref}(x_t, t, c)\|^2}\right)$$
  donde $\lambda > 0$ controla la intensidad. La penalización deja las actualizaciones sin restricciones siempre que $\|v_\theta\| \leq \|v_{ref}\|$, penalizando únicamente el exceso de norma.
• Compatibilidad: Debido a que la penalización actúa en el mismo espacio de salida de la velocidad que la pérdida base, se compone naturalmente con los objetivos de post-entrenamiento existentes sin requerir restricciones separadas en el espacio de parámetros.

Resultados Experimentales
Los autores evaluaron NormGuard en dos modelos base (SD3.5-Medium, FLUX.2-4B), tres métodos de post-entrenamiento (NFT, AWM, DPO) y dos modelos de recompensa (PickScore, HPSv2).

• Calidad Perceptual: Utilizando jueces MLLM (Qwen3.5 y GPT-4.1) para evaluar la calidad de la imagen en ejes como fidelidad de textura y detección de artefactos, NormGuard superó consistentemente al baseline no regularizado en todas las siete configuraciones. Las tasas de victoria fueron significativamente más altas (por ejemplo, 64% vs 35% para SD3.5-M + NFT).
• Realismo Forense: NormGuard mejoró las puntuaciones de realismo forense (RealScore) en seis de siete configuraciones, reduciendo artefactos como la iluminación antinatural y el sobre-afilado.
• Preservación de la Recompensa: El método preservó en gran medida las ganancias de recompensa del baseline. Las diferencias en PickScore y HPSv2 fueron mínimas (oscilando entre -0.004 y +0.011), lo que indica que controlar la inflación radial no impide los cambios direccionales responsables de la alineación de la recompensa.
• Inferencia de Pocos Pasos: Los beneficios de NormGuard se amplificaron a medida que disminuía el número de pasos de inferencia. La brecha en las tasas de victoria de MLLM se amplió de un 9% a 28 pasos a un 20% a 4 pasos, lo que sugiere que el regularizador estabiliza el proceso de integración de la ODE contra errores de discretización causados por velocidades infladas.
• Ablaciones: Las mejoras no fueron explicadas por el paro temprano (el modelo regularizado superó a los checkpoints tempranos del baseline) y fueron complementarias a la regularización KL, abordando modos de falla distintos.

Significado y Reivindicaciones
El artículo sostiene que la deriva agregada penalizada por la regularización KL estándar puede descomponerse fructíferamente: la inflación de la norma impulsa los artefactos perceptuales mientras conlleva poca señal de recompensa, mientras que el realineamiento direccional explica la mayor parte de la ganancia de recompensa.

NormGuard ofrece un diagnóstico y una solución dirigida para la sobreoptimización de la recompensa en modelos de flujo-matching. Al suprimir el crecimiento excesivo de la norma de la velocidad durante el entrenamiento, mejora el fotorrealismo y reduce los artefactos sin sacrificar la alineación de la recompensa. Los autores posicionan esto como una intervención simple y efectiva que aborda una firma estructural específica de la deriva por RL, distinta de las restricciones globales de la regularización KL o de las soluciones ineficaces de tiempo de inferencia de CFG.

Limitaciones
Los autores señalan que el origen dinámico de la inflación radial sigue siendo una pregunta abierta. Su afirmación de neutralidad de recompensa se basa en señales de primer orden a nivel de lote; las perturbaciones radiales por muestra sí afectan la recompensa, pero estos efectos se promedian a cero entre los prompts. Además, el análisis actual se aplica a objetivos locales de la velocidad (NFT, AWM, DPO), y la extensión de esta perspectiva a objetivos de trayectoria completa como Flow-GRPO queda para trabajos futuros.