Betty: An Automatic Differentiation Library for Multilevel Optimization

TL;DR

Betty is a new software library that simplifies and accelerates multilevel optimization tasks in machine learning, reducing computational costs and memory usage while supporting large-scale models.

Contribution

Introducing Betty, a novel automatic differentiation library for large-scale multilevel optimization that improves efficiency, scalability, and flexibility compared to existing methods.

Findings

Reduced computational complexity from O(d^3) to O(d^2)

Achieved up to 11% increase in test accuracy

Decreased GPU memory usage by 14% and training time by 20%

Abstract

Gradient-based multilevel optimization (MLO) has gained attention as a framework for studying numerous problems, ranging from hyperparameter optimization and meta-learning to neural architecture search and reinforcement learning. However, gradients in MLO, which are obtained by composing best-response Jacobians via the chain rule, are notoriously difficult to implement and memory/compute intensive. We take an initial step towards closing this gap by introducing Betty, a software library for large-scale MLO. At its core, we devise a novel dataflow graph for MLO, which allows us to (1) develop efficient automatic differentiation for MLO that reduces the computational complexity from O(d^3) to O(d^2), (2) incorporate systems support such as mixed-precision and data-parallel training for scalability, and (3) facilitate implementation of MLO programs of arbitrary complexity while allowing a modular interface for diverse algorithmic and systems design choices. We empirically demonstrate that Betty can be used to implement an array of MLO programs, while also observing up to 11% increase in test accuracy, 14% decrease in GPU memory usage, and 20% decrease in training wall time over existing implementations on multiple benchmarks. We also showcase that Betty enables scaling MLO to models with hundreds of millions of parameters. We open-source the code at https://github.com/leopard-ai/betty.