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

Este artigo avalia sete estratégias de ajuste fino para modelos de fundação de potenciais interatômicos aprendidos por máquina (MLIP) através de diversos benchmarks químicos, revelando que, embora pré-requisitos como a qualidade do modelo de fundação e a inicialização de energia correta sejam primordiais, o ajuste fino ingênuo é ideal para a precisão em sistemas únicos, enquanto o replay multihead preserva unicamente a robustez fora da distribuição para uma implantação mais ampla.

O essencial

Imagine que você tem um mestre chef que passou anos aprendendo a cozinhar refeições perfeitas usando apenas ingredientes inorgânicos como rochas, metais e sais. Este chef é um "Modelo de Fundação". Agora, você quer que este chef cozinhe um novo prato específico, como uma sopa orgânica delicada ou um ensopado biológico, usando uma quantidade muito pequena de novas receitas.

A grande questão é: Como você ensina este chef o novo prato sem fazê-lo esquecer os antigos, ou sem arruinar suas habilidades existentes?

Este artigo é um enorme experimento de cozinha testando sete maneiras diferentes de "ajuste fino" (retreinar) este mestre chef. Os pesquisadores descobriram que o método de ensino importa menos do que três etapas críticas de "pré-cozimento": escolher o chef certo, definir a base certa e ajustar o calor.

Aqui está a divisão de suas descobertas em termos simples:

1. As Três Verificações de "Pré-Voo" (A Parte Mais Importante)

Antes mesmo de começar a ensinar a nova receita, você deve acertar três coisas. Se você errar isso, nenhum método de ensino o salvará.

• Escolha o Chef Certo (Qualidade do Modelo de Fundação):

  • A Analogia: Você não contrataria um chef que só sabe ferver água para lhe ensinar a fazer um suflê.
  • A Descoberta: A qualidade do modelo original importa mais do que a estratégia de ajuste fino. Um modelo treinado em um conjunto de dados vasto e diversificado de materiais inorgânicos (como o modelo "OMat24") é muito melhor em aprender novas químicas estranhas do que um modelo mais antigo e menor. Mesmo que você use o mesmo método de ensino, um modelo de fundação "melhor" sempre produzirá um prato final melhor.

• Defina o Ponto Zero (Energia de Referência Atômica / $E_0$):

  • A Analogia: Imagine medir a altura de um edifício. Se você começar a medir a partir do porão em vez do térreo, seus números estarão errados, e o edifício pode parecer flutuando ou enterrado. Na química, você precisa subtrair o "peso" dos átomos individuais para que o modelo aprenda apenas sobre como eles interagem.
  • A Descoberta: Os pesquisadores descobriram que usar uma maneira inteligente e "consciente do modelo" para definir este ponto zero é crucial. Se você usar um palpite médio e preguiçoso, o modelo se torna instável. Ele pode parecer bom no papel (baixas pontuações de erro), mas desmoronará quando você tentar simular a física do mundo real (como um edifício colapsando em um teste de túnel de vento).

• Diminua o Calor (Hiperparâmetros):

  • A Analogia: Ao aprender uma nova habilidade, você não quer se mover tão rápido que tropece, mas também não quer se mover tão devagar que nunca termine.
  • A Descoberta: Diferentes métodos de ensino precisam de diferentes "taxas de aprendizado". Por exemplo, um método chamado LoRA (que altera apenas uma pequena parte do modelo) pode lidar com uma taxa de aprendizado muito rápida, enquanto um método que ensina duas coisas ao mesmo tempo precisa de um ritmo muito lento e suave.

2. As Sete Estratégias de Ensino

Uma vez passadas as três verificações acima, os pesquisadores testaram sete maneiras de ensinar a nova receita:

1. Ajuste Fino Ingênuo (Naive Fine-Tuning): "Apenas continue cozinhando." Você pega o chef inteiro e continua treinando-o nos novos dados.
   • Resultado: Ótimo para aprender um prato específico perfeitamente. Mas, se você tentar usar este chef para um tipo diferente de comida mais tarde, ele pode ter esquecido suas habilidades antigas (um problema chamado "esquecimento catastrófico").
2. Congelamento de Camadas (Layer Freezing): "Não toque no básico." Você trava o conhecimento do chef sobre habilidades básicas de corte e permite que ele aprenda apenas o novo molho.
   • Resultado: Bom, mas às vezes muito rígido. Isso limita o quão bem o chef pode se adaptar aos novos ingredientes.
3. LoRA (Low-Rank Adaptation): "Adicione uma folha de dicas." Em vez de reescrever todo o livro de receitas, você adiciona um pequeno e eficiente bloco de notas ao avental do chef que cobre apenas as novas regras.
   • Resultado: Muito eficiente e preciso para tarefas específicas, semelhante ao ajuste fino ingênuo.
4. Replay de Múltiplas Cabeças (Multihead Replay): "O Chef de Duas Cabeças." Você dá ao chef dois chapéus. Um chapéu é para o novo prato, e o outro é para os pratos antigos e familiares. Eles praticam ambos ao mesmo tempo.
   • Resultado: Este é o vencedor para segurança. É o único método que consistentemente evita que o chef esqueça suas habilidades antigas. Mantém o chef bom no novo prato e nos antigos.
5. Replay de Pseudorótulo (Pseudolabel Replay): "O Chef Sintético." Em vez de usar receitas antigas reais, você usa as própias previsões do chef sobre receitas antigas para praticar.
   • Resultado: Funciona bem e é flexível porque você não precisa dos dados antigos originais, apenas da memória do chef.
6. Replay + LoRA: Combinando a folha de dicas com os dois chapéos.
   • Resultado: Bom, mas o "Dual Head" sozinho era frequentemente suficiente.

3. As Grandes Conclusões

• Não Reinvente a Roda: Se você precisa de um modelo para uma tarefa específica e estreita (como apenas simular água salgada), o Ajuste Fino Ingênuo é a maneira mais rápida e fácil de obter um ótimo resultado.
• Não Esqueça o Passado: Se você precisa de um modelo que possa lidar com situações novas e estranhas (como um novo tipo de bateria ou uma molécula biológica complexa) sem esquecer seu treinamento original, você deve usar o Multihead Replay. É a única estratégia que manteve o modelo robusto e seguro contra o "esquecimento".
• Qualidade Acima de Truques: O artigo enfatiza que dedicar tempo escolhendo um modelo de fundação de alta qualidade e definindo corretamente as referências de energia é mais importante do que escolher o algoritmo de ajuste fino perfeito. Se a fundação for fraca ou a matemática estiver configurada errada, o melhor método de ensino do mundo não ajudará.

Em resumo: Para obter o melhor IA para a química, comece com uma fundação inteligente, defina suas regras matemáticas corretamente e, se você quiser que a IA seja versátil e não esquecida, ensine-a usando o método de "Duas Cabeças" (Multihead Replay).

Resumo técnico

Resumo Técnico: Ajuste Fino de Modelos de Fundação de MLIP

Declaração do Problema
Modelos de fundação de potenciais interatômicos aprendidos por máquina (MLIP) demonstraram a capacidade de transferir conhecimento entre diversos sistemas químicos, oferecendo um fluxo de trabalho que evita o processo intensivo de recursos de treinar potenciais específicos para cada tarefa do zero. No entanto, a comunidade carece de orientação sistemática sobre como e quando realizar o ajuste fino (fine-tuning) desses modelos. Relatos iniciais sugeriram que o ajuste fino ingênuo frequentemente leva ao "esquecimento catastrófico", promovendo a adoção de técnicas de restrição (ex: congelamento de camadas, Adaptação de Baixo Posto ou LoRA) originalmente desenvolvidas para grandes modelos de linguagem. Este artigo investiga se essas restrições são necessárias ou se as falhas iniciais foram causadas por outros fatores, como modelos de fundação mais fracos, inicialização inadequada da energia de referência atômica ($E_0$) ou procedimentos de treinamento instáveis. O estudo visa caracterizar os principais fatores que moldam os resultados do ajuste fino, especificamente a precisão na tarefa-alvo e a robustez fora da distribuição (OOD).

Metodologia
Os autores avaliam sete estratégias distintas de ajuste fino em cinco benchmarks quimicamente diversos, três gerações de modelos de fundação e conjuntos de treinamento abrangendo cinco ordens de magnitude em tamanho.

• Estratégias de Ajuste Fino Avaliadas:

  1. Ingênua (Naive): Atualizações totais de parâmetros via gradiente descendente contínuo.
  2. Congelamento de Camadas (Variantes): Congelar camadas de embedding/passagem de mensagem enquanto treina os readouts; ou congelar o embedding e a primeira camada de passagem de mensagem.
  3. Adaptação de Baixo Posto (LoRA): Injetar decomposições de baixo posto treináveis em ambas as camadas lineares escalares e equivariantes enquanto congela os pesos pré-treinados.
  4. Replay de Múltiplas Cabeças (Multihead Replay): Otimização simultânea em dados da tarefa alvo e em um conjunto de dados de replay (da pré-treinação ou rotulado por pseudolabels) usando cabeças de leitura separadas.
  5. Replay de Pseudolabels: Uma variante do multihead replay onde os rótulos de replay são gerados pelo próprio modelo de fundação, desacoplando a fonte de replay do corpus original de pré-treinação.
  6. Replay + LoRA: Combinando multihead replay com LoRA.

• Benchmarks: O estudo abrange sistemas com crescente afastamento do domínio de pré-treinação OMat24 (sólido inorgânico periódico):

  • Eletrólitos de argirodita de lítio (sólido inorgânico periódico).
  • NaCl aquoso (solução iônica).
  • Polimorfos de gelo (sólido molecular).
  • Reações SN2 (química reativa em fase gasosa).
  • Biomoléculas SPICE (conformeros orgânicos/biomoleculares).

• Implementações Técnicas: Os autores implementaram três novas capacidades no código MACE:

  • LoRA adaptado para arquiteturas de passagem de mensagem equivariantes (cobrindo tanto camadas lineares escalares quanto equivariantes).
  • Replay de pseudolabels para desacoplar as fontes de dados de replay.
  • Reestimação do modelo de energia de referência atômica ($E_0$) consciente do modelo para alinhar as bases pré-treinadas com os dados alvo.

• Métricas de Avaliação: Além dos erros padrão de energia e força pontual, o estudo investiga comportamentos dinâmicos e extrapolativos, incluindo funções de distribuição radial (RDF) de dinâmica molecular (MD), perfis de reação de Nudged Elastic Band (NEB), testes de estabilidade de MD e Busca de Estrutura Aleatória (RSS) para detectar falhas de repulsão de curto alcance.

Principais Resultados

1. Os Pré-requisitos Dominam a Escolha da Estratégia: O estudo constata que a qualidade do modelo de fundação, a inicialização correta de $E_0$ e hiperparâmetros bem escolhidos são pré-requisitos cujo impacto rotineiramente excede o da estratégia específica de ajuste fino.

   • Qualidade da Fundação: Modelos de fundação mais novos (ex: baseados em OMat24) superam consistentemente os mais antigos (baseados em MPTraj) em transferência OOD, mesmo com receitas de ajuste fino fixas.
   • Inicialização de $E_0$: O uso de valores de $E_0$ "médios" leva a erros significativamente maiores e instabilidade de MD (ex: modelos de gelo falhando em 50 ps). Os $E_0$s "reestimados" (alinhando o ponto zero do modelo pré-treinado aos dados alvo) são críticos para estabilidade e transferibilidade, muitas vezes produzindo melhores resultados do que a escolha do algoritmo de ajuste fino em si.
   • Hiperparâmetros: O ajuste fino ingênuo requer taxas de aprendizado reduzidas e maior decaimento de EMA. O LoRA tolera taxas de aprendizado mais altas. O multihead replay requer taxas de aprendizado substancialmente menores para evitar sinais de atualização conflitantes. O weight decay deve ser definido como zero para evitar que os parâmetros se afastem da solução pré-treinada.

2. Desempenho por Objetivo:

   • Especialização Intra-Distribuição (Sistema Único): Para tarefas estreitas (ex: barreiras SN2, solvatação de NaCl aquoso), a maioria das estratégias (Ingênua, LoRA, Multihead) alcança alta precisão, superando consistentemente modelos treinados do zero. O ajuste fino ingênuo oferece a melhor convergência para aplicações de sistema único.
   • Robustez Fora da Distribuição (OOD): Ao avaliar a transferência para composições relacionadas, mas não vistas (ex: eletrólitos não-argirodita) ou químicas diferentes (ex: biomoléculas), o Multihead Replay (com dados originais ou de pseudolabels) é a única abordagem que mantém consistentemente a robustez OOD. Ele mantém a precisão na distribuição de pré-treinação enquanto aprende a tarefa alvo, prevenindo efetivamente o esquecimento catastrófico.
   • Congelamento e LoRA: Embora eficazes para eficiência de parâmetros, o congelamento de camadas e o LoRA mostraram limitações em adaptar-se a características de solvatação ou manter uma ampla robustez química comparado ao multihead replay nos cenários testados.

Significância e Alegações
O artigo alega que a percepção de fragilidade do ajuste fino ingênuo em MLIPs é amplamente resultado de uma configuração subótima, e não de uma limitação intrínseca do método. Os autores argumentam que:

• O ajuste fino ingênuo é um ponto de partida viável e frequentemente superior para aplicações de sistema único, desde que o modelo de fundação seja de alta qualidade e os $E_0$s sejam corretamente reestimados.
• O multihead replay é a estratégia necessária para implantações mais amplas onde a preservação do comportamento do modelo de fundação fora da distribuição de ajuste fino é exigida.
• O replay de pseudolabels oferece uma vantagem prática ao permitir o uso de qualquer conjunto de dados estruturalmente diverso para o replay, removendo a dependência de acesso ao corpus original de pré-treinação.

Este trabalho estabelece que, para profissionais, investir no modelo de fundação mais forte disponível e garantir o alinhamento correto da energia de referência atômica são escolhas de design mais críticas do que selecionar um algoritmo específico de ajuste fino restrito. O estudo fornece um framework sistemático para implantar modelos de fundação MLIP, movendo o ajuste fino de uma opção de nicho para um ponto de partida padrão para o desenvolvimento de sistemas específicos.