清华大学集成电路学院——面向片上学习系统的 ADAM 算法加速器

关注我们，为您推送更多最新资讯。

文章导读

近年来，片上学习系统成为了重要的研究热点。以自适应矩估计 (Adaptive moment estimation, ADAM) 为代表的二阶优化器在神经网络训练中具有显著的收敛优势。如图1所示，此类优化器的实现需要计算二阶矩，依赖于复杂的平方、开方、除法等运算，因此大幅提高了片上部署的时间、空间复杂度。针对此问题，来自清华大学集成电路学院的张春研究员及其团队在Electronics期刊“Artificial Intelligence Circuits and Systems (AICAS)”栏目中发表了文章，提出了Hardware-ADAM这一包含两种ADAM算法的加速器设计。

图1.ADAM算法原理图。

研究过程与结果

首先，为减小ADAM算法所涉及到的平方、开方运算产生的资源、时间开销，作者提出了Efficient-ADAM (E-ADAM) 优化器，并引入了近似运算，如图2所示。研究中所引入的近似最大误差为5.72%，能够满足神经网络训练过程中的精度需求。同时，如图3所示，对于定点数形式的运算，此近似在低位宽的情况下，能够避免出现绝对误差为100%的计算“盲区”，从而使E-ADAM算法具有比原算法更广的输入范围。在E-ADAM的基础上，为进一步提高系统吞吐率，作者引入快速平方根倒数原理并应用于除法中，提出了Fast-ADAM (F-ADAM)，对系统中的除法运算进行了优化，如图4所示，将除法运算转化为简单运算的组合。

图2.对平方与开方运算的近似公式及证明。

图3.近似算法 (a) 与原始算法 (b)、(c) 和 (d) 的相对误差分析。

图4.根据快速平方根倒数算法原理设计的除法运算单元。

为验证设计的功能可行性与性能水平，作者使用Vivado 2018.3高层次综合工具对所提出的设计进行部署，部署平台为Ultra96-v2开发板。训练环境为PyBullet Gymperium提供的强化学习环境HalfCheetah。图5展示了不同优化器了在强化学习训练的一个Epoch中的收敛情况。图6展示了不同优化器在整个HalfCheetah训练过程中的训练奖励增长情况。从图5与图6中可知，经过近似的E-ADAM与F-ADAM在功能上与原始ADAM算法近似，均好于其他非二阶优化器。图7展示了作者所提出的E-ADAM与F-ADAM相对于原始ADAM算法提升的吞吐率，指标为单个运算模块每秒内能够更新的参数量。图8展示了提出的E-ADAM与F-ADAM资源占用情况，提出的设计相较于原始设计和参考文献均有明显提升。

图5.单次Epoch中不同优化器的收敛情况。

图6.整体训练过程中不同优化器的收敛情况。

图7.所提出优化器性能提升示意。

图8.所提出优化器的资源占用情况。

研究总结

本研究提出了应用于片上学习系统的Hardware-ADAM设计，并验证了所提出设计的功能与性能。其中E-ADAM相比于原始ADAM节约了90%的资源占用，并达到了2.8倍的吞吐率；F-ADAM实现了16.4倍的吞吐率，同大幅节约了寄存器、查找表的资源占用。针对不同需求应用E-ADAM与F-ADAM能够在片上学习任务中有效减少优化器部分产生的时间、资源开销。同时，本研究所引入的运算近似思想，包括对平方、开方算法的改进，以及对快速平方根倒数算法的化用，能够作为其他算法的硬件部署的参考，应用于各种算法的硬件设计。

文章链接：https://www.mdpi.com/2050604