论文理解记录：The Lottery Ticket Hypothesis

论文核心

传统非结构化剪枝虽然较大幅度减少了模型参数量，但是由于非结构化的原因导致网络稀疏化，因此很难对剪枝后的模型进行训练。这是非结构化剪枝相对于结构化剪枝不被看好的原因之一。 而本文中提出可以在任意初始化后的原始模型中找到子模块（即非结构化剪枝后的模型），该初始化后的子模块训练时间不会超过原始模型，并达到或超过原始模型的精度，该子模块在论文中被称作“中彩票”。

注意事项

论文中经过实验发现第一次随机初始化至关重要，中彩票的子模块结构和随机初始化参数有关。实验找到子模块后对子模块重新初始化，此时子模块的性能明显下降。对此，论文给出的理由：子模块初始化的参数与优化算法、数据集和模型有关，如中奖彩票的最初初始化的参数特别适合所选择的优化算法，因此优化效果很好。

不足之处

1. 本论文只在较小的数据集上进行相关研究测试
2. 非结构化修剪是论文作者找到唯一的能找到中奖彩票的方法，未能在结构化剪枝的方法中找到中奖彩票
3. 未能明确解释初始化为何对中奖彩票如此重要
4. 在更深的网络中，迭代修剪无法找到中奖彩票，除非用超参数learning rate warmup来调节

寻找中彩票的子模块算法

算法一：

1. 随机初始化神经网络，保存初始化值，并创建掩码
2. 训练迭代
3. 裁剪参数，并更新掩码
4. 重置裁剪后的神经网络权重，使权重值恢复到初始化值
5. 重复2-4步骤，直到得到裁剪充分的网络

算法二：

1. 随机初始化神经网络，保存初始化值，并创建掩码
2. 训练迭代
3. 裁剪参数，更新掩码
4. 重复2-3步骤直到得到裁剪充分的网络
5. 重置裁剪后的神经网络权重，使权重值恢复到初始化值

论文提出两种算法，差别在于第二种每一轮修剪后使用已训练的权重进行重新训练，而第一种在每次重新训练前重置权重，实验证明算法一在训练速度和测试集精度上都要优于第二种

部分代码

# 创建掩码用于模型裁剪
def make_mask(model):
    global step
    global mask
    step = 0
    for name, param in model.named_parameters(): 
        if 'weight' in name:
            step = step + 1
    mask = [None] * step
    step = 0
    for name, param in model.named_parameters(): 
        if 'weight' in name:
            tensor = param.data.cpu().numpy()
            # 把None元素替换为1
            mask[step] = np.ones_like(tensor)
            step = step + 1
    step = 0

# 百分位数裁剪
def prune_by_percentile(percent, resample=False, reinit=False, **kwargs):
        global step
        global mask
        global model
        
        # Calculate percentile value
        step = 0
        for name, param in model.named_parameters():

            # We do not prune bias term
            if 'weight' in name:
                tensor = param.data.cpu().numpy()
                alive = tensor[np.nonzero(tensor)] # flattened array of nonzero values
                # 求百分位值
                percentile_value = np.percentile(abs(alive), percent)

                # Convert Tensors to numpy and calculate
                weight_dev = param.device
                # 小于百分位数的参数设置为0
                new_mask = np.where(abs(tensor) < percentile_value, 0, mask[step])
                
                # Apply new weight and mask
                param.data = torch.from_numpy(tensor * new_mask).to(weight_dev)
                mask[step] = new_mask
                step += 1
        step = 0

完整代码链接

个人总结

本篇论文重点不在于如何进一步压缩模型大小，更多的是提高模型训练速度。在大模型中找到体积小、容易训练且不降低精度的子模型，这对模型的结构探索有很重要的意义。同时论文提出了一些问题有待解决，如为何原始初始化参数对子模型如此重要、在更深的网络中如何寻找到符合要求的子模型等等。