Sparse Iso-FLOP Transformations for Maximizing Training Efficiency
Sparse Iso-FLOP Transformations for Maximizing Training Efficiency
- 论文链接：https://arxiv.org/pdf/2303.11525.pdf
- 源码链接：https://hub.nuaa.cf/CerebrasResearch/Sparse-IFT
简介
模型尺寸和训练数据的增加导致了很多深度学习的突破（AlexNet、ResNet、Transformer、GPT、AlphaGo等）。因此训练和部署深度神经网络的计算和内存占用呈指数级增长。为了能够部署大模型，已经引入了多种技术（例如蒸馏、量化、剪枝）来减少推理FLOPs。虽然这些技术提高了推理效率（测试精度与推理FLOPs）但相关训练成本仍然过高。本文中关注于提高DNN训练效率（测试精度与训练FLOPs）。近期的研究已经探索了使用权重稀疏性减少训练中花费的FLOPs。 The lottery ticket hypothesis: Finding sparse, trainable neural networks证明稀疏子网络（称为彩票）在初始化时存在，并且可以进行训练以匹配其原始密集网络的精度。动态稀疏训练试图在单个训练运行中找到最优稀疏自网络。虽然这些方法主要目标时通过用更少的FLOPs达到密集的精度来提高训练效率，但它们性能往往比密集的baseline差，或者依赖更长的训练计划（最多2-5次训练迭代）来缩小差距。因此，这些技术有时甚至需要比训练密集模型更多的FLOPs。与之前研究相比，本文专注于通过稀疏性提高精度同时消耗与密集模型相同的训练FLOPs，从而显示训练效率的提高。本文引入了稀疏Iso-FLOP变换（Sparse Iso-FLOP Transformation，Sparse-IFT）家族，这可以作为DNN密集层的简单替代。这些变换增加了层的表示能力并在没有改变层潜在FLOPs情况下促进最优稀疏子网的发现。例如，使得一层更宽但更稀疏增加了维度但由于稀疏度维持了FLOPs。所有的Sparse-IFT成员由单一超参数-稀疏度参数化。
本文方法
基于密集指标训练是FLOP低效的
先前研究已经显示现代化DNN是过度参数化的，其中在每一层学习的特征和权重是稀疏的。近期研究LTH显示稀疏DNN可以训练到与对应密集DNN相同精度，只要以较好的稀疏掩码开始。这些研究表明DNN最优权重集合是稀疏的。因此在训练时将这些权重表示为密集矩阵是低效的，基于稀疏矩阵训练应该更高效。然而在实践中，大多数稀疏训练方法获得比密集baseline相比更差的精度。本文假设这是因为在单一训练周期内搜索彩票网络的低效。
尽管稀疏模型减少了每步需要的FLOPs，本文假设已有的稀疏训练方法使得计算。例如SOTA稀疏训练方法Top-kast: Top-k always sparse training；Rigging the lottery: Making all tickets winners将这些FLOPs节省的投入到更长的训练计划中以缩小精度差距，并弥补在训练早期无法发现最优掩码的问题。这种设置是低效的，因为它最终需要比密集baseline更多的训练FLOPs才能达到相同目标精度。在本文研究中，采用正交方法，并将这些FLOPs节省投入到1. 增加层的表示能力和2. 增加搜索空间，假设这可以促使最优稀疏掩码发现。本文通过用FLOPs等效稀疏变换来实现这一点。本文将这些变换表示为稀疏Iso-FLOPs变换（Sparse-Iso）族。
设定
简洁起见，将本文方法解释为一个全连接层神经网络。定义 N \mathcal{N} N为N层由 Θ N \Theta_{\mathcal{N}} ΘN参数化的DNN网络。定义 Θ N ∈ { θ 1 , … , θ L } \Theta_{\mathcal{N}}\in \{\theta_{1},\ldots,\theta_{L}\} ΘN∈{θ1,…,θL}是DNN的参数。l层输出定义为： z l = σ ( f θ l ( z l − 1 ) ) z_{l}=\sigma(f_{\theta_{l}}(z_{l-1})) zl=σ(fθl(zl−1))。 σ \sigma σ 是激活函数， f θ l f_{\theta_{l}} fθl是前向函数。定义 f θ l ( z l − 1 ) = θ l T z l − 1 f_{\theta_{l}}(z_{l-1})=\theta_{l}^{T}z_{l-1} fθl(zl−1)=θlTzl−1。 f θ l f_{\theta_{l}} fθl需要的总FLOPs表示为 B ⋅ D i n ⋅ D o u t B\cdot D_{in}\cdot D_{out} B⋅Din⋅Dout。
稀疏Iso-FLOP变换
标准设置中，前向函数 f θ l f_{\theta_{l}} fθl将输出特征计算为输入特征的线性变换。从理论角度看，前馈函数可以利用任意的非线性变换。然而在实践中，由于GPU广泛支持，大多数变换都表示为密集矩阵乘法。
如前所述，本文有兴趣通过增强前馈的代表能力来提高DNN训练效率。通过通过堆叠更多的层、增加宽度、专家混合等增加表示能力会增加计算FLOPs。在本文研究中，本文在权重矩阵中使用非结构化稀疏性，并确保变换的FLOPs与密集前馈函数的FLOPs相同。定义 Ψ l \Psi_{l} Ψl是对于特定层l的稀疏Iso-FLOP集合：
Ψ : { ψ l ( s ) , 0 ≤ s ≤ 1 , g ( ψ l ) ≈ g ( f θ l ) } \Psi:\{\psi_{l}(s),0\leq s\leq 1,g(\psi_{l})\approx g(f_{\theta_{l}})\} Ψ:{ψl(s),0≤s≤1,g(ψl)≈g(fθl)}
ψ \psi ψ是一个变换， s s s表示稀疏度。该集合中每个变换满足以下属性：1. 变换 ψ l \psi_{l} ψl计算FLOPs与密集变换 f θ l f_{\theta_{l}} fθl的相同。2. 变换由单一超参数（稀疏度级别）参数化。因为这些变换是对于密集前向函数的Iso-FLOP，本文可以在不影响层的FLOPs情况下使用它们作为临时替换。虽然许多FLOPs等效变换属于稀疏IFT，但本文研究中，详细介绍了四个不同的成员：稀疏宽、稀疏并行、稀疏因子化（Sparse Factorized）和稀疏掺杂（Sparse Doped）。
稀疏IFT成员
稀疏宽
稀疏宽变换通过增加输出特征数量来增强层表示能力，同时保持权重 s s s部分稀疏。当使用这种变换，本文使用相同加宽因子 k s w k_{sw} ksw加宽网络所有 L L L层的输入和输出特征，以避免跨层的特征维度不匹配。定义 θ l s w ∈ R k s w ⋅ D i n ⋅ k s w ⋅ D o u t \theta_{l}^{sw}\in \mathbb{R}^{k_{sw}\cdot D_{in}\cdot k_{sw}\cdot D_{out}} θlsw∈Rksw⋅Din⋅ksw⋅Dout是变换矩阵，s比例的权重是稀疏度。由于非稀疏权重分数由1-s给出，因此该变换所需的FLOPs为 B ⋅ ( k s w ⋅ D i n ) ⋅ ( k s w ⋅ D o u t ) ⋅ ( 1 − s ) B\cdot (k_{sw}\cdot D_{in})\cdot (k_{sw}\cdot D_{out})\cdot (1-s) B⋅(ksw⋅Din)⋅(ksw⋅Dout)⋅(1−s)。将这些设置等于原始密集 f θ l f_{\theta_{l}} fθl 的FLOPs，可以获得加宽因子 k s w = 1 ( 1 − s ) k_{sw}=\sqrt{\frac{1}{(1-s)}} ksw=(1−s)1。如果设置稀疏度s为0，获得 k s w k_{sw} ksw为1并恢复原始密集前向函数。
稀疏并行
稀疏并行变换取代前向函数为 k s p k_{sp} ksp非线性函数的和。定义 θ l s p ∈ { θ l s p , 1 , … , θ l s p , k s p } \theta_{l}^{sp}\in \{\theta_{l}^{sp,1},\ldots,\theta_{l}^{sp,k_{sp}}\} θlsp∈{θlsp,1,…,θlsp,ksp}为这些变换的参数。这时稀疏并行变换表示为 ψ l s p = ∑ j = 1 k s p σ ( ( θ l s p , j ) T z l ) \psi_{l}^{sp}=\sum_{j=1}^{k_{sp}}\sigma((\theta_{l}^{sp,j})^{T}z_{l}) ψlsp=∑j=1kspσ((θlsp,j)Tzl)。实践中 ψ l s p \psi_{l}^{sp} ψlsp由 k s p k_{sp} ksp平行分支的层实现。该变换的计算FLOPs是 k s p ⋅ B ⋅ D i n ⋅ D o u t ⋅ ( 1 − s ) k_{sp}\cdot B\cdot D_{in}\cdot D_{out}\cdot (1-s) ksp⋅B⋅Din⋅Dout⋅(1−s)。设定变换FLOPs与 f θ f_{\theta} fθFLOPs相同，可以获得 k s p = 1 ( 1 − s ) k_{sp}=\frac{1}{(1-s)} ksp=(1−s)1。当 s = 0 s=0 s=0时并行分支数量 k s p k_{sp} ksp设置为1。如果用恒等变换代替非线性函数 σ \sigma σ，可以恢复原始的密集前向变换。
稀疏因子化
前向函数 f θ l f_{\theta_{l}} fθl的变换矩阵记为 θ l ∈ R D i n × D o u t \theta_{l}\in \mathbb{R}^{D_{in}\times D_{out}} θl∈RDin×Dout。多个研究已经探索了矩阵分解技术以表示变换矩阵 θ l \theta_{l} θl为两个矩阵乘积 θ l = U V T \theta_{l}=UV^{T} θl=UVT，其中 U ∈ R D i n × d , V ∈ R D o u t × d U\in \mathbb{R}^{D_{in}\times d},V\in \mathbb{R}^{D_{out}\times d} U∈RDin×d,V∈RDout×d。Initialization and regularization of factorized neural layers和Drone: Data-aware low-rank compression for large nlp models已经探索了低秩因子分解 d ≪ d o u t d\ll d_{out} d≪dout作为一种结构化稀疏性的形式，以提高训练和推理效率。相反，本文使用因子分解增加表示能力，且不降低或增加FLOPs。定义 θ l s f ∈ { U l , V l } \theta_{l}^{sf}\in \{U_{l},V_{l}\} θlsf∈{Ul,Vl}为该变换参数。 U l , V l U_{l},V_{l} Ul,Vl是稀疏分数为s的稀疏矩阵。此时该功能变换是 ψ l s f = V l T σ ( U l T z l ) \psi_{l}^{sf}=V_{l}^{T}\sigma(U_{l}^{T}z_{l}) ψlsf=VlTσ(UlTzl)。此时变换的计算FLOPs是 d s f ⋅ B ⋅ ( D i n + D o u t ) ⋅ ( 1 − s ) d_{sf}\cdot B\cdot (D_{in}+D_{out})\cdot (1-s) dsf⋅B⋅(Din+Dout)⋅(1−s)。设定该FLOPs等于 f θ l f_{\theta_{l}} fθl的FLOPs，可以获得 d s f = D i n D o u t ( D i n + D o u t ) ⋅ ( 1 − s ) d_{sf}=\frac{D_{in}D_{out}}{(D_{in}+D_{out})\cdot (1-s)} dsf=(Din+Dout)⋅(1−s)DinDout。设定稀疏度 s = 0 s=0 s=0时，恢复了一个具有密集矩阵的非线性低秩因子分解。
稀疏掺杂
受到Scatterbrain: Unifying sparse and low-rank attention approximation，Doping: A technique for efficient compression of lstm models using sparse structured additive matrices等启发，稀疏掺杂基于低秩分解和稀疏矩阵的组合估计密集矩阵。本文中稀疏掺杂即将前向函数取代为低秩分解（秩 d s d d_{sd} dsd）和非结构化稀疏权重矩阵（稀疏度s）。定义 U l , V l U_{l},V_{l} Ul,Vl为低秩矩阵与 θ l s d \theta_{l}^{sd} θlsd是非结构化稀疏度的矩阵。函数变换可以表示为： ψ l s d = V l T ( U l T z l ) + σ ( ( θ l s d ) T z l ) \psi_{l}^{sd}=V_{l}^{T}(U_{l}^{T}z_{l})+\sigma((\theta_{l}^{sd})^{T}z_{l}) ψlsd=VlT(UlTzl)+σ((θlsd)Tzl)。与该变换相关的计算FLOPs为 B ⋅ d s d ⋅ ( D i n + D o u t ) + ( 1 − s ) ⋅ B ⋅ D i n ⋅ D o u t B\cdot d_{sd}\cdot (D_{in}+D_{out})+(1-s)\cdot B\cdot D_{in}\cdot D_{out} B⋅dsd⋅(Din+Dout)+(1−s)⋅B⋅Din⋅Dout。设定该FLOPs等于 f θ l f_{\theta_{l}} fθlFLOPs可以获得 d s d = s ⋅ D i n ⋅ D o u t ( D i n + D o u t ) d_{sd}=\frac{s\cdot D_{in}\cdot D_{out}}{(D_{in}+D_{out})} dsd=(Din+Dout)s⋅Din⋅Dout。当 s → 0 , d s d → 0 s\rightarrow 0,d_{sd}\rightarrow 0 s→0,dsd→0，变换的低阶组件消失，可以通过将 σ \sigma σ设置为相等变换。
搜索空间基数
本文假设是通过稀疏IFT增加稀疏性掩码的搜索空间可以使训练更高效。过去的研究也支持了这一假设。What’s hidden in a randomly weighted neural network?证明在随机初始化网络中找到彩票的几率随着网络宽度而增加。The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training与Search spaces for neural model training表明通过增加宽度或深度来增加搜索空间提高了精度。在本文中将搜索空间基数定义为稀疏训练方法可以探索的权重数量。表1描述了稀疏IFT族每个成员的搜索空间的基数。
