[논문 리뷰] Storchastic: A Framework for General Stochastic Automatic Differentiation

Emile van Krieken, Jakub M. Tomczak외 1명|arXiv (Cornell University)|2021. 04. 01.

Domain Adaptation and Few-Shot Learning참고 문헌 45인용 수 8

한 줄 요약

Storchastic는 임의의 순서의 도함수 추정에 대해 편향이 없는 저분산 기울기 추정 방법을 사용할 수 있도록 해주는 일반적인 스토하스틱 자동 미분 프레임워크이다. 이 프레임워크는 추정기의 네 가지 구성 요소—제안 분포, 가중 함수, 기울기 함수, 제어 변수—로 분해함으로써 임의의 순서의 도함수 추정에 대해 증명 가능한 편향 없는 기울기 추정을 제공하며, 측도 기반 파생도 및 SPSA와 같은 고급 기법을 지원한다.

ABSTRACT

Modelers use automatic differentiation (AD) of computation graphs to implement complex Deep Learning models without defining gradient computations. Stochastic AD extends AD to stochastic computation graphs with sampling steps, which arise when modelers handle the intractable expectations common in Reinforcement Learning and Variational Inference. However, current methods for stochastic AD are limited: They are either only applicable to continuous random variables and differentiable functions, or can only use simple but high variance score-function estimators. To overcome these limitations, we introduce Storchastic, a new framework for AD of stochastic computation graphs. Storchastic allows the modeler to choose from a wide variety of gradient estimation methods at each sampling step, to optimally reduce the variance of the gradient estimates. Furthermore, Storchastic is provably unbiased for estimation of any-order gradients, and generalizes variance reduction techniques to higher-order gradient estimates. Finally, we implement Storchastic as a PyTorch library at https://github.com/HEmile/storchastic.

연구 동기 및 목표

• 연속 변수에 국한되거나 고분산 점수 함수 추정기로 인해 문제가 발생하는 기존 스토하스틱 AD 프레임워크의 한계를 해결한다.
• 모델러가 각 샘플링 단계에서 최적의 기울기 추정 방법을 선택하여 분산을 최소화하면서 편향을 유발하지 않도록 한다.
• 제어 변수와 같은 분산 감소 기법을 스토하스틱 계산 그래프 내에서 임의의 순서의 도함수 추정으로 일반화한다.
• 추정기 구성 요소에 대한 잘 정의된 조건 하에서 n차 도함수 추정의 편향 없음을 보장하는 형식적 수학적 기반을 제공한다.
• 실제 적용과 다양한 기울기 추정기 실험을 지원하기 위한 실용적이고 오픈소스의 PyTorch 라이브러리 구현을 제공한다.

제안 방법

• 스토하스틱 기울기 추정기를 네 가지 핵심 구성 요소—제안 분포, 가중 함수, 기울기 함수, 제어 변수—로 분해한다.
• 기울기 전파의 엄밀한 분석을 가능하게 하기 위해 수학적 프레임워크를 사용해 정방향 모드 AD 평가를 형식화한다.
• 네 가지 구성 요소가 특정 조건을 만족할 경우, 결과로 나오는 가짜 손실이 편향 없는 n차 도함수 추정을 제공함을 증명한다.
• 제안된 구성 요소 분해를 통해 측도 기반 파생도(MVD) 및 SPSA와 같은 고급 기울기 추정 기법을 프레임워크에 통합한다.
• 제어 변수 기법을 n차 도함수 추정으로 확장함으로써 고차 도함수의 분산 감소를 가능하게 한다.
• 딥 러닝 워크플로우와 동적 계산 그래프에 쉽게 통합할 수 있도록 PyTorch 라이브러리로 프레임워크를 구현한다.

실험 결과

연구 질문

• RQ1 편향 없는 임의의 순서의 도함수 추정을 보장하면서도 다양한 기울기 추정 방법을 지원할 수 있는 일반적인 스토하스틱 AD 프레임워크를 설계할 수 있는가?
• RQ2 제어 변수와 같은 분산 감소 기법을 스토하스틱 계산 그래프 내에서 고차 도함수 추정으로 체계적으로 확장할 수 있는가?
• RQ3 기울기 추정기의 네 구성 요소에 대해 어떤 형식적 조건이 n차 도함수 추정의 편향 없음을 보장하는가?
• RQ4 측도 기반 파생도 및 SPSA와 같은 고급 기법들이 하나의 조합 가능한 프레임워크에 통합될 수 있는가?
• RQ5 모델러가 각 샘플링 단계에서 기울기 추정기를 자유롭게 선택할 수 있는 정도는 어느 정도이며, 이로 인해 기울기 추정의 정확성이 손상되지 않는가?

주요 결과

• Storchastic는 네 가지 추정기 구성 요소가 특정 수학적 조건을 만족할 경우 임의의 순서의 도함수 추정이 편향 없음을 공식적으로 증명한 최초의 프레임워크이다.
• 이 프레임워크는 재파라미터라이제이션, 점수 함수 추정기, RELAX, MAPO, 순서 없는 집합 추정기, 측도 기반 파생도를 포함한 다양한 기울기 추정 방법을 지원한다.
• Storchastic는 제어 변수 기반 분산 감소를 임의의 순서의 도함수 추정으로 확장할 수 있도록 하여, 분야에 있어 새로운 기여를 한다.
• PyTorch에서의 구현 덕분에 모델러는 단 한 줄의 코드 수정만으로 다양한 추정기를 쉽게 비교할 수 있으며, 이는 VAE 실험에서 메서드 비교를 단순화함을 보여준다.
• MNIST VAE에서의 실험 결과는 프레임워크가 다양한 추정기를 정확히 지원하며, 예상되는 성능 순서가 이전 연구와 일치함을 보여주며, 제한된 학습 에포크와 하이퍼파rameter 튜닝 없이도 성능이 기대에 부합함을 확인한다.
• SPSA와 측도 기반 파생도가 스토하스틱 AD 파이프라인에 성공적으로 통합되어, 이 프레임워크의 일반성과 확장 가능성의 우수성을 입증한다.