最近好多论文开始将 神经架构搜索(NAS) 应用于大模型或 大型语言/视觉语言模型的设计中。
比如: LangVision-LoRA-NAS、Jet-Nemotron、PhaseNAS 等看来NAS又要有一波热度了,所以我来回顾一下NAS的基础技术。
深度学习的成功很大程度上依赖于神经网络架构的精心设计。从AlexNet到ResNet,再到Transformer,每一个里程碑式的架构都凝聚了研究者大量的领域知识和反复试验。这种依赖人工设计的模式存在明显局限性:设计过程耗时费力,且很难保证找到全局最优解。
神经架构搜索(Neural Architecture Search, NAS)正是为了解决这一问题而诞生。NAS将神经网络设计转化为一个可以自动求解的优化问题,通过算法自动搜索最优架构,显著提升了架构设计的效率和性能上限。
本文将深入分析NAS的核心技术原理,重点讨论三种主要搜索策略:强化学习方法、进化算法以及基于梯度的方法,并通过具体的代码实现来展示这些方法的实际应用效果。
NAS的技术框架与核心原理
神经架构搜索本质上是一个自动化的网络设计系统,它将传统的人工设计过程转化为机器可以处理的优化问题。与人工逐层设计不同,NAS算法能够在庞大的架构空间中高效搜索,通过系统性评估找到性能最优的网络结构。
NAS的工作流程可以抽象为一个通用框架,其中搜索空间定义了所有可能的网络架构,搜索策略决定了如何在这个空间中进行探索,而性能评估机制则为搜索过程提供反馈信号。
整个NAS过程包含三个关键环节:
搜索空间定义是NAS的基础,它确定了算法可以构建的所有可能神经网络架构的集合。搜索空间的设计需要在表达能力和计算复杂度之间找到平衡点。
搜索策略选择决定了算法如何在搜索空间中进行探索。不同的搜索策略适用于不同的场景,主要包括强化学习、进化算法、基于梯度的方法、贝叶斯优化以及随机搜索等。
性能评估环节负责评估候选架构的质量。由于完整训练每个候选架构的计算开销巨大,实际应用中通常采用加速技术,如在小规模数据集上训练、使用性能预测模型,或者采用权重共享机制让子网络继承超网络的预训练权重。
双层优化问题的数学表述
NAS的目标是找到能够最小化验证损失的最优架构,这个过程可以表述为一个双层优化问题。在这个框架中,需要同时优化网络架构和模型参数两个层面。
数学上,这个优化问题可以形式化为:
其中
α∗
表示最优架构,
A
是整个搜索空间,
L_val
是验证损失,
w∗(α)
是给定架构
α
下的最优权重参数。
这里
L_train
表示训练损失,
w
是模型权重,
α
是架构参数。
外层优化(公式1)处理架构搜索问题,内层优化(公式2)处理标准的模型训练。外层优化的难点在于架构空间通常是离散的且高维,传统的梯度优化方法难以直接应用。NAS的核心贡献就是提供了有效求解这个双层优化问题的算法框架。
NAS的应用领域扩展
NAS技术已经从最初的图像分类任务扩展到多个重要应用领域。在计算机视觉中,NAS不仅能够发现新的网络架构,还能针对特定任务(如小目标检测)优化专用的损失函数。在自然语言处理领域,NAS帮助设计更适合处理特定语言现象(如稀有词汇、长距离依赖)的网络结构。在医学影像分析中,NAS生成的架构对细微病理特征表现出更高的敏感性,为临床诊断提供了重要支持。
NAS的三种核心搜索策略
搜索策略是NAS系统的核心组件,决定了算法在架构空间中的探索方式和效率。不同的搜索策略有各自的适用场景和技术特点。
基于强化学习的搜索方法
强化学习将架构搜索建模为一个序贯决策问题,其中智能体通过不断试错来学习选择高性能架构的策略。在这个框架中,网络架构的构建过程被分解为一系列决策步骤,每个步骤选择特定的架构组件(如层类型、连接方式等)。
强化学习方法的核心是训练一个控制器网络,该网络学习生成有希望的架构候选。控制器接收当前的架构状态作为输入,输出下一个架构组件的选择概率分布。通过与环境的交互(即评估生成架构的性能),控制器逐步学习到更好的架构生成策略。
这种方法特别适合处理复杂的搜索空间,因为强化学习天然具备处理序贯决策和延迟奖励的能力。同时,通过设计合适的奖励函数,可以很容易地将多个优化目标(如准确率、延迟、模型大小)整合到一个统一的框架中。
进化算法的群体搜索机制
进化算法借鉴生物进化的思想,将网络架构类比为生物个体,通过模拟自然选择过程来搜索最优架构。在这个框架中,维护一个由多个架构组成的群体,通过选择、交叉、变异等操作产生新的架构候选。
进化算法的优势在于其全局搜索能力强,不容易陷入局部最优解。算法维护多个搜索方向,通过群体的多样性保证了搜索的泛化性。进化算法天然支持并行计算,可以同时评估群体中的多个个体,显著提升搜索效率。
对于多目标优化问题,进化算法可以找到一组非支配解(帕累托前沿),为用户提供在不同目标间权衡的多种选择。这在实际应用中非常有价值,因为往往需要在准确性和效率之间找到平衡点。
基于梯度的可微分搜索
基于梯度的方法通过构建可微分的搜索空间,将离散的架构搜索问题转化为连续优化问题,从而可以直接应用梯度下降算法。这类方法的代表是可微架构搜索(Differentiable Architecture Search, DARTS)。
DARTS的核心思想是构建一个包含所有可能操作的超网络(supernet),然后通过学习各个操作的权重来隐式地进行架构搜索。在训练过程中,每个边上的操作是所有候选操作的加权组合,权重通过梯度下降进行优化。
DARTS采用交替优化的策略来处理双层优化问题:首先固定架构参数,优化网络权重;然后固定网络权重,优化架构参数。这种交替优化的方式避免了直接求解双层优化问题的困难。
基于梯度的方法的主要优势是计算效率高,相比强化学习和进化算法,DARTS通常能在更短的时间内找到满意的架构。但这种方法也有局限性,主要适用于相对简单的搜索空间和单目标优化问题。
三种搜索策略的实验对比
为了直观展示不同搜索策略的特点和性能,我们通过一个具体的实验来比较三种方法在循环神经网络架构搜索中的表现。虽然任何类型的网络都可以通过NAS进行优化,但复杂架构通常能从自动化搜索中获得更显著的收益。
有效的搜索空间设计是NAS成功的关键。搜索空间需要为算法提供一套明确的规则和约束,确保生成的架构既有创新性又具有实用性。
在架构层面,需要明确允许使用的构建模块,包括层类型(全连接层、卷积层等)、激活函数以及优化器选择。连接规则定义了这些模块如何组合,例如是否允许跳跃连接、分支结构等。参数范围设置确定了网络规模的边界,如层数范围、每层神经元数量、正则化参数等。
搜索空间设计需要在灵活性和可处理性之间找到平衡:既要足够灵活以包含高性能的新颖架构,又要有适当的约束以保证计算的可行性。
我们定义的搜索空间同时涵盖了架构结构参数和训练超参数,这样NAS算法可以协同优化网络设计和训练策略,寻找整体最优的解决方案。
import torch.nn as nn
import torch.optim as optim
search_space = {
# 架构结构
'num_hidden_layers': [1, 2, 3, 4, 5],
'hidden_layer_size': [32, 64, 128, 256, 512],
'activation_function': ['ReLU', 'LeakyReLU', 'Tanh'],
# 超参数
'learning_rate': [0.1, 0.01, 0.001, 0.0001],
'optimizer': ['Adam', 'SGD', 'RMSprop'],
'dropout_rate': [0.0, 0.2, 0.4, 0.6]
}
activation_map = {
'ReLU': nn.ReLU,
'LeakyReLU': nn.LeakyReLU,
'Tanh': nn.Tanh
}
optimizer_map = {
'Adam': optim.Adam,
'SGD': optim.SGD,
'RMSprop': optim.RMSprop
}
架构评估是NAS流程中的关键环节,它为搜索算法提供了判断架构质量的标准。评估函数需要接收完整的架构描述(包括结构参数和超参数),返回量化的性能指标,使搜索算法能够比较不同架构的优劣。
import torch
def evaluate_architecture(architecture, X_train, y_train, X_val, y_val, num_epochs=50):
# 初始化模型、优化器和损失函数
model = build_model(architecture)
criterion = nn.MSELoss()
optimizer_class = optimizer_map[architecture['optimizer']]
optimizer = optimizer_class(model.parameters(), lr=architecture['learning_rate'])
# 训练模型
model.train()
for _ in range(num_epochs):
optimizer.zero_grad()
outputs = model(X_train)
loss = criterion(outputs, y_train)
loss.backward()
optimizer.step()
# 使用验证数据集验证模型
model.eval()
with torch.no_grad():
val_outputs = model(X_val)
val_loss = criterion(val_outputs, y_val)
return val_loss.item()
强化学习方法通过训练一个策略网络来生成架构候选。控制器网络以循环神经网络为基础,能够处理变长的架构序列,并为每个架构组件生成选择概率。
import torch.nn as nn
class ArchitectureController(nn.Module):
def __init__(self, search_space):
super(ArchitectureController, self).__init__()
self.search_space = search_space
self.keys = list(search_space.keys())
self.vocab_size = [len(search_space[key]) for key in self.keys]
self.num_actions = len(self.keys)
self.rnn = nn.RNN(input_size=1, hidden_size=64, num_layers=1)
self.policy_heads = nn.ModuleList([nn.Linear(64, vs) for vs in self.vocab_size])
def forward(self, input, hidden):
output, hidden = self.rnn(input, hidden)
logits = [head(output.squeeze(0)) for head in self.policy_heads]
return logits, hidden
强化学习的搜索过程采用策略梯度方法,通过最大化期望奖励来优化控制器参数。
import torch
def run_rl_search(
search_space, X_train, y_train, X_val, y_val, num_epochs=10, num_episodes=5
):
# 使用ArchitectureController类初始化控制器
controller = ArchitectureController(search_space)
controller_optimizer = optim.Adam(controller.parameters(), lr=0.01)
# 开始搜索
best_loss = float('inf')
best_architecture = None
for episode in range(num_episodes):
# 梯度置零
controller_optimizer.zero_grad()
# rnn期望输入形状为(batch_size, timesteps, features)
hidden = torch.zeros(1, 1, 64)
# 初始化列表/字典来存储对数概率和架构选择
log_probs = []
architecture = {}
# 测试架构选择
for i, key in enumerate(controller.keys):
# 执行控制器
logits, hidden = controller(torch.zeros(1, 1, 1), hidden)
# 为当前架构选择创建分类分布
dist = torch.distributions.Categorical(logits=logits[i])
# 从分布中采样一个动作
action_index = dist.sample()
# 存储选择的架构值和对数概率
architecture[key] = search_space[key][action_index.item()]
log_probs.append(dist.log_prob(action_index))
# 计算验证损失
val_loss = evaluate_architecture(architecture, X_train, y_train, X_val, y_val, num_epochs=num_epochs)
# 更新最优架构选择
reward = -val_loss
policy_loss = torch.sum(torch.stack(log_probs) * -reward)
policy_loss.backward()
controller_optimizer.step()
if val_loss < best_loss:
best_loss = val_loss
best_architecture = architecture
return best_architecture, best_loss
best_arch_rl, best_perf_rl = run_rl_search(
search_space, X_train, y_train, X_val, y_val, num_episodes=5
)
### 进化算法的群体演化搜索
进化算法通过维护一个架构群体并模拟生物进化过程来搜索最优解。该方法在每一代中评估群体中所有个体的适应度,然后通过选择、交叉和变异操作产生下一代群体。
```python
import random
from copy import deepcopy
def run_evolutionary_search(X, y, search_space, population_size=10, num_generations=5):
best_loss = float('inf')
best_architecture = None
# 创建训练和验证数据集
split_idx = int(len(X) * 0.8)
X_train, X_val = X[:split_idx], X[split_idx:]
y_train, y_val = y[:split_idx], y[split_idx:]
# 在群体中开始搜索
population = []
for _ in range(population_size):
# 随机选择要测试的架构
architecture = {key: random.choice(search_space[key]) for key in search_space}
population.append(architecture)
# 迭代所有代(架构选项集合)
for _ in range(num_generations):
fitness = []
for arch in population:
loss = evaluate_architecture(arch, X_train, y_train, X_val, y_val, num_epochs=10)
fitness.append((loss, arch))
if loss < best_loss:
best_loss = loss
best_architecture = arch
# 通过从代中选择'精英'(高性能架构)创建新群体
fitness.sort(key=lambda x: x[0])
new_population = []
num_elites = population_size // 2
elites = [arch for loss, arch in fitness[:num_elites]]
new_population.extend(elites)
# 从新群体创建和变异后代
while len(new_population) < population_size:
parent1 = random.choice(elites)
parent2 = random.choice(elites)
child = deepcopy({})
for key in parent1: child[key] = random.choice([parent1[key], parent2[key]])
mutation_key = random.choice(list(search_space.keys()))
child[mutation_key] = random.choice(search_space[mutation_key])
new_population.append(child)
population = new_population
return best_architecture, best_loss
best_arch_ea, best_perf_ea = run_evolutionary_search(
search_space, population_size=10, num_generations=5
)
基于梯度的方法通过构建可微分的搜索空间来直接优化架构参数。我们实现了一个简化版的DARTS框架,展示其核心思想。
import torch.nn as nn
# 定义单元
class Cell(nn.Module):
def __init__(self, in_features, out_features, ops):
super(Cell, self).__init__()
self.ops = nn.ModuleList([
nn.Sequential(nn.Linear(in_features, out_features), op()) for op in ops
])
def forward(self, x, weights):
return sum(w * op(x) for w, op in zip(weights, self.ops))
# 定义目标模型
class Model(nn.Module):
def __init__(self, search_space):
super(Model, self).__init__()
self.ops_list = [activation_map[name] for name in search_space['activation_function']]
self.num_ops = len(self.ops_list)
self.num_hidden_layers = max(search_space['num_hidden_layers'])
self.hidden_layer_size = search_space['hidden_layer_size'][0]
self.alphas = nn.Parameter(torch.randn(self.num_hidden_layers, self.num_ops, requires_grad=True))
self.layers = nn.ModuleList()
self.layers.append(nn.Linear(1, self.hidden_layer_size))
for _ in range(self.num_hidden_layers - 1):
self.layers.append(Cell(self.hidden_layer_size, self.hidden_layer_size, self.ops_list))
self.output_layer = nn.Linear(self.hidden_layer_size, 1)
def forward(self, x):
architecture_weights = nn.functional.softmax(self.alphas, dim=-1)
output = x
for i, layer in enumerate(self.layers):
if isinstance(layer, nn.Linear):
output = layer(output)
elif isinstance(layer, Cell):
output = layer(output, architecture_weights[i-1])
return self.output_layer(output)
def discretize(self):
architecture = {
'num_hidden_layers': self.num_hidden_layers,
'hidden_layer_size': self.hidden_layer_size,
'learning_rate': 0.001,
'optimizer': 'Adam',
'dropout_rate': 0.0
}
best_op_indices = self.alphas.argmax(dim=-1)
best_ops = [self.ops_list[i].__name__ for i in best_op_indices]
architecture['activation_function'] = best_ops[0]
return architecture
DARTS采用交替优化策略,分别更新网络权重和架构参数:
import torch.nn as nn
import torch.optim as optim
def run_gradient_based_search(search_space, X_train, y_train, X_val, y_val, num_epochs=50):
# 定义模型、损失函数和优化器
model = Model(search_space)
criterion = nn.MSELoss()
arch_params = [model.alphas]
optimizer_alpha = optim.Adam(arch_params, lr=0.001)
arch_param_ids = {id(p) for p in arch_params}
weight_params = [p for p in model.parameters() if p.requires_grad and id(p) not in arch_param_ids]
optimizer_w = optim.Adam(weight_params, lr=0.01)
# 开始搜索
for epoch in range(num_epochs):
# 梯度置零
optimizer_w.zero_grad()
# 前向传播
outputs = model(X_train)
# 优化
loss_w = criterion(outputs, y_train)
loss_w.backward()
optimizer_w.step()
# 反向传播
optimizer_alpha.zero_grad()
val_outputs = model(X_val)
loss_alpha = criterion(val_outputs, y_val)
loss_alpha.backward()
optimizer_alpha.step()
best_architecture = model.discretize()
final_loss = evaluate_architecture(best_architecture, X_train, y_train, X_val, y_val, num_epochs=50)
return best_architecture, final_loss
best_arch_gb, best_perf_gb = run_gradient_based_search(
search_space, X_train, y_train, X_val, y_val
)
实验结果分析与方法对比
结果
进化算法(EA)方法是在此任务中找到最优架构的最有效方法,达到了最低的最佳验证MSE为0.1498。
以下是结果的细分。
1. 强化学习(RL)
RL方法找到的最佳验证MSE为0.2744。损失和奖励值在五个回合中变化显著,最好的结果出现在最后一个回合。
运行了五个回合:
- 回合1:损失 = 1.1483,奖励 = -1.1483
- 回合2:损失 = 3.2017,奖励 = -3.2017
- 回合3:损失 = 4.0062,奖励 = -4.0062
- 回合4:损失 = 2.5762,奖励 = -2.5762
- 回合5:损失 = 0.2744,奖励 = -0.2744
找到的最佳架构:
- 隐藏层数:4
- 隐藏层大小:64
- 激活函数:Tanh
- 学习率:0.1
- 优化器:RMSprop
- Dropout率:0.2
最佳验证MSE:0.2744
2. 进化算法(EA)
EA成功地在五代中最小化了损失,导致最佳总体验证MSE为0.1498。
这个架构在搜索的第二代中被找到。
用十个个体的群体搜索了五代:
- 第1/5代 —— 此代中的最佳损失:0.4558
- 第2/5代 —— 此代中的最佳损失:0.1498
- 第3/5代 —— 此代中的最佳损失:0.3062
- 第4/5代 —— 此代中的最佳损失:0.4200
- 第5/5代 —— 此代中的最佳损失:0.3125
找到的最佳架构:
- 隐藏层数:5
- 隐藏层大小:512
- 激活函数:Tanh
- 学习率:0.1
- 优化器:SGD
- Dropout率:0.2
最佳验证MSE:0.1498
3. 基于梯度的方法
尽管运行了50个epoch,这种方法的表现最差。它导致了最高的最佳验证MSE为3.6725。
用50个epoch进行搜索:
- Epoch 10/50:训练损失 = 0.0938,架构损失 = 2.1598
- Epoch 20/50:训练损失 = 0.0509,架构损失 = 1.6185
- Epoch 30/50:训练损失 = 0.0338,架构损失 = 1.7296
- Epoch 40/50:训练损失 = 0.0184,架构损失 = 0.4939
- Epoch 50/50:训练损失 = 0.0114,架构损失 = 0.2417
找到的最佳架构:
- 隐藏层数:5
- 隐藏层大小:32
- 激活函数:LeakyReLU
- 学习率:0.001
- 优化器:Adam
- Dropout率:0.0
最佳验证MSE:3.6725
强化学习方法的性能表现
强化学习方法经过5个episode的训练,最终找到的最优验证MSE为0.2744。从训练过程可以看出,该方法存在较大的性能波动,这反映了策略学习过程中的探索-利用权衡问题。
训练过程中的损失变化如下:Episode 1的损失为1.1483,Episode 2达到3.2017,Episode 3进一步上升到4.0062,Episode 4回落至2.5762,最终在Episode 5获得最佳结果0.2744。这种波动模式说明强化学习需要充分的探索过程才能收敛到较好的解。
最终发现的最优架构配置为:4个隐藏层,每层64个神经元,使用Tanh激活函数,学习率设置为0.1,采用RMSprop优化器,dropout率为0.2。这个配置在网络复杂度和泛化能力之间取得了较好的平衡。
进化算法的优异表现
进化算法在5代进化过程中展现了稳定且优秀的搜索性能,最终获得了0.1498的最佳验证MSE。值得注意的是,这个最优架构在第2代就被发现,说明进化算法能够快速识别出有前景的架构模式。
各代的最佳损失演化轨迹为:第1代0.4558,第2代0.1498,第3代0.3062,第4代0.4200,第5代0.3125。从这个轨迹可以看出,进化算法通过维护群体多样性,即使在后续代中出现性能波动,但始终保持了全局最优解。
最优架构配置包含5个隐藏层,每层512个神经元,使用Tanh激活函数,学习率0.1,SGD优化器,dropout率0.2。相比强化学习发现的架构,进化算法倾向于选择更大的网络容量,这可能与其群体搜索机制有关。
基于梯度方法的局限性分析
尽管运行了50个epoch,基于梯度的方法表现最差,验证MSE达到3.6725。这个结果暴露了DARTS在某些场景下的局限性,特别是在搜索空间设计不够精细或者超参数设置不当的情况下。
训练过程显示了典型的双层优化收敛模式:训练损失从0.0938(epoch 10)逐步下降到0.0114(epoch 50),而架构损失则从2.1598下降到0.2417。虽然两个损失都在下降,但最终的泛化性能仍然不佳。
发现的架构配置为:5个隐藏层,每层32个神经元,LeakyReLU激活函数,学习率0.001,Adam优化器,无dropout。这个配置相对保守,可能是梯度优化倾向于选择稳定但不够激进的架构导致的。
方法特性的深层分析
三种方法的表现差异反映了它们在搜索机制上的根本不同。进化算法通过群体多样性和并行搜索,能够更好地避免局部最优并发现全局最优解。强化学习虽然有一定的探索能力,但单一智能体的搜索路径限制了其搜索效率。基于梯度的方法虽然计算效率高,但在处理复杂、非凸的架构搜索空间时容易陷入局部最优。
这些结果也说明了搜索空间设计和超参数调优在NAS中的重要性。不同的方法可能需要针对性的搜索空间设计才能发挥最佳性能。
总结
神经架构搜索技术为自动化神经网络设计提供了强有力的工具,显著降低了高性能架构设计的门槛。通过算法自动化搜索,NAS能够发现超越人工设计的新颖架构,特别是在需要最前沿性能的关键应用中展现出重要价值。
当前的NAS技术仍面临两个核心挑战。首先是计算成本问题:早期的NAS方法往往需要数千GPU小时的计算资源,这种巨大的计算开销限制了技术的普及应用。其次是泛化能力不足:大多数NAS方法针对特定任务进行优化,缺乏跨任务的迁移能力,导致每个新问题都需要重新进行昂贵的搜索过程。
从长远来看,NAS技术正在向更加智能化和实用化的方向发展,特别是在大语言模型领域已经开始展现出巨大潜力。近期的研究如LangVision-LoRA-NAS、Jet-Nemotron等工作表明,研究者们已经开始探索用NAS来优化大语言模型的架构设计,包括注意力机制、前馈网络结构、以及模型压缩策略等关键组件。这种趋势预示着NAS将从传统的小规模网络设计扩展到大规模预训练模型的自动化优化,最终实现真正意义上的端到端神经网络自动设计。这不仅会进一步降低深度学习应用的技术门槛,也将为发现适应特定任务和资源约束的新型大模型架构提供强有力的工具。
https://avoid.overfit.cn/post/a4cb8686e30e47b0912a78487ba813f9
作者:Kuriko Iwai