[논문 리뷰] Don't Unroll Adjoint: Differentiating SSA-Form Programs

Michael Innes|arXiv (Cornell University)|2018. 10. 18.

Parallel Computing and Optimization Techniques참고 문헌 22인용 수 61

한 줄 요약

이 논문은 SSA-form 프로그램을 활용한 역방향 모드 AD 접근법을 도입하고 Julia의 Zygote로 구현되어 복잡한 언어 특징의 차별화 및 컴파일러와의 상호작용을 효율적으로 가능하게 한다.

ABSTRACT

This paper presents reverse-mode algorithmic differentiation (AD) based on source code transformation, in particular of the Static Single Assignment (SSA) form used by modern compilers. The approach can support control flow, nesting, mutation, recursion, data structures, higher-order functions, and other language constructs, and the output is given to an existing compiler to produce highly efficient differentiated code. Our implementation is a new AD tool for the Julia language, called Zygote, which presents high-level dynamic semantics while transparently compiling adjoint code under the hood. We discuss the benefits of this approach to both the usability and performance of AD tools.

연구 동기 및 목표

• 표현력이 풍부하고 추적 기반 방법을 넘어서는 고성능 AD의 필요성을 동기화한다.
• 기존 컴파일러와 통합되어 최적화된 미분 코드를 생성하는 SSA-form 기반 AD를 제안한다.
• Zygote 구현과 Julia의 컴파일러 및 런타임과의 상호작용을 시연한다.
• 이 접근 방식이 제어 흐름, 변이 및 고차 미분과 같은 언어 특징을 어떻게 지원하는지 보여준다.

제안 방법

• J 연산자와 역전파를 사용하여 SSA-form IR을 통해 그래디언트를 전파하는 미분 표현.
• 역방향의 기본 블록과 phi 노드를 사용해 제어 흐름 및 데이터 의존성을 adjoint에서 처리한다.
• 역방향 미분 시 SSA 반복 정의를 관리하기 위해 스택에 alpha 값을 저장하고 해석한다.
• 연관된 adjoint pullback를 통해 cons 셀, 박스, 클로저 등 복잡한 데이터 구조와 변이를 다룬다.
• pullback를 인라이닝하고 효율적인 코드를 생성하기 위해 Julia의 타입 추론 및 컴파일러 최적화를 활용한다.
• 고수준의 미분 가능 프로그래밍과 기존의 컴파일러 최적화 사이의 다리 역할로 위치시킨다.

실험 결과

연구 질문

• RQ1 SSA-form IR에 대한 역방향 AD가 제어 흐름, 변이 및 고차 함수 등 전체 언어 특징을 포괄할 수 있는가?
• RQ2 SSA-form 프로그램의 미분이 이를 호스트 컴파일러와 어떻게 상호작용하여 효율적인 차별화 코드를 생성하는가?
• RQ3 추적 기반 AD 시스템과 비교했을 때 성능 및 사용성의 트레이드오프는 어떠한가?
• RQ4 분기와 루프를 통과하는 올바른 데이터 흐름을 보존하도록 adjoint 코드를 어떻게 생성할 수 있는가?
• RQ5 이 접근 방식이 실제 모델과 컴파일러 백엔드(예: LLVM, TPUs)에 확장 가능한가?

주요 결과

• Zygote는 Julia용으로 작동하는 SSA-form AD 도구를 제공하며, 매우 효율적인 차별화 코드를 산출한다.
• adjoint 구성은 phi 노드와 alpha 추적을 사용하여 역 CFG를 통해 올바른 데이터 흐름을 보존한다.
• 풀백은 고차 미분 및 변이와 데이터 구조를 포함한 복잡한 언어 특징을 지원하도록 구성될 수 있다.
• 실험은 손으로 작성된 미분에 비해 여러 경우에서 경쟁력 있는 성능과 추적 방식에 대한 우수한 벤치마크를 보여준다.
• 이 접근 방식은 표현력을 희생하지 않으면서 기존 컴파일러 최적화 및 향후 가속기 활용을 가능하게 한다.
• SSA 기반 미분은 현대 컴파일러 스택과 통합되어 다양한 언어에 걸친 차 differentiable 프로그래밍을 가능하게 한다.