LPZero: Language Model Zero-cost Proxy Search from Zero（未更新完预览版本）

LPZero代码

摘要

神经架构搜索 (NAS) 有助于自动执行有效的神经网络搜索，同时需要大量的计算资源，尤其是对于语言模型。零样本 NAS 利用零成本 (ZC) 代理来估计模型性能，从而显着降低计算需求。然而，现有的 ZC 代理严重依赖于深入的专家知识和重复的反复试验成本。此外，大多数现有的 ZC 代理都无法超越简单基线（参数数量）的性能。为了应对这些挑战，我们引入了一个称为 LPZero（语言模型零成本代理搜索）的新型框架。它旨在自动设计高效的语言模型 ZC 代理，并实现更高的排名一致性。具体来说，我们首先将现有的 ZC 代理设计整合到一个统一的框架中作为搜索空间，然后应用进化算法启发式地识别语言模型的新代理候选者。为了提高搜索过程的效率，我们引入了预测剪枝策略 (PPS)。该策略旨在预先消除没有希望的代理，从而降低代理降级的风险。在 FlexiBERT 和 GPT-2 搜索空间上进行的大量实验证明了我们算法的有效性。值得注意的是，我们的方法实现的性能排名一致性显著超过了当前代理的一致性。

1 Introduction

Zero-shot NAS 预测神经网络架构的性能，而无需实际训练，使用随机初始化的模型。这种方法能够快速高效地估计架构性能，从而消除训练过程中通常消耗的时间和资源。此外，我们观察到大多数代理无法超越基于参数数量的基准性能，如图 1 所示。这个限制提出了一个基本但至关重要的问题：如何高效、自动地为语言模型设计新的代理？

为了解答这个问题，我们将其分解为两个步骤：

（1）设计一个统一的搜索空间来包含现有的 ZC 代理；

对于第一个步骤，我们重新审视现有的 ZC 代理，如表 1 所示，设计了一个综合性的搜索空间，涵盖了当前的 ZC 代理。具体而言，这些代理根据输入类型分为六类：激活（A）、Jacobs（J）、梯度（G）、头（H）、权重（W）和 Softmax（S），如图 2 所示。在这个统一的框架中，我们从这些类别中选择两种输入，记为 θ。每个输入通过 n 种一元运算 f(·) 进行转换，结果通过二元运算 g(·) 进行组合。这个过程生成了一个候选代理 φ(f, g, θ)，并进入我们的搜索空间。更多细节请参见附录 A。

（2）使用进化算法发现新的代理。

对于第二个步骤，我们提出了一个新颖的 LPZero 框架，代表了从零开始的语言模型代理搜索。正如图 3 所示，我们最初选择 p 个候选代理来建立初始种群，并评估它们在 FlexiBERT 搜索空间中的排名一致性。通过锦标赛选择，我们识别出两个有前景的父代理（φn,m）。接下来，我们进行交叉和变异操作，生成后代代理 φq。为了评估其排名一致性 Spearman ρq，我们使用该代理对每个架构 Ωi 进行评分，并将结果与各自的真实值 gti（例如，平均准确率）进行比较。鉴于搜索空间的稀疏性，我们倡导采用预测修剪策略（PPS），旨在消除无效的代理，从而提高搜索效率。

我们的主要贡献如下：

我们设计了一个全面且高质量的搜索空间，涵盖了大部分现有的适用于语言模型的 ZC 代理。
我们引入了从零开始的语言模型代理搜索（LPZero）框架，并结合预测修剪策略（PPS），防止代理降级，从而提高搜索效率。
在 FlexiBERT 和 GPT-2 上进行的实验验证了我们 LPZero 框架中识别的代理的优越性，表明我们提出的方法具有有效性。

2 Related Work

Zero-shot NAS

作为一种评估候选神经网络架构准确性并且无需进行大量训练的低成本策略，已经获得了越来越多的关注。与传统的One-shot NAS方法相比，这种方法提供了更高效的计算替代方案。Zero-shot NAS的核心是其准确度排名代理，这对其有效性至关重要。尽管大多数现有的代理已经为计算机视觉（CV）任务开发，但在自然语言处理（NLP）任务中，相关的研究探索相对较少。
NWOT（Mellor et al., 2021）利用不同图像的局部Jacobian值，构建了一种基于输入Jacobian相关性的模型排名指标。类似地，ZenNAS（Lin et al., 2021）通过使用输入图像的梯度范数作为排名标准来评估候选架构。此外，Zero-cost NAS（Abdelfattah et al., 2021）受到“最佳大脑损伤”（Optimal Brain Damage）原则（LeCun et al., 1989）的启发，引入了基于剪枝的度量作为零成本代理。这些度量包括GradNorm（Abdelfattah et al., 2021）、Plain（Abdelfattah et al., 2021）、SNIP（Lee et al., 2019）、GraSP（Wang et al., 2020）、Fisher（Turner et al., 2020）和Synflow（Tanaka et al., 2020）。这些代理评估网络参数的重要性，并聚合层级值来估算整体性能。

Zero-cost Proxies for Transformer

最近的研究（Serianni和Kalita, 2023）重新激发了零成本代理在基于Transformer的网络中的应用，这标志着该领域的重要进展。LiteTransformerSearch（Javaheripi et al., 2022）观察到，零成本代理在计算机视觉任务中表现出色，但在应用于自然语言处理时，并未超越基准方法，即解码器的参数数量。Serianni和Kalita（2023）通过FlexiBERT基准重新唤起了零成本代理在RNN和BERT-based Transformer模型中的重要性。研究提出了一系列代理，如突触多样性、突触显著性、激活距离、Jacobian余弦、注意力置信度和头部重要性，强调了它们在无需大量训练的情况下优化架构搜索过程的潜力。

Automatic Search for ZC Proxies

一些研究探索了如何自动搜索零成本代理，特别是EZNAS（Akhauri et al., 2022）和EMQ（Dong et al., 2023）。EZNAS为卷积网络引入了一个专门的搜索空间，并在多个基准测试中取得了令人称赞的表现（Ying et al., 2019；Dong和Yang, 2020）。然而，当将其应用于基于Transformer的网络时，其有效性显著下降。另一方面，EMQ（Dong et al., 2023）开发了一个专门针对混合精度量化代理的搜索空间，但它并未优化用于Transformer-based网络。相比之下，我们的LPZero框架专门为语言模型设计，特别是为Transformer架构优化，展示了更优越且更有前景的表现。

总结：本文提出的LPZero框架在设计零成本代理并自动化搜索过程中展示了卓越的效果，尤其是在Transformer架构上，较现有的方法表现出更强的潜力和优势。

3 Methodology

在本节中，我们设计了一个搜索空间，并详细介绍了LPZero的进化框架及其分析。

3.1 LPZero Search Space Design

大多数 AutoML 方法的搜索空间（如 Real 等人，2020；Liu 等人，2019）通常是为特定目的而专门设计的，并不适用于代理（proxy）搜索。先前的一些自动损失搜索方法（Li 等人，2021b,a；Gu 等人，2022）使用网络输出 y和真实标签 y标量）作为输入，这相对容易处理。然而，这些方法的搜索空间本身较为初级，和我们的最相似。但对于 ZC（零成本）代理搜索问题，我们需要更多可以处理标量、向量和矩阵输入的运算，从而可能导致形状不匹配的问题。

值得注意的是，f20和 f21是两个特殊的一元运算：

f20 表示“直通”，即不对输入进行任何修改，直接返回输入；
f21表示“剪枝”操作，它会移除该分支，相当于返回空。

3.2 Search Algorithm

受 AutoML (He et al., 2021; Li et al., 2019) 的启发，进化算法成为我们搜索算法设计的核心机制。进化算法属于遗传算法的一个分支，它通过对种群中的个体进行生成、评估和选择，来模拟自然选择的过程，以解决优化问题。图 3 展示了我们 LPZero 框架的搜索流程。

在初始化时，我们从搜索空间中 均匀采样 p个零成本（ZC）代理来构建初始种群。然后，我们在搜索空间上测量这些代理的排名相关性，以衡量每个代理的预测能力。接下来，在每次迭代中，我们采用锦标赛选择（tournament selection）从种群中选出占比 R的代理（默认 R=10%）作为潜在的候选者，然后从中随机选出两个个体 ϕn,ϕm 作为“父代理”（parents）。随后，这些父代理将以概率 Cr 和 Mr 分别执行交叉和变异操作，产生后代。

为了验证后代的有效性，我们从搜索空间中采样 S个候选网络结构，并计算它们在“真实得分”和“代理得分”之间的排名相关性。由于搜索空间非常稀疏，包含大量无效或缺乏潜力的 ZC 代理，我们提出了 Early-Stopping 策略 来过滤掉这些候选代理。

Crossover and Mutation（交叉与变异）

每个算法表达式（AE）由 两个分支 和 一个聚合节点 组成。这些分支代表代理架构中的各个组成部分或操作，而聚合节点则将这两个分支的输出结合起来，形成最终的代理得分。正如图 4 所示，我们展示了交叉和变异操作的示意图。

在 交叉操作 中，会从两个父 AE 中交换其“遗传信息”，以生成后代。具体而言，这个过程通过交换父 AE 的某些片段，组合出新的操作及结构。
与之相对，变异操作 会对单个 AE 的“遗传信息”进行随机改动，从而有可能将全新的结构引入种群。

Predictive-Pruning Strategy（预测式剪枝策略）

LPZero 框架中的 预测式剪枝策略（Predictive-Pruning Strategy）在应对庞大且稀疏的搜索空间时发挥了关键作用。它能够迅速识别并舍弃那些无潜力或无效的零成本（ZC）代理，从而节省计算资源并加快寻找最优解的进程。通过使用附录 B 中所述的预定义标准，该策略会评估候选代理的可行性；凡是不满足指定标准的代理都会被从种群中移除，进而缩小搜索空间，让计算资源更多地集中于更有潜力的候选者。总体而言，这种策略性筛选过程提高了 LPZero 框架的效率和有效性，加快了高质量代理架构的发现。

Searched ZC Proxy（搜索到的零成本代理）

LPZero 框架与目标
- 该框架的核心目标是找到一个零成本代理，能够最大化其与真实网络性能排名之间的 Spearman 相关系数。
- 代理 ϕ以符号化表达式的形式存在，由一元运算和二元运算组成。
- 在搜索空间中，允许使用多种基础运算（包括可能什么也不做的“直通”以及剪枝操作），使得搜索具有极高的灵活度和可扩展性。
搜索空间设计
- 输入包括 6 类特征（Activation、Jacobs、Gradients、Head、Weight、Softmax），对应不同网络内部信息。
- 一元运算 21 个，二元运算 4 个，这些运算可以被自由组合成丰富的代理表达式。
- 理论上可组合出 26,460 种潜在代理结构，搜索空间十分庞大。
进化式搜索策略：LPZero 算法
- 通过进化算法迭代搜索：在每一代中，选取优秀代理（前 P），然后随机进行交叉（crossover）和变异（mutation）来生成新代理。
- 新生成的代理如果满足有效性（PPS 验证），则被保留到种群中，否则重新生成。
- 每代都对新产生的代理进行评估并更新种群，最终迭代到给定轮次后，输出 Spearman 值最高的代理。