Source-linked AI summary
Using learned optimizers to make models robust to input noise
Luke Metz, Niru Maheswaranathan, Jonathon Shlens, Jascha Sohl-Dickstein, Ekin D. Cubuk
TL;DR
Vision models can perform exceptionally on images from their training distribution but often fail on perturbed images, and training on some corruptions may not transfer to others. This paper meta-trains learned optimizers using corrupted validation performance while training the inner model on clean data. The learned optimizer generalizes across Gaussian noise scales, while results on novel corruption types vary by setting.
Problem
Vision models can perform exceptionally on training-matched images, but slight perturbations often sharply degrade performance, and training on specific perturbations leaves vulnerability to other noise distributions.
Source: §1 · p. 1Method
Results
The learned optimizer can outperform Adam on Gaussian noise at scales beyond outer-training, while results on novel corruptions are mixed and depend on the inner-training setting.
Source: §1 · p. 1, §4.1 · p. 3, §4.2 · p. 4Takeaways & Limitations
The results present meta-learning as an initial approach to producing robust classifiers, while indicating that optimizer parameterization and inductive biases warrant further study.
Source: §5 · p. 4Meta-learning requires a corruption distribution, and the nine existing benchmark corruptions differ substantially, making generalization to unseen corruptions challenging.
Source: §5 · p. 4
Abstract
from arXiv · showhide
State-of-the art vision models can achieve superhuman performance on image classification tasks when testing and training data come from the same distribution. However, when models are tested on corrupted images (e.g. due to scale changes, translations, or shifts in brightness or contrast), performance degrades significantly. Here, we explore the possibility of meta-training a learned optimizer that can train image classification models such that they are robust to common image corruptions. Specifically, we are interested training models that are more robust to noise distributions not present in the training data. We find that a learned optimizer meta-trained to produce models which are robust to Gaussian noise trains models that are more robust to Gaussian noise at other scales compared to traditional optimizers like Adam. The effect of meta-training is more complicated when targeting a more general set of noise distributions, but led to improved performance on half of held-out corruption tasks. Our results suggest that meta-learning provides a novel approach for studying and improving the robustness of deep learning models.
1. Introduction
Vision models can perform exceptionally on images drawn from their training distribution, yet even slight perturbations can sharply degrade performance. Training with perturbations can improve robustness to those perturbations, but often leaves models vulnerable to other noise distributions; this work investigates learned optimizers as an approach to robustness.
1. Introduction
- Models can achieve superhuman image-classification performance on training-matched distributions, but even slight image perturbations often cause catastrophic degradation.Examples include geometric changes and shifts in brightness or contrast.Source: §1 · p. 1
- Training with image perturbations can improve robustness to the specific perturbations used, while leaving models vulnerable to most other noise distributions.The paper explores meta-learning an optimizer to produce models that perform well on corrupted images, including some additional noise distributions.Source: §1 · p. 1
- The study investigates whether meta-trained optimizers can produce classifiers robust to corruptions, including noise distributions not used to train the inner model.The inner optimizer is trained using clean-data gradients, while the outer objective evaluates performance on corrupted images.Source: §1 · p. 1
2. Methods
The methods train task-specific learned optimizers for a four-layer CIFAR-10 classifier. The optimizer is meta-trained against noisy validation loss, while the inner model is trained on clean data.
2. Methods
- The inner problem trains a four-layer CNN on CIFAR-10 minibatches using a task-specific learned optimizer.The classifier has logits f(x; w), with parameters w.Source: §2 · p. 1
- The update rule can be a learned function of gradients and additional features, parameterized by outer-parameters θ.The SGD example is U_sgd(g; α) = αg, with learning rate α as its single outer-parameter; the learned update functions are more complex.Source: §2 · p. 1
- The outer objective evaluates cross-entropy on noise-injected validation images, while the inner model is trained using clean-data gradients.Noise is applied to validation inputs for outer evaluation, not to inner-model training in this setup.Source: §2 · pp. 1–2
- Outer-parameters are optimized for noised-validation performance using truncated evolutionary strategies, and the optimizer is parameterized by a small fully connected network.The network operates on each inner-parameter independently, with some cross-parameter normalization; its inputs include features such as weight value and timestep.Source: §2 · p. 2
3. Related Work
Prior work documents neural networks’ vulnerability to common image distortions and examines how learned optimizers can target objectives that differ from the inner training loss.
3. Related Work
- CNNs rely more on texture than shape than humans, and style-transfer augmentation can improve robustness on the Common Corruptions Benchmark (Geirhos et al., 2018; Hendrycks & Dietterich, 2019).Other studies report worse ANN performance than human performance on distorted images and limited correlation between their errors (Elsayed et al., 2018).Source: §3 · p. 2
- Research connects robustness to commonplace corruptions with worst-case corruptions, including adversarial examples (Gilmer et al., 2018; Fawzi et al., 2018; Ford et al., 2019).Cubuk et al. (2017) report a universal functional form for input-distortion sensitivity across models, attributed to low correlation between class outputs.Source: §3 · p. 2
- Learned optimizers have progressed from early studies to complex parametric update rules, with prior work also targeting validation-loss objectives distinct from training loss (Bengio et al., 1990; 1992; Metz et al., 2018).Related objective settings include unsupervised learning and reinforcement learning (Metz et al., 2019; Houthooft et al., 2018).Source: §3 · p. 2
4. Experiments
The experiments compare learned optimizers outer-trained for robustness with Adam and a clean-target learned-optimizer baseline across Gaussian noise and benchmark corruptions. The learned optimizer performs well on Gaussian noise at other scales, while results on held-out corruption types vary and depend strongly on whether inner-training uses corrupted data.
4. Experiments
- The study tests Gaussian noise and Hendrycks–Dietterich benchmark corruptions using a four-layer CNN on CIFAR-10, reporting cross-entropy loss on test images.The benchmark experiment divides nine corruptions, excluding JPEG, into seven outer-training and two outer-test corruptions.Source: §4 · p. 2, §4.2 · p. 3
- The comparison includes learned optimizers targeting noisy validation data, a clean-validation learned-optimizer baseline, and learning-rate-tuned Adam; models are evaluated after inner-training on clean or noised data.Adam's learning rate is tuned for the target noise distribution, rather than the outer-training noise distribution.Source: §4 · p. 2, Table 1 · p. 3
4.1. Gaussian Noise
- Inner-training the learned optimizer on 0.05-noised data yields considerable robustness improvements and outperforms all other models, but its advantage over noisy-trained Adam holds at 0.05 noise, not higher levels.Source: §4.1 · p. 3
4.2. Novel corruption types
- For frosted glass blur and shot noise, the learned optimizer generalizes beyond outer-training severities 1–3 and beats clean-trained Adam, but corrupted-data inner-training performs better.In Figure 2, the outer-test corruption is frosted glass blur and the outer-train corruption is shot noise; inner-training uses outer-train corruptions, not the corruption being evaluated.Source: Figure 2 · p. 4
- On held-out corruption tasks, the robustness-targeted learned optimizer beats clean-trained Adam and the clean-target learned-optimizer baseline except on brightness; inner-training on outer-train corruptions substantially improves performance.For fog and brightness, the appendix results favor the clean-target learned-optimizer baseline over both the robustness-targeted optimizer and Adam.Source: §4.2 · pp. 3–4
5. Discussion
The discussion presents meta-learning as a way to outer-train optimizers for robust classifiers, while emphasizing the study's scale and optimizer-design constraints. It also identifies the need for corruption distributions and the difficulty of generalizing across dissimilar corruptions.
5. Discussion
- The study demonstrates outer-training optimizers to produce robust classifiers, but describes its results as a small-scale first step toward real-world settings.The authors contrast an MLP learned optimizer with learning-rate-tuned Adam, whose tested parameterization cannot use learning-rate decay or regularizers such as weight decay.Source: §5 · p. 4
- The experiments assume clean-data inner-training; training on corrupted data instead substantially improves performance for both Adam and learned optimizers.The authors propose future work that inner-trains on one corruption distribution, targets another during meta-training, and outer-tests on a third.Source: §5 · p. 4
- Meta-learning requires a corruption distribution, and the varied nature of the nine tested corruptions makes generalization to unseen corruptions challenging.The authors suggest meta-unsupervised learning as a possible way to build heuristic corruption types for training.Source: §5 · p. 4
A. Optimizer Details
The appendix describes an optimizer based on prior learned-optimizer work and specifies the inner model and its training data. Outer-training uses clean CIFAR-10 data, while evaluation can inner-train on noised data.
A. Optimizer Details
- The learned optimizer is similar to the optimizer used by Metz et al. (2018).The appendix introduces these implementation details before specifying the inner model.Source: A. Optimizer Details · p. 7
A.1. Inner-Model
- The inner model is a four-layer convolutional network with ReLU activations, hidden sizes 32, 32, 64, and 64, and a 10-unit output projection.Layers use kernel size 3 and strides 2, 2, 1, and 1; spatially meaned final-layer activations feed the projection, and cross-entropy loss is used.Source: §A.1 · p. 7
- Outer-training uses CIFAR-10 images rescaled to 0–1 without added noise, though evaluation can inner-train the model on noised data.Source: §A.1 · p. 7
A.2. Learned optimizer architecture
The learned optimizer processes per-unit features with a shared one-hidden-layer MLP. It produces two outputs that encode an update direction and a log step length.
A.2. Learned optimizer architecture
- The shared optimizer MLP receives features for each unit, including gradients, multi-timescale momentum, weights, and log absolute weights.Feature values are normalized by their second moment across each tensor.Source: §A.2 · p. 7
- Time-based features use sinusoids of the current inner-training iteration, with timescales 2, 10, 20, 100, 200, and 1000.The inputs also include the log norm of each tensor value and the log number of units in the tensor.Source: §A.2 · p. 7
- A 32-unit hidden layer produces outputs a and b, combined as ΔW = 0.001aexp(0.001b) to encode an update direction and log step length.The factor 0.001 is used to stabilize the initial step size and avoid initializing in an unstable outer-loss regime.Source: §A.2 · p. 7
A.3. Outer-training details
Outer-training used asynchronous distributed workers with increasing, jittered truncations; evolutionary strategies were chosen for memory efficiency, and seeds were selected by outer-train loss.
A.3. Outer-training details
- Training ran asynchronously across 256 workers with batch size 256, using truncations that grew from 100 to 10k over 5k outer-iterations and were jittered by 20%.Workers sent gradient information to a centralized learner and synchronized weights after each partial truncation; experiments never trained to completion.Source: §A.3 · p. 7
- Outer-gradient estimation used evolutionary strategies for lower memory use; the authors expected a combined estimator to speed outer-training.Antithetical sampling with shared randomness was used where possible to reduce variance.Source: §A.3 · p. 7
- The outer-train loss selected the best model from 4 random seeds for corruption experiments and 3 for Gaussian-noise experiments.Source: §A.3 · p. 7
A.4. Outer-training task distribution: Gaussian experiments
The Gaussian-noise experiments used validation CIFAR-10 images corrupted with Gaussian noise as the outer-objective.
A.4. Outer-training task distribution: Gaussian experiments
- The Gaussian-noise outer-objective used validation CIFAR-10 images with 0.05 Gaussian noise added.Source: §A.4 · p. 7
A.5. Outer-training task distribution: Corruption experiments
The corruption experiments sampled severity and corruption type for the outer-objective, drawing a new augmentation at each inner-training step.
A.5. Outer-training task distribution: Corruption experiments
- Each outer-objective sample used severity 1, 2, or 3 and one of seven corruptions, with a new augmentation sampled for each inner-training.The corruption set comprised Gaussian, shot, and impulse noise; defocus and zoom blur; brightness; and contrast.Source: §A.5 · p. 7
B. Additional Corruption Plots
Figure 3 compares optimizer evaluations and inner-training outcomes across fog and brightness corruptions, including tests at different target-corruption severities.
B. Additional Corruption Plots
- Figure 3 contrasts outer-test fog with outer-train brightness, showing optimizer evaluations during outer-training in column a.Source: Figure 3 · p. 8
- In column b, solid and dashed lines distinguish inner-training on clean data from training on corruptions matching the outer-training distribution.Source: Figure 3 · p. 8
- Column c compares test performance after clean-data or matching-noise inner-training across different severities of the target corruption.Source: Figure 3 · p. 8