Resumen Técnico: Coincidencia de Distribuciones Contrastivas para Monte Carlo Secencial Amortizado en Difusión Discreta

Enunciado del Problema

Los modelos de difusión discreta se han convertido en marcos potentes para generar datos categóricos estructurados, como texto, secuencias de ADN y proteínas. Una aplicación central de estos modelos es la alineación por recompensa, donde el objetivo es muestrear desde una distribución objetivo que sesgue el prior preentrenado hacia valores más altos de una función de recompensa aguas abajo (por ejemplo, preferencia humana o diseñabilidad de proteínas).

Aunque el Monte Carlo Secencial (SMC) retorcido ofrece un método asintóticamente exacto para muestrear desde dichas distribuciones sesgadas por recompensa, su aplicación en espacios de estados discretos enfrenta un cuello de botella fundamental: la estimación de la función de retorcimiento óptima.

• En dominios continuos, la función de retorcimiento puede aproximarse eficientemente utilizando la fórmula de Tweedie o guía basada en gradientes.
• En dominios discretos, la ausencia de una relación análoga obliga a depender de aproximaciones costosas de Monte Carlo (MC) para estimar la función de retorcimiento en cada paso de inferencia.
• Esta dependencia crea una sobrecarga computacional severa, particularmente cuando la propia función de recompensa es costosa de evaluar (por ejemplo, simulaciones de plegamiento de proteínas o puntuación de preferencia de LLM basada en API). Escalar el número de muestras de MC ($M$) mejora la precisión pero aumenta linealmente el tiempo de inferencia, haciendo que el enfoque sea poco práctico para muchas aplicaciones del mundo real.

Metodología

Los autores proponen la Coincidencia de Distribuciones Contrastivas (CDM), un marco diseñado para amortizar el costo de la inferencia SMC aprendiendo una función de retorcimiento parametrizada mediante un objetivo de aprendizaje contrastivo.

1. Objetivo de Coincidencia de Distribuciones Contrastivas

En lugar de aprender la función de retorcimiento mediante regresión directa (minimizando el Error Cuadrático Medio frente a estimaciones de MC), CDM minimiza la divergencia Kullback-Leibler (KL) hacia adelante entre la distribución objetivo $p^*_t$ y una distribución parametrizada $p^\phi_t$ modulada por la función de retorcimiento aprendida $\psi^\phi_t$.

El objetivo se formula como:
$$\mathcal{L}_{\text{CDM}}(\phi) = \mathbb{E}_{t}\left[ \mathcal{D}_{\text{KL}}\left( p^*_t(\mathbf{x}_t) \big\| p^\phi_t(\mathbf{x}_t) \right) \right]$$

El gradiente de esta pérdida exhibe una estructura contrastiva:
$$-\nabla_\phi \mathcal{L} = \mathbb{E}_{t}\left[ \underbrace{\mathbb{E}_{p^*_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Término Positivo}} - \underbrace{\mathbb{E}_{p^\phi_t}[\nabla_\phi \log \psi^\phi_t]}_{\text{Término Negativo}} \right]$$

• Término Positivo: Aumenta el peso del log-retorcimiento en muestras extraídas de la distribución objetivo óptima (regiones de alta recompensa).
• Término Negativo: Disminuye el peso del log-retorcimiento en muestras extraídas de la aproximación actual (regiones subóptimas).

2. Esquema de Entrenamiento Eficiente

Para abordar la ineficiencia de muestrear desde la distribución objetivo $p^*_t$ para el término positivo, los autores aprovechan una propiedad única de los modelos de difusión: el kernel hacia adelante de forma cerrada.

• Estimación de Gradientes Basada en el Proceso Hacia Adelante: En lugar de muestrear independientemente desde cada objetivo intermedio $p^*_t$, el método primero muestrea muestras positivas limpias ($\mathbf{x}_0$) desde la distribución objetivo utilizando SMC. Luego, utiliza el kernel hacia adelante de forma cerrada $p^{\text{base}}(\mathbf{x}_t | \mathbf{x}_0)$ para generar múltiples muestras positivas en cualquier paso de tiempo intermedio $t$ a un costo negligible.
• Entrenamiento Basado en Búfer: Se mantiene un búfer de muestras positivas limpias que se reutiliza a través de múltiples pasos de tiempo y actualizaciones de gradiente, reduciendo significativamente el número de evaluaciones de recompensa costosas requeridas durante el entrenamiento.
• Muestreo Negativo: Las muestras negativas se extraen mediante Muestreo por Importancia (IS) utilizando la función de retorcimiento parametrizada actual, lo cual es computacionalmente eficiente.

3. Parametrización

Para asegurar que la función de retorcimiento aprendida no introduzca sobrecarga en la inferencia, la función de retorcimiento $\psi^\phi_t$ se parametriza como una cabeza escalar ligera adjunta a la capa de características final del backbone de difusión preentrenado. Esto permite que el modelo produzca tanto logits como la estimación del retorcimiento en un solo paso hacia adelante, añadiendo menos del 5% de sobrecarga computacional en comparación con el modelo base.

Contribuciones Clave

1. Marco CDM: Introducción de un nuevo marco de amortización SMC para difusión discreta que reduce el costo de aplicar la función de retorcimiento en tiempo de inferencia a una operación de tiempo constante.
2. Entrenamiento Nativo de Difusión: Diseño de un esquema de entrenamiento que explota el proceso hacia adelante de forma cerrada de los modelos de difusión para permitir un entrenamiento eficiente con funciones de recompensa costosas, evitando la ineficiencia de muestreo de los enfoques basados en regresión estándar.
3. Versatilidad y Sinergia: Demostración de que CDM es agnóstico a la distribución de propuesta, permitiendo que se combine con propuestas ajustadas (por ejemplo, d1, DRAKES) para obtener ganancias de rendimiento sinérgicas.
4. Objetivo Superior: Validación de que el objetivo de aprendizaje contrastivo produce un rendimiento superior y una convergencia más rápida en comparación con los objetivos de retorcimiento basados en regresión estándar (Soft Value).

Resultados Experimentales

Los autores evaluaron CDM en cuatro aplicaciones diversas: generación de texto tóxico, diseño de secuencias de ADN regulatorio, diseñabilidad de proteínas y alineación de modelos de lenguaje grandes de difusión (dLLM).

• Rendimiento de Escalado: En todas las tareas, CDM superó consistentemente a las líneas base (Best-of-N, SMC estándar con $M$ variable y Soft Value) cuando se igualó el tiempo de reloj.
• Recompensas Costosas: En la tarea de diseñabilidad de proteínas, donde la evaluación de la recompensa es el cuello de botella dominante, CDM logró un escalado superior en comparación tanto con Best-of-N como con Soft Value, mientras que el SMC estándar no pudo escalar efectivamente debido a la sobrecarga de inferencia.
• Compatibilidad con Ajuste Fino: Cuando se integró con propuestas ajustadas (d1 y DRAKES), CDM mejoró aún más el rendimiento y mitigó problemas de colapso modal a menudo observados en modelos ajustados, manteniendo la diversidad mientras lograba recompensas altas.
• Eficiencia de Entrenamiento: CDM demostró una convergencia más rápida en tiempo de reloj en comparación con Soft Value, atribuido a la efectividad del objetivo de aprendizaje contrastivo con muestras negativas.

Significado y Afirmaciones

El artículo afirma que CDM aborda el cuello de botella principal en la inferencia SMC retorcida para difusión discreta: la dependencia de la estimación costosa de Monte Carlo. Al aprender la función de retorcimiento mediante un objetivo contrastivo que aprovecha el proceso hacia adelante de la difusión, el método permite una alineación de recompensas eficiente y escalable.

Los autores enfatizan que su enfoque es ortogonal al ajuste fino de propuestas; puede combinarse con métodos de ajuste fino existentes para lograr mayores ganancias de rendimiento. Además, el método sigue siendo aplicable incluso cuando la función de recompensa es no diferenciable o computacionalmente costosa, un escenario donde la guía basada en gradientes y el SMC estándar tienen dificultades. Los resultados sugieren que CDM establece una nueva frontera de Pareto para la compensación entre la calidad de la recompensa y el tiempo de inferencia en el modelado generativo discreto.