Resumen Técnico: Empujando las Fronteras de Utilidad-Diversidad Biomolecular con Optimización Relativa de Política de Supergroup

1. Declaración del Problema

La generación biomolecular tiene como objetivo producir candidatos (moléculas pequeñas o proteínas) que satisfagan objetivos químicos o biológicos específicos. Aunque el Aprendizaje por Refuerzo (RL) es un marco estándar para el post-entrenamiento de generadores preentrenados con el fin de maximizar la utilidad (por ejemplo, similitud con fármacos, puntuaciones de acoplamiento, estabilidad), la optimización exclusiva de la utilidad a menudo conduce a un colapso de modos, donde el generador concentra la masa de probabilidad en una familia estrecha de candidatos de alta puntuación. Por el contrario, mantener la diversidad es un desafío porque la diversidad es una propiedad a nivel de conjunto (definida sobre colecciones de muestras), mientras que la optimización estándar de políticas actualiza trayectorias individuales.

Los métodos existentes de RL que promueven la diversidad a menudo dependen de proxies indirectos, como penalizaciones de novedad relativas a muestras históricas o filtros basados en memoria. Estos enfoques pueden penalizar en exceso modos útiles de alta densidad o inducir una deriva distribucional. El desafío central abordado por este artículo es: ¿Cómo podemos optimizar directamente la diversidad del conjunto de muestras como un objetivo de primera clase, asignando al mismo tiempo crédito útil a los candidatos generados individualmente?

2. Metodología: Optimización Relativa de Política de Supergroup (SGRPO)

Los autores proponen la Optimización Relativa de Política de Supergroup (SGRPO), un marco flexible inspirado en la Optimización Relativa de Política de Grupo (GRPO) que optimiza directamente el equilibrio entre la utilidad a nivel de trayectoria y la diversidad a nivel de conjunto.

Mecanismo Central

SGRPO opera muestreando múltiples conjuntos de candidatos (grupos) bajo la misma condición y comparándolos para redistribuir las recompensas de diversidad. El proceso implica cuatro pasos clave:

1. Muestreo de Supergroup bajo la Misma Condición:
   Para una condición dada $C$ (por ejemplo, un bolsillo de proteína específico o ninguna condición), la política muestrea $M$ grupos, cada uno conteniendo $K$ trayectorias generadas independientemente. Esta colección $S(C) = \{G_1, \dots, G_M\}$ forma un "supergroup" de $N=MK$ candidatos. Las comparaciones se restringen a este supergroup para evitar confundir la calidad de la política con la dificultad intrínseca de diferentes condiciones.

2. Puntuación de Diversidad a Nivel de Grupo:
   A cada grupo $G_m$ se le asigna una puntuación de diversidad $R_m = D(G_m)$ basada en una métrica especificada por el usuario (por ejemplo, distancia Tanimoto interna para moléculas, similitud Levenshtein para proteínas). La señal de diversidad relativa del grupo se calcula mediante una comparación leave-one-out (dejar uno fuera) contra otros grupos en el supergroup:
   $$A_m^{grp} = R_m - \frac{1}{M-1}\sum_{h \neq m} R_h$$
   Un $A_m^{grp}$ positivo indica que el grupo es más diverso que sus alternativas.

3. Redistribución Consciente del Conjunto (Contribución Leave-One-Out):
   Para cerrar la brecha entre las señales a nivel de conjunto y las actualizaciones de políticas individuales, SGRPO redistribuye la recompensa de diversidad del grupo a las trayectorias individuales basándose en su contribución a la diversidad del grupo.

   • Para cada trayectoria $x_{m,i}$, se calcula una contribución leave-one-out: $c_{m,i} = D(G_m) - D(G_m \setminus \{x_{m,i}\})$.
   • Estas contribuciones se estandarizan y se utilizan para formar pesos softmax conscientes de la señal.
   • Si un grupo es altamente diverso ($A_m^{grp} > 0$), la señal positiva se concentra en candidatos con altas contribuciones de diversidad. Si un grupo es menos diverso ($A_m^{grp} < 0$), la señal negativa se concentra en candidatos con bajas contribuciones.
   • Esto asegura que la recompensa promedio del grupo se preserve mientras se asigna crédito a las moléculas específicas que impulsan la diversidad.

4. Actualización de Política Relativa al Supergroup:
   La recompensa final para una trayectoria combina su utilidad individual $r_{m,i}$ y la recompensa de diversidad redistribuida $\tilde{R}_{m,i}$:
   $$\bar{r}_{m,i} = (1-\lambda)r_{m,i} + \lambda \tilde{R}_{m,i}$$
   donde $\lambda$ controla el equilibrio utilidad-diversidad. La función de ventaja se calcula como una línea base leave-one-out sobre todas las trayectorias en el supergroup, y la política se actualiza utilizando un objetivo de estilo PPO recortado con regularización KL.

3. Contribuciones Clave

• Optimización Directa a Nivel de Conjunto: SGRPO es el primer marco que optimiza directamente la diversidad a nivel de conjunto como un objetivo principal en la generación biomolecular, en lugar de depender de proxies de novedad indirectos o penalizaciones dependientes del historial.
• Diseño Desacoplado: El método está desacoplado de arquitecturas de generadores específicas (autoregresivas o de difusión), recompensas de utilidad o métricas de diversidad, permitiendo su instanciación con varios optimizadores al estilo GRPO (por ejemplo, GRPO estándar para modelos autoregresivos, GRPO Acoplado para difusión discreta).
• Mecanismo de Asignación de Crédito: La introducción de un mecanismo de contribución de diversidad leave-one-out permite que el sistema asigne crédito a candidatos individuales que apoyan genuinamente la diversidad del conjunto, resolviendo el problema de optimizar una propiedad del conjunto con actualizaciones individuales.
• Robustez ante el Tamaño del Grupo: El método sigue siendo efectivo incluso con tamaños de grupo pequeños, lo que lo hace computacionalmente viable para el diseño biomolecular a gran escala.

4. Resultados Experimentales

Los autores evaluaron SGRPO en tres configuraciones distintas de generación biomolecular:

1. Diseño de Moléculas Pequeñas De Novo: Utilizando GenMol (difusión discreta).
2. Diseño de Moléculas Pequeñas Basado en Bolsillos: Utilizando GenMol-P (difusión condicionada por bolsillo).
3. Diseño de Proteínas De Novo: Utilizando ProGen2 (autoregresivo).

Métricas de Rendimiento:
La evaluación se centró en la Frontera de Pareto Utilidad-Diversidad, medida por:

• Hypervolume (HV): El área dominada por los puntos operativos no dominados.
• Distancia al Punto Ideal (DIP): Distancia al mejor punto teórico de utilidad y diversidad.
• Indicador R2: Déficit promedio ponderado del peor caso.

Hallazgos:

• Expansión de la Frontera: En todas las tareas y barridos de decodificación, SGRPO expandió consistentemente la frontera de Pareto utilidad-diversidad alcanzable en comparación con generadores preentrenados, GRPO estándar y GRPO asistido por memoria.
• Preservación de la Diversidad: En regímenes de alta utilidad donde las líneas base (especialmente GRPO) sufrieron un colapso severo de la diversidad, SGRPO mantuvo una diversidad significativamente mayor. Por ejemplo, en el diseño basado en bolsillos, SGRPO desplazó la trayectoria del punto operativo hacia afuera, reteniendo una diversidad marcadamente mayor en niveles de utilidad comparables.
• Estudios de Ablación: Eliminar la recompensa de diversidad revertió el rendimiento a niveles de GRPO estándar. Eliminar la asignación de crédito leave-one-out debilitó el rendimiento, confirmando la importancia de redistribuir las señales de diversidad a contribuyentes específicos.
• Dinámicas Distribucionales: Las visualizaciones de las distribuciones de secuencias mostraron que, mientras que GRPO y GRPO asistido por memoria se contraían en regiones estrechas o derivaban, SGRPO exploró y preservó múltiples clústeres, indicando una cobertura más amplia del espacio de secuencias.

5. Significado y Afirmaciones

El artículo afirma que SGRPO ofrece un principio de post-entrenamiento ampliamente aplicable para expandir los puntos operativos disponibles para los modelos de generación biomolecular. Al recompensar directamente conjuntos generados diversos, el método empuja la frontera utilidad-diversidad hacia afuera, permitiendo que los modelos logren una mejor utilidad en niveles de diversidad fijos (y viceversa) sin sacrificar la cobertura.

Los autores enfatizan que este enfoque aborda una limitación fundamental en el diseño biomolecular basado en RL actual: la incapacidad de tratar la diversidad como un objetivo de primera clase durante la optimización de políticas. Los resultados sugieren que optimizar directamente la diversidad a nivel de conjunto es una estrategia práctica y robusta para mitigar el colapso de modos en tareas de generación tanto de moléculas pequeñas como de proteínas. El código se pone a disposición para facilitar investigaciones futuras y su instanciación en diferentes dominios.