EquiFiLM: Charge-Conditioned Equivariant Force Fields via Feature-wise Linear Modulation

Resumen Técnico: EquiFiLM

Planteamiento del Problema

Los campos de fuerza basados en aprendizaje automático (MLFF) de base, como MACE-MP-0 y UMA, han alcanzado una precisión cercana a la de la teoría del funcional de la densidad (DFT) en amplios espacios químicos. Sin embargo, las formulaciones estándar se condicionan exclusivamente a las posiciones atómicas y las especies, asumiendo una física de estado fundamental en equilibrio. Carecen de un mecanismo nativo para manejar cambios inducidos externamente en el estado electrónico, tales como la carga total, campos aplicados o la excitación electrónica. En consecuencia, estos modelos no logran generalizar a procesos impulsados como la fotoexcitación o la inyección de carga.

Los enfoques existentes para incorporar el condicionamiento externo enfrentan compromisos significativos:

1. Especialistas por estado: Entrenar modelos separados para cada estado de carga produce una precisión aceptable para la carga de entrenamiento, pero falla al generalizar a cargas no vistas.
2. Bases fundamentadas en la física: Modelos como MACE-POLAR-1-M o 4G-HDNNP están construidos con conciencia de carga explícita (por ejemplo, equilibración de carga o maquinaria de Coulomb por átomo). Aunque son efectivos, requieren ser entrenados desde cero en corpus masivos (del orden de $10^8$ estructuras) y a menudo involucran una supervisión de carga atómica dedicada o un equilibrado iterativo complejo.
3. Degradación: Aplicar un modelo de base estándar a sistemas cargados resulta en una rápida degradación de la precisión, con errores de fuerza que ascienden a cientos de meV/Å.

Metodología

Los autores proponen EquiFiLM, un adaptador ligero y agnóstico al backbone que añade un eje de condicionamiento externo continuo a cualquier MLFF fundacional equivariante.

Mecanismo Central: ChargeFiLMBlock

El método utiliza Modulación Lineal por Características (FiLM), adaptada específicamente para redes equivariantes.

• Entrada: Un escalar de condicionamiento por grafo $c$ (por ejemplo, la carga total $q$ normalizada por el número de átomos).
• Arquitectura: Dos pequeñas redes MLP (Perceptrón Multicapa) mapean $c$ a los parámetros de compuerta $\gamma(c)$ y $\beta(c)$.
• Modulación: El bloque se inserta entre el tensor de mensaje y la proyección lineal en cada capa de interacción del backbone. Modula únicamente los canales escalares ($\ell=0$) de las características de los nodos:
  $$m'^{(t, \ell=0)}_i = (1 + \gamma(c)) \odot m^{(t, \ell=0)}_i + \beta(c)$$
  Las características equivariantes de orden superior ($\ell \ge 1$) pasan sin cambios.
• Preservación de la Equivarianza: Dado que los parámetros de compuerta dependen solo de un escalar por grafo (el cual es invariante ante rotación y traslación) y la modulación actúa únicamente sobre canales escalares, el bloque preserva exactamente la equivarianza E(3).
• Inicialización: Las capas de salida de las MLPs de compuerta se inicializan en cero, asegurando que al inicio del entrenamiento, el adaptador actúe como una función identidad, haciendo que el modelo sea bit-idéntico al backbone no condicionado.

Detalles de Implementación

• Backbone: El método se demuestra utilizando MACE-MatPES como el backbone.
• Datos de Entrenamiento: El modelo se entrena sobre trayectorias de dinámica molecular ab initio (AIMD) de agua líquida en cuatro cargas totales ($q \in \{0, 6e, 10e, 16e\}$) usando VASP al nivel meta-GGA r2SCAN.
• Eficiencia: El adaptador añade aproximadamente un 15% de parámetros respecto al backbone, pero no incurre en un costo de inferencia mensurable por paso en comparación con el backbone no condicionado.

Resultados Clave

El modelo resultante, E-MACE, fue evaluado en agua líquida cargada a través de cargas de entrenamiento, interpolación y extrapolación.

Precisión y Generalización

• Cargas de Entrenamiento: E-MACE logra una reducción de 3.1 veces en el RMSE de la fuerza (de 21.3 a 6.96 meV/Å) y una reducción de 61 veces en el RMSE de energía por átomo (de 6.1 a 0.1 meV/átomo) comparado con un modelo base MACE-MatPES ajustado con los mismos datos sin EquiFiLM.
• Comparación con Bases Grandes: E-MACE iguala la precisión de fuerza de MACE-POLAR-1-M, un modelo de base diseñado para la carga que es ~20 veces más grande (15.2M vs 0.76M de parámetros) y entrenado desde cero en conjuntos de datos masivos.
• Generalización: El modelo generaliza a cargas de interpolación ($q \in \{2e, 4e, 8e, 12e, 14e\}$) y de extrapolación ($q \in \{18e, 20e\}$). El RMSE de la fuerza permanece dentro de 18–61 meV/Å, y el RMSE de la energía dentro de 0.7–5.4 meV/átomo.
• Eficiencia de Datos: Se logra una alta precisión con datos mínimos. El RMSE de la energía se satura con solo el 25% de los datos de entrenamiento (1,600 fotogramas), y el RMSE de la fuerza se satura al 50% (3,200 fotogramas). Esto contrasta con las $\approx 10^8$ estructuras requeridas para entrenar bases conscientes de la carga desde cero.

Consistencia Física

• Conservación de la Energía: Las simulaciones de dinámica molecular (MD) utilizando E-MACE demuestran una conservación de energía estable en el ensamble microcanónico (NVE) a través de todo el rango de carga, con derivas medias por debajo de 0.01 meV/átomo. Esto confirma que el bloque FiLM preserva la naturaleza conservativa de las fuerzas.
• Respuesta Estructural: El modelo predice con éxito la respuesta dependiente de la carga de la función de distribución de pares (PDF) reducida, $\Delta G(r; q)$. Captura el desplazamiento hacia afuera de la primera capa de coordinación O-O y la emergencia de picos H-H a cargas altas (indicando la formación reductiva de H$_2$), consistente con observables de difracción de electrones ultrarrápidos (UED).

Estudios de Ablación

• Tipo de Modulación: El desplazamiento aditivo $\beta$ por sí solo captura la mayor parte de la ganancia de rendimiento, lo que sugiere que el efecto principal de la carga es un cambio en la línea base de la energía. La compuerta multiplicativa $\gamma$ se retiene por generalidad (por ejemplo, para el escalamiento de temperatura/presión) y añade un costo insignificante.
• Inyección de Entrada: Simplemente concatenar la carga como un token de entrada (sin modulación por capa) no mejora la precisión, confirmando que la modulación por capa es necesaria para que la red ajuste las características internas en cada profundidad.

Significado y Reivindicaciones

El artículo afirma que EquiFiLM establece un nuevo paradigma para extender los MLFF de base a procesos impulsados:

1. Problema del Adaptador: El condicionamiento de un MLFF de base se plantea como un "problema de adaptador", análogo a la adaptación de tareas en lenguaje y visión. Requiere solo un pequeño corpus de ajuste fino (miles de fotogramas) que abarque el eje de condicionamiento, en lugar de reentrenar la base desde cero.
2. Agnóstico al Backbone y al Condicionamiento: La receta se aplica sin cambios arquitectónicos a cualquier MLFF equivariante con canales de capa de interacción escalares y puede manejar cualquier variable externa escalar continua (por ejemplo, carga, temperatura, presión, fracción de dopaje).
3. Eficiencia vs. Rendimiento: Logra una precisión comparable a las bases masivas diseñadas específicamente para la carga con una fracción del conteo de parámetros y de datos de entrenamiento, manteniendo la velocidad de inferencia del backbone original.
4. Validez Física: Al preservar estrictamente la equivariancia E(3) y actuar solo sobre canales escalares, el método asegura que las fuerzas resultantes sean gradientes exactos de una energía dependiente de la posición, permitiendo simulaciones de MD estables de larga escala temporal.

Los autores conclen que este enfoque hace que sea factible añadir grados de libertad controlables (como carga o temperatura) a los modelos de base existentes para aplicaciones en electroquímica, fotoquímica y diseño de materiales dopados sin el costo prohibitivo de entrenar nuevas bases.