0 | module NN.Optimisers.Instances
  1 |
  2 | import Data.Materialise
  3 | import Data.Container.Additive
  4 | import Data.Num
  5 | import NN.Optimisers.Definition
  6 |
  7 |
  8 | ||| Gradient ascent optimiser. Has trivial state
  9 | ||| @lr is the learning rate
 10 | public export
 11 | GA : Num pType => (mon : ComMonoid pType) => FromDouble pType =>
 12 |   {default 0.001 lr : pType} -> Optimiser (Const pType) Unit
 13 | GA = MkOptimiser
 14 |   (!% \(p, ()) => (p ** \p' => (p + lr * p', ())))
 15 |   (pure ())
 16 |
 17 | ||| Gradient descent optimiser. Has trivial state
 18 | ||| @lr is the learning rate
 19 | public export
 20 | GD : Neg pType =>
 21 |   (mon : ComMonoid pType) => FromDouble pType =>
 22 |   {default 0.001 lr : pType} -> Optimiser (Const pType) Unit
 23 | GD = GA {lr = negate lr}
 24 |
 25 | namespace Momentum
 26 |   public export
 27 |   momentumUpdate : Neg pType =>
 28 |     {lr : pType} ->
 29 |     (gamma : pType) ->
 30 |     (p : pType) ->
 31 |     (s : pType) ->
 32 |     (p' : pType) ->
 33 |     (pType, pType)
 34 |   momentumUpdate gamma p s p' = let s' = gamma * s + p'
 35 |                                 in (p - lr * s', s')
 36 |
 37 |   public export
 38 |   lookAhead : Num pType =>
 39 |     (gamma, p, s : pType) ->
 40 |     pType
 41 |   lookAhead gamma p s = p + gamma * s
 42 |   
 43 |   ||| Gradient Descent with momentum, optionally with Nesterov acceleration
 44 |   public export
 45 |   GDMomentum : Neg pType =>
 46 |    (mon : ComMonoid pType) =>
 47 |    FromDouble pType =>
 48 |    {default False nesterov : Bool} ->
 49 |    {default 0.001 lr : pType} ->
 50 |    {default 0.9 gamma : pType} ->
 51 |    Optimiser (Const pType) pType
 52 |   GDMomentum = MkOptimiser
 53 |     (!% \(p, s) => (if nesterov then lookAhead gamma p s else p
 54 |                    ** momentumUpdate {lr} gamma p s))
 55 |     (pure 0)
 56 |   
 57 | namespace Adam
 58 |   ||| Adam step. The moments are parameter-shaped; the bias-correction powers
 59 |   ||| beta^t are the same scalar at every coordinate, so they are `Double`
 60 |   public export
 61 |   adamUpdate : Neg pType => Fractional pType => Sqrt pType =>
 62 |     FromDouble pType => Materialise pType =>
 63 |     {lr : pType} ->
 64 |     (beta1 : Double) ->
 65 |     (beta2 : Double) ->
 66 |     (epsilon : pType) ->
 67 |     (p : pType) ->
 68 |     (m : pType) ->
 69 |     (v : pType) ->
 70 |     (b1p : Double) ->
 71 |     (b2p : Double) ->
 72 |     (g : pType) ->
 73 |     (pType, pType, pType, Double, Double)
 74 |   adamUpdate beta1 beta2 epsilon p m v b1p b2p g =
 75 |     let g' = materialise g
 76 |         m' = materialise (fromDouble beta1 * m + fromDouble (1 - beta1) * g')
 77 |         v' = materialise (fromDouble beta2 * v + fromDouble (1 - beta2) * g' * g')
 78 |         b1p' = b1p * beta1
 79 |         b2p' = b2p * beta2
 80 |         mHat = fromDouble (1 / (1 - b1p')) * m'
 81 |         vHat = fromDouble (1 / (1 - b2p')) * v'
 82 |     in (p - lr * mHat / (sqrt vHat + epsilon), m', v', b1p', b2p')
 83 |
 84 |   ||| Adam optimiser
 85 |   ||| State: the two moments and the two scalar bias-correction powers
 86 |   ||| @lr is the learning rate
 87 |   ||| @beta1 is the exponential decay rate for the first moment estimate
 88 |   ||| @beta2 is the exponential decay rate for the second moment estimate
 89 |   ||| @epsilon is a small constant for numerical stability
 90 |   public export
 91 |   Adam : Neg pType =>
 92 |    (mon : ComMonoid pType) =>
 93 |    FromDouble pType => Materialise pType =>
 94 |    Fractional pType => Sqrt pType =>
 95 |    {default 0.001 lr : pType} ->
 96 |    {default 0.9 beta1 : Double} ->
 97 |    {default 0.999 beta2 : Double} ->
 98 |    {default 1.0e-8 epsilon : pType} ->
 99 |    Optimiser (Const pType) (pType, pType, Double, Double)
100 |   Adam = MkOptimiser
101 |     (!% \(p, (m, v, b1p, b2p)) =>
102 |       (p ** adamUpdate {lr} beta1 beta2 epsilon p m v b1p b2p))
103 |     (pure (0, 0, 1, 1))