Technische Samenvatting: Een Mean-Field Framework voor Distributionele Controle tijdens de Inferentie van Diffusiemodellen

Probleemstelling

Diffusiemodellen worden steeds vaker gebruikt als controleerbare samplers, waarbij generatieprocessen tijdens de inferentie worden gestuurd om aan specifieke doelstellingen te voldoen. Hoewel bestaande methoden effectief omgaan met puntgewijze beloningen (gedefinieerd op individuele monsters, $r: \mathbb{R}^d \to \mathbb{R}$), vereisen veel kritieke toepassingen distributionele beloningen (gedefinieerd op de gegenereerde maat, $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$). Deze distributionele doelstellingen omvatten het kalibreren van populaties aan experimentele observaties, het stimuleren van diversiteit of het balanceren van modi.

Huidige benaderingen voor distributionele sturing integreren vaak gradiënttermen van de beloning in de inferentiedynamiek (bijv. repulsieve potentialen of ensemble-likelihoods). Echter, vergelijkbaar met vroege methoden voor puntgewijze sturing, missen deze heuristieken theoretische garanties met betrekking tot de gesamplede distributie. Het simpelweg toevoegen van een gradiëntterm zorgt er namelijk niet voor dat de deeltjes de gewenste getilte maat $\mu^*$ volgen, die de KL-geregulariseerde doelstelling maximaliseert:
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
In tegenstelling tot het puntgewijze geval, waar de getilde doelmaat een expliciete exponentiële tilt is ($\mu^* \propto p_1 e^r$), resulteert het distributionele geval in een impliciete doelmaat waarbij het tilt-potentiaal afhankelijk is van de maat zelf. Er bestaat momenteel geen theoretisch gefundeerd framework om de dynamiek en gewichten te corrigeren om deze impliciete distributie te targeten.

Methodologie

De auteurs stellen een Mean-Field Framework voor om distributionele controle tijdens de inferentie te formaliseren en op te lossen.

1. Theoretische Formulering

Het probleem wordt geformuleerd als het targeten van een getilde maat $\mu^*_t$ op elk tijdstip $t \in [0,1]$, die impliciet wordt gedefinieerd door:
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
waarbij $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ de eerste variatie (functionele afgeleide) is van de tijdsafhankelijke beloning $R_t$. Dit creëert een zelfconsistentierelatie waarbij de doelmaat afhankelijk is van het potentiaal, wat op zijn beurt weer afhankelijk is van de maat.

2. Gewogen McKean-Vlasov Dynamiek

Om dit pad te volgen, leiden de auteurs een Gewogen McKean-Vlasov systeem af. In tegenstelling tot standaard gradiëntsturing induceert de sturingsterm $\nabla_x \Psi_t(X_t; \mu_t)$ interactie tussen deeltjes omdat deze afhankelijk is van de evoluerende empirische maat $\mu_t$.

De geïdealiseerde dynamiek voor een deeltje $X_t$ en zijn log-gewicht $A_t$ is:
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
Hierbij is $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ een cruciale correctieterm. In tegen afhankelijk van het puntgewijze geval waar $\dot{\Psi}_t$ simpelweg $\partial_t r$ is, betreft dit hier de tijdsafhankelijke evolutie van de maat $\mu_t$. De auteurs tonen aan dat $\dot{\Psi}_t$ een impliciete integraalvergelijking bevredigt die de tweede variatie van de beloningsfunctionaal betreft.

3. Eindig-Deeltjes Algoritme

Om dit te implementeren, stellen de auteurs een gewogen interagerend deeltjesschema voor (Algoritme 1):

• Positionele Updates: Deeltjes worden bijgewerkt met behulp van de basisdrift, score en de gradiënt van de beloning met betrekking tot de huidige empirische maat $\hat{\mu}_t$.
• Log-Gewicht Updates: Gewichten worden bijgewerkt om het verschil te corrigeren tussen de gestuurde dynamiek en het doelgetilde pad. Dit vereist het berekenen van $\dot{\Psi}_t$.
• Oplossen voor $\dot{\Psi}_t$: Het artikel biedt twee methoden om de impliciete tijdsafgeleide term te berekenen:
  1. Impliciet Lineair Systeem: Het oplossen van een lineair systeem afgeleid uit de integraalvergelijking (Propositie 3.3).
  2. Fixed-Point Iteratie: Een lichtgewicht Picard-iteratieschema dat de gewichtsupdates benadert.
• Resampling: Standaard Sequential Monte Carlo (SMC) resampling wordt toegepast om gewichtsdegeneratie te voorkomen.

Het framework herstelt de standaard puntgewijze beloningssturing als een speciaal geval waarbij $R(\mu) = \int r(x) d\mu(x)$.

Belangrijkste Bijdragen

1. Formalisering van Distributionele Sturing: De auteurs formaliseren sturing op distributieniveau als het targeten van een getilde maat waarbij de tilt afhankelijk is van de maat zelf, waarmee ze een distributioneel analoog vaststellen voor puntgewijze reward-getilde targets.
2. Afleiding van Gewogen McKean-Vlasov Dynamiek: Ze leiden een principieel dynamisch systeem af met correctieve Feynman-Kac herweging die garandeert dat de gewogen interagerende deeltjes het gewenste getilde pad volgen.
3. Theoretische Herstel en Fundament: Het framework herstelt de standaard puntgewijze sturing als een speciaal geval en biedt een theoretische fundering voor bestaande heuristische batch-niveau sturingsmethoden (bijv. diversiteits-sampling, experimentele guidance).
4. Convergentiegaranties: Het artikel bewijst dat het eindige deeltjessysteem convergeert naar de geïdealiseerde dynamiek naarmate het aantal deeltjes $N \to \infty$.

Experimentele Resultaten

De auteurs valideren het framework in zowel laagdimensionale als hoogdimensionale settings:

• Laagdimensionaal (1D Mengsel):

  • Taak: Een bimodale Gaussische mengeling tilden naar een doelverdeling met andere modus-gewichten met behulp van een kwadratische Maximum Mean Discrepancy (MMD) beloning.
  • Bevindingen: Standaard gradiëntsturing duwt deeltjes richting de doelverdeling maar faalt in het nauwkeurig matchen van de precieze vorm van de getilde distributie $\mu^*$. In contrast hiermee target de mean-field methode $\mu^*$ nauwkeurig, waarbij de optimaliteitskloof $J(\mu) - J(\mu^*)$ wordt geminimaliseerd wanneer de sturingssterkte $\lambda$ overeenkomt met de objectieffunctieparameter. Dit houdt stand bij diverse ruis-schema's.

• Hoogdimensionaal (Eiwitconformatie):

  • Taak 1 (HIV-1 Protease): Het sturen van Boltz-2 gegenereerde ensembles om overeen te komen met experimenteel waargenomen residu-afstandverdelingen (DEER-data). Het basismodel was bevoordeeld richting de "gesloten" staat; het doel was om de "semi-open" staat proporties te herstellen.
    • Resultaat: Mean-field sturing mat de empirische residu-afstandverdeling nauwkeurig, terwijl gradiënt-alleen sturing een distributie produceerde met een licht afwijkende vorm en sterkere pieken.
  • Taak 2 (Adenylaat Kinase): Sturen naar een 2D-verdeling van bindingshoeken afgeleid van Moleculaire Dynamica simulaties.
    • Resultaat: Mean-field sturing produceerde een sterkere verschuiving naar de doelverdeling vergeleken met gradiënt-alleen methoden.
  • Taak 3 (Röntgenkristallografie): Het sturen van Protenix-generaties om te passen bij een elektronen-dichtheidskaart (4OLE structuur) die heterogene conformaties vertoont.
    • Resultaat: Voor dezelfde sturingssterkte bereikte mean-field sturing een hogere cosinus-alignment met de doel-dichtheid en plaatste het meer monsters in de alternatieve conformatie (Conformatie B) vergeleken met gradiënt-alleen sturing, wat suggereert dat gradiënt-alleen methoden ondersturen ten opzichte van het ware inverse-probleem target.

Betekenis en Claims

Het artikel beweert de kloof tussen puntgewijze en distributionele inferentie-tijd controle te overbruggen door een principieel theoretisch fundament te bieden voor de laatste.

• Correctheid: De primaire claim is dat de voorgestelde procedure exact de voorgeschreven reward-getilde distributie $\mu^*$ target in de limiet van oneindig veel deeltjes, waarbij de discrepanties inherent aan heuristische gradiëntsturing worden gecorrigeerd.
• Interpreteerbaarheid: Het framework verheldert dat bestaande batch-niveau sturingsmethoden effectief gradiëntstappen uitvoeren maar de correctieve herweging missen die nodig is om uit de ware posterior/getilde maat te samplen.
• Praktische Bruikbaarheid: Hoewel de methode extra computationele kosten introduceert (het oplossen van $\dot{\Psi}_t$), merken de auteurs op dat dit vaak verwaarloosbaar is vergeleken met netwerk-evaluaties in hoogdimensionale settings (bijv. eiwitten) en dat de fixed-point solver efficiënt is.
• Beperkingen: De auteurs erkennen dat de afhankelijkheid van resampling een beperking kan zijn voor modaliteiten zoals afbeeldingen (mogelijk leidend tot duplicaten van monsters) en dat de methode differentieerbare beloningsfuncties vereist. Ze suggereren dat het combineren met variantieverlagende technieken (bijv. Poisson vergelijking solvers) de resampling-problemen zou kunnen verzachten.

Dit werk positioneert zich niet als een vervanging voor alle bestaande methoden, maar als een rigoureuze onderbouwing die verklaart waarom bepaalde heuristieken werken en biedt een mechanisme om te garanderen dat ze de juiste distributie targeten.