Technische Samenvatting: Calibratie van Generatieve Modellen aan Distributie-Constraints

Probleemstelling

Generatieve modellen lijden vaak aan miscalibratie, waarbij de statistieken van de steekproefverdeling (bijvoorbeeld klassfrequenties, verhoudingen van secundaire structuren of genderrepresentatie) systematisch afwijken van de gewenste waarden. Dit probleem doet zich voor in verschillende domeinen: beeldmodellen vertonen mode-collapse; taalmodellen versterken maatschappelijke vooroordelen; en eiwit/DNA-modellen genereren structuren met atypische frequenties van substructuren. Deze fouten vloeien vaak voort uit onevenwichtige datasets, suboptimale trainingsdynamiek of post-hoc aanpassingen zoals bemonstering bij lage temperatuur.

De auteurs formuleren calibratie als een geconstrueerd optimalisatieprobleem. Gegeven een basis-generatief model $p_{\theta_{base}}(x)$, een statistiek $h(x)$ en een doelverwachting $h^*$, is het doel een nieuwe verdeling $p_\theta$ te vinden die voldoet aan $E_{p_\theta}[h(x)] = h^*$, terwijl de Kullback-Leibler (KL)-divergentie met het basismodel wordt geminimaliseerd:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
Deze formulering zoekt het model dat het dichtst bij de basisverdeling ligt en toch voldoet aan de calibratieconstraint. Het exact oplossen hiervan is echter onuitvoerbaar, omdat zowel de doelstelling als de constraint verwachtingen bevatten die moeilijk te berekenen zijn voor hoogdimensionale generatieve modellen.

Methodologie

Om de onuitvoerbaarheid aan te pakken, introduceert het artikel twee surrogaatdoelstellingen voor fijnafstemming, genaamd CGM-relax en CGM-reward. Beide methoden vertrouwen op stochastische optimalisatie met gebruik van onbevooroordeelde gradiëntschattingen afgeleid van steekproeven.

1. CGM-relax (Relax Loss)

Deze benadering vervangt de harde constraint door een strafterm. De doelstellingsfunctie is:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Constraint-overtreding ($L_{viol}$): Het gekwadrateerde $\ell_2$-norm van het verschil tussen de empirische statistiek en het doel. De auteurs leiden een onbevooroordeelde schatter voor deze term af met behulp van een formule voor bias-correcte steekproefvariantie.
• KL-divergentie ($L_{KL}$): Gemeten via Monte Carlo-steekproeven.
• Optimalisatie: De methode maakt gebruik van importance sampling om onbevooroordeelde gradiënten te berekenen. Specifiek wordt er gestoken uit het huidige model $p_\theta$ (met stop-gradient) en worden steekproeven herweging toegepast om gradiënten van de KL-term en de constraint-overtreding te schatten. Een hyperparameter $\lambda$ maakt een afweging tussen het voldoen aan de constraint en dicht bij het basismodel blijven.

2. CGM-reward (Reward Loss)

Deze benadering koppelt het calibratieprobleem aan het Maximum Entropy-principe. De oplossing voor het maximum entropy-probleem (minimaliseren van KL-divergentie met het basismodel onderhevig aan momenten-constraints) is een exponentiële kanteling van het basismodel: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Schatting van $\alpha^*$: Omdat de ware maximum entropy-verdeling onuitvoerbaar is, schat de methode eerst de optimale reward-parameters $\alpha^*$ door een duaal optimalisatieprobleem op te lossen met $N$ steekproeven uit het basismodel.
• Verliesfunctie: Het doel voor fijnafstemming is de KL-divergentie tussen het huidige model $p_\theta$ en de geschatte maximum entropy-verdeling $p_{\hat{\alpha}_N}$:
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  Dit komt overeen met een reward-fijnafstemmingsprobleem waarbij de "reward" wordt afgeleid uit de constraints.
• Optimalisatie: Net als bij CGM-relax worden hier score-function gradiëntschatters gebruikt met leave-one-out baselines om de variantie te verlagen.

Gradiëntschatting

Beide methoden maken gebruik van score-function gradiëntschatters (REINFORCE-stijl) in combinatie met importance sampling. Om de afhankelijkheid van de verwachting van $\theta$ te hanteren, gebruiken de auteurs een "stop-gradient"-truc waarbij steekproeven worden getrokken uit een vaste versie van het huidige model, en gradiënten worden berekend via de dichtheidsverhouding (score functie). Variantievermindering wordt bereikt met behulp van leave-one-out baselines.

Belangrijkste Bijdragen

1. Formulering: Het formuleren van calibratie van generatieve modellen als een geconstrueerd KL-minimalisatieprobleem.
2. Algoritmen: Voorstellen van twee praktische fijnafstemmingsalgoritmen (CGM-relax en CGM-reward) die de oplossing van het geconstrueerde probleem benaderen met stochastische optimalisatie.
3. Schaalbaarheid: Aantonen dat deze methoden honderden gelijktijdige constraints en modellen met tot 9 miljard parameters kunnen verwerken.
4. Theoretische Connectie: Het leggen van de link tussen calibratie, maximum entropy en reward-fijnafstemming, waarbij wordt aangetoond dat CGM-reward de maximum entropy-oplossing targett, terwijl CGM-relax een geregulariseerde versie targett.

Resultaten

De auteurs evalueren de methoden in drie diverse toepassingen:

1. Eiwitontwerp (Genie2 en ESM3):

   • Doel: Calibreren van het gehalte aan secundaire structuren (alpha-helices en beta-strands) om te matchen met statistieken van natuurlijke eiwitten (CATH-dataset).
   • Constraints: Tot 99 gelijktijdige constraints (bivariate CDF's).
   • Resultaat: CGM-relax verkleinde de afstand tot natuurlijke eiwitverdelingen met meer dan 4x voor Genie2 en meer dan 2x voor ESM3. CGM-reward boekte bescheidener verbeteringen, deels vanwege de moeilijkheid om $\alpha^*$ te schatten met hoogdimensionale constraints. Beide methoden behielden biofysische plausibiliteit, zij het met een lichte toename in "ontwerpfouten".

2. Beeldgeneratie (TarFlow):

   • Doel: Balanceren van klassenproporties binnen de categorie "wildlife" van de AFHQ-dataset (bijvoorbeeld het gelijkstellen van leeuwen, tijgers, vossen en wolven).
   • Constraints: 5 gelijktijdige constraints.
   • Resultaat: CGM-relax verkleinde de totale variatieafstand van klassenproporties van 0,306 naar 0,101. Hoewel de realisme van de steekproeven (FID) licht afnam, produceerde het gekalibreerde model diversere en gebalanceerdere generaties. CGM-reward slaagde er niet in om de constraint-overtreding significant te verminderen.

3. Taalmodelleren (Gemma-2-9B-IT):

   • Doel: Verminderen van genderbias in gegenereerde kinderboeken voor diverse beroepen (bijvoorbeeld het garanderen van een gelijke kans op mannelijke/vrouwelijke advocaten).
   • Constraints: 8 beroeps-voorwaardelijke constraints.
   • Resultaat: Beide methoden verkleinden de genderonevenwichtigheid. CGM-relax verkleinde de onevenwichtigheid met meer dan een factor drie ten opzichte van de baseline, terwijl CGM-reward een statistisch significante maar kleinere verbetering liet zien. De fijnafgestemde modellen behielden de kwaliteit van het verhaal (geen verslechtering in log-kans).

Synthetische Experimenten:

• Zeldzame Gebeurtenissen: CGM slaagde er succesvol in om zeldzame modes in Gaussische mengsels op te wegen, zelfs wanneer de basiswaarschijnlijkheid zo laag was als $10^{-3}$.
• Hoge Dimensies: CGM-relax verwerkte effectief tot $10^3$ constraints. Daarentegen werd CGM-reward onuitvoerbaar voor constraints >30 vanwege de moeilijkheid om het empirische maximum entropy-probleem op te lossen (Vergelijking 7).

Betekenis en Beweringen

Het artikel beweert dat CGM een praktisch kader biedt voor het kalibreren van hoogdimensionale generatieve modellen om te voldoen aan distributie-constraints zonder dat er wijzigingen nodig zijn in de modelarchitectuur of hertraining vanaf nul.

• Robuustheid: CGM-relax blijkt robuuster dan CGM-reward, vooral in settings met hoogdimensionale constraints waar het schatten van de maximum entropy-parameters onuitvoerbaar wordt.
• Afwegingen: De methoden bieden een hanteerbare afweging tussen calibratie-accuraatheid en trouw aan het basismodel (gecontroleerd door $\lambda$).
• Beperkingen: De auteurs erkennen dat er residu-fouten blijven, vooral voor zeer zeldzame gebeurtenissen. Bovendien vertrouwt het kader momenteel op modellen met hanteerbare waarschijnlijkheidsfuncties (diffusiemodellen, normalizing flows, autoregressieve modellen), wat de directe toepassing beperkt tot impliciete modellen zoals GAN's.

Het werk positioneert calibratie als een distinct en vruchtbaar onderzoeksrichting, gescheiden van standaard reward-fijnafstemming, door gebruik te maken van de constraints zelf als het drijvende signaal voor distributie-uitlijning.