Generación de conformadores moleculares 3D a través de la atención agregada y de grano grueso equivalente: el blog de investigación de inteligencia artificial de Berkeley




Figura 1: Arquitectura CoarsenConf.

La generación de confórmeros moleculares es una tarea fundamental en la química computacional. El objetivo es predecir estructuras moleculares 3D estables de baja energía, conocidas como confórmeros, dada la molécula 2D. Las conformaciones moleculares precisas son cruciales para varias aplicaciones que dependen de cualidades espaciales y geométricas precisas, incluido el descubrimiento de fármacos y el acoplamiento de proteínas.

Presentamos CoarsenConf, un codificador automático variacional jerárquico (VAE) equivalente a SE (3) que agrupa información de coordenadas atómicas de grano fino en una representación de nivel de subgráfico de grano grueso para una generación de confórmero autorregresivo eficiente.

Fondo

La granularidad gruesa reduce la dimensionalidad del problema, lo que permite la generación autorregresiva condicional en lugar de generar todas las coordenadas de forma independiente, como se hizo en trabajos anteriores. Al condicionar directamente las coordenadas 3D de los subgrafos generados anteriormente, nuestro modelo se generaliza mejor a través de subgrafos química y espacialmente similares. Esto imita el proceso de síntesis molecular subyacente, en el que pequeñas unidades funcionales se unen para formar grandes moléculas similares a fármacos. A diferencia de los métodos anteriores, CoarsenConf genera conformadores de baja energía con la capacidad de modelar coordenadas atómicas, distancias y ángulos de torsión directamente.

La arquitectura CoarsenConf se puede dividir en los siguientes componentes:
(I) El codificador $q_\phi(z| X, \mathcal{R})$ toma el confórmero de verdad fundamental de grano fino (FG) $X$, el confórmero aproximado RDKit $\mathcal{R}$ y el confórmero de grano grueso (CG ) conformador $\mathcal{C}$ como entradas (derivado de $X$ y una estrategia de CG predefinida), y genera una representación de CG equivalente de longitud variable a través del paso de mensajes equivalentes y convoluciones de puntos.
(II) Se aplican MLP equivalentes para conocer la media y la varianza logarítmica de las distribuciones anterior y posterior.
(III) El posterior (entrenamiento) o anterior (inferencia) se muestrea y se introduce en el módulo de selección de canales, donde se utiliza una capa de atención para aprender el camino óptimo de la estructura CG a FG.
(IV) Dado el vector latente FG y la aproximación RDKit, el decodificador $p_\theta(X |\mathcal{R}, z)$ aprende a recuperar la estructura FG de baja energía a través del paso de mensajes equivariantes autorregresivos. Todo el modelo se puede entrenar de extremo a extremo optimizando la divergencia KL de las distribuciones latentes y el error de reconstrucción de los conformadores generados.

Formalismo de tareas MCG

Formalizamos la tarea de Molecular Conformer Generation (MCG) como el modelado de la distribución condicional $p(X|\mathcal{R})$, donde $\mathcal{R}$ es el conformador aproximado generado por RDKit y $X$ es el óptimo confórmero(s) de baja energía. RDKit, una biblioteca de Cheminformatics de uso común, utiliza un algoritmo económico basado en la geometría de la distancia, seguido de una optimización económica basada en la física, para lograr aproximaciones de confórmero razonables.

de grano grueso




Figura 2: Procedimiento de granulado grueso.
(I) Ejemplo de granulado grueso de longitud variable. Las moléculas de grano fino se dividen a lo largo de enlaces giratorios que definen ángulos de torsión. Luego son de grano grueso para reducir la dimensionalidad y aprender una distribución latente a nivel de subgrafo. (II) Visualización de un confórmero 3D. Los pares de átomos específicos se resaltan para las operaciones de paso de mensajes del decodificador.

El granulado grueso molecular simplifica la representación de una molécula al agrupar los átomos de grano fino (FG) en la estructura original en perlas individuales de grano grueso (CG) $\mathcal{B}$ con un mapeo basado en reglas, como se muestra en la Figura 2 (I). El granulado grueso se ha utilizado ampliamente en el diseño molecular y de proteínas y, de manera análoga, la generación a nivel de fragmentos o de subgráficos ha demostrado ser muy valiosa en diversas tareas de diseño de moléculas 2D. Desglosar los problemas generativos en partes más pequeñas es un enfoque que se puede aplicar a varias tareas de moléculas 3D y proporciona una reducción de la dimensionalidad natural para permitir trabajar con grandes sistemas complejos.

Observamos que, en comparación con trabajos anteriores que se centran en estrategias de CG de longitud fija, en las que cada molécula se representa con una resolución fija de perlas CG de $N$, nuestro método utiliza CG de longitud variable por su flexibilidad y capacidad para admitir cualquier elección de cuentas de CG de longitud variable. técnica de granulado. Esto significa que un solo modelo de CoarsenConf puede generalizarse a cualquier resolución de grano grueso, ya que las moléculas de entrada pueden asignarse a cualquier número de perlas CG. En nuestro caso, los átomos que consisten en cada componente conectado resultante de la ruptura de todos los enlaces giratorios se engrosan en una sola cuenta. Esta elección en el procedimiento CG obliga implícitamente al modelo a aprender sobre los ángulos de torsión, así como las coordenadas atómicas y las distancias interatómicas. En nuestros experimentos, utilizamos GEOM-QM9 y GEOM-DRUGS, que en promedio poseen 11 átomos y 3 perlas CG, y 44 átomos y 9 perlas CG, respectivamente.

SE(3)-Equivarianza

Un aspecto clave cuando se trabaja con estructuras 3D es mantener una equivarianza adecuada. Las moléculas tridimensionales son equivariantes bajo rotaciones y traslaciones, o equivalencia SE(3). Aplicamos la equivalencia SE(3) en el codificador, el decodificador y el espacio latente de nuestro modelo probabilístico CoarsenConf. Como resultado, $p(X | \mathcal{R})$ permanece sin cambios para cualquier rototraducción del confórmero aproximado $\mathcal{R}$. Además, si $\mathcal{R}$ se gira 90° en el sentido de las agujas del reloj, esperamos que el $X$ óptimo muestre la misma rotación. Para obtener una definición y una discusión detalladas sobre los métodos para mantener la equivarianza, consulte el documento completo.

Atención Agregada




Figura 3: Backmapping de grueso a fino de longitud variable a través de Aggregated Attention.

Presentamos un método, que llamamos Atención agregada, para aprender el mapeo de longitud variable óptimo desde la representación CG latente hasta las coordenadas FG. Esta es una operación de longitud variable, ya que una sola molécula con $n$ átomos puede asignarse a cualquier número de $N$ cuentas CG (cada cuenta está representada por un solo vector latente). El vector latente de una sola cuenta CG $Z_{B}$ $\in R^{F \times 3}$ se usa como clave y valor de una operación de atención de cabeza única con una dimensión de incorporación de tres para que coincida con la x, coordenadas y, z. El vector de consulta es el subconjunto del conformador RDKit correspondiente a la cuenta $B$ $\in R^{ n_{B} \times 3}$, donde $n_B$ es de longitud variable ya que sabemos a priori cuántos átomos FG corresponden a una determinada cuenta CG. Aprovechando la atención, aprendemos de manera eficiente la combinación óptima de características latentes para la reconstrucción de FG. Llamamos a esto Atención Agregada porque agrega segmentos 3D de información FG para formar nuestra consulta latente. La atención agregada es responsable de la traducción eficiente de la representación CG latente a coordenadas FG viables (Figura 1 (III)).

Modelo

CoarsenConf es un VAE jerárquico con un codificador y decodificador equivalente a SE(3). El codificador opera sobre características atómicas invariantes SE(3) $h \in R^{ n \times D}$, y coordenadas atomísticas equivalentes SE(3) $x \in R^{n \times 3}$. Una sola capa de codificador se compone de tres módulos: de grano fino, de agrupación y de grano grueso. Las ecuaciones completas para cada módulo se pueden encontrar en el documento completo. El codificador produce un tensor CG equivalente final $Z \in R^{N \times F \times 3}$, donde $N$ es el número de perlas y F es el tamaño latente definido por el usuario.

El papel del decodificador es doble. El primero es convertir la representación gruesa latente nuevamente en espacio FG a través de un proceso que llamamos selección de canal, que aprovecha la atención agregada. El segundo es refinar la representación de grano fino de forma autorregresiva para generar las coordenadas finales de baja energía (Figura 1 (IV)).

Hacemos hincapié en que, mediante la granularidad gruesa mediante la conectividad del ángulo de torsión, nuestro modelo aprende los ángulos de torsión óptimos sin supervisión, ya que la entrada condicional al decodificador no está alineada. CoarsenConf garantiza que cada siguiente subgráfico generado se gire correctamente para lograr un error de distancia y coordenadas bajo.

Resultados experimentales




tabla 1: Calidad de los conjuntos de confórmeros generados para el conjunto de pruebas GEOM-DRUGS ($\delta=0.75Å$) en términos de Cobertura (%) y RMSD promedio ($Å$). CoarsenConf (5 épocas) se limitó a utilizar el 7,3 % de los datos utilizados por Torsional Diffusion (250 épocas) para ejemplificar un régimen de datos limitados y de bajo cálculo.

El error promedio (AR) es la métrica clave que mide el RMSD promedio para las moléculas generadas del conjunto de prueba apropiado. La cobertura mide el porcentaje de moléculas que se pueden generar dentro de un umbral de error específico ($\delta$). Presentamos las métricas media y máxima para evaluar mejor la generación robusta y evitar el sesgo de muestreo de la métrica mínima. Hacemos hincapié en que la métrica mínima produce resultados intangibles, ya que a menos que se conozca a priori el confórmero óptimo, no hay forma de saber cuál de los confórmeros generados 2L para una sola molécula es el mejor. La Tabla 1 muestra que CoarsenConf genera el promedio más bajo y el peor de los casos de error en todo el conjunto de pruebas de moléculas de DROGAS. Además, mostramos que RDKit, con una optimización económica basada en la física (MMFF), logra una mejor cobertura que la mayoría de los métodos basados ​​en el aprendizaje profundo. Para obtener definiciones formales de las métricas y discusiones adicionales, consulte el documento completo vinculado a continuación.

Para más detalles sobre CoarsenConf, lea el artículo en arXiv.

BibTex

Si CoarsenConf inspira su trabajo, considere citarlo con:

@article{reidenbach2023coarsenconf,
      title={CoarsenConf: Equivariant Coarsening with Aggregated Attention for Molecular Conformer Generation},
      author={Danny Reidenbach and Aditi S. Krishnapriyan},
      journal={arXiv preprint arXiv:2306.14852},
      year={2023},
}