Optuna NAS: 40 Trials to Match Hand-Tuned Architecture

The Manual Tuning Trap

Spent three weeks manually tweaking layer counts and dropout rates on a ResNet variant. 47 experiments. Meticulous spreadsheet tracking. Final accuracy: 91.2% on CIFAR-10.

Then I ran Optuna for 40 trials overnight. 91.4%.

That stung. But it also freed me from architecture obsession. Here’s what I learned about making Optuna’s architecture search actually work — because the default settings will waste your GPU hours.

Why Most NAS Tutorials Mislead You

Most Optuna tutorials show you the happy path: define an objective, call study.optimize(), get magic results. What they skip is the part where your first 10 runs OOM, your trials take 3 hours each because you forgot pruning, and your final “optimal” architecture is actually just the first thing that didn’t crash.

The core insight that changed my approach: Optuna isn’t magic Bayesian optimization. It’s infrastructure for structured experimentation with smart sampling. The “smart” part — TPESampler’s Tree-structured Parzen Estimator — only kicks in after enough exploration. The first $n$ trials (default 10) are essentially random.

So if your search space is huge and your budget is 20 trials, you’re mostly doing random search with extra steps.

Setting Up a Search Space That Won’t Explode

Here’s a real architecture search for a CNN classifier. Not a toy MNIST example — something closer to what you’d actually ship:

import optuna
from optuna.pruners import HyperbandPruner
from optuna.samplers import TPESampler
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import gc

# Optuna 3.4+ recommended — the suggest_* API stabilized here
# Earlier versions had subtle bugs with categorical sampling in multi-objective

class DynamicCNN(nn.Module):
    def __init__(self, trial, input_channels=3, num_classes=10):
        super().__init__()

        n_layers = trial.suggest_int("n_conv_layers", 2, 5)

        layers = []
        in_ch = input_channels

        for i in range(n_layers):
            out_ch = trial.suggest_int(f"conv_{i}_channels", 16, 128, step=16)
            kernel = trial.suggest_categorical(f"conv_{i}_kernel", [3, 5])

            layers.append(nn.Conv2d(in_ch, out_ch, kernel, padding=kernel//2))

            # Activation choice per layer — this matters more than people think
            act_name = trial.suggest_categorical(f"conv_{i}_act", ["relu", "gelu", "silu"])
            if act_name == "relu":
                layers.append(nn.ReLU())
            elif act_name == "gelu":
                layers.append(nn.GELU())
            else:
                layers.append(nn.SiLU())

            # BatchNorm or not — surprisingly, skipping it sometimes helps
            use_bn = trial.suggest_categorical(f"conv_{i}_bn", [True, False])
            if use_bn:
                layers.append(nn.BatchNorm2d(out_ch))

            # Pool every other layer to avoid spatial dimension explosion
            if i % 2 == 1:
                layers.append(nn.MaxPool2d(2))

            in_ch = out_ch

        self.features = nn.Sequential(*layers)
        self.adaptive_pool = nn.AdaptiveAvgPool2d(1)

        # MLP head
        hidden = trial.suggest_int("fc_hidden", 64, 512, step=64)
        dropout = trial.suggest_float("dropout", 0.1, 0.5)

        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_ch, hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.adaptive_pool(x)
        return self.classifier(x)

Notice the step=16 for channel counts. This isn’t arbitrary — it prevents the search space from exploding. Without stepping, suggest_int("channels", 16, 128) gives you 113 options per layer. With 5 layers, that’s $113^5 \approx 1.8 \times 10^{10}$combinations just for channels. The <!–MD_CODE_BLOCK_13–> reduces it to 8 options per layer: DOLLAR_AMOUNT_1^5 = 32768$. Still large, but tractable.

The Objective Function: Where Most People Mess Up

The training loop matters more than the architecture definition:

def objective(trial):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    try:
        model = DynamicCNN(trial).to(device)
    except Exception as e:
        # This catches weird architecture combinations that don't compile
        # Happens more often than you'd expect with dynamic shapes
        print(f"Architecture failed to build: {e}")
        raise optuna.TrialPruned()

    # Count params — useful for analysis later
    n_params = sum(p.numel() for p in model.parameters())
    trial.set_user_attr("n_params", n_params)

    # Reject absurdly large models early
    if n_params > 10_000_000:  # 10M param limit
        raise optuna.TrialPruned()

    lr = trial.suggest_float("lr", 1e-4, 1e-2, log=True)
    weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-3, log=True)

    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=30)
    criterion = nn.CrossEntropyLoss()

    # Assume train_loader, val_loader are defined globally
    # In practice you'd pass these or use a closure

    best_val_acc = 0.0
    patience_counter = 0

    for epoch in range(30):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()

            # Gradient clipping — essential for architecture search
            # Some architectures explode without this
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

        scheduler.step()

        # Validation
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for data, target in val_loader:
                data, target = data.to(device), target.to(device)
                output = model(data)
                _, predicted = output.max(1)
                total += target.size(0)
                correct += predicted.eq(target).sum().item()

        val_acc = correct / total

        # CRITICAL: Report intermediate values for pruning
        trial.report(val_acc, epoch)

        if trial.should_prune():
            raise optuna.TrialPruned()

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            patience_counter = 0
        else:
            patience_counter += 1
            if patience_counter >= 5:
                break  # Early stop this trial

    return best_val_acc

The trial.report() and trial.should_prune() pattern is what makes architecture search bearable. Without it, every trial runs the full 30 epochs. With Hyperband pruning, most bad trials die by epoch 5.

How TPESampler Actually Makes Decisions

The Tree-structured Parzen Estimator splits your trial history into “good” and “bad” based on a threshold $\gamma$ (default 0.25, so top 25% are “good”). Then for each parameter:

$$EI(x) = \frac{l(x)}{g(x)}\gamma + (1 – \gamma)$$

Where $l(x)$ is the density estimate from “good” trials and $g(x)$ from “bad” ones. TPE maximizes this Expected Improvement ratio. The practical implication: TPE needs at least 4 trials (with $\gamma=0.25$) before it can meaningfully separate good from bad.

I recommend setting n_startup_trials=20 for architecture search:

sampler = TPESampler(
    n_startup_trials=20,  # More exploration before exploitation
    multivariate=True,    # Models parameter correlations
    seed=42
)

The multivariate=True flag models correlations between parameters. This matters for architecture search because choices interact: a 5-layer network with small channels performs differently than a 2-layer network with large channels. Univariate TPE misses this.

HyperbandPruner: The Settings Nobody Gets Right

Hyperband is based on successive halving with multiple brackets. The math:

$$n_i = \left\lfloor \frac{n}{s^i} \right\rfloor, \quad r_i = r \cdot s^i$$

Where $n$ is the initial number of configurations, $r$ is the minimum resource (epochs), and $s$ is the reduction factor (default 4). In practice:

pruner = HyperbandPruner(
    min_resource=3,       # Minimum epochs before pruning decision
    max_resource=30,      # Total epochs if not pruned
    reduction_factor=3    # More aggressive pruning (default is 4)
)

I switched from reduction_factor=4 to 3 after noticing good architectures were getting killed too early. The default assumes your performance metric is relatively stable by epoch N/4. For architecture search where random init variance is high, you want gentler pruning.

A subtlety that bit me: Hyperband expects your trial.report(value, step) to use integer steps. I was reporting with float epoch fractions (like 2.5 for mid-epoch). Silent failure — pruning never triggered. The docs don’t emphasize this enough.

GPU Memory: The Actual Bottleneck

Architecture search chews through GPU memory because you’re constantly allocating new models. PyTorch doesn’t always clean up immediately:

def objective_with_cleanup(trial):
    result = None
    try:
        result = objective(trial)
    finally:
        # Aggressive cleanup between trials
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
    return result

But here’s what most guides miss: torch.cuda.empty_cache() only releases cached memory back to CUDA’s pool. It doesn’t free memory held by Python objects. If your model or intermediate tensors are still referenced somewhere — a closure, a global, a callback — that memory stays pinned.

On an RTX 3080 (10GB VRAM), I could run ~50 trials before hitting fragmentation issues. The symptom: OOM errors on trials that should fit, with nvidia-smi showing 2GB “free” but PyTorch unable to allocate a 500MB tensor. Solution: restart the study with load_if_exists=True.

The Database Backend Decision

SQLite works for single-process studies. For distributed search across multiple GPUs:

# Single GPU
study = optuna.create_study(
    study_name="cnn_nas",
    storage="sqlite:///optuna_study.db",
    load_if_exists=True,
    sampler=sampler,
    pruner=pruner,
    direction="maximize"
)

# Multi-GPU (separate processes)
# Use PostgreSQL or MySQL — SQLite locks under concurrent writes
storage = "postgresql://user:pass@localhost/optuna"
study = optuna.create_study(
    study_name="cnn_nas_distributed",
    storage=storage,
    load_if_exists=True,
    sampler=sampler,
    pruner=pruner,
    direction="maximize"
)

I’ve seen SQLite corrupt with 4+ workers hitting it simultaneously. The errors are cryptic: database is locked or worse, silently dropped trials. PostgreSQL handles concurrent access properly. If you’re doing serious NAS, set it up.

Optuna vs Ray Tune vs NNI: When Each Wins

I’ve used all three for production architecture search. Quick breakdown:

Optuna wins when you want Python-native flexibility and don’t need distributed search beyond a single multi-GPU machine. The API is the cleanest. If you covered LoRA vs Full Fine-Tuning: Cost-Accuracy Trade-offs, Optuna’s approach to hyperparameter search complements fine-tuning experiments well.

Ray Tune wins for cluster-scale search. It integrates with Ray’s distributed runtime, so scaling from 1 GPU to 100 GPUs across machines is (theoretically) seamless. The API is clunkier — lots of callback ceremony — but if you’re already in the Ray ecosystem, stick with it.

NNI (Microsoft’s Neural Network Intelligence) wins for graph-based NAS where you’re searching over architecture graphs, not just hyperparameters. Its ENAS and DARTS implementations are more mature. But the documentation quality dropped after 2023, and some features feel abandoned.

For most practitioners: start with Optuna. Move to Ray Tune if you hit scaling limits.

A Surprising Result: Shallow Networks Often Win

After 200 trials on CIFAR-10, the top 5 architectures all had 3 conv layers. Not 4. Not 5. Three.

The intuition I had wrong: more layers = more capacity = better performance. The reality for moderately-sized datasets: depth increases training instability during architecture search. The good deep architectures exist, but they’re harder to find because more things can go wrong during training.

This aligns with lottery ticket hypothesis findings (Frankle and Carlin, 2019) — there are winning subnetworks, but initialization luck matters. With limited search budget, shallower networks are more reliably trainable.

Analyzing Search Results

# After study.optimize() completes

# Best trial
print(f"Best accuracy: {study.best_trial.value:.4f}")
print(f"Best params: {study.best_trial.params}")
print(f"Params in best model: {study.best_trial.user_attrs['n_params']:,}")

# Parameter importance — this is gold
from optuna.importance import get_param_importances
importances = get_param_importances(study)
for param, importance in sorted(importances.items(), key=lambda x: -x[1])[:5]:
    print(f"{param}: {importance:.3f}")

Output from my run (your mileage varies):

Best accuracy: 0.9142
Best params: {'n_conv_layers': 3, 'conv_0_channels': 48, 'conv_0_kernel': 3, ...}
Params in best model: 847,234

lr: 0.412
n_conv_layers: 0.287
dropout: 0.156
conv_0_channels: 0.089
weight_decay: 0.041

Learning rate dominates. This is typical — architecture decisions matter less than training dynamics for medium-complexity tasks. The implication: if you’re budget-constrained, fix a reasonable architecture and spend your trials on learning rate and regularization.

When Architecture Search Isn’t Worth It

I’m not entirely sure why, but architecture search rarely pays off for:

1. Transfer learning scenarios — The pretrained backbone dominates. Search over the head, not the feature extractor.
2. Small datasets (<5000 samples) — Overfitting noise in validation accuracy. Your “optimal” architecture just memorized the val set.
3. Well-studied domains — For ImageNet classification, ResNet/EfficientNet variants have been searched to death. You won’t beat them with 100 trials.

Architecture search shines for:

1. Novel modalities — Time series, sensor data, domain-specific inputs where pretrained models don’t exist.
2. Compute-constrained deployment — Finding the smallest architecture that meets your accuracy bar.
3. Multi-objective optimization — Trading off latency vs accuracy, or memory vs performance.

Multi-Objective Search: Accuracy vs Params

This is where Optuna gets genuinely useful beyond basic hyperparameter tuning:

def objective_multiobjective(trial):
    # ... same architecture construction ...

    # Return (accuracy, -n_params)
    # Negative because we maximize both objectives
    # But smaller models are better, so negate
    return best_val_acc, -n_params

study = optuna.create_study(
    directions=["maximize", "maximize"],  # Both are maximize
    sampler=TPESampler(seed=42),
    pruner=HyperbandPruner()
)

study.optimize(objective_multiobjective, n_trials=100)

# Get Pareto front
pareto_trials = study.best_trials
for t in pareto_trials:
    acc, neg_params = t.values
    print(f"Accuracy: {acc:.4f}, Params: {-neg_params:,}")

The Pareto front gives you the efficient boundary: for any point on it, you can’t improve accuracy without adding parameters, and vice versa. This is what deployment decisions should be based on, not single-objective “best” models.

The CUDA Version Trap

Here’s a specific technical surprise: PyTorch 2.0+ with CUDA 12.x handles dynamic model construction differently than PyTorch 1.x with CUDA 11.x. Specifically, the memory allocator changed.

I hit this error on PyTorch 2.1, CUDA 12.1:

RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call

The “asynchronously reported” part is key — the OOM happened on a previous trial, but CUDA reported it lazily. The fix was adding torch.cuda.synchronize() between trials. Without it, you get misleading error attribution.

Practical Budget Guidelines

On an RTX 3080 (10GB VRAM), training a ~1M parameter CNN on CIFAR-10 for 30 epochs takes ~90 seconds per trial. With Hyperband pruning, average trial time drops to ~35 seconds (most get pruned early).

Beyond 200 trials, I’ve rarely seen meaningful improvements. The search space is either well-explored, or you’re overfitting to validation noise. One exception: multi-objective search benefits from higher trial counts because the Pareto front needs density.

Debugging these long runs at 2am? Dark Chocolate Espresso Beans are the only thing that kept me awake during a failed 500-trial run that I should have pruned at 100.

What I’d Do Differently Next Time

Start with parameter importance analysis on a small pilot study (20-30 trials). If learning rate dominates, you’re not doing architecture search — you’re doing hyperparameter tuning with extra steps. Freeze the architecture at something reasonable (ResNet-18, 3-layer CNN, whatever) and just tune training dynamics.

If architecture parameters actually rank high in importance, then expand the search. But most of the time, they don’t.

And always, always use load_if_exists=True. Studies crash. GPUs OOM. Your SSH session dies. The ability to resume from the database has saved me dozens of hours.

FAQ

Q: How many trials does Optuna need before TPESampler outperforms random search?
Generally 30-50 trials. TPE needs enough “good” and “bad” examples to fit meaningful density estimates. Below 20 trials, the performance difference is negligible. For quick experiments under 20 trials, random sampling is fine and arguably more transparent.

Q: Can I use Optuna for transformer architecture search (attention heads, layers, etc.)?
Yes, but be careful with the search space explosion. Transformers have more interacting hyperparameters than CNNs. I’d recommend fixing embedding dimension and searching over depth/heads separately. Also, transformer NAS benefits hugely from multi-objective search (accuracy vs inference latency) since attention scales quadratically with sequence length.

Q: Why does my Optuna study show 0% pruned trials when I’m using HyperbandPruner?
Most likely your trial.report() calls aren’t happening, or they’re using float steps instead of integers. Add print statements to verify the pruner is receiving intermediate values. Also check that your objective function doesn’t catch and suppress the TrialPruned exception — this silently disables pruning.

For most NAS projects: Optuna with TPESampler and HyperbandPruner, 50-100 trials, PostgreSQL backend if distributed. The framework stays out of your way and the results are reproducible. That’s about as much as you can ask from AutoML.

What I’m still figuring out: efficient NAS for very deep networks (50+ layers) where gradient flow becomes the dominant problem. The search space isn’t just architecture — it’s also initialization and normalization strategies. There’s probably a good multi-task learning formulation here, but I haven’t cracked it yet.