Calibrating Generative Models to Feature Distributions with MMD Finetuning

Resumen Técnico: Calibración de Modelos Generativos a Distribuciones de Características mediante el Ajuste Fino con MMD (kCGM)

Planteamiento del Problema
Los modelos generativos suelen producir muestras que son individualmente plausibles, pero que exhiben un desalineamiento distributivo con respecto a características clave de interés. Por ejemplo, un modelo preentrenado en un espacio químico amplio puede generar moléculas que difieren significablemente en la distribución de sus características de un tipo terapéutico específico, como los antibióticos. Corregir esto mediante un ajuste fino (finetuning) directo sobre el conjunto objetivo es problemático: conlleva el riesgo de sobreajustarse a conjuntos de datos pequeños, requiere muestras completas del objetivo (que pueden ser escasas) y no logra especificar qué características concretas deben coincidir. Además, los métodos existentes suelen tener dificultades con características de valores discretos o de caja negra (por ejemplo, huellas dactilares de Morgan) que no son diferenciables. El artículo identifica una brecha en algoritmos capaces de calibrar modelos generativos para que coincidan con una distribución de características objetivo utilizando únicamente supervisión a nivel de características, sin requerir la diferenciabilidad del mapa de características o del procedimiento de muestreo.

Metodología: kCGM
Los autores presentan los Modelos Generativos de Calibración por Kernel (kCGM), un método de ajuste fino que minimiza la Discrepancia de la Media Máxima (MMD) entre la distribución de características generada y una distribución de características objetivo, regularizada por la divergencia de Kullback-Leibler (KL) respecto al modelo preentrenado.

• Función Objetivo: El método optimiza una función objetivo $L_\theta$ compuesta por dos términos:

  1. Pérdida MMD: Un estimador insesgado de muestra finita del cuadrado de la MMD entre la distribución de características inducida por el modelo $P$ (donde $y = h(x)$ y $x \sim p_\theta$) y la distribución empírica objetivo $Q$ (proveniente de las características objetivo $y^{targ}_{1:N}$).
  2. Regularización KL: Un término $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ para asegurar que el modelo ajustado permanezca cerca de la distribución preentrenada, evitando el olvido catastrófico o el colapso de modo.

• Selección de Kernel: A diferencia de enfoques previos que podrían depender de kernels de producto escalar (que solo igualan medias), kCGM utiliza kernels que definen una métrica sobre distribuciones de probabilidad (cero solo para distribuciones idénticas).

  • Para características euclidianas continuas, se utiliza un kernel de distancia de energía.
  • Para características de vectores de bits discretos (por ejemplo, huellas dactilares de Morgan), se emplea un kernel de Tanimoto/Jaccard.
  • Estas elecciones evitan la introducción de hiperparámetros de ancho de banda adicionales.

• Estimación de Gradientes: Dado que el mapa de características $h$ y el proceso de muestreo pueden no ser diferenciables, kCGM emplea gradientes de función de puntuación (estilo REINFORCE).

  • Línea Base Leave-One-Out (LOO): Una contribución técnica crítica es la derivación de una novedosa línea base LOO para la pérdida MMD. Las líneas base LOO estándar promedian los coeficientes de otras muestras en un lote, pero los coeficientes de MMD están acoplados entre las muestras. Los autores derivan un estimador "MMD-LOO" específico que aproxima el coeficiente promedio que habría tenido una muestra específica si no hubiera aparecido en el mini-lote, reduciendo la varianza sin introducir sesgo.
  • Algoritmo: El proceso implica muestrear del modelo actual (sin seguimiento de gradientes), calcular los vectores de características, calcular los coeficientes de MMD y KL con líneas base LOO, y actualizar los pesos del modelo mediante gradientes de función de puntuación.

Contribuciones Clave

1. Formulación: El artículo formula la coincidencia de la distribución de características como un ajuste fino de MMD regularizado por KL utilizando únicamente muestras de características objetivo, eliminando la necesidad de muestras completas del objetivo en el espacio de muestra original.
2. Estimador Insesgado: Los autores derivan un estimador de gradiente de función de puntuación insesgado para el objetivo MMD, permitiendo el uso de mapas de características no diferenciables y muestreadores no reparametrizables (por ejemplo, modelos autorregresivos, difusión discreta).
3. Generalización: kCGM generaliza a los Modelos Generativos de Calibración (CGM). Mientras que CGM iguala las medias de las características objetivo, kCGM iguala la distribución completa mediante MMD. El artículo señala que el uso de un kernel de producto escalar en kCGM recupera el objetivo de relajación de CGM.

Resultados Experimentales
Los autores evalúan kCGM en tres dominios: generación de moléculas pequeñas, generación de estructuras de proteínas y generación de ADN regulador.

• Antibióticos de Moléculas Pequeñas:

  • Tarea: Calibrar un modelo de grafo autorregresivo G2PT para que coincida con la distribución de características de 174 antibióticos conocidos.
  • Características: Huellas dactilares de Morgan (subestructura local), andamios de Murcko (topología) y descriptores fisicoquímicos.
  • Hallazgos: kCGM logra un compromiso superior entre la coincidencia con la distribución de características objetivo (medida por MMD) y la distancia del modelo preentrenado en comparación con CGM y el ajuste fino directo. Crucialmente, mientras que el ajuste fino directo sobre el pequeño conjunto de antibióticos condujo a una alta proporción de moléculas químicamente inválidas, kCGM mejoró la coincidencia de características al mismo tiempo que aumentó la validez química. kCGM también demostró transferencia de características cruzadas, mejorando la coincidencia de andamios y descriptores incluso cuando se entrenó solo con huellas dactilares.

• Diversidad de Estructura de Proteínas:

  • Tarea: Ajustar finamente a Genie 2 (un modelo de difusión de proteínas) para que coincida con la distribución de estructura secundaria (hélices alfa frente a láminas beta) del conjunto de datos CATH.
  • Hallazgos: El modelo preentrenado producía mayoritariamente hélices alfa. kCGM redujo el desacuerdo con la distribución objetivo en aproximadamente 2 veces en comparación con CGM. kCGM mostró una curva de compromiso suave con el peso $\lambda$ de KL, mientras que el rendimiento de CGM era menos sensible a $\lambda$, lo que dificultaba la optimización.

• Generación de ADN Regulador:

  • Tarea: Ajustar finamente un modelo de difusión discreta para generar secuencias de ADN con perfiles de actividad celular específicos y realistas (predichos por AlphaGenome).
  • Hallazgos: kCGM logró coincidir con la distribución completa de los perfiles de actividad para 47 temas de regulación cis diferentes, superando a CGM (que solo igualaba las medias). El modelo ajustado generó secuencias con perfiles de actividad similares a los de potenciadores (enhancers) reales, mientras que el modelo preentrenado generaba secuencias con baja actividad.

Significado y Reivindicaciones
El artículo afirma que kCGM proporciona una alternativa robusta al ajuste fino directo, particularmente cuando los conjuntos de datos objetivo son pequeños o cuando solo se necesita preservar aspectos específicos de una distribución. Su principal significación radica en su capacidad para:

• Alinear modelos generativos a distribuciones de características especificadas por el usuario utilizando únicamente supervisión a nivel de características.
• Manejar características de caja negra y no diferenciables (por ejemplo, huellas dactilares, salidas de predictores) que son comunes en dominios científicos pero incompatibles con el ajuste fino basado en gradientes estándar.
• Mantener o mejorar la calidad de la muestra (validez) mientras se corrige el desalineamiento distributivo, un modo de fallo común del ajuste fino directo.
• Ser aplicable a diversas arquitecturas, incluyendo modelos autorregresivos, difusión en espacio continuo y modelos de difusión discreta.

Los autores señalan limitaciones, como la potencial ineficacia para características de muy alta dimensión debido a la maldición de la dimensionalidad en las pruebas de hipótesis MMD, y el requisito de probabilidades logarítmicas de trayectoria de muestreo tratables (excluyendo modelos de flow matching). Sugieren que trabajos futuros podrían extender este marco a entornos donde solo existen restricciones distributivas sin muestras completas en el espacio generativo original.