Commun. Chem. | FRATTVAE：基于片段树 Transformer 的变分自编码器

Leveraging Tree-Transformer VAE with Fragment Tokenization for Highperformance Large Chemical Model Generation

期刊: Communications Chemistry 链接: https://doi.org/10.1038/s42004-025-01640-w 代码: https://github.com/slab-it/FRATTVAE 简介: 本文提出了 Fragment Tree-Transformer based VAE（FRATTVAE），旨在解决现有分子生成模型处理大型复杂化合物时的局限性，其创新点在于将分子视为以片段为节点的树结构，结合 Transformer 架构和片段 token 化技术。FRATTVAE 通过预处理将分子分解为片段并组织成树结构，利用树位置编码和 ECFP 嵌入片段特征，再通过 Transformer-based VAE 进行编码和解码，同时结合条件变体实现定向生成。实验使用了 ZINC250K、MOSES、GuacaMol、Polymer、SuperNatural3 等数据集，结果显示 FRATTVAE 在重建精度、FCD 等指标上超越现有方法，能高效处理大型复杂分子，在性质优化和目标导向任务中表现优异，且计算效率和可扩展性强。该论文证明了 FRATTVAE 是一种强大且通用的分子生成与优化解决方案，为化学信息学和药物发现领域提供了新工具。

在药物发现与材料科学领域，分子生成模型作为探索庞大化学空间的核心工具，其性能直接决定了候选化合物的发现效率与质量。传统方法中，基于SMILES字符串的化学语言模型（CLM）在处理大型复杂分子时，常面临结构准确性丢失的问题；而基于图结构的模型虽在分子有效性上有所提升，却受限于序列处理机制导致的计算效率瓶颈。针对这些挑战，近日，发表于《Communications Chemistry》的研究论文提出了Fragment Tree-Transformer based VAE（FRATTVAE），通过创新性地融合树结构表示、片段 token 化与Transformer架构，实现了大分子生成在准确性、效率与可扩展性上的三重突破。

分子生成模型的现状与局限

当前分子生成模型主要分为两类：基于字符串表示的模型与基于图结构的模型。前者以SMILES为代表，借助LSTM或Transformer等自然语言处理架构学习分子语法，但其对数据集规模的强依赖性及SMILES表示的敏感性（微小字符变化可能导致分子结构剧变），使其在复杂分子生成中稳定性不足。后者如Junction Tree VAE（JTVAE）和Hierarchical VAE（HierVAE），通过将分子分解为片段并构建树结构，减少了无效结构的生成，但依赖LSTM的序列处理机制导致计算速度慢，且难以适应超大规模数据集。

变分自编码器（VAE）作为主流框架，通过构建连续的化学潜在空间实现分子的编码与解码，但其在处理立体异构体、含盐或溶剂分子时，传统图表示往往失效。此外，现有模型在平衡重建精度与生成多样性上存在固有矛盾：例如SMILES-based Transformer VAE虽能达到高重建精度，却因有效性低（仅26.35%）而失去实用价值；而部分图模型虽保证了有效性，却在分布学习指标（如Fréchet ChemNet Distance，FCD）上表现不佳。

FRATTVAE的架构创新与技术细节

FRATTVAE的核心设计在于将分子视为以化学片段为节点的树结构，并利用Transformer的并行处理能力突破序列依赖瓶颈。其架构包含三个关键模块：

预处理与片段化：采用BRICS规则进行分子分解，将分子断裂为具有 retrosynthetic化学意义的片段，每个片段通过 dummy 原子标记连接位点，同时保留手性等立体化学信息。片段被组织为树结构，根节点通过RDKit的CANGEN算法确定，确保结构唯一性。这一过程不仅保留了片段固有的药理活性与理化性质，还使模型能自然处理金属离子、溶剂等传统图表示难以涵盖的成分。

树结构编码与嵌入：为使Transformer理解树的层级关系，FRATTVAE采用基于栈式向量的树位置编码，每个节点的位置信息由其父节点位置与自身在子节点中的排序共同决定，形成维度为n×k（n为节点最大子节点数，k为树深度）的特征向量。片段的化学特征则通过ECFP（Extended-Connectivity Fingerprints）表示，其能捕捉局部化学环境信息，为模型提供丰富的特征输入。

Transformer-based VAE核心：编码器通过“超级根节点”（〈super root〉）聚合树结构信息，将其映射为潜在变量z；解码器则基于z与树位置编码，通过掩码多头注意力机制逐层生成子节点，最终重建分子结构。训练过程采用重建损失（交叉熵）与KL散度的加权和作为目标函数，确保潜在空间既反映化学特性，又保持分布多样性。条件变体C-FRATTVAE通过在编码和解码阶段引入分子属性（如分子量、logP），实现了多条件约束下的定向生成。

该图展示了 FRATTVAE 的模型架构和分子生成过程。其中，图 1a 呈现了模型的整体架构，编码器和解码器均采用多头注意力和前馈层处理潜在变量以进行分子生成，过程中融入了树位置编码以及分子量（MolWt）、logP 等条件，通过注意力机制处理复杂依赖关系，确保大型复杂分子结构的准确生成。图 1b 描绘了 FRATTVAE 中的分子生成流程，训练条件变分自编码器时，在〈super root〉之前提供条件，从潜在变量出发，解码器按顺序生成子节点，保证分子结构的正确重建。

该图展示了 FRATTVAE 的模型架构和分子生成过程。其中，图 1a 呈现了模型的整体架构，编码器和解码器均采用多头注意力和前馈层处理潜在变量以进行分子生成，过程中融入了树位置编码以及分子量（MolWt）、logP 等条件，通过注意力机制处理复杂依赖关系，确保大型复杂分子结构的准确生成。图 1b 描绘了 FRATTVAE 中的分子生成流程，训练条件变分自编码器时，在〈super root〉之前提供条件，从潜在变量出发，解码器按顺序生成子节点，保证分子结构的正确重建。

性能验证：从基准测试到实际应用

在五个代表性数据集（ZINC250K、MOSES、GuacaMol、Polymer、SuperNatural3）上的测试表明，FRATTVAE在关键指标上全面超越现有模型：

该表对比了不同方法在 MOSES、GuacaMol、Polymer、SuperNatural3 和 ZINC250K 五个基准数据集上的分布学习指标，包括重建精度（Recon）、相似度（Similar）、有效性（Valid）、唯一性（Unique）、新颖性（Novelty）、Fréchet ChemNet 距离（FCD）等。结果显示，FRATTVAE 在多个数据集上表现优异，尤其在 FCD 等关键指标上超越现有方法，在处理大型分子的数据集上也展现出高重建精度和良好的性质分布匹配度，部分模型因无法处理特定数据集或性能不佳而表现较差。

该表对比了不同方法在 MOSES、GuacaMol、Polymer、SuperNatural3 和 ZINC250K 五个基准数据集上的分布学习指标，包括重建精度（Recon）、相似度（Similar）、有效性（Valid）、唯一性（Unique）、新颖性（Novelty）、Fréchet ChemNet 距离（FCD）等。结果显示，FRATTVAE 在多个数据集上表现优异，尤其在 FCD 等关键指标上超越现有方法，在处理大型分子的数据集上也展现出高重建精度和良好的性质分布匹配度，部分模型因无法处理特定数据集或性能不佳而表现较差。

分布学习能力：在MOSES数据集上，FRATTVAE的重建精度达94.87%，有效性为100%，FCD分数0.8654，显著高于JTVAE（FCD 0.7933）和MoLeR（FCD 0.8525）；在含大分子的GuacaMol数据集上，其FCD达0.8242，是唯一能稳定处理该数据集的模型（JTVAE因计算成本过高无法运行）。对于天然产物数据集（SuperNatural3），FRATTVAE的重建相似度达77.74%，NP-likeness指标（1D Wasserstein距离）为0.1480，远超SMIVAE的0.1837，证明其对结构异质性分子的捕捉能力。

该图展示了在单条件和多条件下生成分子的性质分布结果。图 2a 为单条件生成的分子性质分布，红色代表训练数据集 ZINC250K 的分布，尽管 ZINC250K 的化合物性质范围有限，条件生成存在挑战，但总体上能按条件生成分子，还能生成该数据集中不常见的大分子量和低 QED 值分子。图 2b 为多条件生成的分子性质分布，包含 SA 分数为纵轴、分子量（MolWt）、logP 或 QED 为横轴的二维核密度估计图，显示生成的分子分布符合各条件要求。

该图展示了在单条件和多条件下生成分子的性质分布结果。图 2a 为单条件生成的分子性质分布，红色代表训练数据集 ZINC250K 的分布，尽管 ZINC250K 的化合物性质范围有限，条件生成存在挑战，但总体上能按条件生成分子，还能生成该数据集中不常见的大分子量和低 QED 值分子。图 2b 为多条件生成的分子性质分布，包含 SA 分数为纵轴、分子量（MolWt）、logP 或 QED 为横轴的二维核密度估计图，显示生成的分子分布符合各条件要求。

属性优化与定向生成：在ZINC250K的PlogP优化任务中，FRATTVAE生成分子的PlogP值达16-17，C-FRATTVAE更达20-21，远超PSVAE（5-9）和MoLeR（8-9）。这得益于片段级操作使模型能一次性添加疏水环与长碳链，而原子级模型难以实现此类跃变式优化。在GuacaMol的20项目标导向任务中，FRATTVAE在12项任务中表现最优，尤其在药物重发现（如Celecoxib rediscovery得分0.835）和多属性优化（如Osimertinib MPO得分0.899）中优势显著，且生成分子的结构警报通过率（78.2%）高于MoLeR（65.2%），合成可行性（SA score 3.126）更优。

该表展示了以各性质为目标函数时，使用 MSO（分子群优化）方法进行性质优化的结果，列出了在 QED 和惩罚 logP（PlogP）方面发现的排名前三的分子。其中，C-FRATTVAE 通过将各性质作为条件生成分子，FRATTVAE 及其条件变体在 PlogP 优化中表现突出，生成分子的 PlogP 值远高于 PSVAE 和 MoLeR 等基线模型，在 QED 优化中也能达到较高水平。

该表展示了以各性质为目标函数时，使用 MSO（分子群优化）方法进行性质优化的结果，列出了在 QED 和惩罚 logP（PlogP）方面发现的排名前三的分子。其中，C-FRATTVAE 通过将各性质作为条件生成分子，FRATTVAE 及其条件变体在 PlogP 优化中表现突出，生成分子的 PlogP 值远高于 PSVAE 和 MoLeR 等基线模型，在 QED 优化中也能达到较高水平。

该表呈现了 FRATTVAE 和 MoLeR 在 GuacaMol 目标导向优化任务中的性能比较，包含每个任务的得分、平均得分、质量评级和 SA 分数等。FRATTVAE 在 20 个任务中的 12 个表现优于 MoLeR，在药物重发现、相似度和多属性优化（MPO）任务中优势明显，平均得分更高，生成分子的结构警报通过率（质量）和合成可行性（SA 分数）更优，同时还列出了数据集最佳值以及 GraphGA 和 CReM 的文献报道值作为参考。

该表呈现了 FRATTVAE 和 MoLeR 在 GuacaMol 目标导向优化任务中的性能比较，包含每个任务的得分、平均得分、质量评级和 SA 分数等。FRATTVAE 在 20 个任务中的 12 个表现优于 MoLeR，在药物重发现、相似度和多属性优化（MPO）任务中优势明显，平均得分更高，生成分子的结构警报通过率（质量）和合成可行性（SA 分数）更优，同时还列出了数据集最佳值以及 GraphGA 和 CReM 的文献报道值作为参考。

计算效率与可扩展性：FRATTVAE的预处理速度（ZINC250K数据集上1.7 ms/分子）和训练速度（10.8 ms/分子）远超PSVAE（74.1 ms/分子）和NPVAE（无法在大规模数据集上运行）。其Transformer架构支持GPU并行计算，成功训练出含10.3亿参数的FRATTVAE-large模型，在1200万分子（ChEMBL+DrugBank+PubChem10M）上实现稳定收敛，展示了处理超大规模化学数据的潜力。

该表比较了 FRATTVAE 与基于图的基线模型在 ZINC250K、MOSES 和 Polymer 数据集上的预处理、训练和采样速度（毫秒 / 分子）。结果显示，FRATTVAE 在预处理和训练阶段速度远快于其他模型，在生成阶段，虽在小分子数据集上略慢于 MoLeR，但仍快于 PSVAE，且在处理大型分子的 Polymer 数据集时，各阶段速度均超过所有模型，体现出其高效的计算效率和良好的可扩展性。

该表比较了 FRATTVAE 与基于图的基线模型在 ZINC250K、MOSES 和 Polymer 数据集上的预处理、训练和采样速度（毫秒 / 分子）。结果显示，FRATTVAE 在预处理和训练阶段速度远快于其他模型，在生成阶段，虽在小分子数据集上略慢于 MoLeR，但仍快于 PSVAE，且在处理大型分子的 Polymer 数据集时，各阶段速度均超过所有模型，体现出其高效的计算效率和良好的可扩展性。

结论与展望

FRATTVAE通过片段树结构与Transformer的创新性结合，首次实现了分子生成在准确性、效率与可扩展性上的统一。其核心优势在于：片段化保留了化学合理性，树结构捕捉了层级关系，Transformer并行机制突破了计算瓶颈。该模型不仅在基准测试中全面领先，更在天然产物、聚合物等复杂分子生成中展现出独特价值，为药物发现（如天然产物衍生药物设计）和材料科学（如高分子材料优化）提供了全新工具。

未来工作可进一步拓展训练数据规模、引入更多条件约束（如合成可及性），并探索与强化学习的结合，以推动FRATTVAE在实际药物开发流程中的落地应用。对于从事计算化学、AI药物发现的研究者而言，该模型的架构设计与片段处理策略具有重要的借鉴意义，其开源代码（https://github.com/slab-it/FRATTVAE）也为后续研究提供了坚实基础。