Technische Samenvatting: NoiseTilt: Noise-Tilted Reverse Kernels voor Diffusion Reward Alignment

Probleemstelling

Inference-time scaling is een krachtig mechanisme geworden voor het afstemmen van vooraf getrainde diffusiemodellen op gebruikersvoorkeuren zonder kostbare hertraining. Wanneer voorkeuren worden uitgedrukt als beloningsfuncties (reward functions), biedt het iteratieve denoising-proces van diffusiemodellen een model-agnostische manier om sampling te sturen naar regio's met een hoge beloning. Bestaande methoden kampen echter met een fundamentele afruil tussen gradiëntsturing (gradient guidance) en ruis-compatibiliteit (noise compatibility):

1. Mean-Shifted Methoden (Gradiëntsturing): Benaderingen zoals DPS [10] injecteren beloningsgradiënten direct in het gemiddelde van de reverse Gaussische kernel. Hoewel dit effectief is bij het sturen van generatie, verschuift het de gesamplede toestanden buiten het "ruis-compatibele regime"—de nauwe regio in de latente ruimte waar het vooraf getrainde model op is getraind. Deze verschuiving zorgt ervoor dat tussenliggende toestanden buiten de distributie (out-of-distribution) raken, wat de kwaliteit van de generatie verslechtert en het systeem vatbaar maakt voor reward hacking.
2. Search-gebaseerde Methoden: Benaderingen zoals SVDD [38] behouden het basis reverse-gemiddelde door $K$ kandidaten te trekken uit de stochastische kernel en de beste te selecteren. Hoewel dit de ruis-compatibiliteit en de kwaliteit van de samples behoudt, laat het directe gradiëntsturing achterwege en vertrouwt het op willekeurige sampling om richtingen te vinden die de beloning verbeteren. Dit is computationeel inefficiënt en vereist een evenredig groot aantal samples om sterke afstemming te bereiken.

Geen enkele voorafgaande methode bereikt simultaan gradiëntsturing en ruis-compatibiliteit.

Methodologie: Noise-Tilted Reverse Kernel (NTRK)

De auteurs stellen Noise-Tilted Reverse Kernels (NTRK) voor, een sampler die deze afruil oplost door het reverse-gemiddelde vast te houden en beloningsinformatie uitsluitend via de ruisterm te routeren.

Kernmechanisme

De standaard reverse-stap wordt gedefinieerd als $\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \bm{\epsilon}_t$, waarbij $\bm{\epsilon}_t \sim \mathcal{N}(\bm{0}, \bm{I})$.
NTRK wijzigt de ruisterm naar:
$$\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \tilde{\bm{\epsilon}}_t$$
waarbij de geleide ruis $\tilde{\bm{\epsilon}}_t$ een mengeling is van een gewitte beloningsrichting (whitened reward direction) en ongebiaste stochastiek:
$$\tilde{\bm{\epsilon}}_t = \sqrt{\rho_t} \, \mathcal{W}(\nabla_{\bm{x}_t} r(\hat{\bm{x}}_{0|t})) + \sqrt{1-\rho_t} \, \bm{\epsilon}_t$$
Hierbij controleert $\rho_t \in [0, 1]$ de sterkte van de sturing, en $\hat{\bm{x}}_{0|t}$ is de Tweedie-schatting van de schone afbeelding.

De Whitening Operator ($\mathcal{W}$)

De centrale uitdaging is dat een ruwe beloningsgradiënt gestructureerd en deterministisch is, wat de statistische eigenschappen van typische Gaussische ruis schendt die vereist zijn door het vooraf getrainde model. Het direct injecteren hiervan zou "atypische ruis" creëren, wat leidt tot artefacten.

Om dit op te lossen, introduceren de auteurs een whitening operator $\mathcal{W}$ die de beloningsgradiënt transformeert naar een "typische" ruisvector voordat deze wordt geïnjecteerd. Deze operator zorgt ervoor dat de perturbatie binnen het ruis-compatibele regime blijft terwijl het directionele signaal van de gradiënt behouden blijft. De operator wordt geconstrueerd via een reeks projecties op hoog-betrouwbare sets afgeleid van standaard Gaussische statistieken:

1. Two-Level Order-Statistic (2OS) Projectie: De inputvector wordt in tegels (tiles) geherstructureerd. Waarden worden gesorteerd binnen elke tegel en vervolgens over de tegels heen. De resulterende dubbel-gesorteerde statistieken worden afgeklemd tot 99,99% betrouwbaarheidsintervallen afgeleid van de Beta-distributie. Dit onderdrukt gelokaliseerde gestructureerde artefacten terwijl de globale histogram wordt behouden.
2. Tile-Wise Statistieken: De operator beperkt verder het gemiddelde en de gecentreerde energie van elke tegel om overeen te komen met de verwachte distributies van standaard Gaussische ruis.
3. Orthogonale Transformatiedomeinen: Dezelfde betrouwbaarheidsinterval-projecties worden toegepast in getransformeerde domeinen (bijv. compacte Fourier- en Hadamard-achtige mixing) om resterende globale structuren te onderdrukken die niet zichtbaar zijn in het ruimtelijke domein.

Cruciaal is dat voor inputs die al typische Gaussische ruis zijn, $\mathcal{W}$ hen nagenoeg onveranderd laat, terwijl gestructureerde inputs substantieel worden gewit (whitened).

Belangrijkste Bijdragen

1. Resolutie van de Afruil: NTRK is de eerste methode die simultaan gradiëntgebaseerde sturing en ruis-compatibiliteit bereikt. Het behoudt het vooraf getrainde reverse-gemiddelde (in tegen tegenstelling tot mean-shifted methoden) terwijl het slechts één sample per stap gebruikt (in tegenstelling tot search-gebaseerde methoden).
2. De Whitening Operator: Het artikel identificeert dat de kwaliteit van de beloningsafstemming direct verbonden is met de kwaliteit van de whitening operator. De voorgestelde 2OS-gebaseerde projectie blijkt superieur aan eenvoudige norm-projectie of harde gelijkheidsbeperkingen (zoals WGNC [24]), omdat het een balans vindt tussen het onderdrukken van structuur en het behoud van natuurlijke ruisvariabiliteit.
3. Efficiëntie: Door gradiënten via de ruisterm te routeren, bereikt NTRK state-of-the-art prestaties met aanzienlijk minder function evaluations (NFE).

Experimentele Resultaten

De auteurs evalueerden NTRK over drie primaire taken met gebruik van FLUX [35] voor afbeeldingen en Wan2.1 [70] voor video:

• Esthetische Afbeeldingsgeneratie: NTRK behaalde de hoogste doelbeloning (Aesthetic Score) van alle baselines. Opvallend genoeg overtrof NTRK met slechts 25 NFEs de beloning van de beste baseline (SVDD) die 500 NFEs gebruikte, wat een 20$\times$ reductie in compute betekent. Het behield ook of verbeterde heldere metrieken zoals ImageReward en HPSv2.
• Tekst-gealigneerde Afbeeldingsgeneratie: Bij complexe compositionele prompts (T2I-CompBench++) presteerde NTRK beter dan baselines in PickScore (doelbeloning) en heldere metrieken zoals VQA Score en HPSv2.
• Preference-Aligned Videogeneratie: Gebruikmakend van VideoReward als doel, genereerde NTRK video's met betere tekstalignering, bewegingskwaliteit en visuele kwaliteit vergeleken met DPS en FreeDoM.
• Aanvullende Toepassingen: De methode werd succesvol toegepast op kwantiteit-bewuste generatie (tellen van objecten) en VLM-gebaseerde beloningsafstemming (horizon en diepte positionering), wat een brede toepasbaarheid aantoont.
• Integratie met Fine-Tuning: NTRK werd getoond orthogonaal te zijn aan fine-tuning (bijv. MixGRPO [37]), waarbij extra winst werd geboekt wanneer het bovenop gefinetunede modellen werd toegepast.

Betekenis en Claims

Het artikel claimt dat NTRK een eenvoudig, breed toepasbaar mechanisme biedt voor beloningsafstemming tijdens de inference-fase dat de vooraf getrainde sampling-dynamiek intact laat. Door de beloningsgradiënt te behandelen als een perturbatie die gewit moet worden in plaats van een verschuiving van het gemiddelde, vermijdt NTRK de out-of-distribution problemen die gradiëntgebaseerde sturing teisteren.

De auteurs benadrukken dat de effectiviteit van de methode sterk afhangt van de whitening operator. Hun ablatie-studies tonen aan dat zwakkere operators er niet in slagen schadelijke structuur te verwijderen, terwijl te rigide operators de genuante ruis verstoren. De voorgestelde 2OS-gebaseerde aanpak navigeert succesvol door deze afruil, waardoor generatie met een hoge beloning mogelijk is zonder de inherente kwaliteit van het vooraf getrainde model op te offeren.

Het werk suggereert dat toekomstige inference-time scaling verder kan gaan dan brute-force search of destructieve mean shifts, door in plaats daarvan de statistische eigenschappen van ruis te benutten om generatie efficiënt te sturen.