- 人工智能
- 预训练
- 深度学习
- 科研
- 科学计算
【免费下载链接】AIDD
🔬 如果你关注晶体结构预测(CSP)或材料生成模型,AIDD 仓库中的UniPath项目值得一读:它用可学习时间流匹配(Learnable-Time Flow Matching)统一建模晶体结构与能量,并用SPD 流形(对称正定矩阵黎曼几何)来刻画晶格,在 MPT-S52 等基准上表现出色。本文面向新手,用最少公式讲清楚它"为什么胜出",并附快速上手步骤与源码导读。
1 分钟看懂:UniPath 要解决什么问题
晶体结构预测的典型任务是:给定化学式(元素种类与数量),生成能量最低的可能结构。UniPath 的论文标题直白地概括了它的野心——"One Path to Model Them All":一条生成路径同时搞定两件事。
| 任务 | 内容 | 数据形式 |
|---|---|---|
| CSP(晶体结构预测) | 生成晶格 + 分数坐标 + 元素类型 | {train,val,test}.csv+_ori.pt |
| MD(分子动力学) | 预测能量、力 | 轨迹数据 |
传统做法中,CSP 用 DiffCSP、CDVAE 等专用生成模型,MD 用 CHGNet 等专用势函数模型,两套代码各管一摊。UniPath 则让同一个 backbone 同时输出 CSP 生成向量和 MD 能量/力头,对应配置见 model_pl_uni.py,CSP 数据集目录需同时放 CSP + MD 两类数据(见 README.md 的 "Expected data layout" 一节)。
核心设计一:晶格为什么要放在 SPD 流形上
这是 UniPath 最有辨识度的设计,也是标题里"SPD 流形设计为何胜出"的答案。
直觉:晶格不是普通向量
晶格由 6 个参数(3 个边长 + 3 个夹角)定义。如果直接把这 6 个数当成欧式空间里的向量做流匹配,会出现尴尬的约束问题:
- 边长必须为正;
- 夹角必须满足体积为正的几何约束;
- 两个"接近"的晶格,用参数差衡量可能并不"接近"(比如 89° 和 91° 在参数差里是 2,但物理意义上只是微小倾斜)。
UniPath 的做法是把晶格矩阵 L 映射为SPD(对称正定)矩阵:
P = LᵀLP 一定对称、正定,天然满足所有物理约束;而且 P 对晶格的等价基变换(选不同的晶胞方向)不敏感——代码文件末尾还专门验证了这一点。相关实现在 src/flowmm/rfm/manifolds/spd.py:
- lattice_matrix_to_spd_matrix():
LᵀL映射; - geodesic():用矩阵对数/指数计算的黎曼测地线
A·exp(t·log(A⁻¹B)),而不是简单线性插值; - closest_isotropic():把任意 SPD 矩阵"压"到各向同性点,作为带条件先验的采样中心。
为什么这能"胜出"
- 永不越界:SPD 矩阵永远正定,生成的晶格永远是合法晶格,不需要额外的裁剪/投影修复;
- 度量合理:
spd_riemanian_geo配置下,距离、内积、归一化都用黎曼度量(见 arch_uni.py 中metric_normalized的判定),"相似晶体"在损失函数里真的更近; - 条件先验自然:
SPDGivenN类以"原子数 N → 期望体积"构造各向同性均值(N/密度)^{2/3}·I,流匹配从该均值出发生成,物理意义清晰(见 spd.py)。
项目把晶格流形做成了可插拔选项,方便你做消融实验:
non_symmetric / spd_euclidean_geo / spd_riemanian_geo / lattice_params / lattice_params_normal_base见 manifold_getter.py。
核心设计二:可学习时间流匹配(Learnable-Time Flow Matching)
从"固定时钟"到"自适应时钟"
流匹配(Flow Matching)训练一个向量场:给定当前样本 xₜ 和时间 t,预测把噪声推向真实数据的"速度" uₜ。传统实现里 t 是训练时随机抽的、推理时人为指定的"固定时钟"。
UniPath 走得更远:把时间 t 也变成网络要预测的量。核心改动集中在 arch_uni.py:
- 当
t_mode属于T_LEARN_LIST(learn、learn_V2、learn_V3.0、learn_V3.1,定义于 tlearn_util.py)时,模型会额外挂一个时间预测头 time_out; - 推理时若未给定 t,模型自己先跑一遍
csp_t_forward/md_t_forward推断当前所处的时间,再据此生成结构(见 infer_forward); learn_V3系列还会用一个可学习的时间掩码嵌入t_mask_emb处理"时间未知"的状态(见 CSPBackbone)。
训练侧的时间预测损失很简单——MSE 回归真实 t:
loss_t = F.mse_loss(t, t_pred) # learn_V3 版本会拼接 0 分支 loss = w_time * loss_t + (1 - w_time) * loss_flow见 model_pl_uni.py 与 L1080-L1082。
为什么这有价值?推理阶段模型能"自己知道走到哪一步了",配合推理时的退火机制(inference_anneal_slope,见 model_pl_uni.py),可以控制坐标/晶格各部分"收敛的先后节奏",生成更稳定的结构。
多阶段训练:PathRefine 是如何炼成的
UniPath 发布的 4 个 checkpoint(mpts52_PathRefine、mpts52_uns_PathRefine、mp20_PathRefine、mp20_uns_PathRefine)都不是单步训练出来的,而是多阶段流水线的产物,每阶段 warm-start 上一阶段的权重。以 mpts52 为例(完整命令见 README.md "Training from Scratch" 一节,脚本在 UniPath/script/uni_train/main/ 下):
| 阶段 | 做什么 | 关键参数 |
|---|---|---|
| ① CSP+MD 联合预训练 | 冻结部分参数学主向量场 | 150 epochs |
| ② 无监督流路径精修 | 引入 unsupervised flow 损失 | --w_uns 0.1 |
| ③ + 一致性损失 | 加 consistency loss 收尾 | --w_uns 0.1 --w_consist 0.05 |
两个值得新手注意的细节:
w_uns/w_consist的取值很敏感:mpts52 用 0.1/0.05,mp20_uns 却用 0.02/0.5,不同数据集需单独调参;- 阶段必须按顺序执行:warm-start 靠的是"部分权重加载"(
partial_ckpt_load),父 checkpoint 按精确文件名引用,自己重训时 step 数会变,需要更新引用脚本。
快速上手:从克隆到跑通采样
环境准备(详见 README.md):
git clone https://gitcode.com/BAAI/AIDD cd AIDD/UniPath conda env create -f environment.yml && conda activate crystal_uniug pip install -e remote/cdvae remote/DiffCSP-official remote/riemannian-fm -e .然后下载 checkpoint(~870 MB)与数据集(~22 GB,需 4 类 CSP + MD 数据齐全),设置DATA_DIR。
采样四连击(单卡 A100 上reconstruct约 1.5~3.5 小时):
CKPT=./ckpt_download/ckpt/mp20_PathRefine/every_n_epochs/epoch=499-step=190587.ckpt python scripts_model/evaluate.py reconstruct ${CKPT} --subdir reproduce_1 \ --inference_anneal_slope 10.0 --stage test --num_evals 1 \ --batch_size 1024 --sample_way flow --single_gpu && \ python scripts_model/evaluate.py consolidate ${CKPT} --subdir reproduce_1 && \ python scripts_model/evaluate.py old_eval_metrics ${CKPT} --subdir reproduce_1 --stage test && \ python scripts_model/evaluate.py lattice_metrics ${CKPT} --subdir reproduce_1 --stage test⚠️ 两个新手高频坑:
- 不要把
.ckpt单独移出去——模型加载时会向上找同目录的.hydra/config.yaml推断架构,移走就报Failed to find hydra config!; --inference_anneal_slope 10.0是复现论文的关键参数,默认配置里是 0(见 default_uni.yaml)。
评测指标主要看match_rate(匹配率)与rms_dist(RMS 距离),输出在old_eval_metrics_reconstruct_*.json。
源码导读:按这个顺序读效率最高
| 顺序 | 文件 | 看点 |
|---|---|---|
| 1 | UniPath/README.md | 环境、下载、采样、训练全流程 |
| 2 | src/flowmm/rfm/manifold_getter.py | 三大部件(类型/坐标/晶格)的流形组装 |
| 3 | src/flowmm/rfm/manifolds/spd.py | SPD 流形:测地线、采样、对数概率 |
| 4 | src/flowmm/model/solvers.py | 带投影的 ODE 积分器(Euler 步 +projx) |
| 5 | UniPath/arch_uni.py | 统一模型:CSP backbone + MD 头 + 时间预测头 |
| 6 | UniPath/model_pl_uni.py | 训练损失拼装与推理退火调度 |
| 7 | scripts_model/conf/vectorfield/rfm_cspnet_uni.yaml | 默认超参:hidden_dim 512、6 层、cutoff 7.0 |
总结:UniPath 胜出的三个关键词
- 统一(One Path):CSP 生成与 MD 能量预测共用一个 backbone,一条路径建模全部任务;
- SPD 流形:把晶格放进对称正定矩阵的黎曼几何里,约束天然满足、度量物理合理、采样永不越界;
- 可学习时间的流匹配:时间不再是外部旋钮,而是模型自己预测的内部状态,配合推理退火让生成过程更可控。
想深入细节?建议从 spd.py 的geodesic函数读起,再对照 model_pl_uni.py 的损失拼装,半小时就能建立完整心智模型。🚀
- 人工智能
- 预训练
- 深度学习
- 科研
- 科学计算
【免费下载链接】AIDD
相关推荐
突破欧氏空间限制:Geomstats流形机器学习完全指南
突破欧氏空间限制:Geomstats流形机器学习完全指南 你是否还在为非欧几里得数据(如旋转矩阵、概率分布、图形结构)的机器学习建模而困扰?传统欧氏空间算法在处
高效多图像查看器MulimgViewer:专业图像比较与批量处理的完整解决方案
高效多图像查看器MulimgViewer:专业图像比较与批量处理的完整解决方案 MulimgViewer是一个功能强大的开源多图像查看器,专为需要同时查看、比较
深度学习重构流体力学:5大技术突破实现CFD计算革命
深度学习重构流体力学:5大技术突破实现CFD计算革命 当传统计算流体动力学遭遇算力瓶颈,AI技术如何重新定义流体模拟的边界?DeepCFD项目通过深度卷积神经网
人工智能深度学习科学计算
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考