[논문 리뷰] Zygote: A Differentiable Programming System to Bridge Machine Learning and Scientific Computing

Mike Innes, Alan Edelman외 5명|arXiv (Cornell University)|2019. 07. 17.

Parallel Computing and Optimization Techniques참고 문헌 39인용 수 63

한 줄 요약

Zygote는 제너릭 프로그램 구조, 특히 제어 흐름, 재귀, 가변성 등을 포함한 일반적인 프로그램 아키텍처에서 자동 미분을 가능하게 하는 줄리아(Julia)로 구현된 미분 가능 프로그래밍 시스템이다. 사용자 간섭 없이 작동하며 고성능 스테이징 계산을 지원하고 과학 계산 라이브러리와 원활하게 통합되어 기계 학습과 과학 계산을 위한 통합 프레임워크를 제공한다. 표현력 있고 미분 가능한 프로그래밍을 통해 실현된다.

ABSTRACT

Scientific computing is increasingly incorporating the advancements in machine learning and the ability to work with large amounts of data. At the same time, machine learning models are becoming increasingly sophisticated and exhibit many features often seen in scientific computing, stressing the capabilities of machine learning frameworks. Just as the disciplines of scientific computing and machine learning have shared common underlying infrastructure in the form of numerical linear algebra, we now have the opportunity to further share new computational infrastructure, and thus ideas, in the form of Differentiable Programming. We describe Zygote, a Differentiable Programming system that is able to take gradients of general program structures. We implement this system in the Julia programming language. Our system supports almost all language constructs (control flow, recursion, mutation, etc.) and compiles high-performance code without requiring any user intervention or refactoring to stage computations. This enables an expressive programming model for deep learning, but more importantly, it enables us to incorporate a large ecosystem of libraries in our models in a straightforward way. We discuss our approach to automatic differentiation, including its support for advanced techniques such as mixed-mode, complex and checkpointed differentiation, and present several examples of differentiating programs.

연구 동기 및 목표

• 일반적인 프로그램 구조(제어 흐름, 재귀, 가변성 등)를 통해 기계 학습과 과학 계산을 연결하기 위해 미분 가능 프로그래밍을 가능하게 하기.
• 리팩터링이나 계산 스테이징이 필요 없이 최소한의 사용자 오버헤드로 줄리아에서 자동 미분을 지원하기.
• 통합적이고 표현력 있는 프로그래밍 모델을 통해 과학 계산 라이브러리를 기계 학습 모델에 원활하게 통합하기.
• 혼합 모드, 복소수, 체크포인팅 미분과 같은 고급 기법으로 자동 미분를 확장하기.
• 미분 가능 프로그래밍이 실세계의 과학적 및 기계 학습 워크로드를 지원할 수 있도록 성능과 일반성 모두 확보하기

제안 방법

• Zygote는 줄리아의 다형 디스패치와 메타프로그래밍 기능을 활용하여 효율적인 기울기 계산을 생성함으로써 고성능 코드를 컴파일한다.
• 제어 흐름, 재귀, 가변 상태를 포함한 임의의 줄리아 프로그램을 자동으로 미분하기 위해 소스 투 소스 변환 기법을 사용한다.
• 역방향 자동 미분 엔진을 활용하여 프로그램의 추상 구문 트리(abstract syntax tree)를 기반으로 기호적 조작을 통해 혼합 모드 및 복소수 미분을 지원한다.
• 중간 값들을 선택적으로 재계산함으로써 역전파 동안 메모리 사용량을 줄이기 위해 체크포인팅을 구현한다.
• Zygote는 기존 줄리아 생태계의 수치 및 과학 계산 라이브러리와 통합되어 복잡한 과학 계산 워크로드의 엔드 투 엔드 기울기 계산을 가능하게 한다.
• 확장 가능하도록 설계되어 사용자가 최소한의 부가 코드로 사용자 정의 기울기 계산 연산을 정의할 수 있다.

실험 결과

연구 질문

• RQ1 미분 가능 프로그래밍 시스템이 코드 리팩터링 없이도 제어 흐름, 재귀, 가변성과 같은 일반적인 프로그래밍 요소를 지원할 수 있는가?
• RQ2 복잡한 과학 계산 워크로드에 대해 미분 가능 프로그래밍 시스템이 기울기 컴파일 및 실행을 얼마나 효율적으로 수행할 수 있는가?
• RQ3 일반 프로그래밍 언어 환경에서 혼합 모드 및 복소수 미분을 자동 미분으로 확장할 수 있는 정도는 어느 정도인가?
• RQ4 미분 가능 프로그래밍을 통해 과학 계산 라이브러리를 기계 학습 모델에 원활하게 통합할 수 있는가?
• RQ5 고수준 언어인 줄리아에서 전체 프로그램 기울기 계산을 활성화할 경우 발생하는 성능 및 사용성의 상충 관계는 무엇인가?

주요 결과

• Zygote는 복잡한 제어 흐름과 가변성을 포함한 일반 줄리아 프로그램을 사용자 간섭 없이도 성공적으로 기울기 계산한다.
• 대부분의 경우 런타임 오버헤드 없이 효율적인 저수준 코드를 컴파일함으로써 고성능을 달성한다.
• 혼합 모드 및 복소수 기울기 계산은 프로그램의 추상 구문 트리에 대한 기호적 조작을 통해 내장 지원된다.
• 체크포인팅이 효과적으로 통합되어 역전파 동안 메모리 소비를 줄이면서도 계산 효율성을 유지한다.
• ODE 해법기 및 선형 대수 라이브러리와 같은 과학 계산 라이브러리를 직접 기울기 계산 가능한 모델 내에서 사용할 수 있어 표현력과 재사용성이 향상된다.
• Zygote는 기계 학습과 과학 계산을 하나의 성능이 뛰어나고 사용자 우아한 시스템으로 통합할 수 있음을 입증한다.