​​摘要：本文用权重共享的 one-shot 的 NAS 方式对 BERT 做 NAS 搜索。

本文分享自华为云社区《[NAS论文][Transformer][预训练模型]精读NAS-BERT》，作者：苏道 。

NAS-BERT:Task-Agnostic and Adaptive-Size BERT Compression with Neural ArchitectureSearch

简介：

论文代码没有开源，但是论文写得挺清晰，应该可以手工实现。BERT 参数量太多推理太慢(虽然已经支持用 tensorRT8.X 取得不错的推理效果，BERT-Large 推理仅需 1.2 毫秒)，但是精益求精一直是科研人员的追求，所以本文用权重共享的 one-shot 的 NAS 方式对 BERT 做 NAS 搜索。

涉及到的方法包括 block-wise search, progressive shrinking，and performanceapproximation

讲解：

1、搜索空间定义

搜索空间的 ops 包括深度可分离卷积的卷积核大小[3/5/7]，Hidden size 大小【128/192/256/384/512】MHA 的 head 数[2/3/4/6/8]，FNN[512/768/1021/1536/2048]、和 identity 连接，也就是跳层了，一共 26 个 op，具体可见下图：

注意这里的 MHA 和 FNN 是二选一的关系，但是可以比如说第一层选 MHA 第二层选 FNN，这样构成一个基本的 Transformer 块，可以说这个方法打破的定式的 Transformer 块的搜索又包含了 Transformer 和 BERT 的结构，不同层间也是链式链接，每层只选择一个 op，如下图

2、超网络的训练方式

【 Block-Wise Training + Knowledge Distillation、分块训练+KD 蒸馏】

（1）首先把超网络等分成 N 个 Blocks

（2）以原始的 BERT 作为 Teacher 模型，BERT 也同样分为 N 个 Blocks

（3）超网络（Student）中第 n 个块的输入是 teacher 模型第 n-1 个块的输出，来和 teacher 模型的第 n 个块的输出做均方差来作为 loss，来预测 teacher 模型中这第 n 个 block 的输出

（4）超网络的训练是单架构随机采样训练

（5）由于 student 块的隐藏大小可能与 teacher 块中的 hiddensize 不同，能直接利用教师块隐藏的输入，和输出作为学生块的训练数据。为了解决这个问题，需要在学生块的输入和输出处使用一个可学习的线性变换层来转换每个 hidden size，以匹配教师块的大小，如下图所示

【 Progressive Shrinking】

搜索空间太大，超网络需要有效的训练，可以借助 ProgressiveShrinking 的方式来加速训练和提高搜索效率，以下简称为 PS。但是不能简单粗暴的剔除架构，因为大架构再训练初期难收敛，效果不好，但是并不能代表其表征能力差，所以本文设置了一个 PS 规则：

其含义，a^t 表示超网络中最大的架构，p(▪)表示参数量大小，l(▪)表示 latency 大小，B 表示设置 B 个区间桶，b 表示当前为第几个区间。如果一个架构 a 不满足 p_b>p(a)>pb_1 并且 l_b>l(a)>l_b-1 这个区间，就剔除。

PS 的过程就是从每个 B 桶中抽 E 个架构，过验证集，剔除 R 个最大 loss 的架构，重复这个过程直到只有 m 个架构在每个桶中

3、Model Selection

建一个表，包括 latency 、loss、 参数量 和结构编码，其中 loss 和 latency 是预测评估的方法，评估方法具体可以看论文，对于给定的模型大小和推理延迟约束条件，从满足参数和延迟约束的表中选择最低 loss 的 T 个架构，然后把这个 T 个架构过验证集，选取最好的那个。

实验结果

1、和原始 BERT 相比在 GLUE Datasets 上都有一定的提升：

2、和其他变种 BERT 相比效果也不错：

消融实验

1、PS 是否有效？

如果不用 PS 方法，需要巨大的验证上的时间(5min vs 50hours)，并且超网络训练更难收敛，影响架构排序：

2、是 PS 架构还是 PS 掉 node

结论是 PS 掉 node 太过粗暴，效果不好：

3、二阶段蒸馏是否有必有？

本文蒸馏探究了预训练阶段和 finetune 阶段，也就是 pre-trainKD 和 finetune KD，结论是：

1、预训练蒸馏效果比 finetune 时候蒸馏好

2、两阶段一起蒸馏效果最好

点击关注，第一时间了解华为云新鲜技术~