SinAE：一种用于跨域原子系统的单架构流匹配自编码器

# 用于跨域原子系统的单架构流匹配自动编码器
来源：https://arxiv.org/html/2607.12380
Yuxuan Ren Department of Computer Science National University of Singapore yuxuan\.ren@nus\.edu\.sg &Fan Yang Tencent AI for Life Science Lab fionafyang@tencent\.com &Jianhua Yao Tencent AI for Life Science Lab jianhuayao@tencent\.com &Yatao Bian Department of Computer Science National University of Singapore ybian@nus\.edu\.sg

###### 摘要

小分子、晶体和蛋白质都可简化为三维空间中的原子，但其生成流程在各个领域仍然碎片化，每个领域都有自己的图、等变或基于框架的架构。跨域训练可以缓解每个领域的数据稀缺问题，但在三维坐标空间中直接生成很难轻松处理所有三个领域的异构结构先验，并且没有先前的潜在自动编码器能同时在所有三个领域实现无损和架构通用性。我们引入SinAE，一个用于分子、晶体和蛋白质的单架构流匹配自动编码器，采用标准Transformer编码器和解码器，没有等变算子、图算子或特定于领域的算子。SinAE不要求编码器捕捉精细几何结构，而是将重建负担转移到迭代流匹配解码器，在各领域实现近乎无损的重建，并将重建误差比先前的潜在基线降低多个数量级。相同的逐令牌潜在表示支持标准的扩散Transformer先验，在分子、晶体和蛋白质生成基准上达到强劲性能。分子-晶体联合训练严格提升了两个领域，为通过共享原子潜在表示实现跨域迁移提供了直接证据。代码可在https://anonymous.4open.science/r/SinAE_anonymous-523D/获取。

## 1 引言

三维原子结构生成支撑着药物发现、材料设计和蛋白质工程领域的进展。尽管在规模和领域惯例上存在差异，但底层数据是相同的：原子放置在R3\\mathbb\{R\}^\{3\}中，具有离散的元素类型，共享相同的成键和位阻排斥物理规律。然而，生成流程已沿着社区路线各自发展。分子生成器基于等变消息传递网络（Gilmer等人，2017（https://arxiv.org/html/2607.12380#bib.bib26）；Schütt等人，2017（https://arxiv.org/html/2607.12380#bib.bib27）；Thomas等人，2018（https://arxiv.org/html/2607.12380#bib.bib28））；晶体生成器还必须处理周期性边界条件和晶格-基矢不变性（Xie等人，2022（https://arxiv.org/html/2607.12380#bib.bib23）；Jiao等人，2023（https://arxiv.org/html/2607.12380#bib.bib22））；蛋白质生成器依赖残基级别的参数化，如SE(3)框架（Yim等人，2023（https://arxiv.org/html/2607.12380#bib.bib47）；Watson等人，2023（https://arxiv.org/html/2607.12380#bib.bib46）；Bose等人，2024（https://arxiv.org/html/2607.12380#bib.bib56））或对表示上的三角注意力（Jumper等人，2021（https://arxiv.org/html/2607.12380#bib.bib29）；Lin等人，2023（https://arxiv.org/html/2607.12380#bib.bib38））。每个领域的三维数据集都很小：QM9（Ramakrishnan等人，2014（https://arxiv.org/html/2607.12380#bib.bib37））包含134134k个分子，MP-20（Xie等人，2022（https://arxiv.org/html/2607.12380#bib.bib23））包含4545k个晶体，按现代标准都属于低数据范围。一个涵盖所有三个领域的统一架构将允许每个领域为其他领域提供训练信号，并重用共享的几何规律性（键长分布、配位模式、位阻排斥等）。

表1：SinAE与潜在自动编码器基线的能力比较。✓：支持；✗：不支持。*标准Transformer*：主干网络不使用等变算子或图/消息传递。*联合类型+坐标*：模型联合生成离散原子类型或序列标签与连续坐标。*亚埃级重建*：在主要基准上重建RMSD低于10−3Å。

跨域原子生成沿着两条路线引起了越来越多的兴趣。*直接生成*路线避免了自动编码器，在坐标空间中或通过空间令牌化生成结构（Morehead等人，2026（https://arxiv.org/html/2607.12380#bib.bib52）；Lu等人，2026（https://arxiv.org/html/2607.12380#bib.bib58）；Zhang等人，2025（https://arxiv.org/html/2607.12380#bib.bib59））。这因其简单性而引人注目，但给单个模型带来了沉重负担：生成器必须同时学习有效的组成成分、精确的三维几何以及每个领域的独特结构先验（晶体的周期性和晶格等价性、蛋白质的残基级主链几何、分子的键连接性和价规则）。在实践中，将这些异构先验容纳在一个生成器中需要针对每种模态的令牌化方法和架构选择（Lu等人，2026（https://arxiv.org/html/2607.12380#bib.bib58））、尚未验证可用于无条件从头采样的分阶段序列-坐标生成（Zhang等人，2025（https://arxiv.org/html/2607.12380#bib.bib59）），或者将领域覆盖限制为分子和晶体（Morehead等人，2026（https://arxiv.org/html/2607.12380#bib.bib52））。*潜在自动编码器*路线提供了一种替代方案：自动编码器在其编码器和解码器中吸收每个领域的几何复杂性，先验作用于平滑的潜在编码而不是原始的三维结构（Joshi等人，2025（https://arxiv.org/html/2607.12380#bib.bib6））。单领域工作表明，近乎无损的潜在重建是可以实现的（Luo等人，2026（https://arxiv.org/html/2607.12380#bib.bib3）；Li等人，2026（https://arxiv.org/html/2607.12380#bib.bib44）），但将其扩展到跨域设置引入了现有方法尚未共同解决的两个挑战。

第一个是*重建保真度*。自动编码器中的重建误差直接限制了生成质量：先验无法纠正有损解码器引入的几何误差。在原子系统中，这个限制尤为严格，因为亚埃级的坐标漂移可能将键长和键角推至物理有效范围之外，并对应于势能的大幅变化。迄今为止唯一的跨域潜在自动编码器在QM9上实现了94.6％的结构匹配率，在MP-20上实现了84.5％（Joshi等人，2025（https://arxiv.org/html/2607.12380#bib.bib6）），接受这一损失作为领域通用性的代价。第二个是*训练复杂性*。现有的自动编码器依赖特定于领域的损失组合：蛋白质的结构令牌化结合了FAPE、距离和冲突损失以及各自的权重调整（Hayes等人，2025（https://arxiv.org/html/2607.12380#bib.bib30））；跨域管道使用单独的坐标和原子类型损失以及每个领域的权重（Joshi等人，2025（https://arxiv.org/html/2607.12380#bib.bib6））；分子自动编码器添加了键类型和成键距离损失以及多路权重向量（Luo等人，2026（https://arxiv.org/html/2607.12380#bib.bib3））。每增加一个领域就意味着重新平衡整个损失景观，使得跨域扩展成为一个损失工程问题，而不是数据扩展问题。

据我们所知，没有现有的跨域自动编码器同时具备无损、架构通用且可统一目标训练的特点（表1（https://arxiv.org/html/2607.12380#S1.T1））。

参见图注 图1：SinAE概述。一个用于分子、晶体和蛋白质的流匹配自动编码器，具有共享的逐令牌潜在表示（对于周期性系统为N+3个令牌，其中三个晶格令牌）。同一潜在表示上的DiT先验可实现从头生成。底部：每个领域的重建RMSD。

受此差距启发，我们引入了SinAE，一个用于分子、晶体和蛋白质的流匹配自动编码器，其编码器和解码器都是标准Transformer，没有等变算子、图算子或特定于领域的算子。遵循DiTo（Chen等人，2025b（https://arxiv.org/html/2607.12380#bib.bib18））的发现——单个扩散损失可以匹配或超越多损失VAE训练——我们用迭代流匹配解码器替换单次解码器，该解码器在所有领域中使用一个统一的重建目标进行训练。解码器通过多个ODE步骤细化结构，早期解析原子放置，后期细化成键几何和位阻（第5.2节（https://arxiv.org/html/2607.12380#S5.SS2）），在没有辅助损失或每个领域权重调整的情况下达到亚毫埃级RMSD。添加一个新领域需要添加训练数据，而不是添加损失项。

由于解码器处理每个领域的几何复杂性，所得的潜在空间近似高斯分布，并且不受原始结构空间对称性约束的限制。一个标准的DiT（Peebles和Xie，2023（https://arxiv.org/html/2607.12380#bib.bib51））足以作为生成先验，无需等变层、三角注意力或晶格感知模块，从而产生比直接在坐标空间中生成的方法更简单、更高效的管道。在推理时，生成的潜在表示以与训练重建相同的精度解码，无需事后松弛。表1（https://arxiv.org/html/2607.12380#S1.T1）总结了由此产生的能力概况。

总之，SinAE做出了三项贡献：

1. 1. 跨越三个领域的原子结构令牌化的单一架构。一个标准Transformer流匹配自动编码器在一个架构下处理分子、晶体和蛋白质主链，没有等变算子、图算子或特定于领域的算子。
2. 2. 所有三个领域的近乎无损重建。在分子和晶体上实现亚毫埃级全原子Kabsch对齐RMSD（QM9上0.0002Å，MP-20上0.0017Å）且蛋白质主链CA-RMSD约为0.02Å，分子和晶体匹配率达到100%，使用所有领域统一的流匹配重建目标实现。
3. 3. 跨域训练产生正迁移；潜在生成具有竞争力。QM9 + MP-20联合训练在重建和生成上严格改善了单领域基线，为共享原子潜在表示有益于两类化学种群提供了直接证据。相同的潜在表示支持一个标准DiT先验，该先验在硬指标上达到强劲的晶体生成以及有竞争力的分子和蛋白质生成。

## 2 背景与相关工作

分子。用于三维分子的早期深度学习方法基于分子图上的等变消息传递（Gilmer等人，2017（https://arxiv.org/html/2607.12380#bib.bib26）；Schütt等人，2017（https://arxiv.org/html/2607.12380#bib.bib27）；Thomas等人，2018（https://arxiv.org/html/2607.12380#bib.bib28））。扩散和流匹配生成器如EDM（Hoogeboom等人，2022（https://arxiv.org/html/2607.12380#bib.bib9））、EQGAT-diff（Le等人，2024（https://arxiv.org/html/2607.12380#bib.bib8））、SemlaFlow（Irwin等人，2024（https://arxiv.org/html/2607.12380#bib.bib4））和FlowMol（Dunn和Koes，2024（https://arxiv.org/html/2607.12380#bib.bib7））直接以专门化的等变架构在三维坐标空间中操作。在潜在侧，GeoLDM（Xu等人，2023（https://arxiv.org/html/2607.12380#bib.bib31））将分子压缩为SE(3)等变VAE编码，但其编码器-解码器瓶颈在GEOM-Drugs上留下0.25Å RMSD，这会在生成时传播到键长和键角误差。UAE-3D（Luo等人，2026（https://arxiv.org/html/2607.12380#bib.bib3））通过使用关系Transformer编码器和SO(3)增强实现近乎无损的分子重建来解决这个问题，但仅限于小分子，不扩展到晶体或蛋白质。

晶体。晶体生成还必须尊重周期性边界条件和晶格-基矢不变性（Xie等人，2022（https://arxiv.org/html/2607.12380#bib.bib23）；Jiao等人，2023（https://arxiv.org/html/2607.12380#bib.bib22））。CDVAE（Xie等人，2022（https://arxiv.org/html/2607.12380#bib.bib23））开创了基于VAE的晶体潜在表示；FlowMM（Miller等人，2024（https://arxiv.org/html/2607.12380#bib.bib21））在晶体晶格流形上引入了黎曼流匹配；FlowLLM（Sriram等人，2024（https://arxiv.org/html/2607.12380#bib.bib41））使用大语言模型作为基础分布。ADiT（Joshi等人，2025（https://arxiv.org/html/2607.12380#bib.bib6））通过共享的VAE + DiT管道统一了分子和晶体，但其单次VAE解码器在两个领域都留下了显著的重建损失。同时期的Zatom-1（Morehead等人，2026（https://arxiv.org/html/2607.12380#bib.bib52））采取相反的方法，消除了自动编码器，直接在R3中进行流匹配；它不产生可重用的潜在表示，也不处理蛋白质。SinAE表明，使用流匹配解码器，潜在路线在保留平滑潜在空间的同时，在硬生成指标上超越了这两条线（第4.4节（https://arxiv.org/html/2607.12380#S4.SS4））。

蛋白质。蛋白质主链生成主要被SE(3)框架扩散方法主导，如FrameFlow（Yim等人，2023（https://arxiv.org/html/2607.12380#bib.bib47））和RFdiffusion（Watson等人，2023（https://arxiv.org/html/2607.12380#bib.bib46）），它们直接对残基框架进行操作；它们的架构与蛋白质几何紧密耦合，无法与分子或晶体编码器共享。在潜在侧，ProteinAE（Li等人，2026（https://arxiv.org/html/2607.12380#bib.bib44））通过扩散自动编码器压缩主链，并在重建保真度上领先先前的潜在方法，但它是单领域的，且不联合生成残基类型序列与坐标。SinAE使用与分子+晶体模型相同的标准Transformer架构，在蛋白质重建上实现了10-20倍的更严格精度，同时额外支持序列-结构协同设计（第4.5节（https://arxiv.org/html/2607.12380#S4.SS5））。

## 3 方法

### 3.1 问题公式化与概述

在所有领域中，我们将原子结构表示为G=(X,A)\\mathcal\{G\}=(\\mathbf\{X\},\\mathbf\{A\})，其中X∈RN×3\\mathbf\{X\}\\in\\mathbb\{R\}^\{N\\times 3\}是三维坐标，A∈{1,...,K}N\\mathbf\{A\}\\in\\\{1,\\dots,K\\\}^\{N\}是离散的原子类型。对于周期性系统（晶体），我们还包含一个晶格矩阵L∈R3×3\\mathbf\{L\}\\in\\mathbb\{R\}^\{3\\times 3\}；对于非周期性系统（分子、蛋白质），则不使用晶格。

SinAE针对重建和生成进行训练。给定G\\mathcal\{G\}，自动编码器产生一个潜在编码z\\mathbf\{z\}，由此可以以近乎精确的保真度恢复G\\mathcal\{G\}；然后一个学习的先验p(z)p(\\mathbf\{z\})能够从头采样新结构。这些能力由两个按顺序训练的组件提供：

1. 1.流匹配自动编码器。一个Transformer编码器将G\\mathcal\{G\}映射到紧凑的潜在编码z\\mathbf\{z\}，而一个非等变的Transformer解码器通过条件流匹配从z\\mathbf\{z\}重建G\\mathcal\{G\}。解码器（而不是编码器）负责重建保真度，使得编码器保持简单且领域通用。
2. 2.潜在流匹配先验。一个标准的扩散Transformer（DiT）学习编码器输出上分布p(z)p(\\mathbf\{z\})。在生成时，一个新的潜在编码从该先验中采样，并通过冻结的解码器产生