[논문 리뷰] Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory

Takashi Matsubara, Yuto Miyatake외 1명|arXiv (Cornell University)|2021. 02. 19.

Model Reduction and Neural Networks참고 문헌 36인용 수 5

한 줄 요약

이 논문은 대칭성 보존 적분기를 사용하여 당김 시스템을 해결함으로써 신경 미분방정식(Neural ODEs)에서 정확한 기울기를 계산하는 새로운 방법인 대칭성 당김 방법(symplectic adjoint method)을 제안한다. 역전파(backpropagation)는 신경망 $ f $ 의 각 사용에 대해서만 적용하고 대칭성 보존 적분기를 활용함으로써, 단계 수와 네트워크 크기 비례하는 메모리 소비로 정확한 기울기를 달성한다. 이는 역전파나 체크포인팅 기법보다 훨씬 낮으며, 표준 당김 방법보다도 속도와 반올림 오차에 대한 강건성에서 뛰어나다.

ABSTRACT

A neural network model of a differential equation, namely neural ODE, has enabled the learning of continuous-time dynamical systems and probabilistic distributions with high accuracy. The neural ODE uses the same network repeatedly during a numerical integration. The memory consumption of the backpropagation algorithm is proportional to the number of uses times the network size. This is true even if a checkpointing scheme divides the computation graph into sub-graphs. Otherwise, the adjoint method obtains a gradient by a numerical integration backward in time. Although this method consumes memory only for a single network use, it requires high computational cost to suppress numerical errors. This study proposes the symplectic adjoint method, which is an adjoint method solved by a symplectic integrator. The symplectic adjoint method obtains the exact gradient (up to rounding error) with memory proportional to the number of uses plus the network size. The experimental results demonstrate that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, performs faster than the adjoint method, and is more robust to rounding errors.

연구 동기 및 목표

• 신경 미분방정식 훈련에서 역전파 및 체크포인팅 기법의 높은 메모리 소비 문제를 해결하기 위해.
• 신경 미분방정식의 기울기 계산에서 계산 비용을 줄이고 수치 오차에 대한 강건성을 향상시키기 위해.
• 작은 단계 크기 없이도 이산 시간에서 정확한 기울기 계산을 가능하게 하기 위해.
• 당김 방법 수준의 메모리 효율성을 유지하면서도 높은 정확도와 속도를 확보하기 위해.
• 물리적 편미분방정식(PDEs)과 같은 강성(stiff) 또는 고차수 시스템에서 효과적으로 작동하는 강건하고 확장 가능한 훈련 솔루션을 제공하기 위해.

제안 방법

• 방법은 시간에 따라 거꾸로 당김 시스템을 해결하기 위해 대칭성 보존 적분기를 사용하여 기하학적 구조를 유지하고, 반올림 오차 수준 이내의 정확한 기울기를 가능하게 한다.
• 역전파 적용 범위는 전체 계산 그래프가 아니라 각각의 신경망 $ f $ 사용에 국한되며, 이로 인해 메모리 소비가 감소한다.
• 적절한 체크포인트를 사용하여 중간 상태 $ x_n $ 와 내부 단계 값 $ X_{n,i} $ 를 저장함으로써 메모리 오버헤드를 최소화한다.
• 이 방법은 어떤 룬게-쿠타 방법과도 호환되며, 명시적 및 암시적 적분기 모두 지원한다.
• 대칭성 구조는 수치적 안정성을 보장하고 후방 적분 중 오차 누적도 줄인다.
• 표준 당김 방법과 달리, 당김 적분에서 작은 단계 크기를 필요로 하지 않아 표준 당김 방법보다 더 빠른 계산이 가능하다.

실험 결과

연구 질문

• RQ1 작은 메모리 오버헤드로 대칭성 보존 적분기를 사용하여 신경 미분방정식에서 정확한 기울기를 계산할 수 있는가?
• RQ2 제안된 방법의 메모리 소비는 역전파, 체크포인팅, 표준 당김 방법과 비교해 어떻게 되는가?
• RQ3 대칭성 당김 방법은 정확한 기울기 유지 조건에서 표준 당김 방법보다 더 빠른 계산을 달성하는가?
• RQ4 기존의 역전파 및 당김 접근 방식과 비교해 반올림 오차에 대해 얼마나 강건한가?
• RQ5 KdV 방정식 및 케인-힐리아르 방정식과 같은 강성 또는 고차수 시스템을 효과적으로 처리할 수 있는가?

주요 결과

• 대칭성 당김 방법은 정방향 적분과 동일한 단계 크기로도 반올림 오차 수준 이내의 정확한 기울기를 달성한다. 이는 표준 당김 방법이 오차 제어를 위해 더 작은 단계 크기를 필요로 하는 것과 대비된다.
• 메모리 소비는 $ O(MN + s) $ 로 표현되며, 여기서 $ M $ 은 성분 수, $ N $ 은 시간 단계 수, $ s $ 는 내부 단계 수이다. 이는 역전파의 $ O(MNsL) $ 와 비교해 극적으로 낮으며, 당김 방법과 유사한 수준이다.
• KdV 방정식에 대한 실험에서, 메모리 소비 79.8 MiB, MSE $ 1.61 \pm 4.00 \times 10^{-3} $ 으로, 1회 반복당 162 ms로 표준 당김 방법(276 ms/itr)보다 빠른 속도를 기록했다.
• Cahn–Hilliard 시스템에선 MSE $ 5.47 \pm 1.46 \times 10^{-6} $, 메모리 80.3 MiB, 1회 반복당 568 ms로, 강성 시스템에서의 강건성과 효율성을 입증했다.
• 장기적인 계산 그래프를 거쳐 오차를 누적시키는 대신 각 적분 단계에서 기울기를 계산하므로, 역전파 및 체크포인팅 기법보다 반올림 오차에 더 강건하다.
• 표준 당김 방법과 달리, 수치 오차를 억제하기 위해 작은 단계 크기를 필요로 하지 않기 때문에 실질적으로 표준 당김 방법보다 더 빠르다.