Technische Samenvatting: TMPO voor Diverse en Efficiënte Diffusie-Alignering

1. Probleemstelling

Versterkt Leren (RL) heeft aanzienlijk potentieel aangetoond in het aligneren van diffusie- en flow-matching-modellen met downstream-taken. Echter, bestaande methoden vertrouwen voornamelijk op scalar beloningsmaximalisatie, wat de auteurs inherent mode-zoekend identificeren. Deze optimalisatiedoelstelling zorgt ervoor dat modellen waarschijnlijkheidsmassa concentreren op een paar hoog-beloonde trajecten, wat leidt tot twee kritieke faalmodi:

1. Verslechtering van Generatieve Diversiteit: Het model stort in tot een smal scala aan composities en stijlen (visuele mode-collapse), waarbij het de anisotrope verdeling van plausibele output ignoreert die tijdens pre-training is geleerd.
2. Versterking van Onbetrouwbare Beloningen: Door specifieke kenmerken die beloningen proxyeren (bijv. OCR-nauwkeurigheid of specifieke esthetische scores) te overversterken, versterkt het model spurious attributen, vaak ten koste van de algehele generatiekwaliteit.

Huidige mitigatiestrategieën, zoals KL-regularisatie, falen om de fundamentele mode-zoekende aard van de doelstelling te veranderen, terwijl bestaande Generative Flow Network (GFlowNet)-benaderingen vaak vertrouwen op onuitvoerbare staat-stroom-schattingen of lokale beperkingen die fouten accumuleren over lange denoising-horizonten.

2. Methodologie: Traject Matching Policy Optimization (TMPO)

TMPO verschuift fundamenteel het optimalisatieparadigma van scalar beloningsmaximalisatie naar verdeling van beloning op trajectniveau.

2.1 Softmax Traject Balans (Softmax-TB)

In plaats van de verwachte beloning te maximaliseren, streeft TMPO ernaar de trajectprobabiliteiten van het beleid af te stemmen op een door beloningen geïnduceerde doelforverdeling.

• Doelstelling: Voor een groep van $K$ trajecten gegenereerd vanuit dezelfde prompt, stemt de methode de beleidswaarschijnlijkheden $P_\theta(\tau)$ af op een Boltzmann-verdeling $q(\tau) \propto \exp(\beta R(\tau))$.
• Partition-vrije Formulering: Door probabiliteiten en beloningen te normaliseren binnen de groep van $K$ trajecten, wordt de onuitvoerbare globale partitiefunctie geëlimineerd. De methode definieert een Softmax-TB-voordeel:
  $$A_i = \log q_i - \log p_i = \log \frac{\exp(\beta R_i)}{\sum_j \exp(\beta R_j)} - \log \frac{P_\theta(\tau_i)}{\sum_j P_\theta(\tau_j)}$$
• Theoretische Garantie: De auteurs bewijzen dat het minimaliseren van deze doelstelling equivalent is aan het minimaliseren van de forward KL-divergentie ($D_{KL}(q \parallel p)$) binnen de groep. In tegenstelling tot reverse KL (gebruikt in standaard policy gradients), is forward KL mode-dekkend: het legt een oneindige straf op als de beleidswaarschijnlijkheid $p_i$ naar nul nadert voor elk traject $i$ dat een positieve doelgewicht $q_i$ heeft. Dit garandeert theoretisch de behoud van diversiteit zonder aanvullende KL-regularisatie.

2.2 Dynamische Stochastische Boom Sampling

Om de rekenkosten van het sample van meerdere onafhankelijke trajecten voor grootschalige flow-matching-modellen aan te pakken, introduceert TMPO een boomgestructureerde samplingstrategie:

• Prefix Delen: Trajecten delen een deterministische ODE-prefix (gemeenschappelijke denoising-stappen) en vertakken zich op dynamisch geplande stochastische punten.
• Dynamische Vertakking: Vertakkingspunten zijn niet vast; ze volgen een curriculum dat verschuift van vroege hoog-ruisgebieden (grofkorrelige diversiteit) naar latere laag-ruisgebieden (fijnkorrelige variaties) naarmate de training vordert. Vertakkingsposities worden stochastisch verstoord met een Beta-verdeling om overfitting op een vaste boomgeometrie te voorkomen.
• Efficiëntie: Deze aanpak reduceert redundante berekening door de forward pass voor de gedeelde prefix te delen, wat de trainingstijd aanzienlijk verlaagt terwijl effectieve exploratie behouden blijft.

3. Belangrijkste Bijdragen

1. Identificatie van de Oorzaak: Het paper identificeert de mode-zoekende aard van beloningsmaximalisatie als de primaire oorzaak van reward hacking en verslechtering van diversiteit in diffusie-RL, en stelt een verschuiving naar verdelingsmatching voor.
2. Softmax-TB Doelstelling: Introductie van een partition-vrije, op trajectniveau gebaseerde verdelingsmatching-doelstelling die de mode-dekkende eigenschappen van forward KL-divergentie erft, en zo op natuurlijke wijze generatieve diversiteit behoudt.
3. Dynamische Stochastische Boom Sampling: Een schaalbaar samplingmechanisme dat de trainingstijd voor grootschalige modellen vermindert door het delen van denoising-prefixen en het dynamisch aanpassen van vertakkingsplanningen.

4. Experimentele Resultaten

De auteurs evalueerden TMPO op de FLUX.1-dev-backbone over drie aligneringstaken: compositie-afbeeldingsgeneratie (GenEval), visuele tekstweergave (OCR) en alignering op menselijke voorkeur (PickScore).

• Diversiteit: TMPO verbeterde de generatieve diversiteit met gemiddeld 9,1% in vergelijking met state-of-the-art methoden (Flow-GRPO, MixGRPO, TreeGRPO, GARDO). Het was de enige methode die consequent positieve Log Geometric Mean Distance (LGMD)-scores behield, wat wijst op gezonde diversiteit, terwijl baselines vaak negatieve LGMD vertoonden (mode-collapse).
• Prestaties: TMPO behaalde concurrerende of superieure prestaties op taakspecifieke metrieken (GenEval-nauwkeurigheid, OCR-nauwkeurigheid) en scores voor menselijke voorkeur (HPS-v2.1, ImageReward, PickScore).
• Efficiëntie: Door gebruik te maken van prefix-delen en boom-sampling, verlaagde TMPO de iteratietrainingstijd met maximaal 27% in vergelijking met bestaande methoden.
• Trade-offs: Pareto-analyse toonde aan dat TMPO de optimale afweging bereikt tussen beloningsmaximalisatie, behoud van diversiteit en trainingsefficiëntie.

5. Betekenis en Claims

Het paper claimt dat TMPO een fundamentele oplossing biedt voor het dilemma van reward hacking in de alignering van generatieve AI. Door de doelstelling te herformuleren om beloningsgeïnduceerde verdelingen af te stemmen in plaats van scalar beloningen te maximaliseren, dwingt TMPO structureel behoud van diversiteit af zonder te vertrouwen op externe regularisatie of heuristische herweging. De methode toont aan dat het mogelijk is hoge aligneringsscores te bereiken terwijl de rijke, multimodale outputverdeling die tijdens pre-training is geleerd, behouden blijft, en dit allemaal terwijl de rekenkosten die typisch gepaard gaan met multi-traject RL-methoden worden verminderd. De auteurs positioneren TMPO als een schaalbaar framework dat geschikt is voor grootschalige flow-matching-modellen, met potentiële uitbreidingen naar videogeneratie en robotische actiegeneratie.