[논문 리뷰] Forward-Mode Automatic Differentiation in Julia

Jarrett Revels, Miles Lubin외 1명|ENLIGHTEN (Jurnal Bimbingan dan Konseling Islam)|2016. 07. 26.

Distributed and Parallel Computing Systems인용 수 100

한 줄 요약

ForwardDiff.jl은 Julia에서 forward-mode 자동 미분을 스택 할당된 chunked 이중수(dual numbers)를 사용하여 효율적인 그래디언트와 고차 도함수를 제공하며, C++ 구현과 경쟁력 있는 성능 및 JuMP 통합을 제공합니다.

ABSTRACT

We present ForwardDiff, a Julia package for forward-mode automatic differentiation (AD) featuring performance competitive with low-level languages like C++. Unlike recently developed AD tools in other popular high-level languages such as Python and MATLAB, ForwardDiff takes advantage of just-in-time (JIT) compilation to transparently recompile AD-unaware user code, enabling efficient support for higher-order differentiation and differentiation using custom number types (including complex numbers). For gradient and Jacobian calculations, ForwardDiff provides a variant of vector-forward mode that avoids expensive heap allocation and makes better use of memory bandwidth than traditional vector mode. In our numerical experiments, we demonstrate that for nontrivially large dimensions, ForwardDiff's gradient computations can be faster than a reverse-mode implementation from the Python-based autograd package. We also illustrate how ForwardDiff is used effectively within JuMP, a modeling language for optimization. According to our usage statistics, 41 unique repositories on GitHub depend on ForwardDiff, with users from diverse fields such as astronomy, optimization, finite element analysis, and statistics. This document is an extended abstract that has been accepted for presentation at the AD2016 7th International Conference on Algorithmic Differentiation.

연구 동기 및 목표

• Julia에서 경쟁력 있는 성능으로 forward-mode 자동 미분을 동기화하고 가능하게 한다.
• 고차 도함수와 커스텀 수 타입을 지원하는 다차원 이중수(Dual) 표현을 도입한다.
• epsilon 벡터를 스택에 할당하고 청크(chunked) 계산을 가능하게 하여 forward-mode의 메모리 할당을 줄인다.
• JuMP와의 실용적 통합 및 최적화에서의 forward-over-reverse 해essian-벡터 곱을 시연한다.

제안 방법

• 입력 값과 스택 할당된 부분 도함수 벡터를 포함하는 다차원 이중수 타입 Dual{N,T}를 정의한다.
• 입력에 직교적 epsilon 구성요소를 추가하여 방향 도함수를 추적하는 벡터 forward 모드를 구현한다.
• 메모리 대역폭과 계산 비용의 균형을 맞추기 위해 입력을 파티션하고 청크 단위로 그래디언트를 계산하는 Chunk 모드를 도입한다.
• 표준 함수(예: sin, cos, -)를 통해 자동 미분이 가능하도록 Dual에 대한 Julia Base 메서드를 오버로드한다.
• 중첩된 Dual 타입(Dual{M, Dual{N,T}} 등)을 사용하여 한 번의 평가에서 M x N의 2차 도함수를 얻도록 고차 도함수를 지원한다.
• JuMP 내에서 forward-over-reverse Hessian-벡터 곱을 계산하고 닫힌 형식 표현에 포함된 사용자 정의 함수의 미분화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 Julia의 chunk-mode forward-mode AD가 대규모 그래디언트 계산에서 실제로 역방향 모드 AD를 능가할 수 있는가?
• RQ2 중첩된 Dual 타입이 한 번의 평가 패스에서 고차 도함수를 효율적으로 가능하게 하는가?
• RQ3 ForwardDiff를 JuMP와 얼마나 효과적으로 통합하여 최적화 모델 내의 사용자 정의 함수의 미분을 가능하게 하는가?
• RQ4 메모리 대역폭과 함수 평가 비용 관점에서 서로 다른 청크 크기의 성능 트레이드오프는 무엇인가?
• RQ5 Julia의 forward-mode AD가 다중 스레드와 SIMD를 활용하여 성능을 더 향상시킬 수 있는가?

주요 결과

• ForwardDiff는 보고된 벤치마크에서 C++ 구현에 비해 경쟁력 있는 그래디언트 소요 시간을 달성하며, 청크 크기가 성능에 영향을 준다.
• Rosenbrock 및 Ackley 함수에서 입력 크기 k=12000일 때 ForwardDiff의 시간은 특정 청크 크기에서 감소하여 메모리 대 대역폭 트레이드오프를 보여준다.
• Python의 autograd(역방향 모드)와 비교할 때, 표 2의 예시 크기에서 ForwardDiff가 현저한 속도 우위를 보인다(예: 최대 k=10000).
• 다중 스레드 ForwardDiff는 4스레드를 사용할 때 단일 스레드 버전에 비해 약 2배의 속도 향상을 보인다.
• 청크 모드는 더 크거나 작은 그래디언트 계산 청크를 가능하게 하여 캐시 정렬 및 메모리 성능에 영향을 주며 JuMP의 forward-over-reverse Hessian-벡터 계산에 중요한 역할을 한다.
• ForwardDiff는 중첩(Dual{M,Dual{N,T}} 등)을 통한 고차 도함수를 지원하여 한 번의 패스로 M x N의 2차 도함수를 가능하게 한다.