Optimizer Design

Most published optimizers win on a narrow slice of architectures and lose elsewhere; a per-workload-tuned baseline is easy to beat on one model and very hard to beat on ten at once. This task asks for genuine optimization research: one optimizer class and one hyperparameter configuration that converge quickly across transformers, convolutional networks, plain MLPs, embedding-heavy recommenders, and contrastive encoders, and that generalize to architecture families it was never tuned on.

The bar is not vanilla AdamW. The frozen per-workload reference step counts were calibrated offline from a portfolio of strong controls (Nesterov SGD, normalized and shape-routed Muon hybrids, adaptive schedules, table-aware Muon/Adam variants, bounded SOAP and Shampoo, and Schedule-Free AdamW), with each denominator accepted only from repeated confirmation runs on the exact task hardware.

The agent must create /app/custom_optimizer.py defining CustomOptimizer(torch.optim.Optimizer) and /app/optimizer_config.json, whose object is passed to the constructor as **kwargs. The frozen training loop drives it through the ordinary construct / zero_grad / step contract.

• Seven visible workloads for development: nano-GPT on WikiText-103, ResNet-18 on CIFAR-100, a masked atom-set transformer for QM9 dipole-moment regression, next-item prediction on MovieLens-1M, a ViT and a 12-layer plain MLP on CIFAR-10, and a SimCSE contrastive encoder on AG News. Three held-out workloads from different architecture families are scored at verification.
• One class and one config for all ten workloads. Branching on workload or model class names is forbidden; adapting behavior based on parameter shape is explicitly allowed; that is legitimate optimizer design.
• The optimizer must be self-contained: imports limited to torch, numpy, scipy, and a fixed set of standard-library modules that deliberately excludes filesystem, network, and OS access.
• The training infrastructure (train_workload.py, run_visible.py, workloads/) is read-only; runs against modified copies are rejected.

A complete experiment loop on an H100: the seven visible workload definitions in /app/workloads/ (exact architectures, datasets, loss functions, and per-workload target losses), the frozen training loop, and python3 /app/run_visible.py to run one workload or all seven, saving per-step loss curves, speedups, and timing to /app/runs/ for comparison across experiments.

The starter custom_optimizer.py is a simple AdamW implementation provided only to demonstrate the interface; it is not the reference portfolio and has no promised speedup. Datasets are baked into the image from hash-pinned sources; the machine is offline for task purposes, and a sandbox-timer CLI reports the remaining budget.

A clean-room verifier runs in a fresh container on the captured /app, first checking every frozen training file against a SHA-256 manifest, then training all ten workloads (seven visible plus three held-out) with the submitted optimizer. The scorer never imports candidate code: submitted optimizer code runs alone in a chrooted, privilege-dropped, seccomp-confined process that sees only mirrored parameters and gradients, while trusted root-owned processes own the models, data, forward/backward passes, step counting, checkpoints, and all scored quantities.

• Per workload, the verifier records the first validation checkpoint at which the EMA validation loss crosses the unchanged target: a crossing earns baseline_steps / candidate_steps against the frozen reference denominator; a miss earns capped loss-ratio partial credit (at most 1.0). Reference optimizers are not re-run; only the frozen scalars are read.
• The ten speedups combine by geometric mean G, mapped asymptotically to reward: 0 at or below the 1.0x reference baseline, 1 − 1/G above it. G = 2 (half the aggregate optimizer updates) maps to 0.5, with no finite full-credit target.
• A crashed or diverged workload is floored near zero before the geometric mean, dragging the aggregate down hard: uniform gains beat one dramatic outlier.
• Disallowed imports, tampered frozen files, and branching on workload identity are rejected outright; candidate stdout is diagnostic only and cannot influence the score.