0 | module NN.Optimisers.Definition
 1 |
 2 | import Data.Tensor
 3 | import Data.Container.Additive
 4 | import NN.Utils
 5 |
 6 | -- We can make a choice in optimisation:
 7 | -- a) either the parameter can depend on the input, in which case we can't use the hom representation of a learner
 8 | -- b) or it doesn't, meaning we can curry the learner and treat it as something we optimise
 9 |
10 |
11 | ||| Dependent stateful optimiser, modelled as a dependent lens
12 | ||| Dependent version of section 8.1.3 of https://arxiv.org/abs/2403.13001
13 | ||| Because we use dependent Para, optimiser can depend on the input
14 | public export
15 | record DOptimiser {inputTy : Type}
16 |   (paramCont : inputTy -> AddCont)
17 |   (stateTy : Type)
18 |   where
19 |   constructor MkDOptimiser
20 |   ||| Notably this produces an ordinary dependent lens, not an additive one
21 |   opt : (x : inputTy) ->
22 |     (Const (paramCont x).Shp >< Const stateTy) =%> UC (paramCont x)
23 |
24 |
25 | -- todo upgrade to something like HasIO instead of IO?
26 | public export
27 | record Optimiser
28 |   (paramCont : AddCont)
29 |   (stateTy : Type)
30 |   where
31 |   constructor MkOptimiser
32 |   opt : Const paramCont.Shp >< Const stateTy =%> UC paramCont
33 |   ||| Initialising the state; parameters are initialised by the model
34 |   initState : IO stateTy
35 |
36 | ||| An optimiser is a lens: forward reads the parameter, backward updates (parameter, state)
37 | public export
38 | (.fwd) : Optimiser p s -> (p.Shp, s) -> p.Shp
39 | (.fwd) (MkOptimiser opt _) = opt.fwd
40 |
41 | public export
42 | (.bwd) : (opt : Optimiser pCont stateTy) ->
43 |   (ps : (pCont.Shp, stateTy)) ->
44 |   (pCont.PosSet (opt.fwd ps)) ->
45 |   (pCont.Shp, stateTy)
46 | (.bwd) (MkOptimiser opt _) = opt.bwd
47 |
48 | ||| From 8.1.3. "Can we compose optimisers?" of https://arxiv.org/abs/2403.13001
49 | ||| Not used yet
50 | public export
51 | composeParallel : Optimiser pCont s ->
52 |   Optimiser qCont t -> 
53 |   Optimiser (pCont >*< qCont) (s, t)
54 | composeParallel (MkOptimiser o1 initS) (MkOptimiser o2 initT) = MkOptimiser
55 |   (!% \((p, q), (s, t)) => ((o1.fwd (p, s), o2.fwd (q, t)) **
56 |     \(p', q') => let (pUpdated, sUpdated) = o1.bwd (p, s) p'
57 |                      (qUpdated, tUpdated) = o2.bwd (q, t) q'
58 |                  in ((pUpdated, qUpdated), (sUpdated, tUpdated))))
59 |   [| (initS, initT) |]