Technische Samenvatting: Proximal Policy Optimization voor Amortized Discrete Sampling

Probleemstelling

De paper behandelt de taak van het samplen uit een kansverdeling $p(x) = R(x)/Z$ over een eindige discrete ruimte $\mathcal{X}$, waarbij $R(x)$ een positieve beloningsfunctie is en $Z$ een onbekende normaliserende constante. Dit probleem komt voor in domeinen zoals statistische fysica, Bayesiaanse statistiek en computationele biologie. De auteurs richten zich op Generative Flow Networks (GFlowNets), die het genereren van objecten $x$ modelleren als een sequentieel proces over een Directed Acyclic Graph (DAG). Het doel is om een stochastische forward policy $P_F$ te trainen die deze DAG navigeert zodanig dat de distributie van terminale toestanden overeenkomt met $p(x)$.

Hoewel de training van GFlowNets gedomineerd is door waarde-gebaseerde objectives (bijv. Trajectory Balance, Detailed Balance), merken de auteurs op dat policy gradient-methoden, in het bijzonder Proximal Policy Optimization (PPO), ondanks hun dominantie in grootschalige Reinforcement Learning (RL) en diffusion model post-training, onderbelicht blijven.

Methodologie

Theoretische Fundering: GFlowNets als Soft RL

Dit werk bouwt voort op de theoretische verbinding die in eerdere literatuur [42, 7] is vastgesteld tussen GFlowNets en entropie-geregulariseerde RL. Door de discount factor $\gamma=1$ en de entropie-coëfficiënt $\alpha=1$ in te stellen, demonstreren de auteurs dat het GFlowNet sampling-probleem equivalent is aan het maximaliseren van de soft value function in een RL-setting.

• Beloningsstructuur: Cruciaal is dat de per-stap beloning wordt gedefinieerd als $r(s_t, s_{t+1}) = \log P_B(s_t | s_{t+1})$ voor niet-terminale transities, waarbij $P_B$ een backward policy is, en $r(s_T) = \log R(s_T)$ bij terminale toestanden.
• Soft Policy Improvement: De optimale policy voor deze soft-RL objective komt overeen met de doelstelling voor de samplingverdeling $R(x)/Z$.

Ent-PPO: Een Nieuwelijk Policy Gradient Algoritme

De paper leidt Ent-PPO (Entropic Proximal Policy Optimization) af, een variant van PPO die is aangepast aan de soft-RL formulering van GFlowNets. De afleiding benadrukt twee kritieke componenten die vaak ontbreken in naïeve toepassingen van PPO op GFlowNets:

1. Correcte Beloningsformulering: De tussenliggende beloningen moeten $\log P_B$ zijn (zoals hierboven gedefinieerd) om de soft-RL optima te laten samenvallen met de sampling target. Zonder dit reduceert de objective tot standaard RL met sparse rewards, wat optimaliseert voor $\arg\max_x R(x)$ (mode collapse) in plaats van de volledige distributie.
2. Analytische KL Trust Region: In tegenstelling tot standaard PPO, dat een heuristische entropie bonus $c_2 H(\pi)$ toevoegt, leidt Ent-PPO een analytische KL-penalty af direct vanuit de soft policy improvement theorem.
   • De surrogate objective wordt afgeleid door de cross-entropy term (uit de soft advantage definitie) en de entropie term te combineren.
   • Dit resulteert in een per-stap loss: $\ell_t(\theta) = \text{PPOClip}(\rho_t, \hat{A}_t) - \text{KL}(\pi_\theta(\cdot|s_t) \parallel \pi_{\text{old}}(\cdot|s_t))$.
   • De KL-term fungeert als een trust region die intrinsiek is aan de soft-RL formulering, waardoor de noodzaak voor een aparte, instelbare entropie-coëfficiënt vervalt.

Variatiereductie

De auteurs evalueren systematisch variatiereductie technieken voor de Vanilla Policy Gradient (VPG) baseline binnen dit framework:

• Estimators: Ze vergelijken de simpelste policy gradient, reward-to-go, reward-to-go met een geleerde value baseline, Generalized Advantage Estimation (GAE), en Subtrajectory Evaluation Balance (SubEB-GAE).
• Bevindingen: GAE met een geleerde value baseline wordt geïdentificeerd als de meest effectieve strategie, die een gunstige bias-variance trade-off biedt en de convergentiesnelheid en stabiliteit aanzienlijk verbetert vergeleken met simpelere estimators.

Belangrijkste Bijdragen

1. Afleiding van Standaard Policy Gradients voor GFlowNets: De auteurs leiden standaard policy gradient objectives af voor discrete samplers onder het entropie-geregulariseerde RL framework, en valideren experimenteel het belang van variatiereductie technieken (baselines en advantage estimation) voor stabiele training.
2. Ent-PPO Algoritme: Ze stellen Ent-PPO voor, de eerste succesvolle toepassing van PPO op GFlowNets. De afleiding identificeert expliciet de noodzakelijke tussenliggende beloningen ($\log P_B$) en de analytische KL trust region.
3. Empirische Validatie: De paper demonstreert dat Ent-PPO gevestigde off-policy GFlowNet objectives (DB, TB, SubTB) en naïeve PPO varianten verslaat op een reeks benchmarks, van synthetische grids tot complexe moleculaire graafgeneratie.

Experimentele Resultaten

De auteurs evalueren hun methoden op verschillende omgevingen:

• Synthetisch/Kleinschalig: Hypergrid, TFBind8 (DNA-sequenties), en String QM9.
• Grootschalig: sEH (fragment-gebaseerde moleculaire grafen) en QM9 (atoom-voor-atoom moleculaire grafen).

Belangrijkste Bevindingen:

• Superioriteit over Baselines: Ent-PPPO bereikt consistent lagere Total Variation (TV) afstand op kleine problemen en een betere Evidence Lower Bound (ELBO) / Evidence Upper Bound (EUBO) op grote moleculaire problemen vergeleken met DB, TB en SubTB.
• Sample Efficiëntie: Ent-PPO profiteert significant van meerdere update epochs ($K > 1$) per batch, wat een verbeterde sample efficiëntie en convergentiesnelheid laat zien, een eigenschap die gebruikelijk is in grootschalige RL maar voorheen onontdekt was in GFlowNets.
• Noodzaak van Correcties:
  • KL Penalty: Het verwijderen van de KL-term (terugkeren naar naïeve PPO) zorgt ervoor dat training instort of instabiel wordt, vooral bij meerdere update epochs. De KL-term is essentieel om te voorkomen dat de policy te ver afwijkt van de behavior policy terwijl de correcte soft-RL objective behouden blijft.
  • Clipping: Het verwijderen van de PPO clipping mechanisme terwijl de KL-term behouden blijft, leidt tot divergentie, wat aangeeft dat de analytische KL alleen onvoldoende is om de fluctuaties van de importance ratio te begrenzen.
• Backward Policy Learning: Hoewel Ent-PPO niet van nature joint forward-backward learning ondersteunt zoals off-policy objectives, tonen de auteurs aan dat het gebruiken van Trajectory Likelihood Maximization (TLM) om de backward policy apart te leren, de convergentie aanzienlijk versnelt.

Betekenis en Claims

De paper claimt dat Ent-PPO een significante vooruitgang vormt in GFlowNet training door de kloof te overbruggen tussen moderne, schaalbare policy gradient methoden en de specifieke vereisten van discrete sampling.

• Theoretische Rigor: Het werk corrigeert eerdere pogingen om PPO op GFlowNets toe te passen door rigoureus de noodzakelijke beloningsstructuur en trust region af te leiden vanuit de eerste principes van entropie-geregulariseerde RL.
• Praktische Impact: De resultaten suggereren dat policy gradient methoden, wanneer ze correct zijn aangepast, de dominante waarde-gebaseerde GFlowNet objectives kunnen verslaan, met name wat betreft convergentiesnelheid en data-efficiëntie.
• Schaalbaarheid: Het vermogen van Ent-PPO om meerdere update epochs te benutten en te schalen naar complexe moleculaire generatietaken (QM9), suggereert een pad naar het toepassen van GFlowNets op grotere, complexere problemen waar traditionele waarde-gebaseerde methoden moeite kunnen hebben met credit assignment of partition function schatting.

De auteurs concluderen dat hoewel hun methode momenteel focust op discrete domeinen, het soft-RL framework een veelbelovende richting biedt voor het aanpassen van policy gradient methoden naar continue sampling domeinen in de toekomst.