Neural Architecture Search: автоматический дизайн нейросетей
opensourceainasneural-networksauto-mlmachine-learningit
Введение: Neural Architecture Search
Что такое NAS?
Neural Architecture Search (NAS) = автоматический поиск архитектуры нейросети
Традиционный подход:
Человек → Дизайн архитектуры → Обучение → Оценка
NAS подход:
Пространство поиска → Автоматический поиск → Лучшая архитектура
Цель NAS:
Найти архитектуру, которая превосходит человеческий дизайн
Факт: Google's NASNet (2017) превзошла Inception v3 на ImageNet с 3% меньшей ошибкой.
Основные подходы к NAS
1. Reinforcement Learning-based NAS
Zoph & Le (2017): NAS с RL
Controller (RNN) → Генерирует архитектуру → Обучается → Accuracy → Reward
Цикл:
1. Controller генерирует описание ячейки (cell)
2. Обучается child network на CIFAR-10
3. Accuracy = reward
4. RL обновляет Controller (Policy Gradient)
Использовал:
- LSTM Controller
- REINFORCE algorithm
- Sharing weights для speedup
class NASController:
"""RL Controller for NAS"""
def __init__(self, hidden_size=400, num_layers=2):
self.controller = torch.nn.LSTM(
input_size=embedding_dim,
hidden_size=hidden_size,
num_layers=num_layers
)
self.output_layer = torch.nn.Linear(hidden_size, num_actions)
self.optimizer = torch.optim.Adam(self.parameters(), lr=2.5e-3)
def sample_architecture(self):
"""Sample architecture from policy"""
start_token = torch.tensor([START_IDX])
outputs = []
log_probs = []
input_emb = self.embedding(start_token)
hiddens = None
for t in range(max_length):
output, hiddens = self.controller(input_emb, hiddens)
logits = self.output_layer(output)
probs = torch.softmax(logits, dim=-1)
action = torch.multinomial(probs, 1)
outputs.append(action.item())
log_probs.append(torch.log(probs))
input_emb = self.embedding(action)
return self.decode_architecture(outputs), torch.stack(log_probs)
def train_step(self, architecture, accuracy, log_probs):
"""Policy Gradient update"""
reward = torch.tensor([accuracy - BASELINE])
baseline = 0.9 * self.baseline + 0.1 * accuracy
self.baseline = baseline
# REINFORCE with baseline
loss = -torch.sum(log_probs * reward.detach())
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item(), baseline
2. Evolutionary-based NAS
Evolutionary NAS (Real et al., 2019): EfficientNAS
Популяция архитектур эволюционирует:
1. Инициализация случайной популяции
2. Оценка fitness (accuracy) каждой
3. Selection: лучшие выбирают родителей
4. Crossover: комбинируем архитектуры
5. Mutation: случайные изменения
6. Повторить
Преимущества:
- Параллелизуем легко
- Не нужны градиенты
- Работает с non-differentiable metrics
class EvolutionaryNAS:
"""Evolutionary algorithm for architecture search"""
def __init__(self, pop_size=50, mutation_rate=0.1):
self.pop_size = pop_size
self.mutation_rate = mutation_rate
self.population = []
self.fitnesses = []
def initialize_population(self, search_space):
"""Create random initial population"""
self.population = [
self.random_architecture(search_space)
for _ in range(self.pop_size)
]
self.fitnesses = [0.0] * self.pop_size
def evaluate(self, architecture, dataset):
"""Train and evaluate architecture"""
accuracy = self.train_and_eval(architecture, dataset)
return accuracy
def evolve(self, generation=100):
"""Run evolutionary loop"""
for gen in range(generation):
for i, arch in enumerate(self.population):
self.fitnesses[i] = self.evaluate(arch, dataset)
sorted_idx = np.argsort(self.fitnesses)[::-1]
self.population = [self.population[i] for i in sorted_idx]
self.fitnesses = [self.fitnesses[i] for i in sorted_idx]
parents = self.population[:self.pop_size // 5]
new_population = parents.copy()
while len(new_population) < self.pop_size:
p1, p2 = random.sample(parents, 2)
child = self.crossover(p1, p2)
if random.random() < self.mutation_rate:
child = self.mutate(child)
new_population.append(child)
self.population = new_population
return self.population[0], self.fitnesses[0]
3. Gradient-based NAS (DARTS)
DARTS (2019): Differentiable Architecture Search
Ключевая идея: релаксация дискретного поиска в непрерывный
Search space:
Каждый edge = weighted sum всех операций
e_j = Σ_{op ∈ O} α_op · op(x)
α_op — learnable mixing coefficients
Optimization:
L(α, w) = loss(train_data, architecture(α, w))
Обновляем:
w ← w - ∇_w L (weights)
α ← α - ∇_α L (arch parameters)
class DARTSCell(torch.nn.Module):
"""Differentiable NAS cell"""
OPERATIONS = {
'none': lambda x: x,
'conv3x3': lambda x: torch.nn.functional.relu(
torch.nn.functional.conv2d(x, kernel_size=3, padding=1)
),
'conv1x1': lambda x: torch.nn.functional.relu(
torch.nn.functional.conv2d(x, kernel_size=1)
),
'maxpool3x3': lambda x: torch.nn.functional.max_pool2d(
x, kernel_size=3, padding=1
),
'avgpool3x3': lambda x: torch.nn.functional.avg_pool2d(
x, kernel_size=3, padding=1
),
'skip': lambda x: x,
'conv5x5': lambda x: torch.nn.functional.relu(
torch.nn.functional.conv2d(x, kernel_size=5, padding=2)
),
}
def __init__(self, C_in, C_out, stride=1):
super().__init__()
self.operations = nn.ModuleList()
self.alphas = nn.ModuleList()
for op_name in self.OPERATIONS:
self.operations.append(self.OPERATIONS[op_name])
self.alphas.append(nn.Parameter(torch.zeros(1)))
self.downsample = nn.Conv2d(C_in, C_out, 1, stride=stride) if C_in != C_out or stride != 1 else None
self.batch_norm = nn.BatchNorm2d(C_out)
def forward(self, x, alpha_weights):
"""Forward pass with architecture weights"""
if self.downsample:
x = self.downsample(x)
weighted_features = []
for i, op in enumerate(self.operations):
alpha = torch.nn.functional.softmax(alpha_weights[i], dim=0)
# Soft selection of operations
feature = op(x)
weighted_features.append(alpha.max() * feature)
out = torch.stack(weighted_features).sum(dim=0)
out = self.batch_norm(out)
return torch.nn.functional.relu(out)
class DARTS(torch.nn.Module):
"""Full DARTS network"""
def __init__(self, C=16, num_cells=8, num_classes=10):
super().__init__()
self.stem = nn.Sequential(
nn.Conv2d(3, C, 3, padding=1, bias=False),
nn.BatchNorm2d(C)
)
self.cells = nn.ModuleList()
for i in range(num_cells):
self.cells.append(DARTSCell(C, C))
self.global_pool = nn.AdaptiveAvgPool2d(1)
self.classifier = nn.Linear(C, num_classes)
# Architecture parameters
self.arch_alpha = nn.ParameterList([
nn.Parameter(torch.zeros(len(DARTSCell.OPERATIONS)))
for _ in range(num_cells)
])
def forward(self, x, alpha_weights=None):
x = self.stem(x)
for i, cell in enumerate(self.cells):
if alpha_weights is not None:
x = cell(x, alpha_weights[i])
else:
x = cell(x, self.arch_alpha[i])
x = self.global_pool(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
def get_architecture(self):
"""Extract discrete architecture from continuous relaxation"""
architecture = []
for i, alpha in enumerate(self.arch_alpha):
ops = torch.nn.functional.softmax(alpha, dim=0)
selected = ops.argmax().item()
architecture.append(list(DARTSCell.OPERATIONS.keys())[selected])
return architecture
4. One-Shot NAS
One-Shot NAS: Weight Sharing
Идея: обучаем одну supernet, извлекаем подсети
Supernet:
Все операции делят веса через momentum
w = Σ α_op · w_op
При инференсе: выбираем op с max α_op
Преимущества:
- Один раз обучаем supernet
- Быстро извлекаем архитектуры
- Меньше GPU дней
class OneShotNAS(torch.nn.Module):
"""One-shot NAS with weight sharing"""
def __init__(self, num_operations=7):
super().__init__()
self.cells = nn.ModuleList([
SharedCell(num_operations) for _ in range(8)
])
self.classifier = nn.Linear(128, 10)
def forward(self, x, architecture):
"""Forward pass with specific architecture"""
for cell, arch in zip(self.cells, architecture):
x = cell(x, arch)
x = x.mean(dim=(2, 3))
return self.classifier(x)
def extract_and_eval(self, dataset, topk=5):
"""Extract top-k architectures and evaluate"""
candidates = []
for _ in 1000:
# Sample architecture from current alphas
arch = []
for cell in self.cells:
ops = [torch.nn.functional.softmax(cell.alpha, dim=0).argmax()
for _ in cell.operations]
arch.append(ops)
# Evaluate on validation set
acc = self.evaluate_on_dataset(dataset, arch)
candidates.append((acc, arch))
candidates.sort(reverse=True)
return candidates[:topk]
Search Space Design
Типы search space
Search Space = множество возможных архитектур
1. Cell-based (наиболее популярен)
Normal Cell:
x₂ = A₁(x₁, x₀) + A₂(x₁, x₀) + A₃(x₂, x₁) + A₄(x₂, x₀)
Reduction Cell:
xₖ = A₁(xₖ₋₁, xₖ₋₂) + ...
Cell → Stacking → Full network
2. Layer-level
- Выбор типа слоя (Conv, LSTM, Attention)
- Количество фильтров
- Kernel size
3. Hyperparameter
- Learning rate
- Batch size
- Optimizer params
Навигация по search space
Навигация = выбор лучшей архитектуры
Methods:
1. Grid Search: полный перебор (медленно!)
2. Random Search: случайная выборка (удивительно эффективно!)
3. Bayesian Optimization: гауссовские процессы
4. Hyperband: early stopping
5. Population-based: эволюция
class NASNavigator:
"""Navigate search space efficiently"""
def __init__(self, search_space, budget=100):
self.search_space = search_space
self.budget = budget
self.results = []
def random_search(self):
"""Random search - surprisingly effective!"""
best_acc = 0
best_arch = None
for _ in range(self.budget):
arch = self.search_space.sample()
acc = self.evaluate(arch)
self.results.append((arch, acc))
if acc > best_acc:
best_acc = acc
best_arch = arch
return best_arch, best_acc
def bayesian_optimization(self):
"""Bayesian optimization with GP surrogate"""
from skopt import gp_minimize
def objective(params):
arch = self.search_space.decode(params)
acc = self.evaluate(arch)
return -acc # minimize negative accuracy
result = gp_minimize(
objective,
self.search_space.bounds,
n_calls=self.budget,
acq_func="EI" # Expected Improvement
)
return self.search_space.decode(result.x), -result.fun
def hyperband(self, max_iter=100, eta=3):
"""Hyperband: aggressive early stopping"""
budget_per_iter = self.budget
configurations = {}
for s in range(int(np.floor(np.log(max_iter, eta)))):
n = (budget_per_iter / (eta ** s)) * (1 / (s + 1))
r = max_iter * (eta ** (-s))
# Generate random configurations
for i in range(int(n)):
arch = self.search_space.sample()
configurations[(s, i)] = (arch, r)
# Progressive halving
for s in reversed(range(int(np.floor(np.log(max_iter, eta))))):
n = int((budget_per_iter / (eta ** s)) * (1 / (s + 1)))
r = max_iter * (eta ** (-s))
# Evaluate current configurations
results = {}
for key in list(configurations.keys()):
if key[0] == s:
arch, _ = configurations[key]
results[key] = self.evaluate(arch, r)
# Keep top half
sorted_results = sorted(results.items(), key=lambda x: x[1], reverse=True)
survivors = dict(sorted_results[:len(sorted_results) // 2])
# Increase budget for survivors
for key, acc in survivors.items():
configurations[key] = (configurations[key][0], r * eta)
best_key = max(configurations,
key=lambda k: self.evaluate(configurations[k][0]))
return configurations[best_key]
NAS для LLM
Поиск оптимальной архитектуры Transformer
NAS для трансформеров:
Search Space:
- Количество attention heads: 4, 8, 12, 16
- Dimension: 256, 512, 768, 1024
- FFN expansion: 2x, 4x, 8x
- Depth: 6, 12, 24, 32
- Attention pattern: full, sparse, linear
Objective:
accuracy ≥ threshold
latency ≤ budget
params ≤ limit
class TransformerNAS:
"""NAS for Transformer architecture"""
SEARCH_SPACE = {
'n_heads': [4, 8, 12, 16],
'd_model': [256, 512, 768, 1024],
'ffn_dim': [1024, 2048, 3072, 4096, 8192],
'n_layers': [6, 12, 24, 32],
'attention_type': ['full', 'sparse', 'linear'],
'activation': ['relu', 'gelu', 'swiglu'],
}
def __init__(self, max_latency_ms=100, max_params=100_000_000):
self.max_latency = max_latency_ms
self.max_params = max_params
def evaluate(self, config):
"""Evaluate architecture with latency constraint"""
# Build model
model = Transformer(**config)
# Check parameter count
n_params = sum(p.numel() for p in model.parameters())
if n_params > self.max_params:
return -float('inf')
# Train briefly
accuracy = self.train_and_eval(model)
# Measure latency
latency = self.measure_latency(model)
if latency > self.max_latency:
return -float('inf')
return accuracy
def search(self, method='random'):
"""Run NAS search"""
if method == 'random':
return self._random_search()
elif method == 'bayesian':
return self._bayesian_search()
def _random_search(self, iterations=100):
best_acc = 0
best_config = None
for _ in range(iterations):
config = {
k: random.choice(v) for k, v in self.SEARCH_SPACE.items()
}
acc = self.evaluate(config)
if acc > best_acc:
best_acc = acc
best_config = config
return best_config, best_acc
AutoGLM: автоматический дизайн LLM
AutoGLM Pipeline:
1. Cell Search:
- Attention head count
- Feedforward dimension
- Activation function
2. Depth Search:
- Сколько слоёв нужно
- Где добавить reduction
3. Sparsity Search:
- Which attention patterns
- Sparse vs dense
4. Compression:
- Pruning unimportant heads
- Quantization-aware search
Практические инструменты
Google Optuna
import optuna
def objective(trial):
# Sample hyperparameters
n_layers = trial.suggest_int('n_layers', 1, 12)
d_model = trial.suggest_categorical('d_model', [128, 256, 512])
n_heads = trial.suggest_int('n_heads', 2, 16)
dropout = trial.suggest_float('dropout', 0.0, 0.5)
model = Transformer(
n_layers=n_layers,
d_model=d_model,
n_heads=n_heads,
dropout=dropout
)
# Train and return validation accuracy
accuracy = train_and_eval(model)
return accuracy
study = optuna.create_study(direction='maximize')
study.optimize(objective, n_trials=100)
print(f"Best config: {study.best_params}")
print(f"Best accuracy: {study.best_value}")
NNi (NNI by Microsoft)
# nnimanager.yaml
searchAlgorithm:
name: pbt # Population Based Training
maxTrialNumber: 100
trialConcurrency: 4
displayInterval: 100
trialCommand: python train.py
trialCodeDir: .
trialGpuNumber: 1
searchSpaceFile: search_space.json
{
"n_layers": {"_type": "randint", "_value": [1, 12]},
"d_model": {"_type": "choice", "_value": [128, 256, 512]},
"learning_rate": {"_type": "loguniform", "_value": [1e-5, 1e-2]}
}
Заключение
NAS — мощный инструмент для автоматического дизайна нейросетей:
- RL-based NAS — точный, но дорогой
- Evolutionary NAS — хорошо параллелизуется
- Gradient-based NAS (DARTS) — быстрый, но может overfit
- One-Shot NAS — компромисс скорость/качество
Для LLM:
- Optuna — лучший инструмент для поиска гиперпараметров
- NNi (Microsoft) — распределённый NAS
- AutoGLM — автоматический дизайн LLM
Ключевой инсайт: Random Search часто превосходит сложные методы при достаточном budget!