N2N Learning: Network To Network Compression via Policy Gradient Reinforcement Learning 笔记

N2N Learning: Network To Network Compression via Policy Gradient Reinforcement Learning

这篇论文是一篇关于使用增强学习算法来进行模型压缩的论文，作者提出了一种使用增强学习来实现在给定的大的网络(teacher)的情况下，实现一个小的高准确率的小的网络(student)。
作者提出了一个假设，将一个teacher网络转化成一个student网络的一系列的步骤是Markov Decision Process(MDP)。简单来说就是，当前的步骤只和有限的之前几步有关系，通常是一步。MDP是增强学习算法的基本假设。
作者将网络的架构作为状态。这样就使得状态空间非常得大，因为任何一个大的网络，比它小的网络都会非常得多。
状态转移：由动作决定，例如：移除某一个卷积，减少fc层的大小等。每一个动作都会将当前模型结构转换为另一个。
增强学习的过程就是学习一个最优的策略来最优化reward，这个reward基于模型的准确率，和压缩率。
根据之前描述的动作，例如改变某一层，就会导致动作空间非常巨大，并且和层数呈指数关系。因此，作者提出了一种解决办法，引入了一种two-stage动作选择机制。首先选择一个macro-scale "layer removal"动作，接着选择一个micro-scale "layer shrinkage"动作。
每一个通过agent的动作选择得到的模型结构，都trained with Knowledge Distillation。
Knowledge Distillation: knowledge distillation指的是训练一个小的网络（student）来模仿大网络（teacher），并且能够在表现上与大网络接近。
如下图：

模型

方法

MDP 马尔可夫过程

作者将在寻找压缩模型一系列过程看作是MDP问题。通常MDP由五元组表示

• state: 状态空间，包含所有能够由teacher网络压缩得到的网络的一个有限集合。
• Action: 动作空间指的是能够将网络转化成另一个网络的动作。在这篇论文中，只有两种类型的动作：layer removal actions 和 layer parameter reduction actions。
• 转移函数: 在这篇论文中，转移函数是一个deterministic函数，也就是说不是随机的，给定当前状态，动作，可到的是一个固定的状态，概率是1.
• Discount Factor，本文折扣因子是1
• Reward：中间状态的reward是0.

Student-Teacher Reinforcement Learning

根据上面描述的MDP，增强学习算法的目的就是学习一个策略能够最优化下面公式：

total reward

中间状态的reward是0的话，这个公式应该就是优化最终状态的reward。
动作的选择分为两个部分：首先选择一组动作，用来判断判读是否留下或者去掉每一层；接着，使用另一个策略选择一组动作，用来决定剩下每一层减少，压缩的程度。
算法描述如下：

算法

Layer Removal

层的去留由一个Bidirectional LSTM policy来决定，如下图

removal

策略：

其中

是hidden state

，

表示层的类型，

表示kernel size，

表示stride，

表示padding，

表示输出层的个数。

Layer Shrinkage

shrinkage

如上图

Reward Function

模型压缩的目的是最大化压缩率的同时，保证一定程度的准确率。
reward公式如下：

reward

表示student模型的relative compression ratio，

表示student模型在validation上的准确率，

表示teacher模型的准确率。

- (#params(student)) / (#params(teacher))
也就是1减去student的参数个数除以teacher的参数个数

Optimization

作者使用了REINFORCE来进行策略的梯度更新，具体公式如下：

REINFORCE

可以使用带baseline 的来减少variance

Knowledge Distillation

Lkd

note更多细节参考论文。。。