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
Insilico and Liquid adapt LFM2 on 400+ chemistry tasks with SFT+GRPO. The 2.6B model beats GPT-5.5 and TxGemma-27B on optimization, FGBench, and several ADMET endpoints.
General-purpose LLMs are unreliable on drug-discovery work. Scaling the model or adding reasoning tokens barely moves the numbers. Specialist systems still win on ADMET, retrosynthesis, and 3D generation, but each one covers a narrow slice. There is no single assistant that can sit across medicinal chemistry, synthesis planning, and early DMPK.
Insilico Medicine and Liquid AI treat drug discovery as a pile of verifiable reasoning tasks. They build MMAI Gym, more than 400 tasks, then run supervised fine-tuning plus reinforcement learning on Liquid's LFM2. LFM2 is a hybrid: gated short convolutions for most layers, a few grouped-query attention layers for global mixing. Two sizes are adapted, a 2.6B dense model and a 24B MoE with about 2B active parameters.
The gym spans six buckets: 2D molecules, 3D molecules, 2D proteins, 3D proteins, drug-gene interactions, and cross-modal work. Sources include TDC, MOSES, FGBench, MuMO-Instruct, GEOM, ZINC, and CrossDocked2020, plus in-house data. Batches sample uniformly across the six buckets first, then pick a task, so large corpora cannot drown small ones.
Molecules are not locked to one string syntax. Inputs randomly convert among SMILES, SELFIES, and IUPAC names, and SMILES strings are rewritten with non-canonical atom orders. The tokenizer adds chemistry-specific tokens and wraps molecules in tags such as <smiles>CC(=O)O</smiles>. A 3D small molecule is a SMILES string followed by per-atom coordinates; a protein is laid out residue by residue with amino-acid and atom-name tokens. Hydrogens are dropped and rebuilt later. The aim is molecular concepts, not one serialization.
Training has two stages. SFT runs for about 100,000 steps at roughly 3.15 million tokens per step, AdamW at 1e-5. Online RFT uses GRPO: normalized absolute error for regression, a binary hit for classification, RDKit validity for generation, plus a <think> format reward and a thinking-length reward. Rewards share one scale; KL regularization sits at 0.4 so tasks do not fight. Generation tasks train for about 1,000 steps at temperature 1.4. They train both a multi-task generalist and a single-category specialist. The ADMET category already contains 22 property tasks, so "single-task" here means one family, not one endpoint.
Evaluation paraphrases the instruction, augments the molecule encoding, and samples several times. Regression takes the median, classification majority vote. Closed models are few-shot and never trained on MMAI Gym. That comparison tests whether the recipe works. It does not test equal training budgets.
On MuMO-Instruct, the model must edit a molecule, improve conflicting properties, and stay similar to the starting structure. The raw 2.6B scores 38.0 success on BDP. After MMAI multi-task training it reaches 86.0 / 91.6 / 98.8 on BDP / BDQ / BPQ, ahead of GPT-5.5 at 83.2 / 91.6 / 92.6 and ahead of the specialist GeLLM³O line. On the four-property BDPQ split GPT-5.5 still leads at 77.0; the 2.6B model is at 57.8 and the 24B MoE at 60.4.
FGBench asks what happens to a property after a named functional-group edit. The 2.6B single-category model hits 0.840 accuracy and 55.95 RMSE on single-group questions. GPT-5.5 sits at 0.772 and 187.4. Llama-3.1 70B is at 0.683 accuracy. Thinking helps numeric questions; boolean questions barely need it.
Single-step retrosynthesis is the sharpest before/after. Base LFM2 scores 0% unique valid reactants on URSA-expert-2026 and zeros on ChemCensor. After training, the 2.6B model reaches 94% unique with a max score of 1.45. The 24B model is at 70% unique, max 1.73, and average Top-3 1.38, above Gemini 3 Flash at 1.21. Diversity is far above NatureLM (27%) and RetroDFM-R (12%).
On TDC ADMET, 2.6B lipophilicity MAE is 0.434, better than TxGemma-27B at 0.538 and slightly better than the TDC specialist SOTA at 0.456. DILI AUROC is 0.937 versus TxGemma 0.886. BBB is 0.924, matching TDC SOTA. Volume of distribution and clearance still belong to specialist regressors. The untuned 2.6B had 0.953 MAE on lipophilicity and 0.366 AUROC on BBB.
| Task | 2.6B-MMAI | Baseline |
| MuMO BPQ success | 98.8 | GPT-5.5 92.6 |
| FGBench single-group Acc | 0.840 | GPT-5.5 0.772 |
| Lipophilicity MAE (lower better) | 0.434 | TxGemma-27B 0.538 |
| URSA unique | 94% | base 0% |
| UniDock (CrossDocked) | -7.3 | GPT-5.5 -4.7 |
For unconditional 3D generation, 2.6B JS divergences on bond lengths, angles, and dihedrals are 0.116 / 0.094 / 0.101, tighter than MolDiff and nach0. In pocket-conditioned generation the CrossDocked UniDock score is -7.3, the best among language models and close to PocketXMol at -8.1. PoseBusters intramolecular validity is only 56% against GPT-5.5 at 95%, and intermolecular pocket fit is 0%. Docking looks good. Geometry often does not.
On tasks with automatic graders, a 2.6B model plus a chemistry-shaped SFT+GRPO recipe can beat ten-times-larger general models and a 27B therapeutic LLM on several public boards. The gym, not the parameter count, is what moves the numbers. The dense 2.6B even beats the 24B MoE on most 2D tasks.
Almost nobody can use the weights. There is no public download, only an access form. Prompts, reasoning traces, and in-house data stay private. The paper frames this as dual-use control. For anyone who wants to reproduce the system, it is a wall.
Checkpoints and the full training environment are withheld. The models have not been red-teamed. All numbers are computational; nothing was synthesized or assayed. Pocket generation scores well on docking and poorly on geometric constraints. The authors blame scarce public 3D data. Serializing coordinates as text is also a hard way to keep bond angles and stereochemistry honest.
The comparison is lopsided. Few-shot GPT-5.5 is a fair test of whether domain data helps, not a fair test of whether this is the new chemistry SOTA. Specialist non-LLM models still win several TDC endpoints. Heavy domain adaptation also dulls the base model's general skills, which the paper treats as an accepted trade.