Page 189 - 《软件学报》2025年第9期

4100 软件学报 2025 年第 36 卷第 9 期

2.1 层重组卡尔曼滤波的交替训练方法
观察 1 中发现的损失函数的带权系数选取问题, 一种很自然的想法是交替更新, 把联合训练拆分成交替训练,
即先按照能量的误差进行权重更新, 再按照原子受力的误差进行一次权重的更新. 交替训练可避免联合训练中敏
感的损失函数系数的选取问题. 尤其是当目标函数所需拟合的任务增加时, 联合训练的损失函数中各个误差项的
权重的组合呈现出指数增长的选取情况时, 交替训练可以有效规避这个问题. 与此同时, 我们发现在层重组卡尔曼
滤波优化器的训练过程中, 交替训练能比联合训练的误差降得更低, 如图 3(b) 所示. 图中测试数据为铜体系, 分别
在联合训练和交替训练的情况下, 前 30 个 epoch 能量和受力的均方根误差的值.
交替训练带来精度上的收益的同时, 带来了一次额外的权重更新的开销, 进而导致同一样本的训练时间变长.
公式 (7) 表示交替训练比联合训练多花的时间在一次权重更新上 T w , 单个样本联合训练和交替训练的具体时
update
间构成如公式 (8) 和公式 (9) 所示. 接下来介绍一个样本一次完整的计算时间的组成, 整体时间由 3 部分计算过程
构成: (1) 前向计算的时间, 用 T fwd 表示, 结合其上标 E+F, E, F 分别表示模型正向过程计算出能量和力、能量、力
的时间. 且不同上标的 T fwd 可由各自的模型前向计算的浮点操作次数/机器计算效率得到, 计算量由 C fwd 表示, 机
器计算效率记为 FLOPS (由机器唯一确定); (2) 层重组卡尔曼滤波的计算时间, 用 T kf 表示, 结合其上标 E+F, E, F

分别表示基于能量和力、能量、力进行卡尔曼滤波算法计算的时间. 同样的, 不同上标的 T kf 可由各自的任务的卡
尔曼滤波更新操作的浮点操作次数/机器计算效率得到, 计算量由 C k f 表示; (3) 权重更新的时间 T w .

update
∆T 交替-联合 = T 交替 −T 联合 = T w (7)
update
E
C E +C F C +C F
w
E+F
T 联合 = T fwd +T kf E+F +T w update = fwd fwd + kf kf +T update (8)
FLOPS FLOPS
C E C E C F C F
w
F
F
E
T 交替 = T E fwd +T +T w update +T fwd +T +T update = fwd + kf +T w update + fwd + kf +T w update (9)
kf
kf
FLOPS FLOPS FLOPS FLOPS
总体来说, 交替训练能有效规避联合训练各部分损失的带权相加问题, 同时不可避免地引入一次额外的权重
更新的时间消耗. 鉴于神经网络力场训练的特殊性, 能量的预测是前向的神经网络计算过程, 基于能量进行层重组
的卡尔曼滤波优化器的权重更新涉及一阶导数的计算, 单原子受力的预测值包含一阶导数的计算 (因原子受力由
原子总能量对单个原子位置求导得到), 因此基于力场进行层重组的卡尔曼滤波优化器的权重更新涉及二阶导数
的计算. 高阶导数的计算开销通常更大 (因为二阶导数的计算需在一阶导数的基础上再进行求导操作). 因此公式 (9)
T , 联合训练相比于交替训练引发的一次额外的权重的更新操作是整体时间的很小一部
F
的主导的时间开销来自 kf
分, 也即交替训练的额外时间开销较小. 实际上, 在铜数据集下, 基于层重组卡尔曼滤波优化器的联合训练耗费
20 s/epoch, 交替训练需要 23 s/epoch, 此数据是在 NVIDIA A100 机器上测试得到.

2.2 层重组卡尔曼滤波的分块阈值建模
观察 2 是一阶优化器普遍性存在的问题. 卡尔曼滤波优化器天然具有参数少的特性, 值得一提的是, 层重组卡
尔曼滤波算法在卡尔曼滤波算法的基础上增加了分块大小的调整. 为了减少分块大小的经验性设置 (在 Hu 等人 [35]
的工作中暂未给出分块大小的选取方案), 本文对分块大小建立了模型.
给定一个神经网络, 不妨把第 p 层的神经网络的权重数量记为 n p . 卡尔曼滤波算法用于神经网络权重更新时,
需把神经网络每一层的权重分别张成一个列向量 (如图 4(a) 和 (b) 子图中最左侧橙色块所示). 假设神经网络共
1 层, 不失一般性, 我们把神经网络的各层权重数量, 按从前往后的顺序依次记为: ( n 1 ,...,n s ,...,n e ,...,n p ,...,n l ), 总
的参数量用 N 表示. 全局的卡尔曼滤波需要把权重聚合成一个 (N ×1) 的列向量, 在利用卡尔曼滤波算法进行权重
(N × N) 的稠密的误差协方差矩阵 P N 过大时, P 矩阵本身
更新时, 一个 会参与计算和迭代. 当神经网络的参数量
及算法计算过程中 P 矩阵的中间变量的读写均会引发过高的内存占用. 为了减少由稠密 P 矩阵及其中间变量带
来的内存占用压力, 层重组的卡尔曼滤波算法中的 P 矩阵 (图 4(b) 所示) 对全局卡尔曼滤波的稠密 P 矩阵做了一
N b ), 分别进行聚合和拆分操作.
个对角块的解耦合操作. 具体地, 按照给定的分块阈值 (blocksize, 记为
聚合操作: 对于聚合连续若干层 ( n s ,...,n e ), 满足聚合后的权重数量小于预设的块大小的阈值, 且如果再多聚