↓Skip to main content
  1. Posts/

论文解读：径向核注意力（RKA）——破局聚合物机器学习力场的全局特征交换难题

·10 mins
Table of Contents
近日，深圳鼎犀智创（Rhinovate Technology） 联合 香港中文大学计算机科学与工程系、北京大学材料科学与工程学院 在 arXiv 上发表了最新研究成果：Radial Kernel Attention for Global Feature Exchange in Polymer Machine Learning Force Fields（arXiv:2610.41675v3）。针对高分子聚合物原子级建模中长期存在的“局域截断视野盲区”与“全局注意力平方级算力开销”之间的核心矛盾，我们提出了 径向核注意力（Radial Kernel Attention, RKA） 机制。

在 OPoly26 官方全量验证集（201,865 个高分子团簇结构）上的评测表明：

  • 在相同的 5M 级 EquiformerV3 主干网络中嵌入 RKA（参数量仅从 5.48M 微增至 5.56M），即可将单原子能量 MAE 大幅降低 67.4%，原子力 MAE 降低 36.3%；
  • 完整版 20M + RKA 模型（20.62M 参数）以 不到 1/14 的参数规模，在总能量、单原子能量和原子力三项核心指标上全面超越 Meta FAIR 发布的 UMA-S-1.2 与 UMA-S-1.2.1（291.48M 参数），其中对分子动力学最为关键的原子力预测误差降低了 27.6%（达到 2.619 meV/Å）。

1. 核心痛点：为什么高分子力场不能只靠“局域图网络”？ #

在原子尺度模拟聚合物（Polymers）体系时，链构象的动态演变不仅会改变局部的共价键合几何，还会使原本在序列上相距甚远的链段发生空间接触，并不断重组高分子链与周围分子链、溶剂分子以及离子之间的非键相互作用。

当前主流的机器学习力场（MLFF）——从基于连续距离滤波的 SchNet，到高阶等变消息传递网络 NequIP、MACE 以及 EquiformerV3——大多将计算组织在以截断半径 \(r_c\)（通常为 \(5\sim 6\text{ Å}\)）定义的**局域邻域图（Local Neighbor Graph）**上。这种局域化设计虽然有效控制了单层计算复杂度，但在面对真实的高分子体系时却暴露出两个根本性瓶颈：

  1. 空间尺度远超局域截断半径： 我们对官方 OPoly26 训练集全部 5,902,827 个结构（包含不同化学组分、链拓扑、聚合度及溶剂化环境，单结构最多含 360 个原子）进行了空间跨度分析，统计了每个结构内任意原子对的最大欧氏距离： $$D_g = \max_{i,j \in g} \|r_i - r_j\|_2$$ 统计结果显示，\(D_g\) 的中位数高达 \(23.0\text{ Å}\)，90 分位数更是达到 \(30.7\text{ Å}\)，远远超出了局域主干网络所使用的 \(6\text{ Å}\) 截断半径。仅靠有限层数的局域消息传递，节点感受野既难以覆盖整个团簇，更无法在非连通的子图分量之间建立通信路径。
  2. 纯局域加性模型的“跨片段失明”： 从物理本质上看，考虑一个纯局域加性能量模型（Purely local additive model）作用于两个内部几何构型固定、但彼此所有跨片段原子距离均大于截断半径 \(r_c\) 的非连通片段（例如未成键的链段或溶剂-离子团簇）。此时平移其中一个片段，所有原子的局域邻域输入均保持完全不变，导致模型预测的总能量对两个片段的相对空间排布完全不敏感。然而在真实的物理化学体系中，这类长程相对排布的变化显然会影响体系能量与受力。

要在不引入 \(O(N_g^2)\) 全原子对显式注意力矩阵的前提下打破这一限制，核心挑战在于：能否设计一种兼具几何不变性、线性计算复杂度，且能与等变方向特征无缝协同的紧凑全局交互机制？

2. 方法创新：径向核注意力（Radial Kernel Attention, RKA） #

为了解决上述矛盾，我们提出了 径向核注意力（RKA）。其核心设计哲学是：利用结构质心（Cluster Centroid）作为随刚体变换同步移动的公共参考系，通过可学习的一维径向高斯核重叠（Radial Feature Overlap）定义几何核，并与内容相关的 Query–Key 注意力进行张量因子化融合。

图 1：径向核注意力（RKA）原理示意。(a) 有限团簇定义公共质心与原子相对半径；(b) 一维可学习重叠高斯基函数将半径映射为共享径向特征；(c) 按径向通道累加全局 Key-Value 统计矩阵，各节点通过自身径向权重与 Query 线性读取并输出门控标量残差。

2.1 基于质心的可学习径向特征 #

给定包含 \(N_g\) 个原子的结构 \(g\)，各原子三维坐标为 \(r_i \in \mathbb{R}^3\)，不变标量特征为 \(h_i \in \mathbb{R}^C\)。首先在线性时间内计算结构的公共质心 \(c_g\) 及每个原子的缩放半径 \(\rho_i\)（取固定长度尺度 \(s = 10\text{ Å}\)）：

$$c_g = \frac{1}{N_g}\sum_{j\in g} r_j, \qquad \rho_i = \frac{\|r_i - c_g\|_2}{s}$$

质心提供了一个不依赖任何特定方向偏好的公共参考点。随后，我们使用 \(M\) 个在一维半径轴上可学习的高斯基函数（中心参数 \(\mu_a\)，宽度参数 \(\eta_a\)，实验中取 \(M=16\)）将每个原子映射至共享径向空间并做归一化：

$$\psi_{ia} = \exp\left[-\operatorname{softplus}(\eta_a)(\rho_i - \mu_a)^2\right], \qquad \phi_{ia} = \frac{\psi_{ia}}{\max\left(\sum_{a'=1}^M \psi_{ia'}, \epsilon\right)}$$

由此，任意两个原子 \(i\) 和 \(j\) 之间的几何交互核（Geometric Kernel）自然定义为它们径向特征向量的内积：

$$K_{\text{rad}}(i, j) = \sum_{a=1}^M \phi_{ia}\phi_{ja}$$

这一设计具有两个关键物理特性：

  1. 化学环境区分度：处于相同质心半径壳层的原子虽然拥有相同的几何权重 \(\phi_i\)，但会在下文的注意力机制中通过内容相关的 Query 和 Key 加以区分；
  2. 全局结构响应：由于质心 \(c_g\) 取决于团簇中的所有原子，当体系中任意远端片段发生相对位移时，质心与所有原子的径向分布随之改变，从而在无需局域边连通的情况下建立起全局几何感知通路。

2.2 内容相关的因子化聚合与节点读取 #

为了实现 \(O(N_g)\) 线性复杂度，RKA 改变了传统注意力“先算 \(N_g \times N_g\) 权重矩阵、再加权求和”的求值顺序。首先对输入标量特征做 RMSNorm 归一化以生成非负的查询向量 \(q_i\) 和键向量 \(k_i\)，同时对未归一化特征做线性投影以保留值向量 \(v_i\) 的幅度信息（实验中取 \(D_q = 16, D_v = 32\)）：

$$\bar{h}_i = \operatorname{RMSNorm}(h_i), \quad q_i = \operatorname{softplus}(W_q \bar{h}_i), \quad k_i = \operatorname{softplus}(W_k \bar{h}_i), \quad v_i = W_v h_i$$

随后分两步完成全局特征交换：

  1. 全局统计矩阵聚合：针对结构 \(g\) 和每一个径向通道 \(a \in \{1,\dots,M\}\)，遍历所有节点累加形成一个维度仅为 \(D_q \times D_v\) 的共享 Key–Value 矩阵： $$T_{ga} = \sum_{j \in g} \phi_{ja} k_j v_j^\top \in \mathbb{R}^{D_q \times D_v}$$
  2. 节点自适应读取与门控残差更新：每个原子 \(i\) 利用自身的查询向量 \(q_i\) 和径向权重 \(\phi_{ia}\) 从共享统计矩阵中读取全局信息，并通过节点局域缩放因子 \(d_i\) 归一化： $$y_i = \frac{1}{d_i} \sum_{a=1}^M \phi_{ia} T_{ga}^\top q_i, \qquad d_i = \max\left(\sum_{a=1}^M \sum_{b=1}^{D_q} \phi_{ia} k_{ib}, \epsilon\right)$$

如果将上式代数展开，可以得到其等价的全原子对形式：

$$y_i = \sum_{j \in g} A_{ij} v_j, \qquad A_{ij} = \frac{q_i^\top k_j}{d_i} K_{\text{rad}}(i, j)$$

可以看到，RKA 清晰地将化学内容相似性（\(q_i^\top k_j\)）与径向空间几何核（\(K_{\text{rad}}(i,j)\)）解耦融合，但在实际执行中永远不需要构造显式的全原子对矩阵 \(A\)。其隐式交互矩阵的秩满足 \(\operatorname{rank}(A) \le \min(N_g, M D_q) = \min(N_g, 256)\)，这意味着对于 256 个原子以内的体系它甚至允许满秩交互——线性复杂度完全得益于因子化求值顺序，而非低秩近似妥协。

最后，读取结果通过零初始化的线性投影 \(W_o\) 和可学习通道门控 \(\gamma = 0.25\operatorname{sigmoid}(\lambda)\)（初始值为 \(0.05\)）以残差形式写回标量通道：

$$h'_i = h_i + \gamma \odot W_o y_i$$

3. 架构集成：等变局域主干与严格保守力场 #

我们将 RKA 实例化于 EquiformerV3 等变图神经网络主干中，构建了端到端能量守恒（Conservative）的聚合物机器学习力场。

图 2：RKA 在等变原子势函数中的集成架构。每个局域 EquiformerV3 块更新多阶几何表征后，RKA 读取其 l=0 标量通道与原子坐标，并在进入下一局域块前写回标量残差。总能量结合元素参考能与 ZBL 屏蔽排斥势，通过对坐标求全导数获得严格保守的原子力。
  1. 局域与全局的互补交错： 在每一个截断半径为 \(6\text{ Å}\) 的 EquiformerV3 局域等变块之后插入一个独立参数的 RKA 分支。插入点处，RKA 仅更新 \(\ell = 0\) 的旋转不变标量通道 \(X_{i,0,:}^{(t+1)}\)，而 \(\ell > 0\) 的高阶球谐方向特征保持不变直接透传；在随后的局域等变层与归一化操作中，经过全局更新的标量特征会进一步与高阶方向表征发生耦合。
  2. 严格的对称性与保守力求导（Proposition 1）： 在整体刚体变换 \(r'_i = R r_i + t\) 下，质心同步变换为 \(c'_g = R c_g + t\)，\(\|r'_i - c'_g\|_2 = \|r_i - c_g\|_2\) 严格不变，从而保证了输出能量的平移与旋转不变性。体系总能量与原子力定义为： $$U_\theta(g) = \sum_{i \in g} (u_i + b_{Z_i}) + U_{\text{rep}}(g), \qquad F_{i,\theta} = -\frac{\partial U_\theta(g)}{\partial r_i}$$ 其中 \(b_{Z_i}\) 为固定元素参考能，\(U_{\text{rep}}(g)\) 为 Ziegler–Biersack–Littmark (ZBL) 屏蔽短程排斥势。尤为关键的是，质心 \(c_g\) 与径向特征 \(\rho_i\) 完整保留在自动微分计算图中（不施加任何 stop-gradient），且 RKA 分支在混合精度训练中强制使用 FP32 精度，确保原子力监督信号能够同时、准确地训练局域等变主干与全局径向通路。

4. 实验验证：OPoly26 基准上的“四两拨千斤” #

我们在公开的大规模高分子密度泛函理论（DFT）数据集 OPoly26 上进行了系统评估。所有模型均在官方训练集（5,902,827 个结构）上训练 12 个 Epoch，直接取第 12 个 Epoch 结束时的最终权重，在官方完整验证集（201,865 个结构，零跳过）上评估以下三项平均绝对误差（MAE）：

  • 总能量误差（\(\text{MAE}_U\)，meV/structure）
  • 单原子能量误差（\(\text{MAE}_e\)，meV/atom）
  • 原子力分量误差（\(\text{MAE}_F\)，meV/Å）

4.1 5M 同构主干受控消融：+1.5% 参数带来 67.4% 能量误差下降 #

为了严格剥离并验证 RKA 模块本身的贡献，我们在完全共享 4 层 EquiformerV3 主干（\(\ell_{\max}=2\)，128 通道，包含完全相同的电荷/自旋条件化与 ZBL 短程排斥势）的条件下，对比了 5M baseline（5.48M 参数）与 5M + RKA（5.56M 参数）：

评估指标 (OPoly26 官方全量验证集) 5M Baseline (5.48M) 5M + RKA (5.56M) 误差降幅 (Reduction)
单原子能量 \(\text{MAE}_e\) (meV/atom) 0.4665 0.1521 ↓ 67.4%
总能量 \(\text{MAE}_U\) (meV/structure) 59.44 21.31 ↓ 64.1%
原子力 \(\text{MAE}_F\) (meV/Å) 6.632 4.223 ↓ 36.3%

实验结果令人振奋：仅仅增加了 0.08M（约 8 万）参数，引入 RKA 全局特征交换通路便将单原子能量 MAE 从 0.4665 meV/atom 断崖式降低至 0.1521 meV/atom（降幅达 67.4%），总能量 MAE 降低 64.1%，同时将坐标导数（原子力）MAE 降低了 36.3%。这直接证实了打通超越局域图连通性的全局信息交换对聚合物体系能量与受力预测的关键价值。

4.2 完整 20M 模型对比 Meta FAIR 通用原子大模型 UMA（291M） #

我们进一步构建了 7 层（\(\ell_{\max}=3\)，最大邻居数 300）的完整版 20M + RKA 模型（20.62M 参数），并与 Meta FAIR 发布的两个多数据集联合预训练通用原子势模型 UMA-S-1.2 和 UMA-S-1.2.1（均采用 64 专家混合线性专家 MoLE 架构，总参数量高达 291.48M，且其训练数据包含 OPoly26）在相同的 201,865 个验证集结构上进行对比：

模型名称 参数规模 (M) 单原子能量 \(\text{MAE}_e\) (meV/atom) 总能量 \(\text{MAE}_U\) (meV/structure) 原子力 \(\text{MAE}_F\) (meV/Å)
UMA-S-1.2 (Meta FAIR) 291.48 0.1329 17.03 3.619
UMA-S-1.2.1 (Meta FAIR) 291.48 0.1258 16.09 3.616
20M + RKA (Ours) 20.62 0.1182 15.81 2.619

对比数据清晰表明：

  • 以小博大，全面超越：20M + RKA 以 不到 UMA 1/14 的参数量（20.62M vs. 291.48M），在全部三项指标上优于两个版本的 UMA 模型；
  • 力场精度显著领先：相比 UMA-S-1.2.1，20M + RKA 的单原子能量误差进一步降低 6.0%（较 UMA-S-1.2 降低 11.1%），而在决定分子动力学模拟稳定性的核心指标——原子力 MAE 上，更是从 3.616 meV/Å 大幅降至 2.619 meV/Å，误差降幅高达 27.6%！

5. 总结与展望 #

在高分子机器学习力场中，如何在保持几何对称性与线性计算复杂度的前提下捕获跨越局域截断半径的长程结构信息，是一个兼具理论深度与工程实用价值的核心问题。

径向核注意力（RKA） 提供了一条极其简洁而高效的解决路径：它不需要显式参数化复杂的长程静电或色散修正项，也不需要构造昂贵的 \(O(N_g^2)\) 全原子对注意力矩阵，而是通过基于团簇质心的径向高斯基重叠核与因子化线性注意力，让相距遥远甚至图上不连通的链段、溶剂与离子通过共享径向统计量实现全局通信。无论是 5M 模型中仅用 0.08M 额外参数换取的 67.4% 能量误差骤降，还是 20M 模型对 291M 通用大模型 UMA 的全面反超，都印证了贴合分子体系几何本征的归纳偏置（Inductive Bias）依然是科学智能（AI4M）实现精度与效率双重突破的关键钥匙。

📄 论文信息：
标题：Radial Kernel Attention for Global Feature Exchange in Polymer Machine Learning Force Fields
作者：Wenbo Fu¹, Chenghui Mao¹, Haifeng Lv¹,³, Chengwen Hu¹,*, Zhou Lu¹,²,†
单位：¹ 深圳鼎犀智创科技有限公司（Rhinovate Technology） · ² 香港中文大学计算机科学与工程系 · ³ 北京大学材料科学与工程学院
预印本链接：arXiv:2610.41675v3