[논문 리뷰] Randomized Automatic Differentiation

Deniz Oktay, Nick McGreivy외 3명|arXiv (Cornell University)|2020. 07. 20.

Gaussian Processes and Bayesian Inference참고 문헌 34인용 수 8

한 줄 요약

이 논문은 랜덤화된 자동 미분(RAD)을 소개한다. RAD는 정확한 역방향 자동 미분 대신 편향이 없고 메모리 사용이 적은 기울기 추정치를 도입하여 자동 미분 과정에 랜덤성을 도입하는 프레임워크이다. 메모리 소비를 늘려 분산을 늘리는 대신, RAD는 확률적 최적화에서 더 빠른 수렴을 가능하게 한다. 이는 전방향 및 순환 신경망에서 검증되었으며, 반응형 냉각제 제어를 위한 선형 PDE 문제에 적용되어 높은 메모리 절감 효과를 보였다.

ABSTRACT

The successes of deep learning, variational inference, and many other fields have been aided by specialized implementations of reverse-mode automatic differentiation (AD) to compute gradients of mega-dimensional objectives. The AD techniques underlying these tools were designed to compute exact gradients to numerical precision, but modern machine learning models are almost always trained with stochastic gradient descent. Why spend computation and memory on exact (minibatch) gradients only to use them for stochastic optimization? We develop a general framework and approach for randomized automatic differentiation (RAD), which can allow unbiased gradient estimates to be computed with reduced memory in return for variance. We examine limitations of the general approach, and argue that we must leverage problem specific structure to realize benefits. We develop RAD techniques for a variety of simple neural network architectures, and show that for a fixed memory budget, RAD converges in fewer iterations than using a small batch size for feedforward networks, and in a similar number for recurrent networks. We also show that RAD can be applied to scientific computing, and use it to develop a low-memory stochastic gradient method for optimizing the control parameters of a linear reaction-diffusion PDE representing a fission reactor.

연구 동기 및 목표

• 기존 자동 미분에서 도출되는 정확한 기울기와 현대 딥 러닝 최적화의 확률적 성격 사이의 괴리 문제를 해결하기 위해.
• 역방향 자동 미분의 메모리 사용량을 줄이기 위해 편향 없는 추정치를 도입하는 랜덤화된 방법을 통해 저장소 요구량을 감소시키기 위해.
• 신경망과 물리 시뮬레이터에 적용 가능한 일반적인 프레임워크를 개발하기 위해.
• 고정된 메모리 예산 하에서 RAD가 작은 배치 학습보다 더 빠른 수렴을 달성할 수 있는지 입증하기 위해.
• 고비용의 메모리 요구량을 수반하는 대규모 과학 계산 문제에 자동 미분의 적용 범위를 확장하기 위해.

제안 방법

• 계산 그래프 경로의 무작위 샘플링을 통해 편향 없는 기울기 추정치를 계산하는 역방향 자동 미분의 랜덤화된 변형을 제안한다.
• 모서리 가중치(국소 도함수)를 무작위로 부분 추출하여 저장소를 줄이는 선형화된 계산 그래프(LCG)에 랜덤화를 적용한다.
• Bauer의 공식을 사용하여 LCG 내 경로 열거를 통해 야코비안을 계산하고, 중간 도함수의 저장소를 줄이기 위해 무작위 경로 선택을 수행한다.
• LCG 내 일부 경로를 선택하는 부분 추출 기법을 도입하여 편향을 유지하면서도 저장소 사용량을 줄인다.
• 비가역적인 그래프에 대해 추정할 수 없는 경우를 대비해 타원형 급수 추정치를 사용하여 복잡한 모델에서 계산 가능한 기울기 추정치를 가능하게 한다.
• RAD를 전방향 및 순환 신경망에 적용하고, 선형 반응-확산 PDE에 적용하여 메모리 효율성을 입증한다.

실험 결과

연구 질문

• RQ1 랜덤화된 자동 미분는 정확한 기울기 추정치를 유지하면서도 역방향 미분의 메모리 사용량을 줄일 수 있는가?
• RQ2 고정된 메모리 예산 하에서 RAD는 작은 배치 학습에 비해 수렴 속도가 빠른가?
• RQ3 고메모리 요구량을 수반하는 과학 계산 문제, 예를 들어 PDE 최적화에 RAD를 효과적으로 적용할 수 있는가?
• RQ4 계산 그래프의 어떤 구조적 특성이 편향 없이 효과적인 랜덤화를 가능하게 하는가?
• RQ5 RAD 추정치의 분산은 표준 자동 미분에 비해 어떻게 되며, 이를 효과적으로 제어할 수 있는가?

주요 결과

• 전방향 신경망의 경우, 동일한 메모리 예산 하에서 RAD는 작은 배치 학습보다 더 적은 반복 수로 수렴한다.
• 순환 신경망의 경우, 메모리 사용량이 감소했음에도 불구하고 RAD는 작은 배치 학습과 유사한 반복 수로 수렴한다.
• RAD는 반응-확산 PDE의 저메모리 최적화를 가능하게 하여 과학 계산 분야에의 적용 가능성을 입증한다.
• 선형화된 계산 그래프에서 경로 부분 추출을 통해 편향 없는 기울기 추정치를 유지하면서도 메모리 요구량을 크게 줄였다.
• 문제의 특정 구조를 활용하여 편향 없이 저장소를 줄일 수 있기 때문에, 메모리 제약 조건 하에서 표준 자동 미분보다 RAD가 더 우수한 성능을 보였다.
• 이 방법은 일반화 가능하며, 비선형성과 루프를 포함한 다른 계산 그래프에 대해서도 랜덤화된 반올림 및 타원형 추정치를 통해 확장 가능하다.