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 |   ||| Procedure for initialising parameters
34 |   initParam : IO paramCont.Shp
35 |   ||| Procedure for initialising state
36 |   initState : IO stateTy
37 |
38 | public export
39 | (.fwd) : Optimiser p s -> (p.Shp, s) -> p.Shp
40 | (.fwd) (MkOptimiser opt _ _) = opt.fwd
41 |
42 | public export
43 | (.bwd) : (opt : Optimiser pCont stateTy) ->
44 |   (ps : (pCont.Shp, stateTy)) ->
45 |   (pCont.Pos (opt.fwd ps)) ->
46 |   (pCont.Shp, stateTy)
47 | (.bwd) (MkOptimiser opt _ _) = opt.bwd
48 |
49 | ||| From 8.1.3. "Can we compose optimisers?" of https://arxiv.org/abs/2403.13001
50 | ||| Not used yet
51 | public export
52 | composeParallel : Optimiser pCont s ->
53 |   Optimiser qCont t -> 
54 |   Optimiser (pCont >< qCont) (s, t)
55 | composeParallel (MkOptimiser o1 initP initS) (MkOptimiser o2 initQ initT) = MkOptimiser
56 |   (!% \((p, q), (s, t)) => ((o1.fwd (p, s), o2.fwd (q, t)) **
57 |     \(p', q') => let (pUpdated, sUpdated) = o1.bwd (p, s) p'
58 |                      (qUpdated, tUpdated) = o2.bwd (q, t) q'
59 |                  in ((pUpdated, qUpdated), (sUpdated, tUpdated))))
60 |   [| (initP, initQ) |]
61 |   [| (initS, initT) |]