Overview of the Robustness Library
The robustness library, developed by researchers at MIT’s Madry Lab, is a Python-based toolkit designed to streamline the training, evaluation, and experimentation with adversarially robust neural networks. It serves as a foundational resource for researchers and practitioners working on adversarial machine learning, offering pretrained robust models, dataset utilities, and a flexible framework for testing and red teaming. The library is widely adopted as a baseline in robustness research and is actively used across multiple projects within the Madry Lab, including studies on adversarial training, representation learning, and benchmarking subpopulation shifts.
The library’s development began in 2019 and has since evolved to support a range of functionalities, from command-line interfaces for model training to programmatic APIs for custom experimentation. The MIT License permits broad usage, making it accessible for both academic and industry applications.
How It Works
The robustness library provides a modular architecture that simplifies the process of training, evaluating, and manipulating neural networks. Users can interact with the library through a command-line interface (CLI) or by importing it as a Python package. The CLI allows for the training and evaluation of standard and robust models across various datasets and architectures, with support for custom datasets and model configurations. For example, a typical CLI command to train a ResNet18 model on the CIFAR dataset without adversarial training might look like this:
python -m robustness.main --dataset cifar --data /path/to/cifar \
--adv-train 0 --arch resnet18 --out-dir /logs/checkpoints/dir/
Programmatic usage involves importing specific modules such as modelutils, datasets, and train to construct and train models. Users can define custom loss functions, optimization parameters, and data loading pipelines, making it adaptable to diverse research needs.
The library also facilitates input manipulation, including the generation of adversarial examples, representation inversion, and feature visualization. For instance, generating an adversarial example using a ResNet50 model involves specifying attack parameters such as the norm constraint (e.g., L-infinity or L2), epsilon (ε) value, step size, and number of iterations.
from robustness.datasets import CIFAR
from robustness.modelutils import makeandrestoremodel
ds = CIFAR('/path/to/cifar')
model, _ = makeandrestoremodel(arch='resnet50', dataset=ds,
'iterations': 100, # Number of PGD steps
'customloss': None # Use default cross-entropy loss
The library supports a variety of optimization techniques, including real or estimated gradients, Fourier or pixel-based perturbations, and custom loss functions. This flexibility allows researchers to experiment with different adversarial training strategies and evaluate their effectiveness.
Use Cases
The robustness library has been instrumental in numerous research projects conducted by the Madry Lab. Some notable applications include:
- Adversarial Training and Robust Representations: The library was used in the project Learning Perceptually-Aligned Representations via Adversarial Robustness (arXiv:1906.00945), where researchers explored how adversarial robustness could improve feature representations in neural networks.
- Transferability of Robust Models: The study Do Adversarially Robust ImageNet Models Transfer Better?
- Benchmarking Subpopulation Shifts: The BREEDS benchmark (arXiv:2008.04859) incorporated the robustness library to evaluate model performance under distribution shifts, providing insights into the robustness of neural networks in real-world scenarios.
Beyond these specific projects, the library serves as a general-purpose tool for training and evaluating neural networks, whether or not they involve adversarial robustness. Its modular design allows researchers to extend its functionality for custom datasets, architectures, and training protocols.
Key Features and Capabilities
The robustness library offers several key features that enhance its utility in adversarial machine learning research:
- Pretrained Robust Models: The library provides a collection of pretrained models for datasets such as CIFAR-10 and ImageNet, trained under various norms (L2 and L-infinity) and epsilon (ε) values. These models are evaluated using 20-step and 100-step Projected Gradient Descent (PGD) attacks to assess their robustness. The library also includes models trained under standard conditions (ε = 0.0) for comparison.
- Dataset Support: The library includes built-in support for datasets like CIFAR-10, CIFAR-100, and ImageNet, along with utilities for loading and preprocessing data. It also supports custom datasets, enabling researchers to adapt the library to their specific needs.
- The library supports different optimization methods, such as PGD, and allows customization of attack parameters like the norm constraint, epsilon value, and step size.
- Flexible Training Framework: The library provides a comprehensive training framework that supports adversarial training, custom loss functions, and distributed training across multiple devices.
- Extensibility: The robustness library is designed to be easily extendable. Researchers can add custom datasets, model architectures, and training protocols by leveraging its modular structure. The API documentation and walkthroughs provide guidance on integrating new components, making it accessible for both beginners and advanced users.
The pretrained models released by the Madry Lab are evaluated using standardized protocols. For CIFAR-10 with an L2 norm, models are tested at ε-test values of 0.0, 0.25, 0.5, and 1.0, with robust accuracies reported for both 20-step and 100-step PGD attacks. These evaluations provide a benchmark for comparing the performance of different adversarial training strategies.
Technical Specifications and Requirements
The robustness library is implemented in Python and requires PyTorch with CUDA support for full functionality.
- Dependencies: PyTorch (version 1.1.0 or later), Dill (version 0.2.9), and the cox package for experiment logging.
- Pretrained Model Formats: Pretrained models are saved as PyTorch checkpoint files (.pt) and are compatible with specific versions of PyTorch and Dill.
- API and Documentation: The library includes extensive API documentation and walkthroughs, covering topics such as getting started with the library, customizing training protocols, and creating custom datasets.