NAS 学习笔记(十七)- SNAS

  |     |   本文总阅读量:

版权声明:本文原创,转载请留意文尾,如有侵权请留言, 谢谢

引言

  本文分享一篇商汤发表在 ICLR 2019 上关于 NAS 的论文 [1],文中提出了 SNAS 算法,它解释了 ENAS 利用强化学习去搜索带来的收敛慢的原因,通过对NAS进行重新建模,直接通过梯度优化NAS的目标函数。

Motivation

   作者首先解释了 NAS 和 ENAS 利用 MDP 结合强化学习收敛慢的原因,因为 NAS 是一个确定环境中的完全延迟奖励的任务,这里的延迟奖励指的是 RL 的 controller 要等整个架构 search 完毕然后在验证集上 evaluate 才能得到一个 accuracy 作为一个 reward,也就是:

\[ r_{t}=\left\{\begin{array}{ll} 0, & s_{t} \neq \text {terminal} \\ \text {acc,} & s_{t}=\text {terminal} \end{array} .\right. \]

  延迟奖励会指数级延长TD的收敛需要的更新次数,延迟奖励会给指数级多的状态的MC价值评估带来抖动。

SNAS

Search Space and Architecture Sampling

  如上图所示,搜索空间可以通过一个 onehot 矩阵来表示,每条边上的 op 可以随机产生,计算一个新的节点时可以通过下面的公式,和 DARTS 有点类似:

\[ x_{j}=\sum_{i<j} \tilde{\boldsymbol{O}}_{i, j}\left(x_{i}\right)=\sum_{i<j} \boldsymbol{Z}_{i, j}^{T} \boldsymbol{O}_{i, j}\left(x_{i}\right) \]

  \(Z_{i,j}\) 表示在边 \((i,j)\) 上的 one-hot 随机变量,从母网络中产生子网络,可以通过在母网络的每一条边的所有可能神经变换的结果后乘上一个 one-hot 向量来实现。而对于子网络的采样,就因此自然转化为了对一系列one-hot随机变量的采样。
  同时为了解决 MDP 中延迟奖励的问题,作者用 loss function 来替代 accuracy,不需要额外拟合一个 score function,NAS 问题的 score 就已经不是一个来自环境的常数而是一个可微函数了,这可以大幅提高 NAS 的搜索效率,又因为损失函数和准确率都可以表达一个网络学习的结果,这一替换并没有在本质上改变 NAS 问题原本的优化网络结构分布以使得它们的期望性能最好的目标:

\[ \mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}[R(\boldsymbol{Z})]=\mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}\left[L_{\boldsymbol{\theta}}(\boldsymbol{Z})\right] \]

Reparameterization and Differentiable

  之前提到的 \(Z_{i,j}\) 是离散的,作者利用了 reparameterization 和 softmax 来近似化可微:

\[ \begin{aligned} \boldsymbol{Z}_{i, j}^{k} &=f_{\boldsymbol{\alpha}_{i, j}}\left(\boldsymbol{G}_{i, j}^{k}\right) \\ &=\frac{\exp \left(\left(\log \boldsymbol{\alpha}_{i, j}^{k}+\boldsymbol{G}_{i, j}^{k}\right) / \lambda\right)}{\sum_{l=0}^{n} \exp \left(\left(\log \boldsymbol{\alpha}_{i, j}^{l}+\boldsymbol{G}_{i, j}^{l}\right) / \lambda\right)} \end{aligned} \]

  为了实现离散分布,作者用到了 Gumbel-max,它先采样与 one-hot vector 维度相同数量的 uniform distribution 的随机变量,将他们经过 Gumbel 变换转为 Gumbel 随机变量,并从中选择最大的那一维度为1,其他维度为 0,这样采样的随机变量的分布与该离散分布相同,而离散分布的参数也就转化为了 Gumbel max 中的参数,实现了对该离散分布的重参数化。\(\lambda\) 为 softmax 的 temperature,当它趋近于 0 时,该方法产生的随机变量趋近于该离散分布。

Resource Constraint

  作者还通过在目标函数里添加模型复杂度的惩罚项来限制模型的规模,尽可能的搜到比较稀疏的模型:

\[ \mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}\left[L_{\boldsymbol{\theta}}(\boldsymbol{Z})+\eta C(\boldsymbol{Z})\right]=\mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}\left[L_{\boldsymbol{\theta}}(\boldsymbol{Z})\right]+\eta \mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}[C(\boldsymbol{Z})] \]

  惩罚项的量值可以包括参数量、浮点计算数(FLOPs)以及需要的内存,采样出的子网络的这些值的总量计算与每一条边有关,因而相较于在每一条输入边上优化一个全局的网络前向传播的,我们只需要优化每条边上自己对时延的贡献量。如果回到之前贡献分配的语境,全局的时延惩罚被线性分配到了每一条边的决策 \(Z_{i,j}\) 上,这有利于提高收敛效率:

\[ C(\boldsymbol{Z})=\sum_{i, j} C\left(\boldsymbol{Z}_{i, j}\right)=\sum_{i, j} \boldsymbol{Z}_{i, j}^{T} C\left(\boldsymbol{O}_{i, j}\right) \]

  又因为上式是一个线性的变换,我们既可以用重参数化计算的 \(\mathbb{E}_{z \sim p_{\alpha}}[C(z)]\) 期望,也可以用策略梯度的方法。

与 DARTS 的区别

  这篇文章的里有许多数学推导,我并没有细看,我看到中间最大的疑问是它和 DARTS 的最大区别是啥?作者在论文里和知乎解读里也给出了解答,这里直接摘抄。   DARTS,不同于 SNAS 中通过完整的概率建模来提出新方法,DARTS 将网络结构直接近似为确定性的连续权重,类似于 attention。在搜索过程中,表达这个 softmax 连续权重的参数 \(\alpha\) 与网络神经变换的参数 \(\theta\) 同时被更新,完全收敛之后选择 \(\alpha\) 的 argmax 搭建子网络,再重新训练 \(\theta\) 。   因为 SNAS 直接优化 NAS 的目标,作者从 SNAS 的建模出发,对 DARTS 的这一近似作出了概率建模下的解释:这种连续化的近似可以被理解为是将中 \(\mathbb{E}_{z \sim p_{\alpha}}\left[L_{\theta}(z)\right]\) 的全局期望:

\[ \mathbb{E}_{\boldsymbol{Z} \sim p_{\boldsymbol{\alpha}}(\boldsymbol{Z})}\left[L_{\boldsymbol{\theta}}\left(\sum_{m>i} \boldsymbol{Z}_{j, m}^{T} \boldsymbol{O}_{j, m}\left(\boldsymbol{Z}_{i, j}^{T} \boldsymbol{O}_{i, j}\left(x_{i}\right)\right)\right]\right. \]

  直接分解到每一条输入边上,计算了一个解析的期望:

\[ L_{\boldsymbol{\theta}}\left(\sum_{m>j} \mathbb{E}_{p_{\boldsymbol{\alpha}_{j, m}}}\left[\boldsymbol{Z}_{j, m}^{T} \boldsymbol{O}_{j, m}\left(\mathbb{E}_{p_{\boldsymbol{\alpha}_{i, j}}}\left[\boldsymbol{Z}_{i, j}^{T} \boldsymbol{O}_{i, j}\left(x_{i}\right)\right]\right)\right]\right) \]

  如果说 \(L\) 对于每一个 \(Z\) 都是线性的,上面两个式子就是等价的。但是因为设计了 ReLU-Conv-BN 的堆叠,带来了非线性,这两个目标并不等价。也就是说,在 DARTS 的连续化近似中带来了很大的偏差(bias)。这一方面带来了最终优化的结果并没有理论保证的问题,使得一阶梯度的结果不尽人意;另一方面因为连续化近似并没有趋向离散的限制,最终通过删除较低权重的边和神经变换产生的子网络将无法保持训练时整个母网络的精度,DARTS 提出用二阶梯度通过基于梯度的元学习来解决第一个问题,但是对于第二个问题,并没有给出一个自动化的解法,而是人工定义了一些规则来挑选边和神经变换,构建子网络,再重新训练。

Conclusion

  非常硬核的一篇文章,解答了我对 NAS 中 MDP 问题的许多问题,但是还是有些地方没有看懂,等有时间再仔细研究研究。

Refer

  • [1] S. Xie, H. Zheng, C. Liu, and L. Lin, “SNAS: Stochastic neural architecture search,” 7th Int. Conf. Learn. Represent. ICLR 2019, pp. 1–17, 2019.
  • https://zhuanlan.zhihu.com/p/53920376

坚持原创技术分享,您的支持将鼓励我继续创作,π(3.14)元就够啦!



文章目录
  1. 1. 引言
  2. 2. Motivation
  3. 3. SNAS
    1. 3.1. Search Space and Architecture Sampling
    2. 3.2. Reparameterization and Differentiable
    3. 3.3. Resource Constraint
    4. 3.4. 与 DARTS 的区别
  4. 4. Conclusion
  5. 5. Refer
您是第 位小伙伴 | 本站总访问量 | 已经写了 611.2k 字啦

载入天数...载入时分秒...