[论文解读] Betty: An Automatic Differentiation Library for Multilevel Optimization

Sang Keun Choe, Willie Neiswanger+2 more|arXiv (Cornell University)|Jul 5, 2022

Metaheuristic Optimization Algorithms Research被引用 5

一句话总结

Betty 是一个用于多级优化（MLO）的新型自动微分库，通过引入数据流图抽象，实现了通过最佳响应雅可比矩阵的链式规则组合来高效、可扩展地计算梯度。它将计算复杂度从 𝒪(d³) 降低至 𝒪(d²)，支持混合精度和数据并行训练，并实现了端到端的复杂 MLO 流水线，相比先前系统，测试准确率最高提升 11%，GPU 显存使用量降低 14%，训练速度加快 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.

研究动机与目标

• 为解决基于梯度的多级优化（MLO）在实现和扩展过程中面临的主要瓶颈，包括复杂的手动梯度组合以及高昂的内存和计算成本。
• 通过支持高效、模块化和可扩展的 MLO 实现，弥合理论 MLO 研究与实际软件系统之间的差距。
• 支持超越双层优化的任意 MLO 程序结构，包括三层及更深的层级结构。
• 提供系统级特性，如混合精度和数据并行训练，以减少内存和计算开销。
• 使具有数亿参数的大规模 MLO 应用成为可能，这些应用在现有库中此前难以实现。

提出的方法

• 提出一种新型的 MLO 问题数据流图表示方法，其中优化问题被组织为具有有向依赖关系的分层节点。
• 通过反向遍历数据流图来计算梯度，利用迭代相乘最佳响应雅可比矩阵，高效应用链式法则。
• 通过避免冗余的海森矩阵计算并优化雅可比矩阵组合，将自动微分的计算复杂度从 𝒪(d³) 降低至 𝒪(d²)。
• 设计了一个模块化的软件栈，将算法选择（如 AID、ITD）与系统优化（如混合精度、数据并行性）解耦。
• 将单步梯度下降更新抽象为每个问题的“step”函数，支持可组合且可重用的 MLO 程序构建。
• 与 JAX 和 PyTorch 生态系统集成，支持现代深度学习系统，包括自动微分和分布式训练。

实验结果

研究问题

• RQ1 能否利用数据流图抽象，通过链式规则的雅可比矩阵组合，系统且高效地计算多级优化中的梯度？
• RQ2 如何优化 MLO 中的自动微分，以将计算复杂度从 𝒪(d³) 降低至 𝒪(d²)？
• RQ3 能否设计一个软件库，支持超越双层优化的任意 MLO 程序结构，同时保持模块化和高性能？
• RQ4 混合精度和数据并行训练等系统级优化在多大程度上能提升 MLO 的可扩展性并减少内存使用？
• RQ5 所提出的库能否实现对具有数亿参数的 MLO 模型的训练，这些模型此前在现有工具中无法实现？

主要发现

• Betty 通过反向遍历数据流图，高效组合最佳响应雅可比矩阵，将 MLO 中自动微分的计算复杂度从 𝒪(d³) 降低至 𝒪(d²)。
• 与现有实现相比，该库在基准任务上实现了最高 11% 的测试准确率提升，表明优化稳定性与收敛性得到改善。
• 由于优化的梯度计算和系统级特性，该库在多个基准上将 GPU 显存使用量降低最多 14%，训练耗时减少最多 20%。
• 该库成功将 MLO 扩展至具有数亿参数的模型，实现了此前无法实现的大规模训练场景。
• Betty 支持模块化集成多种最佳响应雅可比矩阵算法和系统优化技术（如混合精度和数据并行训练），而无需修改底层代码。
• Betty 以 Apache-2.0 许可证开源，包含所有实验的完整实现和示例 MLO 程序，确保可复现性与社区可扩展性。