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

Resumo Técnico: NormGuard: Restrições de Norma Preservadoras de Recompensa em Aprendizado por Reforço de Fluxo-Correspondência (Flow-Matching)

Definição do Problema
O pós-treinamento por aprendizado por reforço (RL) tornou-se uma abordagem padrão para alinhar modelos generativos baseados em fluxo com as preferências humanas. No entanto, esse processo frequentemente leva à "sobre-otimização de recompensa", onde os modelos exploram imperfeições no modelo de recompensa (proxy), resultando em degradação da qualidade perceptual (ex: excesso de nitidez, desvios de cor, iluminação não natural) que o modelo de recompensa falha em penalizar. As estratégias de mitigação existentes, como o early stopping e a regularização KL, tratam o desvio do pós-treinamento como uma quantidade agregada única. Elas restringem o desvio total da velocidade do modelo em relação a uma referência, mas falham em distinguir entre os componentes do desvio que carregam sinais de recompensa e aqueles que geram artefatos visuais.

Observações Chave e Diagnósticos
Os autores identificam uma assinatura estrutural específica deste desvio: o ajuste fino por RL aumenta consistentemente a norma da velocidade por passo ($\|v_\theta\|$) em 5% a 15% em relação à referência pré-treinada ($v_{ref}$) em diferentes métodos (NFT, AWM, DPO) e modelos (SD3.5, FLUX).

O artigo investiga duas soluções potenciais antes de propor uma nova:

1. Renormalização em Tempo de Inferência: Inspirados pelo Classifier-Free Guidance (CFG), onde o redimensionamento da velocidade na inferência mitiga artefatos, os autores testaram o redimensionamento da velocidade ajustada por RL de volta à norma de referência durante a inferência. Eles descobriram que isso falha para modelos treinados por RL. Diferente do CFG, onde a inflação é uma combinação explícita de cabeças de velocidade na inferência, a inflação induzida por RL é co-adaptada nos pesos do modelo. Redimensionar na inferência preserva a recompensa, mas falha em corrigir os artefatos visuais (excesso de nitidez, iluminação não natural).
2. Análise de Sensibilidade Adjunta: Os autores analisaram se a supressão da magnitude da velocidade durante o treinamento removeria um sinal de recompensa coerente. Usando uma análise de sensibilidade adjunta em 6.400 amostras, descobriram que a mudança de recompensa de primeira ordem sob o escalonamento multiplicativo da velocidade exibe altas razões sinal-ruído (3× a 100×) e heterogeneidade de sinal entre os prompts. A média ao nível de lote (batch) permanece próxima de zero, sugerindo que suprimir a inflação da norma dificilmente removeria sistematicamente os ganhos de recompensa.

Metodologia: NormGuard
Com base na descoberta de que correções em tempo de inferência são ineficazes e que a supressão da norma não carrega um custo de recompensa consistente, os autores propõem o NormGuard, um regularizador de tempo de treinamento.

• Mecanismo: O NormGuard é uma penalidade de dobradiça (hinge penalty) unilateral que é ativada apenas quando a norma da velocidade atual excede a norma do modelo de referência. Ela é adicionada de forma aditiva a qualquer perda de base local à velocidade (como as usadas em NFT, AWM e DPO).
• Formulação: Para um estado amostrado $(x_t, t, c)$, o regularizador é definido 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)$$
  onde $\lambda > 0$ controla a intensidade. A penalidade deixa as atualizações sem restrições enquanto $\|v_\theta\| \leq \|v_{ref}\|$, penalizando apenas o "excesso" de norma.
• Compatibilidade: Como a penalidade atua no mesmo espaço de saída de velocidade da perda base, ela se compõe naturalmente com os objetivos de pós-treinamento existentes sem exigir restrições separadas no espaço de parâmetros.

Resultados Experimentais
Os autores avaliaram o NormGuard em dois modelos base (SD3.5-Medium, FLUX.2-4B), três métodos de pós-treinamento (NFT, AWM, DPO) e dois modelos de recompensa (PickScore, HPSv2).

• Qualidade Perceptual: Usando juízes MLLM (Qwen3.5 e GPT-4.1) para avaliar a qualidade da imagem em eixos como fidelidade de textura e detecção de artefatos, o NormGuard superou consistentemente o baseline não regularizado em todas as sete configurações. As taxas de vitória foram significativamente maiores (ex: 64% vs 35% para SD3.5-M + NFT).
• Realismo Forense: O NormGuard melhorou as pontuações de realismo forense (RealScore) em seis das sete configurações, reduzindo artefatos como iluminação não natural e excesso de nitidez.
• Preservação de Recompensa: O método preservou amplamente os ganhos de recompensa do baseline. As diferenças em PickScore e HPSv2 foram mínimas (variando de -0,004 a +0,011), indicando que controlar a inflação radial não impede as mudanças direcionais responsáveis pelo alinhamento da recompensa.
• Inferência de Poucos Passos: Os benefícios do NormGuard aumentaram conforme o número de passos de inferência diminuía. A lacuna nas taxas de vitória de MLLM ampliou de 9% em 28 passos para 20% em 4 passos, sugerindo que o regularizador estabiliza o processo de integração de ODE contra erros de discretização causados por velocidades infladas.
• Ablações: As melhorias não foram explicadas pelo early stopping (o modelo regularizado superou os checkpoints iniciais do baseline) e foram complementares à regularização KL, abordando modos de falha distintos.

Significância e Alegações
O artigo afirma que o desvio agregado penalizado pela regularização KL padrão pode ser frutiferamente decomposto: a inflação da norma impulsiona artefatos perceptuais enquanto carrega pouco sinal de recompensa, enquanto o realinhamento direcional responde pela maior parte do ganho de recompensa.

O NormGuard oferece um diagnóstico e uma correção direcionados para a sobre-otimização de recompensa em modelos de fluxo-correspondência. Ao suprimir o crescimento excessivo da norma de velocidade durante o treinamento, ele melhora o fotorrealismo e reduz artefatos sem sacrificar o alinhamento de recompensa. Os autores posicionam o método como uma intervenção simples e eficaz que aborda uma assinatura estrutural específica do desvio de RL, distinta das restrições globais da regularização KL ou das correções ineficazes de tempo de inferência do CFG.

Limitações
Os autores observam que a origem dinâmica da inflação radial permanece uma questão aberta. Sua alegação de neutralidade de recompensa baseia-se em sinais de primeira ordem ao nível de lote; perturbações radiais por amostra afetam a recompensa, mas esses efeitos se anulam em média entre os prompts. Além disso, a análise atual aplica-se a objetivos locais à velocidade (NFT, AWM, DPO), e a extensão desta perspectiva para objetivos de nível de trajetória, como o Flow-GRPO, é deixada para trabalhos futuros.