[논문 리뷰] MALI: A memory efficient and reverse accurate integrator for Neural ODEs

Juntang Zhuang, Nicha C. Dvornek외 2명|arXiv (Cornell University)|2021. 05. 03.

Model Reduction and Neural Networks참고 문헌 38인용 수 7

한 줄 요약

MALI는 기존 방법이 시간에 따라 증가하는 것과는 달리 통합 시간 동안 일정한 메모리 사용을 달성하기 위해 异步 레프트로그(ASYNCHRONOUS LEAPFROG, ALF) 솔버를 활용하는 메모리 효율적이고 역방향 정확도가 높은 Neural ODE용 통합기이다. 이는 역방향 시간에서 궤적의 정밀도를 유지함으로써 정확한 기울기 추정을 가능하게 하며, ImageNet 학습, 시계열 모델링, 연속적 생성 모델링에서 기존 방법을 능가한다.

ABSTRACT

Neural ordinary differential equations (Neural ODEs) are a new family of deep-learning models with continuous depth. However, the numerical estimation of the gradient in the continuous case is not well solved: existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory, while the naive method and the adaptive checkpoint adjoint method (ACA) have a memory cost that grows with integration time. In this project, based on the asynchronous leapfrog (ALF) solver, we propose the Memory-efficient ALF Integrator (MALI), which has a constant memory cost w.r.t integration time similar to the adjoint method, and guarantees accuracy in reverse-time trajectory (hence accuracy in gradient estimation). We validate MALI in various tasks: on image recognition tasks, to our knowledge, MALI is the first to enable feasible training of a Neural ODE on ImageNet and outperform a well-tuned ResNet, while existing methods fail due to either heavy memory burden or inaccuracy; for time series modeling, MALI significantly outperforms the adjoint method; and for continuous generative models, MALI achieves new state-of-the-art performance. We provide a pypi package: https://jzkay12.github.io/TorchDiffEqPack

연구 동기 및 목표

• 기존 Neural ODE 학습 방법에서 야기되는 높은 메모리 비용과 역방향 시간 궤적 추정의 부정확성 문제를 해결하기 위해.
• 통합 시간과 관계없이 일정한 메모리 사용을 유지하는 통합기를 개발하여, 인접 방법의 효율성과 동일시하기 위해.
• 역방향 시간 통합에서 수치 정확도를 확보함으로써 기울기 추정의 신뢰성을 향상시키기 위해.
• ImageNet과 같은 대규모 데이터셋에서 기존 방법이 메모리 또는 정확도 제약으로 실패하는 상황에서도 Neural ODE의 실현 가능한 학습을 가능하게 하기 위해.
• 연속적 생성 모델링 및 시계열 모델링 작업에서 최고 수준의 성능를 달성하기 위해.

제안 방법

• MALI는 낮은 메모리 오버헤드로 안정적이고 정확한 시간 통합을 가능하게 하는 이방향 레프트로그(ASYNCHRONOUS LEAPFROG, ALF) 솔버에 기반한다.
• 이 방법은 역방향 시간 통합에 필요한 중간 상태만 저장하므로, 통합 기간과 관계없이 메모리 비용이 일정하게 유지된다.
• MALI는 인접 계산 동안 일관된 시간단계 동기화를 유지함으로써 역방향 시간 궤적 정확도를 강제한다.
• 에러 누적을 방지하기 위해 메모리 효율적이고 수치적으로 안정적인 체크포인팅 전략을 사용한다.
• 정확도를 유지하면서도 적응형 시간 단계를 지원하므로 복잡한 동역학에 적합하다.
• MALI는 PyTorch에 구현되어 있으며 공개 패키지로 배포되었다: https://jzkay12.github.io/TorchDiffEqPack

실험 결과

연구 질문

• RQ1 긴 통합 시간 동안 메모리 사용이 일정하고 역방향 정확도를 손상시키지 않는 Neural ODE 통합기를 설계할 수 있는가?
• RQ2 MALI는 기존 방법이 메모리 또는 기울기 부정확성으로 실패하는 ImageNet과 같은 대규모 데이터셋에서 Neural ODE의 실현 가능한 학습을 가능하게 하는가?
• RQ3 MALI는 시계열 모델링 작업에서 인접 방법과 비교해 기울기 정확도와 성능 측면에서 어떻게 다른가?
• RQ4 MALI는 연속 정규화 흐름 및 기타 연속 생성 모델링 작업에서 최고 수준의 성능를 달성할 수 있는가?
• RQ5 역방향 시간 궤적 정확도는 다양한 기계학습 작업에서 최종 모델 성능에 어떤 영향을 미치는가?

주요 결과

• MALI는 ImageNet에서 Neural ODE의 첫 실현 가능한 학습을 가능하게 하여, 잘 튜닝된 ResNet을 능가하는 성능을 달성했다.
• MALI는 시계열 모델링 작업에서 인접 방법을 능가하며, 뛰어난 예측 정확도를 입증했다.
• MALI는 특히 정규화 흐름 기반 밀도 추정에서 연속 생성 모델링 분야에서 새로운 최고 수준의 성능를 달성했다.
• ACA 및 단순 방법과는 달리, MALI는 통합 시간에 관계없이 일정한 메모리 비용을 유지한다.
• MALI는 높은 정확도를 유지한 채 역방향 시간 궤적을 제공하여 신뢰할 수 있고 정밀한 기울기 추정을 가능하게 한다.
• PyPI에 배포된 오픈소스 구현은 연구자들이 다양한 연속 깊이 모델에 대해 MALI를 쉽게 채택하고 확장할 수 있도록 한다.