Resumen Técnico: Calibración de Modelos Generativos a Restricciones Distribucionales

Enunciado del Problema

Los modelos generativos sufren frecuentemente de mala calibración, donde las estadísticas de la distribución de muestreo (por ejemplo, frecuencias de clases, proporciones de estructura secundaria o representación de género) se desvían sistemáticamente de los valores deseados. Este problema se manifiesta en diversos dominios: los modelos de imágenes exhiben colapso de modos; los modelos de lenguaje refuerzan sesgos sociales; y los modelos de proteínas/ADN generan estructuras con frecuencias de subestructura atípicas. Estos errores a menudo surgen de desequilibrios en los conjuntos de datos, dinámicas de entrenamiento subóptimas o ajustes post-hoc como el muestreo a baja temperatura.

Los autores enmarcan la calibración como un problema de optimización con restricciones. Dado un modelo generativo base $p_{\theta_{base}}(x)$, una estadística $h(x)$ y una expectativa objetivo $h^*$, el objetivo es encontrar una nueva distribución $p_\theta$ que satisfaga $E_{p_\theta}[h(x)] = h^*$ mientras se minimiza la divergencia de Kullback-Leibler (KL) con respecto al modelo base:
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{sujeto a} \quad E_{p_\theta}[h(x)] = h^*$$
Esta formulación busca el modelo más cercano a la distribución base que satisface la restricción de calibración. Sin embargo, resolver esto exactamente es intratable porque tanto el objetivo como la restricción involucran expectativas que son difíciles de calcular para modelos generativos de alta dimensión.

Metodología

Para abordar la intratabilidad, el artículo introduce dos objetivos sustitutos para el ajuste fino, denominados CGM-relax y CGM-reward. Ambos métodos dependen de la optimización estocástica utilizando estimaciones de gradiente imparciales derivadas de muestras.

1. CGM-relax (Pérdida de Relajación)

Este enfoque reemplaza la restricción dura con un término de penalización. La función objetivo es:
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Violación de la Restricción ($L_{viol}$): La norma $\ell_2$ al cuadrado de la diferencia entre la estadística empírica y el objetivo. Los autores derivan un estimador imparcial para este término utilizando una fórmula de varianza de muestra corregida por sesgo.
• Divergencia KL ($L_{KL}$): Medida mediante muestreo Monte Carlo.
• Optimización: El método utiliza muestreo por importancia para calcular gradientes imparciales. Específicamente, muestrea desde el modelo actual $p_\theta$ (con gradiente detenido) y repondera las muestras para estimar los gradientes del término KL y la violación de la restricción. Un hiperparámetro $\lambda$ equilibra entre satisfacer la restricción y mantenerse cerca del modelo base.

2. CGM-reward (Pérdida de Recompensa)

Este enfoque conecta el problema de calibración con el principio de Máxima Entropía. La solución al problema de máxima entropía (minimizar la KL con respecto al modelo base sujeto a restricciones de momentos) es una inclinación exponencial del modelo base: $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Estimación de $\alpha^*$: Dado que la distribución de máxima entropía verdadera es intratable, el método primero estima los parámetros óptimos de recompensa $\alpha^*$ resolviendo un problema de optimización dual utilizando $N$ muestras del modelo base.
• Función de Pérdida: El objetivo de ajuste fino es la divergencia KL entre el modelo actual $p_\theta$ y la distribución de máxima entropía estimada $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$$
  Esto es equivalente a un problema de ajuste fino por recompensa donde la "recompensa" se deriva de las restricciones.
• Optimización: Similar a CGM-relax, utiliza estimadores de gradiente de función de puntuación con líneas base de "dejar uno fuera" para reducir la varianza.

Estimación de Gradientes

Ambos métodos emplean estimadores de gradiente de función de puntuación (estilo REINFORCE) combinados con muestreo por importancia. Para manejar la dependencia de la expectativa en $\theta$, los autores utilizan un truco de "gradiente detenido" donde las muestras se extraen de una versión fija del modelo actual, y los gradientes se calculan mediante la relación de densidad (función de puntuación). La reducción de varianza se logra utilizando líneas base de "dejar uno fuera".

Contribuciones Clave

1. Formulación: Enmarcar la calibración de modelos generativos como un problema de minimización de KL con restricciones.
2. Algoritmos: Proponer dos algoritmos prácticos de ajuste fino (CGM-relax y CGM-reward) que aproximan la solución al problema con restricciones utilizando optimización estocástica.
3. Escalabilidad: Demostrar que estos métodos pueden manejar cientos de restricciones simultáneas y modelos con hasta 9 mil millones de parámetros.
4. Conexión Teórica: Establecer el vínculo entre calibración, máxima entropía y ajuste fino por recompensa, mostrando que CGM-reward apunta a la solución de máxima entropía mientras que CGM-relax apunta a una versión regularizada.

Resultados

Los autores evalúan los métodos en tres aplicaciones diversas:

1. Diseño de Proteínas (Genie2 y ESM3):

   • Objetivo: Calibrar el contenido de estructura secundaria (hélices alfa y láminas beta) para que coincida con las estadísticas naturales de proteínas (conjunto de datos CATH).
   • Restricciones: Hasta 99 restricciones simultáneas (CDFs bivariadas).
   • Resultado: CGM-relax redujo la distancia a las distribuciones naturales de proteínas en más de 4x para Genie2 y más de 2x para ESM3. CGM-reward logró mejoras más modestas, en parte debido a la dificultad de estimar $\alpha^*$ con restricciones de alta dimensión. Ambos métodos mantuvieron la plausibilidad biofísica, aunque con un ligero aumento en "fallos de diseño".

2. Generación de Imágenes (TarFlow):

   • Objetivo: Equilibrar las proporciones de clases dentro de la categoría "vida silvestre" del conjunto de datos AFHQ (por ejemplo, igualar leones, tigres, zorros y lobos).
   • Restricciones: 5 restricciones simultáneas.
   • Resultado: CGM-relax redujo la distancia de variación total de las proporciones de clases de 0.306 a 0.101. Aunque el realismo de la muestra (FID) disminuyó ligeramente, el modelo calibrado produjo generaciones más diversas y equilibradas. CGM-reward no logró reducir significativamente la violación de restricciones.

3. Modelado de Lenguaje (Gemma-2-9B-IT):

   • Objetivo: Reducir el sesgo de género en historias infantiles generadas para diversas profesiones (por ejemplo, asegurar una probabilidad igual de abogados masculinos/femeninos).
   • Restricciones: 8 restricciones condicionales a la profesión.
   • Resultado: Ambos métodos redujeron el desequilibrio de género. CGM-relax redujo el desequilibrio en más de tres veces en relación con la línea base, mientras que CGM-reward mostró una mejora estadísticamente significativa pero menor. Los modelos ajustados mantuvieron la calidad de la historia (sin degradación en la probabilidad logarítmica).

Experimentos Sintéticos:

• Eventos Raros: CGM logró aumentar el peso de modos raros en mezclas gaussianas, incluso cuando la probabilidad base era tan baja como $10^{-3}$.
• Altas Dimensiones: CGM-relax manejó efectivamente hasta $10^3$ restricciones. En contraste, CGM-reward se volvió inviable para restricciones $>30$ debido a la dificultad de resolver el problema empírico de máxima entropía (Ecuación 7).

Significado y Afirmaciones

El artículo afirma que CGM proporciona un marco práctico para calibrar modelos generativos de alta dimensión para cumplir con restricciones distribucionales sin requerir cambios en la arquitectura del modelo o reentrenamiento desde cero.

• Robustez: Se demuestra que CGM-relax es más robusto que CGM-reward, particularmente en configuraciones de restricciones de alta dimensión donde estimar los parámetros de máxima entropía se vuelve intratable.
• Compensaciones: Los métodos ofrecen una compensación navegable entre la precisión de la calibración y la fidelidad al modelo base (controlada por $\lambda$).
• Limitaciones: Los autores reconocen que persisten errores residuales, especialmente para eventos muy raros. Además, el marco actualmente depende de modelos con verosimilitudes tratables (modelos de difusión, flujos normalizadores, modelos autoregresivos), lo que limita su aplicación directa a modelos implícitos como GANs.

El trabajo posiciona la calibración como una dirección de investigación distinta y fructífera, separada del ajuste fino estándar por recompensa, al utilizar las propias restricciones como la señal impulsora para la alineación distribucional.