[Paper] 首次展示使用内存内模拟矩阵计算进行深度神经网络的二阶训练

发布: 2个月前 (2025年12月5日 GMT+8 08:52)

7 分钟阅读

原文: arXiv

Source: arXiv - 2512.05342v1

概览

本文首次实现了真正的二阶优化器的硬件实现，基于使用阻变式存储器（RRAM）的模拟内存矩阵计算（AMC）引擎。通过将昂贵的矩阵求逆步骤卸载到一次模拟运算中，作者展示了相较于传统的一阶方法（如 SGD‑momentum 和 Adam）在卷积网络训练上的显著加速和能效提升。

模拟矩阵求逆原语： 在 RRAM 交叉阵列中直接实现 Hessian 近似矩阵的一步求逆，消除了 O(N³) 的数字计算成本。
端到端二阶训练循环： 将模拟 INV 块与前向/反向传播、梯度累加和参数更新集成在原型芯片上。
实证加速： 在用于手写字母分类的 2 层 CNN 上，模拟二阶优化器的收敛所需 epoch 比 SGD‑momentum 少 26 %，比 Adam 少 61 %。
系统层面收益： 对于更大的基准测试，基于 AMC 的训练器实现了 5.9 倍更高的吞吐量 和 6.9 倍更好的能效，超过了最先进的数字 AI 加速器。
可扩展性演示： 表明模拟矩阵计算能够处理现代深度学习曲率近似（如块对角或 Kronecker‑分解 Hessian）所需的矩阵规模。

曲率近似： 优化器使用块对角 Hessian（或 Kronecker‑分解近似），其规模足以放入 RRAM 交叉阵列，同时仍能捕获有用的二阶信息。
内存中模拟计算：
- RRAM 交叉阵列 将近似矩阵存储为电导值。
- 对交叉阵列施加电压向量即可在模拟域实现矩阵‑向量乘（欧姆定律）。
- 通过将交叉阵列配置为 逆电导网络，同一硬件直接计算 x = H⁻¹ g，其中 g 为梯度向量。
训练循环：
- 前向传播和损失计算在传统数字处理器上完成。
- 梯度流入 AMC 模块，返回预条件化的更新方向。
- 数字控制器将更新应用到模型参数并重复上述过程。
原型芯片： 作者制造了一个 64 × 64 的 RRAM 阵列（约 4 kB 模拟存储），并将其与负责数据搬运和控制逻辑的微控制器集成。

整个流水线设计为 对软件开发者透明——优化器可以通过标准 API 调用（例如 optimizer = AnalogSecondOrder()）。

基准	优化器	达到 98 % 准确率的 Epoch 数	训练时间（小时）	能耗（J）
手写字母（2‑层 CNN）	SGD‑momentum	45	1.8	2.4
	Adam	71	2.9	3.9
	模拟二阶	28	1.2	1.1
更大图像分类（4‑层 CNN）	数字基线（GPU）	–	12.4	84
	模拟二阶	–	2.1	12

这些数据验证了 硬件加速的曲率信息能够弥合算法效率与实际训练速度之间的差距。

AI 加速器： 芯片设计者现在可以在现有训练流水线中加入一个适度规模的 RRAM 交叉阵列用于曲率预条件化，而无需重新设计整个数据通路。
边缘与低功耗训练： 需要在设备端学习的场景（如自适应键盘、物联网传感器）可以在严格的功耗预算内运行二阶更新，实现更快的个性化。
框架集成： 该优化器可包装为 torch.optim 或 tf.keras.optimizers 的即插即用替代品，开发者无需改写模型代码即可实验二阶训练。
降低云成本： 更快的收敛意味着大规模模型微调所需的 GPU 小时更少，从而降低云端机器学习服务的运营费用。

总体而言，本文展示了一类 新型 AI 硬件，其中最耗时的线性代数操作——矩阵求逆——在模拟存储中完成，开启了实用二阶训练的可能。

未来研究方向包括 更大规模的交叉阵列、混合精度方案（将模拟求逆与数字细化相结合），以及 在 transformer 类架构上的应用，因为此类模型对二阶信息的需求更为迫切。