Zephon Blog
0%

《MEST:Accurate and Fast Memory-Economic Sparse Training Framework on the Edge》

发表于 更新于 分类于
本文字数： 5.8k 阅读时长 ≈ 5 分钟

译：《MEST：边缘上的准确快速内存经济的稀疏训练框架》

论文地址：MEST:Accurate and Fast Memory-Economic Sparse Training Framework on the Edge

从内存占用和数据高效方面进行稀疏训练

Motivation

为了推广深度学习在边缘设备上应用，有很多的研究致力于为加速网络推理而移除神经网络上过量的参数，已经存在的各种工作如：基于启发式方法的剪枝、基于正则化剪枝、神经网络结构搜索等。

而最近的一个新的趋势是探索稀疏性作为神经网络的训练加速以实现更好的边缘训练范式。在这个方向首先的工作方向是使用在初始时剪枝的方法如SNIP、GraSP等，这些方法先获得一个固定的稀疏模型结构，然后使用传统训练方法进行训练；但这样整个过程仍然是计算和内存密集型的，因此，不适合端到端边缘训练范式，这种具有预先固定结构的稀疏训练方法也面临精度折衷问题。

此外，使用动态稀疏掩码的稀疏训练如SET、DeepR、DSR、RigL、SNFS等，对于端到端到边缘训练则有巨大潜力。具体地，它们通过从未训练的密集模型中直接选择一个稀疏模型结构，然后随着稀疏模型训练，在给定稀疏率的情况下启发式探索各种稀疏拓扑。稀疏训练的基本原则是使用与密集训练一样的训练轮数，但每次训练迭代速度有很大改善，因为稀疏性减少了计算量。

和推理加速一样，作者发现稀疏训练所采用的稀疏方案密切相关，如非结构化的、结构化的、细粒度结构化的等待，这些可能导致稀疏训练的精确度、训练速度和内存占用性能的不同。

基于这些，作者提出了他们的Memory-Economic Sparse Training(MEST)框架，并总结论文贡献如下：

  • 提出一个新的MEST框架，并且使用Elastic Mutation(EM)和soft memory bound来增强MEST，实现在边缘上的准确性和快速执行。
  • 系统地研究了稀疏方案在真实边缘设备上的计算精度、内存占用和训练速度，为未来边缘训练模式提供指导
  • 探索了稀疏训练场景下训练数据的效率，提出两阶段稀疏训练方法，在不影响最终精度的情况下，对训练过程中信息量小的样本进行原地识别与去除
  • 通过实验证明了MEST的作用

相关工作

稀疏方案

神经网络剪枝在推理加速上已经实现了很好的研究，这个方向的主要工作使用的”预训练-剪枝-再训练“流，这对于边缘设备训练不兼容。按照所采用的稀疏方案，那些工作可以被分为：非结构化、结构化、细粒度结构化，包括基于模式(pattern-based)的和基于块(block-based)的

尽管这些稀疏方案主要是为加速推理提出的，在稀疏训练中，它们在准确率、内存占用和训练速度方面也起到了重要使用。对于内存占用，重点关注的是随稀疏方案变化的两个主要组件：模型表示(model representation)和在训练过程中产生的梯度(gradient)

稀疏训练

稀疏训练工作主要被分为两组：固定掩码稀疏训练和动态掩码稀疏训练，除此以外，也有一些工作是在早期的训练阶段剪枝密集网络的，但是它们不在边缘稀疏训练的范围内。

固定掩码稀疏训练

固定掩码方法提出将剪枝与训练解耦，使得剪枝后的稀疏模型训练可以在边缘设备上进行。其中：

  • SNIP基于连接敏感性保留剪枝后的损失
  • GraSP根据连接在网络梯度流中的作用进行剪枝
  • SynFlow提出迭代的突触流剪枝，避免了层坍塌并保留了整个网络中突触强度的总流

由于这些剪枝算法没有结合反向传播，它们可以在初始化时进行全局剪枝，而不需要数据。基于非结构化的SNIP目标，3SP进一步引入了剪枝分数的计算感知权重，这样通过移除更多计算密集型通道来主动偏向修剪，这些通道要么是对损失影响很小的，要么是需要显著计算成本的。

然而，初始化时的剪枝过程仍然是计算和存储密集型的，因此与端到端连续稀疏训练不兼容。并且这些工作除了3SP外，都采用了非结构化稀疏方案。

动态掩码稀疏训练

为了在整个训练阶段中降低计算和内存占用，稀疏训练还被研究出很多可以在训练期间调整稀疏拓扑并保持较低内存占用的方法。其中：

  • Sparse Evolutionary Training(SET)采用的是每轮训练结束时基于幅度(magnitude-based)剪枝和随机增长的方法
  • DeepR则是将动态稀疏训练参数化与随机参数更新相结合进行训练，适用于相对较小和较浅的网络
  • DSR使用一种在在训练稀疏深度残差网络时实现高参数效率的动态重新参数化的方法
  • SNFS使用一种稀疏动量算法，这是一种使用指数平滑的梯度(动量)来识别能有效降低误差的层和权重的算法
  • RigL提出通过仅在更新步骤计算密集梯度，在训练期间迭代更新稀疏模型拓扑

其中，尽管SNFS和RigL是稀疏训练，但他们实际上只涉及与剪枝和非零权重相关的所有梯度计算，因此他们的内存占用和密集训练是等价的。

方法简介

在边缘设备上稀疏训练的稀疏方案

稀疏训练工作普遍使用训练FLOP数表示训练速度性能，实际上FLOPS并不能考虑由稀疏数据操作引起的实际执行开销。例如：非结构化稀疏表现出不规则的内存访问模式，会导致显著的执行开销，并且，密集模型可以采用Winograd来显著加速计算速度，这有可能适用于稀疏计算。因为，在经过调查后，作者执行编译器级别的优化，利用基于计算图的编译方法，包括对计算图本身的优化、张量优化等。

作者通过在实际的三星S20智能手机上测量稀疏卷积的总的前向和反向执行时间，然后相对于相应的密集层进行归一化评估加速率。

如图可以看到，当稀疏率低于70%和80%时，基于块稀疏和非结构稀疏方案不可以实现加速，甚至会减速，因此选择一个合适的稀疏方案是稀疏训练加速的重要因素。

提出的Memory-Economic Sparse Traning(MEST)框架

为方便在边缘设备上稀疏训练，MEST根据以下目标设计：

  1. 考虑边缘设备资源限制的端到端内存经济地训练
  2. 在维持高精度的情况下利用稀疏方案实现高稀疏训练加速

基于这些，提出了MEST方法(普通版)：在训练过程中周期性去除不太重要的非零权重，并重新生成零权重，并将这称为变异(mutation)。通过这种方法，探索具有特定稀疏方案和比例的期望中的稀疏拓扑。

之前的工作通常是直接使用权重大小作为重要性指标，但这并不是一个理想方法，因为权重大小可能在训练过程中波动很大。因此，在MEST中，结合了权重当前梯度作为其变化趋势的指标来估计其重要性，定义重要性分数为： \[ Scr_{w_\tau} = |w_\tau|+|\lambda \frac{\partial \ell (W_{\tau-1},D)}{\partial w_{\tau-1}}| \] 其中\(D\)表示训练数据，\(\ell (W_{\tau-1},D)\)和\(\frac{\partial \ell (W_{\tau-1},D)}{\partial w_{\tau-1}}\)分别表示第\(\tau\)轮的损失和梯度，\(\lambda\)是梯度的系数。这样一来，有三种权重被认为是相对重要的：

  1. 权重值大但梯度小的
  2. 权重值小但梯度大的
  3. 权重值大且梯度大的

不同于使用密集模型的梯度来找到要重新生长的权重的方法如RigL，MEST只使用稀疏梯度来区分并移除不那么重要的权重，然后随机选择要重新生长的权重，这样MEST就可以在整个稀疏训练过程中严格保持权重和梯度的稀疏性。这对于在资源有限的边缘设备上进行节省内存的稀疏训练至关重要。此外，与以往仅使用非结构化稀疏的动态稀疏训练方法不同，MEST在变异策略中考虑了不同稀疏方案的约束

弹性突变(Elastic Mutation, EM)

作者提出弹性突变(Elastic Mutation, EM)来随着训练过程逐渐降低变异比例，并将其称为Memory-Economic Sparse Training with Elastic Mutation (MEST+EM)。这主要是基于两个考虑：

  1. 在动态稀疏训练过程中一个大的变异比例(mutation ratio)会提供更大的搜索空间
  2. 网络结构的突然变化可能会影响收敛

因此提出EM，在保持足够的搜索空间的同时，使得稀疏模型平滑稳定到优化的结构

软内存约束弹性突变(Soft memory bound Elastic Mutation, EM&S)

作者提出Soft Memory Bound针对软内存约束增强称为进一小提升准确率的选择，并称为Soft-boundedElastic Mutation (MEST+EM&S)。

EM&S允许将新增长的权重添加到现有权重中，然后进行训练，再从包括新生长的权重在内的所有权重中选择不太重要的权重。如果模型中的现有权重比新增长的权重更重要，这样可以避免强制移除这些权重，这可以被认为是在变异过程中增加了一种“撤销”机制。

算法简单表示

\(W\in \mathbb R^N\)是整个网络的权重，第\(l\)层的权重\(W^l\)的数量是\(N^l\)，目标稀疏比是\(s \in (0,1)\)，\(W^l\)的变异因子是\(p \in (0,1)\)，总的训练轮数是\(\tau_{end}\)，然后对第一个\(\tau_{stop}(<\tau_{end})\)，频率为\(\Delta \tau\)执行权重变异

算法1展示的MEST+EM和MEST+EM&S，二者主要不同在于 ArgRemoveTo(·) and ArgGrowTo(·)的执行顺序。

MEST+EM是按照以下步骤第\(\Delta \tau\)轮执行权重变异：

  1. 首先使用\(ArgRemoveTo(W^l,s+p)\)从总的\(s\times N^l\)非零权重移除\(p\times N^l\)不重要的权重；
  2. 然后使用 \(ArgGrowTo(W^l, s)\)将模型恢复到稀疏\(s\)，这随机激活\(p\times N^l\)个零权重
  3. 新激活的权重虽然是零，但将被视为稀疏模型的一部分并进行训练

在整个训练过程中，模型稀疏性严格以\(s\)为界，从而保持了一个硬内存限制

而MEST+EM&S，对于每\(\Delta \tau\)轮：

  1. 首先使用\(ArgGrowTo(W^l, s-p)\)生长模型以降低稀疏度比到\(s-p\)
  2. 然后训练\(\Delta \tau\)轮
  3. 然后通过\(ArgRemoveTo(W^l,s)\)去除权重来增加稀疏度到\(s\)

也会在给定的范围衰变\(p\)，直到\(\tau_{stop}\)，在整个训练过程中，以稀疏度比\(s−p\)训练权值，通过\(p\)的衰减逐渐增加到目标稀疏率。

此外，变异过程实际上是在索引上操作的，以便于在边缘设备上实现

研究稀疏训练中的数据效率

数据效率在传统训练中已经被研究许多了，已经证明了每个训练样本提供给网络的信息量不同，学习样本的难度也不同。一些训练样本很容易在早期训练阶段就学会，而一旦学会了一些样本，它们就再也不会被“遗忘”(即错误分类)。从训练数据集中移除那些简单的、不那么重要的样本不会导致最终模型的准确率降低。但是，并没有研究稀疏训练场景中的数据效率。由于稀疏训练中的动态稀疏掩码生成仍然不清楚数据效率是否可以进一步加速稀疏训练。因此，作者提出需要先弄清楚模型稀疏性在可移除样本数上的影响，然后讨论利用数据效率加速稀疏训练的可能性。

模型稀疏性对数据集效率的影响

在之前的研究中，有人提出使用训练样本在整个训练过程中的遗忘事件数来表示该训练样本的信息量，而遗忘事件则被定义为个体训练样本在训练过程中从正确分类到不正确分类的转变，这就可以认为是因为这个样本被网络遗忘了，一个样本可能会被遗忘多次。一个不可遗忘的样本表示一个一旦正确分类就永远不会被忘记的样本，它就会被认为对网络而言信息量较少，容易学习。

为研究数据效率是否可以用于加速稀疏训练，作者首先考虑在不同稀疏比下，稀疏训练过程后能识别出的”难忘“样本的数量，结果如Figure 3(a)，可以看到，难忘样本数随模型稀疏度的增加而减少，且与模型精度 呈正相关，在高稀疏情况下，网络泛化性能下降，使得一些简单的样本更难记忆。此外，可以看出更好的稀疏训练算法(如MEST+EM&S)会产生更多”难忘“的训练样本，这说明了通过删除更多训练样本来获得更高加速的可能性。

突变会导致遗忘吗？

为解决这一问题，作者通过变异前后”难忘“样本的数量进行评估，结果如Figure 3(b)所示，不管是MEST+EM还是MEST+EM&S的变异都不会导致明显的遗忘增加，这是因为变异的权重都是不重要的，这对于模型精度的影响很小。

基于边缘的数据高效稀疏训练

为了识别更少信息的训练样本，之前的工作是通过整个训练过程收集遗忘事件数据，然后使用压缩后的数据集从头开始训练网络，很明显，这并不高效，对于边缘设备上训练是不可负担的方法。

不同于之前的工作，这篇论文采用的是将识别更少信息样本和数据高效稀疏训练结合到一个训练过程，目标是在相同训练轮数时获得类似的最终精度。因此，作者提出了一个数据高效的训练方法(DE)，将训练过程分成两个阶段，第一阶段，将完整数据集用于稀疏训练，计数每个样本的遗忘事件，这一阶段需要进行一定的轮数获得一个稳定的识别结果；而第二阶段，部分训练样本就会被移除，得到一个压缩的训练集，将这个压缩训练集用于剩下的训练轮数。而被移除的样本数量取决于达到遗忘事件阈值的样本数量。这里要说明的是只有很少数次被遗忘的样本可能也是相对容易学习的，也可能可以移除而不会损伤正确率。设整个训练数据集\(D\)，压缩后训练数据集为\(\hat D\)，则： \[ \hat D = \{x_i|x_i\in D\quad and \quad f(x_i)\le th\} \] 其中\(x_i\)和\(f(x_i)\)分别表示在第一阶段中完全训练集中第\(i\)个训练样本和它的遗忘事件数量，而\(th\)则是给定阈值。

下图显示了两阶段训练方法在不同稀疏率、稀疏方案下获得的最终精度状态

可以看出：

  1. 与密集模型相比，稀疏模型需要更久来确定一组良好的可移除(信息量较少)的样本
  2. 要被移除的样本越多，在第一阶段需要的训练轮数越多
  3. 非结构化稀疏方案比细粒度稀疏方案需要更少的轮数
  4. MEST+EM和MEST+EM&S需要类似的轮数
  5. 除了”难忘“的样本外，遗忘数很少的样本也是可移除而不会损害最终精度的