首次在智能手机上训练BERT和ResNet，能耗降35%

来源 | 机器之心

【导读】 研究者表示，他们将边缘训练看作一个优化问题，从而发现了在给定内存预算下实现最小能耗的最优调度。

目前，智能手机和嵌入式平台等边缘设备上已经广泛部署深度学习模型来进行推理。其中，训练仍然主要是在具有 GPU 等高通量加速器的大型云服务器上完成。集中式云训练模型需要将照片和按键等敏感数据从边缘设备传输到云端，从而牺牲了用户隐私并导致了额外的数据移动成本。

图注：推特 @Shishir Patil

因此，为了使用户在不牺牲隐私的情况下个性化他们的模型，联邦学习等基于设备的训练方法不需要将数据整合到云端，也能执行本地训练更新。这些方法已被部署在谷歌 Gboard 键盘上以个性化键盘建议，也被 iPhones 手机用来提升自动语音识别。同时，当前基于设备的训练方法不支持训练现代架构和大模型。在边缘设备上训练更大的模型不可行，主要是有限的设备内存无法存储反向传播激活。ResNet-50 的单次训练迭代所需的内存是推理的 200 多倍。

以往工作提出的策略包括分页到辅助内存和重新实现，以减少云端训练的内存占用。但是，这些方法会显著增加整体能耗。与分页方法相关的数据传输通常需要比重计算数据更多的能量。随着内存预算的缩减，重新实现会以 O(n^2 ) 的速度增加能耗。

在 UC 伯克利最近的一篇论文中，几位研究者表明分页和重新实现是高度互补的。通过对简单操作重新实现，同时将复杂操作的结果分页到闪存或 SD 卡等辅助存储器上，他们能够以最小的能耗扩展有效的内存容量。并且，通过这两种方法的结合，研究者还证明了 在移动级边缘设备上训练 BERT 等模型是可能的 。通过将边缘训练看作一个优化问题，他们发现了在给定内存预算下实现最小能耗的最优调度。

论文地址：https://arxiv.org/pdf/2207.07697.pdf

项目主页：https://poet.cs.berkeley.edu/

GitHub 地址：https://github.com/shishirpatil/poet

研究者提出了 POET（Private Optimal Energy Training），这是一种在内存受限边缘设备上对现代神经网络进行能量最优训练的算法，其架构如下图 1 所示。鉴于为反向传播缓存所有激活张量的成本极高，POET 对激活进行优化分页和重新实现，因而可以将内存消耗最高减少两倍。他们将边缘训练问题重新表述为整数线性程规划（ILP），发现可以通过求解器在 10 分钟内将其求解到最优。