Resumen Técnico: Un Marco de Campo Medio para el Control Distribucional en el Tiempo de Inferencia de Modelos de Difusión

Planteamiento del Problema

Los modelos de difusión se utilizan cada vez más como muestreadores controlables, donde los procesos de generación se orientan en el tiempo de inferencia para satisfacer objetivos específicos. Mientras que los métodos existentes manejan eficazmente las recompensas puntuales (definidas sobre muestras individuales, $r: \mathbb{R}^d \to \mathbb{R}$), muchas aplicaciones críticas requieren recompensas distribucionales (definidas sobre la medida generada, $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$). Estos objetivos distribucionales incluyen calibrar poblaciones con observaciones experimentales, fomentar la diversidad o equilibrar modos.

Los enfoques actuales para la orientación distribucional suelen incorporar términos de gradiente de la recompensa en la dinámica de inferencia (por ejemplo, potenciales repulsivos o verosimilitudes de conjuntos). Sin embargo, al igual que los primeros métodos de orientación puntual, estos heurísticos carecen de garantías teóricas respecto a la distribución muestreada. Específicamente, el simple hecho de añadir un término de gradiente no garantiza que las partículas tracen la medida inclinada deseada $\mu^*$, que maximiza el objetivo regularizado por KL:
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
A diferencia del caso puntual, donde el objetivo inclinado es una inclinación exponencial explícita ($\mu^* \propto p_1 e^r$), el caso distribucional resulta en un objetivo implícito donde el potencial de inclinación depende de la medida misma. Actualmente no existe un marco teóricamente fundamentado para corregir la dinámica y los pesos para apuntar a esta distribución implícita.

Metodología

Los autores proponen un Marco de Campo Medio para formalizar y resolver el control distribucional en el tiempo de inferencia.

1. Formulación Teórica

El problema se plantea como el apuntar a una medida inclinada $\mu^*_t$ en cada paso de tiempo $t \in [0,1]$, definida implícitamente por:
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
donde $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ es la primera variación (derivada funcional) de la recompensa dependiente del tiempo $R_t$. Esto crea una relación de autocoherencia donde la medida objetivo depende del potencial, el cual, a su vez, depende de la medida.

2. Dinámica de McKean-Vlasov Ponderada

Para trazar este camino, los autores derivan un sistema de McKean-Vlasov ponderado. A diferencia de la orientación por gradiente estándar, el término de orientación $\nabla_x \Psi_t(X_t; \mu_t)$ induce interacción entre las partículas porque depende de la medida empírica evolutiva $\mu_t$.

La dinámica idealizada para una partícula $X_t$ y su log-peso $A_t$ es:
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
Aquí, $\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ es un término de corrección crucial. A diferencia del caso puntual donde $\dot{\Psi}_t$ es simplemente $\partial_t r$, aquí involucra la evolución temporal de la medida $\mu_t$. Los autores demuestran que $\dot{\Psi}_t$ satisface una ecuación integral implícita que involucra la segunda variación del funcional de recompensa.

3. Algoritmo de Partículas Finitas

Para implementar esto, los autores proponen un esquema de partículas interactuantes ponderadas (Algoritmo 1):

• Actualizaciones de Posición: Las partículas se actualizan utilizando la deriva base, el score y el gradiente de la recompensa con respecto a la medida empírica actual $\hat{\mu}_t$.
• Actualizaciones de Log-Peso: Los pesos se actualizan para corregir la discrepancia entre la dinámica orientada y el camino inclinado objetivo. Esto requiere calcular $\dot{\Psi}_t$.
• Resolución de $\dot{\Psi}_t$: El artículo proporciona dos métodos para calcular el término de derivada temporal implícito:
  1. Sistema Lineal Implícito: Resolviendo un sistema lineal derivado de la ecuación integral (Proposición 3.3).
  2. Iteración de Punto Fijo: Un esquema de iteración de Picard ligero que aproxima las actualizaciones de peso.
• Remuestreo: Se emplea el remuestreo estándar de Monte Carlo Secuencial (SMC) para prevenir la degeneración de pesos.

El marco recupera la orientación de recompensa puntual estándar como un caso especial donde $R(\mu) = \int r(x) d\mu(x)$.

Contribuciones Clave

1. Formalización de la Orientación Distribucional: Los autores formalizan la orientación a nivel de distribución como el apuntar a una medida inclinada donde la inclinación depende de la medida misma, estableciendo un análogo distribucional a los objetivos de recompensa inclinados puntuales.
2. Derivación de la Dinámica de McKean-Vlasov Ponderada: Derivan un sistema dinámico principiado con un reponderado Feynman-Kac correctivo que asegura que las partículas interactuantes ponderadas tracen el camino inclinado deseado $\mu^*_t$.
3. Recuperación Teórica y Fundamentación: El marco recupera la orientación puntual estándar como un caso especial y proporciona una base teórica para los métodos heurísticos de orientación a nivel de lote existentes (por ejemplo, muestreo de diversidad, guía experimental).
4. Garantías de Convergencia: El artículo demuestra que el sistema de partículas finitas converge a la dinámica idealizada cuando el número de partículas $N \to \infty$.

Resultados Experimentales

Los autores validan el marco tanto en entornos de baja como de alta dimensionalidad:

• Baja Dimensionalidad (Mezcla 1D):

  • Tarea: Inclinar una mezcla Gaussiana bimodal hacia una distribución objetivo con pesos de modo diferentes usando una recompensa de la Distancia de Máximo Medio Cuadrático (MMD) al cuadrado.
  • Hallazgos: La orientación por gradiente estándar empuja las partículas hacia el objetivo pero falla en coincidir con la forma precisa de la distribución inclinada $\mu^*$. En contraste, el método de campo medio apunta con precisión a $\mu^*$, minimizando la brecha de optimalidad $J(\mu) - J(\mu^*)$ cuando la fuerza de orientación $\lambda$ coincide con el parámetro del objetivo. Esto se mantiene a través de varios esquemas de ruido.

• Alta Dimensionalidad (Conformación de Proteínas):

  • Tarea 1 (Proteasa HIV-1): Orientar conjuntos generados por Boltz-2 para coincidir con las distribuciones de distancia de residuos observadas experimentalmente (datos DEER). El modelo base estaba sesgado hacia el estado "cerrado"; el objetivo era recuperar las proporciones del estado "semi-abierto".
    • Resultado: La orientación de campo medio coincidió con precisión con la distribución de distancia de residuos empírica, mientras que la orientación solo por gradiente produjo una distribución con una forma ligeramente diferente y picos más fuertes.
  • Tarea 2 (Adenilato Kinasa): Orientar hacia una distribución 2D de ángulos de enlace derivada de simulaciones de Dinámica Molecular.
    • Resultado: La orientación de campo medio produjo un desplazamiento más fuerte hacia la distribución objetivo en comparación con los métodos de solo gradiente.
  • Tarea 3 (Cristalografía de Rayos X): Orientar las generaciones de Protenix para ajustarse a un mapa de densidad electrónica (estructura 4OLE) que exhibe conformaciones heterogéneas.
    • Resultado: Para la misma fuerza de orientación, la orientación de campo medio logró una mayor alineación de coseno con la densidad objetivo y colocó más muestras en la conformación alterna (Conformación B) en comparación con la orientación de solo gradiente, lo que sugiere que los métodos de solo gradiente sub-orientan respecto al verdadero objetivo del problema inverso.

Significado y Reivindicaciones

El artículo afirma cerrar la brecha entre el control de inferencia puntual y el distribucional proporcionando una base teórica sólida para este último.

• Corrección: La principal reivindicación es que el procedimiento propuesto apunta exactamente a la distribución inclinada por la recompensa prescrita $\mu^*$ en el límite de partículas infinitas, corrigiendo las discrepancias inherentes a la orientación por gradiente heurística.
• Interpretabilidad: El marco clarifica que los métodos de orientación por lote existentes están realizando efectivamente pasos de gradiente pero carecen del reponderado correctivo necesario para muestrear de la verdadera medida posterior/inclinada.
• Practicidad: Aunque el método introduce un costo computacional adicional (resolver para $\dot{\Psi}_t$), los autores señalan que esto suele ser insignificante en comparación con las evaluaciones de red en entornos de alta dimensionalidad (por ejemplo, proteínas) y que el resolvedor de punto fijo es eficiente.
• Limitaciones: Los autores reconocen que la dependencia del remuestreo puede ser una limitación para modalidades como las imágenes (potencialmente llevando a muestras duplicadas) y que el método requiere funciones de recompensa diferenciables. Sugieren que la combinación con técnicas de reducción de varianza (por ejemplo, resolvedores de la ecuación de Poisson) podría mitigar los problemas de remuestreo.

Este trabajo se posiciona no como un reemplazo para todos los métodos existentes, sino como un fundamento riguroso que explica por qué ciertos heurísticos funcionan y proporciona un mecanismo para asegurar que apunten a la distribución correcta.