[논문 리뷰] Automatic Differentiation of Programs with Discrete Randomness

Gaurav Arya, Moritz Schauer외 2명|arXiv (Cornell University)|2022. 10. 16.

Data Stream Mining Techniques인용 수 12

한 줄 요약

이 논문은 이산적 난수를 포함하는 프로그램을 위한 재매개변수화 기반 자동 미분 방법을 소개한다. 이는 무한소 변화를 확률적 경로를 따라 전파시켜 편향이 없고 분산이 낮은 기울기 추정을 가능하게 한다. 이 방법은 마르코프 체인, 에이전트 기반 모델, 입자 필터와 같은 확률적 시뮬레이션에서 전진 및 역방향 모드 AD를 모두 지원하며, 절단 및 스무딩 기법을 통해 O(1)의 계산 오버헤드를 달성한 오픈소스 자바스크립트 구현을 제공한다.

ABSTRACT

Automatic differentiation (AD), a technique for constructing new programs which compute the derivative of an original program, has become ubiquitous throughout scientific computing and deep learning due to the improved performance afforded by gradient-based optimization. However, AD systems have been restricted to the subset of programs that have a continuous dependence on parameters. Programs that have discrete stochastic behaviors governed by distribution parameters, such as flipping a coin with probability $p$ of being heads, pose a challenge to these systems because the connection between the result (heads vs tails) and the parameters ($p$) is fundamentally discrete. In this paper we develop a new reparameterization-based methodology that allows for generating programs whose expectation is the derivative of the expectation of the original program. We showcase how this method gives an unbiased and low-variance estimator which is as automated as traditional AD mechanisms. We demonstrate unbiased forward-mode AD of discrete-time Markov chains, agent-based models such as Conway's Game of Life, and unbiased reverse-mode AD of a particle filter. Our code package is available at https://github.com/gaurav-arya/StochasticAD.jl.

연구 동기 및 목표

• 베르누이 시행과 같은 이산적 확률적 행동을 갖는 프로그램의 미분 문제를 해결하기 위해, 기존 자동 미분가 비가역적이고 불연속적인 결과로 인해 실패하는 이유를 다루는 것.
• 확률적 기대값을 위한 기울기 계산이 가능하도록 하면서도, 이산적 구조를 유지하는 방법을 개발하는 것.
• 연속 경로 기반 AD와 유사한 편향 없고 분산이 낮은 기울기 추정기를 달성하면서, 연속적 완화나 조정 가능한 파rameter에 의존하지 않는 것.
• 마르코프 체인과 입자 필터를 포함한 복잡한 확률적 시스템에서 전진 및 역방향 모드 AD를 효율적으로 지원하면서 최소한의 계산 오버헤드를 유지하는 것.
• 기존 AD 도구와 원활하게 통합되며, 확률적 모델링에서 기울기 기반 최적화를 가능하게 하는 조합 가능하고 종단 간(end-to-end)으로 미분 가능한 프레임워크를 제공하는 것.

제안 방법

• 무한소 매개변수 변화에 의한 결과 확률의 비율 변화를 추적하는 '확률적 도함수'(stochastic derivative) 개념을 도입한다. 이는 경로 기반 추정기이다.
• 재매개변수화를 통해 이산적 난수 변수를 균일 난수 변수와 매개변수의 결정적 함수로 표현함으로써, 확률적 경로를 따라 기울기 전파가 가능하도록 한다.
• 이산적 사건을 다루기 위해 '부드러운 확률적 도함수'(smoothed stochastic derivative) 기법을 적용하여, 관련된 샘플 경로를 통해 편향 없는 기울기 추정이 가능하도록 한다.
• 온라인 절단 전략을 사용하여 전진 모드 AD에서 비활성화되거나 영향력이 낮은 경로를 동적으로 제거함으로써 O(1)의 계산 오버헤드를 유지한다.
• 스무딩과 경로 기반 추정을 조합하여 역방향 모드 AD 메커니즘을 도출하며, 직선 통과 기울기 추정기(straight-through gradient estimator)를 특수 케이스로 회복한다.
• ForwardDiff.jl, Zygote.jl, ChainRulesCore.jl와 통합된 StochasticAD.jl이라는 자바스크립트 패키지로 이 방법을 구현하여 전진 및 역방향 미분을 원활하게 지원한다.

실험 결과

연구 질문

• RQ1 기존 도함수가 정의되지 않는 이산적 난수 성분(예: 베르누이 또는 포아송 추출)을 포함하는 프로그램으로 자동 미분를 확장할 수 있는가?
• RQ2 이산적 난수를 포함하는 확률적 프로그램의 기대값을 위한 편향 없고 분산이 낮은 기울기 추정기는 어떻게 구성할 수 있는가?
• RQ3 에이전트 기반 모델이나 마르코프 체인과 같은 복잡한 확률적 시뮬레이션에 대해 전진 모드 AD를 고비용 없이 효율적으로 적용할 수 있는가?
• RQ4 연속 경로 기반 방법의 성능을 따라잡는 이산적 난수 프로그램을 위한 역방향 모드 AD 메커니즘을 유도할 수 있는가?
• RQ5 이산적 난수 프로그램의 기울기 추정기에서 분산을 최소화하기 위해 관련 샘플 경로를 체계적으로 생성하는 방법은 무엇인가?

주요 결과

• 제안된 방법은 이산적 난수 프로그램의 기대값에 대해 편향 없는 기울기 추정을 달성하며, 변화량이 0에 수렴할수록 분산이 유한하게 유지된다.
• 이 프레임워크는 절단 기법을 통해 O(1)의 계산 오버헤드를 유지하면서도 이산 시간 마르코프 체인과 콘웨이의 생명 게임과 같은 에이전트 기반 모델의 전진 모드 AD를 지원한다.
• 이 방법은 입자 필터의 종단 간 역방향 모드 AD를 가능하게 하며, 이는 이전 연구에서 알려진 기법을 특수 케이스로 회복한다.
• 부드러운 확률적 도함수 기법은 전진 및 역방향 미분을 모두 가능하게 하여 전통적인 스코어 함수 추정기의 한계를 극복한다.
• 실증 평가 결과, 이 방법은 유한 차분보다 분산이 낮고, 조정 가능한 또는 학습된 파rameter가 필요 없이 연속적 완화 방법과 경쟁 가능하다.
• 오픈소스 StochasticAD.jl 패키지는 다양한 응용 분야에서의 실용성을 입증하며, 확률적 시뮬레이션과 확률적 추론에 적용 가능하다.