主页（http://www.175shouji.com）：首次在智能手机上训练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 分钟内将其求解到最优。
图注：POET 在边缘设备上对 SOTA 机器学习模型的训练进行优化。
对于部署在真实世界边缘设备上的模型，当边缘设备出现空闲并可以计算周期时就会进行训练，例如谷歌 Gboard 会在手机充电时安排模型更新。因此，POET 也包含了严格的训练限制。给定内存限制和训练 epoch 的数量，POET 生成的解决方案也能满足给定的训练截止期限。此外，研究者还利用 POET 开发了一个全面的成本模型，并证明它在数学上是保值的（即不做近似），适用于现有的开箱即用架构。
论文一作 Shishir Patil 在演示视频中表示，POET 算法可以在智能手机等商用边缘设备上训练任何需要极大内存的 SOTA 模型。他们也成为了首个展示在智能手机和 ARM Cortex-M 设备上训练 BERT 和 ResNet 等 SOTA 机器学习模型的研究团队。
集成分页和重新实现
重新实现和分页是降低大型 SOTA ML 模型内存消耗的两种技术。在重新实现中，一旦不再需要激活张量就会被删除，最常见的是在前向传播期间。从而释放了宝贵的内存，可用于存储后续层的激活。当再次需要删除的张量时，该方法会根据谱系的规定从其他相关的激活中重新计算。而分页，也称为 offloading，是一种减少内存的补充技术。在分页中，不是立即需要的激活张量从主存储器调出到二级存储器，例如闪存或 SD 卡。当再次需要张量时，将其分页。
图 2 显示了一个八层神经网络的执行时间表。沿着 X 轴，每个单元对应神经网络的每一层（共 8 层 L8）。Y 轴表示一个 epoch 内的逻辑时间步长。图中占用的单元（用颜色填充）表示在相应的时间步执行的操作（前向 / 后向传播计算、重新实现或分页）。
例如，我们可以看到 L1 的激活是在第一个时间步 ( T1 ) 计算的。在 T2 和 T3 时刻，分别计算 L2 和 L3 的激活量。假设层 L2 和 L3 恰好是内存密集型但计算成本较低的运算，例如非线性 ( tanH、ReLU 等 ) ，那么重新实现就成为了最佳选择。我们可以删除激活 ( {T3, L2}， {T4, L3} ) 来释放内存，当后向传播过程中需要这些激活时，可以再重新实现它们 ( {T14, L3}， {T16, L2} ) 。
假设 L5 和 L6 层是计算密集型运算，例如卷积、密集矩阵乘法等。对于此类运算，重新实现将导致运行时间和能量的增加，并且这种方式是次优的。对于这些层，最好将激活张量分页到辅助存储（{T6，L5}，{T7，L6}），并在需要时分页到（{T10，L6}，{T11，L5 }）。
分页的一个主要优点是，根据内存总线的占用情况，可以进行 pipelin 处理，以隐藏延迟。这是因为现代系统具有 DMA（直接内存访问）特性，它可以在计算引擎并行运行时将激活张量从辅助存储移动到主内存。例如，在时间步 T7，可以同时将 L6 调出并计算 L7。但是，重新实现是计算密集型的，不能并行化，这导致运行时间增加。例如，我们必须将时间步 T14 用于重新计算 L3，从而延迟其余反向传播执行。
POET
该研究提出了 POET，这是一种用于深度神经网络的图形级编译器，它重写了大型模型的训练 DAG，以适应边缘设备的内存限制，同时保持高能效。
