Calibrating Generative Models to Feature Distributions with MMD Finetuning

Technische Samenvatting: Kalibreren van Generatieve Modellen naar Feature-distributies met MMD Finetuning (kCGM)

Probleemstelling
Generatieve modellen produceren vaak monsters die individueel plausibel zijn, maar een distributionele mismatch vertonen met betrekking tot belangrijke kenmerken (features) van belang. Bijvoorbeeld: een model dat is voorgetraind op een brede chemische ruimte kan moleculen genereren die qua feature-distributie aanzienlijk verschillen van een specifieke therapeutische klasse, zoals antibiotica. Het corrigeren hiervan via directe finetuning op de doelset is problematisch: het riskeert overfitting op kleine datasets, vereist volledige doel-monsters (die schaars kunnen zijn) en faalt in het specificeren van welke specifieke features gematcht moeten worden. Bovendien hebben bestaande methoden vaak moeite met discrete features of black-box features (bijv. Morgan fingerprints) die niet-differentieerbaar zijn. Het artikel identificeert een gat in algoritmen die in staat zijn om generatieve modellen te kalibreren om een doel-feature-distributie te matchen met behulp van uitsluitend feature-niveau supervisie, zonder dat differentieerbaarheid van de feature-map of het samplingsproces vereist is.

Methodologie: kCGM
De auteurs introduceren kernel Calibrating Generative Models (kCGM), een finetuning-methode die de Maximum Mean Discrepancy (MMD) minimaliseert tussen de gegenereerde feature-distributie en een doel-feature-distributie, geregulariseerd door de Kullback-Leibler (KL) divergentie ten opzichte van het voorgetrainde model.

• Doelfunctie: De methode optimaliseert een doelstelling $L_\theta$ die bestaat uit twee termen:

  1. MMD Loss: Een unbiased finite-sample estimator van de gekwadrateerde MMD tussen de door het model geïnduceerde feature-distributie $P$ (waarbij $y = h(x)$ en $x \sim p_\theta$) en de empirische doel-distributie $Q$ (afkomstig van doel-features $y^{targ}_{1:N}$).
  2. KL Regularisatie: Een term $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ om ervoor te zorgen dat het gefinetunede model dicht bij de voorgetrainde distributie blijft, wat catastrofale vergetelheid (catastrophic forgetting) of mode collapse voorkomt.

• Kernelselectie: In tegen tegenstelling tot eerdere benaderingen die mogelijk vertrouwen op dot-product kernels (die alleen gemiddelden matchen), maakt kCGM gebruik van kernels die een metriek definiëren op waarschijnlijkheidsdistributies (die alleen nul zijn voor identieke distributies).

  • Voor continue Euclidische features wordt een energy-distance kernel gebruikt.
  • Voor discrete bit-vector features (bijv. Morgan fingerprints) wordt een Tanimoto/Jaccard kernel ingezet.
  • Deze keuzes vermijden het introduceren van extra bandwidth-hyperparameters.

• Gradiëntschatting: Omdat de feature-map $h$ en het samplingsproces mogelijk niet-differentieerbaar zijn, gebruikt kCGM score-function gradiënten (REINFORCE-stijl).

  • Leave-One-Out (LOO) Baseline: Een cruciale technische bijdrage is de afleiding van een nieuwe LOO-baseline voor de MMD-loss. Standaard LOO-baselines middelen coëfficiënten van andere monsters in een batch, maar MMD-coëfficiënten zijn gekoppeld over monsters heen. De auteurs leiden een specifieke "MMD-LOO" estimator af die de gemiddelde coëfficiënt benadert die zou zijn ontstaan als een specifiek monster niet in de mini-batch aanwezig was, wat de variantie vermindert zonder bias te introduceren.
  • Algoritme: Het proces omvat het samplen uit het huidige model (zonder gradiënttracking), het berekenen van feature-vectoren, het berekenen van MMD- en KL-coëfficiënten met LOO-baselines, en het updaten van de modelgewichten via score-function gradiënten.

Belangrijkste Bijdragen

1. Formulering: Het artikel formuleert feature-distributie matching als KL-geregulariseerde MMD finetuning met uitsluitend doel-feature-monsters, waardoor de noodzaak voor volledige doel-monsters in de oorspronkelijke sample-ruimte wordt geëlimineerd.
2. Unbiased Estimator: De auteurs leiden een unbiased score-function gradiënt estimator af voor de MMD-doelstelling, wat het gebruik van niet-differentieerbare feature-maps en niet-reparameteriseerbare samplers (bijv. autoregressieve modellen, discrete diffusie) mogelijk maakt.
3. Generalisatie: De methode generaliseert Calibrating Generative Models (CGM). Waar CGM de doel-feature gemiddelden matcht, matcht kCGM de volledige distributie via MMD. Het artikel merkt op dat het gebruik van een dot-product kernel in kCGM de CGM-relax doelstelling herstelt.

Experimentele Resultaten
De auteurs evalueren kCGM over drie domeinen: generatie van kleine moleculen, eiwitstructuur-generatie en regulatoire DNA-generatie.

• Kleine Moleculen (Antibiotica):

  • Taak: Het kalibreren van een G2PT autoregressief graafmodel om de feature-distributie van 174 bekende antibiotica te matchen.
  • Features: Morgan fingerprints (lokale substructuur), Murcko scaffolds (topologie) en fysisch-chemische descriptoren.
  • Bevindingen: kCGM bereikt een superieure afweging (tradeoff) tussen het matchen van de doel-feature-distributie (gemeten via MMD) en de afstand tot het voorgetrainde model vergeleken met CGM en directe finetuning. Cruciaal is dat terwijl directe finetuning op de kleine antibioticaset leidde tot een hoog aandeel chemisch ongeldige moleculen, kCGM de feature-matching verbeterde terwijl het de chemische validiteit verhoogde. kCGM demonstreerde ook cross-feature transfer, waarbij scaffold- en descriptor-matching werd verbeterd, zelfs wanneer er alleen op fingerprints werd getraind.

• Eiwitstructuur Diversiteit:

  • Taak: Het finetunen van Genie 2 (een eiwitdiffusiemodel) om de secundaire structuur-distributie (alfa-helices versus bèta-stranden) van de CATH-dataset te matchen.
  • Bevindingen: Het voorgetrainde model produceerde voornamelijk alfa-helices. kCGM verminderde de discrepantie met de doel-distributie met ongeveer factor 2 vergeleken met CGM. kCGM vertoonde een vloeiende tradeoff-curve met de KL-gewicht $\lambda$, terwijl de prestaties van CGM minder gevoelig waren voor $\lambda$, wat optimalisatie bemoeilijkte.

• Regulatoire DNA-generatie:

  • Taak: Het finetunen van een discreet diffusiemodel om DNA-sequenties te genereren met realistische celtype-specifieke activiteitsprofielen (voorspeld door AlphaGenome).
  • Bevindingen: kCGM slaagde erin de volledige distributie van activiteitsprofielen voor 47 verschillende cis-regulatorische topics te matchen, waarbij het CGM (dat alleen gemiddelden matcht) overtrof. Het gefinetunede model genereerde sequenties met activiteitsprofielen die lijken op echte enhancers, terwijl het voorgetrainde model onderactieve sequenties genereerde.

Betekenis en Claims
Het artikel claimt dat kCGM een robuust alternatief biedt voor directe finetuning, met name wanneer doeldatasets klein zijn of wanneer slechts specifieke aspecten van een distributie behouden moeten blijven. De primaire betekenis ligt in het vermogen om:

• Generatieve modellen te aligneren met door de gebruiker gespecificeerde feature-distributies met behulp van uitsluitend feature-niveau supervisie.
• Black-box en niet-differentieerbare features (bijv. fingerprints, predictor outputs) te verwerken, die veel voorkomen in wetenschappelijke domeinen maar incompatibel zijn met standaard gradiënt-gebaseerde finetuning.
• De sample kwaliteit (validiteit) te behouden of te verbeteren terwijl de distributionele miskalibratie wordt gecorrigeerd, een veelvoorkomend falen bij directe finetuning.
• Toepasbaar te zijn op diverse architecturen, inclusief autoregressieve modellen, continue-space diffusie en discrete diffusiemodellen.

De auteurs noemen beperkingen, zoals de potentiële ineffectiviteit voor zeer hoog-dimensionale features vanwege de vloek van dimensionaliteit in MMD hypothese-testen, en de vereiste voor tractabele sampling-traject log-waarschijnlijkheden (uitgesloten van flow matching modellen). Ze suggereren dat toekomstig werk dit framework kan uitbreiden naar settings waar enkel distributionele restricties bestaan zonder volledige samples in de oorspronkelijke generatieve ruimte.