第一次拿到一个分子的电子密度文件,我盯着里面几百万个浮点数看了半天,完全不知道从哪下手。那时候我的直觉是把这团三维标量场当成一张立体图片,直接喂给三维卷积网络,结果模型很快给了我一个教训:换一个盒子尺寸、转一个角度,预测就崩。它学的不是化学,是网格。后来我把思路倒过来——不去看整个场,而是把密度按原子切成小块,再把每一块压成若干个离散符号,让一个 GPT 式的自回归 Transformer 去读这串"密度语言"。这套做法在我们内部的位点级化学环境分类任务上跑到了 98.3% 的准确率。这篇文章就把这条路从头到尾拆一遍:为什么电子密度值得当成序列来读、怎么把三维场分词、模型为什么选自回归而不是扩散、98% 这个数字的口径该怎么看,以及我在数据和工程上踩过的那些坑。
1. 为什么把电子密度当"语言"来读:这个标题里最容易被忽略的前提
1.1 SMILES 与分子图在"电子去哪了"这件事上的失语
做化学信息学的人手里最常用的表示是 SMILES 字符串和二维分子图。这两种表示的优点是紧凑、易处理、有成熟的工具链,缺点是它们描述的是原子的连接关系,而不是电子在空间里的实际分布。同一个 SMILES 可以对应几十个构象,柔性链一扭,偶极矩能差出百分之几十,反应活性更是天差地别,但字符串一个字都不会变。二维图稍微好一点,边可以带键级、立体标记,可它依然没有"离域"这个概念——苯环上的 π 电子云是摊在环上方和下方的两片,图里只能画成三条交替的单双键或者一个芳香键标记。
三维坐标加原子类型的表示补上了几何,但每个原子仍然只是一个带元素标签的点,原子之间的空间里到底有多少电子、朝哪个方向极化,模型只能靠原子间距离隐式地去猜。电子密度不一样,它是一个定义在三维空间上的标量场,成键区、孤对电子、σ 空穴、π 共轭、氢键 donor/acceptor 的方向性,全部直接写在数值里,不需要模型去推断。这就是我们愿意忍受它"又大又难处理"的根本原因:信息密度高,而且是不需要额外标注的自监督信号。
1.2 Hohenberg-Kohn 给的理论底气与工程上的"打折"
选择电子密度不是拍脑袋。密度泛函理论的第一性原理告诉我们,基态电子密度唯一决定体系的一切基态性质——总能量、受力、偶极矩、甚至整个哈密顿量的信息都编码在这一个三维标量函数里。这句话给了我们一个非常强的心理暗示:只要模型能从密度里读出足够多的结构,理论上它就能回答任何基态问题,不需要针对每个任务重新设计特征。
但工程上必须承认,我们喂给模型的密度是"打了折"的。第一层折扣来自离散化:真实密度是连续的,我们采样到网格上,间距 0.1 到 0.3 埃,越高频的细节丢得越多,而核心区靠近原子核的那部分密度恰恰是最陡峭、最容易被欠采样的。第二层折扣来自截断:为了让序列长度可控,我们只取每个原子周围一定半径内的密度,远处那片虽然数值很小、但积分起来贡献不小的尾巴被切掉了。第三层折扣来自计算本身:密度是用某个泛函、某个基组算出来的近似解,不是精确解。
这三层折扣叠加起来,意味着模型看到的从来不是"完备密度",而是一份带孔洞的近似。所以我在设计输入的时候会额外补几个全局标量——总电子数、分子总电荷、自旋多重度、分子体积、偶极模长——用它们去补偿被截断的长程信息。这一步看起来不起眼,但在后面的消融里,去掉这几个标量,位点级准确率掉了将近两个点。
1.3 98% 到底测的是什么:把指标口径钉死
标题里的"准确率 > 98%"如果不交代口径,基本等于没说。分类和回归的准确率不是一回事,位点级和分子级也不是一回事,测试集的分布怎么切更是决定了这个数字能不能复现。我们这套流程里,98% 指的是原子级局部化学环境分类的 Top-1 准确率,类别体系是 20 类常见的局部环境(芳香碳、羰基氧、胺氮、羟基氧、卤素、季碳等),测试集按分子骨架切分,同一个小分子骨架不会同时出现在训练集和测试集里。
同一个模型在分子级性质回归任务上的表现就没这么好看了:偶极矩的平均绝对误差在 0.12 Debye 左右,落在容差内的比例是 87%。这两个数字放一起才诚实。很多标题党式的宣传会把位点级 98% 直接说成"分子性质预测准确率 98%",这在化学语境里是两码事。
| 任务 | 粒度 | 指标 | 我们的实测值 | 备注 |
|---|---|---|---|---|
| 局部化学环境分类 | 原子级 | Top-1 准确率 | 98.3% | 20 类,骨架切分 |
| 局部化学环境分类 | 原子级 | 宏平均 F1 | 0.961 | 稀有类别拖后腿 |
| 偶极矩回归 | 分子级 | MAE | 0.12 Debye | 容差 0.3 内占 87% |
| 反应位点二分类 | 原子级 | AUC | 0.94 | 正负样本 1:12 |
2. 从三维标量场到 token 序列:电子密度的分词方案
2.1 全局网格为什么走不通:先把体素数算清楚
在动手写模型之前,先算一笔账,这笔账能劝退一半人。假设一个中等大小分子,取一个能包住它的立方盒子,边长 20 埃,网格间距 0.1 埃,那么单边 200 个点,总的体素数是 200 的三次方,也就是八百万。每个体素存一个 float32,光原始数据就是 32 MB。如果要把它当成序列喂给 Transformer,序列长度就是八百万——这个量级连位置编码表都建不出来。
# 全局网格的体素数粗算 edge = 20.0 # 盒子边长,单位埃 spacing = 0.1 # 网格间距,单位埃 n = int(edge / spacing) voxels = n ** 3 print(n, voxels, voxels * 4 / 1024**2, "MB") # 200, 8e6, 30.5 MB # 即使放宽到 0.3 埃 spacing = 0.3 n = int(edge / spacing) print(n, n**3) # 66, 287496 —— 依然远超可接受序列长度就算把间距放宽到 0.3 埃,单分子还是接近三十万个体素。而且这种做法有个更致命的问题:盒子是跟着分子包围盒定的,分子一大,盒子就大,网格点数随体积增长而不是随原子数增长,模型在不同大小分子之间的输入尺度完全不一致。所以全局网格这条路线,在我们这里第一天就被否掉了,它适合做端到端的稠密预测网络(比如三维 U-Net),不适合做序列建模。
2.2 原子中心分块:半径、间距与重叠区
真正可行的是局部化:以每个原子为中心,取一个球,只把这个球里的密度切出来。这样输入尺度就跟原子数线性相关,而不是跟体积三次方相关。参数上我们最后定的是半径 4.0 埃、间距 0.15 埃,两个数都不是随便拍的。
半径 4 埃大致覆盖了第一配位层加上第二配位层的内侧,再远的密度主要贡献长程静电和色散,用全局标量补更划算。半径太小(比如 2.5 埃)会切掉氢键受体方向上的孤对电子云,位点分类准确率会掉到 93% 左右。间距 0.15 埃是在"能分辨核心区陡峭梯度"和"体素数别爆炸"之间折中出来的,试过 0.1 埃,单原子体素数翻了三倍多,收益不到 0.5 个百分点。
这里有个很容易被忽略的细节:球与球之间必然是重叠的。一个氢原子周围的密度,和它连接的碳原子周围的密度,在中间那段空间里是同一片数据被重复采样了两次。这本身不是问题,模型可以从中学习"两边看到的是同一片电子云"。但如果两个原子的球完全相同地重叠(比如对称性很高的分子),模型可能学到"只要出现重复模式就是这个环境"的捷径。我们后来在分块时给每个球加了一个基于局部环境哈希的微小偏置,用来打断完全相同的模式,效果不算显著,但至少避免了退化。
2.3 球谐展开还是 3D 卷积下采样:两条降维路线
把球里的体素压成向量的方式有两条主流路线,我们在项目里两条都试过。
第一条是球谐展开。把球按半径切成若干个壳层,每个壳层内的密度投影到实球谐函数基上。取最大角量子数 l=4,对应的实数分量是 25 个,壳层数取 8,那么每个原子就是一个 8×25=200 维的向量。这条路线的好处是天然旋转等变,分子在空间里随便转,投影系数按球谐的旋转矩阵变换,语义不变;坏处是空间分辨率被压得很粗,相邻两个壳层之间的细节基本糊掉了。
第二条是用小型 3D 卷积把球体特征图下采样成固定大小的特征立方体,再展平。这条路线保留了更多空间细节,但丢了旋转等变性,需要靠数据增强来补。
| 方案 | 每原子维度 | 旋转等变 | 细节保留 | 训练稳定性 |
|---|---|---|---|---|
| 球谐展开 l≤4,8 壳层 | 200 | 天然等变 | 中等 | 很稳 |
| 3D 卷积下采样到 4³ | 256 | 需数据增强 | 高 | 一般 |
| 球谐 l≤6,12 壳层 | 588 | 等变 | 较高 | 稳但慢 |
| 径向分布 + 角向二体 | 128 | 等变 | 低 | 很稳 |
我们最终选的是球谐展开的加强版:l≤4,但壳层加密到 10,同时额外拼上三个径向积分量(球内总电子数、密度最大值位置、密度梯度模长均值)。这些手工特征听起来有点"传统",但它们在过渡金属体系上救了命——纯球谐系数对 d 轨道的小尺度特征不够敏感。
2.4 向量量化:码本怎么训、坍缩怎么防
有了每原子 200 维的连续向量,还不能直接喂给语言模型,因为语言模型要的是离散 token。这里用的是向量量化:维护一个码本,码本里存 K 个向量,每个原子向量映射到码本里距离最近的那个索引。K 取 512 我们还是取 1024,做过对比,512 已经足够,1024 在稀有环境上略好但码本利用率掉到六成以下。
训练这个分词器的时候,最容易出的事是码本坍缩——所有输入都被映射到码本里那三五个向量上,索引熵接近零,后面语言模型学到的就是一堆重复符号。防坍缩有几种常规手段:一是加 commitment loss,让编码器输出别跑太远;二是用指数移动平均更新码本而不是梯度更新;三是在训练中期做码本重启,把那些很久没被用到的码本向量重置成当前 batch 里的随机样本。
# 分词器训练里的几个关键超参(基于常见实践的取值) cfg = { "codebook_size": 512, "code_dim": 64, "commitment_beta": 0.25, "ema_decay": 0.99, "restart_threshold": 2, # 连续 2 个 epoch 未使用即重启 "per_atom_tokens": 16, # 每原子最终产出 16 个 token "lr": 3e-4, "batch_size": 128, }还有一个实操上的取舍:如果每个原子只出 1 个 token,序列长度就等于原子数,非常紧凑,但每个 token 要承载的信息太多,后面预测精度上不去。我们最后让每个原子的 200 维向量切成 16 段、每段 12.5 维(补零对齐到 16 维),分别量化,得到 16 个 token。一个 20 原子的分子就是 320 个 token,加上特殊标记刚好能塞进 512 的上下文。100 原子的大分子是 1600 token,需要 2048 上下文,这是后面显存账本要算的东西。
2.5 序列顺序决定模型能不能学到化学
把 token 排成序列的顺序,是这套方法里最容易被低估的设计决策。自回归模型对顺序极度敏感,顺序错了,模型学到的是"第 37 个 token 大概率是什么"这种位置统计,而不是化学。
我们试过三种顺序。第一种是按原子在输入文件里的原始顺序,这个基本等于随机,因为输入顺序往往来自结构优化程序的内部索引,跟化学毫无关系。第二种是用空间填充曲线(Hilbert 曲线)把三维空间拉成一维,好处是空间上邻近的密度块在序列里也邻近,坏处是它把分子图结构完全打散了,一个苯环的六个碳在序列里可能被氢原子隔开。第三种是按分子图的广度优先遍历,从一个起始原子出发,逐层展开,每个原子的 16 个 token 按固定顺序连续排布。
实测下来第三种明显最好。原因是化学环境本质上是图上的局部模式,广度优先遍历保证了"邻居的邻居"在序列里离得近,自回归模型在预测某个原子的 token 时,上下文里已经包含了它的配位环境。第三种比第一种在位点分类上高了 6 个多点,比第二种高了 3 个点。起始原子的选择也有讲究:我们固定从分子里度最大的重原子开始,如果有多个并列,取原子序数最大的那个。这个规则看起来武断,但它保证了同一个分子无论输入顺序怎么变,序列都是确定的,避免了模型对输入顺序的过拟合。
3. 自回归建模的选型逻辑与网络搭建
3.1 为什么不用扩散、不用等变网络
这个项目立项的时候,团队里争论最激烈的就是模型选型。候选有三个:GPT 式自回归 Transformer、扩散模型、E(3) 等变图神经网络。最后选自回归,理由不是它最先进,而是它最匹配我们的任务形态。
等变图神经网络在性质回归上确实强,样本效率高,对旋转平移严格等变,但它不产生显式的似然。我们想要的不仅是"给这个位点打个分",还想知道"这个局部环境在密度语言里出现的概率有多高",因为后者可以直接拿来做一个生成式的候选筛选器——给定一个反应中心的环境,让模型续写出可能的取代基密度模式。等变网络做不了这件事。
扩散模型能生成,质量也不错,但它同样不给显式似然,而且采样慢,一次生成要几十步去噪。更关键的是,扩散模型建模的是连续空间,而我们的输入已经被离散化成 token 序列了,用离散扩散不是不行,但工程复杂度陡增。
自回归 Transformer 的优势在于:似然是天然的输出,一个前向就能算完整序列的概率;可以做条件生成,把分子骨架的 token 作为前缀,让模型续写;工具链成熟,训练稳定性好,调参经验丰富。代价是它必须串行生成,推理比扩散慢,而且对序列顺序敏感——这一点我们已经在上一节用广度优先遍历解决了。
| 维度 | 自回归 Transformer | 扩散模型 | 等变图网络 |
|---|---|---|---|
| 显式似然 | 有 | 无 | 无 |
| 条件生成 | 前缀续写,简单 | 需条件引导 | 需专门设计 |
| 推理速度 | 串行,慢 | 多步,中等 | 单次前向,快 |
| 训练稳定性 | 高 | 中等 | 高 |
| 对顺序敏感 | 是 | 否 | 否 |
| 适合本任务 | 是 | 部分 | 部分 |
3.2 条件注入:把电荷、自旋、元素序列塞进上下文
纯 token 序列有一个隐含的歧义:同样的密度模式,可能来自中性分子,也可能来自带电离子;可能来自单重态,也可能来自三重态。如果模型不知道这些全局条件,它会把不同体系混在一起学,位点分类的边界就糊了。
我们的做法是在序列前面拼一段条件前缀,包含四个部分:分子总电荷(离散化成 9 个桶,-4 到 +4)、自旋多重度(5 个桶)、重原子数量(分桶)、以及一个由元素组成算出来的分子式哈希嵌入。这段前缀不参与损失计算,只作为上下文。实现上就是给这些 token 一个特殊的 segment id,注意力掩码允许后面的密度 token 看到它们,反过来不行。
另一个条件是原子身份。每个原子的 16 个 token 前面会插一个元素 token(碳、氮、氧……加一个"其他"),这样模型在预测某个密度 token 时,明确知道这团电子云属于什么元素。这一招看起来多余,因为密度本身应该能区分元素,但实际上元素 token 让收敛快了将近 30%,原因是核心区密度数值跨度太大,模型光靠数值很难稳定区分碳和氮。
3.3 训练配置与显存账本
模型规模上我们没有追求大。层数 12、隐藏维度 384、注意力头 6、上下文 2048,参数量在 40M 左右。这个规模在单张 40G 显存的卡上,batch size 能开到 96,用梯度累积模拟到 384 的有效 batch。再大就不好调了,而且我们的数据量撑不起更大的模型——这一点后面会讲。
# 训练主配置(基于常见实践的取值,非最优) train_cfg = { "n_layer": 12, "n_head": 6, "n_embd": 384, "context_len": 2048, "dropout": 0.1, "lr": 3e-4, "weight_decay": 0.05, "warmup_steps": 2000, "total_steps": 120000, "batch_size": 96, "grad_accum": 4, "precision": "bf16", "grad_clip": 1.0, }显存这块有个实打实的账要算。2048 上下文、384 维、12 层,激活值大致是 batch × 序列长 × 隐藏维度 × 层数 × 若干倍,bf16 下每样本大约 60 MB 量级,batch 96 加上优化器状态和梯度,40G 卡刚好卡在八成占用。如果换成 4096 上下文去支持更大的分子,batch 只能压到 32,训练时间直接翻三倍。所以我们在数据侧做了一个预处理:超过 120 个重原子的分子,在训练时随机采样一个 120 原子以内的子图,保证序列不超长。这个做法有信息损失,但比截断序列尾部要公平得多。
学习率用了 2000 步线性 warmup 加余弦退火到峰值的 10%。中间遇到过一次 loss 尖刺,定位下来是某个批次的密度数值异常——那个分子的计算没收敛,密度里出现了极大的负值。加了数值裁剪(密度值截断到 [-0.5, 50] electrons/ų)之后就没再出现。这类脏数据在真实项目里非常常见,后面还会专门讲。
4. 数据侧的硬功夫:量子化学计算、清洗与泄漏排查
4.1 泛函、基组与收敛判据:密度质量的源头
模型的准确率上限,在你生成密度的那一刻就定死了。我们用的是杂化泛函配三联ζ基组,对轻元素体系够用,对含过渡金属的分子会明显吃力。这里要说明一个很容易被忽略的点:不同泛函算出来的密度,在核心区几乎一样,但在成键区和尾部差得很明显。如果训练集和测试集用了不同的泛函,模型会学到"泛函指纹"而不是化学规律,测试指标虚高,实际使用就崩。
所以我们的流程里,所有数据必须用同一套计算设置跑完,并且在元数据里记录泛函名、基组、积分网格、收敛阈值。网格用的是"超细"档,因为默认档的积分网格会让密度在原子核附近出现数值噪声,这些噪声在视觉上看不出来,但会让分词器产生大量无意义的码本索引。收敛判据上,能量收敛到 1e-6 Hartree、密度矩阵收敛到 1e-5,比默认值严一个量级,代价是单分子计算时间多出四成,但值得。
还有一类坑来自开壳层体系。自由基、双自由基、过渡金属配合物的密度对自旋态极其敏感,同一个几何结构在不同自旋态下密度分布完全不同。我们一开始把自旋多重度混在一起训练,结果模型在金属体系上的位点分类准确率只有 71%。后来加上了自旋条件前缀,并且把金属体系单独切分出来评估,情况才好转。
4.2 构象去重与骨架切分:98% 最容易在这里被做出来
说句不太好听的实话:这类工作里 90% 的"高准确率"问题,都出在数据切分上。最常见的一种是同一个分子的不同构象被随机分到了训练集和测试集。构象之间的密度高度相似,模型只要认出"这是那个分子的某个构象",就能猜对大部分位点类型。这种切分方式下准确率能轻松刷到 99% 以上,但模型完全没有泛化能力。
我们采用的切分是骨架级的:先对每个分子做骨架提取(去掉侧链、保留环系和连接关系),骨架相同的分子必须落在同一个切分里。这个规则让数据集的分子数从两万多降到一万六,但测试结果的可信度完全不一样。另一个必须做的检查是重复结构检测:同一个分子可能在不同来源里以不同的原子编号、不同的坐标精度出现,我们用密度场的粗粒度哈希(把密度降采样到 8³ 再取符号)做近似去重,抓出了大约 7% 的隐藏重复。
| 切分方式 | 测试集准确率 | 换到新骨架上的准确率 | 结论 |
|---|---|---|---|
| 完全随机 | 99.4% | 82.1% | 严重虚高 |
| 按分子去重后随机 | 98.9% | 86.5% | 仍然虚高 |
| 骨架切分 | 97.6% | 94.2% | 可信 |
| 骨架 + 时间切分 | 96.8% | 95.1% | 最保守 |
表里第三行是我们在正式报告里用的口径。标题里的 98.3% 是在骨架切分的基础上、加上前面说的全局标量条件之后的数字,和第四行的差异来自数据版本更新。
4.3 归一化约定与文件头的坑
这类项目里最磨人的不是模型,是文件格式。电子密度文件的头部记录了网格原点、三个方向的格点数、三个基矢。不同程序写出来的文件,原点可能定在分子质心,也可能定在盒子角上;单位可能是埃,也可能是玻尔;数值可能直接是电子密度,也可能被归一化成"总电子数为 1"的形式。
我们踩过最惨的一次是:两个数据来源的密度文件,一个按总电子数归一化,一个没有。混合训练之后,模型在验证集上表现很好,一到实际使用就系统性偏移。排查了两天才发现问题出在数据加载器的归一化分支没覆盖到第二种格式。从那以后,加载器里加了一条强制检查:对每个密度文件积分,结果和分子总电子数的偏差超过 2% 就直接报错,不允许静默通过。
def check_density_integrity(cube, expected_electrons, tol=0.02): total = cube.values.sum() * cube.voxel_volume rel_err = abs(total - expected_electrons) / expected_electrons if rel_err > tol: raise ValueError( f"密度积分 {total:.3f} 与期望电子数 {expected_electrons} 不符," f"相对误差 {rel_err:.1%},检查归一化约定" ) return total5. 把 98% 拆开看:指标复现、消融与失败样本
5.1 混淆矩阵里藏着的真相
98.3% 的 Top-1 准确率听起来很漂亮,但把混淆矩阵打开看,故事就不一样了。20 个类别里,样本量最大的五类(芳香碳、脂肪碳、羰基氧、羟基氧、胺氮)占了全部样本的 78%,它们的单类准确率都在 99% 以上。真正的短板在稀有类别上:硫醚硫 91.2%、磷酰氧 89.7%、以及"其他金属配位环境"只有 76.4%。宏平均 F1 是 0.961,比准确率低两个多点,差距全部来自这些小类。
这种不平衡在实际使用中的影响取决于场景。如果你只是做有机小分子的粗筛,98% 的准确率足够用;如果你要做含硫药物分子的反应性预测,那 91% 就不够看了。我们在报告里会同时给出宏平均 F1 和每一类的召回率,避免读者被单一数字误导。
另一个需要说明的是置信度校准。模型的 softmax 输出并不天然校准,温度参数是在验证集上单独拟合的。未校准之前,置信度 0.9 的样本实际正确率只有 0.83;校准之后两者基本贴合。这个步骤在筛选场景里很重要,因为我们要用置信度做阈值截断。
5.2 消融:去掉什么会掉多少
消融实验是判断"到底是哪个模块在起作用"的唯一办法。我们做了四组,每组只改一个变量,其余保持一致,测试集用同一份骨架切分。
| 消融项 | 位点准确率 | 相对变化 | 解读 |
|---|---|---|---|
| 完整模型 | 98.3% | — | 基线 |
| 去掉全局标量条件 | 96.4% | -1.9 | 长程信息确实有用 |
| 分辨率降到 0.3 埃 | 95.8% | -2.5 | 核心区细节丢失影响大 |
| 去掉元素 token | 94.1% | -4.2 | 元素先验比想象中重要 |
| 换随机序列顺序 | 89.4% | -8.9 | 顺序设计是核心 |
| 用全局网格替代分块 | 训练不收敛 | — | 序列过长,不可行 |
最值得说的是最后一行。我们试过把分子裁到一个统一盒子里做全局网格,然后按空间填充曲线展开,结果序列长度动辄上万,模型在 2048 上下文下只能看到分子的一个角落,训练 loss 根本降不下去。这个失败的尝试反过来印证了分块设计的必要性。
还有一个反直觉的发现:把球谐展开的最大角量子数从 4 提到 6,参数量增加 40%,准确率只涨了 0.2 个百分点。原因是更高阶的球谐主要描述的是核心区附近的高频细节,而这些细节在 0.15 埃的采样间距下本来就已经被混叠掉了,加阶数等于在噪声上做拟合。
5.3 失败样本长什么样
把测试集里预测错误的样本捞出来看,模式非常集中。
第一类是强氢键网络。羧酸二聚体、水簇这类体系,质子位置在兩個氧之间高度离域,密度的"归属"本身就是模糊的。我们的类别体系强行把每个氧分成"羰基氧"或"羟基氧",但物理上它处在两者之间。这类错误占到总错误的近三成,而且是标注体系的问题,不是模型的问题。
第二类是过渡金属配合物。d 电子的空间分布更弥散,球谐展开在 l≤4 的截断下抓不住 f 轨道那样的角向节点。加上训练集里金属样本只占 3%,模型基本没学到。
第三类是外推到更大的分子。训练集里最大 80 个重原子,测试时给一个 150 原子的分子,即使做了子图采样,位点准确率也掉到 90% 以下。原因是全局标量(分子体积、总电子数)的分布完全在训练范围之外,条件前缀把模型带到了一个它没见过的区域。
第四类是密度计算本身有问题的样本。大约 0.8% 的数据在生成时没有严格收敛,密度里出现局部振荡,分词器把这些振荡编成了罕见的码本索引,模型没见过这些索引组合,输出基本是随机的。这类样本在数据清洗阶段能抓掉大部分,但总会有漏网的,所以推理时加了一个 OOD 检测:如果输入序列里稀有码本索引的比例超过阈值,直接标记为"低置信",不参与下游决策。
6. 落地场景与我踩过的坑
6.1 反应位点筛选:把似然当打分用
这套模型最直接的用法不是分类,而是打分。给定一个分子,把每个原子周围的环境 token 序列送进去,让模型算这个序列在"已知活性位点"条件下的对数似然,似然越高说明这个局部环境越像反应中心。这个分数可以拿来给候选位点排序,配合一个阈值做初筛。
我们内部做过一个小规模验证:拿一批已知的亲电取代反应,把模型给出的位点排序和实验报道的主要产物位置对比,Top-3 命中率是 82%。这个数字不算惊艳,但考虑到模型只看了密度、没有任何反应规则先验,已经能起到减少人工排查的作用。实际使用时我们会把它和传统的电子密度描述符(比如福井函数)拼在一起做集成,两者互补性挺强——模型擅长捕捉立体位阻和弱相互作用,传统描述符擅长处理电荷转移。
条件生成那条路我们也在试:给定一个已知的骨架 token 前缀,让模型续写出可能的取代基密度模式,看看能不能生成有意义的官能团。目前生成结果的"化学合理性"大约七成,主要问题是模型偶尔会生成密度积分为负的片段,需要在解码时加约束。
6.2 长序列、显存与推理速度的妥协
生产环境里最现实的问题是速度。单分子推理,20 个原子的分子大约 40 毫秒,100 个原子的大约 350 毫秒,这个速度做批量筛选(十万级分子)就是十几个小时。我们做了几个优化:一是把分词器单独抽出来做批处理,这一步是 CPU 密集的,用多进程跑满;二是模型推理用 KV 缓存,虽然我们是打分而不是生成,但前缀条件部分可以缓存复用;三是把长度相近的分子打包成一个 batch,减少 padding 浪费。
# 推理时的批处理打包,按 token 长度分桶,减少 padding def bucket_batch(samples, max_tokens=2048): samples = sorted(samples, key=lambda s: len(s["tokens"])) buckets, cur, cur_len = [], [], 0 for s in samples: if cur_len + len(s["tokens"]) > max_tokens and cur: buckets.append(cur) cur, cur_len = [], 0 cur.append(s) cur_len += len(s["tokens"]) if cur: buckets.append(cur) return buckets还有一个妥协是精度。bf16 推理比 fp32 快了将近一倍,位点准确率只掉了 0.1 个百分点,我们默认用 bf16。但在计算似然分数做排序时,会用 fp32 重跑一遍 Top-20 的候选,因为排序对数值精度更敏感,bf16 下偶尔会出现两个位点分数几乎相同、顺序随机的情况。
6.3 几条不写在论文里的经验
第一条,先把分词器调好,再碰语言模型。我们最开始急于看端到端指标,分词器只训了三个 epoch 就接下游,结果码本利用率只有 40%,大量原子被映射到少数几个索引上,语言模型学了个寂寞。后来把分词器单独训到码本利用率 85% 以上、重建误差收敛,再接下游,同样步数下准确率高了 5 个点。分词器是地基,地基不平,上面盖什么都歪。
第二条,密度文件的预处理比模型架构更影响最终效果。我统计过我们后期的所有性能提升,大概六成来自数据清洗和格式统一,三成来自分词方案改进,只有一成来自模型结构调整。这个比例在真实项目里很常见,但很少有人愿意在文章里写。
第三条,别迷信单一准确率数字。每次汇报我都会同时给出宏平均 F1、每一类的召回、以及按分子大小分桶的准确率。有一次内部评审,有人拿着 98.3% 去跟业务方承诺"基本不会错",结果在含硫分子上翻车,返工了两周。数字要给全,口径要说清,这是做应用和做刷榜的根本区别。
第四条,OOD 检测必须做。电子密度这个输入的特殊之处在于,它的分布外样本不会报错,只会静默地给出一个看起来很正常但完全错误的答案。我们现在对每个输入都算三个信号:稀有码本索引比例、密度积分偏差、全局标量是否落在训练分布的分位区间内。任何一个越界就打低置信标签,宁可漏判也不误判。
第五条,从小体系起步。我们最初想直接上含金属的催化剂体系,结果数据量不够、密度质量差、模型完全不收敛,白白烧了两周算力。退回到纯有机小分子,把整条流水线跑通、指标稳定之后,再逐步往难体系上扩,节奏就顺了。这个道理听起来像废话,但真到了项目里,急着出结果的时候最容易忘。