MMAI Gym for Science: Training Liquid Foundation Models for Drug Discovery
Maksim Kuznetsov, Zulfat Miftahutdinov, Rim Shayakhmetov, Mikolaj Mizera, Roman Schutski, Bogdan Zagribelnyy, Ivan Ilin, Nikita Bondarev, Thomas MacDougall, Mathieu Reymond, Mihir Bafna, Kaeli Kaymak-Loveless, Eugene Babin, Maxim Malkov, Mathias Lechner, Ramin Hasani, Alexander Amini, Vladimir Aladinskiy, Alex Aliper, Alex Zhavoronkov
cs.LG, cs.AI, cs.CL
2026-03-04
英矽智能与 Liquid AI 用 400 余项药化任务对 LFM2 做 SFT 和 GRPO。2.6B 模型在分子优化、官能团推理和多项 ADMET 上压过 GPT-5.5 与 27B TxGemma，逆合成 unique 从 0 拉到 94%。
通用大模型做药化任务不稳。放大参数、加推理 token，增益很小。专精模型在 ADMET 预测、逆合成、3D 生成上更强，但一个任务一个模型，没法当跨任务助手。
英矽智能和 Liquid AI 把药物发现拆成 400 多项可验证的推理任务，用监督微调加强化学习，把通用因果语言模型改成药化通才。底座是 Liquid 的 LFM2：以门控短卷积为主、穿插少量分组查询注意力的混合架构，专门为长上下文和低延迟设计。他们训了两个尺寸，2.6B 稠密版和 24B（激活约 2B）的 MoE 版。
MMAI Gym 覆盖六类任务：2D 分子、3D 分子、2D 蛋白、3D 蛋白、药物-基因相互作用、跨模态。数据来自 TDC、MOSES、FGBench、MuMO-Instruct、GEOM、ZINC、CrossDocked2020 等公开集，再加内部数据。组 batch 时先在六类里均匀抽，再抽具体任务，避免大集淹没小集。
分子表示故意不锁死一种语法。输入随机在 SMILES、SELFIES、IUPAC 名称之间转换，SMILES 还做非规范遍历。词表加了化学专用 token，分子用标签包起来，乙酸会变成 <smiles>CC(=O)O</smiles>。3D 小分子先写图结构再按原子顺序列坐标；蛋白按残基铺氨基酸 token 和原子名。氢原子省略，推理后再用化学工具补。这样做是为了让模型学分子概念，避免只背某一种字符串格式。
训练分两段。SFT 约 10 万步，每步合计约 314 万 token，AdamW 学习率 1×10⁻⁵。随后用 GRPO 做在线强化学习：回归看归一化绝对误差，分类看对错，生成看 RDKit 能否解析成合法分子，再加 <think> 格式奖励和思考长度奖励。多任务时所有奖励压到同一尺度，KL 系数 0.4，防止任务互相冲。生成类任务大约跑 1000 步，采样温度 1.4。他们同时训了多任务通才和单类专才。ADMET 这一类里就有 22 个性质预测任务，单类并不等于单终点。
评测时每条测试样本会换指令模板、换分子编码、多次采样，回归取中位数，分类取多数票。闭源大模型走 few-shot，没有在 MMAI Gym 上微调。这个对照能说明领域配方有没有用，不能直接当成同等训练预算下的模型对决。
分子多目标优化（MuMO-Instruct）要求改结构、同时改善互相冲突的性质、还要跟原分子像。未适配的 2.6B 在 BDP 成功率只有 38.0。MMAI 多任务版把 BDP 拉到 86.0、BDQ 91.6、BPQ 98.8。三属性任务整体压过 GPT-5.5 的 83.2 / 91.6 / 92.6，也压过任务专精的 GeLLM³O。四属性 BDPQ 上 GPT-5.5 仍是 77.0，2.6B 只有 57.8，24B MoE 到 60.4，没追上。
官能团推理看的是「改掉某个基团，性质怎么变」。FGBench 上 2.6B 单任务版单基团准确率 0.840、数值 RMSE 55.95；GPT-5.5 是 0.772 和 187.4，70B Llama 准确率只有 0.683。开推理对数值题更有用，布尔题关掉推理几乎不掉点。
单步逆合成是反差最大的一项。底座 LFM2 在 URSA-expert-2026 上 unique 为 0，ChemCensor 全零。训完后 2.6B unique 94%、最高分 1.45；24B unique 70%、最高分 1.73，平均 Top-3 为 1.38，超过 Gemini 3 Flash 的 1.21。多样性明显高于 NatureLM 的 27% 和 RetroDFM-R 的 12%。
ADMET 上，2.6B 脂溶性 MAE 0.434，好过 27B TxGemma 的 0.538，也略好过 TDC 榜单专精模型的 0.456。药物性肝损伤 AUROC 0.937，压过 TxGemma 的 0.886。血脑屏障 0.924，与 TDC SOTA 持平。分布容积、肝微粒体清除率这些回归终点，任务专精模型仍领先。底座 2.6B 在脂溶性上 MAE 是 0.953，在血脑屏障上 AUROC 只有 0.366，领域训练把这两个数翻了过来。
| 任务 | 2.6B-MMAI | 对照 |
| MuMO BPQ 成功率 | 98.8 | GPT-5.5 92.6 |
| FGBench 单基团 Acc | 0.840 | GPT-5.5 0.772 |
| 脂溶性 MAE（越低越好） | 0.434 | TxGemma-27B 0.538 |
| URSA unique | 94% | 底座 0% |
| UniDock（CrossDocked） | -7.3 | GPT-5.5 -4.7 |
无条件 3D 生成里，2.6B 的键长、键角、二面角 JS 散度是 0.116 / 0.094 / 0.101，好过扩散模型 MolDiff 和化学 LLM nach0。口袋条件生成上，CrossDocked 的 UniDock 分数 -7.3，语言模型里最好，接近扩散模型 PocketXMol 的 -8.1；PoseBusters 分子内合法性只有 56%，GPT-5.5 是 95%，分子间口袋契合是 0%。对接分数好看，几何经常对不齐。
药化这种答案可自动打分的领域，2.6B 加上任务对齐的 SFT 和 GRPO，可以在若干公开基准上打过十倍大的通用模型和治疗向 27B 模型。拉开数字的是 MMAI Gym 的任务配方，不是把模型做大。2.6B 稠密版在多数 2D 任务上强过 24B MoE。
能直接拿来用的人很少。权重不公开下载，要走申请；训练用的提示、推理轨迹和内部数据也不放。论文把这写成双用途风险管控，对想复现的人就是硬墙。
权重和完整训练环境不公开，模型没做过对抗红队。评估全是计算基准，没有合成，也没有湿实验。口袋生成里对接分数高、几何约束差，论文归因于公开 3D 数据不够。文本坐标表示本身也很难稳住键角和立体化学。
对照不对称是方法学上最需要打折的地方。把 few-shot 的 GPT-5.5 当基线，能证明领域配方有效，证明不了「这个 2.6B 是新的药化 SOTA」。若干 TDC 终点上，非 LLM 专精模型仍明显更好。强领域适配还会伤底座的通用能力，论文把这写成刻意取舍。