DeepMD-kit多任务预训练模型微调中的原子类型映射问题解析

问题背景

在使用DeepMD-kit进行分子动力学模拟时，研究人员经常会遇到需要基于预训练模型进行微调(fine-tuning)的场景。近期有用户报告了一个典型问题：在使用基于2024Q1版本的多任务预训练模型进行微调后，通过不同方式评估模型时出现了不一致的结果。

现象描述

用户基于OpenLAM 2.1.0的27头多任务预训练模型(分支为Domains_OC2M)进行微调，训练数据使用CP2K 2024.1计算得到。微调过程顺利完成，但在评估阶段发现：

1. 使用dp --pt test命令测试时，DFT计算结果与DP预测结果表现出良好的一致性
2. 使用Python API中的DeepEval模块进行重新标记(relabel)时，却发现与原始DFT标记数据存在显著差异

问题根源分析

经过深入调查，发现问题出在原子类型映射(atype mapping)的处理上。在多任务预训练模型中，原子类型映射必须与预训练阶段保持一致，即使用118种元素的完整type_map，而不是用户在input.json中自定义的type_map。

这与单任务预训练模型的行为不同。在单任务场景下，用户可以使用自定义的type_map，但在多任务模型中，必须保持与预训练一致的原子类型定义，否则会导致评估结果不一致的问题。

解决方案

对于使用多任务预训练模型进行微调的用户，需要特别注意以下几点：

1. 保持原子类型映射一致性：必须使用与预训练模型相同的118元素type_map，不能随意修改
2. 评估方法选择：不同评估方法可能对输入数据处理方式不同，需要确认评估流程是否正确处理了原子类型映射
3. 数据预处理：确保训练数据和评估数据中的原子类型都正确映射到预训练模型定义的原子类型空间

经验总结

这个问题揭示了DeepMD-kit在多任务学习场景下的一个重要特性。多任务预训练模型通过共享底层表示学习不同元素体系的通用特征，因此原子类型的全局一致性至关重要。用户在进行微调时，必须尊重预训练阶段建立的原子类型体系，不能像单任务模型那样自由定义原子类型映射。

这一发现不仅解决了当前用户的问题，也为后续使用多任务预训练模型的研究人员提供了重要参考，避免了类似问题的发生。同时，这也提示DeepMD-kit开发团队需要在文档中更明确地强调多任务模型与单任务模型在使用上的这一关键区别。