基于强化学习的化学大模型后训练：提升药物分子生成质量

1. 项目缘起：当大模型遇上药物设计

最近几年，大语言模型（LLM）在文本生成、代码编写等领域展现出的惊人能力，让很多领域的研究者都在思考一个问题：能不能把这种能力“嫁接”到我们自己的专业领域？作为一名长期关注AI在药物研发领域应用的从业者，我自然也在琢磨这件事。药物设计，特别是小分子药物的发现，是一个典型的“大海捞针”过程。化学空间极其庞大，理论上可能存在的小分子数量远超宇宙中的原子总数。传统的基于经验、基于规则或者基于分子对接的虚拟筛选方法，虽然有效，但往往效率有限，且容易陷入局部最优。

于是，一个很自然的想法出现了：能否训练一个专门针对化学分子“语言”的LLM，让它像写文章一样，“生成”出具有理想性质的候选药物分子？这个想法并不新鲜，已经有不少研究团队在尝试。然而，我发现在实践中，直接使用预训练好的化学分子生成模型，往往存在一个共性问题：模型生成的分子在结构上可能很新颖、很合理，但在关键的药物属性上，比如与靶点蛋白的结合亲和力、类药性、合成可行性等方面，表现并不稳定，甚至很差。这就像让一个精通语法但不懂药理的人去设计药物，他能写出语法正确的“分子句子”，但句子本身可能毫无疗效，甚至有毒。

这正是“后训练”需要介入的关键点。预训练让模型学会了化学的“语法”（原子、键、官能团的连接规则），但要让模型真正理解“药理学语义”（什么样的结构对应什么样的活性、安全性），就需要更精细的引导。而强化学习，恰恰提供了一种将“语义目标”直接反馈给模型进行优化的强大框架。我最近完成的一个项目，核心就是探索如何利用强化学习对化学领域的LLM进行后训练，显著提升其在面向特定靶点的药物设计任务中的“命中率”。这不仅仅是调几个参数，而是为模型注入“设计意图”的过程。

2. 核心架构拆解：从分子生成器到“有目标的探索者”

要实现基于强化学习的LLM后训练，我们需要构建一个完整的交互式学习循环。这个架构可以理解为将一个原本“自由创作”的分子生成模型，改造成一个在特定“目标”指导下进行“定向探索”的智能体。整个系统主要由四个核心组件构成：智能体（Agent）、环境（Environment）、奖励函数（Reward Function）和策略（Policy）。下面我们来逐一拆解它们在药物设计场景下的具体形态。

2.1 智能体：化学分子LLM

智能体就是我们后训练的对象——一个经过化学领域预训练的大语言模型。这里的关键在于模型的“输入-输出”形式。我们通常采用基于SMILES（简化分子线性输入系统）字符串的模型。SMILES用一串ASCII字符唯一地表示一个分子的二维结构，非常适合作为LLM的“语言”。

• 模型选择：可以选择专门为化学设计的预训练模型，如ChemBERTa、MolT5，或者使用通用的Transformer架构（如GPT-2）在大量SMILES字符串上从头预训练。在我们的项目中，我们基于一个开源的GPT-2架构，在ZINC、ChEMBL等大型化合物数据库的SMILES序列上进行了预训练，让模型学会了生成语法上基本正确的分子。
• 智能体的动作：在强化学习框架下，模型生成一个完整的SMILES字符串的过程，被分解为一系列“动作”。每个动作就是在当前序列的末尾，从词汇表（包含原子、键类型、环的开闭符号等）中选择下一个token（字符）。因此，生成一个分子就是执行一个动作序列（一个Episode）。

2.2 环境：分子评估模拟器

环境负责接收智能体生成的分子（SMILES字符串），并返回一个对该分子的“评价”。在真实的药物研发中，这个评价需要经过复杂的生物实验。在计算模拟中，我们用一个计算评估管道来模拟环境。

这个管道通常包含多个步骤：

1. 语法与合法性检查：首先检查生成的SMILES是否能被正确解析为一个有效的、在化学上合理的分子结构。无效分子直接给予极低的负奖励。
2. 属性计算：对有效分子，调用一系列计算工具来预测其关键性质。这通常包括：
   • 类药性：使用如RDKit计算QED（定量估计类药性）、Lipinski五规则等。
   • 合成可行性：使用如SA Score（合成可及性分数）来评估分子合成的难易程度。
   • 靶点亲和力：这是最核心的一环。对于特定靶点蛋白，使用分子对接软件（如AutoDock Vina,GNINA）或更快速的机器学习打分函数，来预测生成分子与靶点的结合自由能（docking score）。分数越低（负值越大），通常意味着结合可能越强。
3. 多样性惩罚：为了避免模型陷入“模式坍塌”，反复生成少数几个高分分子，环境还需要评估生成分子与已有分子库或近期生成分子之间的相似性（如基于分子指纹的Tanimoto系数），对过于相似的分子进行奖励惩罚。

2.3 奖励函数：定义“好药物”的指挥棒

奖励函数是强化学习成功的灵魂，它直接将我们的设计目标——找到一个高活性、易合成、类药的分子——量化为一个标量数值。设计一个好的奖励函数需要权衡多个目标，通常采用加权和的形式：

R(molecule) = w1 * R_docking + w2 * R_qed + w3 * R_sa + w4 * R_novelty + R_penalty

• R_docking: 对接得分的转化。例如，将对接得分（负值）归一化到[0,1]区间，或使用Sigmoid函数进行映射。得分越好，奖励越高。
• R_qed和R_sa: 类药性和合成可行性分数，本身就在[0,1]之间，值越高越好。
• R_novelty: 新颖性奖励，鼓励生成与已知活性分子结构差异较大的分子。
• R_penalty: 惩罚项，用于处理无效SMILES、违反关键化学规则（如存在反应性过高的官能团）等情况。
• w1, w2, w3, w4: 权重系数。这是需要反复调试的艺术。初期可以更侧重R_docking以快速提升活性，后期可以增加R_sa和R_novelty的权重以优化其他属性。

注意：奖励函数的形状（而非绝对值）对训练稳定性影响巨大。实践中，我们经常会对每个奖励分量进行标准化处理（例如，减去均值除以标准差），或者使用PPO等算法内置的奖励归一化机制，以防止某个奖励分量主导更新过程导致训练崩溃。

2.4 策略优化：PPO算法与梯度更新

策略就是智能体根据当前状态（已生成的SMILES前缀）选择下一个token的概率分布。我们的目标是通过与环境的交互，优化这个策略，使其生成高奖励分子的概率最大化。

我们选择近端策略优化（PPO）算法作为优化器，这是目前与LLM结合最成熟、最稳定的强化学习算法之一。其核心优势在于通过“裁剪”机制限制了每次策略更新的幅度，避免了训练中的剧烈震荡，非常适合像LLM这样参数庞大的模型。

具体到我们的LLM，策略优化过程可以概括为：

1. 采样：用当前的分子生成模型（策略）生成一批分子（例如，1024个SMILES字符串）。
2. 评估：将这批分子送入环境，计算每个分子对应的奖励R。
3. 计算优势：使用广义优势估计（GAE）等方法，计算每个生成步骤（每个token）的“优势值”A。A衡量了在某个状态下采取某个动作，比平均情况好多少。
4. 构造损失函数：PPO的损失函数包含三部分：
   • 策略损失：鼓励增加高优势值动作的概率，抑制低优势值动作的概率，并通过裁剪防止更新过大。
   • 价值函数损失：训练一个价值函数网络（通常是一个小的MLP，以模型隐藏状态为输入）来估计状态的价值，用于更准确地计算优势。
   • 熵奖励：在损失中加入策略熵的奖励，鼓励探索，防止策略过早收敛到单一模式。
5. 反向传播与更新：计算总损失，通过反向传播更新LLM的策略网络参数（以及价值函数网络参数）。

这个过程反复迭代。随着训练的进行，你会观察到模型生成的分子平均奖励逐渐上升，这意味着模型越来越倾向于生成符合我们多目标优化要求的候选药物分子。

3. 实操流程与工程化细节

理论清晰后，落地实施是关键。下面我将以我们项目的具体实践为例，拆解从环境搭建到训练调优的全流程，并分享其中遇到的坑和解决方案。

3.1 环境搭建与数据准备

基础框架选择：我们使用PyTorch作为深度学习框架。对于强化学习部分，虽然可以手动实现PPO，但强烈建议使用成熟的RL库，如Stable-Baselines3或Ray RLlib。我们最终选择了Ray RLlib，因为它对分布式训练支持更好，且与自定义环境、大模型的集成相对灵活。

分子评估环境封装：这是工程的重点。我们需要将前述的“计算评估管道”封装成一个符合OpenAI Gym或Farama Foundation Gymnasium接口的环境。

import gym from gym import spaces import numpy as np from rdkit import Chem from rdkit.Chem import QED, Crippen import subprocess # 用于调用外部对接软件 class DrugDesignEnv(gym.Env): def __init__(self, target_pdb_path, reward_weights): super(DrugDesignEnv, self).__init__() # 动作空间：词汇表大小 self.action_space = spaces.Discrete(VOCAB_SIZE) # 状态空间：可以用当前生成的token ID序列表示，这里简化处理 self.observation_space = spaces.Box(low=0, high=VOCAB_SIZE, shape=(MAX_SEQ_LEN,), dtype=np.int32) self.target_pdb = target_pdb_path self.reward_weights = reward_weights self.current_smiles = "" self.step_count = 0 def reset(self, seed=None, options=None): # 重置环境，开始一个新的分子生成 self.current_smiles = "" self.step_count = 0 # 返回初始状态，例如一个特殊的开始符`[START]`的编码 return np.array([START_TOKEN_ID]), {} def step(self, action): # action是一个token ID token = id_to_token[action] self.current_smiles += token self.step_count += 1 terminated = False truncated = False reward = 0.0 # 判断是否生成结束（遇到结束符或达到最大长度） if token == END_TOKEN or self.step_count >= MAX_SEQ_LEN: terminated = True # 分子生成完毕，计算最终奖励 mol = Chem.MolFromSmiles(self.current_smiles) if mol is None: # 无效分子 reward = -10.0 else: # 计算各项分数 qed_score = QED.qed(mol) sa_score = calculate_sa_score(mol) # 需要实现 docking_score = run_docking(mol, self.target_pdb) # 调用外部对接，需要实现 novelty = calculate_novelty(mol) # 需要实现 # 加权合成总奖励 reward = (self.reward_weights['w1'] * transform_docking_score(docking_score) + self.reward_weights['w2'] * qed_score + self.reward_weights['w3'] * (1 - sa_score) + # SA分数越低越好，需转换 self.reward_weights['w4'] * novelty) else: # 在生成过程中，可以给予稀疏奖励或零奖励 reward = 0.0 # 构建新的状态（当前已生成的序列） next_state = encode_smiles(self.current_smiles) return next_state, reward, terminated, truncated, {} # ... 其他辅助函数，如 run_docking, calculate_sa_score 等

关键依赖安装：

• RDKit：化学信息学核心库，用于分子处理、描述符计算。
• OpenBabel或PyMol：用于分子格式转换，为对接准备配体文件。
• 分子对接软件：如AutoDock Vina，需要单独安装并确保命令行可调用。对于大规模生成，对接是计算瓶颈，可以考虑使用GPU加速的对接程序（如GNINA）或预先训练好的快速打分神经网络。

3.2 训练循环与关键参数配置

将自定义环境与RLlib和我们的LLM模型整合起来，是另一个工程难点。RLlib通常期望智能体是一个相对简单的神经网络，而我们的LLM则非常庞大。我们需要定义一个自定义的TorchModel，并将其包装为RLlib可以识别的策略。

from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.utils.annotations import override import torch.nn as nn class LLMAsTorchModel(TorchModelV2, nn.Module): def __init__(self, obs_space, action_space, num_outputs, model_config, name): TorchModelV2.__init__(self, obs_space, action_space, num_outputs, model_config, name) nn.Module.__init__(self) # 加载我们预训练好的化学LLM self.llm = load_pretrained_chem_llm() # 价值函数头：一个简单的线性层，输入是LLM最后一个隐藏层的[CLS] token或平均池化 self.value_head = nn.Linear(self.llm.config.hidden_size, 1) @override(TorchModelV2) def forward(self, input_dict, state, seq_lens): # input_dict["obs"] 是token ID序列 obs = input_dict["obs"] # 通过LLM获取最后一层的隐藏状态 llm_outputs = self.llm(input_ids=obs, output_hidden_states=True) last_hidden_state = llm_outputs.hidden_states[-1] # 取[CLS] token的状态作为策略和价值的输入 state_representation = last_hidden_state[:, 0, :] # 策略logits：通过一个线性层映射到动作空间（词汇表大小） logits = self.policy_head(state_representation) # 价值估计 value = self.value_head(state_representation).squeeze(-1) return logits, state @override(TorchModelV2) def value_function(self): return self._value

配置训练时，有几个参数对稳定性和效果至关重要：

# 一个简化的RLlib PPO配置示例 algo_config: algorithm: PPO env: DrugDesignEnv model: custom_model: LLMAsTorchModel lr: 1e-6 # 学习率必须非常小！LLM微调的标准操作。 train_batch_size: 4000 # 较大的批次有助于稳定训练 sgd_minibatch_size: 500 num_sgd_iter: 10 # 每次采样后的优化迭代次数 clip_param: 0.2 # PPO裁剪参数 vf_clip_param: 10.0 # 价值函数裁剪 entropy_coeff: 0.01 # 熵系数，鼓励探索 gamma: 0.99 # 折扣因子 lambda: 0.95 # GAE参数 kl_coeff: 0.0 # 通常PPO不使用KL惩罚，裁剪机制已足够 num_workers: 4 # 并行环境工作者数量，加速采样 num_gpus: 1 # 使用GPU framework: torch

启动训练后，你需要密切监控几个关键指标：

1. episode_reward_mean：平均回合奖励，这是最直接的优化目标，应该呈现上升趋势。
2. policy_loss和vf_loss：策略损失和价值损失，应该逐渐收敛并保持较小波动。
3. entropy：策略熵，初期应保持一定水平，随着训练会缓慢下降，但不应骤降至零（模式坍塌）。
4. kl：新旧策略之间的KL散度，PPO算法会将其控制在一定范围内。

3.3 奖励塑形与课程学习：让训练更平滑

直接使用最终对接得分作为奖励，往往因为信号过于稀疏（只有分子生成完才有奖励）和噪声大（对接打分本身有误差）而导致训练困难。这里分享两个非常有效的技巧：

奖励塑形：在分子生成过程中提供中间奖励。例如，当模型生成了一个有效的子结构（如某个特定的药效团），即使分子未完成，也可以给予一个小额正奖励。这能更及时地引导模型。在我们的实现中，我们定义了一个“子结构匹配”函数，在每一步检查当前SMILES是否包含已知活性片段，如果包含，则给予一个小的正向奖励增量。

课程学习：不要一开始就让模型挑战高难度的多目标优化。可以设计一个由易到难的训练课程：

• 阶段一：奖励函数主要奖励生成有效的SMILES字符串（语法正确）。权重w1（对接）可以设为零或很小。
• 阶段二：在模型能稳定生成有效分子后，逐步引入类药性（QED）和合成可行性（SA）的奖励。
• 阶段三：最后，再显著提高对接得分的权重w1，让模型在保持前两个属性的基础上，优化对靶点的亲和力。

这种分阶段训练的策略，能极大提高训练的稳定性和最终效果。我们通过一个外部的调度器，每训练一定步数，就动态调整reward_weights字典中的权重值，实现了自动化课程学习。

4. 效果评估、常见问题与优化策略

训练完成后，我们如何判断模型是否真的变“聪明”了？除了看训练曲线，更重要的是对模型生成的结果进行离线评估。

4.1 多维度评估指标体系

我们从以下几个维度对后训练前后的模型进行对比评估：

1. 生成质量：

   • 有效性：随机生成一批分子，计算能被RDKit成功解析的比例。预训练后通常能达到95%以上，强化学习后训练不应显著降低此比例。
   • 独特性：在生成的一批分子中，唯一分子的比例。避免模型陷入“模式坍塌”，反复生成同一个或几个分子。
   • 新颖性：生成的分子与训练数据（如ZINC库）或已知活性分子之间的最大相似度（Tanimoto系数）。值越低，新颖性越高。

2. 药物属性：

   • 类药性分布：统计生成分子的QED分数分布，并与已知药物（如ChEMBL中的分子）的分布进行比较。理想情况下，分布应向右偏移（更高QED）。
   • 合成可行性分布：统计SA Score，理想情况下应向左偏移（更低SA，更易合成）。
   • 理化性质：计算分子量、脂水分配系数（LogP）、氢键供受体数等，并绘制在“药物化学空间”中，看是否更集中于类药区域。

3. 靶向活性（核心）：

   • 对接得分分布：对生成分子进行批量对接，计算其对接得分的分布。与随机从ZINC库中抽取的分子、或者与仅做预训练的模型生成的分子进行对比。我们期望后训练模型生成的分子，其对接得分的平均值更低（结合更好），且高分（低对接值）分子的比例显著提升。
   • 虚拟筛选富集率：这是一个更严格的测试。假设我们有一个包含已知活性分子和诱饵分子的测试集。用模型生成一批分子，与测试集混合，然后根据对接得分排序。计算在前1%、5%等位置，活性分子被富集的比例（EF值）。EF值越高，说明模型“定向设计”的能力越强。

在我们的实验中，经过RL后训练的模型，在针对某个激酶靶点的任务中，其生成分子中对接打分优于-10 kcal/mol的比例，从预训练模型的约2%提升到了15%以上，且这些高分分子在QED和SA Score上依然保持良好水平。

4.2 实战中踩过的坑与解决方案

坑一：奖励函数设计不当导致训练崩溃

• 现象：训练初期，episode_reward_mean急剧下降至很大的负值，然后不再变化，模型“学废了”，只生成无效字符串。
• 根因：无效分子的惩罚（R_penalty）设置得过于严苛（例如-50），而有效分子的正奖励范围在[0,1]。模型很快发现，只要生成一个无效分子，就会获得巨大的负奖励，远高于辛辛苦苦生成一个有效分子可能得到的微小正奖励的期望值。因此，最优策略变成了“尽快结束回合并接受固定惩罚”，即生成一个非法字符立刻终止。
• 解决方案：平衡奖励尺度。确保生成一个“中等偏上”的有效分子所获得的正奖励，其绝对值显著大于生成无效分子的惩罚。可以将惩罚设置为一个相对较小的负值（如-1），同时调整有效分子各项奖励的权重和转换函数，使优秀分子的总奖励能达到+5或+10。也可以采用奖励标准化，让每个批次的奖励均值为0，方差为1，这是RL中的常用稳定技巧。

坑二：模型“遗忘”预训练知识

• 现象：训练后，模型生成的分子有效性、多样性暴跌，变得语法不通，仿佛忘记了化学规则。
• 根因：强化学习的梯度更新强度太大，或者学习率太高，覆盖了模型在预训练阶段学到的化学语言模型基础分布。这被称为“灾难性遗忘”。
• 解决方案：
  • 大幅降低学习率：LLM的微调学习率通常在1e-6到1e-5量级，RL后训练时可能需要更低（如5e-7）。
  • 在损失函数中加入KL散度惩罚：除了PPO的裁剪，可以显式地在损失中加入当前策略与原始预训练模型策略之间的KL散度惩罚项，强制新策略不要偏离原始策略太远。这相当于给模型一个“正则化”，让它不忘本。
  • 混合训练：在每个训练批次中，不仅包含RL采样得到的数据，也混入一部分预训练数据（从ZINC等库中采样SMILES），让模型同时进行传统的语言模型损失（交叉熵）和RL策略损失的计算。

坑三：计算瓶颈与效率优化

• 现象：80%以上的时间花在分子对接上，训练速度极慢。
• 解决方案：
  • 并行化：充分利用RLlib的多worker架构，每个worker独立运行环境实例，并行进行分子生成和评估。
  • 缓存：对生成的SMILES字符串进行哈希，建立奖励缓存。如果同一个分子被多次生成（在探索过程中很常见），直接返回缓存中的奖励值，避免重复计算。
  • 使用代理模型：训练一个快速的神经网络（如图神经网络GNN）来近似对接打分函数。先用对接软件计算一部分分子作为训练集，训练这个代理模型。在RL环境评估时，用代理模型预测打分，其速度比物理对接快几个数量级。定期用真实对接软件更新代理模型的训练数据。这是工业级应用中的常见做法。

4.3 超越PPO：更前沿的探索方向

当基本流程跑通后，可以尝试一些更前沿的优化，以进一步提升性能或效率：

• 离线强化学习：如果我们已经拥有一个包含分子及其属性（对接分数、QED等）的现有数据集，可以尝试离线RL算法（如CQL, IQL）。这些算法可以从静态数据集中学习，而不需要昂贵的在线交互，更适合利用历史实验数据。
• 分层强化学习：将分子生成过程分为两层。高层策略决定分子的宏观结构特征（如骨架类型、主要官能团），低层策略负责在高层指定的框架下进行原子和键的细节生成。这可以更好地控制生成分子的结构多样性，并可能加速探索。
• 多目标优化：使用基于帕累托前沿的多目标RL算法，如MO-PPO。这样可以直接优化多个目标的帕累托最优解集，而不是通过手动调权重将其合并为单一目标，能更系统地探索活性、类药性、可合成性之间的权衡空间。

这个项目让我深刻体会到，将LLM与强化学习结合用于药物设计，不是一个简单的“调包”过程，而是一个涉及深度学习、强化学习、计算化学和工程优化的系统性工程。从奖励函数设计的“艺术”，到稳定训练的“技巧”，再到工程效率的“优化”，每一步都需要细致的考量和反复的迭代。但看到模型从漫无目的地“造句”，到后来能稳定地“写出”具有潜力的药物分子草图时，那种成就感是巨大的。这不仅仅是AI能力的提升，更是为我们提供了一种全新的、数据驱动的分子发现范式，其探索化学空间的效率和指向性，是传统方法难以比拟的。