Technische Zusammenfassung: Kalibrierung generativer Modelle an Verteilungsbeschränkungen

Problemstellung

Generative Modelle leiden häufig unter Fehlkalibrierung, bei der die Statistiken der Stichprobenverteilung (z. B. Klassenhäufigkeiten, Anteile sekundärer Strukturen oder Geschlechterrepräsentation) systematisch von den gewünschten Werten abweichen. Dieses Problem zeigt sich in verschiedenen Domänen: Bildmodelle weisen einen Mode-Collapse auf; Sprachmodelle verstärken gesellschaftliche Vorurteile; und Protein-/DNA-Modelle erzeugen Strukturen mit untypischen Häufigkeiten von Unterstrukturen. Diese Fehler resultieren oft aus Ungleichgewichten in den Datensätzen, suboptimalen Trainingsdynamiken oder nachträglichen Anpassungen wie dem Sampling bei niedrigen Temperaturen.

Die Autoren fassen Kalibrierung als ein Optimierungsproblem mit Nebenbedingungen auf. Gegeben ein Basis-Generativmodell $p_{\theta_{base}}(x)$, eine Statistik $h(x)$ und ein Ziel-Erwartungswert $h^*$, besteht das Ziel darin, eine neue Verteilung $p_\theta$ zu finden, die $E_{p_\theta}[h(x)] = h^*$ erfüllt und gleichzeitig die Kullback-Leibler-Divergenz (KL) zum Basismodell minimiert:
$$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^*$$
Diese Formulierung sucht das Modell, das der Basisverteilung am nächsten liegt und die Kalibrierungsbedingung erfüllt. Die exakte Lösung dieses Problems ist jedoch undurchführbar, da sowohl die Zielfunktion als auch die Nebenbedingung Erwartungswerte beinhalten, die für hochdimensionale generative Modelle schwer zu berechnen sind.

Methodik

Um die Undurchführbarkeit zu adressieren, führt die Arbeit zwei surrogate Zielfunktionen für das Fine-Tuning ein, die als CGM-relax und CGM-reward bezeichnet werden. Beide Methoden basieren auf stochastischer Optimierung unter Verwendung unverzerrter Gradientenschätzer, die aus Stichproben abgeleitet werden.

1. CGM-relax (Relax-Loss)

Dieser Ansatz ersetzt die harte Nebenbedingung durch einen Strafterm. Die Zielfunktion lautet:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Verletzung der Nebenbedingung ($L_{viol}$): Die quadrierte $\ell_2$-Norm der Differenz zwischen der empirischen Statistik und dem Ziel. Die Autoren leiten einen unverzerrten Schätzer für diesen Term mittels einer bias-korrigierten Stichprobenvarianzformel her.
• KL-Divergenz ($L_{KL}$): Gemessen über Monte-Carlo-Sampling.
• Optimierung: Die Methode verwendet Importance Sampling, um unverzerrte Gradienten zu berechnen. Konkret wird aus dem aktuellen Modell $p_\theta$ (mit Stop-Gradient) gesampelt und die Stichproben neu gewichtet, um die Gradienten des KL-Terms und der Verletzung der Nebenbedingung zu schätzen. Ein Hyperparameter $\lambda$ steuert den Trade-off zwischen der Erfüllung der Nebenbedingung und der Nähe zum Basismodell.

2. CGM-reward (Reward-Loss)

Dieser Ansatz verknüpft das Kalibrierungsproblem mit dem Maximum-Entropy-Prinzip. Die Lösung des Maximum-Entropy-Problems (Minimierung der KL-Divergenz zum Basismodell unter Momentenbedingungen) ist eine exponentielle Neigung des Basismodells: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Schätzung von $\alpha^*$: Da die wahre Maximum-Entropy-Verteilung undurchführbar ist, schätzt die Methode zunächst die optimalen Reward-Parameter $\alpha^*$, indem sie ein duales Optimierungsproblem unter Verwendung von $N$ Stichproben aus dem Basismodell löst.
• Verlustfunktion: Das Fine-Tuning-Ziel ist die KL-Divergenz zwischen dem aktuellen Modell $p_\theta$ und der geschätzten Maximum-Entropy-Verteilung $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$$
  Dies entspricht einem Reward-Fine-Tuning-Problem, bei dem der „Reward" aus den Nebenbedingungen abgeleitet wird.
• Optimierung: Ähnlich wie bei CGM-relax werden Score-Function-Gradientenschätzer (REINFORCE-Style) mit Leave-One-Out-Baselines verwendet, um die Varianz zu reduzieren.

Gradientenschätzung

Beide Methoden verwenden Score-Function-Gradientenschätzer (im Stil von REINFORCE) in Kombination mit Importance Sampling. Um die Abhängigkeit des Erwartungswerts von $\theta$ zu handhaben, verwenden die Autoren einen „Stop-Gradient"-Trick, bei dem Stichproben aus einer fixierten Version des aktuellen Modells gezogen werden und die Gradienten über das Dichteverhältnis (Score-Function) berechnet werden. Eine Varianzreduktion wird durch Leave-One-Out-Baselines erreicht.

Hauptbeiträge

1. Formulierung: Die Fassung der Kalibrierung generativer Modelle als ein Optimierungsproblem mit KL-Minimierung unter Nebenbedingungen.
2. Algorithmen: Die Vorstellung zweier praktischer Fine-Tuning-Algorithmen (CGM-relax und CGM-reward), die die Lösung des Problems mit Nebenbedingungen unter Verwendung stochastischer Optimierung approximieren.
3. Skalierbarkeit: Der Nachweis, dass diese Methoden Hunderte simultaner Nebenbedingungen und Modelle mit bis zu 9 Milliarden Parametern bewältigen können.
4. Theoretische Verbindung: Die Herstellung des Zusammenhangs zwischen Kalibrierung, Maximum Entropy und Reward-Fine-Tuning, wobei gezeigt wird, dass CGM-reward auf die Maximum-Entropy-Lösung abzielt, während CGM-relax auf eine regularisierte Version abzielt.

Ergebnisse

Die Autoren bewerten die Methoden über drei verschiedene Anwendungen hinweg:

1. Proteindesign (Genie2 und ESM3):

   • Ziel: Kalibrierung des Gehalts an sekundären Strukturen (Alpha-Helices und Beta-Stränge), um natürliche Proteinstatistiken (CATH-Datensatz) zu entsprechen.
   • Nebenbedingungen: Bis zu 99 simultane Nebenbedingungen (bivariate CDFs).
   • Ergebnis: CGM-relax reduzierte den Abstand zu natürlichen Proteinverteilungen um mehr als das 4-Fache für Genie2 und um mehr als das 2-Fache für ESM3. CGM-reward erzielte bescheidenere Verbesserungen, teilweise aufgrund der Schwierigkeit, $\alpha^*$ bei hochdimensionalen Nebenbedingungen zu schätzen. Beide Methoden behielten die biophysikalische Plausibilität bei, wenngleich mit einer leichten Zunahme von „Design-Fehlern".

2. Bildgenerierung (TarFlow):

   • Ziel: Ausgleich der Klassenanteile innerhalb der Kategorie „Wildtiere" des AFHQ-Datensatzes (z. B. Gleichstellung von Löwen, Tigern, Füchsen, Wölfen).
   • Nebenbedingungen: 5 simultane Nebenbedingungen.
   • Ergebnis: CGM-relax reduzierte die totale Variationsdistanz der Klassenanteile von 0,306 auf 0,101. Während die Realitätsnähe der Stichproben (FID) leicht abnahm, produzierte das kalibrierte Modell diversere und ausgewogenere Generierungen. CGM-reward konnte die Verletzung der Nebenbedingung nicht signifikant reduzieren.

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

   • Ziel: Reduzierung des Geschlechterbias bei generierten Kindergeschichten für verschiedene Berufe (z. B. Sicherstellung gleicher Wahrscheinlichkeit für männliche/weibliche Anwälte).
   • Nebenbedingungen: 8 berufsbedingte Nebenbedingungen.
   • Ergebnis: Beide Methoden reduzierten das Geschlechterungleichgewicht. CGM-relax reduzierte das Ungleichgewicht im Vergleich zur Basislinie um mehr als das Dreifache, während CGM-reward eine statistisch signifikante, aber geringere Verbesserung zeigte. Die feinabgestimmten Modelle behielten die Qualität der Geschichten bei (keine Verschlechterung der Log-Wahrscheinlichkeit).

Synthetische Experimente:

• Seltene Ereignisse: CGM konnte seltene Modi in Gaußschen Mischungen erfolgreich hochgewichten, selbst wenn die Basiswahrscheinlichkeit so niedrig wie $10^{-3}$ war.
• Hohe Dimensionen: CGM-relax bewältigte effektiv bis zu $10^3$ Nebenbedingungen. Im Gegensatz dazu wurde CGM-reward für Nebenbedingungen $>30$ unpraktikabel aufgrund der Schwierigkeit, das empirische Maximum-Entropy-Problem (Gleichung 7) zu lösen.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass CGM einen praktischen Rahmen für die Kalibrierung hochdimensionaler generativer Modelle an Verteilungsbeschränkungen bietet, ohne Änderungen an der Modellarchitektur oder ein Neutrainieren von Grund auf zu erfordern.

• Robustheit: Es wird gezeigt, dass CGM-relax robuster ist als CGM-reward, insbesondere in Szenarien mit hochdimensionalen Nebenbedingungen, wo die Schätzung der Maximum-Entropy-Parameter undurchführbar wird.
• Trade-offs: Die Methoden bieten einen navigierbaren Trade-off zwischen Kalibrierungsgenauigkeit und Treue zum Basismodell (gesteuert durch $\lambda$).
• Einschränkungen: Die Autoren räumen ein, dass Restfehler verbleiben, insbesondere bei sehr seltenen Ereignissen. Darüber hinaus stützt sich das Framework derzeit auf Modelle mit berechenbaren Likelihoods (Diffusionsmodelle, Normalizing Flows, autoregressive Modelle), was seine direkte Anwendung auf implizite Modelle wie GANs einschränkt.

Die Arbeit positioniert Kalibrierung als eine eigenständige und fruchtbare Forschungsrichtung, die sich vom standardmäßigen Reward-Fine-Tuning unterscheidet, indem sie Nebenbedingungen selbst als treibendes Signal für die Verteilungsanpassung verwendet.