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

本文评估了针对机器学习原子间势（MLIP）基础模型在多种化学基准测试下的七种微调策略，结果表明，虽然基础模型质量和正确的能量初始化等先决条件至关重要，但对于单一系统的准确性而言，朴素微调是最优选择，而多头回放则能独特地保持用于更广泛部署的分布外鲁棒性。

要点

想象一下，你有一位大师级厨师，他花费多年时间，仅使用无机原料（如岩石、金属和盐）来学习烹饪完美的佳肴。这位厨师就是一个“基础模型”（Foundation Model）。现在，你想让这位厨师烹饪一道特定的新菜肴，比如一道精致的有机汤或生物炖菜。

核心问题是：你该如何教这位厨师学习这道新菜，既不让他忘记旧有的菜谱，又不破坏他原有的技能？

这篇论文是一场大规模的厨房实验，测试了七种不同的“微调”（Fine-tuning，即重新训练）这位大师级厨师的方法。研究人员发现，教学方法的重要性其实不如三个关键的“预烹饪”步骤：选择合适的厨师、设定正确的基准以及调节火力。

以下是他们研究结果的通俗解读：

1. 三项“飞行前检查”（最重要的部分）

在你开始教授新食谱之前，必须确保三件事万无一失。如果你在这上面出错，任何教学方法都无法挽救。

• 挑选合适的厨师（基础模型的质量）：

  • 类比： 你不会雇佣一位只会烧开水的厨师来教你如何制作舒芙蕾。
  • 研究发现： 原有模型的质量比微调策略更重要。一个在庞大且多样化的无机材料数据集（如“OMat24”模型）上训练过的模型，比一个旧的、规模较小的模型，在学习新的、奇特的化学性质方面要出色得多。即使使用相同的教学方法，一个“更好”的基础模型永远会产生更好的成品。

• 设定零点（原子参考能量 / $E_0$）：

  • 类比： 想象你在测量一栋建筑的高度。如果你从地下室而不是一楼开始测量，你的数据就会出错，建筑看起来可能会像是在漂浮或者被埋在地下。在化学中，你需要减去单个原子的“重量”，这样模型才能只学习它们是如何“相互作用”的。
  • 研究发现： 研究人员发现，使用一种智能的、“模型感知型”的方式来设定这个零点至关重要。如果你使用一个粗略的平均值，模型会变得不稳定。它在纸面上看起来可能表现很好（误差得分很低），但当你尝试模拟现实世界的物理现象（比如建筑物在风洞测试中的表现）时，它会崩溃。

• 调低火力（超参数）：

  • 类比： 在学习一项新技能时，你不希望因为动作太快而绊倒，也不希望因为动作太慢而永远无法完成。
  • 研究发现： 不同的教学方法需要不同的“学习率”。例如，一种叫做 LoRA（仅改变模型极小部分）的方法可以承受非常快的学习率，而一种同时教授两件事的方法则需要非常缓慢、温柔的节奏。

2. 七种教学策略

一旦通过了上述三项检查，研究人员测试了七种教授新食谱的方法：

1. 朴素微调（Naive Fine-Tuning）： “只管继续烹饪。” 你对整个厨师进行训练，让他学习新数据。
   • 结果： 对于完美掌握某一道特定菜肴非常有效。但如果你以后想让这位厨师做其他类型的食物，他可能会忘记旧有的技能（这被称为“灾难性遗忘”问题）。
2. 层冻结（Layer Freezing）： “不要触碰基础。” 你锁住厨师关于基础刀工的知识，只让他学习新的酱汁。
   • 结果： 效果不错，但有时过于僵化。这限制了厨师适应新食材的能力。
3. LoRA（低秩自适应）： “加一张小抄。” 你不是重写整本食谱，而是在厨师的围裙上贴一张高效的小便条，只涵盖新的规则。
   • 结果： 对于特定任务非常高效且准确，类似于朴素微调。
4. 多头重放（Multihead Replay）： “双头厨师。” 你给厨师戴上两顶帽子。一顶帽子用于新菜，另一顶帽子用于旧的、熟悉的菜肴。他们同时练习这两者。
   • 结果： 这是胜出者，因为它最安全。 这是唯一一种能持续防止厨师忘记旧技能的方法。它让厨师既能做好新菜，也能保持旧菜的水准。
5. 伪标签重放（Pseudolabel Replay）： “合成厨师。” 你不是使用真实的旧食谱，而是利用厨师对旧食谱的自我预测来进行练习。
   • 结果： 效果良好且具有灵活性，因为你不需要原始的旧数据，只需要厨师的记忆即可。
6. 重放 + LoRA（Replay + LoRA）： 将“小抄”与“双头帽子”结合起来。
   • 结果： 表现良好，但单纯使用“双头”模式通常就已经足够了。

3. 核心总结

• 不要重复造轮子： 如果你需要一个针对特定、狭窄任务的模型（比如仅仅模拟盐水），朴素微调是获得极佳结果最快、最简单的方法。
• 不要忘记过去： 如果你需要一个能够处理奇特、新颖情况（比如一种新型电池或复杂的生物分子）且不丢失原有训练能力的模型，你必须使用 多头重放（Multihead Replay）。这是唯一能让模型保持稳健并防止“遗忘”的策略。
• 质量胜于技巧： 论文强调，花时间挑选高质量的基础模型并正确设置能量参考，比选择完美的微调算法更为重要。如果基础薄弱或数学设置错误，世界上最好的教学策略也无济于事。

简而言之： 要获得最好的化学 AI，请从一个聪明的基础模型开始，正确设置你的数学规则，并且如果你希望 AI 具备多功能性且不易遗忘，请使用“双头”教学法（Multihead Replay）。

技术摘要

技术摘要：机器学习间原子势（MLIP）基础模型的微调

问题陈述
机器学习间原子势（MLIP）基础模型已展现出跨越不同化学体系的迁移能力，提供了一种避免从头开始训练特定任务势能的高资源消耗过程的工作流。然而，社区仍缺乏关于如何以及何时进行微调的系统性指导。早期报告表明，朴素微调（naive fine-tuning）往往会导致“灾难性遗忘”，从而促使人们采用最初为大型语言模型开发的约束技术（如层冻结、低秩自适应 LoRA）。本文研究了这些约束是否必要，或者早期的失败是否是由其他因素（如较弱的基础模型、不恰当的原子参考能量 $E_0$ 初始化以及不稳定的训练程序）导致的。本研究旨在表征塑造微调结果的主要因素，特别是目标任务的准确性和分布外（OOD）鲁棒性。

方法论
作者在五个具有化学多样性的基准测试、三代基础模型以及跨越五个数量级的训练集上，评估了七种不同的微调策略。

• 评估的微调策略：

  1. 朴素（Naive）： 通过持续梯度下降进行全参数更新。
  2. 层冻结（变体）（Layer Freezing (Variants)）： 冻结嵌入层/消息传递层并训练读出层（readout）；或冻结嵌入层和第一层消息传递层。
  3. 低秩自适应（LoRA）： 在冻结预训练权重的同时，向标量和等变线性层中注入可训练的低秩分解。
  4. 多头重放（Multihead Replay）： 使用独立的读出头，在目标数据和来自预训练或伪标签的重放数据集上进行同步优化。
  5. 伪标签重放（Pseudolabel Replay）： 多头重放的一种变体，其中重放标签由基础模型自身生成，从而将重放来源与原始预训练语料库解耦。
  6. 重放 + LoRA（Replay + LoRA）： 将多头重放与 LoRA 相结合。

• 基准测试： 研究涵盖了与 OMat24 预训练领域（周期性无机体相）偏离程度递增的体系：

  • 锂银硫锗电解质（无机周期性固体）。
  • 水合 NaCl（离子溶液）。
  • 冰多晶型（分子固体）。
  • SN2 反应（气相反应化学）。
  • SPICE 生物分子（有机/生物分子构象）。

• 技术实现： 作者在 MACE 代码库中实现了三项新功能：

  • 适配于等变消息传递架构（涵盖标量和等变线性层）的 LoRA。
  • 用于解耦重放数据源的伪标签重放。
  • 模型感知的原子参考能量（$E_0$）重估，以使预训练基准与目标数据对齐。

• 评估指标： 除了标准的点状能量和力误差外，研究还探测了动态和外推行为，包括来自分子动力学（MD）的径向分布函数（RDF）、努格尔弹性带（NEB）反应剖面、MD 稳定性测试以及用于检测短程排斥失效的随机结构搜索（RSS）。

关键结果

1. 前提条件决定策略选择： 研究发现，基础模型的质量、正确的 $E_0$ 初始化以及精心选择的超参数是前提条件，其影响通常超过了特定微调策略本身。

   • 基础模型质量： 较新的基础模型（例如基于 OMat24 的模型）在 OOD 迁移方面始终优于旧模型（基于 MPTraj 的模型），即使使用固定的微调方案也是如此。
   • $E_0$ 初始化： 使用“平均化”的 $E_0$ 会导致显著更高的误差和 MD 不稳定性（例如，冰模型在 50 ps 内失效）。“重估”后的 $E_0$（将预训练模型的零点与目标数据对齐）对于稳定性和可迁移性至关重要，其效果通常优于微调算法的选择。
   • 超参数： 朴素微调需要降低学习率并增加 EMA 衰减。LoRA 可以容忍较高的学习率。多头重放需要大幅降低学习率，以避免相互竞争的更新信号。权重衰减应设置为零，以防止参数偏离预训练解。

2. 按目标分类的性能：

   • 分布内专业化（单一体系）： 对于窄任务（如 SN2 能垒、水合 NaCl 溶剂化），大多数策略（朴素、LoRA、多头）都能实现极高的准确度，且一致性地超越了从头开始训练的模型。对于单体系应用，朴素微调提供了最佳的收敛性。
   • 分布外（OOD）鲁棒性： 当评估向相关但未见成分（如非银硫锗电解质）或不同化学性质（如生物分子）的迁移时，多头重放（使用原始数据或伪标签数据）是唯一能够持续保持 OOD 鲁棒性的方法。它在学习目标任务的同时，保持了对预训练分布的准确性，有效地防止了灾难性遗忘。
   • 冻结与 LoRA： 虽然在参数效率方面表现出色，但层冻结和 LoRA 在适应溶剂化特征或维持广泛的化学鲁棒性方面，与测试场景中的多头重放相比显示出了局限性。

意义与主张
本文声称，MLIP 中观察到的朴素微调的脆弱性很大程度上是由于次优的设置，而非该方法的内在局限性。作者认为：

• 朴素微调是一个可行且通常更优的起点，用于单体系应用，前提是基础模型质量高且 $E_0$ 已正确重估。
• 多头重放是更广泛部署时的必要策略，即在需要保留基础模型在微调分布之外的行为时。
• 伪标签重放通过允许使用任何结构多样化的数据集进行重放，提供了实际优势，消除了对访问原始预训练语料库的依赖。

这项工作为从业者建立了标准：投资于最强大的可用基础模型并确保正确的原子参考能量对齐，比选择特定的约束微调算法更为关键的设计选择。该研究为部署 MLIP 基础模型提供了系统性框架，将微调从一种小众选择转变为针对特定体系开发的默认起点。