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

この記事は[FOLIO Advent Calendar 2018](https://adventar.org/calendars/3424)、21日目の記事です。昨日はhajipionによる[大規模デザインカンファレンス「デザインシップ」の着想から開会まで、超ダッシュで振り返る10ヶ月！ #Designship2018](https://hajipion.com/2792.html)でした。

この記事は近年話題に上がることが多くなってきた**自動微分**を実際に実装することでより理解を深めようというものです。前回の[自動微分を実装して理解する（前編）](https://qiita.com/lotz/items/39c52f08cc9b5d8439ca)では自動微分の定義とフォワードモードの実装を見てきました。今回は**リバースモード**の自動微分の実装にチャレンジしてみましょう。

## リバースモード
現実世界（特に機械学習）の関数は一般的に出力の次元より入力の次元のほうが圧倒的に多いので、フォワードモードよりリバースモードの方が効率面で優れている場面が多々あります。しかし一般的なリバースモードの実装方法では数式を組み立てた順番を計算グラフとして保持する必要があるため、その実装はフォワードモードの時よりも複雑になってしまいます。前回と同様にソースコード変換[^1]による実装や演算子のオーバーロードによる実装[^2]ももちろん可能ですが、ここでは最近話題になった[^3][^4]シンプルなリバースモードの実装を紹介したいと思います。

### 双方向の計算による実装
リバースモードを使った実装が複雑になるのは、数式を組み立てたあとに逆順に導関数を計算するという双方向の計算を実装する必要があるからです。そこで数式を組み立てるのと同時に逆向きに導関数を組み立てることでリバースモードの自動微分を実現することを考えます。

この方法では順方向に値を計算する関数`x -> y`と、`x`と`y`におけるこれから合成される関数の微分係数`dfdy`を受け取って`x`における関数全体の微分係数`dx`を返す関数`x -> dfdy -> dfdx`を考えます（`dfdy`は$\frac{\partial f}{\partial y}$、`dfdx`は$\frac{\partial f}{\partial x}$のつもりです）。つまり

```hs
(x -> y, x -> dfdy -> dfdx)
```

という型を考えます。これを使っていくつかの関数を実装してみましょう。

```hs
pow' :: Int -> (Double -> Double, Double -> Double -> Double)
pow' n = (\x -> x^n, \x dfdy -> dfdy * n * x^(n-1))

sin' :: (Double -> Double, Double -> Double -> Double)
sin' = (sin, \x dfdy -> dfdy * cos x)

exp' :: (Double -> Double, Double -> Double -> Double)
exp' = (exp, \x dfdy -> dfdy * exp x)
```

実装したのはべき乗と三角関数$\sin$と指数関数です。例えばこれらの関数をうまく合成して、

```math
f = e^{\sin x^2}
```

という数式を

```math
\begin{matrix}
f &=& e^w \\
w &=& \sin u \\
u &=& x^2 \\
\end{matrix}
```

という風に組み立てることができたとします。すると導関数の計算は、まず最初に`exp`の導関数部分が計算されるでしょう。引数には`x`と、`dfdy`として

```math
\frac{\partial f}{\partial f} = 1
```

が代入され、計算結果は`exp x`となります。
次に`sin`の導関数部分が計算されます。引数には`x`と、`dfdy`として先程計算した

```math
\frac{\partial f}{\partial f}\frac{\partial f}{\partial w} = e^x
```

が代入され`exp x * cos x`が計算されます。
最後に`pow`の導関数部分が計算されます。引数には`x`と、`dfdy`として先程計算した

```math
\left(\frac{\partial f}{\partial f}\frac{\partial f}{\partial w}\right)\frac{\partial w}{\partial u} = e^x\cos x
```

が代入され`exp x * cos x * n * x^(n-1)`が計算されます。こうして無事に導関数を計算することができました。この計算順序を見てみると、

```math
\left(\left(\frac{\partial f}{\partial f}\frac{\partial f}{\partial w}\right)\frac{\partial w}{\partial u}\right)\frac{\partial u}{\partial x}
```

という順番になっています。実はフォワードモードの自動微分は

```math
\frac{\partial f}{\partial f}\left(\frac{\partial f}{\partial w}\left(\frac{\partial w}{\partial u}\frac{\partial u}{\partial x}\right)\right)
```

という順番で計算する手法だったので、リバースモードは導関数の計算順序が逆になっているのがわかると思います。

それでは本当に計算をこのような順番となるようにうまく組み立てることができるのでしょうか？今考えている型をもう一度見てみましょう

```hs
(x -> y, x -> dfdy -> dfdx)
```

この型、どこかで見覚えがありませんか？そう、1つ目の要素をGetter、2つ目の要素をSetterだと思えばまさにLensの型ですね！ということで自動微分の型を van Laarhoven Lens として表してみましょう。

```hs
type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens sa sbt afb s = sbt s <$> afb (sa s)

type AD x dfdx y dfdy = Lens x dfdx y dfdy

type AD' x dfdx = AD x dfdx x dfdx
```

`Lens`と`lens`の定義は[microlens](http://hackage.haskell.org/package/microlens)ライブラリから拝借してきました。これを使えば先程の関数は

```hs
pow' :: Int -> AD' Double Double
pow' n = lens (\x -> x^n) (\x dfdy -> dfdy * (fromIntegral n) * x^(n-1))

sin' :: AD' Double Double
sin' = lens sin (\x dfdy -> dfdy * cos x)

exp' :: AD' Double Double
exp' = lens exp (\x dfdy -> dfdy * exp x)
```

のように実装することができます。Lensになったおかげで関数の合成には通常の合成関数の演算子`.`が使えるので、

```math
f = e^{\sin x^2}
```

という数式は

```hs
f :: AD' Double Double
f = pow' 2 . sin' . exp'
```

のように実装することができます。
この`AD`から関数の値を取り出す関数にはLensとして値を取り出す関数[`view`](http://hackage.haskell.org/package/microlens-0.4.10/docs/Lens-Micro-Extras.html#v:view)を使うことができます。

```hs
> view f 0
1.0

> view f (sqrt $ pi / 2)
2.718281828459045
```

これは

```math
e^{\sin 0} = e^0 = 1
```

そして

```math
e^{\sin \left(\sqrt{\frac{\pi}{2}}\right)^2} = e^{\sin\frac{\pi}{2}} = e
```

なので合っていますね。

一方、導関数の値を取り出す関数は`dfdy`として$1$をセットすればいいので

```hs
grad :: Num b => AD' a b -> a -> b
grad l = set l 1
```

と実装することができます。

```hs
> grad f (sqrt $ pi)
-3.544907701811034
```

これは

```math
\frac{\partial f}{\partial x} = 2x\cos\left(x^2\right)e^{\sin x^2}
```

であり

```math
\begin{matrix}
2{\sqrt \pi}\cos\left(\left({\sqrt \pi}\right)^2\right)e^{\sin \left({\sqrt \pi}\right)^2} &=& 2{\sqrt \pi}\cos\pi e^{\sin \pi} \\
 &=& 2{\sqrt \pi}\times(-1) \times e^0 \\
 &=& -2{\sqrt \pi} \\
 &=& -3.544\dots
\end{matrix}
```

となりちゃんと計算できていることが分かりました :clap: 

このように`AD`をLensとして表現することで関数合成の演算子`.`を使って数式を構築し、リバースモードの自動微分が実装できることが分かりました。ただこの方法は値ではなく関数を拡張しているため、自動微分したい関数をポイントフリースタイルで書かないといけないのが欠点です。文献[4][^4]にはこの方法を利用してニューラルネットワークを実装した例も紹介されています。

### ADはLensなのか

ところでLensには満たすべき法則として以下のようなLens則[^5]がありました。

```
> view l (set l v s) == v

> set l (view l s) s == s

> set l v' (set l v s) == set l v' s
```

しかし先程実装した`f`で実験してみると

```hs
> view f (set f 5 1)
1.0168416238956406

> set f (view f 1) 1
5.815127313951733
```

そして

```hs
> set f 1 (set f 1 1)
5.016872534352483

> set f 1 1
2.506761534986894
```

となり`AD`はLens則を満たしていません。

```hs
(x -> y, x -> dfdy -> dfdx)
```

元の型を思い出してみるとLensのGetterに対応する部分は関数、Setterに対応する部分は導関数であり、何か値を取り出したり詰めたりするための参照のような性質は持っていないので、Lensとしての性質を持たないのは当たり前のような気もします。合成関数の演算子`.`を使うためにLensという概念を形式的に導入しましたが、`AD`にLensが持つような抽象的な構造や性質は期待しないほうが良さそうです。

### 追記: 2021/12/15

Lens則は必ずしもLensが満たすべき性質とは思われてないようです。["Categories of Optics"](https://arxiv.org/abs/1809.00738)によるとLens則を満たすLensは　一般的なLensと区別してlawfulなものと呼ばれています。

この自動微分をLensとして構築する考え方を応用した話を以下の記事に書いたので気になる人は参考にしてください。

https://zenn.dev/lotz/articles/14458f024674e14f4134

## 探求: 使い勝手の良いリバースモードの実装
（ここからは私独自の試みです）
双方向の計算によるリバースモードの実装は数式の実装をポイントフリースタイルで書かなければいけないため数式の記述が煩雑になってしまいました。どうにか演算子のオーバーロードの時のように元のプログラミング言語の演算子を使って数式を組み立てるようにはできないでしょうか。ここでは今までの知識を総動員して

1. EDSLから変換する先として`AD`を利用する
2. 1で実装したEDSLに対して演算子のオーバーロードを使う

という方針で使い勝手の良いリバースモードの実装を検討していきたいと思います。

### EDSLから変換する先として`AD`を利用する
フォワードモードをソースコード変換で実装した時にEDSLをEDSLに変換するという手法を用いました。ここではEDSLをEDSLに変換するのではなく、値の組み合わせによって関数の組み合わせを表現することで、EDSLを`AD`に変換することを考えます。さっそくEDSLを定義してみましょう。

```hs
data Expr = X             -- 変数
          | Lit Double    -- 定数
          | Neg Expr      -- -x
          | Abs Expr      -- |x|
          | Sig Expr      -- sign x
          | Add Expr Expr -- x + y
          | Mul Expr Expr -- x * y
          deriving Show
```

後々で`Num`のインスタンスにすることを念頭にそれぞれの値コンストラクタを用意しています。このEDSLを`AD`に変換する前に`AD`の関数をいくつか用意しておきましょう。

```hs
pair' :: Lens' a b -> Lens' c d -> Lens' (a, c) (b, d)
pair' l1 l2 = lens (\(a, c) -> (view l1 a, view l2 c))
                   (\(a, c) (b, d) -> (set l1 b a, set l2 d c))

dup' :: Num a => Lens' a (a,a)
dup' = lens (\x -> (x, x)) (\x (dfdx, dfdy) -> dfdx + dfdy)

lit' :: Num a => a -> Lens' a a
lit' x = lens (\_ -> x) (\_ _ -> 0)

neg' :: Num a => Lens' a a
neg' = lens (\x -> -x) (\_ dfdy -> -dfdy)

abs' :: Num a => Lens' a a
abs' = lens (\x -> abs x) (\x dfdy -> dfdy * signum x)

sig' :: Num a => Lens' a a
sig' = lens (\x -> signum x) (\_ _ -> 0)

add' :: Num a => Lens' (a,a) a
add' = lens (\(x,y) -> x + y) (\_ dfdy -> (dfdy, dfdy))

mul' :: Num a => Lens' (a,a) a
mul' = lens (\(x,y) -> x * y) (\(x, y) dfdz -> (dfdz * y, x * dfdz))
```

`pair'`は実装より型を見ると何をしたいか分かりやすいと思います。`dup'`は入力を2つに分岐する関数で、このあとの実装を見ると必要性がわかると思います。あとは定数を表現する`lit'`と`Num`のメソッドに対応するLensたちです。これらのLensを使ってEDSLから`AD`への変換を実装してみましょう。

```hs
eval :: (Expr -> Expr) -> AD' Double Double
eval f = eval' (f X)
  where
    eval' :: Expr -> AD' Double Double
    eval' X           = id
    eval' (Lit x)     = lit' x
    eval' (Neg e)     = eval' e . neg'
    eval' (Abs e)     = eval' e . abs'
    eval' (Sig e)     = eval' e . sig'
    eval' (Add e1 e2) = dup' . pair' (eval' e1) (eval' e2) . add'
    eval' (Mul e1 e2) = dup' . pair' (eval' e1) (eval' e2) . mul'
```

`X`が基底部となるように`AD`への変換がうまく再帰的に書けているのがわかるかと思います。特に`Add`や`Mul`の変換が`dup'`と`pair'`を使って綺麗に配線できました。

### EDSLに対して演算子のオーバーロードを使う
それでは作った`Expr`を`Num`のインスタンスにしてみましょう。

```hs
instance Num Expr where
  e1 + e2       = e1 `Add` e2
  e1 * e2       = e1 `Mul` e2
  negate e      = Neg e
  abs e         = Abs e
  signum e      = Sig e
  fromInteger n = Lit (fromInteger n)
```

難しいことは無く、単純に用意したEDSLに変換しているだけですね。以上の準備で演算子のオーバーロードを使ったリバースモードの自動微分を実装することができました。実際に使ってみましょう

```hs
> f = eval $ \x -> 3 * x * x + x * 1

> :type f
f :: Functor f => (Double -> f Double) -> Double -> f Double

> f' = grad f

> f' 1
7.0

> f' 2
13.0

> f' 3
19.0
```

`f`として

```math
f = 3x^2 + x + 1
```

という関数を想定しています。`f`の実装にはHaskell標準の演算子が使われており、ポイントフリースタイルの制約は取り払われています。$f$の導関数は

```math
\frac{df}{dx} = 6x + 1
```

となり

```math
\begin{matrix}
\frac{df}{dx}(1) &=& 7 \\
\frac{df}{dx}(2) &=& 13 \\
\frac{df}{dx}(3) &=& 19
\end{matrix}
```

となるのでうまく実装できてそうですね :clap:

以上のように、演算子のオーバーロード・ソースコード変換・双方向の計算による実装を全て駆使して使い勝手の良いリバースモードの自動微分（1変数ですが:sweat_smile:）を実装することができました。このあとの課題として、型族を使って多変数関数に拡張したり、Expression Problemを解決しながらより多くの数値的な型クラスのインスタンスを作ったり、行列やベクトルに対応することなどが考えられますが、体力が尽きたので今回の実験はここまでにしておきます。

この記事を通して少しでも皆さんの自動微分への理解に貢献できれば幸いです。

### 追記: 2024/12/02
EDSLを使ったリバースモード自動微分について、Haskellのライブラリ [ad](https://hackage.haskell.org/package/ad), [backprop](https://hackage.haskell.org/package/backprop) で実際に使用されている より効率的な実装方法をだめぽ氏が解説してくれた記事があるのでそちらも参考にしてください！

https://zenn.dev/mod_poppo/articles/haskell-dsl-atomicmodifyioref

[^1]: [Reverse mode Automatic Differentiation - The H2 Wiki](http://h2.jaguarpaw.co.uk/posts/reverse-mode-automatic-differentiation/)
[^2]: Haskellの有名な自動微分ライブラリである[ad](http://hackage.haskell.org/package/ad)や[backprop](http://hackage.haskell.org/package/backprop)は演算子のオーバーロードを使ってリバースモードの自動微分を実装しています
[^3]: [[1804.00746] The simple essence of automatic differentiation](https://arxiv.org/abs/1804.00746)
[^4]: [Reverse Mode Differentiation is Kind of Like a Lens II - Hey There Buddo!](http://www.philipzucker.com/reverse-mode-differentiation-is-kind-of-like-a-lens-ii/)
[^5]: <https://artyom.me/lens-over-tea-2#lens-laws>