[论文解读] On-FPGA Training with Ultra Memory Reduction: A Low-Precision Tensor Method

Kaiqi Zhang, Cole Hawkins+3 more|arXiv (Cornell University)|Apr 7, 2021

Tensor decomposition and applications参考文献 23被引用 10

一句话总结

本文提出了一种低精度、秩自适应的张量化神经网络训练框架，通过结合张量-Train分解、贝叶斯秩自适应与低精度计算，实现了高达292倍的内存缩减，支持在FPGA上进行端到端训练，同时保持极低的精度损失。该方法在Xilinx MPSoC上相较嵌入式CPU实现了59倍的加速与123倍的能效提升，验证了在资源受限边缘设备上实现高效、实时训练的可行性。

ABSTRACT

Various hardware accelerators have been developed for energy-efficient and real-time inference of neural networks on edge devices. However, most training is done on high-performance GPUs or servers, and the huge memory and computing costs prevent training neural networks on edge devices. This paper proposes a novel tensor-based training framework, which offers orders-of-magnitude memory reduction in the training process. We propose a novel rank-adaptive tensorized neural network model, and design a hardware-friendly low-precision algorithm to train this model. We present an FPGA accelerator to demonstrate the benefits of this training method on edge devices. Our preliminary FPGA implementation achieves $59 imes$ speedup and $123 imes$ energy reduction compared to embedded CPU, and $292 imes$ memory reduction over a standard full-size training.

研究动机与目标

• 在资源受限的边缘设备（如FPGA）上实现端到端、节能且实时的神经网络训练。
• 通过实现数量级的内存缩减，解决端到端训练中高内存与计算成本的挑战。
• 开发一种硬件友好的一次性训练方法，避免迭代式秩调优或多轮训练。
• 展示在FPGA上实现端到端训练，所有模型参数均存储于片上，实现隐私保护与低延迟学习。
• 在保持竞争力模型精度的同时，实现高速度与低能耗的高效训练。

提出的方法

• 该方法采用张量-Train矩阵（TTM）分解，通过一系列低秩张量紧凑表示神经网络权重，显著减少参数量。
• 引入一种基于贝叶斯先验的秩自适应训练模型，可在训练过程中自动确定最优张量秩，无需手动选择秩。
• 框架集成低精度优化（定点数运算），进一步降低训练期间的内存与计算开销。
• 设计了一种定制FPGA加速器，包含三个处理单元（PE）：PE1负责沿第一维与最后一维的并行张量收缩，PE2处理中间维度的收缩，PE3以16路并行方式计算外积。
• 硬件设计采用128个乘加（MAC）单元并行处理，通过将张量维度对齐为16的倍数，优化内存访问，提升数据重用效率。
• 训练流水线使用HLS实现，目标为Xilinx MPSoC，支持所有模型参数片上存储，实现FPGA上的全端到端训练。

实验结果

研究问题

• RQ1 具有秩自适应压缩的张量化神经网络是否能实现极低内存占用的FPGA端到端训练并保持高效率？
• RQ2 与标准全精度训练相比，采用低精度张量化训练可实现多大程度的内存缩减？
• RQ3 与嵌入式CPU相比，基于FPGA的训练在端到端学习中能带来多大的加速与能效提升？
• RQ4 所提出的贝叶斯秩自适应方法是否能在无需手动超参数调优或多轮训练的情况下实现高模型精度？
• RQ5 在不造成显著精度下降的前提下，低精度计算在张量化训练中可实现多大程度的集成？

主要发现

• 与标准全精度训练相比，该方法实现了292倍的内存缩减，FashionMNIST测试精度仅下降4.5%。
• FPGA实现相较在1.2GHz ARM处理器的Raspberry Pi 3B上训练，实现了59倍加速与123倍能效提升。
• 模型参数完全存储于片上，实现无需外部存储权重的端到端设备内训练。
• 基于贝叶斯的秩自适应先验可实现自动、一次性模型压缩，无需组合式秩搜索或超参数调优。
• 硬件设计使用了79.55%的LUT、21.37%的FF、77.22%的DSP与17.82%的BRAM，表明资源利用高效。
• 实际中实现了294倍的内存缩减，略高于摘要中提到的292倍，验证了内存效率的稳健性。