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来决定,如下图
策略: shrinkage
如上图 reward
REINFORCE
可以使用带baseline 的来减少variance
Knowledge Distillation
Lkdnote更多细节参考论文。。。
网友评论