[논문 리뷰] Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE

Juntang Zhuang, Nicha C. Dvornek외 4명|arXiv (Cornell University)|2020. 06. 03.

Model Reduction and Neural Networks참고 문헌 41인용 수 34

한 줄 요약

본 논문은 Adaptive Checkpoint Adjoint (ACA) 방법을 제시하여 Neural ODEs에서 forward와 reverse 궤적을 정렬하고 그래프 깊이를 줄이며 adaptive solvers를 가능하게 하여, 더 우수한 정확도와 효율성을 달성합니다.

ABSTRACT

Neural ordinary differential equations (NODEs) have recently attracted increasing attention; however, their empirical performance on benchmark tasks (e.g. image classification) are significantly inferior to discrete-layer models. We demonstrate an explanation for their poorer performance is the inaccuracy of existing gradient estimation methods: the adjoint method has numerical errors in reverse-mode integration; the naive method directly back-propagates through ODE solvers, but suffers from a redundantly deep computation graph when searching for the optimal stepsize. We propose the Adaptive Checkpoint Adjoint (ACA) method: in automatic differentiation, ACA applies a trajectory checkpoint strategy which records the forward-mode trajectory as the reverse-mode trajectory to guarantee accuracy; ACA deletes redundant components for shallow computation graphs; and ACA supports adaptive solvers. On image classification tasks, compared with the adjoint and naive method, ACA achieves half the error rate in half the training time; NODE trained with ACA outperforms ResNet in both accuracy and test-retest reliability. On time-series modeling, ACA outperforms competing methods. Finally, in an example of the three-body problem, we show NODE with ACA can incorporate physical knowledge to achieve better accuracy. We provide the PyTorch implementation of ACA: \url{https://github.com/juntang-zhuang/torch-ACA}.

연구 동기 및 목표

• NODE 그래디언트 추정치가 기존 방법으로 편향되거나 비효율적인 이유를 설명한다.
• 앞으로의 순방향 궤적과 역방향 궤적을 정렬하여 정확한 adjoint 그래디언트를 얻기 위한 적응형 체크포인팅 전략을 개발한다.
• 계산 그래프에서 중복 구성 요소를 제거하여 계산 비용을 낮춘다.
• ACA가 이미지 분류 및 시계열 작업에서 오차율을 줄이고 학습 속도를 높임을 보여준다.
• ACA를 이용한 NODE가 다이나믹 시스템에서 물리적 지식을 반영하여 정확도를 개선할 수 있음을 보인다.

제안 방법

• Adaptive Checkpoint Adjoint (ACA)를 도입하여 역전파 시 순방향 궤적을 기록함으로써 그래디언트 정확성을 보장한다.
• trajectory checkpointing을 적용하여 순방향 계산과 역방향 계산을 정렬한다.
• 중복 구성 요소를 제거하여 더 얕은 계산 그래프를 생성한다.
• ACA 프레임워크 내에서 adaptive ODE solver를 지원한다.
• 재현 가능성을 위해 ACA의 PyTorch 구현을 제공한다.

실험 결과

연구 질문

• RQ1 ACA가 전통적인 adjoint 및 naive 방법에 비해 Neural ODE에서 그래디언트의 수치적 정확성을 향상시킬 수 있는가?
• RQ2 ACA가 이미지 분류 작업에서 정확도 유지 또는 향상을 보장하면서 학습 시간을 단축하는가?
• RQ3 ACA가 시계열 모델링 작업에서 경쟁 방법과 비교하여 어떤 성능을 보이는가?
• RQ4 ACA를 사용하여 NODE에 물리적 지식을 반영하여 다이나믹 시스템에서 정확도를 높일 수 있는가?

주요 결과

• ACA는 이미지 분류 작업에서 adjoint 및 naive 방법에 비해 약 절반의 오차율로 학습 시간이 절반 정도 소요된다.
• ACA로 학습된 NODE는 정확도와 테스트 재현성에서 ResNet보다 우수하다.
• ACA는 시계열 모델링에서 경쟁 방법에 비해 우수한 성능을 보인다.
• 삼체 문제 예제에서 ACA를 적용한 NODE가 물리적 지식을 더 잘 활용하여 정확도를 높인다.