Resumen Técnico: Transporte de Recompensa – Control de Propiedades en el Ajuste de Flujo mediante la Alineación en el Espacio de Ruido

1. Planteamiento del Problema

En el Ajuste de Flujo (Flow Matching, FM) y modelos generativos relacionados basados en difusión, el acoplamiento —la regla que determina qué vector de ruido se empareja con qué punto de datos durante el entrenamiento— se trata típicamente como un heurístico computacional. Los enfoques estándar utilizan un emparejamiento independiente por simplicidad o transporte óptimo (OT) por minilote para enderezar las trayectorias y reducir la varianza. En estos paradigmas, el acoplamiento se optimiza únicamente para facilitar el entrenamiento, no para codificar propiedades estructurales específicas en el campo de flujo aprendido.

Consecuentemente, lograr la generación controlable (por ejemplo, generar moléculas con puntuaciones específicas de logP o QED) suele requerir mecanismos explícitos en el momento de la inferencia, tales como:

• Guía libre de clasificador (Classifier-free guidance): Requiere un valor de propiedad objetivo como entrada.
• Guía por gradiente (Gradient guidance): Computa los gradientes de un predictor de propiedades en cada paso de muestreo.
• Aprendizaje por Refuerzo (RL): Ajusta finamente el generador contra un oráculo, lo que a menudo requiere consultas al oráculo durante el muestreo.

Estos métodos introducen una sobrecarga de inferencia, dependencia de modelos de recompensa o de oráculos externos, o procesos de entrenamiento complejos. El artículo plantea si el propio acoplamiento puede servir como una interfaz de alineación, integrando la estructura controlable directamente en el campo de flujo durante el entrenamiento, de modo que el control esté disponible en la inferencia sin computación o llamadas al oráculo adicionales.

2. Metodología: Transporte de Recompensa

Los autores proponen el Transporte de Recompensa, un método que redefine el acoplamiento ruido-datos como un mecanismo para inyectar la estructura de la tarea en el campo de flujo.

2.1 Acoplamiento de Transporte Óptimo Alineado con la Propiedad

En lugar de emparejar el ruido y los datos de forma independiente, el Transporte de Recompensa construye un acoplamiento monótono basado en una propiedad molecular objetivo $y(x)$ (por ejemplo, logP, QED).

1. Claves Escalares: Para cada vector de ruido $z$ (una secuencia de embeddings de tokens), se computa una clave de ordenamiento escalar $s(z)$. La elección principal es la norma euclidiana de la media agrupada del ruido: $s(z) = \|\bar{z}\|_2$.
2. Emparejamiento Basado en Rangos: Dado un lote de vectores de ruido $\{z_n\}$ y moléculas $\{x_n\}$ con propiedades $\{y_n\}$, ambos conjuntos se ordenan por sus respectivas claves ($s$ y $y$). Luego se emparejan rango por rango: el vector de ruido con el $k$-ésimo $s$ más pequeño se empareja con la molécula que tiene el $k$-ésimo $y$ más pequeño.
3. Objetivo de Entrenamiento: El modelo de ajuste de flujo $v_\theta$ se entrena para transportar estos pares alineados. La pérdida combina un MSE estándar sobre la predicción del extremo (en el espacio de embedding continuo) y una pérdida de Entropía Cruzada (CE) para asegurar que los tokens decodificados formen moléculas válidas.

2.2 Condicionamiento de Dirección

Para asegurar que el modelo pueda acceder a la coordenada del espacio de ruido $s$ durante la inferencia, se introduce un MLP de Dirección. Este MLP de dos capas mapea el escalar normalizado $s$ a un embedding de dimensión $d$, el cual se añade a la entrada del transformador junto con los embeddings de tiempo y posición. Esto permite que el modelo condicione su campo de velocidad a la coordenada de ruido específica utilizada en el emparejamiento monótono.

2.3 Mecanismo de Inferencia

En la inferencia, el usuario especifica una señal de dirección escalar $s^*$ (una coordenada del espacio de ruido, no un valor de propiedad objetivo). El modelo genera moléculas integrando el campo de flujo comenzando desde vectores de ruido donde $s(z) \geq s^*$.

• Garantía Teórica (Proposición 1): En el límite donde el flujo aprendido preserva perfectamente la correspondencia de rango entre $s$ y $y$ (límite de preservación de acoplamiento), el condicionamiento en $s \geq s^*$ induce una distribución equivalente a un paso de selección del Método de la Entropía Cruzada (CEM). Específicamente, trunca la distribución generada al $\epsilon$-fracción superior de la distribución de la propiedad, donde $\epsilon$ está determinado por el umbral $s^*$.
• Sobrecarga Cero: Este control no requiere oráculo, ni modelo de recompensa, ni computación de gradientes, ni muestreo de rechazo en el momento de la inferencia.

3. Contribuciones Clave

1. El Acoplamiento como Interfaz de Alineación: El artículo redefine el acoplamiento ruido-datos en el ajuste de flujo de un heurístico de entrenamiento a una interfaz de alineación a nivel de distribución. Demuestra que elegir el acoplamiento basado en una propiedad organiza el campo de flujo aprendido de tal manera que las coordenadas del espacio de ruido corresponden a los valores de la propiedad.
2. Dirección Molecular de un Solo Botón (Single-Knob): El método logra un control monótono sobre las propiedades moleculares (logP y QED) utilizando una única coordenada de ruido escalar. Crucialmente, el mismo botón produce respuestas estructurales opuestas para diferentes objetivos: aumentar $s$ genera moléculas más grandes para logP (lipofilicidad) pero moléculas más pequeñas para QED (similitud con fármacos). Esto descarta un sesgo de tamaño genérico y confirma que el acoplamiento escribe la estructura específica del objetivo en el flujo.
3. Transporte de Recompensa para Secuencias Discretas: Los autores implementan esto para secuencias de tokens SELFIES de longitud variable, identificando y abordando dos modos de fallo específicos en los modelos de flujo que generan datos discretos:
   • Borrado de Señal Implícito: En arquitecturas Pre-LayerNorm, la información de magnitud (codificada en $\|\bar{z}\|$) se borra. El MLP de Dirección resuelve esto inyectando un patrón de dirección aprendido.
   • Confusión de Longitud (Length Confound): Sin una pérdida MSE no enmascarada sobre los tokens de relleno (padding), el modelo aprende un atajo donde un $s$ alto simplemente aumenta la longitud de la secuencia para satisfacer la restricción de la norma. El MSE no enmascarado obliga al modelo a regredir el ruido bruto en las posiciones de relleno, anclando la señal de magnitud y evitando este atajo.

4. Resultados Experimentales

El método fue evaluado en los conjuntos de datos ZINC-250K y GuacaMol.

• Desempeño de Control:
  • logP: En ZINC-250K, el barrido de $s$ de -3 a +7 desplazó el logP medio de 1.50 a 5.44 (+137%), con una correlación de Spearman $\rho=1.000$ entre $s$ y la propiedad media. El tamaño del efecto (Cohen's $d$) alcanzó 1.63.
  • QED: El mismo botón desplazó el QED de forma monótona (por ejemplo, +0.059 en $s=5$ en ZINC), con $\rho=1.000$.
  • Especificidad Estructural: Para logP, aumentar $s$ aumentó el número de átomos pesados (12 $\to$ 23). Para QED, aumentar $s$ disminuyó el número de átomos pesados (23 $\to$ 16).
• Generalización: Los resultados se replicaron en GuacaMol (1.6M de moléculas), mostrando tamaños de efecto aún más fuertes para logP ($d=2.31$) y una monotonicidad consistente para QED dentro del rango válido.
• Comparación con Baselines:
  • Comparado con el Ajuste de Flujo Condicional (donde los valores de propiedad son entradas explícitas), el Transporte de Recompensa logra una monotonicidad similar ($\rho=1.000$) pero opera mediante la truncación de la distribución en lugar de alcanzar un objetivo explícito.
  • Comparado con la difusión de predicción de $\epsilon$ ($\epsilon$-prediction diffusion), el método muestra que la alineación a nivel de acoplamiento está estructuralmente ausente en la predicción de $\epsilon$, ya que el objetivo Bayes-óptimo es independiente del acoplamiento, atenuando la señal del gradiente.
• Dinámica de Entrenamiento: La preservación del rango (la correlación entre $s$ y la $y$ generada) emerge rápidamente durante el ajuste fino (dentante de un epoch) y se estabiliza por debajo de 1.0 (por ejemplo, 0.57 para logP), lo cual es consistente con la brecha de "aproximación de distribución" descrita en la Proposición 1.

5. Significado y Reivindicaciones

El artículo afirma que el Transporte de Recompensa establece el acoplamiento en el ajuste de flujo como un canal de control distinto a nivel de distribución que puede portar la estructura de alineación por sí mismo.

• Eficiencia: Elimina la necesidad de llamadas al oráculo, modelos de recompensa o optimización iterativa en el momento de la inferencia, reduciendo significamente el costo computacional en comparación con métodos de RL o de guía.
• Mecanismo: Demuestra que la "capacidad de dirección" (steerability) de los modelos generativos no necesita depender de entradas de condicionamiento explícitas; puede integrarse en la geometría del modelo a través del acoplamiento de entrenamiento.
• Limitaciones y Alcance: Los autores son modestos sobre el alcance. El método proporciona un botón distributivo unidimensional (truncando la distribución) en lugar de una asignación de valor exacto por muestra. Funciona bien para propiedades escalares correlacionadas con características moleculares globales (como tamaño o polaridad), pero falla para propiedades topológicas como la Accesibilidad Sintética (SA), donde una clave escalar 1D no puede distinguir estructuras complejas de tamaño similar. Además, se muestra que el método es incompatible con los objetivos de predicción de $\epsilon$, resaltando que la interfaz depende del objetivo de predicción de la red.

En resumen, este trabajo desplaza el paradigma de la generación controlable de "guiar la trayectoria en la inferencia" a "alinear la distribución de entrenamiento", ofreciendo un mecanismo de control de propiedades con sobrecarga cero y basado en principios para el ajuste de flujo.