基于CVAE与Transformer的多约束条件AI分子生成技术实践
1. 项目概述当AI成为“药物化学家”最近几年AI在药物研发领域的渗透速度远超预期。从早期的虚拟筛选辅助到如今直接参与分子设计与优化AI正从一个“计算工具”演变为一个具备初步“化学直觉”的“合作者”。我接触这个领域也有几年了从最初用现成的模型做性质预测到现在尝试构建端到端的分子生成优化流程踩过的坑不少但收获的惊喜更多。今天想和大家深入聊聊的就是“AI药物分子优化”中的一个核心且极具挑战性的环节如何从一个简单的分子线性表示SMILES出发生成出同时满足多个、甚至相互冲突的约束条件的新分子。这听起来像是个“既要、又要、还要”的难题。在真实的药物研发项目中一个理想的候选分子需要同时具备高活性对靶点蛋白有效、良好的成药性比如合适的溶解度、代谢稳定性、低毒性、以及可合成性化学家能在实验室里以合理的成本合成出来。传统的方法往往是一个一个条件去筛选耗时费力且容易陷入局部最优。而基于深度学习的生成模型为我们提供了一种全新的思路将所有这些约束条件“编码”进模型的训练或生成过程中让AI自己去探索浩瀚的化学空间寻找那个能满足所有苛刻条件的“全能选手”。这个过程的核心输入是SMILESSimplified Molecular Input Line Entry System它是一种用ASCII字符串精确描述分子结构的语言。比如阿司匹林就是“CC(O)Oc1ccccc1C(O)O”。对AI来说一个SMILES字符串就是一个“句子”分子生成就变成了“写句子”的任务。而“多约束条件生成”就是要求AI写出既符合语法化学价键规则、又满足特定主题和要求如“logP在2到3之间”、“不含警示结构”的“好句子”。接下来我会结合自己的实践经验拆解从SMILES处理到构建多约束条件生成模型的完整链条分享其中的关键设计、实操细节以及那些只有亲手做过才会知道的“坑”。2. 核心思路与方案选型为什么是“条件生成”当我们面对“多约束条件生成”这个需求时第一个要回答的问题是用什么模型架构主流的分子生成模型如RNN循环神经网络、Transformer、VAE变分自编码器和GAN生成对抗网络各有优劣。但针对“条件生成”我们需要模型不仅能学会化学空间的分布还要能接受外部条件的指导。2.1 主流模型架构的横向对比在实践中我主要对比和尝试过以下几种思路条件变分自编码器CVAE这是最直观的思路之一。在标准的VAE中编码器将分子SMILES压缩成一个连续的低维潜向量latent vector解码器再从这个向量重建分子。CVAE则在这个潜向量中拼接concatenate上条件向量比如我们希望的血浆蛋白结合率值。这样解码器在重建或生成时就会同时受到潜向量和条件向量的影响。优点概念清晰训练相对稳定生成分子的多样性好。缺点对条件控制的精确性有时不足尤其是在条件维度较高或条件间存在复杂关系时。生成的分子可能在“满足条件A”和“满足条件B”之间摇摆。强化学习RL微调先训练一个无条件的生成模型例如一个GPT风格的Transformer作为“策略网络”。然后我们定义一系列“奖励函数”每个函数对应一个约束条件例如预测的活性值越高奖励越大预测的毒性值越高惩罚越大。通过策略梯度等RL方法引导模型生成能获得高总奖励的分子。优点非常灵活可以处理极其复杂、非可微的奖励函数比如基于规则的结构过滤器。缺点训练不稳定需要精心设计奖励函数和平衡权重否则模型容易崩溃或陷入单一模式。计算成本也较高。条件Transformer如MolGPT, Chemformer在Transformer的解码过程中将条件信息作为额外的输入。例如在生成每个token之前将条件向量与当前的上下文表示进行融合。这类似于在文本生成中给定“写作风格”或“主题”。优点得益于Transformer强大的序列建模能力对长程依赖和复杂条件建模效果好。缺点需要大量的训练数据且条件信息的注入方式需要精巧设计否则容易被模型忽略。流模型Flow-based Models与扩散模型Diffusion Models这是近两年的新热点。它们通过定义可逆的变换将简单的噪声分布转化为复杂的分子分布。条件信息可以在变换过程中被引入。优点理论上能建模更复杂的分布生成质量高。缺点模型复杂训练和推理速度较慢在分子生成领域的应用还不够成熟工具链不如前几种丰富。实操心得对于大多数从零开始的团队或项目我强烈建议从CVAE或条件Transformer入手。它们有相对成熟的代码库如DeepChem、PyTorch Geometric、Hugging Face Transformers社区支持好更容易快速搭建原型并看到效果。RL方法更像是一把“手术刀”适合在已有不错的基础模型上进行精细的、目标明确的优化但不适合作为起点。2.2 多约束条件的融合策略加权求和还是分层控制确定了模型骨架下一个难题是如何处理“多”个约束条件。假设我们有四个条件活性pIC50 8、溶解性LogS -4、类药性QED 0.6、合成可及性SA Score 4.5。我们如何让模型同时考虑它们加权求和法最常见为每个条件计算一个得分或损失然后乘以一个权重后相加得到总得分或总损失。总得分 w1 * 活性得分 w2 * 溶解性得分 w3 * QED得分 - w4 * SA得分。关键权重的设置是艺术也是科学。权重过大模型可能只优化单一指标权重过小该条件可能被忽略。我通常的做法是先进行单目标优化观察每个目标单独能带来的得分变化范围然后根据业务重要性手动调整权重使各个目标在总损失中的贡献量级相当。这是一个需要反复实验的过程。分层/序贯控制法先让模型满足最核心的1-2个条件如活性和类药性生成一批候选分子。然后用这些分子作为起点在第二个阶段引入其他条件如溶解性和合成难度进行进一步优化或筛选。这可以通过串联不同的模型或在RL框架中分阶段调整奖励函数来实现。适用场景当约束条件之间存在明显的优先级或因果关系时。例如必须先有活性再谈其他。帕累托优化思想不追求一个在所有指标上都“最优”的分子这种分子往往不存在而是寻找一批“非支配”分子。即在这批分子中你无法找到一个分子在所有指标上都比另一个分子好。这更适合在生成大量分子后进行后处理和分析时使用。踩坑记录早期我曾试图让模型同时优化5个以上的约束结果生成的分子要么结构怪异要么干脆无法解析。后来意识到约束不是越多越好。有些强相关的约束如LogP和溶解性可以合并考虑有些可以通过后处理过滤如排除含有硝基的分子。在模型生成阶段聚焦于3-4个最关键、最难以通过规则后处理实现的约束往往能取得更好的效果。3. 从数据到模型构建端到端流程理论说再多不如一行代码。下面我以一个简化但完整的流程为例展示如何构建一个基于CVAE的多条件分子生成器。3.1 数据准备与SMILES标准化数据质量决定模型天花板。我们通常从ChEMBL、PubChem等公开数据库获取分子及其生物活性数据。import pandas as pd from rdkit import Chem from rdkit.Chem import Descriptors, Crippen, QED from rdkit.Chem.Scaffolds import MurckoScaffold import numpy as np # 1. 加载数据 df pd.read_csv(chembl_akt1_activity.csv) # 假设有Akt1靶点的活性数据 df[SMILES] df[SMILES].astype(str) # 2. SMILES标准化与清洗 def standardize_smiles(smi): try: mol Chem.MolFromSmiles(smi) if mol is None: return None # 去盐、标准化价态、生成规范SMILES mol Chem.RemoveHs(mol) # 移除氢原子简化表示 return Chem.MolToSmiles(mol, isomericSmilesTrue) except: return None df[canonical_smiles] df[SMILES].apply(standardize_smiles) df df.dropna(subset[canonical_smiles]).drop_duplicates(subset[canonical_smiles]) # 3. 计算条件属性标签 def calculate_properties(smiles): mol Chem.MolFromSmiles(smiles) if mol is None: return None props {} props[mw] Descriptors.MolWt(mol) # 分子量 props[logp] Crippen.MolLogP(mol) # 脂水分配系数 props[qed] QED.qed(mol) # 类药性 props[tpsa] Descriptors.TPSA(mol) # 极性表面积 # 这里可以添加更复杂的预测模型如用预训练模型预测pIC50 return props property_list [] for smi in df[canonical_smiles]: prop calculate_properties(smi) property_list.append(prop) prop_df pd.DataFrame(property_list) df pd.concat([df.reset_index(dropTrue), prop_df], axis1) # 4. 定义条件目标 # 假设我们希望分子量500, logP在2-3之间 QED0.6 df[condition_1] (df[mw] 500).astype(int) df[condition_2] ((df[logp] 2) (df[logp] 3)).astype(int) df[condition_3] (df[qed] 0.6).astype(int) # 条件向量可以是一个多热编码也可以是一个连续值如logP的具体值 df[condition_vector] df[[condition_1, condition_2, condition_3]].values.tolist() print(f清洗后数据量{len(df)}) print(df[[canonical_smiles, mw, logp, qed, condition_vector]].head())注意事项标准化至关重要同一个分子可能有多个有效的SMILES表示。不进行标准化模型会认为“CC(O)Oc1ccccc1C(O)O”和“OC(O)c1ccccc1OC(O)C”是两个不同的分子严重干扰学习。处理无效分子RDKit的MolFromSmiles可能失败约1-5%必须过滤掉这些无效数据。条件的设计是将条件设为二分类标签是/否还是连续值具体数值取决于你的生成目标。对于“优化”连续值能提供更细粒度的指导对于“筛选”二分类标签更简单。3.2 构建条件变分自编码器CVAE我们将使用PyTorch构建一个简单的CVAE。核心思想是编码器将SMILES和条件向量一起编码成潜向量z的均值和方差解码器从z和条件向量重建SMILES。import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader # 首先需要构建词汇表并将SMILES转化为索引序列 from collections import Counter all_smiles df[canonical_smiles].tolist() tokens [] for smi in all_smiles: # 简单的字符级分词更优的方案是使用基于子词BPE的分词 tokens.extend(list(smi)) vocab Counter(tokens) # 添加特殊令牌 special_tokens [pad, sos, eos, unk] vocab special_tokens sorted([token for token in vocab.keys()]) token_to_idx {token: i for i, token in enumerate(vocab)} idx_to_token {i: token for i, token in enumerate(vocab)} vocab_size len(vocab) class SMILESDataset(Dataset): def __init__(self, df, token_to_idx, max_len100): self.df df self.token_to_idx token_to_idx self.max_len max_len def __len__(self): return len(self.df) def __getitem__(self, idx): smi self.df.iloc[idx][canonical_smiles] # 将SMILES转化为索引序列并添加起止符 seq [token_to_idx[sos]] [token_to_idx.get(c, token_to_idx[unk]) for c in smi] [token_to_idx[eos]] seq seq[:self.max_len] # 填充 if len(seq) self.max_len: seq seq [token_to_idx[pad]] * (self.max_len - len(seq)) seq torch.tensor(seq, dtypetorch.long) # 条件向量 cond torch.tensor(self.df.iloc[idx][condition_vector], dtypetorch.float) return seq, cond # 定义CVAE模型 class CVAE(nn.Module): def __init__(self, vocab_size, embedding_dim, hidden_dim, latent_dim, condition_dim, max_len): super().__init__() self.vocab_size vocab_size self.max_len max_len self.embedding nn.Embedding(vocab_size, embedding_dim, padding_idxtoken_to_idx[pad]) # 编码器Bi-LSTM self.encoder_lstm nn.LSTM(embedding_dim condition_dim, hidden_dim, bidirectionalTrue, batch_firstTrue) self.hidden2mu nn.Linear(hidden_dim * 2, latent_dim) # 双向LSTM输出维度是hidden_dim*2 self.hidden2logvar nn.Linear(hidden_dim * 2, latent_dim) # 解码器LSTM self.latent2hidden nn.Linear(latent_dim condition_dim, hidden_dim) self.decoder_lstm nn.LSTM(embedding_dim condition_dim, hidden_dim, batch_firstTrue) self.outputs2vocab nn.Linear(hidden_dim, vocab_size) def encode(self, x, cond): # x: (batch, seq_len), cond: (batch, cond_dim) batch_size x.size(0) embedded self.embedding(x) # (batch, seq_len, emb_dim) # 将条件向量扩展到每个时间步 cond_expanded cond.unsqueeze(1).expand(-1, embedded.size(1), -1) # (batch, seq_len, cond_dim) encoder_input torch.cat([embedded, cond_expanded], dim-1) _, (hidden, _) self.encoder_lstm(encoder_input) # hidden: (2, batch, hidden_dim) - (batch, 2*hidden_dim) hidden hidden.permute(1, 0, 2).contiguous().view(batch_size, -1) mu self.hidden2mu(hidden) logvar self.hidden2logvar(hidden) return mu, logvar def reparameterize(self, mu, logvar): std torch.exp(0.5 * logvar) eps torch.randn_like(std) return mu eps * std def decode(self, z, cond, target_seqNone, teacher_forcing_ratio0.5): # z: (batch, latent_dim), cond: (batch, cond_dim) batch_size z.size(0) # 初始化解码器隐藏状态 decoder_input torch.full((batch_size, 1), token_to_idx[sos], devicez.device, dtypetorch.long) decoder_hidden self.latent2hidden(torch.cat([z, cond], dim-1)).unsqueeze(0) # (1, batch, hidden_dim) decoder_cell torch.zeros_like(decoder_hidden) outputs [] for t in range(self.max_len - 1): # 最多生成max_len-1个token不包括sos embedded self.embedding(decoder_input) # (batch, 1, emb_dim) cond_expanded cond.unsqueeze(1) # (batch, 1, cond_dim) decoder_input_step torch.cat([embedded, cond_expanded], dim-1) lstm_out, (decoder_hidden, decoder_cell) self.decoder_lstm(decoder_input_step, (decoder_hidden, decoder_cell)) output self.outputs2vocab(lstm_out.squeeze(1)) # (batch, vocab_size) outputs.append(output.unsqueeze(1)) # (batch, 1, vocab_size) # 决定下一个输入是真实值还是预测值 teacher_force torch.rand(1).item() teacher_forcing_ratio if teacher_force and target_seq is not None: decoder_input target_seq[:, t].unsqueeze(1) # 使用真实的下一个token else: top1 output.argmax(1) # (batch,) decoder_input top1.unsqueeze(1) # (batch, 1) outputs torch.cat(outputs, dim1) # (batch, seq_len-1, vocab_size) return outputs def forward(self, x, cond): mu, logvar self.encode(x, cond) z self.reparameterize(mu, logvar) recon_x self.decode(z, cond, target_seqx) return recon_x, mu, logvar # 损失函数重构损失 KL散度 def loss_function(recon_x, x, mu, logvar): # recon_x: (batch, seq_len-1, vocab_size), x: (batch, seq_len) # 忽略pad token target x[:, 1:] # 去掉sos recon_loss F.cross_entropy(recon_x.reshape(-1, recon_x.size(-1)), target.reshape(-1), ignore_indextoken_to_idx[pad]) kl_loss -0.5 * torch.sum(1 logvar - mu.pow(2) - logvar.exp()) return recon_loss 0.001 * kl_loss # KL权重需要调参关键点解析条件注入位置注意在编码器和解码器的每一步都拼接了条件向量。这确保了条件信息能充分影响整个编码-解码过程。如果只在潜向量z处拼接控制力可能会减弱。KL散度权重这是一个超参数。权重太大模型会倾向于忽略输入数据潜空间会坍缩成一个简单的分布如标准正态导致生成分子多样性差权重太小模型会退化成普通自编码器潜空间不规则不利于插值和条件生成。通常从0.001开始尝试。教师强制Teacher Forcing在训练解码器时以一定概率将上一个时间步的真实标签作为当前输入而不是模型自己的预测。这能加速训练但可能导致推理时不使用教师强制性能下降。常用策略是随着训练进行逐渐降低教师强制比率。3.3 模型训练与条件生成训练过程是标准的PyTorch流程。训练完成后我们就可以进行条件生成了。def generate_with_condition(model, condition_vector, num_samples10, max_len100): 根据给定条件向量生成分子 model.eval() with torch.no_grad(): cond torch.tensor(condition_vector, dtypetorch.float).unsqueeze(0) # (1, cond_dim) cond cond.to(device) # 从标准正态分布采样潜向量 z torch.randn(num_samples, model.latent_dim).to(device) cond cond.expand(num_samples, -1) # 扩展到每个样本 generated_smiles [] decoder_input torch.full((num_samples, 1), token_to_idx[sos], devicedevice, dtypetorch.long) decoder_hidden model.latent2hidden(torch.cat([z, cond], dim-1)).unsqueeze(0) decoder_cell torch.zeros_like(decoder_hidden) for t in range(max_len - 1): embedded model.embedding(decoder_input) cond_expanded cond.unsqueeze(1) decoder_input_step torch.cat([embedded, cond_expanded], dim-1) lstm_out, (decoder_hidden, decoder_cell) model.decoder_lstm(decoder_input_step, (decoder_hidden, decoder_cell)) output model.outputs2vocab(lstm_out.squeeze(1)) top1 output.argmax(1) decoder_input top1.unsqueeze(1) # 收集生成的token for i in range(num_samples): if t len(generated_smiles[i]): continue token idx_to_token[top1[i].item()] if token eos: # 已经生成结束用pad填充后续输入以避免无效计算 decoder_input[i] torch.tensor([token_to_idx[pad]], devicedevice) else: if len(generated_smiles) i: generated_smiles.append([]) generated_smiles[i].append(token) # 将token列表转换为SMILES字符串 final_smiles [] for token_list in generated_smiles: # 过滤掉起止符和填充符 smiles_str .join([t for t in token_list if t not in [sos, eos, pad, unk]]) final_smiles.append(smiles_str) return final_smiles # 示例生成满足“分子量500, logP在2-3之间 QED0.6”的分子 target_condition [1, 1, 1] # 对应我们之前定义的三个二分类条件 generated generate_with_condition(model, target_condition, num_samples5) for i, smi in enumerate(generated): print(f生成分子 {i1}: {smi}) # 可以用RDKit验证生成分子的属性 mol Chem.MolFromSmiles(smi) if mol: print(f 分子量: {Descriptors.MolWt(mol):.2f}, LogP: {Crippen.MolLogP(mol):.2f}, QED: {QED.qed(mol):.2f})4. 进阶优化与实战技巧基础模型跑通只是第一步。要让它在真实项目中发挥作用还需要大量的“精雕细琢”。4.1 提升生成分子的有效性与多样性一个常见的痛点是生成的SMILES字符串无法被RDKit解析成有效的分子无效或者生成的分子结构千篇一律模式坍塌。对抗无效SMILES语法约束在解码的每一步只允许从符合化学价键规则的token中进行采样。这需要维护一个分子构建的状态机实现较复杂但效果显著。后处理过滤这是更简单直接的方法。生成大量分子比如10000个然后用RDKit快速过滤掉无效的、重复的分子。虽然粗暴但在计算资源充足时很有效。使用更鲁棒的分词器字符级分词对模型学习化学语法要求很高。改用基于子词如BPE的分词让模型学习“CO”、“c1ccccc1”苯环这样的常见化学子结构能大幅提升生成有效性和学习效率。对抗模式坍塌调整KL散度权重适当增加KL权重鼓励潜空间更接近标准正态分布从而采样到更多样的潜向量。引入多样性奖励在训练或生成过程中计算生成分子集合的内部相似度如基于分子指纹的Tanimoto相似度并惩罚相似度过高的分子鼓励多样性。使用更先进的架构如前所述流模型或扩散模型在建模复杂多峰分布上具有理论优势。4.2 处理连续值与多目标权衡我们的例子中使用了二分类条件。但很多时候我们希望精确控制一个连续值比如“生成LogP在2.5左右的分子”。连续条件编码直接将连续值如2.5作为条件向量的一部分输入模型。但要注意数值缩放Normalization。将不同量纲的属性如分子量200-500 LogP -2到6缩放到相近的范围如0-1或标准正态分布有助于模型稳定训练。多目标帕累托前沿采样我们无法让一个分子所有属性都最优。一种高级技巧是不在潜空间采样一个点而是采样一条线或一个面。例如我们可以让条件向量在“高活性”和“高溶解性”两个目标之间线性插值从而生成一系列位于这两个目标权衡边界上的分子供化学家进一步选择。# 示例在“偏向活性”和“偏向溶解性”之间插值 cond_active [1.0, 0.0] # 假设条件向量是二维连续值 cond_soluble [0.0, 1.0] for alpha in [0, 0.2, 0.5, 0.8, 1.0]: interpolated_cond [(1-alpha)*a alpha*b for a, b in zip(cond_active, cond_soluble)] molecules generate_with_condition(model, interpolated_cond, num_samples5) # 分析这批分子在活性和溶解性上的分布4.3 与外部预测模型联动我们模型中的条件如QED, LogP是用RDKit的简单描述符计算的。但对于“生物活性”这种复杂属性我们需要更精确的预测模型如基于图的神经网络QSAR模型。两阶段流程生成阶段使用基于简单、快速描述符如分子量、LogP、QED的CVAE生成大量初步候选分子。筛选阶段用训练好的高精度活性预测模型对生成的分子进行打分和排序选出Top-N。端到端流程将高精度预测模型作为奖励函数集成到强化学习框架中。生成模型每产生一个分子就送去预测模型打分并将分数作为奖励反馈给生成模型引导其向高活性区域探索。这种方法更强大但实现和调参更复杂。5. 常见问题与排查实录在实际操作中你一定会遇到各种各样的问题。下面是我总结的一些典型问题及其解决思路。问题现象可能原因排查与解决思路生成的分子100%无效1. SMILES标准化不一致训练和生成时处理方式不同。2. 词汇表缺失关键字符。3. 模型完全未学会语法训练不收敛。1.检查数据管道确保生成时使用的token_to_idx和标准化函数与训练时完全一致。打印出前几个生成序列肉眼检查是否像SMILES。2.检查词汇表确保它包含了训练集中所有出现过的字符。可以添加一个unk令牌处理未见字符。3.检查训练损失观察重构损失是否在下降。如果KL损失权重过大可能导致“KL消失”模型忽略输入。尝试降低KL权重。生成分子多样性极差1. 模式坍塌Mode Collapse。2. KL损失权重过小潜空间没有正则化。3. 训练数据本身多样性不足。1.检查潜空间对一批训练数据的潜向量z进行PCA或t-SNE可视化看它们是否聚集在一个点或一个小区域。如果是说明模式坍塌。2.调整KL权重逐步增大KL权重观察生成多样性的变化。3.引入多样性惩罚在损失函数中加入基于分子指纹的多样性项。4.检查数据分析训练数据中骨架的分布是否被少数几个优势骨架主导条件控制失灵1. 条件信息在模型中被淹没。2. 条件之间的权重设置不合理。3. 条件本身定义模糊或冲突。1.可视化条件影响固定潜向量z逐渐改变某个条件值如LogP从1到5观察生成分子的变化趋势。如果没变化说明条件未起作用。2.调整条件注入方式尝试将条件向量与潜向量z拼接后再通过一个全连接层变换而不是简单拼接。3.重新审视条件用散点图分析训练数据中各个条件之间的关系。是否存在“高活性必然导致低溶解性”这种强负相关如果是要求模型同时达到两者最优是不现实的。模型训练震荡不收敛1. 学习率过高。2. 梯度爆炸/消失。3. 批次内数据差异过大。1.使用学习率预热和衰减。2.梯度裁剪在loss.backward()之后optimizer.step()之前加入torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)。3.检查数据确保批次内的SMILES长度不要差异过大可通过DataLoader的collate_fn进行填充。4.使用更稳定的架构如果RNN/VAE不稳定可以尝试Transformer-based的模型。最后一点个人体会AI分子生成不是一个“一劳永逸”的魔法黑箱。它更像是一个强大的“创意加速器”。最有效的工作流是“AI生成 - 化学家评估 - 反馈给AI”。化学家的经验、直觉和对复杂性质的判断如代谢位点、合成路线目前仍是AI难以完全替代的。因此在设计条件时多与药物化学家沟通将他们的经验规则如“避免这个警示结构”、“倾向于这个骨架”转化为模型可以理解的约束条件才能让AI真正成为得力的合作伙伴在浩瀚的化学星海中更高效地指引我们找到那颗希望之星。