[논문 리뷰] Betty: An Automatic Differentiation Library for Multilevel Optimization

Sang Keun Choe, Willie Neiswanger외 2명|arXiv (Cornell University)|2022. 07. 05.

Metaheuristic Optimization Algorithms Research인용 수 5

한 줄 요약

Betty는 다수준 최적화(MLO)를 위한 혁신적인 자동 미분 라이브러리로, 체인 룰 기반의 최적 반응 자코비안 조합을 통해 효율적이고 확장 가능한 기울기 계산을 가능하게 하는 데이터 플로우 그래프 추상화를 도입한다. 이는 계산 복잡도를 𝒪(d³)에서 𝒪(d²)로 감소시키며, 혼합 정밀도 및 데이터 병렬 학습을 지원하고, 이전 시스템보다 최대 11% 높은 테스트 정확도, 14% 낮은 GPU 메모리 사용량, 20% 빠른 학습 속도를 기록한 종단 간 복잡한 MLO 파이프라인을 구현할 수 있게 한다.

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 응용 프로그램을 가능하게 하기 위해.

제안 방법

• 최적화 문제를 방향성 있는 의존 관계를 가진 계층적 노드로 구조화한 MLO 문제의 새로운 데이터 플로우 그래프 표현을 도입한다.
• 데이터 플로우 그래프의 역방향 순회를 통해 체인 룰을 효율적으로 적용하여 최적 반응 자코비안의 반복적 곱셈을 통해 기울기를 계산한다.
• 중복된 헤시안 계산을 피하고 자코비안 조합을 최적화하여 자동 미분의 계산 복잡도를 𝒪(d³)에서 𝒪(d²)로 감소시킨다.
• 알고리즘 선택(예: AID, ITD)과 시스템 최적화(예: 혼합 정밀도, 데이터 병렬성)를 분리한 모듈러한 소프트웨어 스택을 설계한다.
• 한 단계 기울기 경사 하강 업데이트를 각 문제별로 '스텝' 함수로 추상화하여 조합 가능하고 재사용 가능한 MLO 프로그램 구축을 가능하게 한다.
• JAX 및 PyTorch 생태계와 통합하여 현대적인 딥러닝 시스템, 즉 자동 미분 및 분산 학습을 지원한다.

실험 결과

연구 질문

• RQ1 체인 룰 기반 자코비안 조합을 통해 데이터 플로우 그래프 추상화를 사용해 다수준 최적화에서 체계적이고 효율적으로 기울기를 계산할 수 있는가?
• RQ2 MLO에서 자동 미분을 어떻게 최적화하여 계산 복잡도를 𝒪(d³)에서 𝒪(d²)로 낮출 수 있는가?
• RQ3 이중 최적화를 초월한 임의의 MLO 프로그램 구조를 지원하면서도 모듈성과 성능을 유지할 수 있는 소프트웨어 라이브러리를 설계할 수 있는가?
• RQ4 혼합 정밀도 및 데이터 병렬 학습과 같은 시스템 수준 최적화가 MLO의 확장성 향상과 메모리 사용 감소에 얼마나 기여하는가?
• RQ5 제안된 라이브러리는 기존 도구로는 불가능했던 수억 개의 파라미터를 가진 MLO 모델의 학습을 가능하게 하는가?

주요 결과

• Betty는 데이터 플로우 그래프의 역방향 순회를 통해 최적 반응 자코비안을 효율적으로 조합함으로써 MLO에서 자동 미분의 계산 복잡도를 𝒪(d³)에서 𝒪(d²)로 감소시킨다.
• 기존 구현 대비 벤치마크 작업에서 최대 11% 높은 테스트 정확도를 기록하여 최적화 안정성과 수렴 성능 향상을 입증한다.
• 최적화된 기울기 계산과 시스템 수준 기능 덕분에 여러 벤치마크에서 GPU 메모리 사용량을 최대 14% 감소시키고 학습 월타임을 최대 20% 단축시킨다.
• 라이브러리는 기존에는 불가능했던 수억 개의 파라미터를 가진 모델까지 MLO를 확장하여 이전에는 실현 불가능했던 대규모 학습 시나리오를 가능하게 한다.
• Betty는 혼합 정밀도 및 데이터 병렬 학습과 같은 다양한 최적 반응 자코비안 알고리즘과 시스템 최적화를 모듈러하게 통합할 수 있으며, 저수준 코드 수정 없이도 가능하다.
• Apache-2.0 라이선스 하에 오픈소스로 공개된 Betty에는 모든 실험의 완전한 구현과 예제 MLO 프로그램이 포함되어 있어 재현 가능성과 커뮤니티 확장성을 보장한다.