自動微分を実装して理解する（前編）

近年、機械学習への応用により自動微分の技術が再び脚光を浴び始めています¹。例えばDeepLearningを微分な可能な関数の組み合わせ論へと発展させた微分可能プログラミング²では、プログラミングにより実装した関数の導関数を求める操作が基本的な役割を果たしています。この記事では自動微分に対する理解をより深めるために、自動微分のアルゴリズムをゼロから実装していきたいと思います。

自動微分

自動微分はプログラムによって定義された数学的な関数が与えられたときに、その導関数をアルゴリズムによって機械的に求める手法です。

例えば

f(a, b) = ab + \sin b

という関数が与えられたときにその導関数を計算することを考えます。まずこの関数の計算グラフは以下のようになっています。

\begin{matrix}
f &=& w + u \\
w &=& ab \\
u &=& \sin b \\
\end{matrix}

この計算グラフに沿って導関数を下から

\begin{matrix}
u' &=& \cos b \cdot b' \\
w' &=& a'b + ab' \\
f' &=& w' + u' \\
\end{matrix}

と計算する方法はフォワードモードの自動微分と呼ばれています。例えば変数$a$に関する偏導関数を求めたい場合は

\begin{matrix}
a' &=& \frac{\partial a}{\partial a} &=& 1 \\
b' &=& \frac{\partial b}{\partial a} &=& 0 \\
w' &=& \cos b \cdot b' &=& 0\\
u' &=& a'b + ab' &=& b \\
f' &=& w' + u' &=& b \\
\end{matrix}

と計算することができ、導関数は

f' = \frac{\partial f}{\partial a} = b

となることがわかります。

逆に計算グラフに沿って導関数を上から

\begin{matrix}
\frac{\partial f}{\partial w} &=& 1 \\
\frac{\partial f}{\partial u} &=& 1 \\
\frac{\partial f}{\partial a} &=& \frac{\partial f}{\partial w}\frac{\partial w}{\partial a} &=& b \\
\frac{\partial f}{\partial b} &=& \frac{\partial f}{\partial w}\frac{\partial w}{\partial b} + \frac{\partial f}{\partial u}\frac{\partial u}{\partial b} &=&  a + \cos b \\
\end{matrix}

と計算する方法はリバースモードの自動微分と呼ばれています。

リバースモードの自動微分では変数$a, b$についての偏導関数が既に求まっていることがわかります。一般に関数の出力の次元に比べて入力の次元が大きい時はリバースモードが効率よく、反対に出力の次元に比べて入力の次元が小さい時はフォワードモードが効率が良いことが知られています。

この記事では実装に焦点を当てるため自動微分についての詳しい解説は行いませんでしたが興味がある人には以下の文献が参考になると思います。

• 自動微分 - Wikipedia
• 自動微分の基礎 | PwCあらた有限責任監査法人
• 自動微分で遊ぼう
• ゼロから作るDeep Learning ❸ ―フレームワーク編

フォワードモード

それではフォワードモードの自動微分を実装してみましょう。

ソースコード変換

前述したフォワードモードの説明は数式によるものでしたが、我々が今から考えたいのは関数がプログラムによって与えらたときに導関数を求める方法です。例えば以下のようなプログラムが与えられたとしましょう³。

double f(double x1 , double x2) {
  double w3 , w4 , w5 ;
  w3 = x1 * x2 ;
  w4 = sin(x1);
  w5 = w3 + w4 ;
  return w5 ;
}

これを自動微分するにはソースコードをパースしてその導関数となる以下のようなソースコードを出力することが考えられます。

double df(double x1 ,double x2 ,double dx1, double dx2) {
  double dw3, dw4, dw5;
  dw3 = dx1 * x2 + x1 * dx2;
  dw4 = cos(x1) * dx1;
  dw5 = dw3 + dw4;
  return dw5;
}

このソースコード変換による自動微分を実装してみましょう。とは言っても今からレキサーとパーサーを用意するのは骨が折れますから、今回は数式を記述する独自言語をHaskellのEDSLとして用意し、それを自動微分してみることにしましょう。

data Expr
   = X              -- 変数 x
   | Lit Float      -- 定数
   | Neg Expr       -- -a
   | Add Expr Expr  -- a + b
   | Mul Expr Expr  -- a * b
   | Sin Expr       -- sin a
   | Cos Expr       -- cos a
   deriving Show

このExprは1変数の数式を表現できるような非常に簡単なEDSLです。

例えば

f(x) = x^2 + \sin x

という関数は

f :: Expr
f = X `Mul` X `Add` Sin X

のように実装する事ができます。
そして導関数を求める関数は以下のように再帰的に実装することができます。

grad :: Expr -> Expr
grad X           = Lit 1
grad (Lit _)     = Lit 0
grad (Neg e)     = Neg (grad e)
grad (Add e1 e2) = grad e1 `Add` grad e2
grad (Mul e1 e2) = (grad e1 `Mul` e2) `Add` (e1 `Mul` grad e2)
grad (Sin e)     = Cos e `Mul` grad e
grad (Cos e)     = Neg (Sin e) `Mul` grad e

足し算には微分の線形性を、掛け算にはライプニッツ則を、三角関数の微分には合成関数の微分法を用いているのがわかります。ghciで実行してみましょう。

> f
Add (Mul X X) (Sin X)
> grad f
Add (Add (Mul (Lit 1.0) X) (Mul X (Lit 1.0))) (Mul (Cos X) (Lit 1.0))

grad fをよく読むと

1 \times x + x \times 1 + \cos x \times 1

となっており

f'(x) = 2x + \cos x

と期待通りの答えになっていることがわかります。

EDSLを用いたソースコード変換によって簡単に自動微分を実装できました。Exprを評価するeval :: Expr -> Double -> Doubleを作ったり、Exprを多変数に拡張したり、できることはまだありそうですがここでは先を急ぐことにして、これらの実装は読者への課題としておきましょう

演算子のオーバーロード

EDSLの変換を用いた自動微分の実装は簡単でしたが、実装した数式の見た目が煩雑になってしまいました。できればAddやMulではなく標準で用意されている+, *等の演算子を使って記述したいですよね。HaskellではNum等の数値を扱う型クラスのインスタンスにすればこれらの演算子を使うことができるので、あとはインスタンスを定義する適切なデータ構造を見つけてやれば良さそうです。実は二重数という概念を使えば演算子のオーバーロードを使った自動微分を簡単に実装することができます。実装の前に二重数の数学的な性質を少し見ていきましょう。

複素数が二乗すると-1になるような元$i$を実数に加えてできる数だったように、二重数は二乗すると0になるような元$\epsilon$（ただし$\epsilon \neq 0$）を実数に加えてできる数です⁴。二重数全体の集合$D$に含まれる元は以下のよう形に書くことができます。

D = \{a + b\epsilon\ |\ a, b \in {\mathbb R}\}

今、実数を実数に写す微分可能な関数$f$が与えられたときに、以下のような二重数を二重数に写す関数$f^*$を考えることができます。

f^*(a + b\epsilon) = f(a) + f'(a)b\epsilon

例えば$x+\epsilon$を$f^*$に与えると

f^*(x + \epsilon) = f(x) + f'(x)\epsilon

のように実部に関数、二重数の部分に導関数が対応することがわかります。イメージとしては微小変分$\epsilon$を考えて$x$周りでのテイラー展開をして二次以上の項は無視していると考えるとわかりやすいかと思います。この$f$を$f^*$に対応させる写像を$\phi$と書くことにしましょう。すなわち

\phi(f) = f^*

と表します。

実はこの$\phi$は元の関数の加減乗除の演算と合成関数の操作を保つことがわかります。例えば関数の足し算を考えてみると

\phi(f + g) = (f + g) + \left(f' + g'\right)\epsilon

であり、一方で

\begin{matrix}
\phi(f) &=& f + f'\epsilon \\
\phi(g) &=& g + g'\epsilon \\
\phi(f) + \phi(g) &=& (f + g) + \left(f' + g'\right)\epsilon
\end{matrix}

となるので

\phi(f + g) = \phi(f) + \phi(g)

となり、足し算してから写しても写してから足し算しても結果が変わらないことが分かります。左辺の足し算は実数上の足し算ですが右辺の足し算は二重数上の足し算であることに注意してください。

同様に関数の掛け算を考えてみましょう。

\phi(fg) = fg + \left(f'g + fg'\right)\epsilon

であり、一方で

\begin{matrix}
\phi(f) &=& f + f'\epsilon \\
\phi(g) &=& g + g'\epsilon \\
\phi(f)\phi(g) &=& fg + f'g\epsilon + fg'\epsilon + f'g'\epsilon^2 \\
&=& fg + \left(f'g + fg'\right)\epsilon
\end{matrix}

となり、

\phi(fg) = \phi(f)\phi(g)

となることがわかります。式の途中では$\epsilon^2=0$であることを用いました。（割り算についても同様の関係が成り立つかを考えるのは面白い問題なのでぜひ挑戦してみてください。）

最後に関数の合成を考えましょう。

\phi(f \circ g)(a+b\epsilon) = f(g(a)) + f'(g(a))g'(a)b\epsilon

となり、一方で

\begin{matrix}
\phi(f)(a + b\epsilon) &=& f(a) + f'(a)b\epsilon \\
\phi(g)(a + b\epsilon) &=& g(a) + g'(a)b\epsilon \\
(\phi(f)\circ \phi(g))(a + b\epsilon) &=& \phi(f)(\phi(g)(a + b\epsilon)) \\
&=& \phi(f)(g(a) + bg'(a)\epsilon) \\
&=& f(g(a)) + f'(g(a))g'(a)b\epsilon
\end{matrix}

となるので

\phi(f\circ g) = \phi(f)\circ \phi(g)

となることがわかります。途中式が少々煩雑になりましたが、わかりやすさのため要素を明示する形で書きました。

以上よりわかるのは$\phi$により二重数の部分が導関数に対応するという関係は、写した先で四則演算や合成関数を作っても変わらないということです⁵。この性質のおかげで二重数の関数の合成を考えればもとの実数の関数の導関数が手に入ることが保証されています。

それでは実際に二重数を使って自動微分を実装していきましょう。

data D a = D a a deriving Show

real, grad :: D a -> a
real (D x _) = x
grad (D _ y) = y

Dは二重数を表す型で値コンストラクタは実部と2重数部分の2つの値を受け取ります。

instance Num a => Num (D a) where
  (D x x') + (D y y') = D (x + y) (x' + y')
  (D x x') * (D y y') = D (x * y) (x' * y + x * y')
  negate (D x x')     = D (negate x) (negate x')
  abs (D x x')        = D (abs x) (x' * (signum x))
  signum (D x x')     = D (signum x) 0
  fromInteger n       = D (fromInteger n) 0

instance Fractional a => Fractional (D a) where
  recip (D x x') = D (recip x) (-1 * x' * (recip (x * x)))
  fromRational x = D (fromRational x) 0

型クラスNumとFractionalのインスタンスを作成しました。二重数部分には導関数が対応すると考えるとすんなり理解できるかと思います。

instance Floating a => Floating (D a) where
  pi             = D pi 0
  exp   (D x x') = D (exp x)   (x' * exp x)
  log   (D x x') = D (log x)   (x' / x)
  sin   (D x x') = D (sin x)   (x' * cos x)
  cos   (D x x') = D (cos x)   (- x' * sin x)
  asin  (D x x') = D (asin x)  (x' / (sqrt(1 - x ** 2)))
  acos  (D x x') = D (acos x)  (- x' / (sqrt(1 - x ** 2)))
  atan  (D x x') = D (atan x)  (x' / (1 + x ** 2))
  sinh  (D x x') = D (sinh x)  (x' * cosh x)
  cosh  (D x x') = D (cosh x)  (x' * sinh x)
  asinh (D x x') = D (asinh x) (x' / (sqrt(1 + x ** 2)))
  acosh (D x x') = D (acosh x) (x' / (sqrt(x ** 2 - 1)))
  atanh (D x x') = D (atanh x) (x' / (1 - x ** 2))

型クラスFloatingには様々な数学的な関数が定義されています。二重数のインスタンスは前述した$\phi$により二重数に拡張された関数を定義していると考えると良いでしょう。

ソースコード変換と同じように関数

f(x) = x^2 + \sin x

を実装してみましょう。

f :: Floating a => a -> a
f x = x^2 + sin x

元の数式とほとんど同じ見た目に実装することができました。多相関数を使っているので直接的に二重数のデータ構造Dに依存していないのも嬉しいポイントです。EDSLを用いた場合のように導関数の式の形を直接確認することはできませんが、値を直接評価してうまく動いているか確かめてみましょう。

> grad $ f (D 0 1)
1.0

> grad $ f (D pi 1)
5.283185307179586

$f$の導関数は

f'(x) = 2x + \cos x

なので

\begin{matrix}
f'(0) &=& \cos 0 &=& 1 \\
f'(\pi) &=& 2\pi + \cos \pi \\
&=& 6.28\dots - 1 &=& 5.28\dots \\
\end{matrix}

となり微分係数が正しく求められていることを確認できました

この二重数による実装は多変数関数への拡張が容易です。冒頭の章で例に上げた

f(a, b) = ab + \sin b

という数式は

f :: Floating a => a -> a -> a
f a b = a * b + sin b

のように実装することができます。$a, b$それぞれの原点での微分係数は

> grad $ f (D 0 1) (D 0 0)
0.0

> grad $ f (D 0 0) (D 0 1)
1.0

となり理論的な値と一致していることが確かめられます。二重数を用いて微分する時は微分したい変数の二重数部分を1にしてそれ以外の変数の二重数部分を0にするのがポイントです。

二重数のアイデアを拡張して3乗したら0になる数を実数に加えると二階微分までを表現できる数を作ることもできます。さらに遅延評価を利用すれば無限回微分まで表現できるデータ構造を作ることも可能です⁶。また同様のアイデアを拡張して四元数関数の自動微分を考えることも可能です（『四元数ニューラルネットワークとGHR微積分』）。

以上、フォワードモードの自動微分の実装と解説でした。後編ではリバースモードの自動微分について実際に実装して解説していきたいと思います。お楽しみに

※2018/12/21
後編を公開しました

1. [1502.05767] Automatic differentiation in machine learning: a survey ↩

2. 教師あり学習の学習モデルを微分可能プログラミングと純粋関数型言語で記述する - Qiita ↩

3. 参考: Havard Berland, "Automatic differentiation" ↩

4. 環を割ったりClifford代数だったりの話はここではしません ↩

5. 逆に二重数に対してこのような準同型性を満たすような対応はライプニッツ則を満たすので微分とみなせるということも知られています 微分環と双対数 ↩

6. Conal Elliott » Beautiful differentiation ↩