神经网络架构搜索 (NAS) 面向高效大语言模型

量化和剪枝等技术改进现有模型，而神经网络架构搜索 (NAS) 采取了更根本的方法。NAS不是修改预定义的架构，而是自动化了设计过程本身，寻求从一开始就高效的新颖大语言模型结构。这与事后优化不同，其目的是找出针对特定目标（如降低延迟、参数量或计算成本 (FLOPs)）进行优化的架构，通常同时保持高准确性。

NAS三要素：搜索空间、搜索策略和评估方法

其核心是，NAS基于三个相互作用的组成部分运行：

1. 搜索空间： 这界定了NAS可以检查的可能架构范围。对于大语言模型而言，搜索空间可能极其庞大且复杂。它可能涉及以下选择：

   • 微观架构：模块的内部结构（例如，注意力头维度、前馈网络扩展因子、激活函数、归一化层类型、卷积层）。
   • 宏观架构：整体网络拓扑（例如，层数、层宽度、模块间的连接模式、混合专家模型中的路由决策）。
   • 运算单元选择：选择不同运算的实现方式（例如，不同的注意力变体）。 设计一个有效的搜索空间很要紧。它必须具有足够的表达能力以包含高效和准确的模型，但又要有足够的约束以使搜索可行。

2. 搜索策略： 这是用于遍历搜索空间的算法。常用策略包括：

   • 强化学习 (RL)： 一个代理（控制器）学习策略以生成有潜力的架构。它根据生成架构的性能获得奖励。
   • 进化算法 (EA)： 维护一个架构种群，迭代应用变异（小的随机改变）和交叉（结合优秀架构的部分），并根据适应度（性能）进行选择。
   • 基于梯度的方法： 像可微分架构搜索 (DARTS) 这样的技术将离散的架构选择放松到连续空间，允许通过梯度下降进行优化。这通常涉及一个包含所有可能路径/操作的“超网”，学习与每个选择相关的权重。这些方法可以快得多，但有时会遇到找到退化架构或偏向无参数操作的问题。 策略的选择影响搜索效率和所获架构的质量。对于大语言模型而言，即使只评估一个候选架构的巨大成本使得样本效率高的策略非常可取，尽管基于梯度的方法需要仔细实施以处理规模和复杂性。

3. 性能评估策略： 通过完整训练评估每个采样架构的真实性能对于大语言模型来说计算成本过高。因此，NAS依赖于评估策略：

   • 低保真度训练： 以更少的轮数、在更小的数据集上或使用更小的模型维度来训练候选架构。
   • 代理任务： 在计算成本更低的、更简单的相关任务上评估性能。
   • 权重共享 / 一次性模型： 训练一个单一的、大型超网，该超网涵盖搜索空间中的所有架构。候选架构通过继承超网中的权重进行评估，避免了单独训练。
   • 性能预测器： 训练一个替代模型（例如，一个小型神经网络），利用先前评估过的架构作为训练数据，根据其规格预测架构的性能。 每种评估策略都在评估成本和准确性之间引入了权衡。性能估计的可靠性直接影响搜索的成功。

使用NAS优化效率

将NAS应用于大语言模型的一个主要动因是直接优化效率指标，同时兼顾任务性能。这通常被视为一个多目标优化问题。搜索策略的目标函数可能如下所示：

$$\text{最小化 } (\text{任务损失} + \lambda_1 \cdot \text{延迟} + \lambda_2 \cdot \text{参数数量} + \dots)$$

这里，$\text{任务损失}$ 可以是困惑度或下游任务的损失，而其他项则表示由超参数 $\lambda_i$ 加权的效率约束。

硬件感知NAS

精密的NAS方法将硬件特性直接纳入搜索过程。硬件感知NAS (HW-NAS) 不仅仅是最小化FLOPs等通用指标，而是针对特定目标硬件（例如，特定的GPU、CPU或移动NPU）的实际延迟或能耗进行优化。这可以通过以下方式实现：

• 构建用于性能评估的硬件性能模型（分析或学习）。
• 直接在搜索空间中包含硬件特定的约束（例如，根据硬件支持限制内存占用或运算单元选择）。 相比于硬件无关的搜索，HW-NAS可以获得明显更适合部署目标的架构。

下面的图表说明了典型的NAS工作流程：

神经网络架构搜索的典型循环。搜索策略从已定义的搜索空间中提出架构，然后对这些架构进行评估（通常使用代理），同时考虑性能和潜在的硬件约束。此反馈指导策略生成更优架构。

NAS应用于大语言模型的挑战

尽管其具有潜力，但将NAS应用于大语言模型面临着重大障碍：

• 计算成本： NAS所需的资源巨大，通常比训练单个模型高出几个数量级。即使采用高效的评估策略，在大语言模型的庞大架构空间中进行搜索也需要大量的计算基础设施。
• 搜索空间设计： 设计一个既足够丰富以产生新颖、高效的设计，又足够受限以实现可行检查的搜索空间并非易事。空间过大则搜索难以处理；过窄则可能错过最优解。
• 性能评估差距： 使用代理方法评估的性能与全面训练后的真实性能之间可能存在较大差距。在低保真度评估下表现良好的架构，在规模扩大后可能并非最优。
• 稳定性和可复现性： 基于梯度的NAS方法尤其对超参数和初始条件敏感，有时会导致结果复现困难或难以实现稳定收敛。

NAS与其他优化方法的结合

NAS与之前讨论的其他技术并非相互排斥。通过NAS获得的架构可以作为后续量化、剪枝或知识蒸馏的有效起点。例如：

• NAS找到的架构可能已经具备使其更适合激进量化的特性（如特定的激活分布或层敏感性）。
• 寻找本身产生稀疏激活模式的架构可以增强剪枝的有效性。
• 使用NAS设计一个较小的学生模型架构，专门通过蒸馏来模仿更大的教师模型，可以比简单地缩小标准架构获得更好的结果。

尽管其功能强大，但在单个NAS循环中同时优化架构以及量化或剪枝策略会大幅增加复杂性，并且仍是一个活跃的研究方向。

总而言之，NAS通过自动化架构发现，代表了一种设计高效大语言模型的方法。尽管计算需求高且面临独特挑战，尤其是在现代大语言模型的规模下，但与优化固定的、手动设计的架构相比，它提供了获得从根本上更高效结构的潜力。随着研究的进展，NAS技术，特别是那些融入硬件感知的技术，在开发真正高效的大语言模型中可能会变得越来越重要。