Riepilogo Tecnico: AEGPO – Ottimizzazione della Politica Guidata dall'Entropia Adattiva per Modelli di Diffusione

1. Definizione del Problema

Il Reinforcement Learning from Human Feedback (RLHF), in particolare il Group Relative Policy Optimization (GRPO), è diventato un paradigma standard per l'allineamento di modelli di diffusione e flow-matching con le preferenze umane. Tuttavia, le varianti esistenti di GRPO soffrono di due inefficienze fondamentali dovute alla loro dipendenza da euristiche statiche:

1. Allocazione inefficiente dei campioni: I metodi standard assegnano un numero uniforme di rollout a ogni prompt, ignorando la significativa variazione del "valore di apprendimento" tra i diversi campioni. Alcuni prompt inducono cambiamenti minimi nella politica, mentre altri guidano miglioramenti sostanziali.
2. Esplorazione statica dei timestep: Le strategie di esplorazione (ad esempio, il branching nei rollout basati su alberi) seguono tipicamente programmi fissi e predefiniti. Queste non riescono ad adattarsi alla natura dinamica del processo di denoising, dove i momenti critici per l'esplorazione variano significativamente tra i prompt e le fasi di addestramento.

Queste limitazioni comportano uno spreco di risorse computazionali su campioni a basso valore e un'esplorazione subottimale nei timestep non critici, rallentando la convergenza e limitando le prestazioni finali di allineamento.

2. Metodologia: AEGPO

Gli autori propongono AEGPO (Adaptive Entropy-Guided Policy Optimization), un framework plug-and-play che sfrutta i segnali di attenzione intrinseci per guidare dinamicamente sia la selezione dei campioni che l'esplorazione. L'intuizione centrale è che l'Entropia dell'Attenzione funge da potente segnale duale proxy.

2.1 Segnali Core

Il metodo definisce due metriche intrinseche derivate dalle mappe di attenzione del modello:

• $\Delta$Entropy (Segnale a livello di campione): La variazione relativa della dispersione dell'attenzione tra la politica attuale e la politica base.
  • Definizione: $\Delta\text{Entropy} = \frac{1}{T} \sum_{t=1}^{T} |\text{Entropy}_{\theta}(t) - \text{Entropy}_{\theta_{base}}(t)|$.
  • Funzione: Agisce come un robusto proxy del valore di apprendimento del campione. Un alto $\Delta$Entropy indica che un prompt costringe il modello a riorganizzare significativamente la propria strategia di attenzione, implicando un alto potenziale di apprendimento.
• Entropy($t$) (Segnale a livello di timestep): La dispersione assoluta dell'attenzione a un determinato step di denoising $t$.
  • Definizione: La media dell'entropia di Shannon dei pesi di attenzione attraverso i token di testo per le feature dell'immagine al passo $t$.
  • Funzione: I picchi in Entropy($t$) identificano i timestap critici per l'esplorazione. Questi picchi corrispondono a momenti di alta incertezza o di ampia integrazione semantica (spesso occorrenti nelle fasi iniziali di struttura e in quelle tardive di raffinamento), dove l'esplorazione produce la massima varianza e diversità di reward.

2.2 Strategia Adattiva a Doppio Livello

AEGPO utilizza questi segnali per riallocare la computazione senza aumentare il budget totale di rollout:

1. Allocazione Adattiva Globale:

   • Il batch di prompt viene suddiviso in tier "ad alto valore" e "a basso valore" in base alla mediana di $\Delta$Entropy.
   • I prompt ad alto valore (alto $\Delta$Entropy) ricevono un budget di rollout maggiore ($r_{high}$), mentre i prompt a basso valore ne ricevono meno ($r_{low}$).
   • Vincolo: $r_{high} + r_{low} = 2r_{avg}$, garantendo che la computazione totale rimanga costante.

2. Esplorazione Adattiva Locale:

   • Invece di programmi di branching fissi, il metodo identifica dinamicamente i top-$K$ timestep con i picchi più alti di Entropy($t$) per ogni rollout.
   • Il branching (esplorazione) viene eseguito esclusivamente in questi identificati timestep ad alta dispersione.
   • Ciò assicura che l'esplorazione sia concentrata sui momenti più informativi della traiettoria di denoising.

3. Contributi Chiave

• Scoperta del Dual-Signal: Il documento stabilisce l'Entropia dell'Attenzione come un proxy intrinseco duale, dove la variazione relativa ($\Delta$Entropy) indica il valore di apprendimento del campione e i picchi assoluti (Entropy($t$)) identificano i timestep critici per l'esplorazione.
• Framework AEGPO: Introduzione di un framework adattivo a doppio livello che migliora la selezione dei campioni e l'esplorazione dei timestep utilizzando meccanismi leggeri guidati dall'entropia.
• Generalizzabilità: Dimostrazione che AEGPO è model-agnostic e framework-agnostic, migliorando con successo diversi metodi di ottimizzazione della politica (DanceGRPO, BranchGRPO, FlowGRPO, DiffusionNFT) e modelli base (FLUX.1-dev, SD3.5-M).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sui dataset HPDv2.1 e Pick-a-Pic utilizzando FLUX.1-dev e SD3.5-M come modelli base.

• Velocità di Convergenza: AEGPO accelera significativamente l'ottimizzazione della politica. Raggiunge una convergenza 2 volte più veloce su DanceGRPO e una convergenza 5 volte più veloce su DiffusionNFT rispetto alle varianti standard, raggiungendo reward finali superiori in meno step.
• Metriche di Performance:
  • Su FLUX.1-dev, BranchGRPO + AEGPO ha raggiunto un nuovo stato dell'arte per il punteggio HPS-v2.1 di 0.374 (rispetto a 0.363) e un Pick Score di 0.232 (rispetto a 0.229).
  • Su SD3.5-M, DiffusionNFT + AEGPO ha migliorato l'HPS-v2.1 a 0.343 e l'ImageReward a 1.57.
• Trade-off Reward-KL: I modelli AEGPO dominano costantemente la frontiera Pareto Reward-KL rispetto ai baseline, ottenendo reward di allineamento significativamente più alti per ogni livello di deviazione della politica (divergenza KL).
• Studi di Ablazione: Sia l'Allocazione Globale (G) che l'Esplorazione Locale (L) forniscono benefici indipendenti, con l'approccio combinato (G+L) che produce le prestazioni più elevate.
• Overhead Computazionale: Il metodo introduce un modesto overhead (un aumento di circa l'11,1% nel tempo per step e un aumento trascurabile della VRAM) dovuto all'estrazione delle mappe di attenzione intermedie, il che è giustificato dalla drastica riduzione degli step di addestramento totali necessari per raggiungere la performance di picco.

5. Significato e Affermazioni

Il documento afferma che AEGPO affronta le inefficienze centrali del campionamento statico nell'allineamento dei modelli di diffusione. Passando da un'allocazione basata su euristiche a un'adattazione basata su segnali intrinseci, il metodo assicura che le risorse computazionali siano concentrate sui campioni più informativi e sui momenti più critici del processo di generazione.

Gli autori sottolineano che AEGPO è un modulo generalizzabile e plug-and-play che non richiede nuove funzioni di perdita o complessi modelli di reward. Invece, sfrutta la dinamica dell'attenzione interna del modello per guidare l'ottimizzazione, offrendo un percorso più efficiente verso un allineamento di alta qualità attraverso diverse architetture di diffusione e framework RLHF. Il lavoro suggerisce che i segnali intrinseci come l'entropia dell'attenzione sono risorse sottoutilizzate per migliorare l'efficienza dell'ottimizzazione della politica nei modelli generativi.