Title: PA&DA: Jointly Sampling PAth and DAta for Consistent NAS
Paper: https://arxiv.org/pdf/2302.14772.pdf
Code: https://github.com/ShunLu91/PA-DA
导读
基于权重共享机制,One-shot NAS方法训练一个supernet,然后通过继承supernet的权重来对sub-network进行评估排序,在很大程度上大大降低了搜索成本。然而,一些研究指出,在训练过程中,共享的权重存在不同的梯度下降方向。论文进一步发现在supernet训练过程中出现较大的梯度方差,降低了supernet排名的一致性。为了缓解这一问题,论文通过联合优化path和data(PA&DA)的抽样分布,来明确地最小化supernet训练的梯度方差。论文从理论上推导了梯度方差与两种抽样分布之间的关系,揭示了最优采样概率与路径和训练数据的归一化梯度范数成正比。
论文的方法在优化路径和数据的采样分布时可以忽略计算成本,可在supernet训练中实现了较低的梯度方差,使得supernet具有更好的泛化性能,从而得到更一致的NAS。论文在各种搜索空间中与其他改进的方法进行了全面的比较。结果表明,该方法具有更可靠的排名性能和更高的搜索架构精度,优于其他方法,表明了该方法的有效性。
贡献
论文使用CIFAR-10在NAS-Bench-201上进行了实验,采用SPOS算法对supernet进行训练,并逐步增加supernet各edge的候选操作,论文记录了训练过程中所有候选操作参数的平均梯度方差,并通过测量相同的64个sub-network的排名结果来评估supernet的子网排名一致性。
如上图所示,supernet的子模型越多,梯度方差越大,排序一致性越差。这些结果表明,在训练过程中,子网较大的梯度方差会损害supernet排名的一致性,通过使用归一化梯度范数作为重要指标,并在supernet训练过程中对路径和数据采用重要抽样策略,可以减少超网梯度方差,提高其排序一致性。
论文的主要贡献如下:
论文验证了supernet训练的权重共享机制导致了较大的梯度方差,损害了supernet的性能,恶化了其排名一致性。
通过推导supernet梯度方差与采样分布之间的关系,论文提出了在supernet训练过程中通过联合优化路径和数据采样分布来明确地最小化梯度方差。论文发现最优采样概率与路径和数据的归一化梯度范数成正比,并在supernet训练过程中对它采用了重要抽样。
论文的方法只需要可忽略不计的计算来执行路径和数据的重要性采样,并且不需要繁琐的超参数调优。在NAS-Bench-201上获得了最高的 Kendall’s Tau 0.713,并在DARTS和ProxylessNAS搜索空间上获得了优越的性能。
方法
Sampling-based One-Shot NAS
Sampling-based One-Shot NAS一般分为Supernet训练和Sub-network搜索两个阶段:
Stage2(Searching stage): 不断地从训练好的Supernet中采样抽取出Sub-network进行评估,得到其在验证数据集上的性能。这里可以使用启发式搜索算法搜索最优子模型
论文试图在训练过程中减少Supernet的梯度方差,以提高Supernet的收敛性和排名的一致性。论文提出在Supernet训练过程中联合优化抽样分布
其中,d (p)和d (q)是关于路径和数据采样分布的梯度方差函数。下面,将介绍如何推导它们之间的关系,并交替地优化这两个抽样分布。
Path Importance Sampling
在第i ii步训练中,以概率 pi从路径抽样分布p(A) 中采样一个子模型α i ,得到的随机梯度为:
论文期望最小化上式中的梯度方差,通过优化抽样分布p:
可以发现**,所以随机梯度公式的约束可以重新表述为:
为了求解上述约束优化问题,使用拉格朗日数乘法,转换为无条件极值问题进行求解:
并进一步推导出最优抽样分布 p ∗
可以得出最优路径抽样概率 与子模型 的归一化梯度范数成正比,即对具有较大梯度范数的子模型进行抽样可以减少Supernet训练的梯度方差。
在实际应用中,论文将子模型 的梯度范数测量为其包含的候选操作的梯度范数的和,并将每个候选操作的归一化梯度范数作为其抽样概率。
论文计算每个常规backward的梯度范数,并在每个epoch后更新候选操作的采样概率。因此,论文对路径采样分布p的优化只需要微不足道的计算,而且特别有效
Data Importance Sampling
根据归一化梯度范数对训练数据进行采样有助于减少深度模型训练的梯度方差,可以形式表示为:
论文最后一层∇ L的预激活输出的损失函数的梯度来逼近每个训练数据的梯度范数的上界,即:
这样,我们就可以通过访问每个训练数据的上界来轻松地衡量它们的重要性,例如在图像分类任务中最后一层通常为softmax,当使用交叉熵损失时,可以预先推导出∇ L 的梯度表达式,并在训练过程中方便地计算出它,如下所示:
Importance Sampling NAS
论文的方法旨在通过减少训练过程中的梯度方差来提高Supernet排名的一致性。提出基于路径重要性采样和数据重要性采样的联合优化:
上述计算只需要额外的代码行,可以以mini-batch方式有效地执行。因此,论文使用这个近似来估计训练数据的重要性,并采用归一化的结果来更新每个epoch后的采样分布q。
实验
Evaluation of Supernet Ranking Consistency
如上图所示,PA&DA只需要比SPOS多0.2个GPU小时,与其他模式相比达到最高的KT和P@Top5%,这表明论文的训练模式是有效的,有利于提高超级网排名的一致性。
Search Performance on CIFAR-10
如上图所示,论文的方法达到了最高的平均测试精度97.52 ± 0.07,超过了原版的DARTS及其高级变体。与其他改进的one-shot NAS方法,如NSAS、Few-Shot-NAS、GM和CLOSE相比,论文的方法始终以最小的搜索成本优于它们。
Search Performance on ImageNet
如上图所示,PA&DA在多一点点FLOPs情况下,超过了DA-NAS、FairNAS-A和SUMNAS-M,与SPOS、ProxylessNAS、MAGIC-AT、Few-Shot NAS和GM相比,PA&DA搜索的体系结构更小,获得了最高的top-1精度77.3,足以证明论文的方法的有效性
消融实验
Effect of batch size
较大的batch size通常可以稳定具有较低梯度方差的深度模型的训练。上图可以观察到随着batch size的增大,GV减小,KT单调增加,batch size 512得到最好的KT 0.670 ± 0.029
Effect of schedules for smoothing parameters
每个epoch后更新DA的采样概率,使用样本级分布和线性增加τ会产生最好的结果。
Effect of DA and PA
这两个模块一起使用会得到最好的结果。此外,PA比DA贡献了更多的性能提高。
总结
论文通过联合优化路径和数据采样分布来减少supernet训练的梯度方差,以提高supernet的排名一致性。论文推导了梯度方差和抽样分布之间的关系,并使用归一化梯度范数来更新这两个分布。大量的实验证明了该方法的有效性。在未来,研究者将进一步探索更有效的方法来减少梯度方差的Supernet训练。