Bob.Learner

1module Bob.Learner
2
3import Data.Bytes
4import Data.Float32Bits
5import Model.Architecture
6import Representation.Schema
7import Std.Float
8import Std.Natural
9
10family BobLearningSemantics : Type 0
11constructor BobLearningSemanticsValue
12field unrestricted bobLearningObjectiveIdentity : Bytes
13field unrestricted bobLearningObjective : (family ModelObjective)
14field unrestricted bobLearningLearnerIdentity : Bytes
15field unrestricted bobLearningLearner : (family LearnerContract)
16
17end-family
18
19def bobObjectiveIdentity : Bytes =
20  b"exact-full-256-cross-entropy"
21
22def bobObjective : (family ModelObjective) =
23  (constructor ModelObjective ModelExactVocabulary)
24
25def bobLearnerIdentity : Bytes =
26  b"adamw"
27
28def bobLearner : (family LearnerContract) =
29  (constructor LearnerContract LearnerAdamW)
30
31def bobLearningSemantics : (family BobLearningSemantics) =
32  (record
33    BobLearningSemantics
34    (bobLearningObjectiveIdentity = bobObjectiveIdentity)
35    (bobLearningObjective = bobObjective)
36    (bobLearningLearnerIdentity = bobLearnerIdentity)
37    (bobLearningLearner = bobLearner))
38
39-- AdamW hyperparameters.  Everything the optimizer kernel reads per step is
40-- computed from these four values and the step number.
41def bobLearningRate : F32 = 0.05
42def bobBeta1 : F32 = 0.9
43def bobBeta2 : F32 = 0.999
44def bobEpsilon : F32 = 1e-8
45def bobWeightDecay : F32 = 0.0
46
47def bobF32One : F32 = 1.0
48
49def bobF32Power =
50  (lambda unrestricted base : F32 .
51    (lambda unrestricted exponent : Nat .
52      (nat-eliminate
53        (lambda unrestricted current : Nat . F32)
54        bobF32One
55        (lambda unrestricted predecessor : Nat .
56          (lambda unrestricted induction : F32 . (f32-multiply induction base)))
57        exponent)))
58
59-- The nine scalars the AdamW SM86 kernel reads, in its ABI order: gradient
60-- scale, beta1, 1 - beta1, beta2, 1 - beta2, the decay multiplier
61-- 1 - lr * weightDecay, the bias-corrected step size
62-- lr * sqrt(1 - beta2^t) / (1 - beta1^t), the epsilon scaled by the same
63-- sqrt(1 - beta2^t), and the constant one.  `step` counts from 1.
64--
65-- Everything is binary32, the kernel's own precision.  1 - beta2^t loses
66-- about five significant digits to cancellation at small t (beta2^t is close
67-- to one), so the step size carries a relative error near 1e-5; the kernel
68-- applies the same tolerance to its running moments, and the value is what the
69-- learner declares rather than a number computed elsewhere.
70def bobAdamStepScalars =
71  (lambda unrestricted step : Nat .
72    (let unrestricted rootSecond =
73      (f32-sqrt (f32-subtract bobF32One (bobF32Power bobBeta2 step)))
74    in
75      (bytes-builder-build
76        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobF32One))
77        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta1))
78        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta1)))
79        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta2))
80        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta2)))
81        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One (f32-multiply bobLearningRate bobWeightDecay))))
82        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE
83          (f32-divide (f32-multiply bobLearningRate rootSecond)
84                      (f32-subtract bobF32One (bobF32Power bobBeta1 step)))))
85        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-multiply bobEpsilon rootSecond)))
86                              (bytes-builder-chunk (dataFloat32EncodeLE bobF32One)))))))))))))