Proximal Diffusion Neural Sampler

Technische Samenvatting: Proximal Diffusion Neural Sampler (PDNS)

Probleemstelling
Het leren van een op diffusie gebaseerde neurale sampler om steekproeven te trekken uit een niet-genormaliseerde doelforverdeling (bijvoorbeeld een Boltzmann-verdeling $\pi(x) \propto e^{-\beta E(x)}$) is een fundamentele taak in computationele statistiek, Bayesiaanse inferentie en statistische fysica. Dit probleem wordt vaak geformuleerd als een Stochastisch Optimale Besturing (SOC) probleem op padmaten, waarbij een neuraal netwerk een besturingsproces parameteriseert om een eenvoudige referentieverdeling naar de doelverdeling te transporteren.

Echter, het trainen van deze samplers staat voor aanzienlijke uitdagingen wanneer de doelforverdeling multimodaal is met hoge energiebarrières die de modi van elkaar scheiden. In dergelijke scenario's lijden standaard op SOC gebaseerde methoden (inclusief varianten van Cross-Entropy en Weighted Denoising Cross-Entropy) vaak aan mode collapse. Vroeg in het trainingsproces vangt het model slechts enkele modi; de globale verliesfunctie versterkt vervolgens deze gevangen modi op basis van de gegenereerde beperkte trajecten, waardoor de waarschijnlijkheidsmassa zich concentreert op een subset van modi terwijl andere worden genegeerd. Dit leidt tot instabiele optimalisatie en een slechte dekking van de doelforverdeling.

Methodologie
De auteurs stellen de Proximal Diffusion Neural Sampler (PDNS) voor, een unificerend raamwerk dat is ontworpen om mode collapse aan te pakken door het trainingsproces te herformuleren als een reeks lokale subproblemen in plaats van een enkele globale minimalisatie.

1. Unificerend Raamwerk: Het artikel unificeert eerst continue en discrete diffusie-samplers onder één SOC-raamwerk. Het definieert de optimale padmaat $P^*$ als de oplossing voor het minimaliseren van een relatieve entropiedoelstelling, onderworpen aan een referentiepadmaat $P_{ref}$ die voldoet aan een "geheugenloze" voorwaarde (onafhankelijkheid tussen initiële en terminale verdelingen).
2. Proximaal Puntenmethode: Om mode collapse te mitigeren, past PDNS de proximaal puntenmethode toe op de ruimte van padmaten. In plaats van de globale doelstelling direct te optimaliseren, lost het algoritme iteratief een lokaal subprobleem op bij stap $k$:
   $$P_{\theta^*_k} = \arg\min_{P_\theta} \left( -E_{P_\theta}[r(X_T)] + KL(P_\theta \| P_{ref}) + \frac{1}{\eta_k} KL(P_\theta \| P_{\theta_{k-1}}) \right)$$
   Hierbij is $\eta_k$ een proximaal stapgrootte, en fungeert de term $KL(P_\theta \| P_{\theta_{k-1}})$ als een regularisator die de huidige iteratie dicht bij de vorige houdt.
3. Theoretische Eigenschappen: De auteurs bewijzen dat de optimale oplossing voor de $k$-de proximaal stap een geometrische interpolatie is tussen de vorige iteratie en de optimale padmaat. Specifiek is de terminale verdeling van de $k$-de stap, $P_k$, een geometrische interpolatie tussen de referentie-terminale verdeling $\nu$ en de doelverdeling $\pi$. Naarmate $k \to \infty$, convergeert de reeks padmaten naar de doelverdeling $P^*$. Deze gefaseerde procedure zorgt ervoor dat het model de toestandsruimte geleidelijk verkent, waardoor een voortijdige toewijzing aan een enkele modus wordt voorkomen.
4. Praktische Implementatie (Proximal WDCE): Voor efficiëntie instantiëren de auteurs PDNS met behulp van een Proximal Weighted Denoising Cross-Entropy (Proximal WDCE) doelstelling. Deze aanpak voorkomt het opslaan van volledige trajecten door online geaccumuleerde belangsgewichten te gebruiken.
   • Continue Geval: De doelstelling omvat het afstemmen van een denoising score (of brug) met een gewichtsterm die is afgeleid van de proximaal stapgrootte $\eta_k$ en de energiefunctie.
   • Discrete Geval: De doelstelling past gemaskeerde discrete diffusie aan, waarbij een vergelijkbaar wegingsschema wordt gebruikt dat is gebaseerd op de Radon-Nikodym-afgeleide van de padmaten.
5. Ontwerp van de Scheduler: Het artikel bespreekt strategieën voor het selecteren van de stapgrootte $\eta_k$ (of het equivalente $\lambda_k$). Een kleinere $\eta_k$ (sterkere regularisatie) behoudt modusdekking maar vertraagt de convergentie, terwijl een grotere $\eta_k$ risico loopt op mode collapse. De auteurs stellen zowel vooraf gedefinieerde schema's (lineair afnemend $\lambda_k$) als een adaptieve strategie voor die $\eta_k$ aanpast op basis van een trust-region-beperking op de KL-divergentie tussen het huidige model en de volgende doelstelling.

Belangrijkste Bijdragen

1. Unificerend Raamwerk: Het artikel integreert continue en discrete op SOC gebaseerde samplers in één formulering van padmaten.
2. Proximal Diffusion Neural Sampler (PDNS): Het introduceert een nieuw raamwerk dat proximaal punteniteraties in de padruimte benut om het leren te stabiliseren en mode collapse in multimodale settings te mitigeren.
3. Proximal WDCE: Het leidt concrete, computatie-efficiënte instantiaties van PDNS af voor zowel continue als discrete domeinen met behulp van proximaal gewogen denoising cross-entropy verliezen.
4. Scheduler Strategieën: De auteurs schetsen principiële strategieën voor het selecteren van proximaal stapgroottes, inclusief een adaptieve scheduler die convergentiesnelheid en stabiliteit in evenwicht brengt.

Experimentele Resultaten
De auteurs evalueren PDNS op een breed scala aan continue en discrete benchmarks:

• Continue Taken:
  • Synthetische Energiefuncties: PDNS presteert beter dan sterke baselines (inclusief SMC, DDS, PIS en AS) op multimodale benchmarks zoals het 32-well potentieel (MW54), Funnel en hoogdimensionale Gaussische mengsels (GMM40). Het bereikt state-of-the-art resultaten op de Student's t-mengsel (MoS) en deeltjesgebaseerde potentialen (Lennard-Jones systemen LJ-13 en LJ-55).
  • Moleculaire Dynamica: Op het Alanine Dipeptide (AD) systeem bereikt PDNS KL-divergenties die vergelijkbaar zijn met de state-of-the-art ASBS-methode en vangt conformationele structuren in Ramachandran-plots nauwkeurig op.
• Discrete Taken:
  • Statistische Fysica: PDNS presteert aanzienlijk beter dan LEAPS en Metropolis-Hastings (MH) bij het bemonsteren van Ising- en Potts-modellen bij lage temperaturen (nabij of onder kritieke punten), waar modusafscheiding ernstig is. Het bereikt lagere fouten in magnetisatie en 2-punts correlaties.
  • Combinatorische Optimalisatie: Toegepast op het Maximum Cut-probleem op Barabási–Albert en Erdős–Rényi grafen, demonstreert PDNS het vermogen om oplossingskwaliteiten te bereiken die vergelijkbaar zijn met ground-truth oplosmiddelen (Gurobi).

Betekenis en Beweringen
Het artikel beweert dat PDNS een robuuste oplossing biedt voor het mode collapse-probleem dat inherent is aan het trainen van op diffusie gebaseerde neurale samplers voor complexe, multimodale verdelingen. Door de globale optimalisatie te ontleden in een reeks proximaal stappen, bevordert de methode grondige exploratie over modi heen en stabiliseert het optimalisatie in hoogdimensionale ruimten.

De auteurs merken twee hoofdbeperkingen op:

1. Het onderzoek richt zich op één instantiatie van PDNS (proximaal WDCE), waardoor er ruimte blijft voor het verkennen van andere proximaal doelstellingen of priors.
2. Experimenten worden uitgevoerd op gecontroleerde benchmarks in plaats van op grote schaal real-world toepassingen.

Het werk suggereert toekomstige richtingen die het verkennen van SOC-problemen met niet-geheugenloze referentiedynamiek omvatten, het uitbreiden van het raamwerk naar meer algemene toestandsruimten, en het ontwikkelen van andere instantiaties van de proximaal doelstelling.