基于CVAE与Transformer的多约束条件AI分子生成技术实践

1. 项目概述：当AI成为“药物化学家”

最近几年，AI在药物研发领域的渗透速度远超预期。从早期的虚拟筛选辅助，到如今直接参与分子设计与优化，AI正从一个“计算工具”演变为一个具备初步“化学直觉”的“合作者”。我接触这个领域也有几年了，从最初用现成的模型做性质预测，到现在尝试构建端到端的分子生成优化流程，踩过的坑不少，但收获的惊喜更多。今天想和大家深入聊聊的，就是“AI药物分子优化”中的一个核心且极具挑战性的环节：如何从一个简单的分子线性表示（SMILES）出发，生成出同时满足多个、甚至相互冲突的约束条件的新分子。

这听起来像是个“既要、又要、还要”的难题。在真实的药物研发项目中，一个理想的候选分子需要同时具备高活性（对靶点蛋白有效）、良好的成药性（比如合适的溶解度、代谢稳定性、低毒性）、以及可合成性（化学家能在实验室里以合理的成本合成出来）。传统的方法往往是一个一个条件去筛选，耗时费力，且容易陷入局部最优。而基于深度学习的生成模型，为我们提供了一种全新的思路：将所有这些约束条件“编码”进模型的训练或生成过程中，让AI自己去探索浩瀚的化学空间，寻找那个能满足所有苛刻条件的“全能选手”。

这个过程的核心输入是SMILES（Simplified 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 主流模型架构的横向对比

在实践中，我主要对比和尝试过以下几种思路：

1. 条件变分自编码器（CVAE）：这是最直观的思路之一。在标准的VAE中，编码器将分子（SMILES）压缩成一个连续的低维潜向量（latent vector），解码器再从这个向量重建分子。CVAE则在这个潜向量中拼接（concatenate）上条件向量（比如，我们希望的血浆蛋白结合率值）。这样，解码器在重建或生成时，就会同时受到潜向量和条件向量的影响。

   • 优点：概念清晰，训练相对稳定，生成分子的多样性好。
   • 缺点：对条件控制的精确性有时不足，尤其是在条件维度较高或条件间存在复杂关系时。生成的分子可能在“满足条件A”和“满足条件B”之间摇摆。

2. 强化学习（RL）微调：先训练一个无条件的生成模型（例如一个GPT风格的Transformer），作为“策略网络”。然后，我们定义一系列“奖励函数”，每个函数对应一个约束条件（例如，预测的活性值越高奖励越大，预测的毒性值越高惩罚越大）。通过策略梯度等RL方法，引导模型生成能获得高总奖励的分子。

   • 优点：非常灵活，可以处理极其复杂、非可微的奖励函数（比如基于规则的结构过滤器）。
   • 缺点：训练不稳定，需要精心设计奖励函数和平衡权重，否则模型容易崩溃或陷入单一模式。计算成本也较高。

3. 条件Transformer（如MolGPT, Chemformer）：在Transformer的解码过程中，将条件信息作为额外的输入。例如，在生成每个token之前，将条件向量与当前的上下文表示进行融合。这类似于在文本生成中给定“写作风格”或“主题”。

   • 优点：得益于Transformer强大的序列建模能力，对长程依赖和复杂条件建模效果好。
   • 缺点：需要大量的训练数据，且条件信息的注入方式需要精巧设计，否则容易被模型忽略。

4. 流模型（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）。我们如何让模型同时考虑它们？

1. 加权求和法（最常见）：为每个条件计算一个得分（或损失），然后乘以一个权重后相加，得到总得分（或总损失）。总得分 = w1 * 活性得分 + w2 * 溶解性得分 + w3 * QED得分 - w4 * SA得分。

   • 关键：权重的设置是艺术也是科学。权重过大，模型可能只优化单一指标；权重过小，该条件可能被忽略。我通常的做法是先进行单目标优化，观察每个目标单独能带来的得分变化范围，然后根据业务重要性，手动调整权重，使各个目标在总损失中的贡献量级相当。这是一个需要反复实验的过程。

2. 分层/序贯控制法：先让模型满足最核心的1-2个条件（如活性和类药性），生成一批候选分子。然后，用这些分子作为起点，在第二个阶段引入其他条件（如溶解性和合成难度）进行进一步优化或筛选。这可以通过串联不同的模型，或在RL框架中分阶段调整奖励函数来实现。

   • 适用场景：当约束条件之间存在明显的优先级或因果关系时。例如，必须先有活性，再谈其他。

3. 帕累托优化思想：不追求一个在所有指标上都“最优”的分子（这种分子往往不存在），而是寻找一批“非支配”分子。即，在这批分子中，你无法找到一个分子在所有指标上都比另一个分子好。这更适合在生成大量分子后，进行后处理和分析时使用。

踩坑记录：早期我曾试图让模型同时优化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, isomericSmiles=True) 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(drop=True), prop_df], axis=1) # 4. 定义条件目标 # 假设我们希望：分子量<500, logP在2-3之间， QED>0.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”和“O=C(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_len=100): 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, dtype=torch.long) # 条件向量 cond = torch.tensor(self.df.iloc[idx]['condition_vector'], dtype=torch.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_idx=token_to_idx['<pad>']) # 编码器：Bi-LSTM self.encoder_lstm = nn.LSTM(embedding_dim + condition_dim, hidden_dim, bidirectional=True, batch_first=True) 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_first=True) 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_seq=None, teacher_forcing_ratio=0.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>'], device=z.device, dtype=torch.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, dim=1) # (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_seq=x) 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_index=token_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_samples=10, max_len=100): """根据给定条件向量生成分子""" model.eval() with torch.no_grad(): cond = torch.tensor(condition_vector, dtype=torch.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>'], device=device, dtype=torch.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>']], device=device) 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之间， QED>0.6”的分子 target_condition = [1, 1, 1] # 对应我们之前定义的三个二分类条件 generated = generate_with_condition(model, target_condition, num_samples=5) for i, smi in enumerate(generated): print(f"生成分子 {i+1}: {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：

  1. 语法约束：在解码的每一步，只允许从符合化学价键规则的token中进行采样。这需要维护一个分子构建的状态机，实现较复杂但效果显著。
  2. 后处理过滤：这是更简单直接的方法。生成大量分子（比如10000个），然后用RDKit快速过滤掉无效的、重复的分子。虽然粗暴，但在计算资源充足时很有效。
  3. 使用更鲁棒的分词器：字符级分词对模型学习化学语法要求很高。改用基于子词（如BPE）的分词，让模型学习“C=O”、“c1ccccc1”（苯环）这样的常见化学子结构，能大幅提升生成有效性和学习效率。

• 对抗模式坍塌：

  1. 调整KL散度权重：适当增加KL权重，鼓励潜空间更接近标准正态分布，从而采样到更多样的潜向量。
  2. 引入多样性奖励：在训练或生成过程中，计算生成分子集合的内部相似度（如基于分子指纹的Tanimoto相似度），并惩罚相似度过高的分子，鼓励多样性。
  3. 使用更先进的架构：如前所述，流模型或扩散模型在建模复杂多峰分布上具有理论优势。

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_samples=5) # 分析这批分子在活性和溶解性上的分布

4.3 与外部预测模型联动

我们模型中的条件（如QED, LogP）是用RDKit的简单描述符计算的。但对于“生物活性”这种复杂属性，我们需要更精确的预测模型（如基于图的神经网络QSAR模型）。

• 两阶段流程：
  1. 生成阶段：使用基于简单、快速描述符（如分子量、LogP、QED）的CVAE生成大量初步候选分子。
  2. 筛选阶段：用训练好的高精度活性预测模型对生成的分子进行打分和排序，选出Top-N。
• 端到端流程：将高精度预测模型作为奖励函数，集成到强化学习框架中。生成模型每产生一个分子，就送去预测模型打分，并将分数作为奖励反馈给生成模型，引导其向高活性区域探索。这种方法更强大，但实现和调参更复杂。

5. 常见问题与排查实录

在实际操作中，你一定会遇到各种各样的问题。下面是我总结的一些典型问题及其解决思路。

最后一点个人体会：AI分子生成不是一个“一劳永逸”的魔法黑箱。它更像是一个强大的“创意加速器”。最有效的工作流是“AI生成 -> 化学家评估 -> 反馈给AI”。化学家的经验、直觉和对复杂性质的判断（如代谢位点、合成路线），目前仍是AI难以完全替代的。因此，在设计条件时，多与药物化学家沟通，将他们的经验规则（如“避免这个警示结构”、“倾向于这个骨架”）转化为模型可以理解的约束条件，才能让AI真正成为得力的合作伙伴，在浩瀚的化学星海中，更高效地指引我们找到那颗希望之星。