반양적 미분대수 방정식의 해를 미분하기: VJP 편

개요

이 글에서는 목적 함수가 지수 1(exponentially 1)의 반양적 미분대수 방정식의 해를 통해 초기값이나 파라미터에 의존할 때, 그 기울기(gradient)를 Adjoint 법을 이용해 수치적으로 계산하는 방법을 유도합니다. 대상이 되는 경우는 목적 함수가 $x_1 ext{은 } ext{R}^{N_x} ext{이고 } y_1, ext{과 } ext{dot } y_1 ext{은 } ext{R}^{N_y} ext{의 함수로}$

$J = J(x_1, y_1, ext{dot } y_1)$

으로 주어지고, $x_1, y_1, ext{dot } y_1$이 지수 1의 반양적 미분대수 방정식 초기값 문제
egin{align*} ext{dot } y &= f(t, x, y; ext{alpha}) \ g(t, x, y; ext{alpha}) &= 0 \ y(t_0) &= y_0 ext{

} ext{의 해인 경우입니다. 여기서 } ext{dot } y = dy/dt ext{입니다. 또한 } ext{alpha} ext{는 } ext{R}^M ext{에 속하는 미분대수 방정식의 파라미터입니다. 지수 1이란, } g = 0 ext{을 } x ext{에 대해 풀 수 있다는 것, 즉 } g ext{의 } x ext{에 대한 야코비 행렬이 정칙(non-singular)하다는 것을 의미합니다. 방정식에 시간 미분 } ext{dot } y ext{가 포함되는 } y ext{는 미분적인 변수이고, 포함되지 않는 } x ext{는 대수적인 변수입니다. 미리 미분적인 변수와 대수적인 변수의 구분이 되어 있는 상황을 가정합니다.

미분대수 방정식을 통한 기울기 계산의 어려움과 해결 방안

기울기 계산의 구체적인 내용

이 글의 목적은 미분대수 방정식의 초기값 $y_0$나 파라미터 $ ext{alpha}$에 관한 목적 함수 $J$의 기울기를 계산하는 것입니다. 그렇다면, 목적 함수 $J$와 $y_0$ 또는 $ ext{alpha}$를 어떻게 연결해야 할까요? 한쪽은 $x_1, y_1, ext{dot } y_1$이 수치적으로 구해진 것을 전제로 하는 함수 $J(x_1, y_1, ext{dot } y_1)$이고, 다른 한쪽은 미분대수 방정식
egin{align*} ext{dot } y &= f(t, x, y; ext{alpha}) \ g(t, x, y; ext{alpha}) &= 0 ext{

}가 정의되어 있으며, 보통 전용 솔버에 초기값 $y(t_0) = y_0$를 주어 수치적으로 $x_1 = x(t_1), y_1 = y(t_1), ext{dot } y_1 = ext{dot } y(t_1)$을 계산합니다. 함수 $J$의 계산과 미분대수 방정식의 계산을 통해 $J$와 $y_0$ 또는 $ ext{alpha}$가 관련되어 있다는 것은 분명하지만, 그 기울기 계산 방법은 자명하지 않습니다.

함수의 기울기와 미분대수 방정식의 기울기를 연결하기 위해서는 연쇄 법칙(chain rule)을 사용합니다. 미분대수 방정식을 푸는 과정을 하나의 함수로 간주하면, 그 해 $x_1, y_1, ext{dot } y_1$은 초기값과 파라미터, 시작점, 끝점의 함수로 나타낼 수 있습니다.
egin{align*} x_1 &= ext{Phi}{x}(t_0, y_0, t_1, ext{alpha}) \ y_1 &= ext{Phi}{y}(t_0, y_0, t_1, ext{alpha}) \ ext{dot } y_1 &= ext{Phi}_{yp}(t_0, y_0, t_1, ext{alpha}) ext{. 이것과 연쇄 법칙을 사용하면, } J ext{의 기울기를 다음과 같이 쓸 수 있습니다.}

이하의 각 항은 $\partial J/\partial x_1$과 같은 $J$의 편미분과, $\partial \Phi_x/\partial \alpha$와 같은 미분대수방정식을 푸는 과정 자체를 나타내는 함수의 야코비 행렬의 행렬곱으로 구성되어 있습니다. $\partial J/\partial x_1$과 같은 벡터와 야코비 행렬의 곱으로 표현되는 형태를 Vector-Jacobian product (VJP)라고 부릅니다. 본 기사의 목표는 이 VJP를 계산하는 방법을 유도하는 것입니다.

경사(Gradient) 계산의 어려움

$J$의 편미분은 $J$의 형태만 알면 계산할 수 있지만, 얻어진 $J$의 편미분에 대해 VJP를 계산하는 것은 매우 어렵습니다. 미분대수방정식 (DAE)의 초기값 문제 자체를 미분해야 하기 때문입니다.

VJP를 계산하는 순수한 방법으로는 수치 미분(numerical differentiation)과 미분대수방정식 솔버의 계산 과정 자동미분(automatic differentiation)이 있습니다. 둘 다 계산량이나 메모리 사용 면에서 문제가 있습니다. 수치 미분에서는 초기값 또는 매개변수를 하나씩 아주 작은 양만큼 이동시키면서 DAE를 풀고, 그 변화율을 계산합니다. 하지만 매개변수의 수만큼 DAE를 풀어야 하므로, 매개변수가 많으면 계산량이 엄청나게 커집니다.

미분대수방정식 솔버의 계산 과정 자체를 자동미분하는 방법은, 계산 그래프가 거대해지면서 메모리가 압박될 우려가 있습니다. 솔버가 DAE를 수치적으로 풀 때는 시간 $t$를 $\Delta t$씩 변화시키면서 행렬 연산이나 비선형 함수의 근을 구하여 $x(t+\Delta t)$, $y(t + \Delta t)$, $\dot y(t + \Delta t)$를 계산해 나갑니다. 이 구조는 심층 신경망 (deep neural network)과 거의 같습니다. 따라서 $t_1$부터 $t_0$ 방향으로 각 단계에 연쇄 법칙(chain rule)을 순차적으로 적용하면, VJP도 자동미분으로 계산할 수 있어야 합니다. 하지만 수백에서 수천 개의 시간 스텝을 거친 해의 자동미분을 하는 것은 같은 층 수의 심층 신경망 출력의 자동미분에 해당하며, 계산 중에 거대한 계산 그래프가 구축됩니다.

해결책

이러한 순수한 방법들과 달리, Adjoint 방식은 Adjoint 방정식이라고 불리는 미분대수방정식을 $J$의 편미분을 초기값으로 하여 시간 역방향으로 단 한 번 푸는 것만으로 VJP를 직접 계산합니다. 미분대수방정식의 계산은 원래의 DAE와 Adjoint 방정식, 두 번만 필요하며, 원래의 DAE를 풀 때 계산 그래프를 구축할 필요도 없습니다. 본 기사에서는 Adjoint 방정식을 유도하고, 초기값을 설정하는 방법과 Adjoint 방정식의 해로부터 경사를 계산하는 방법을 제시합니다.

한편, VJP를 계산할 때 $J$를 $x_1$, $y_1$, $\dot y_1$의 선형 함수라고 가정해도 일반성을 잃지 않습니다. $J$의 편미분을 $w_{x,1}^t = \partial J/\partial x_1$, $w_{y,1}^t = \partial J/\partial y_1$, $w_{yp,1}^t = \partial J/\partial \dot y_1$로 두면, VJP는 벡터 $w_{x,1}$, $w_{y,1}$, $w_{yp,1}$과 야코비 행렬의 곱들의 합이 되며, 이는 다음 $J$의 VJP와 일치하기 때문입니다.

J = w_{x,1}^t x_1 + w_{y,1}^t y_1 + w_{yp,1}^t \dot y_1\n여기서, v^t는 행렬・벡터의 전치(transpose)를 나타내며, $w_{x,1}, w_{y,1}, w_{yp,1}$은 상수 벡터입니다. 이후로는 이 선형화된 J를 사용하여 Adjoint 방정식을 유도해 나갑니다.

VJP 계산식의 도출

목적 함수와 라그랑주 범함수(Lagrangian Functional)

VJP 계산식과 Adjoint 방정식을 도출하기 위해 반양적 미분대수 방정식(semi-analytic differential algebraic equation)을 제약 조건으로 하는 다음 함수의 기울기(gradient)를 고려합니다.

\begin{align*} \nJ &= w_{x,1}^t x_1 + w_{y,1}^t y_1 + w_{yp,1}^t \dot y_1 \text{s.t.} &\quad y_1 = y(t_1) \dot y_1 &= f(t_1, x_1, y_1; \alpha) \g(t_1, x_1, y_1; \alpha) &= 0 \dot y &= f(t, x, y; \alpha) \g(t, x, y; \alpha) &= 0 \y(t_0) &= y_0 \end{align*}
함수 J의 기울기를 구하기 위해 각 제약 조건에 미정 계수(Lagrange multiplier)를 곱하여 목적 함수에 더한 라그랑주 범함수 L을 고려합니다.

\begin{align*} \nL&[t_0, t_1, y_0, x(t), y(t), x_1, y_1, \dot y_1, \alpha, \mu_{y_0}, \mu_{y_1}, \mu_{f_1}, \mu_{g_1}, \lambda_f(t), \lambda_g(t)] \&= w_{x,1}^t x_1 + w_{y,1}^t y_1 + w_{yp,1}^t \dot y_1 \&\quad + \mu_{y_1}^t \left( y_1 - y(t_1) \right)\n&+ \mu_{f_1}^t \left( \dot y_1 - f(t_1, x_1, y_1; \alpha) \right)\n&+ \mu_{g_1}^t g(t_1, x_1, y_1; \alpha) \&\quad + \mu_{y_0}^t \left( y_0 - y(t_0) \right)\n&+ \int_{t_0}^{t_1} \lambda_f^t \left( \frac{d y}{d t} - f(t, x, y; \alpha) \right) dt\n&+ \int_{t_0}^{t_1} \lambda_g^t g(t, x, y; \alpha) dt\n\end{align*}
제약 조건이 모두 충족된다면 괄호 안의 내용과 $g$는 모두 0이므로 L = J가 되며, 미정 계수 $\mu_{y_1}, \mu_{f_1}, \mu_{g_1}, \mu_{y_0}, \lambda_f(t), \lambda_g(t)$를 어떻게 선택하든 L의 변화는 J의 변화와 같습니다. 각 계수는 $\lambda_f(t)$와 $\lambda_g(t)$가 구간 $[t_0, t_1]$의 각 시점에서의 미분 방정식과 구속식에, $\mu_{f_1}$와 $\mu_{g_1}$이 시간 $t_1$에서 $\dot y_1, x_1$을 $y_1$에 연결하는 식에, $\mu_{y_1}$와 $\mu_{y_0}$가 변수 $y_1, y_0$와 궤적 $y(t)$의 양 끝과의 연결에 대응합니다. 이 계수들의 자유도를 사용하여 J의 기울기를 읽기 쉬운 형태로 L을 정리해 나갑니다.

전변분(Total Variation) 계산

L의 인자들을 모두 미소하게 움직였을 때의 전변분 $\delta L$을 계산하면 다음과 같습니다.

\begin{align*}
\delta L &= w_{x,1}^t \delta x_1 + w_{y,1}^t \delta y_1 + w_{yp,1}^t \delta \dot y_1 \
&\quad + \delta \mu_{y_1}^t \left( y_1 - y(t_1) \right)

• \mu_{y_1}^t \left( \delta y_1 - \delta y(t_1) - \frac{d y}{d t}(t_1) \delta t_1 \right) \
  &\quad + \delta \mu_{f_1}^t \left( \dot y_1 - f \right)
• \mu_{f_1}^t \left(
  \delta \dot y_1

• \frac{\partial f}{\partial t} \delta t_1
• \frac{\partial f}{\partial x} \delta x_1
• \frac{\partial f}{\partial y} \delta y_1
• \frac{\partial f}{\partial \alpha} \delta \alpha
  \right) \
  &\quad + \delta \mu_{g_1}^t g

• \mu_{g_1}^t \left(
  \frac{\partial g}{\partial t} \delta t_1
• \frac{\partial g}{\partial x} \delta x_1
• \frac{\partial g}{\partial y} \delta y_1
• \frac{\partial g}{\partial \alpha} \delta \alpha
  \right) \
  &\quad + \delta \mu_{y_0}^t \left( y_0 - y(t_0) \right)
• \mu_{y_0}^t \left( \delta y_0 - \delta y(t_0) - \frac{d y}{d t}(t_0)\delta t_0 \right) \
  &\quad + \int_{t_0}^{t_1} \delta \lambda_f^t \left( \frac{d y}{dt} - f \right) dt
• \int_{t_0}^{t_1} \lambda_f^t \left(
  \frac{d}{dt} \delta y

• \frac{\partial f}{\partial x} \delta x
• \frac{\partial f}{\partial y} \delta y
• \frac{\partial f}{\partial \alpha} \delta\alpha
  \right) dt \
  &\quad + \lambda_f^t(t_1) \left( \frac{dy}{dt}(t_1) - f(t_1, x(t_1), y(t_1)) \right) \delta t_1
• \lambda_f^t(t_0) \left( \frac{dy}{dt}(t_0) - f(t_0, x(t_0), y(t_0)) \right) \delta t_0 \
  &\quad + \int_{t_0}^{t_1} \delta \lambda_g^t g dt

• \int_{t_0}^{t_1} \lambda_g^t \left(
  \frac{\partial g}{\partial x} \delta x
• \frac{\partial g}{\partial y} \delta y
• \frac{\partial g}{\partial \alpha} \delta \alpha
  \right) dt \
  &\quad + \lambda_g^t(t_1) g(t_1, x(t_1), y(t_1); \alpha) \delta t_1

• \lambda_g^t(t_0) g(t_0, x(t_0), y(t_0); \alpha) \delta t_0 \
  &= w_{x,1}^t \delta x_1 + w_{y,1}^t \delta y_1 + w_{yp,1}^t \delta \dot y_1

• \mu_{y_1}^t \left( \delta y_1 - \delta y(t_1) - \dot y_1 \delta t_1 \right) \
  &\quad + \mu_{f_1}^t \left(
  \delta \dot y_1

첫 번째 등호의 우변에서는, 적분 구간의 끝점 $t_0, t_1$도 움직이고 있기 때문에, 적분 상단/하단의 이동으로부터 $\delta t_1$, $\delta t_0$에 비례하는 항이 나타납니다. 또한, $y(t_1)$은 함수 $y$와 시각 $t_1$ 양쪽을 통해 변화하므로, 그 변분은 $\delta y(t_1) + \frac{dy}{dt}(t_1)\delta t_1$이 됩니다. $y(t_0)$도 마찬가지입니다. 두 번째 등호에서는, 제약 조건이 충족된다는 것을 사용하여 식을 간단히 했습니다. $\delta\mu$, $\delta\lambda$를 포함하는 항은 제약의 괄호나 $g$가 0이기 때문에 사라지고, 적분 구간의 끝점 이동으로부터 발생한 항도 $\frac{dy}{dt} = f$와 $g = 0$으로부터 사라집니다. 게다가 $\frac{dy}{dt}(t_1) = \dot y_1$과, $\frac{dy}{dt}(t_0) = \dot y_0$라고 놓은 기호를 사용하여 다시 작성했습니다.

적분 안에 남아있는 $\frac{d}{dt}\delta y$는 $\delta y$와 형태가 다르기 때문에 계수를 모을 수 없습니다. 부분 적분을 통해 $\delta y$의 형태로 고칩니다.

\begin{aligned}
\int_{t_0}^{t_1} &\lambda_f^t \left( \frac{d}{dt} \delta y \right) dt \ &= \left[ \lambda_f^t \delta y(t) \right]{t_0}^{t_1} - \int{t_0}^{t_1} \frac{d\lambda_f^t}{dt} \delta y dt \ &= \lambda_f^t (t_1) \delta y(t_1) - \lambda_f^t (t_0) \delta y(t_0) - \int_{t_0}^{t_1} \frac{d \lambda_f^t}{dt} \delta y dt
\end{aligned}
부분 적분의 경계항으로서, 구간 양 끝에서의 변분 $\delta y(t_1), \delta y(t_0)$가 나타납니다. 이 결과를 $\delta L$에 대입하고, 변분별로 계수를 모으면 다음과 같습니다.

egin{aligned}
onumber \
onumber\
onumber
egin{aligned}
onumber \
onumber\
onumber
ext{앞의 3줄은 } ext{변분 } ext{}\ ext{$ackslash$x_1, $ackslash$y_1, $ackslash$dot y_1와, 궤적의 변분 } ext{$ackslash$y(t_1), $ackslash$y(t_0), $ackslash$y(t), $ackslash$x(t)의 항이며, 뒤의 3줄은 } ext{$ackslash$t_1, $ackslash$t_0, $ackslash$y_0, $ackslash$$ ext{alpha}의 항입니다.}

미정계수(Lagrange Multiplier) 선택과 Adjoint 방정식

우리가 알고 싶은 것은 $t_1, t_0, y_0, ext{$ackslash$alpha}를 움직였을 때 L의 변화이므로, $ ext{$ackslash$delta L}$에는 이 변분들만 남기고, 그 외의 변분들의 계수가 0이 되도록 미정계수를 선택합니다. 이렇게 함으로써 남아 있는 계수로부터 직접 VJP(Vector-Jacobian Product)를 읽어낼 수 있습니다. 미정계수를 다음과 같이 정하면, 앞의 3줄은 임의의 $ ext{$ackslash$delta x_1, $ackslash$delta y_1, $ackslash$delta ext{dot y_1}, $ackslash$y(t_1), $ackslash$y(t_0), $ackslash$y(t), $ackslash$x(t)에 대해 0이 됩니다.

egin{align*}
onumber \
onumber\
onumber\
onumber\
onumber\
onumber\
onumber
式(5), (6)은 ext{$ackslash$lambda_f(t), $ackslash$lambda_g(t)에 관한 미분 대수 방정식으로, Adjoint 방정식이라고 불립니다. 식(6)은 $ackslash$lambda_g를 결정하는 대수 방정식이며, 지수 1의 가정으로부터 $rac{\partial g}{\partial x}$가 정칙(regular)하므로 $ackslash$lambda_g에 대해 풀 수 있습니다.

식(1)부터 (4)는 Adjoint 방정식의 초기값을 결정하는 방정식으로, 값이 시점 $t_1$에서 주어집니다. 식(1)로부터 $ackslash$mu_{f_1}이 결정되고, 식(2)는 식(6)과 마찬가지로 $rac{\partial g}{\partial x}$가 정칙하므로 $ackslash$mu_{g_1}에 대해 풀 수 있습니다. 이를 식(3)에 대입하여 $ackslash$mu_{y_1}을 구하고, 식(4)에서 $ackslash$lambda_f($t_1$)로 합니다. 따라서 Adjoint 방정식은 $t_1$부터 $t_0$ 방향으로 시간을 거꾸로 풀어갑니다. 마지막 식(7)은 역방향으로 풀어서 얻은 $ackslash$lambda_f($t_0$)로부터 $ackslash$mu_{y_0}이 결정됨을 보여줍니다.

VJP 계산식

미정 계수(Lagrange multiplier)를 이렇게 선택하면, $ackslash$delta L에는 $ackslash$delta t_1, $ackslash$delta t_0, $ackslash$delta y_0, $ackslash$delta $ackslash$alpha 항만 남게 되며, 그 계수가 우리가 구하고자 하는 VJP가 됩니다.