基于GNN的元学习模型直接预测神经网络参数

深度学习神经网络训练往往需要千万级次梯度迭代才能收敛，参数达到稳定，如何通过元模型预测深度学习模型参数，加速或替代梯度迭代收敛的过程。近期NIPS2021的一篇论文为这一问题提供了一种基于GNN的建模思路并取得了不错的效果，作者来自加拿大圭尔夫和Facebook。

论文题目：

Parameter Prediction for Unseen Deep Architectures

方案

模型结构图数据

神经网络的计算图天然为图结构数据，图中的节点为算子如卷积、全连接和求和等，边为算子直接的连接，由于神经网络的计算有正向和反向，因此可以视为有向图。

图1 神经网络的计算图

模型结构空间

One-Shot神经网络结构搜索算法通过训练一张超网，用于评估子网的性能，从中获取符合搜索目标的子网进行从头训练。这种方法对于网络结构的扩展性差，不能预测任意网络的性能，并且往往关注不同网络结构之间的性能排序而不是预测网络的参数。本文基于DARTS扩展了模型结构空间，覆盖了non-separable 2D conv、SE、Transformer类型算子和LN等15种算子，可以涵盖VGG，RESNET和VIT等模型。通过每次均匀采样所有可能的算子即可构建一个计算单元，计算单元的取值范围由定义好的模型深度、宽度和节点数确定。

图2 模型结构

模型结构图超网

基于超网可以构建一张覆盖模型结构采样空间的大图，每次从图中采样的子图则对应一个神经网络模型结构。用于预测神经网络参数的元模型，包含Gated-GNN、参数生成模块Decoder和神经网络loss计算模块。本文每次采样一个batch的子图进行训练，Gated-GNN图神经网络模拟神经网络前向和后向计算进行消息传播更新子图中每个节点的特征，节点特征初始化为算子类型的one -hot编码和参数shape的编码特征，将图节点特征作为Decoder的输入，输出模型参数，Decoder采用了多层感知机输出3D算子参数，不同类型的算子从输出的最大shape中进行切分获取相应shape的参数，生成的参数用于在ImageNet等数据集上进行Loss计算，并将梯度反传回来。本文加强了Gated-GNN网络远距离消息传播并对不同类型的算子进行相应的正则化。