做化学信息学的人大概都有过这种别扭感手上明明是一个三维的、由原子核和电子云共同构成的实体可喂给模型的却是一串 SMILES 字符或者一张二维拓扑图。信息在第一步就被砍掉了大半——构象、电子离域、空间位阻、静电势分布全都没了。这几年我一直在琢磨怎么把“电子密度”这个东西直接塞进大模型里最近这套基于电子密度场的自回归模型内部都管它叫密度 GPT在一批下游任务上把准确率稳定压到了 98% 以上回归任务的相对误差也进了 2% 以内。这篇就把整个思路、架构取舍、数据管线、训练配置和踩过的坑一次性摊开讲清楚。适合已经有点深度学习基础、又想在分子建模上往前再走一步的人看纯小白也能看懂大方向因为我会尽量用生活化的类比把量子化学那层壳剥掉。1. 为什么把电子密度当作化学的语言来建模1.1 从字符串到三维场分子表示的三次跃迁分子表示这件事过去十几年基本走了三步。第一步是字符串SMILES 一行搞定存储便宜、检索方便但它本质是一维的语法描述连分子有没有手性都得靠 符号硬编码更别说构象了。第二步是图表示原子当节点、键当边消息传递网络在这一层上吃得很好可它有个隐藏假设键是离散的、定域的。共轭体系、过渡金属配合物、氢键网络这些地方键级本身就是个模糊概念你强行离散化模型学到的就是个近似。第三步就是三维场。电子密度 ρ(r) 是定义在连续空间上的标量场它不依赖你怎么画键也不需要你预先指定原子类型之外的任何化学直觉。分子间作用、极化、电荷转移全都写在这个场里。我个人认为这是表示层的一次质变不是简单的“维度变高”。打个比方前两种表示像是给一栋楼画平面示意图第三种是直接给你整栋楼的三维扫描点云——你要问承重墙在哪、管线怎么走答案天然就在里面。这个跃迁的代价也很直接数据量涨了两三个数量级模型输入从几百维变成几十万维训练成本完全不是一个量级。所以这套方案不是拿来替代图网络的而是给那些图网络明显吃力的任务准备的。1.2 Hohenberg-Kohn 定理给的底气敢把宝押在电子密度上底层依据是 Hohenberg-Kohn 定理基态电子密度唯一决定体系的一切基态性质。这句话翻译成人话就是——如果你能完整地拿到 ρ(r)那么能量、偶极矩、能隙、力的分布理论上全都是它的泛函。这就意味着我们不需要为每个下游任务设计一套专门的描述符只要模型能从密度场里把信息抽出来同一个骨干网络可以挂不同的头。当然现实里拿不到“完整”的密度。我们手上的是数值解有限基组、有限格点、有限泛函精度。所以模型学到的不是精确泛函而是一个在特定计算协议下自洽的近似映射。这一点必须提前想明白否则后面评测结果对不上你会怀疑人生。我见过太多人在这一步栽跟头——训练数据用的是一套泛函和基组测试数据换了另一套指标直接掉十几个点还以为是模型泛化不行其实是标签分布漂移了。还有一个容易被忽略的点电子密度天然满足 ∫ρ(r)dr N总电子数。这是免费的物理约束训练时可以作为正则项也可以作为推理后的校验。我们在损失里加了一个很轻的积分一致性项收敛稳定性肉眼可见地提升了。1.3 98% 这个数字到底该怎么读先把话说在前面脱离任务谈准确率是耍流氓。我这里的 98% 指的是几个具体场景下的一致率——元素环境分类判断某个空间格点属于哪种成键环境、反应位点二分类某个原子是否参与成键变化、以及构象能量排序的 top-1 一致率。这几个任务本身类别不平衡不严重随机基线在 50% 到 60% 之间所以 98% 是有含金量的但也别神化。性质回归那边我不敢用“准确率”这个词用的是相对误差。能量、能隙、偶极矩这几个量相对误差在 1.2% 到 1.8% 之间浮动个别小分子体系能压到 1% 以内。力场那块误差稍大原子受力分量的 MAE 大概在 0.03 eV/Å 量级做构象筛选够用做精细的过渡态搜索还差口气。还有一点必须强调划分方式决定了这个数字的可信度。随机划分下 98%和按骨架划分scaffold split下 98%完全是两个概念。我们早期的随机划分能刷到 99% 以上换成骨架划分直接掉到 93%查了半天才发现是训练集和测试集里有大量同源衍生物本质上是记忆而非泛化。后面所有报出来的数字都是在更严格的划分协议下测的。这个坑我下面还会细讲。1.4 这套方案的能力边界在哪不是所有化学问题都值得上密度场模型。如果你的任务只是预测分子量、logP 这种强依赖二维拓扑的经验性质图网络又快又准没必要折腾。反过来下面这几类任务密度场有明显优势涉及电子重排的反应性预测、激发态相关性质、非共价相互作用的强度估计、以及跨元素泛化——尤其是那些训练数据里样本很少的重元素体系。还有一个我特别看好的场景一致性校验。工业界做分子筛选常常几个方法给出互相矛盾的结论。密度模型给出的预测可以和传统量化计算交叉验证两者偏差大的样本挑出来人工看效率提升很明显。这个用法不追求绝对精度追求的是“把可疑样本圈出来”的能力。2. 整体架构设计把体素密度喂进 Transformer2.1 自回归目标怎么定从 GPT 借来的思路语言模型的核心是 next-token prediction放到密度场上最直观的类比是 next-voxel prediction。但直接按格点顺序预测有个致命问题三维空间没有天然的顺序你按 x-y-z 嵌套循环排出来的序列相邻 token 在空间上可能隔得很远模型学到一堆伪相关。我们最后采用的是分块加空间填充曲线Morton 序也就是 Z-order的方案。把整个立方盒切成 8×8×8 的小块每块用一个向量编码器压成一个 token然后按 Z 序把这些块排成序列。Z 序的好处是空间邻近的块在序列里也邻近局部性保住了同时又保留了自回归的因果结构。序列长度从 96³88 万降到 12³1728直接把显存问题解决了。这里有个关键取舍块内编码器是独立的还是共享的。我们用的是共享的 3D 卷积加池化参数量省下来一大截代价是块内细节靠后续的上采样头补。实测下来块大小从 4³ 加到 8³下游精度提升不到 0.5 个点但序列长度翻了三倍不划算。8³ 是我们的甜点。注意块大小和网格分辨率是耦合的。如果你把网格从 96³ 提到 128³ 而块大小不变序列长度会涨到 4096 左右注意力计算成本翻倍还多务必同步评估。2.2 体素化分辨率、边界与动态范围的三重博弈体素化看着简单实际是整个管线里最容易出错的一环。三个参数必须一起调盒子边长、格点间距、截断方式。盒子边长一般取分子最大原子间距离加 6 到 8 Å 的缓冲。缓冲太小密度尾巴被截断积分出来的电子数对不上缓冲太大有效信息占比下降模型大量算力浪费在空白区域。我们的做法是先统计数据集的尺寸分布取 99 分位再留 15% 余量用一个全局固定的盒子尺寸。固定尺寸的好处是可以用批量张量坏处是小分子在盒子里很空——这个后面用稀疏采样缓解。格点间距决定了能分辨的最小特征。0.15 Å 左右能看清成键区域0.25 Å 只能看个大概轮廓。我们起步用的是 0.3 Å后来发现对反应位点这种细节任务不够改到 0.2 Å。以 20 Å 的盒子算就是 100³取整到 96³。动态范围是另一个坑。电子密度在原子核附近能到 10² 到 10³ 量级远离分子快速衰减到 10⁻⁶跨了七八个数量级。直接送进网络梯度会被核附近的极大值主导。我们的处理是两步先做 ρ^(1/3) 变换这招是从密度拟合领域借鉴来的能把动态范围压到两个数量级内再除以总电子数做归一化。实测这样处理后远离区域的信息也不会被完全淹没。2.3 等变性不处理就是给模型埋雷分子在盒子里怎么摆放物理上不应该影响任何预测。但体素化之后转 90 度和不转输入张量完全不同。如果你的训练数据里旋转角度分布有偏模型就会学到伪特征。处理方式有三条路。第一条是数据增强随机旋转后重新体素化。理论上最干净代价是每次增强都要重算插值I/O 和 CPU 开销很大我们把增强放在数据加载的 worker 里异步做勉强能跟上。第二条是直接用 SE(3) 等变网络层数学上优雅但和自回归的因果结构结合时很别扭我们试过一版训练不稳定的问题折腾了两周没解决放弃了。第三条是折中在块编码器里用相对位置编码加上旋转增强做兜底。最后落地的就是第一加第三条。这里有个实操细节——旋转角度别用完全随机的连续值用绕三个轴的离散角度组合每次从 24 个旋转群元素里随机取插值误差可控而且不会因为大角度旋转引入过多重采样模糊。这个小改动把训练稳定性提升了不少。2.4 两阶段训练与多任务头整个训练分两段。第一段是自监督预训练任务是掩码密度建模——随机遮住 15% 到 20% 的块让模型预测被遮住块的密度特征。这个任务比纯自回归好训因为它同时看上下文不用一步步滚。我们试过纯自回归预训练收敛慢而且长程一致性差掩码版本在同样算力下下游效果好 3 个点左右。第二段是下游微调。骨干网络冻结前若干层后面接任务头。任务头设计上有个小心机所有任务共享一个中间投影层然后分叉。这样多任务联合训练时梯度能互相借力。我们同时挂了五个头——三个分类、两个回归损失权重按不确定性加权用可学习的 log 方差省去了手调权重的麻烦。3. 数据管线从量化计算结果到可训练张量3.1 结构来源与清洗脏数据比少数据更致命数据来源主要是几个公开的小分子构象库加上自己跑的一批。清洗环节我列几个必须做的动作。去重不能只看 SMILES 字符串。同一个分子不同构象、不同质子化状态、互变异构体字符串可能一样也可能不一样得用 InChI 或者规范化的图哈希来判。我们用的方案是先按规范图哈希去重再按 RMSD 阈值 0.5 Å 做构象层面的去重。这一步砍掉了将近 18% 的样本但测试集指标反而更可信了。剔除计算不收敛的样本。判断标准是看 SCF 是否在最大迭代内收敛、频率计算有没有虚频、以及密度积分和总电子数的偏差是否超过 10⁻³。这三条筛下来又去掉 6% 左右。别心疼这些样本留着只会让模型学噪声。还有一个慢工出细活的环节元素覆盖度统计。如果你的数据集里 90% 是 C/H/O/N模型对 S、P、卤素的预测会明显退化。我们的做法是分层采样保证稀有元素在训练集里有最低样本数不够的用相似环境的替换结构补。这个操作有点脏但对跨元素泛化的帮助很实在。3.2 计算参数设定一致压倒一切所有密度必须用同一套协议算出来。泛函、基组、格点密度、收敛阈值一个都不能变。我们用的是杂化泛函配三重 zeta 基组加弥散函数格点用最密的积分网格。这套配置算一个中等大小的有机分子大概几分钟到十几分钟几万个样本跑下来机时成本不低但可控。提示如果实在要混用不同协议的数据务必在输入里加一个协议编码one-hot让模型知道这条样本的出处。我们做过对照实验加了编码比不加跨协议测试误差低 40% 左右。还有个细节值得单独说密度的输出格式。很多量化程序默认输出 cube 格式格点是均匀的但原点和步长各样本不一样。你在读入的时候必须重采样到统一的网格上。重采样用三线性插值就够别用高阶插值容易在核附近产生振铃。3.3 体素化的可复现实现下面这段是我实际在用的核心逻辑做了简化去掉了和内部框架耦合的部分。import numpy as np def voxelize_density(coords, values, box_size20.0, grid96, buffer_ratio0.15, power1.0/3): coords: (N,3) 已重采样到统一盒子的格点坐标 values: (N,) 对应的电子密度值 返回: (grid, grid, grid) 的归一化密度张量 spacing box_size / (grid - 1) idx np.rint((coords box_size / 2) / spacing).astype(np.int32) valid np.all((idx 0) (idx grid), axis1) idx, values idx[valid], values[valid] vol np.zeros((grid, grid, grid), dtypenp.float32) np.add.at(vol, (idx[:, 0], idx[:, 1], idx[:, 2]), values) # 动态范围压缩先开三次方再按总电子数归一 vol np.sign(vol) * np.abs(vol) ** power total vol.sum() if total 0: vol / total return vol def blockify(vol, block8): 把 (G,G,G) 切成 (G/block)^3 个块每块展平成 token 输入 g vol.shape[0] // block blocks vol.reshape(g, block, g, block, g, block) blocks blocks.transpose(0, 2, 4, 1, 3, 5).reshape(-1, block ** 3) return blocks def z_order_index(g): 生成 Z 序索引保证空间邻近的块在序列里也邻近 order [] for i in range(g): for j in range(g): for k in range(g): order.append((i, j, k)) def key(t): x y z 0 for b in range(8): x | ((t[0] b) 1) (3 * b) y | ((t[1] b) 1) (3 * b 1) z | ((t[2] b) 1) (3 * b 2) return x | y | z return sorted(range(len(order)), keylambda i: key(order[i]))这段代码有三个地方值得多说两句。np.add.at那一步是把散点密度累加到网格上别用花式索引赋值会覆盖而不是累加这个 bug 我调了大半天才找出来。三次方变换放在归一化之前顺序反了效果差很多。Z 序那段用位交错实现看起来绕但比直接按莫顿码排序快一个数量级处理几万个样本时差别很明显。3.4 存储与加载别小看 I/O一个 96³ 的 float16 张量大概 1.7 MB十万个样本就是 170 GB。用 float32 直接翻倍。我们的做法是磁盘上存 float16 加压缩用 zstd压缩比能到 3 倍左右加载时解压并转 float32。数据加载器的 worker 数开到 CPU 核数的 70%每个 worker 里做旋转增强和块切分。这里有个反直觉的经验增强放在 worker 里做比放在 GPU 上做整体吞吐高。因为重采样是纯 CPU 密集GPU 做这个会打断计算流。我们实测下来worker 内增强的方案单卡吞吐提升约 25%。4. 训练实操配置、显存与收敛诊断4.1 起手配置与超参选择预训练阶段用的是 12 层 Transformer隐藏维度 768注意力头 12 个序列长度 1728。参数量在 1.1 亿左右单卡 80G 显存能跑 batch size 16配合梯度累积到等效 128。优化器用 AdamW学习率峰值 3e-4warmup 占前 5% 步数之后余弦衰减到 1e-5。权重衰减设 0.05这个值比 NLP 那边常用的 0.1 低因为我们的 token 维度小、层数少衰减太狠会欠拟合。微调阶段学习率降到 5e-5回归头用 Huber 损失delta 设 0.5分类头用带标签平滑的交叉熵平滑系数 0.05。多任务权重用不确定性加权自动学初始 log 方差设 0训练几百步之后会自己分化。有个细节预训练用的掩码比例别太高。我们试过 30%模型学出来的表示在下游任务上反而不如 15% 的版本因为上下文太少退化成局部插值了。4.2 显存与吞吐的几个实用开关混合精度是必开的bf16 比 fp16 稳不用折腾 loss scale。注意力那块如果序列再长可以考虑 FlashAttention 的实现但我们 1728 的长度还不至于。梯度检查点打开能省 40% 显存代价是训练慢 20% 左右看你的卡够不够。还有一个省显存的土办法块的编码器输出维度压到 256进了 Transformer 之后再投影到 768。这样块编码器的计算量降下来了而且我们发现对小分子任务块的局部特征本来也不需要太高的维度。4.3 收敛曲线怎么读预训练的 loss 曲线有几个典型形态值得记住。正常收敛是前 2000 步快速下降然后进入长尾缓慢下降全程抖动幅度不超过 5%。如果 loss 在 3000 步左右突然翘起来八成是学习率没降下来或者某批数据有问题先查数据。如果 loss 卡在某个值不动先看掩码位置。我遇到过一次模型把所有被遮住的块都预测成接近零的密度loss 稳定在一个不低的值原因是归一化的时候用了整样本的总电子数而被遮住的区域恰好是空白区预测零反而是最优解。后来改成按可见区域的统计量做归一化问题就没了。微调阶段看的是下游指标而不是 loss。多任务联合训练时经常出现某个任务指标先涨后跌这是任务间梯度冲突。我们的处理是给每个任务设一个梯度范数上限超过就裁剪简单粗暴但有效。4.4 微调时的冻结策略冻结多少层取决于你的下游数据量。数据量在几千条量级冻结前 8 层只调后面 4 层加任务头上万条的话全参微调效果更好。我们做过一组对照全参微调在小数据上过拟合很明显训练集指标 99%验证集掉到 91%。冻结之后验证集反而稳在 95% 以上。还有个偏门技巧微调初期用较小的学习率跑几百步再放开让任务头先适配骨干的表示空间能减少初期的大梯度冲击。这个操作我们叫“预热对齐”听起来玄实际就是分两段设学习率。5. 常见问题与排查速查表5.1 精度上不去的六种典型原因现象最可能的原因排查动作训练集指标高、验证集低数据泄漏同源构象跨集按骨架重新划分检查图哈希重叠所有任务都卡在同一水平骨干没学到东西预训练失败看掩码重建的密度可视化是否模糊只有回归任务差目标值分布偏斜未做变换检查标签分布考虑对数或分位变换稀有元素预测差样本覆盖不足分层统计各元素样本数补采样不同盒子尺寸样本混训效果差归一化不一致统一固定盒子或加尺寸编码训练后期指标抖动学习率衰减不足或 batch 太小降学习率、加梯度累积这张表是我从无数次 debug 里攒出来的按出现频率排序。第一行那个数据泄漏最阴险因为它表现得很像“模型泛化能力不行”很多人第一反应是加正则、加数据方向完全错了。判断方法很简单从训练集里随机抽一批样本算它们和验证集样本的图相似度如果大量样本相似度超过阈值那就是泄漏。5.2 数据侧的隐性 bug密度积分对不上总电子数是最常见也最容易被忽略的问题。我建议把这项检查写进数据管线的单元测试里每个样本读进来第一件事就是算积分偏差超过 10⁻³ 直接打日志丢弃。我们早期有一批数据因为 cube 文件头里的格点数写错了导致重采样时缩放错误模型学着学着就把某个元素的所有样本都预测偏了查了三天才发现是数据问题。还有一个坑是原子坐标和密度网格的对齐。有些程序输出的坐标是相对原点的密度网格是相对盒子的混用会导致分子在盒子里偏到一边。可视化一批样本的密度切片如果分子总是偏在某几个角就是这个问题。5.3 工程侧的性能陷阱数据加载是第一个瓶颈。如果你发现 GPU 利用率长期低于 60%先看数据加载是不是卡住了。用torch.utils.data.DataLoader的persistent_workersTrue加prefetch_factor4能缓解不少。第二个瓶颈在体素化的重采样。三线性插值用纯 Python 写会慢到怀疑人生一定要向量化或者直接用 scipy 的map_coordinates。我们早期用自己写的循环版本一个样本要 200 ms改成向量化后降到 8 ms。第三个是检查点保存。模型加优化器状态有几个 G每分钟存一次会把 I/O 打满。改成每 N 步异步存、只保留最近三个训练稳定性明显改善。5.4 评测协议里的三个陷阱第一个陷阱是随机划分。前面提过了不再重复总之报数字之前先确认划分方式。第二个陷阱是回归指标的选择。MAE 对离群值不敏感RMSE 敏感两个都报才全面。我们内部还额外看一个分位数误差90 分位绝对误差因为工业场景更关心最差的那批样本有多差。第三个陷阱是基线不公平。你拿一个训练了几个月的密度模型去和一个默认超参的图网络比赢了也没什么说服力。我们的做法是所有基线都做一轮轻量超参搜索报最好的结果。这样比较虽然费事但结论站得住。6. 落地场景与影响力边界6.1 现在就能替换掉的环节最直接能替换的是构象筛选里的能量排序。传统流程要跑一轮半经验或者低精度量化计算再挑前若干名做高精度验证。密度模型可以直接给出相对能量排序把候选集从几千缩到几十后面的高精度计算量降一个数量级。我们实测 top-20 的召回率在 97% 左右够用。第二个是反应位点预筛。给定反应物和试剂预测哪些原子最可能参与成键变化把组合爆炸的候选反应路径砍掉大部分。这个用法不要求绝对准确只要能保证真阳性不被漏掉太多。第三个是跨元素的性质外推。传统描述符方法遇到训练集里没有的元素基本歇菜密度场模型因为有物理约束打底外推表现好不少。我们做过一个实验训练集里完全没有含硫化合物测试含硫体系相对误差比图网络低一半左右。6.2 还碰不了的环节过渡态搜索和反应路径优化目前还不行。密度模型给的是静态快照的映射路径上的势能面曲率信息它抓不住。我们试过把相邻帧的密度拼成序列喂进去效果一般误差比传统方法大一个量级。激发态和非绝热过程也别指望。训练数据全是基态密度你让它预测激发态性质本质上是外推到了分布外指标看着还行但物理上不可靠。还有一个现实约束是成本。跑一次密度模型推理虽然比量化计算快几个数量级但比图网络还是慢不少。如果任务本身对精度要求不高用便宜的模型就够了别为了用新技术而用。我个人在实际操作中的体会是这套东西最大的价值不在于替代谁而在于给化学问题提供了一个统一的信息底座。以前做不同的任务要换不同的表示、不同的模型、不同的评估方式现在骨干是共用的换任务只是换个头。这种复用带来的工程效率提升比单纯刷几个点的精度要实在得多。当然前提是你得把数据管线做扎实——我见过太多项目死在数据上而不是模型上。
