Toward Better Geometric Representations for Molecule Generative Models¶
会议: NeurIPS 2026
arXiv: 2605.07693
OpenReview: https://openreview.net/forum?id=5vRYljseiu
领域: 计算生物
关键词: 3D分子生成, 几何表示学习, 扩散模型, 感知损失, 层次化特征对齐
一句话总结¶
针对基于几何表示引导的 3D 分子生成模型中预训练表征空间不平滑且多层语义未充分利用的瓶颈，提出 LENSEs 框架，通过可学习层级池化变分表征头、分子感知损失与节点级特征对齐，实现了更平滑的隐空间与 SOTA 的分子生成质量。
研究背景与动机¶
三维分子生成是现代 AI 制药与 de novo 药物设计的核心基础。为了应对大分子和复杂药物结构中精细的立体几何与拓扑空间建模挑战，近年来涌现了大量基于等变扩散模型（如 EDM）与流匹配（Flow Matching）的生成方法。其中，几何表征条件生成范式（GeoRCG）通过将流程解耦为两阶段展现出极高效率：第一阶段利用表征扩散模型（RDM）在预训练分子编码器（如 UniMol、Frad）提取的隐空间中采样分子级向量；第二阶段则以该表征为条件，指导等变生成器反解出 3D 分子构象。
然而，这一两阶段解耦范式的生成上限直接受制于预训练分子表征空间的内在质量。现有预训练编码器主要依赖于坐标去噪或力场恢复等原子级目标，其池化得到的分子级向量在几何流形上极其不平滑——微小的坐标扰动会导致隐空间向量产生剧烈波动（对应高达数千的局部 Lipschitz 常数），迫使前序方法不得不依靠人工调优的启发式加噪来维持训练稳定性。更为关键的是，预训练编码器各层蕴含着从局部基团到全局骨架的层次化语义，但现有管线仅将其浅显地当作单一条件提供者，深层与浅层的结构语义未被有效传递给生成器，迫使生成器只能从头学习分子结构规律。
为了打破表征与生成之间的单向割裂，本文的核心 insight 是让表征空间在生成任务的驱动下动态重构与对齐，构建高质量的双向协同机制。核心 idea：提出 LENSEs 框架，通过引入可学习层级池化的变分表征头平滑潜变量空间，并结合宏观分子感知损失与微观节点级表征对齐（REPA），使生成模型在训练中主动重塑并利用编码器的多层结构语义。
方法详解¶
整体框架¶
LENSEs（Latent Enhancement for Non-smooth Structural Encodings）整体分为三个阶段： 1. 阶段 I（联合训练）：冻结预训练编码器，将其多层图表征输入到带可学习层级池化与高斯重参数化技巧的表征头中，生成紧凑且平滑的随机隐向量 \(\bm{z}\)，作为等变分子生成器的条件；同时，引入分子感知损失与节点级 REPA 对齐损失，联合优化表征头与生成器。 2. 阶段 II（表征扩散模型训练）：在表征头输出的均值空间 \(\bm{\mu}\) 上训练轻量级表征扩散模型（RDM），学习目标分子的隐空间先验分布。 3. 阶段 III（两阶段推理生成）：先由 RDM 采样目标分子的隐表征 \(\hat{\bm{\mu}}\)，再输入到分子生成器中直接去噪生成符合几何与化学规律的 3D 分子构象 \(\hat{\mathcal{M}}\)。
整体流程如图所示：
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入分子 ℳ<br/>3D 坐标 + 原子类型"] --> B["冻结预训练编码器<br/>提取多层特征 H^(l)"]
B --> C["可学习层级池化变分表征头<br/>softmax 自适应聚合与 KL 正则化"]
C -->|采样平滑隐向量 z| D["等变分子生成器<br/>以 z 为条件进行去噪生成"]
D --> E["去噪分子预测 ℳ^"]
E --> F["分子感知损失 (L_perc)<br/>Straight-Through 离散梯度回传"]
D -.->|中间层节点特征| G["节点级表征投影对齐 (L_repa)<br/>Proj 映射与余弦相似度最大化"]
B -.->|编码器节点特征| G
E --> H["下游 3D 分子输出<br/>低能量高稳定性构象"]关键设计¶
1. 可学习层级池化变分表征头：自适应融合多层语义并平滑隐空间流形 针对预训练编码器顶层表征过分稀疏且对坐标扰动过敏（Lipschitz 常数过大）的缺陷，LENSEs 不直接使用最终层池化结果，而是引入自适应层级加权机制。设编码器共有 \(L\) 层，各层图级表征为 \(\bm{H}^{(l)}(\mathcal{M})\)，表征头通过带 softmax 归一化的可学习权重 \(\bm{w}\) 聚合多层表征： $\(\bm{g}(\mathcal{M}) = \sum_{l=1}^{L} \alpha_{l} \bm{H}^{(l)}(\mathcal{M}), \quad \alpha_{l} = \mathrm{Softmax}(\bm{w})_{l}\)$ 聚合后的向量 \(\bm{g}(\mathcal{M})\) 随后输入均值网络 \(f_{\mu}\) 与截断对数方差网络 \(f_{\log\sigma^2}\)，通过重参数化采样 \(\bm{z} = \bm{\mu} + \bm{\sigma} \odot \bm{\epsilon}\)（其中 \(\bm{\epsilon} \sim \mathcal{N}(\bm{0}, \bm{I})\)）输入生成器。同时引入针对先验正态分布的 KL 散度约束 \(\mathcal{L}_{\mathrm{KL}}\)。该设计在保留第 8 层宏观骨架信息的同时，恢复了前几层（占约 18% 权重）对羟基、氨基等局部官能团的强判别性，并构建了连续平滑的高斯潜在流形。
2. 分子感知损失：通过编码器特征空间提供宏观语义监督 传统的分子生成模型仅在坐标与原子类型空间计算均方误差或交叉熵损失，容易陷入局部欧氏距离陷阱，忽视分子的宏观化学拓扑与物理整体性。LENSEs 借鉴图像领域的感知损失，将去噪生成的分子预测 \(\hat{\mathcal{M}}\) 重新送入冻结的预训练编码器，提取其池化图表征 \(\bm{g}(\hat{\mathcal{M}})\)，并与真实分子表征 \(\bm{g}(\mathcal{M})\) 进行对齐： $\(\mathcal{L}_{\mathrm{perc}} = w(t) \left\| \mathrm{sg}(\bm{g}(\mathcal{M})) - \bm{g}(\hat{\mathcal{M}}) \right\|_{2}^{2}\)$ 其中 \(\mathrm{sg}(\cdot)\) 为梯度截断操作。针对生成的原子类型为离散符号无法求导的问题，采用 Straight-Through Estimator（STE）将离散 one-hot 预测与连续 logits 残差结合，使梯度能沿编码器顺利回传至生成器。同时，引入余弦时间权重调度器 \(w(t)\)，在去噪后期分子结构趋于清晰时施加更强的语义约束。
3. 节点级表征投影对齐（REPA）：消除编码器与生成器特征割裂 为了解决生成器必须在训练过程中“从零重新发明分子拓扑轮廓”的学习负担，LENSEs 在生成器内部隐藏状态与预训练编码器之间建立了微观层面的桥梁。选取编码器第 \(l_e\) 层节点特征 \(\bm{H}_{\mathrm{enc}}^{(l_e)} \in \mathbb{R}^{N \times d_e}\) 与生成器第 \(l_g\) 层节点特征 \(\bm{H}_{\mathrm{gen}}^{(l_g)} \in \mathbb{R}^{N \times d_g}\)，利用轻量线性映射 \(\mathrm{Proj}: \mathbb{R}^{d_g} \to \mathbb{R}^{d_e}\) 映射特征空间，并在所有有效原子集合 \(\mathcal{A}\) 上最大化两者的余弦相似度： $\(\mathcal{L}_{\mathrm{repa}} = -\frac{1}{|\mathcal{A}|} \sum_{i \in \mathcal{A}} \left\langle \frac{\mathrm{Proj}(\bm{H}_{\mathrm{gen}, i}^{(l_g)})}{\|\mathrm{Proj}(\bm{H}_{\mathrm{gen}, i}^{(l_g)})\|_2}, \frac{\mathrm{sg}(\bm{H}_{\mathrm{enc}, i}^{(l_e)})}{\|\mathrm{sg}(\bm{H}_{\mathrm{enc}, i}^{(l_e)})\|_2} \right\rangle\)$ 采用余弦相似度而非均方误差有效消除了不同网络层级特征尺度的固有偏差，促使生成器各节点直接对齐高质量预训练几何先验，大幅加速训练收敛并提升构象真实度。
损失函数 / 训练策略¶
阶段 I 总训练目标为多任务协同优化： $\(\mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{gen}} + \lambda_{\mathrm{KL}} \mathcal{L}_{\mathrm{KL}} + \lambda_{\mathrm{perc}} \mathcal{L}_{\mathrm{perc}} + \lambda_{\mathrm{repa}} \mathcal{L}_{\mathrm{repa}}\)$ 其中 \(\mathcal{L}_{\mathrm{gen}}\) 为基础流匹配/扩散生成损失。在 GEOM-DRUG 实验中，超参数配置为 \(\lambda_{\mathrm{KL}} = 5 \times 10^{-7}\)，\(\lambda_{\mathrm{perc}} = 1 \times 10^{-3}\)，\(\lambda_{\mathrm{repa}} = 1 \times 10^{-3}\)。优化时冻结编码器 \(f_{\mathrm{enc}}\)，仅更新生成器参数 \(\theta\)、表征头参数 \(\psi\) 以及线性投影头 \(\mathrm{Proj}\)。
实验关键数据¶
主实验¶
在真实复杂药物基准 GEOM-DRUG 上，LENSEs 与主流等变扩散/流匹配方法进行了全面对比（化学有效性、分子稳定性以百分比表示，能量与应变能单位为 kcal/mol）：
| 数据集 | 模型 | 原子稳定性 (%) ↑ | 分子稳定性 (%) ↑ | 化学有效性 (%) ↑ | 能量 (kcal/mol) ↓ | 应变能 (kcal/mol) ↓ |
|---|---|---|---|---|---|---|
| GEOM-DRUG | FlowMol | 99.0 | 67.5 | 51.2 | — | — |
| GEOM-DRUG | MiDi | 99.8 | 91.6 | 77.8 | — | — |
| GEOM-DRUG | EQGAT-diff | 99.8 | 93.4 | 94.6 | 148.8 | 140.2 |
| GEOM-DRUG | SemlaFlow | 99.8 | 97.3 | 93.9 | — | — |
| GEOM-DRUG | CanonFlow | 99.8 | 98.4 | 95.9 | — | — |
| GEOM-DRUG | GeoRCG (基线) | 99.8 | 97.20 | 95.3 | 88.6 | 47.64 |
| GEOM-DRUG | LENSEs (本文) | 99.88 | 98.51 | 97.28 | 85.53 | 43.86 |
在 QM9 基准上，LENSEs 同样达到 99.48% 有效性与 99.68% 分子稳定性，Opt-RMSD 降低至 0.21 Å（NFE=100）。
消融实验¶
在 GEOM-DRUG 上针对三大核心损失组件及超参数的消融表现如下：
| 配置 | \(\lambda_{\mathrm{KL}}\) | \(\lambda_{\mathrm{perc}}\) | \(\lambda_{\mathrm{repa}}\) | 有效性 (%) ↑ | 分子稳定性 (%) ↑ | 能量 (kcal/mol) ↓ | 应变能 (kcal/mol) ↓ | 说明 |
|---|---|---|---|---|---|---|---|---|
| GeoRCG 基线 | – | – | – | 94.20 | 97.20 | 103.33 | 66.95 | 原始非平滑条件生成 |
| 仅 KL 约束 | \(5 \times 10^{-7}\) | – | – | 96.52 | 98.01 | 90.66 | 43.38 | 平滑隐空间显著降应变能 |
| KL + 感知损失 | \(5 \times 10^{-7}\) | \(1 \times 10^{-3}\) | – | 96.60 | 98.70 | 82.63 | 39.98 | 宏观语义对齐大幅压低能量 |
| KL + REPA 对齐 | \(5 \times 10^{-7}\) | – | \(1 \times 10^{-3}\) | 96.40 | 98.00 | 86.87 | 45.83 | 节点级先验加速结构收敛 |
| 完整 LENSEs | \(5 \times 10^{-7}\) | \(1 \times 10^{-3}\) | \(1 \times 10^{-3}\) | 97.28 | 98.51 | 85.53 | 43.86 | 各模块互补，综合指标最优 |
关键发现¶
- 隐空间平滑度剧增：局部 Lipschitz 常数分析显示，在 GEOM-DRUG 上 LENSEs 的平均 Lipschitz 常数从原始 UniMol 的 2659.81 骤降至 570.49（降幅达 4.6 倍），有效秩从 65.32 提升至 88.42，彻底克服了表征维度的维度坍缩问题。
- 少步采样效率突出：在 GEOM-DRUG 上仅需 20 步生成时，LENSEs 的有效性即可达 96.7%、分子稳定性达 96.6%，持平甚至超越了基准模型在 50 至 100 步采样的质量，验证了优质表征条件能大幅削减生成轨迹的探索成本。
- 反哺下游属性预测：冻结 LENSEs 优化后的表征并在 QM9 上微调线性层做性质预测，各项量子化学参数 MAE 相比原始 Frad 编码器全面降低（例如 \(C_v\) 从 0.100 降至 0.048，\(\Delta\varepsilon\) 从 96.67 降至 68.76 meV），证明生成对齐任务本身具备成为分子自监督预训练新范式的潜力。
亮点与洞察¶
- 表征与生成任务的互惠闭环：不仅利用预训练几何表征来指导生成，更利用生成目标反向优化与平滑表征空间，将生成过程转化为分子编码器的无监督精调器。
- 分层语义的合理分配：实验表明编码器前几层对微观基团具有最高线性可分性，而深层编码宏观立体拓扑，可学习池化层自适应将 82% 权重给顶层、18% 权重留给浅层，既稳固了全局骨架又补齐了局部微观模式。
- 跨模态技术向科学智能迁移：成功将计算机视觉中的感知损失（Perceptual Loss）和表征对齐（REPA）无缝迁移至 3D 几何分子生成领域，并通过 STE 解决了离散原子类别反向传播的数学障碍。
局限与展望¶
- 计算开销略有增加：在训练阶段需要对生成分子再次通过编码器进行前向推理以计算感知损失，并计算 Straight-Through 梯度，单步迭代耗时略高于纯坐标扩散基准。
- 对非刚性大环分子的泛化边界：目前主要在小分子及类药分子基准（QM9、GEOM-DRUG）上测试，对于数百原子规模的共价大分子、蛋白质-配体共生成场景下的层级对齐效果有待进一步验证。
- 未来方向：探索与一步或极少步流匹配蒸馏模型结合，同时进一步探索纯生成对齐目标作为分子基础模型主干预训练的可行性。
相关工作与启发¶
- vs GeoRCG: GeoRCG 使用冻结的预训练编码器单向提取特征，遭遇隐空间粗糙与非平滑瓶颈，依赖人工加噪；LENSEs 通过变分头、感知损失与 REPA 构建双向优化闭环，有效性从 95.3% 跃升至 97.28%。
- vs EDM / SemlaFlow: 传统等变模型完全依赖几何坐标与图结构从零扩散；LENSEs 继承两阶段解耦范式，并大幅提升了采样效率，在 20 步下即可超越传统模型 100 步质量。
- vs REPA / REED: 将视觉表征对齐范式扩展到带有连续 3D 欧氏对称性与离散化学键/原子类别的图几何空间，证明了宏观语义感知在分子设计中的关键价值。
评分¶
- 新颖性: ⭐⭐⭐⭐ [首次系统揭示并解决了表征条件生成中编码器流形不平滑与层级语义丢失的痛点]
- 实验充分度: ⭐⭐⭐⭐⭐ [涵盖生成指标、Lipschitz 平滑度、有效秩、少步采样及下游性质预测回归验证]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密，动机与方法阐述自洽，分析实验细致入微]
- 价值: ⭐⭐⭐⭐⭐ [为 3D 分子生成及分子表示学习领域提供了兼具实用性与理论启发的全新范式]