[Reading] Neural Architecture Search with Reinforcement Learning

1. 概述

  Neural Architecture Search with Reinforcement Learning 一文提交于 2016 年,文章给出了一种自动化搜索神经网络架构的方法,称为神经网络架构搜索(Neural Architecture Search,NAS)。这种方法使用循环神经网络(RNN)来生成神经网络模型的方法,这个 RNN 通过强化学习的方式训练,来最大化生成模型在验证集上准确率的期望。通过这种方法,文章分别生成了用于图像识别的卷积神经网络(CNN),以及用于语言建模的 RNN 单元,都获得了很好的效果。

  神经网络架构搜索的过程如 Figure 1 所示。

Figure 1

Figure 1

神经网络的架构可以用一个可变长度的字符串来描述,因此也就可以使用一个循环网络(控制器)来生成。生成的网络称为子网络(child network),在真实数据上训练后,使用验证集上的准确率作为奖励信号,由此计算策略梯度,来更新控制器。在下个迭代,控制器就会有更高的概率输出具有更高准确率的网络架构。

2. 搜索方法

2.1. RNN 控制器

  文章使用 RNN 作为控制器,来生成网络架构相关的一系列参数,Figure 2 展示了生成 CNN 的情况。控制器逐层生成网络的各参数,当总层数到达一定数量时,生成停止,得到一个网络。然后训练这个生成的网络直到收敛,记录在留出的验证集上的准确率,据此更新 RNN 控制器的参数 $\theta_c$,以最大化生成网络的准确率的期望。此外生成的网络层数也会随着 RNN 控制器的训练而逐渐增加。

Figure 2

Figure 2

2.2. 训练方法

  文章将学习生成神经网络架构看成是一个强化学习的问题,将控制器输出的网络结构参数看成是动作 $a_{1:T}$,生成的网络经训练收敛后,在验证集上的准确率作为回报 $R$,训练控制器来最大化期望回报

\begin{equation}
J(\theta_c) = E_{P(a_{1:T} ; \theta_c)}[R]
\end{equation}

$R$ 不可微,使用策略梯度来更新 $\theta_c$

\begin{equation}
\nabla_{\theta_c} J(\theta_c) = \sum_{t=1}^T E_{P(a_{1:T}; \theta_c)} \big[\nabla_{\theta_c} \log P(a_t|a_{(t-1):1}; \theta_c) R \big]
\end{equation}

上式近似为

\begin{equation}
\frac{1}{m} \sum_{k=1}^m \sum_{t=1}^T \nabla_{\theta_c} \log P(a_t|a_{(t-1):1}; \theta_c) R_k
\end{equation}

其中 $m$ 是控制器在一个 batch 中生成的不同架构的数量,$T$ 是控制器生成一个架构所需预测的超参数数量,$R_k$ 是生成的第 $k$ 个网络的验证集准确率。为了降低方差,文章使用历史架构准确率的指数滑动平均作为 baseline $b$,即

\begin{equation}
\frac{1}{m} \sum_{k=1}^m \sum_{t=1}^T \nabla_{\theta_c} \log P(a_t|a_{(t-1):1}; \theta_c) (R_k – b)
\end{equation}

2.3. 增加架构复杂度

  Figure 2 中所示的结构只能生成逐层堆叠的架构,为了让控制器能够预测类似 ResNet 中的 skip connection,在每层之后添加了一个额外的 anchor point,如 Figure 4 所示。

Figure 4

Figure 4

  第 $N$ 层的 anchor point 其中包含了 $N – 1$ 个 sigmoid,来预测是否将之前的某层直接连接过来,即

\begin{equation}
\mathrm{P}(Layer \; j \; is \; an \; input \; to \; layer \; i) = \mathrm{sigmoid} \big( v^T \tanh(W_{prev} * h_j + W_{curr} * h_i) \big)
\end{equation}

其中 $h_j$ 为第 $j$ 层 anchor point 的隐藏状态,$W_{prev}$、$W_{curr}$、$v$ 为可训练的参数。

  建立了 skip connection 之后,某一层可能会有多个输入,此时将多个输入在通道维度上拼接起来。但拼接两层的输出具有相同的尺寸,如果尺寸不同,文章使用补零的方法,补充较小的尺寸,再进行拼接。

  如果网络中的某一层没有与之前的层有任何连接,则这一层作为输入层。在最后一层,将之前所有没有被拼接的输出拼接起来,作为分类器的输入。

  上面的例子只是预测了卷积层,通过在控制器中添加额外的步骤来预测层类型,可以引入池化、BatchNorm 等不同类型的层。此外还可以加入学习率等其他超参数的预测。

2.4. 生成 recurrent cell 架构

  文章给出了生成循环网络中 cell 结构的方法。对于循环网络中的 cell,搜索算法需要生成一个形如 $h_t = f(x_t, h_{t-1})$ 的函数 $f$,如 $h_t = \tanh( W_1 * x_t + W_2 * h_{t-1}$。文章将此类函数的结构看成一棵树,树的叶子节点的输入为 $x_t$ 和 $h_{t-1}$,树的节点代表一个二元函数,根节点输出为 $h_t$,如 Figure 5 左图所示。在 LSTM 中每个 cell 还有额外的输入 $c_{t-1}$ 和额外的输出 $c_{t}$,可以在控制器最后加入两个额外的步骤来预测 $c_{t-1}$ 和 $c_{t}$ 和树中的哪个节点相连,如 Figure 5 中图所示。

Figure 5

Figure 5

  Figure 5 中图所示的网络生成步骤为:

  • 在 Tree Index 0 阶段,控制器预测了 Add 和 Tanh,表示计算 $a_0 = \tanh(W_1 * x_1 + W_2 * h_{t-1})$。
  • 在 Tree Index 1 阶段,控制器预测了 Elem Mult 和 ReLU,表示计算 $a_0 = \mathrm{ReLU}((W_3 * x_1) \odot (W_4 * h_{t-1}))$。
  • 在 Cell Inject 阶段,控制器预测了 Add 和 ReLU,在 Cell indices 阶段,控制器预测的第 2 个元素为 $0$,表示将 $c_{t-1}$ 连接到 $a_0$,即计算 $a_0^{new} = \mathrm{ReLU}(a_0 + c_{t-1})$。Figure 5 左图
  • 在 Tree Index 2 阶段,控制器预测了 Elem Mult 和 Sigmoid,表示计算 $a_2 = \mathrm{sigmoid}(a_0^{new} \odot a_1)$,注意 Figure 5 左图树节点最大索引为 2,则 $h_t = a_2$。
  • 在 Cell indices 阶段,控制器预测的第 1 个元素为 $1$,表示 $c_t$ 为 Tree Index 1 在激活函数前的输出,即 $c_t = (W_3 * x_1) \odot (W_4 * h_{t-1})$。

由此就构建出了一个完成的 cell 结构。注意 Figure 5 左图的树有两个叶子节点,称为 base 2 架构。文章在实验中使用了 base 8 架构,以获取更强的表达能力。

3. 实验结果

  文章在图像识别也语言建模两个任务上,验证了通过 NAS 所得到的网络的性能。

3.1. 生成 CNN

  文章使用了 CIFAR-10 数据集来生成 CNN 网络并验证其性能。使用了两层 RNN(每层 35 个隐藏单元)作为控制器,搜索空间包括卷积层、ReLU、batch norm、和 skip
connection,分别尝试了固定步长为 $1$ 和预测步长( $[1, 2, 3]$)两种步长方式。对于生成的网络,从训练集中随机抽取 $5000$ 个样本作为验证集,剩下的 $45000$ 个样本作为训练集,训练 50 个 epoch 后记录验证集准确率,据此更新控制器参数。控制器在初始阶段只生成较浅的网络(6 层),随着训练的进行(每 1600 个子网络样本),不断增加网络深度(深度 $\times 2$)。

  实验一共训练了 12800 个网络,从中选择最佳验证准确率的网络后,通过网格搜索学习率等超参数,然后进行测试,结果如 Table 1 所示。可见 NAS 生成的网络可以达到很好的性能,其中 39 层的模型到达了最佳的 3.65% 错误率,速度是 DenseNet (L = 100; k = 24) 的 1.05 倍。

Table 1

Table 1

3.2. 生成 recurrent cell

  文章使用了 Penn Treebank 数据集来生成 recurrent cell 结构并验证其性能。使用了类似于 3.1 中的 RNN 控制器,搜索的操作包括 [add, elem_mult],搜索的激活函数包括 [identity, tanh, sigmoid, relu],base 为 8。生成 cell 结构后,用其构造一个两层的子网络并训练 35 个 epoch,并使用验证集的困惑度(perplexity)进行评价。

  类似于生成 CNN 的流程,实验选择了最低验证集困惑度的结构,通过网格搜索学习率等超参数,实验结果如 Table 2 所示。可见 NAS 生成的 cell 结构也都获得了很好的性能,最佳模型达到了 62.4 的困惑度,较之前(66)降低了 3.6;其中困惑度为 66 的模型速度是之前最佳模型的 2 倍。

Table 2

Table 2

  为了验证 NAS 生成的 cell 可以泛化到其他任务上,文章将其用于同一数据集的字符语言建模任务,结果如 Table 3 所示,可见新的 cell 结构可以泛化到其他任务,且性能优于 LSTM。

Table 3

Table 3

4. 总结

  文章提出的 NAS 使用 RNN 作为控制器,生成可变长度网络描述来构建网络,并验证了通过 NAS 可以得到高性能的网络架构,且其性能可以泛化到其他任务。