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

Technische Samenvatting: NormGuard: Beloningsbehoudende Normbeperkingen in Flow-Matching Reinforcement Learning

Probleemstelling
Post-training via reinforcement learning (RL) is een standaardmethode geworden voor het afstemmen van flow-gebaseerde generatieve modellen op menselijke voorkeuren. Dit proces leidt echter regelmatig tot "beloningsoveroptimalisatie" (reward over-optimization), waarbij modellen imperfecties in de beloningsproxy uitbuiten, wat resulteert in degradatie van de perceptuele kwaliteit (bijv. overmatige verscherping, kleurverschuivingen, onnatuurlijke verlichting) die het beloningsmodel niet straft. Bestaande mitigatiestrategieën, zoals vroege stopzetting (early stopping) en KL-regularisatie, behandelen de drift tijdens de post-training als een enkele geaggregeerde grootheid. Ze beperken de totale afwijking van de snelheid van het model ten opzichte van een referentie, maar maken geen onderscheid tussen componenten van de drift die beloningssignalen dragen en componenten die visuele artefacten veroorzaken.

Kernobservaties en Diagnostiek
De auteurs identificeren een specifieke structurele handtekening van deze drift: RL-fijnregeling verhoogt consistent de snelheid van de per-stap norm ($\|v_\theta\|$) met 5% tot 15% ten opzichte van de vooraf getrainde referentie ($v_{ref}$) over verschillende methoden (NFT, AWM, DPO) en modellen (SD3.5, FLUX).

Het artikel onderzoekt twee potentiële oplossingen voordat een nieuwe wordt voorgesteld:

1. Renormalisatie tijdens inferentie: Geïnspireerd door Classifier-Free Guidance (CFG), waarbij het herschalen van de snelheid tijdens inferentie artefacten vermindert, testten de auteurs het herschalen van de RL-fijngeregelde snelheid terug naar de referentienorm tijdens de inferentie. Zij stelden vast dat dit faalt voor RL-getrainde modellen. In tegen tegenstelling tot CFG, waar de inflatie een expliciete combinatie is van snelheidskoppen tijdens inferentie, is de door RL geïnduceerde inflatie gecofteerd in de gewichten van het model. Het herschalen tijdens inferentie behoudt de beloning, maar slaagt er niet in de visuele artefacten (overmatige verscherping, onnatuurlijke verlichting) te herstellen.
2. Adjoint Sensitivity Analysis: De auteurs analyseerden of het onderdrukken van de snelheidsmagnitude tijdens de training een coherent beloningssignaal zou verwijderen. Met behulp van een adjoint sensitivity analyse op 6.400 samples vonden zij dat de eerste-orde beloningsverandering onder multiplicatieve snelheidsvergroting een hoge ruis-signaalverhouding (3× tot 100×) en sign-heterogeniteit vertoont over verschillende prompts. Het gemiddelde op batchniveau blijft dicht bij nul, wat suggereert dat het onderdrukken van norminflatie onwaarschijnlijk systematisch de beloningswinst zal verwijderen.

Methodologie: NormGuard
Op basis van de bevinding dat fixes tijdens de inferentie ineffectief zijn en normonderdrukking geen consistente beloningskosten met zich meebrengt, stellen de auteurs NormGuard voor, een regularisator tijdens de training.

• Mechanisme: NormGuard is een eenzijdige hinge-penalty die alleen wordt geactiveerd wanneer de huidige snelheidsnorm de norm van het referentiemodel overschrijdt. Het wordt additief toegevoegd aan elke snelheid-lokale basisverliesfunctie (zoals gebruikt bij NFT, AWM en DPO).
• Formulering: Voor een gesamplede staat $(x_t, t, c)$ is de regularisator gedefinieerd als:
  $$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)$$
  waarbij $\lambda > 0$ de sterkte regelt. De straf laat updates onbeperkt zolang $\|v_\theta\| \leq \|v_{ref}\|$, en straft alleen de "overschot" aan norm af.
• Compatibiliteit: Omdat de penalty werkt in dezelfde snelheid-outputruimte als de basisverliesfunctie, combineert het op natuurlijke wijze met bestaande post-training doelstellingen zonder dat aparte beperkingen in de parameterruimte nodig zijn.

Experimentele Resultaten
De auteurs evalueerden NormGuard over twee basismodellen (SD3.5-Medium, FLUX.2-4B), drie post-training methoden (NFT, AWM, DPO) en twee beloningsmodellen (PickScore, HPSv2).

• Perceptuele Kwaliteit: Gebruikmakend van MLLM-judges (Qwen3.5 en GPT-4.1) om de beeldkwaliteit te beoordelen op assen zoals textuurgetrouwheid en detectie van artefacten, presteerde NormGuard consequent beter dan de ongeregulariseerde baseline in alle zeven configuraties. De win rates waren aanzienlijk hoger (bijv. 64% vs 35% voor SD3.5-M + NFT).
• Forensisch Realisme: NormGuard verbeterde de forensische realisme-scores (RealScore) in zes van de zeven settings, waardoor artefacten zoals onnatuurlijke verlichting en overmatige verscherping werden verminderd.
• Beloningsbehoud: De methode behield grotendeels de beloningswinsten van de baseline. Verschillen in PickScore en HPSv2 waren minimaal (variërend van -0.004 tot +0.011), wat aangeeft dat het controleren van radiale inflatie de directionele veranderingen die verantwoordelijk zijn voor beloningsafstemming niet hindert.
• Few-Step Inference: De voordelen van NormGuard namen toe naarmate het aantal inferentiestappen afnam. De kloof in MLLM-win rates werd groter van 9% bij 28 stappen naar 20% bij 4 stappen, wat suggereert dat de regularisator het ODE-integratieproces stabiliseert tegen discretiefouten veroorzaakt door opgeblazen snelheden.
• Ablaties: De verbeteringen werden niet verklaard door vroege stopzetting (het geregulariseerde model presteerde beter dan eerdere checkpoints van de baseline) en waren complementair aan KL-regularisatie, waarbij ze verschillende faalmodi aanpakten.

Betekenis en Claims
Het artikel beweert dat de geaggregeerde drift die door standaard KL-regularisatie wordt bestraft, vruchtbaar kan worden gedecomponeerd: norminflatie drijft perceptuele artefacten aan terwijl het weinig beloningssignaal draagt, terwijl directionele heroriëntatie het grootste deel van de beloningswinst verklaart.

NormGuard biedt een gerichte diagnose en oplossing voor beloningsoveroptimalisatie in flow-matching modellen. Door het overschot aan snelheid-normgroei tijdens de training te onderdrukken, verbetert het de fotorealistische kwaliteit en vermindert het artefacten zonder de beloningsafstemming op te offeren. De auteurs positioneren dit als een eenvoudige, effectieve interventie die een specifieke structurele handtekening van RL-drift aanpakt, verschillend van de globale beperkingen van KL-regularisatie of de ineffectieve inferentie-tijd fixes van CFG.

Beperkingen
De auteurs merken op dat de dynamische oorsprong van de radiale inflatie een open vraag blijft. Hun claim van beloningsneutraliteit is gebaseerd op eerste-orde signalen op batchniveau; per-sample radiale perturbaties beïnvloeden de beloning wel degelijk, maar deze effecten middelen uit naar nul over verschillende prompts. Bovendien is de huidige analyse van toepassing op snelheid-lokale doelstellingen (NFT, AWM, DPO), en het uitbreiden van dit perspectief naar traject-niveau doelstellingen zoals Flow-GRPO is voor toekomstig werk gelaten.