📂머신러닝

논문 리뷰: Neural Ordinary Differential Equations

개요 및 요약

「Neural Ordinary Differential Equations」는 Ricky T. Q. Chen 외 3명이 2018년에 발표한 논문으로, 2018 NeurIPS Best Papers에 선정되었다. 다음과 같이 간단한 1계 상미분 방정식인 비자율 시스템을 신경망으로 근사하는 방법을 제안한다.

$$\dfrac{\mathrm{d}y}{\mathrm{d}t} = f(y, t)$$

주목할 점은 신경망이 근사(예측)하는 것은 $y$가 아니라 변화율 $f$라는 것이다. 위 식의 양변을 $0$에서 임의의 $T$까지 적분하면 다음과 같다.

$$y(T) - y(0) = \int\limits_{0}^{T} f(y, t) \mathrm{d}t \implies y(T) = y(0) + \int\limits_{0}^{T} f(y, t) \mathrm{d}t$$

따라서 $f$를 정확히 근사할 수만 있으면, 임의의 $T$에 대해서 $y(T)$를 얻을 수 있다.

1 Introduction

아래의 그림과 같이 균등한 시간 간격인 8개의 데이터 포인트로 이루어진 시계열데이터가 주어져있다고 하자.

이런 상황에서 보통 우리가 하고 싶은 일은 $t_{9}$, $t_{10}$, $t_{11}$, $\dots$에 대한 데이터 $\mathbf{h}_{9}$, $\mathbf{h}_{10}$, $\mathbf{h}_{11}$, $\dots$를 예측하는 것이다. 이를 위해 쉽게 생각할 수 있는 방법 중 하나는 잔차 신경망, RNN, normailizing flow 등을 사용하여 다음과 같이 모델링하는 것이다.

$$\mathbf{h}_{t+1} = \mathbf{h}_{t} + f(\mathbf{h}_{t}; \theta) \tag{1}$$

여기서 $f$는 신경망이고, $\theta$는 $f$의 학습해야할 파라미터(가중치)이다. 이런 방법은 간단하고 직관적이며, 가장 먼저 시도해볼 수 있는 방법 중 하나이지만 다음과 같은 단점이 있다.

• 데이터를 보간하기는 힘들다. $(1)$과 같은 모델로는 $t_{3}$과 $t_{4}$ 사이의 데이터를 예측하기 힘들다. 특히나 임의의 시간 $t_{3.472}$에 대한 데이터를 예측하기는 쉽지않다.
• 시간간격이 균등하지 않은 데이터에 대해서는 적용하기 어렵다. $(1)$에서 $\mathbf{h}_{t+1}$를 업데이트하는 규칙은 $f$를 한 번, 두 번, 정수 횟수로 더해가는 것이기 때문에 시간 간격이 균등하지 않으면 $(1)$의 등식은 성립하지 않는다.

이 논문에서 제안하는 방법은 이산 자율 시스템인 $(1)$을 다음과 같이 연속적인 비자율 시스템으로 바꾸어 생각하는 것이다.

$$\dfrac{\mathrm{d}\mathbf{h}}{\mathrm{d}t} = f(\mathbf{h}(t), t; \theta) \tag{2}$$

"비자율 시스템에서 외력(혹은 속도) $f$를 파라미터가 $\theta$인 인공신경망으로 정의한 $(2)$"를 Neural Ordinary Differential Equations (Neural ODE, NODE)라 한다.

이런 ODE를 풀어주는 솔버(즉 $f$를 적분해주는 것)는 이미 많이 연구되어있고, 성능 좋은 다양한 방법들이 제안되어있다. 따라서 인경신공망을 통해 $f$를 잘 근사했다면 아래의 식을 통해 임의의 시간 $t$에 대한 $\mathbf{h}(t)$를 예측할 수 있다. 양 변을 적분하면,

$$\mathbf{h}(t) = \mathbf{h}(0) + \int\limits_{0}^{t} f(\mathbf{h}(t), s; \theta) \mathrm{d}s \tag{3}$$

우변은 초기값 $\mathbf{h}(0)$와 $f$가 주어져있다면 ODE 솔버를 통해 얻을 수 있는 값이므로, 논문에서는 다음과 같이 표기한다.

$$\mathbf{h}(t_{1}) = \operatorname{ODESolve}(\mathbf{h}(t_{0}), f, t_{0}, t_{1}, \theta)$$

논문에서는 ODE 솔버로 모델을 정의하고 값을 계산하면 얻을 수 있는 장점을 다음과 같이 설명한다.

• Memory efficiency^{메모리 효율}:

  아래의 2장에서 역전파 알고리즘을 사용하지 않고도 손실 함수의 그래디언트를 계산하는 방법을 소개한다. 따라서 모델의 함숫값을 계산할 때 중간 계산 결과를 저장하지 않아도 되며, 이는 메모리 사용량이 모델의 크기에 의존하지 않고 일정하게 유지된다는 것을 의미한다.

• Adaptive computation^{적응형 계산}:

  오일러 메소드 이후로 120년이 지나는 동안 더 정확하고 효율적인 ODE 솔버들이 많이 개발되어왔다. 특히 최근의 ODE 솔버들은 오차 모니터링 및 높은 정확도 달성을 위해 실행 중에 평가 전략^{evaluation strategy}을 동적으로 조정하는 기능을 제공한다.

• Scalable and invertible normalizing flows^{노멀라이징 플로우의 일반화}:

  연속적인 모델의 부수적인 이점 중 하나는 변수 변환 공식을 계산하기 쉽다는 것인데 (4장에서 설명한다), 이는 효율적인 노멀라이징 플로우 모델을 만들 수 있다는 것을 의미한다.

• Continuous time-series models^{연속 시계열 모델}:

  훈련 및 예측 데이터가 균등간격으로 이산화되어있어야하는 RNN과 달리, 연속적으로 정의된 Neural ODE는 임의의 시간 $t$에 대한 예측을 수행할 수 있다. 이 내용은 5장에서 자세히 다룬다.

2 Reverse-mode automatic differentiation of ODE solutions

Neural ODE에 대한 핵심 아이디어와 원리는 Introduction에서 모두 설명하였다. 2장은 학습법에 대해서 다룬다. 저자는 Neural ODE를 구현하는 방법을 개발하여 깃허브에 공개 해두었다. 직접 구현할것이 아니라서 학습에 관한 내용이 크게 궁금하지 않다면 넘어가도 좋다고 본다. 3~5장은 Neural ODE가 어떻게 응용될 수 있는지를 다루기 때문에, 본인의 필드에서 Neural ODE를 활용하고 싶다면 3~5장을 상세히 읽어보는 것을 추천한다. 특히나 4장의 내용은 diffusion model 다음으로 생성모델의 주류를 이끌고있는 flow matching과도 연결된다.

관측값(정답 데이터)은 신경망의 파라미터 $\theta$에 의존하지않는 상수이므로, 손실함수를 다음과 같이 간단히 나타내자.

$$L(\mathbf{z}(t_{1})) = L \left( \mathbf{z}(t_{0}) + \int_{t_{0}}^{t_{1}} f(\mathbf{z}(t), t; \theta) \mathrm{d}t \right) = L \left( \operatorname{OSESolve}(\mathbf{z}(t_{0}), f_{\theta}, t_{0}, t_{1}) \right)$$

논문에서 손실함수의 그래디언트를 계산하기 위해 방법은 adjoint sensitivity method이며, 이는 또다른 ODE를 시간 역순으로 적분하는 것이다. 그 결과만 보면 다음과 같다.

$$\dfrac{\mathrm{d}L}{\mathrm{d}\theta} = \int\limits_{t_{1}}^{t_{0}} \left( \dfrac{\partial L}{\partial \mathbf{z}(t)} \right)^{\mathsf{T}} \dfrac{\partial f(\mathbf{z}(t), t; \theta)}{\partial \theta} \mathrm{d}t$$

자세한 내용은 Appendix B, C, D에서 확인할 수 있다. (다운로드 링크) 또한 아래에서 설명한 예제에 대한 코드와 학습에 관련된 메서드는 torchdiffeq 패키지에 포함되어있으며, 깃허브에 공개되어있다.

3 Replacing residual networks with ODEs for supervised learning

Neural ODE는 $(1)$을 $(3)$과 같이 바꾼 것이라 생각하면, residual network의 일반화로 볼 수 있다.

$$\text{residual: } \mathbf{h}_{t+1} = \mathbf{h}_{t} + f(\mathbf{h}_{t}; \theta) \quad\overset{\text{generalization}}{\implies}\quad \dfrac{\mathrm{d}\mathbf{h}}{\mathrm{d}t} = f(\mathbf{h}(t), t; \theta)$$

그래서 ResNet을 활용한 지도학습에서 Neural ODE를 활용할 수 있다. 저자들은 MNIST 데이터셋의 판별문제에서 ResNet 부분을 그대로 Neural ODE로 치환하여 성능을 비교하였다.

ODE의 초기값 문제를 수치적으로 풀기위해 선택한 솔버는 아담스 메소드이다. scipy.integrate 패키지로 제공되는 암시적 메소드를 사용하였다. 최적화 기법인 adjoint sensitivity method를 pytorch.autograd로 구현하였다.

실험 결과로 ODE-Net은 ResNet에 비해서 $1/3$ 정도의 파라미터를 가짐에도 불구하고 거의 비슷한 성능을 보여주었다. $L$은 ResNet의 레이어 수를 의미하며, ResNet의 계산 시간과 역전파에 필요한 메모리는 $L$에 비례한다. $\tilde{L}$은 ODE 솔버의 number of function evaluations (NFE)인데, 적분 구간을 얼마나 나누어 계산할 것인지를 의미한다. ODE 솔버는 수치적분을 수행하므로 적분 구간을 나누게되는데 NFE는 이 구간의 수를 말한다. 가령 $[0, 1]$구간을 $0.05$의 간격으로 적분을 수행하면 $\text{NFE} = \dfrac{1}{0.05} = 20$이다. 논문에서는 이를 ODE-Net의 layer 수로 해석할 수 있다고 언급한다. NFE가 클수록 정확하게 계산할 수 있지만, 계산시간이 오래 걸리므로 적당한 숫자를 선택하는 것이 중요하다. Neural ODE에서는 NFE를 모델이 알아서 선택한다.

Figure 3a, 3b에서는 각각 순방향 계산에서의 NFE에 따른 오차와 계산 시간을 보여주며, 직관에 잘 들어맞는 결과라 할 수 있다. Figure 3c를 보면 역방향 계산이 순방향 계산에 비해 절반 정도의 수준이라는 것을 알 수 있다. 이는 제안하는 학습 방법 adjoint sensitivity method가 꽤 효율적이라는 것을 보여준다. Figure 3d는 훈련이 진행됨에 따라 NFE가 증가함을 보여주는데, 이는 모델이 훈련되며 점점 복잡해지고 이를 잘 표현하기 위해 정교해지는 것이라 설명한다.

4 Continuous Normalizing Flows

$(1)$과 같은 이산적 모델링은 노멀라이징 플로우^{normalizing flow (NF)} 에서도 나타난다.

$$\mathbf{z}_{1} = f(\mathbf{z}_{0})$$

$$\log p(\mathbf{z}_{1}) = \log p(\mathbf{z}_{0}) - \log \left| \det \dfrac{\partial f}{\partial \mathbf{z}_{0}} \right| \tag{a4}$$

간단히 planar flow를 예로 들면,

$$\mathbf{z}_{t+1} = \mathbf{z}_{t} + \mathbf{u} \tanh (\mathbf{w}^{\mathsf{T}}\mathbf{z}_{t} + b)$$

$$\log p(\mathbf{z}_{t+1}) = \log p(\mathbf{z}_{t}) - \log \left| 1 + \mathbf{u}^{\mathsf{T}} \dfrac{\partial \tanh}{\partial \mathbf{z}} \right|$$

노멀 플로우를 학습하는데 있어서 가장 중요하게 고려해야할 사항은 $(a4)$의 행렬식 계산이다. 이는 $\mathbf{z}$의 차원이나 가중치 개수에 대해서 세제곱의 계산비용을 가진다. 흥미롭게도 이산적 구조 $(1)$에서 연속적구조 $(3)$으로 바꾸면 이와 관련된 계산이 더 단순해진다.

---

$\textbf{Theorem 1 }\text{(Instantaneous Change of Variables).}$

$\mathbf{z}(t)$를 연속확률변수, $p(\mathbf{z}(t))$를 이에 대한 확률밀도함수라 하자. 그리고 비자율 미분방정식이 다음과 같이 주어졌다고 하자.

$$\dfrac{\mathrm{d} \mathbf{z}}{\mathrm{d}t} = f(\mathbf{z}(t), t) \tag{a5}$$

$f$가 변수 $\mathbf{z}$에 대해서 립시츠연속이고, 변수 $t$에 대해서 연속이라고 하자. 그러면 로그확률밀도의 도함수는 다음과 같다.

$$\dfrac{\partial \log p(\mathbf{z}(t))}{\partial t} = -\Tr \left( \dfrac{\partial f}{\partial \mathbf{z}(t)} \right)$$

여기서 $\Tr$은 트레이스이다.

---

$(a4)$의 행렬식 계산이 단순한 트레이스 계산으로 바뀌었다. 또한 표준적인 이산 노멀라이징 플로우에서와 달리, 이 경우에서는 $f$가 전단사일 필요가 없다. 위에서 예로든 planar flow를 위 $\textbf{Theorem 1}$에 맞춰 연속적인 형태로 바꾸면 다음과 같다.

$$\dfrac{\mathrm{d}\mathbf{z}(t)}{\mathrm{d}t} = \mathbf{u} \tanh(\mathbf{w}^{\mathsf{T}}\mathbf{z}(t) + b), \qquad \dfrac{\partial \log p (\mathbf{z}(t))}{\partial t} = -\Tr \left( \mathbf{u} \dfrac{\partial \tanh(\mathbf{w}^{\mathsf{T}}\mathbf{z}(t) + b)}{\partial \mathbf{z}(t)} \right)$$

Using multiple hidden units with linear cost

행렬식은 선형이 아니지만, 트레이스는 선형이므로 $\Tr \sum\limits_{n} J_{n} = \sum\limits_{n} \Tr J_{n}$이 성립한다. 따라서 $(a5)$의 $f$를 선형결합으로 나타내어도, 로그확률밀도함수의 도함수는 다음과 같이 간단히 나타낼 수 있다.

$$\dfrac{\mathrm{d} \mathbf{z}(t)}{\mathrm{d}t} = \sum\limits_{n=1}^{M} f_{n}(\mathbf{z}(t)), \qquad \dfrac{\mathrm{d}\log p(\mathbf{z}(t))}{\mathrm{d}t} = -\sum\limits_{n=1}^{M} \Tr \left( \dfrac{\partial f_{n}}{\partial \mathbf{z}(t)} \right)$$

이는 은닉층의 노드 수 $M$에 대해 계산 비용이 선형적으로 증가함을 의미한다. 기존의 노멀라이징 플로우는 $\mathcal{O}(M^{3})$의 계산 비용 때문에 대부분 한 개의 노드를 층을 많이 쌓는 구조로 사용한다.

Time-dependent dynamics

논문에서는 gating이라 불리는 방법을 사용하여 플로우를 아래와 같이 정의한다.

$$\dfrac{\mathrm{d} \mathbf{z}}{\mathrm{d}t} = \sum\limits_{n} \sigma_{n}(t)f_{n}(\mathbf{z}), \qquad \sigma_{n}(t) \in (0 1)$$

즉 $f(\mathbf{z}, t)$가 변수분리 가능하다고 가정하는 것이다. 여기서 $\sigma_{n}$도 $f_{n}$도 모두 학습해야할 신경망이다. 논문에서는 이를 연속 노멀라이징 플로우^{consinuous normalizing flow (CNF)}라 한다. [플로우 매칭]은 이 CNF를 더 효율적으로 학습할 수 있는 하나의 방법이다.

4.1 Experiments with Continuous Normalizing Flows

CNF와 NF를 비교한다. CNF는 위에서 설명한대로 구현되었으며, Adam 옵티마이저를 사용하여 10,000 반복하여 훈련하였다. NF는 처음 제안된 논문에서와 같이 구현되었으며, RMSprop 옵티마이저를 사용하여 500,000 반복하여 훈련하였다. CNF 쪽이 반복 횟수가 현저히 적다. 결과는 아래 그림에서 확인할 수 있다.

아래의 Figure 5에서 위의 두 행은 CNF, 마지막 행은 NF에 대한 결과로 보인다. CNF는 매끄럽게 변환되는 반면, NF는 그렇지 못하며 Two Moons 데이터에 대해서는 제대로 모델링하지 못했다.

5 A generative latent function time-series model

Neural ODE는 비자율 시스템을 신경망으로 근사하는 방법이니 만큼, 시계열데이터의 예측, 보간, 결측치 생성 등에도 사용할 수 있다. 논문에서 제안하는 훈련 방법은 아래와 같다.

0. 시계열 데이터 $\left\{ \mathbf{x}_{i}, t_{i} \right\}_{i=1}^{N}$이 주어져있다고 하자.
1. RNN을 인코더로 사용하여 잠재변수 $\boldsymbol{\mu}$와 $\boldsymbol{\sigma}$를 추출한다. $$(\boldsymbol{\mu}, \log\boldsymbol{\sigma}^{2}) = \operatorname{RNNencoder}\left( \left\{ \mathbf{x}_{i}, t_{i} \right\} \right)$$
2. 정규분포 $N(\mathbf{0}, I)$에서 $\boldsymbol{\epsilon}$를 추출하고 $\boldsymbol{\mu}$와 $\boldsymbol{\sigma}$를 사용하여 잠재변수의 초기값 $\mathbf{z}_{t_{0}}$를 생성한다. $$\mathbf{z}_{t_{0}} = \exp(0.5 \log\boldsymbol{\sigma}^{2}) \odot \boldsymbol{\epsilon} + \boldsymbol{\mu}$$ 이렇게하면 마치 $\operatorname{RNNencoder}$가 시계열데이터를 잠재공간 $N(\boldsymbol{\mu}, \diag(\boldsymbol{\sigma}^{2}))$로 매핑하는 것과 같은 효과를 주면서도, 역전파가 가능하다.
3. 잠재변수의 초기값 $\mathbf{z}_{t_{0}}$와 Neural ODE를 통해 $\mathbf{z}_{t_{i}}$를 계산한다. $$(\mathbf{z}_{t_{1}}, \mathbf{z}_{t_{2}}, \dots, \mathbf{z}_{t_{N}}) = \operatorname{ODESolve}(\mathbf{z}_{t_{0}}, f, t_{0}, \dots, t_{N}, \theta)$$
4. $\mathbf{z}_{t_{i}}$를 FCN으로 정의한 디코더 입력하여 예측된 시계열 데이터 $\mathbf{x}_{i}$를 얻는다. $$\mathbf{x}_{t_{i}} = \operatorname{FCNdecoder}(\mathbf{z}_{t_{i}})$$
5. ELBO를 최대화 한다. $$\text{ELBO} = \sum\limits_{i} \log p(\mathbf{x}_{t_{i}} | \mathbf{z}_{t_{i}}) + \log p(\mathbf{z}_{t_{0}}) - \log q(\mathbf{z}_{t_{0}} | \left\{ \mathbf{x}_{t_{i}}, t_{i} \right\})$$

5.1 Time-series Latent ODE Experiments

실험에 사용한 신경망은 다음과 같다.

• $\operatorname{RNNencoder}$: 노드가 25개인 RNN
• 잠개공간은 4차원. $(\boldsymbol{\mu}), (\boldsymbol{\sigma}) \in \mathbb{R}^{2 \times 2}$
• NeuralODE $f$: 노드가 20개인 은닉층 하나를 가진 MLP
• $\operatorname{FCNdecoder}$: 노드가 20개인 은닉층 하나를 가진 또 다른 MLP

데이터셋에 대한 조건은 아래와 같다.

• 2차원 나선형 데이터 1,000개 생성.
• 절반은 시계방향, 나머지는 시계반대방향.
• 각 궤적은 서로 다른 초기값에서 시작하여 균등한 시간 간격으로 100개의 점 샘플링.
• 관측치에 가우시안 노이즈 추가함.
• 학습할 때는 각 궤적의 100개의 점 중에서 중복없이 무작위로 30/50/100개를 추출하여 사용함.

아래의 표를 보면 NeuralODE의 성능이 현저히 좋은 것을 알 수 있고, 또한 데이터가 적을 때도 많을 때에 비해서 성능이 크게 내려가지 않는다는 것을 확인할 수 있다.

아래의 그림은 RNN과 NeuralODE의 예측 결과이다. NeuralODE의 결과가 더 매끄러우며, 학습하지않은 바깥 영역에서의 예측 또한 더 낫다는 것을 보여둔다.

아래의 그림은 잠재변수의 궤적을 처음 두 차원에 대해서 그린 것이다. 궤적들은 시계방향의 나선과 시계반대방향의 나선 두 가지 군집으로 나뉘는데, 이는 제안하는 방법이 잠재변수를 잘 묘사한다는 것을 의미한다.