Fine-tuning MLIP foundation models: strategies for accuracy and transferability

Este artículo evalúa siete estrategias de ajuste fino para modelos fundacionales de potenciales interatómicos aprendidos por aprendizaje automático (MLIP) a través de diversos bancos de pruebas químicos, revelando que, si bien requisitos previos como la calidad del modelo fundacional y la inicialización correcta de la energía son primordiales, el ajuste fino ingenuo es óptimo para la precisión en sistemas únicos, mientras que la repetición de múltiples cabezales preserva de manera única la robustez fuera de la distribución para un despliegue más amplio.

Lo esencial

Imagina que tienes a un maestro chef que ha pasado años aprendiendo a cocinar platos perfectos usando solo ingredientes inorgánicos como rocas, metales y sales. Este chef es un "Modelo de Fundación". Ahora, quieres que este chef cocine un plato nuevo y específico, como una delicada sopa orgánica o un estofado biológico, usando una cantidad muy pequeña de nuevas recetas.

La gran pregunta es: ¿Cómo le enseñas al chef este nuevo plato sin que olvide cómo cocinar los anteriores, o sin arruinar sus habilidades existentes?

Este artículo es un experimento de cocina masivo que pone a prueba siete formas diferentes de "ajuste fino" (reentrenar) a este maestro chef. Los investigadores descubrieron que el método de enseñanza importa menos que tres pasos críticos de "pre-cocción": elegir al chef adecuado, establecer la base correcta y regular el calor.

Aquí está el desgido de sus hallazgos en términos sencillos:

1. Los tres controles de "pre-vuelo" (La parte más importante)

Antes de siquiera empezar a enseñar la nueva receta, debes hacer tres cosas bien. Si fallas en esto, ningún método de enseñanza te salvará.

• Elige al Chef Adecuado (Calidad del Modelo de Fundación):

  • La Analogía: No contratarías a un chef que solo sabe hervir agua para enseñarte a hornear un suflé.
  • El Hallazgo: La calidad del modelo original importa más que la estrategia de ajuste fino. Un modelo entrenado con un conjunto de datos enorme y diverso de materiales inorgánicos (como el modelo "OMat24") es mucho mejor para aprender química nueva y extraña que un modelo más antiguo y pequeño. Incluso si usas el mismo método de enseñanza, un "mejor" modelo de fundación siempre producirá un plato final mejor.

• Establece el Punto Cero (Energía de Referencia Atómica / $E_0$):

  • La Analogía: Imagina medir la altura de un edificio. Si empiezas a medir desde el sótano en lugar del piso de la calle, tus números estarán mal, y el edificio podría parecer que flota o que está enterrado. En química, necesitas restar el "peso" de los átomos individuales para que el modelo solo aprenda sobre cómo interactúan.
  • El Hallazgo: Los investigadores descubrieron que usar una forma inteligente y "consciente del modelo" para establecer este punto cero es crucial. Si usas una suposición promedio y perezosa, el modelo se vuelve inestable. Puede que parezca bueno en el papel (puntuaciones de error bajas), pero se desmoronará cuando intentes simular la física del mundo real (como un edificio colapsando en una prueba de túnel de viento).

• Baja el Calor (Hiperparámetros):

  • La Analogía: Al aprender una nueva habilidad, no quieres moverte tan rápido que tropieces, pero tampoco quieres moverte tan lento que nunca termines.
  • El Hallazgo: Diferentes métodos de enseñanza necesitan diferentes "tasas de aprendizaje". Por ejemplo, un método llamado LoRA (que solo cambia una parte diminuta del modelo) puede manejar una tasa de aprendizaje muy rápida, mientras que un método que enseña dos cosas a la vez necesita un ritmo mucho más lento y suave.

2. Las Siete Estrategias de Enseñanza

Una vez pasados los tres controles anteriores, los investigadores probaron siete formas de enseñar la nueva receta:

1. Ajuste Fino Naive (Naive Fine-Tuning): "Solo sigue cocinando". Tomas a todo el chef y lo sigues entrenando con los nuevos datos.
   • Resultado: Excelente para aprender un plato específico a la perfección. Pero si intentas usar a este chef para un tipo diferente de comida más tarde, es posible que haya olvidado sus habilidades anteriores (un problema llamado "olvido catastrófico").
2. Congelación de Capas (Layer Freezing): "No toques lo básico". Bloqueas el conocimiento del chef sobre las habilidades básicas de cuchillo y solo dejas que aprenda la nueva salsa.
   • Resultado: Es bueno, pero a veces es demasiado rígido. Limita qué tan bien puede el chef adaptarse a los nuevos ingredientes.
3. LoRA (Low-Rank Adaptation): "Añade una hoja de trucos". En lugar de reescribir todo el libro de cocina, añades una pequeña y eficiente libreta de notas al delantal del chef que solo cubre las nuevas reglas.
   • Resultado: Muy eficiente y preciso para tareas específicas, similar al ajuste fino Naive.
4. Replay Multi-cabezal (Multihead Replay): "El Chef de Doble Cabezal". Le das al chef dos sombreros. Un sombrero es para el nuevo plato, y el otro es para los platos viejos y familiares. Practica ambos al mismo tiempo.
   • Resultado: Este es el ganador en cuanto a seguridad. Es el único método que evita consistentemente que el chef olvide sus habilidades anteriores. Mantiene al chef bueno tanto en el nuevo plato como en los antiguos.
5. Replay de Pseudocategorías (Pseudolabel Replay): "El Chef Sintético". En lugar de usar recetas viejas reales, usas las propias predicciones del chef sobre recetas viejas para practicar.
   • Resultado: Funciona bien y es flexible porque no necesitas los datos originales antiguos, solo la memoria del chef.
6. Replay + LoRA: Combinar la hoja de trucos con el de doble cabezal.
   • Resultado: Bueno, pero el "Doble Cabezal" por sí solo era a menudo suficiente.

3. Las Grandes Conclusiones

• No reinventes la rueda: Si necesitas un modelo para una tarea específica y estrecha (como simular solo agua salada), el Ajuste Fino Naive es la forma más rápida y fácil de obtener un gran resultado.
• No olvides el pasado: Si necesitas un modelo que pueda manejar situaciones nuevas y extrañas (como un nuevo tipo de batería o una molécula biológica compleja) sin olvidar su entrenamiento original, debes usar el Replay Multi-cabezal. Es la única estrategia que mantuvo al modelo robusto y seguro contra el "olvido".
• Calidad sobre Trucos: El artículo enfatiza que dedicar tiempo a elegir un modelo de fundación de alta calidad y establecer correctamente las referencias de energía es más importante que elegir el algoritmo de ajuste fino perfecto. Si la base es débil o las matemáticas están configuradas incorrectamente, el mejor método de enseñanza del mundo no ayudará.

En resumen: Para obtener el mejor IA para la química, comienza con una base inteligente, establece tus reglas matemáticas correctamente y, si quieres que la IA sea versátil y no olvidadiza, enséñale usando el método de "Doble Cabezal" (Multihead Replay).

Resumen técnico

Resumen Técnico: Ajuste Fino de Modelos Fundacionales de MLIP

Planteamiento del Problema
Los modelos fundacionales de potenciales interatómicos aprendidos por máquina (MLIP) han demostrado la capacidad de transferirse a través de diversos sistemas químicos, ofreciendo un flujo de trabajo que evita el proceso intensivo en recursos de entrenar potenciales específicos para cada tarea desde cero. Sin embargo, la comunidad carece de una guía sistemática sobre cómo y cuándo realizar el ajuste fino (fine-tuning) de estos modelos. Informes preliminares sugirieron que el ajuste fino ingenuo a menudo conduce al "olvido catastrófico", lo que impulsó la adopción de técnicas de restricción (por ejemplo, congelación de capas, Adaptación de Bajo Rango o LoRA) desarrolladas originalmente para los grandes modelos de lenguaje. Este artículo investiga si estas restricciones son necesarias o si los fallos tempranos se debieron a otros factores, como modelos fundacionales más débiles, una inicialización inadecuada de la energía de referencia atómica ($E_0$) o procedimientos de entrenamiento inestables. El estudio tiene como objetivo caracterizar los principales factores que dan forma a los resultados del ajuste fino, específicamente la precisión en la tarea objetivo y la robustez fuera de la distribución (OOD).

Metodología
Los autores evalúan siete estrategias de ajuste fino distintas a través de cinco bancos de pruebas químicamente diversos, tres generaciones de modelos fundacionales y conjuntos de entrenamiento que abarcan cinco órdenes de magnitud en tamaño.

• Estrategias de Ajuste Fino Evaluadas:

  1. Ingenua (Naive): Actualizaciones completas de parámetros mediante descenso de gradiente continuo.
  2. Congelación de Capas (Variantes): Congelar las capas de embedding y de paso de mensajes mientras se entrenan las lecturas (readouts); o congelar el embedding y la primera capa de paso de mensajes.
  3. Adaptación de Bajo Rango (LoRA): Inyectar descomposiciones de bajo rango entrenables tanto en las capas lineales escalares como en las equivariantes, manteniendo congelados los pesos preentrenados.
  4. Replay Multicabezal (Multihead Replay): Optimización simultánea en datos de la tarea objetivo y un conjunto de datos de replay (del preentrenamiento o con etiquetas pseudo-etiquetadas) utilizando cabezales de lectura separados.
  5. Replay de Pseudo-etiquetado (Pseudolabel Replay): Una variante del replay multicabezal donde las etiquetas de replay son generadas por el propio modelo fundacional, desacoplando la fuente de replay del corpus de preentrenamiento original.
  6. Replay + LoRA: Combinar el replay multicabezal con LoRA.

• Bancos de Pruebas (Benchmarks): El estudio abarca sistemas con un aumento en la desviación respecto al dominio de preentrenamiento de OMat24 (sólido inorgánico periódico):

  • Electrolitos de argyrodita de litio (sólido inorgánico periódico).
  • NaCl acuoso (solución iónica).
  • Polimorfos de hielo (sólido molecular).
  • Reacciones SN2 (química reactiva en fase gaseosa).
  • Biomoléculas SPICE (conformeros orgánicos/biomoleculares).

• Implementaciones Técnicas: Los autores implementaron tres nuevas capacidades en el código de MACE:

  • LoRA adaptado para arquitecturas de paso de mensajes equivariantes (cubriendo tanto capas lineales escalares como equivariantes).
  • Replay de pseudo-etiquetado para desacoplar las fuentes de datos de replay.
  • Reestimación del modelo-consciente de la energía de referencia atómica ($E_0$) para alinear las bases preentrenadas con los datos de la tarea objetivo.

• Métricas de Evaluación: Más allá de los errores estándar de energía puntual y fuerza, el estudio indaga en comportamientos dinámicos y extrapolativos, incluyendo funciones de distribución radial (RDF) de dinámica molecular (MD), perfiles de reacción de Nudged Elastic Band (NEB), pruebas de estabilidad de MD y búsqueda de estructuras aleatorias (RSS) para detectar fallos de repulsión de corto alcance.

Resultos Clave

1. Los Prerrequisitos Dominan la Elección de la Estrategia: El estudio encuentra que la calidad del modelo fundacional, la correcta inicialización de $E_0$ y la elección adecuada de los hiperparámetros son prerrequisitos cuyo impacto supera rutinamente al de la estrategia específica de ajuste fino.

   • Calidad del Fundacional: Los modelos fundacionales más nuevos (ej. basados en OMat24) superan consistentemente a los más antiguos (basados en MPTraj) en transferencia OOD, incluso con recetas de ajuste fino fijas.
   • Inicialización de $E_0$: El uso de valores de $E_0$ "promediados" conduce a errores significativamente mayores e inestabilidad en MD (por ejemplo, modelos de hielo fallando en menos de 50 ps). Los $E_0$ "reestimados" (alineando el punto cero del modelo preentrenado con los datos de la tarea) son críticos para la estabilidad y la transferibilidad, ofreciendo a menudo mejores resultados que la elección del algoritmo de ajuste fino en sí.
   • Hiperparámetros: El ajuste fino ingenuo requiere tasas de aprendizaje reducidas y un mayor decaimiento de EMA. LoRA tolera tasas de aprendizaje más altas. El replay multicabezal requiere tasas de aprendizaje sustancialmente más bajas para evitar señales de actualización en competencia. El decaimiento de peso (weight decay) debe establecerse en cero para evitar alejar los parámetros de la solución preentrenada.

2. Rendimiento por Objetivo:

   • Especialización en la Distribución (Sistema Único): Para tareas estrechas (ej. barreras SN2, solvatación de NaCl acuoso), la mayoría de las estrategias (Ingenua, LoRA, Multicabezal) logran una alta precisión, superando consistentemente a los modelos entrenados desde cero. El ajuste fino ingenuo ofrece la mejor convergencia para aplicaciones de sistema único.
   • Robustez Fuera de la Distribución (OOD): Al evaluar la transferencia a composiciones relacionadas pero no vistas (ej. electrolitos no argyrodita) o químicas diferentes (ej. biomoléculas), el Replay Multicabezal (con datos originales o pseudo-etiquetados) es el único enfoque que mantiene consistentemente la robustez OOD. Mantiene la precisión en la distribución de preentrenamiento mientras aprende la tarea objetivo, previniendo eficazmente el olvido catastrófico.
   • Congelación y LoRA: Aunque son efectivos para la eficiencia de parámetros, la congelación de capas y LoRA mostraron limitaciones para adaptarse a características de solvatación o mantener una robustez química amplia en comparación con el replay multicabezal en los escenarios probados.

Significancia y Reivindicaciones
El artículo sostiene que la fragilidad percibida del ajuste fino ingenuo en MLIP es en gran medida el resultado de una configuración suboptimal en lugar de una limitación intrínseca del método. Los autores argumentan que:

• El ajuste fino ingenuo es un punto de partida viable y a menudo superior para aplicaciones de sistema único, siempre que el modelo fundacional sea de alta calidad y los $E_0$ se reestimen correctamente.
• El replay multicabezal es la estrategia necesaria para un despliegue más amplio donde se requiere preservar el comportamiento del modelo fundacional fuera de la distribución de ajuste fino.
• El replay de pseudo-etiquetado ofrece una ventaja práctica al permitir el uso de cualquier conjunto de datos estructuralmente diverso para el replay, eliminando la dependencia del acceso al corpus de preentrenamiento original.

Este trabajo establece que, para los profesionales, invertir en el modelo fundacional más fuerte disponible y asegurar la alineación correcta de la energía de referencia atómica son decisiones de diseño más críticas que seleccionar un algoritmo de ajuste fino restringido específico. El estudio proporciona un marco sistemático para desplegar modelos fundacionales de MLIP, moviendo el ajuste fino de ser una opción de nicho a un punto de partida por defecto para el desarrollo específico de sistemas.