確率的勾配降下法をadで実装してみた

[最急降下法](https://www.wikiwand.com/ja/%E6%9C%80%E6%80%A5%E9%99%8D%E4%B8%8B%E6%B3%95)はある関数（主に誤差関数）を最小にするような入力を関数の勾配（一階導関数）の情報を用いて探索する手法です。[確率的勾配降下法](https://www.wikiwand.com/ja/%E7%A2%BA%E7%8E%87%E7%9A%84%E5%8B%BE%E9%85%8D%E9%99%8D%E4%B8%8B%E6%B3%95)はこれをオンライン学習に改良したもので、最近のDeepLearningの流れに乗ってよく話題になってるのを見かけます。
これらのアルゴリズムを利用するには一階導関数の計算をする必要があるのですが最適化したい関数は複雑になることが多く手計算には限界があります。そこでこれらのアルゴリズムを利用するライブラリの多くは同時に関数の自動微分をサポートしていることがほとんどです。
Haskellには[ad](http://hackage.haskell.org/package/ad)という強力な[自動微分](https://www.wikiwand.com/ja/%E8%87%AA%E5%8B%95%E5%BE%AE%E5%88%86)のためのライブラリがあるので、今回はadを利用してどのように確率的勾配降下法が実装できるのかを試してみました。

ad: Automatic Differentiation
------------------------------
adは自動微分のためのライブラリです。[Numeric.AD](http://hackage.haskell.org/package/ad-4.3.2.1/docs/Numeric-AD.html)を読むと大体どういう事ができるのかはわかると思います。幾つか例を挙げると

```hs
> import Numeric.AD

> sin 0
0.0

> diff sin 0
1.0

> (\[x, y] -> log (x + 2 * y)) [1, 2]
1.6094379124341003

> grad (\[x, y] -> log (x + 2 * y)) [1, 2]
[0.2,0.4]

> diff (\x -> if x < 0 then 0 else x) 2
1

> diff (\x -> if x < 0 then 0 else x) (-2)
0

> import Debug.SimpleReflect

> grad (\[x, y] -> log (x + 2 * y)) [x, y]
[0 + (0 + 1 * (0 + recip (x + 2 * y) * 1)),0 + (0 + 2 * (0 + 1 * (0 + recip (x + 2 * y) * 1)))]
```

この様に導関数を簡単に計算することができます。ifが入ったような式でも微分できるのが自動微分の特徴です。さらに[Debug.SimpleReflect](https://hackage.haskell.org/package/simple-reflect-0.3.2/docs/Debug-SimpleReflect.html)を使えば導関数がどのような形になっているかを確認することも出来ます。
自動微分には幾つか種類があってadではモードとして使い分けれるようになっていますが、基本的には`Numeric.AD`を利用していれば裏で最適なモードを選んで計算してくれるようになっています。モードについては一覧と簡単な説明が[GithubのREADME](https://github.com/ekmett/ad#modules)に載っているので気になる人は読んでみて下さい。
先ほどの例に出てきた`diff`は一変数関数の微分を表す関数です。型は

```hs
diff :: Num a => (forall s. AD s (Forward a) -> AD s (Forward a)) -> a -> a
```

のようになっていて`forall s. AD s (Forward a) -> AD s (Forward a)`のような関数を受け取って導関数`a -> a`を返す関数という風に読めます。`Forward`は自動微分のモードで`a`は`Double`や`Rational`など数値の型を表しています。`forall`で導入されている`s`は実は幽霊型になっていて、微分の際に無限小を間違って組み合わせられないように用意されています[^1]。

問題設定
--------
今回は最適化する問題として[ロジステック回帰分析](https://www.wikiwand.com/ja/%E3%83%AD%E3%82%B8%E3%82%B9%E3%83%86%E3%82%A3%E3%83%83%E3%82%AF%E5%9B%9E%E5%B8%B0)を選びました。ざっくり説明すると多次元の説明変数 $\bf x$ と $\\{0,1\\}$ の2値の値を取る目的変数 $t$ の組の集合が教師データとして与えられた時に $y$ が $1$ になる確率が

```math
{\rm P}\left(t = 1\right) = \frac{1}{1+\exp^{-\phi({\bf x})}}
```
となるようなモデルを仮定し、教師データに対して誤差関数

```math
L(x) = - \sum_{i}t_i\log y_i + (1-t_i)\log(1-y_i)
```

を最小にするような関数 $\phi $ を決定する問題です。ここで $({\bf x}\_i, t\_i)$ は $i$ 番目の教師データで $y\_i = \frac{1}{1+\exp^{-\phi({\bf x}\_i)}}$ です。
今回は説明変数 $x = (x_1, x_2)$ を2次元として、教師データは $\phi = x_1 + x_2 - 1$ として生成します。
Haskellの式で書いていきましょう。

```hs
type Input = [Double] -- 説明変数
type Output = Double -- 目的変数

type Var s = AD s (Kahn Double) -- 自動微分のための変数

predictor :: [Var s] -> Input -> Var s
predictor [w1, w2, w3] [x1, x2] = 1 / (1 + exp (-phi))
  where phi = w1 * auto x1 + w2 * auto x2 + w3
```

今回はadのモードとして[Kahn](http://hackage.haskell.org/package/ad-4.3.2.1/docs/Numeric-AD-Mode-Kahn.html)を選びました。[`auto`](http://hackage.haskell.org/package/ad-4.3.2.1/docs/Numeric-AD-Mode-Kahn.html#v:auto)はadの関数で、ここでは`Double -> AD s (Kahn Double)`として使っています。
これらを使って誤差関数は以下のように定義できます。

```hs
loss :: [Var s] -> (Input, Output) -> Var s
loss w (x, t) =
  let y = predictor w x
   in (-(auto t)) * log y + (auto t - 1) * log (1 - y)
```

教師データを生成する関数とパラメータの初期値を生成する関数も定義しておきましょう

```hs
generateTrainingData :: Int -> IO [(Input, Output)]
generateTrainingData n = do
  sequence . replicate n $ do
    withSystemRandom . asGenIO $ \gen -> do
      x1 <- uniformR (-1, 1) gen :: IO Double
      x2 <- uniformR (-1, 1) gen :: IO Double
      let y = if h x1 x2 > 0 then 1 else 0
      pure ([x1, x2], y)
      where
        h x y = x + y - 1

createInitialParams :: IO Params
createInitialParams = do
  withSystemRandom . asGenIO $ \gen -> do
    sequence . replicate 3 $ (uniformR (-1, 1) gen :: IO Double)
```

乱数生成には[mwc-random](http://hackage.haskell.org/package/mwc-random)を使っています。

最適化
------
いよいよ最適化アルゴリズムを実装してパラメータを推定していきましょう。今回は [勾配降下法の最適化アルゴリズムを概観する](http://postd.cc/optimizing-gradient-descent/) という記事で紹介されている順番に従って実装していきたいと思います。

### バッチ勾配降下法

```math
w_{t+1} = w_t - \alpha \frac{\partial L}{\partial w}
```

という更新式に従ってパラメータを更新していくのが勾配降下法です。今回であればパラメータは $\phi(x_1, x_2) = w_1x_1+w_2x_2+w_3$ の ${\bf w} = (w_1, w_2, w_3)$ に当たります。

Haskellの式で書くと

```hs
type Params = [Double]

gradientDescent :: Double -- 学習率
                -> (forall s. [Var s] -> Var s) -- 最適化する関数
                -> Params -> Params -- パラメータを更新する関数
gradientDescent alpha f ws = ws |-| (alpha *| grad f ws)
```

このような感じです。ここで`|-|`と`*|`は自分で定義した演算子でリスト同士やリストとスカラーの加減乗除を行います。

```hs
(|+|), (|-|) :: Num a => [a] -> [a] -> [a]
(|+|) = zipWith (+)
(|-|) = zipWith (-)
(|*|) = zipWith (*)
(|/|) = zipWith (/)
infixl 6 |+|, |-|
infixl 7 |*|, |/|

(*|) :: Num a => a -> [a] -> [a]
(*|) a xs = map (a*) xs
infixl 7 *|
```

バッチ勾配降下法は与えられた全ての教師データを用いて誤差関数を評価してパラメータを更新します。実装すると

```hs
main :: IO ()
main = do
  trainingData <- generateTrainingData 1000
  initialParams <- createInitialParams
  print initialParams
  let totalLoss w = sum $ map (\td -> loss w td) trainingData
      updatedParams = gradientDescent 0.01 totalLoss initialParams
  print updatedParams
```

のようになります。早速実行してみましょう。

```bash
$ runhaskell Main.hs
[0.7114283274355773,0.7118532073025301,0.7939099694680201]
[0.9673771585755228,1.0960957446062718,-4.758942054117518]
```

次は複数回パラメータを更新してみましょう。

```hs
main :: IO ()
main = do
  trainingData <- generateTrainingData 1000
  initialParams <- createInitialParams
  print initialParams
  let totalLoss w = sum $ map (\td -> loss w td) trainingData
      update = gradientDescent 0.01 totalLoss
      updatedParams = foldr (.) id (replicate 100 update) initialParams
  print updatedParams
```

式中の`update`の型は`Params -> Params`となっているので`replicate`で複製して`foldr (.) id`で合成するだけで複数回パラメータを更新する関数が出来上がります。

```bash
$ runhaskell Main.hs
[-0.4738897546795806,0.8157795481090948,0.48281066822261587]
[7.489960683643498,7.207263388659815,-7.547967942565535]
```

実は今回の問題だと解空間は鞍のような形をしていて方向は $[1,1,-1]$ でも絶対値が大きくなればなるほど最適な解となります。100回繰り返すことでより最適な解に向かっていますね。

### 確率的勾配降下法
バッチ勾配降下法が全てのデータを使って1回更新していたのに対し確率的勾配降下法は一つ一つのデータに対して一回ずつパラメータの更新を行います。式で書くと

```hs
main :: IO ()
main = do
  trainingData <- generateTrainingData 1000
  initialParams <- createInitialParams
  print initialParams

  let sgd td = gradientDescent 0.01 (\w -> loss w td)
      update = foldr (.) id . map sgd <$> shuffleM trainingData
  updatedParams <- ($ initialParams) . foldr (.) id <$> (sequence $ replicate 100 update)

  print updatedParams
```

このようになります。[`shuffleM`](http://hackage.haskell.org/package/random-shuffle-0.0.4/docs/System-Random-Shuffle.html#v:shuffleM)は[random-shuffle](http://hackage.haskell.org/package/random-shuffle)に含まれるリストの中身をランダムに入れ替える関数です。実行すると

```bash
$ runhaskell Main.hs
[0.22438797340555694,-0.1705774027060929,0.5316304718400742]
[7.153499539807709,7.1391607065086555,-7.330240069369628]
```

となりちゃんと最適化されていることがわかります。

### ミニバッチ勾配降下法
ミニバッチ勾配降下法はバッチ勾配降下法と確率的勾配降下法のあいのこのようになっており、1回の更新で複数個のデータを使ってパラメータを更新します。1回の更新で全てのデータを使うとバッチ勾配降下法であり1個のデータを使うと確率的勾配降下法となるのでミニバッチ勾配降下法はそれぞれの勾配降下法の一般化になっているとも考えられます。何はともあれ実装してみましょう。

```hs
miniBatch :: MonadIO m
          => Int -- バッチサイズ
          -> [(Input, Output)] -- 教師データ
          -> ((forall s. [Var s] -> Var s) -> (Params -> m Params))
          -> Params-> m Params -- パラメータを更新する関数
miniBatch size trainingData updateWith w = do
  tds <- liftIO $ shuffleM trainingData
  let chunks = chunksOf size tds
      chunkLoss chunk w = sum $ loss w <$> chunk
      updates = map (\chunk -> updateWith $ chunkLoss chunk) chunks
  foldr (>=>) pure updates $ w
```

[`chunksOf`](http://hackage.haskell.org/package/extra-1.5/docs/Data-List-Extra.html#v:chunksOf)は[Data.List.Extra](http://hackage.haskell.org/package/extra-1.5/docs/Data-List-Extra.html)に含まれる関数で、リストを複数個の組に分けることが出来ます。

```hs
> import Data.List.Extra

> :t chunksOf
chunksOf :: Int -> [a] -> [[a]]

> chunksOf 2 [1..8]
[[1,2],[3,4],[5,6],[7,8]]

> chunksOf 3 [1..8]
[[1,2,3],[4,5,6],[7,8]]
```

`>=>`は理論的にもとても大切な関数でモナドに包まれた値を返すような関数同士を合成することが出来ます。

```hs
> import Control.Monad

> :t (>=>)
(>=>) :: Monad m => (a -> m b) -> (b -> m c) -> a -> m c
```

この`miniBatch`という関数を使ってバッチ勾配降下法と確率的勾配降下法とミニバッチ勾配降下法をそれぞれ実装してみましょう。

```hs
main :: IO ()
main = do
  trainingData <- generateTrainingData 1000
  initialParams <- createInitialParams
  print initialParams

  let update1 = miniBatch (length trainingData) trainingData (\loss w -> pure $ gradientDescent 0.01 loss w)
  updatedParams1 <- (foldr (>=>) pure $ replicate 100 update1) initialParams
  putStrLn $ "バッチ勾配降下法: " ++ show updatedParams1

  let update2 = miniBatch 1 trainingData (\loss w -> pure $ gradientDescent 0.01 loss w)
  updatedParams2 <- (foldr (>=>) pure $ replicate 100 update2) initialParams
  putStrLn $ "確率的勾配降下法: " ++ show updatedParams2

  let update3 = miniBatch 50 trainingData (\loss w -> pure $ gradientDescent 0.01 loss w)
  updatedParams3 <- (foldr (>=>) pure $ replicate 100 update3) initialParams
  putStrLn $ "ミニバッチ勾配降下法: " ++ show updatedParams3
```

```bash
$ runhaskell Main.hs
[-0.1314267134564846,0.47514164081265275,0.9054138942958068]
バッチ勾配降下法: [6.988151826244106,7.241309154445538,-7.404787622071784]
確率的勾配降下法: [6.865422819036771,7.116453501225986,-7.273539018836106]
ミニバッチ勾配降下法: [6.867915104203125,7.1170088242408935,-7.278798196866264]
```

### Momentum(慣性)
過去の更新ベクトルを一定の割合で混ぜて更新する方法です。 

```math
v_t = \gamma v_{t-1} + \alpha \frac{\partial L}{\partial w}
```

```math
w_{t+1} = w_t - v_t
```

Haskellで実装する時は過去の更新ベクトルを参照して今回の更新ベクトルで置き換えるので`State`モナドを使うと良さそうです。

```hs
gdWithMomentum :: Double -> Double -- 学習率, 慣性係数
               -> (forall s. [Var s] -> Var s)-> Params -> State [Double] Params
gdWithMomentum rate gamma f ws = do
  State.modify $ \vs -> gamma *| vs |+| rate *| grad f ws
  State.gets $ \vs -> ws |-| vs
```

びっくりするぐらい定義式と同じように実装できました。
さらに先程用意した`miniBatch`は`MonadIO`を使って定義されているので何の修正を加えること無く今作った`gdWithMomentum`と組み合わせることが出来ます。

```hs
main :: IO ()
main = do
  trainingData <- generateTrainingData 1000
  initialParams <- createInitialParams
  print initialParams

  let update = miniBatch 50 trainingData (\loss w -> hoist generalize $ gdWithMomentum 0.01 0.9 loss w)
  updatedParams <- flip State.evalStateT [0,0,0] $ (foldr (>=>) pure $ replicate 100 update) initialParams
  print updatedParams
```

実は`mtl`で定義されている`State`は単純に`StateT s Identity`のエイリアスになっています。`hoist generalize` は[mmorph](http://hackage.haskell.org/package/mmorph)というHackageに含まれる関数で

```hs
> :t hoist generalize
hoist generalize
  :: (Control.Monad.Morph.MFunctor t, Monad n) =>
     t Data.Functor.Identity.Identity b -> t n b
```

のような型をしていて、ここでは`StateT [Double] Identity a`を`StateT [Double] IO a`に変換するために使っています。
実行してみましょう。

```bash
$ runhaskell Main.hs
[-0.30730664055870904,-2.0091256832789783e-2,0.625856609384243]
[16.04566506149459,15.775838482410673,-16.2650354056726]
```

Nesterovの加速勾配降下法は`gdWithMomentum`を少し修正するだけで実装できます。

```hs
nag :: Double -> Double
    -> (forall s. [Var s] -> Var s) -> Params -> State [Double] Params
nag rate gamma f ws = do
  State.modify $ \vs -> gamma *| vs |+| rate *| grad f (ws |-| gamma *| vs)
  State.gets $ \vs -> ws |-| vs
```

### AdaGrad, AdaDelta, RMSProp, Adam
道具は出揃ったので後の手法は実装だけ載せておきます

```hs
adagrad :: Double
        -> (forall s. [Var s] -> Var s) -> Params -> State [Double] Params
adagrad rate f ws = do
  let vs = grad f ws
  State.modify $ \gs -> gs |+| map (^2) vs
  State.gets $ \gs -> ws  |-| rate *| vs |/| (sqrt <$> (gs |+| eps))
    where
      eps = repeat 1.0e-8


adadelta :: Double
         -> (forall s. [Var s] -> Var s) -> Params -> State [(Double, Double)] Params
adadelta rho f ws = do
  (ds, gs) <- State.gets unzip
  let vs = grad f ws
      gs' = rho *| gs |+| (1-rho) *| map (^2) vs
      deltas = vs |*| (sqrt <$> ((ds |+| eps) |/| (gs' |+| eps)))
      ds' = rho *| ds |+| (1-rho) *| deltas |*| deltas
  State.put $ zip ds' gs'
  pure $ ws |-| deltas
    where
      eps = repeat 1.0e-8


rmsprop :: Double
        -> (forall s. [Var s] -> Var s) -> Params -> State [Double] Params
rmsprop rate f ws = do
  let vs = grad f ws
  State.modify $ \gs -> 0.9 *| gs |+| 0.1 *| map (^2) vs
  State.gets $ \gs -> ws  |-| rate *| vs |/| (sqrt <$> (gs |+| eps))
    where
      eps = repeat 1.0e-8


adam :: Double -> Double -> Double
     -> (forall s. [Var s] -> Var s) -> Params -> State (Int, [(Double, Double)]) Params
adam rate beta1 beta2 f ws = do
  (count, (ms, ds)) <- State.gets (fmap unzip)
  let vs = grad f ws
      ms' = beta1 *| ms |+| (1-beta1) *| vs
      ds' = beta2 *| ds |+| (1-beta2) *| map (^2) vs
      count' = count + 1
      deltas = rate *| (ms' |/| repeat (1-beta1^count')) |/| (sqrt <$> (ds' |/| repeat (1 - beta2^count')) |+| eps)
  State.put $ (count', zip ms' ds')
  pure $ ws |-| deltas
    where
      eps = repeat 1.0e-8
```

終わりに
-------
実装したコードは[gistのほう](https://gist.github.com/lotz84/6c9076e110ff335939ea3e3085128d80)に載せておきました。今回は実装するのが目的だったので学習曲線などは出していませんがWriterモナドなどを使えば簡単に途中経過を記録できると思います。グラフ描画には[Chart](https://github.com/timbod7/haskell-chart/wiki)が便利です。実装はHaskellの説明のため簡潔でわかりやすいコードになるよう心がけたので必ずしも最速な実装になっているわけではありません。gistには簡単なベンチマークを載せてあるのでHaskellでもここまで高速な実装ができるよーみたいな記事を書いてもらえると泣いて喜びます :joy: 

参考
-----
* [勾配降下法の最適化アルゴリズムを概観する](http://postd.cc/optimizing-gradient-descent/)
* [AdaGrad, RMSProp, AdaDelta, Adam, SMORMS3](http://qiita.com/skitaoka/items/e6afbe238cd69c899b2a)

[^1]: http://conway.rutgers.edu/~ccshan/wiki/blog/posts/Differentiation/