Paper Note: The lottery ticket hypothesis: Finding sparse, trainable neural networks

Created2021-01-28Updated2021-01-29Visited

The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks

Jonathan Frankle, Michael Carbin

ICLR 2019

[paper] [code]

Main Idea

文章首先通过一个小实验发现：对原网络进行裁剪，最终的网络越稀疏，该网络学习得就越慢，测试准确率就越低。

在本文中，我们证明了始终存在着较小的子网络，这些子网络从零开始训练，在达到类似测试准确度的同时，其学习速度至少与较大的子网络一样快。据此，我们提出了“彩票假说（the lottery ticket hypothesis）”。

The Lottery Ticket Hypothesis. 随机初始化的密集（dense）神经网络包含一个子网络，该子网络经过初始化以后，在单独训练的情况下可以在训练最多相同数量的迭代次数之后与原始网络的测试精度相近。

更正式地来说，考虑一个密集前向神经网络 f(x;θ)f(x;\theta) ，其中初始化参数为 θ=θ0∼Dθ\theta=\theta_0\sim\mathcal{D}_\theta 。使用SGD对其在训练数据集上优化， ff 在第 jj 次迭代时 minimum validation loss ll 达到了测试准确率 aa。另外，考虑训练 f(x;m⊙θ)f(x;m\odot\theta)，其中 mask m∈{0,1}∣θ∣m\in\{0,1\}^{|\theta|}，故该子模型的初始化为 m⊙θ0m\odot\theta_0。当使用 SGD 在相同的训练数据集上优化该子模型（mm 固定），ff 在第 j′j' 次迭代时 minimum validation loss l′l' 达到了测试准确率 a′a'。“彩票假说”认为 ∃ m\exist\ m for which j′≤jj'\le j (commensurate training time), a′≥aa'\ge a (commensurate accuracy), and ∥m∥0≪∣θ∣\|m\|_0\ll |\theta|(fewer parameters)。

我们发现，一种标准的剪枝技术可以从全连接或卷积前馈网络中自动发现此类可训练的子网络。我们称这种可训练的子网络 f(x;m⊙θ0)f(x;m\odot\theta_0) 为 “中奖彩票（winning ticket）”，因为这些彩票胜过了初始彩票。当“中奖彩票”的参数被重新初始化，即 f(x;m⊙θ0′)f(x;m\odot\theta_0') ，其中 θ0′∼Dθ\theta_0'\sim\mathcal{D}_\theta ，此时我们的中奖彩票不再与原始网络的性能相匹配，这证明了这些较小的网络只有经过适当的初始化才能有效地进行训练。

Identifying winning tickets. 我们通过训练网络并裁剪其 smallest-magnitude 权重来确定“中奖彩票”。其余未裁剪的连接构成了“中奖彩票”的结构。对于我们的工作而言，独特的是，每个未裁剪的连接的值都将在训练之前使用原始网络中的初始值作为现在的初始值。 这构成了我们的中心实验：

  1. 随机初始化网络 f(x;θ0)f(x;\theta_0)，其中 θ0∼Dθ\theta_0\sim\mathcal{D}_\theta ；

  2. 训练该网络 jj 个迭代轮次，参数现为 θj\theta_j；

  3. 裁剪 θj\theta_j 中的参数，裁剪比例为 p%p\%，生成 mask mm；

  4. 重置未裁剪参数的值为 θ0\theta_0 中对应的值，生成“中奖彩票” f(x;m⊙θ0)f(x;m\odot \theta_0)

如上所述，这种剪枝办法是一次性（one-shot）的：网络训练一次，剪掉 p%p\% 的权重，未被剪掉的权重被重置。然而，在本文中，我们关注于迭代剪枝（iterative pruning），这种剪枝办法重复训练、裁剪、重置网络达 nn 个轮次；每个轮次裁剪上一轮余留下权重的 p1n%p^{\frac{1}{n}}\%。我们的结果表明，与一次性裁剪相比，迭代式裁剪可以找到与原始网络的精度相匹配的更小规模（smaller sizes）的中奖彩票。

The Lottery Ticket Conjecture. 我们将假设扩展为一个未经检验的猜想，即SGD会找出并训练一部分初始化良好的权重。密集的、随机初始化的网络比裁剪产生的稀疏网络更容易训练，是因为存在更多可能的可以从中训练恢复出中奖彩票的子网络。

全文结束