MOLEXA: Generative Modeling Enables Molecular Structure Retrieval from Coulomb Explosion Imaging

核心思想

库仑爆炸成像（CEI） 用 X 射线自由电子激光将分子电离成离子碎片，通过测量离子动量分布推断分子结构。但这是一个极其困难的逆问题——从碎片的动量反推原子的原始位置。

之前的方法只能处理 3-4 个原子的简单分子。MOLEXA 用扩散 Transformer 首次将这个能力扩展到 >4 原子的分子，MAE 小于 1 Bohr radius（约半个化学键长度）。

背景知识

什么是库仑爆炸成像

步骤	说明
1. X 射线照射	超短脉冲（~15 飞秒）照射分子
2. 电离	原子失去电子，变成带正电的离子
3. 库仑爆炸	正电荷离子互相排斥，分子”爆炸”
4. 探测	测量每个离子碎片的动量方向和大小
5. 反演	从动量分布反推原始分子结构

为什么反演极其困难

挑战	说明
非线性	多体库仑相互作用高度非线性
量子效应	电离过程涉及量子力学
不完整信息	可能检测不到所有碎片
多体动力学	$N$ 个原子有 $N(N-1)/2$ 对相互作用

为什么用扩散模型

分子结构重建是一对多的逆问题——同一组动量可能对应多种构型 → 扩散模型天然适合建模这种多模态后验分布。

方法详解

1. MOLEXA 网络架构（4 模块）

1.1 嵌入模块

将原子信息转换为特征向量：

• 输入：原子序数 $z_i$、电荷态 $q_i$、分子坐标系动量 $p_i$
• 原子嵌入 + 动量线性投影 → 拼接
• 生成成对特征（原子-原子交互）
• 输出维度：384 通道

1.2 动力学提取模块

6 个 Transformer with Memory（TM） 块：

记忆操作（类似 LSTM）：

• 遗忘门：$f_{ij}$（决定丢弃多少旧信息）
• 更新门：$u_{ij}$（决定存储多少新信息）
• 输出门：$o_{ij}$（决定输出多少当前信息）

\[m_{ij} = u_{ij} \odot \tilde{m}_{ij} + f_{ij} \odot m_{ij}\] \[b_{ij} = o_{ij} \odot \tanh(m_{ij})\]

• 32 个注意力头，每头 12 维
• Transformer 的全局注意力 + LSTM 的记忆机制 → 捕获复杂的多体动力学

1.3 结构去噪模块

实现反向扩散过程：

• 条件位置编码器：Fourier 编码噪声级别 + 分子位置
• 2 个 TM 块处理成对特征
• 自注意力块处理原子级特征
• 位置解码器：输出逐渐更干净的结构

\[x_i \leftarrow c_{\text{skip}} \cdot x_i + c_{\text{out}} \cdot y_i\] \[c_{\text{skip}} = \frac{\sigma_d^2}{\sigma_d^2 + \sigma^2}, \quad c_{\text{out}} = \sqrt{\frac{\sigma_d \cdot \sigma}{\sigma_d^2 + \sigma^2}}\]

5 个扩散采样步骤，4 个中间结构。

1.4 不确定性估计模块

• 2 个 TM 块
• 预定义 200 个不确定性 bin（0 到 9.95 Å，步长 0.05 Å）
• 概率加权的不确定性估计

2. 损失函数

\[\mathcal{L} = c_x \cdot \mathcal{L}_x + c_u \cdot \mathcal{L}_u\]

结构损失：

\[\mathcal{L}_x = \mathbb{E}_i \left[w_i \|G_i^{\text{pred}} - G_i^{\text{gt}}\|_2^2\right], \quad w_i = \frac{\sigma_i^2 + \sigma_{\text{data}}^2}{\sigma_i^2 \cdot \sigma_{\text{data}}^2}\]

不确定性损失（交叉熵）：

\[\mathcal{L}_u = -\mathbb{E}_{n,i} \left[q_n^i \cdot \log(s_n^i)\right]\]

3. 两阶段训练

阶段	数据源	样本数	训练时间	GPU
Stage 1：粗糙	经典库仑近似	~600 万	82h	16×A100
Stage 2：精细	从头算（ab initio）模拟	7.6 万	1h	4×A100

Stage 1：用计算廉价但不精确的经典模型生成大量数据 → 学习通用动力学模式

Stage 2：用计算昂贵但精确的量子力学模拟微调 → 修正物理精度

两阶段训练将预测误差降低了 2×。

4. 坐标系标准化

用 Gram-Schmidt 过程：

• 最重离子碎片飞行方向 → x 轴
• 与 x 轴最不平行的离子 → y 轴（正交化）
• z 轴自动确定

→ 消除了旋转/平移不变性的建模需求。

实验结果

测试集性能（<8 原子）

指标	值
RMSE	1.04 a.u.
MAE	0.52 a.u.
准确率（MAE < 0.6 a.u.）	>75%（低不确定性样本）

1 Bohr radius ≈ 0.529 Å，典型化学键 ~1-1.5 Å → MAE < 1 Bohr = 亚键长精度。

按分子大小

原子数	MAE (a.u.)	准确度
2（双原子）	<0.4	100%
3-4	~0.4-0.5	高
5-6	~0.5-0.6	中
7	~0.5-0.7	较低
8-9（泛化）	~0.6-1.0	显著降低

真实实验数据验证

分子	原子数	RMSE (a.u.)	MAE (a.u.)	数据来源
H₂O	3	0.296	0.198	European XFEL
CF₄	5	0.294	0.238	SQS 反应显微镜
C₂H₆O（乙醇）	9	0.524	0.429	9 离子输入

时间分辨化学：环丁烯开环反应

MOLEXA 成功重建了光化学诱导开环反应的结构快照：

• 基态（S₀）：四元碳环
• 最低能量锥形交叉（MECI）
• 扭曲交叉
• 质子迁移

所有重建与从头算基态几何一致。

个人思考

1. “扩散模型解物理逆问题” 是令人兴奋的跨学科应用：分子结构重建的多模态性与扩散模型的概率性天然匹配。
2. 两阶段训练策略巧妙解决了物理 ML 的数据稀缺问题：600 万廉价近似数据学通用模式 + 7.6 万精确数据微调 → 类似 NLP 的预训练-微调范式。
3. Transformer with Memory 的设计非常适合多体物理系统：Transformer 的全局注意力捕获远程相互作用，LSTM 记忆门维护动力学状态。
4. 亚键长精度（MAE < 0.5 Å）是真正实用的标准——化学家可以据此判断分子构型，而非仅作为定性参考。
5. 局限性坦诚：>10 原子的分子仍然困难（1,3-环己二烯”大偏差”）→ 需要更多大分子的从头算训练数据。