Technische Samenvatting: JAGG (Jacobian-Aggregated Group Gradient)

Probleemstelling

Group Relative Policy Optimization (GRPO) is een leidend reinforcement learning (RL) algoritme geworden voor het afstemmen van generatieve modellen op menselijke voorkeuren, waarbij het succesvol de noodzaak voor een aparte waardenetwerk (value network) voor Large Language Models (LLMs) heeft geëlimineerd. Het uitbreiden van GRPO naar diffusie- en flow matching-modellen (specifiek Diffusion Transformers of DiTs) introduceert echter een ernstige computationele flessenhals.

In diffusie-gebaseerde GRPO wordt het generatieproces behandeld als een multi-step beslissingsprobleem waarbij het beleid (policy) wordt geoptimaliseerd over een sampling-traject van $W$ tijdstappen. Standaard training vereist een volledige forward- en backward-pass door de hoogwaardige DiT-backbone voor elke afzonderlijke tijdstap in het traject om de policy gradient te berekenen. Voor high-resolution text-to-image (T2I) generatie, waarbij trajecten vaak tientallen stappen beslaan, resulteert dit in prohibitief dure trainingskosten. In tegen tegenoverliegende van taalmodellen waar token-stappen lichtgewicht zijn, brengt elke backward-pass door een moderne DiT de volledige kosten van een transformer forward-pass met zich mee, waardoor de per-stap backward-eis de dominante trainingsuitgave wordt.

Hoewel training-vrije inferentie-acceleratiemethoden (bijv. $\Delta$-DiT, ScalingCache) succesvol hebben geëxploiteerd dat de DiT hidden states en velocity-voorspellingen glad en bijna lineair variëren langs het denoising-traject om tussenliggende evaluaties over te slaan, bleef het een open vraag of dezezelfde lineariteit ook gebruikt kon worden om de backward-pass kosten in RL-training te verminderen.

Methodologie: JAGG

De auteurs stellen JAGG (Jacobian-Aggregated Group Gradient) voor, een wiskundig gefundeerde acceleratiemethode die het aantal netwerk backward-passes vermindert van $W$ naar 2 voor een groep van $W$ opeenvolgende denoising-stappen.

Kerninzicht

De methode berust op het inzicht dat de Jacobiaan van het denoising velocity veld $v_\theta(z, t)$ ten opzichte van de modelparameters $\theta$ een benaderde lineariteit vertoont in de inputs $(z, t)$, met name tijdens het high-noise regime van het denoising-proces. Dit maakt het mogelijk om de Jacobiaan van tussenliggende stappen nauwkeurig te benaderen via lineaire interpolatie met behulp van enkel de endpoint Jacobians ($J_0$ en $J_{W-1}$).

Algoritmische Workflow

JAGG verwerkt een groep van $W$ tijdstappen via vier fasen:

1. Forward Passes: Het model voert forward passes uit voor alle $W$ stappen. Echter, gradient tracking is alleen ingeschakeld voor de twee eindpunten (stap 0 en stap $W-1$). De tussenliggende $W-2$ stappen worden uitgevoerd met torch.no_grad() om geheugen en computationele kracht te besparen.
2. Upstream Gradient Computation: Voor elke stap $j$ wordt een gedetacheerde proxy van de modeloutput gemaakt. De per-stap loss $L_j$ (afgeleid van de PPO-stijl objective) wordt berekend, en de upstream gradient $s_j = \partial L_j / \partial \hat{v}_j$ wordt berekend via automatic differentiation. Cruciaal is dat deze stap geen transformer backward pass bevat; de gradient stroomt alleen door de lichtgewicht SDE-stap en de PPO objective.
3. Timestep-Weighted Aggregation: De upstream gradients $s_j$ worden geaggregeerd in twee samengestelde signalen, $s_{base}$ en $s_{corr}$, met behulp van interpolatiegewichten $\alpha_j$ afgeleid van de tijdstap $t_j$:
   $$\alpha_j = \frac{t_0 - t_j}{t_0 - t_{W-1}}$$
   $$s_{base} = \sum_{j=0}^{W-1} (1 - \alpha_j) s_j, \quad s_{corr} = \sum_{j=0}^{W-1} \alpha_j s_j$$
   Deze aggregatie reconstrueert effectief het volledige gradient signaal met behulp van de lineaire interpolatie-aanname.
4. Joint Backward: Een enkele torch.autograd.backward call wordt uitgevoerd op de twee endpoint outputs $[\hat{v}_0, \hat{v}_{W-1}]$ met de geaggregeerde gradients $[s_{base}, s_{corr}]$. Dit triggert exact één FSDP all-gather per module, wat de $W$ afzonderlijke backward passes vervangt die nodig zijn voor standaard training.

Theoretische Garantie

De auteurs leveren een formeel bewijs (Theorem 1) dat als het velocity veld $v_\theta(z, t)$ strikt lineair is in $(z, t)$, de $t$-gewogen Jacobian interpolatie exact is met nul benaderingsfout. Onder deze omstandigheden herstelt de two-endpoint backward pass de exacte group gradient.

Adaptive Routing ($jagg\_frac$)

In het besef dat de lineariteitsaanname degradeert in de low-noise (late) stadia van denoising, introduceren de auteurs een routing-parameter $jagg\_frac$. Deze parameter past JAGG alleen toe op het eerste deel van de denoising-groepen (waar de high-noise lineariteitsaanname standhoudt, empirisch geverifieerd via cosine similarity van gradients) en gebruikt exact per-stap backward passes voor de resterende groepen. Voor de geteste T2I-modellen werd $jagg\_frac$ ingesteld op 0.6.

Belangrijkste Bijdragen

1. Identificatie van de Flessenhals: Het artikel identificeert en karakteriseert formeel de computationele bottleneck in diffusie GRPO training veroorzaakt door per-timestep backward passes over het sampling-traject.
2. JAGG Framework: Het stelt een Jacobian-aggregatie framework voor dat het aantal backward passes vermindert van $W$ naar 2 per groep, gebaseerd op een rigoureuze theoretische analyse van de lineariteit van de velocity Jacobian.
3. Theoretisch Bewijs: De auteurs bewijzen dat onder lineaire velocity condities, de voorgestelde interpolatie exact is, wat een solide theoretische basis biedt voor de benadering.
4. Empirische Validatie: Uitgebreide experimenten op meerdere diffusie backbones (Flux, QwenImage) en RL-algoritmen (DanceGRPO, Flow-GRPO-Fast) demonstreren de algemeenheid en effectiviteit van de methode.

Experimentele Resultaten

De auteurs evalueerden JAGG op T2I benchmarks met behulp van HPSv2.1, PickScore en CLIPScore als metrieken.

• Efficiëntie: JAGG bereikte een consistente wall-clock versnelling van 17.3% tot 18.2% per trainingsstap op standaard DanceGRPO setups (het reduceren van de tijd van ~835s naar ~691s op QwenImage en ~279s naar ~228s op Flux). Wanneer toegepast op Flow-GRPO-Fast, zag de update fase alleen al versnellingen van 29.9% tot 30.0%.
• Kwaliteit: De methode leverde deze versnellingen met verwaarloosbare degradatie in generatiekwaliteit. In veel gevallen volgde JAGG de baseline reward curves nauwgezet en kwam het de baseline bij convergentie gelijk of zelfs iets te boven.
• Ablatie: Een ablatie-variant die backward passes uitvoerde op de endpoints maar zonder Jacobian interpolatie (door de volledige loss aan de endpoints toe te wijzen onafhankelijk van elkaar) presteerde consequent minder goed. Dit bevestigt dat de aggregatie van tussenliggende gradient informatie via Jacobian interpolatie cruciaal is voor stabiele training.
• Gradient Kwaliteit: Empirische analyse toonde aan dat in high-noise groepen (0–2), de cosine similarity tussen de JAGG-geaggregeerde gradient en de exacte gradient hoog was (0.74–0.85), wat de lineariteitsaanname valideert. De similarity daalde in latere groepen, wat de $jagg\_frac$ routing strategie rechtvaardigt.

Betekenis en Claims

Het artikel claimt dat JAGG een belangrijke stap is naar het praktisch maken van grootschalige diffusie GRPO training. Door de bijna-lineariteit van de velocity Jacobians in het high-noise regime te exploiteren, vermindert de methode de computationele last van RL-training zonder de generatiekwaliteit op te offeren.

De auteurs benadrukken dat JAGG niet louter een heuristische shortcut is, maar geworteld is in de wiskundige eigenschappen van het denoising-traject. De methode wordt gepresenteerd als orthogonaal aan bestaande inferentie-acceleratietechnieken, specifiek gericht op de backward-pass bottleneck die uniek is voor diffusie RL. De resultaten suggereren dat JAGG de bredere adoptie van RL-gebaseerde alignment voor visuele generatieve modellen mogelijk maakt door de trainingskostenbarrière aanzienlijk te verlagen.