Learning.Learner.MomentNormalized

1module Learning.Learner.MomentNormalized
2
3import Std.Float
4import Data.Float32Bits
5
6-- A CPU (Std.Float F32) learner that moves each parameter by its bias-corrected
7-- first moment divided by the root of its bias-corrected second moment:
8--
9--   m' = beta1 m + (1 - beta1) g                       firstMomentValue
10--   v' = beta2 v + (1 - beta2) (g g)                   secondMomentValue
11--   p' = p beta                                        momentDecayPowerNext
12--   m^ = m' / (1 - p1'),  v^ = v' / (1 - p2')          correctMomentBias
13--   w' = w - (rate m^) / (sqrt v^ + epsilon)            normalizedStepValue
14--
15-- This is the update known as Adam (Kingma & Ba 2015, Algorithm 1) without
16-- weight decay. It differs from Learning.Checked.AdamW, the SM86 owner, on
17-- purpose: there the moments use fused multiply-adds, the reciprocal square
18-- root is the MUFU approximation and the bias corrections are folded into a
19-- precomputed step size; here every operation is a correctly rounded binary32
20-- add, subtract, multiply, divide or sqrt in the order written above, so a CPU
21-- run is reproducible bit for bit by any IEEE-754 implementation that keeps
22-- that order (systems/manfred checks it against a Nim reference).
23--
24-- The complements 1 - beta are derived here in F32 rather than passed in:
25-- 1 - 0.9 rounds to 0.100000024, not to the literal 0.1, and a learner that
26-- took the complement as a separate hyperparameter could silently disagree
27-- with its own beta. The bias powers beta^t are carried by the caller as
28-- running products (momentDecayPowerNext), which needs no pow.
29
30family MomentNormalizedHyper : Type 0
31constructor MomentNormalizedHyperValue
32field unrestricted momentNormalizedRate : F32
33field unrestricted momentNormalizedBeta1 : F32
34field unrestricted momentNormalizedBeta2 : F32
35field unrestricted momentNormalizedEpsilon : F32
36
37end-family
38
39-- One parameter with its two moments, before or after an update.
40family MomentNormalizedElement : Type 0
41constructor MomentNormalizedElementValue
42field unrestricted momentNormalizedWeight : F32
43field unrestricted momentNormalizedFirst : F32
44field unrestricted momentNormalizedSecond : F32
45
46end-family
47
48-- m' = beta m + (1 - beta) g
49def firstMomentValue =
50  (lambda unrestricted beta : F32 .
51    (lambda unrestricted moment : F32 .
52      (lambda unrestricted gradient : F32 .
53        (f32-add
54          (f32-multiply beta moment)
55          (f32-multiply (f32-subtract stdF32One beta) gradient)))))
56
57-- v' = beta v + (1 - beta) (g g)
58def secondMomentValue =
59  (lambda unrestricted beta : F32 .
60    (lambda unrestricted moment : F32 .
61      (lambda unrestricted gradient : F32 .
62        (f32-add
63          (f32-multiply beta moment)
64          (f32-multiply (f32-subtract stdF32One beta) (f32-multiply gradient gradient))))))
65
66-- beta^(t+1) from beta^t; start the product at one.
67def momentDecayPowerNext =
68  (lambda unrestricted beta : F32 .
69    (lambda unrestricted power : F32 .
70      (f32-multiply power beta)))
71
72-- m / (1 - beta^t), with beta^t already advanced to this update.
73def correctMomentBias =
74  (lambda unrestricted power : F32 .
75    (lambda unrestricted moment : F32 .
76      (f32-divide moment (f32-subtract stdF32One power))))
77
78-- w - (rate m^) / (sqrt v^ + epsilon)
79def normalizedStepValue =
80  (lambda unrestricted rate : F32 .
81    (lambda unrestricted epsilon : F32 .
82      (lambda unrestricted weight : F32 .
83        (lambda unrestricted first : F32 .
84          (lambda unrestricted second : F32 .
85            (f32-subtract
86              weight
87              (f32-divide
88                (f32-multiply rate first)
89                (f32-add (f32-sqrt second) epsilon))))))))
90
91-- One parameter's update, given the bias powers already advanced to this step
92-- (power1 = beta1^t, power2 = beta2^t for the t-th update, t >= 1).
93def momentNormalizedUpdate =
94  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
95    (lambda unrestricted power1 : F32 .
96      (lambda unrestricted power2 : F32 .
97        (lambda unrestricted element : (family MomentNormalizedElement) .
98          (lambda unrestricted gradient : F32 .
99            (eliminate MomentNormalizedHyper
100              (lambda unrestricted current : (family MomentNormalizedHyper) . (family MomentNormalizedElement))
101              hyper
102              (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon .
103                (eliminate MomentNormalizedElement
104                  (lambda unrestricted current : (family MomentNormalizedElement) . (family MomentNormalizedElement))
105                  element
106                  (branch MomentNormalizedElementValue weight first second .
107                    (let unrestricted first1 = (firstMomentValue beta1 first gradient) in
108                      (let unrestricted second1 = (secondMomentValue beta2 second gradient) in
109                        (constructor MomentNormalizedElement MomentNormalizedElementValue
110                          (normalizedStepValue rate epsilon weight
111                            (correctMomentBias power1 first1)
112                            (correctMomentBias power2 second1))
113                          first1
114                          second1))))))))))))
115
116def momentNormalizedWeightOf =
117  (lambda unrestricted element : (family MomentNormalizedElement) .
118    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
119      (branch MomentNormalizedElementValue weight first second . weight)))
120
121def momentNormalizedFirstOf =
122  (lambda unrestricted element : (family MomentNormalizedElement) .
123    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
124      (branch MomentNormalizedElementValue weight first second . first)))
125
126def momentNormalizedSecondOf =
127  (lambda unrestricted element : (family MomentNormalizedElement) .
128    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
129      (branch MomentNormalizedElementValue weight first second . second)))
130
131def momentNormalizedBeta1Of =
132  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
133    (eliminate MomentNormalizedHyper (lambda unrestricted current : (family MomentNormalizedHyper) . F32) hyper
134      (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon . beta1)))
135
136def momentNormalizedBeta2Of =
137  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
138    (eliminate MomentNormalizedHyper (lambda unrestricted current : (family MomentNormalizedHyper) . F32) hyper
139      (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon . beta2)))