Learning.ScaledGradientUpdate

1module Learning.ScaledGradientUpdate
2
3import Component.Pointwise
4import Std.Fin
5import Std.List
6import Std.Natural
7
8-- These two operations compose independently: callers choose the arithmetic
9-- and the source of the gradient. The accumulation state is persistent, and
10-- applying the parameter update consumes the newly accumulated value.
11def accumulateScaledGradient =
12  (lambda erased value : Type 0 .
13    (lambda unrestricted add : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
14      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
15        (lambda unrestricted decay : value .
16          (lambda unrestricted previous : value .
17            (lambda unrestricted gradient : value . (add (multiply decay previous) gradient)))))))
18
19def subtractScaledUpdate =
20  (lambda erased value : Type 0 .
21    (lambda unrestricted subtract : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
22      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
23        (lambda unrestricted rate : value .
24          (lambda unrestricted parameter : value .
25            (lambda unrestricted update : value . (subtract parameter (multiply rate update))))))))
26
27def scaledGradientReadOnly =
28  (lambda erased slots : Nat .
29    (lambda unrestricted writable : (family StdFin slots) .
30      (lambda unrestricted gradient : (family StdFin slots) .
31        (lambda unrestricted decay : (family StdFin slots) .
32          (lambda unrestricted rate : (family StdFin slots) .
33            (naturalAnd
34              (pointwiseSlotsSeparate slots writable gradient)
35              (naturalAnd
36                (pointwiseSlotsSeparate slots writable decay)
37                (pointwiseSlotsSeparate slots writable rate))))))))
38
39def scaledGradientSlotsIndependent =
40  (lambda erased slots : Nat .
41    (lambda unrestricted parameter : (family StdFin slots) .
42      (lambda unrestricted accumulator : (family StdFin slots) .
43        (lambda unrestricted gradient : (family StdFin slots) .
44          (lambda unrestricted decay : (family StdFin slots) .
45            (lambda unrestricted rate : (family StdFin slots) .
46              (naturalAnd
47                (pointwiseSlotsSeparate slots parameter accumulator)
48                (naturalAnd
49                  (scaledGradientReadOnly slots parameter gradient decay rate)
50                  (scaledGradientReadOnly slots accumulator gradient decay rate)))))))))
51
52-- The same equations can produce a typed state program instead of evaluating
53-- numbers immediately. Positions are supplied by the state's owner. Assigning
54-- the accumulator first makes its new value available to the parameter update;
55-- the compiler must preserve that dependency and each separate rounding.
56-- Writable slots are independent of each other and of the read-only inputs;
57-- otherwise the first assignment could overwrite a later equation's input.
58def scaledGradientAssignments =
59  (lambda erased word : Type 0 .
60    (lambda erased slots : Nat .
61      (lambda unrestricted parameter : (family StdFin slots) .
62        (lambda unrestricted accumulator : (family StdFin slots) .
63          (lambda unrestricted gradient : (family StdFin slots) .
64            (lambda unrestricted decay : (family StdFin slots) .
65              (lambda unrestricted rate : (family StdFin slots) .
66                (lambda erased independent : (equal Nat (scaledGradientSlotsIndependent slots parameter accumulator gradient decay rate) 1) .
67                  (let unrestricted input =
68                    (lambda unrestricted slot : (family StdFin slots) .
69                      (constructor PointwiseExpression PointwiseInput word slots slot))
70                    in
71                    (let unrestricted multiply =
72                      (pointwiseBinary word slots (constructor PointwiseBinary PointwiseMultiply))
73                      in
74                      (constructor
75                        StdList
76                        StdListCons
77                        (family PointwiseAssignment word slots)
78                        (constructor
79                          PointwiseAssignment
80                          PointwiseAssign
81                          word
82                          slots
83                          accumulator
84                          (accumulateScaledGradient
85                            (family PointwiseExpression word slots)
86                            (pointwiseBinary word slots (constructor PointwiseBinary PointwiseAdd))
87                            multiply
88                            (input decay)
89                            (input accumulator)
90                            (input gradient)))
91                        (constructor
92                          StdList
93                          StdListCons
94                          (family PointwiseAssignment word slots)
95                          (constructor
96                            PointwiseAssignment
97                            PointwiseAssign
98                            word
99                            slots
100                            parameter
101                            (subtractScaledUpdate
102                              (family PointwiseExpression word slots)
103                              (pointwiseBinary
104                                word
105                                slots
106                                (constructor PointwiseBinary PointwiseSubtract))
107                              multiply
108                              (input rate)
109                              (input parameter)
110                              (input accumulator)))
111                          (constructor StdList StdListEmpty (family PointwiseAssignment word slots))))))))))))))