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