[논문 리뷰] ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs

Amir Gholami, Kurt Keutzer외 1명|arXiv (Cornell University)|2019. 02. 27.

Model Reduction and Neural Networks참고 문헌 26인용 수 77

한 줄 요약

ANODE는 체크포인팅이 있는 discretize-then-optimize adjoint 프레임워크를 도입하여 신경 ODEs에 대한 무조건적으로 정확한 기울기를 계산하고, 메모리를 O(LN_t)에서 O(L)+O(N_t)로 줄이며, 기존의 역시간(reverse-time) 접근법의 불안정성을 피한다.

ABSTRACT

Residual neural networks can be viewed as the forward Euler discretization of an Ordinary Differential Equation (ODE) with a unit time step. This has recently motivated researchers to explore other discretization approaches and train ODE based networks. However, an important challenge of neural ODEs is their prohibitive memory cost during gradient backpropogation. Recently a method proposed in [8], claimed that this memory overhead can be reduced from O(LN_t), where N_t is the number of time steps, down to O(L) by solving forward ODE backwards in time, where L is the depth of the network. However, we will show that this approach may lead to several problems: (i) it may be numerically unstable for ReLU/non-ReLU activations and general convolution operators, and (ii) the proposed optimize-then-discretize approach may lead to divergent training due to inconsistent gradients for small time step sizes. We discuss the underlying problems, and to address them we propose ANODE, an Adjoint based Neural ODE framework which avoids the numerical instability related problems noted above, and provides unconditionally accurate gradients. ANODE has a memory footprint of O(L) + O(N_t), with the same computational cost as reversing ODE solve. We furthermore, discuss a memory efficient algorithm which can further reduce this footprint with a trade-off of additional computational cost. We show results on Cifar-10/100 datasets using ResNet and SqueezeNext neural networks.

연구 동기 및 목표

• 신경 ODE를 훈련할 때의 메모리 문제를 동기 부여하고 분석하며, reverse-time gradient 방법의 실패를 식별한다.
• ANODE를 제안한다, 체크포인팅 기반의 adjoint 프레임워크로 무조건적으로 정확한 기울기를 달성한다.
• ResNet/SqueezeNext 기반 ODE 블록으로 CIFAR-10/100에서 메모리 효율성과 안정성을 입증한다.

제안 방법

• 신경망을 잔차 블록을 갖는 ODE로 모델링하고 forward Euler 또는 기타 이산화 방법을 사용한다.
• 그레이디언트를 계산하기 위해 forward ODE를 역방향으로 되돌리는 것이 수치적 문제와 일관성 문제를 초래함을 보인다.
• 필요한 활성화만 저장하도록 checkpointing을 활용한 discretize-then-optimize 그래디언트 계산을 채택한다.
• 저장된 forward 궤적을 재사용하여 그래디언트를 계산하는 DTO 기반 adjoint 계산을 정의한다.
• 선택적 로그형 checkpointing과 함께 저장을 O(LN_t)에서 O(L)+O(N_t)로 줄이는 메모리 관리 방식 제공.
• Euler 및 RK2 스킴으로 ResNet 및 SqueezeNext 백본을 사용하여 CIFAR-10/100에서 ANODE를 실험한다.

실험 결과

연구 질문

• RQ1 신경 ODE가 메모리 제약 학습 하에서 무조건적으로 정확한 기울기를 허용하는가?
• RQ2 체크포인팅이 있는 discretize-then-optimize adjoint 프레임워크가 일반 NN 블록( ReLU 활성화 및 합성곱 포함)에서 안정적이고 정확한 기울기를 제공하는가?
• RQ3 ANODE가 ODE 레이어를 통한 역시간 방법과 비교했을 때 메모리 vs. 계산 트레이드오프는 어떤가?
• RQ4 이산화 선택(Euler, RK2)이 신경 ODE의 학습 안정성과 정확도에 어떤 영향을 미치는가?

주요 결과

• Reverse-time (backward) ODE 해법은 수치적으로 불안정할 수 있으며 일반 네트워크에서 잘못된 기울기를 생성할 수 있다.
• Optimize-then-discretize (OTD) 기울기는 이산화된 순방향 해법과 일치하지 않아 작은 시간 간격에서 학습 발산을 초래할 수 있다.
• Discretize-then-optimize with checkpointing (ANODE)는 무조건적으로 정확한 기울기와 안정적인 학습을 제공한다.
• ANODE는 forward 계산 비용은 동일하게 유지하면서 저장을 O(LN_t)에서 O(L)+O(N_t)로 줄인다.
• checkpointing이 있는 이산 DTO를 사용하면 안정성을 위해 가중치 노름을 제약하지 않고도 효율적인 그래디언트 계산이 가능하다.
• ResNet/SqueezeNext ODE 블록으로 CIFAR-10/100에서 ANODE가 안정적인 수렴과 역시간 baseline보다 더 나은 성능을 달성하는 것을 보여준다.