main.bend

bend-ml-autograd@0.1.1.0 · 299 lines · 33 definitions

1# bend-ml-autograd: automatic differentiation with a proved law.
2#
3#   import bend-ml-autograd@0.1.1.0/main.bend as AG
4#
5# 1) PROVED MODEL (over Nat): the reverse mode of autodiff gives the same result
6#    as the forward mode (dual numbers). Proved by the kernel, using the sum and
7#    product lemmas of bend-ml-nat-lemmas.
8# 2) SCALAR AUTOGRAD in F32 (micrograd style), with the same structure.
9# 3) LAYERS WITH TENSORS (forward and backward) whose shapes are checked by the type.
10#
11# F32 is not a real number, so the NUMERICAL correctness of F32 is validated by
12# tests against PyTorch (reference/test_autograd.py), not by proof.
13
14import Base
15import bend-ml-nat-lemmas@0.1.0.0/main.bend as NL
16import bend-ml-tensor@0.1.0.0/main.bend as T
17
18# =====================================================================
19# 1) Proved model
20# =====================================================================
21
22# Expressions with one variable NX, over Nat (a commutative semiring).
23type NE is Data:
24  NCst{n: Nat}
25  NX{}
26  NAdd{a: NE, b: NE}
27  NMul{a: NE, b: NE}
28
29# value of the expression at x
30def nval(e: NE, +x: Nat) -> Nat:
31  match e:
32    case NCst{n}:
33      n
34    case NX{}:
35      x
36    case NAdd{a, b}:
37      Nat.add(nval(a, x), nval(b, x))
38    case NMul{+a, +b}:
39      Nat.mul(nval(a, x), nval(b, x))
40
41# FORWARD-mode derivative (dual numbers): derivative of the sum and the product rule
42def nfwd(e: NE, +x: Nat) -> Nat:
43  match e:
44    case NCst{n}:
45      0n
46    case NX{}:
47      1n
48    case NAdd{a, b}:
49      Nat.add(nfwd(a, x), nfwd(b, x))
50    case NMul{+a, +b}:
51      Nat.add(Nat.mul(nfwd(a, x), nval(b, x)), Nat.mul(nval(a, x), nfwd(b, x)))
52
53# REVERSE-mode derivative: g is the gradient arriving from the output; each node passes on
54# g (sum) or g * value-of-the-other-factor (product) to its children and adds up what
55# arrives at NX.
56def nbwd(e: NE, +x: Nat, +g: Nat) -> Nat:
57  match e:
58    case NCst{n}:
59      0n
60    case NX{}:
61      g
62    case NAdd{a, b}:
63      Nat.add(nbwd(a, x, g), nbwd(b, x, g))
64    case NMul{+a, +b}:
65      Nat.add(nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))))
66
67# a * (b + c) == a*b + a*c   (left distributivity; nat-lemmas only has the right one)
68def ndist_l(+a: Nat, +b: Nat, +c: Nat) -> {Nat.mul(a, Nat.add(b, c)) == Nat.add(Nat.mul(a, b), Nat.mul(a, c)) : Nat}:
69  Equal.trans(Nat, Nat.mul(a, Nat.add(b, c)), Nat.mul(Nat.add(b, c), a), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
70    NL.mul_comm(a, Nat.add(b, c)),
71    Equal.trans(Nat, Nat.mul(Nat.add(b, c), a), Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
72      Equal.sym(Nat, Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.mul(Nat.add(b, c), a), NL.mul_dist(b, c, a)),
73      Equal.trans(Nat, Nat.add(Nat.mul(b, a), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(c, a)), Nat.add(Nat.mul(a, b), Nat.mul(a, c)),
74        Equal.cong(Nat, Nat, k => Nat.add(k, Nat.mul(c, a)), Nat.mul(b, a), Nat.mul(a, b), NL.mul_comm(b, a)),
75        Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(a, b), k), Nat.mul(c, a), Nat.mul(a, c), NL.mul_comm(c, a)))))
76
77# The product case, with everything abstracted into numbers:
78# (g*vb)*fa + (g*va)*fb == g * (fa*vb + va*fb)
79def nmul_case(+g: Nat, +fa: Nat, +fb: Nat, +va: Nat, +vb: Nat, +ba: Nat, +bb: Nat, iha: {ba == Nat.mul(Nat.mul(g, vb), fa) : Nat}, ihb: {bb == Nat.mul(Nat.mul(g, va), fb) : Nat}) -> {Nat.add(ba, bb) == Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))) : Nat}:
80  Equal.trans(Nat, Nat.add(ba, bb), Nat.add(Nat.mul(Nat.mul(g, vb), fa), bb), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
81    Equal.cong(Nat, Nat, k => Nat.add(k, bb), ba, Nat.mul(Nat.mul(g, vb), fa), iha),
82    Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(g, vb), fa), bb), Nat.add(Nat.mul(Nat.mul(g, vb), fa), Nat.mul(Nat.mul(g, va), fb)), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
83      Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(Nat.mul(g, vb), fa), k), bb, Nat.mul(Nat.mul(g, va), fb), ihb),
84      Equal.trans(Nat, Nat.add(Nat.mul(Nat.mul(g, vb), fa), Nat.mul(Nat.mul(g, va), fb)), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(Nat.mul(g, va), fb)), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
85        Equal.cong(Nat, Nat, k => Nat.add(k, Nat.mul(Nat.mul(g, va), fb)), Nat.mul(Nat.mul(g, vb), fa), Nat.mul(g, Nat.mul(fa, vb)),
86          Equal.trans(Nat, Nat.mul(Nat.mul(g, vb), fa), Nat.mul(g, Nat.mul(vb, fa)), Nat.mul(g, Nat.mul(fa, vb)),
87            Equal.sym(Nat, Nat.mul(g, Nat.mul(vb, fa)), Nat.mul(Nat.mul(g, vb), fa), NL.mul_assoc(g, vb, fa)),
88            Equal.cong(Nat, Nat, k => Nat.mul(g, k), Nat.mul(vb, fa), Nat.mul(fa, vb), NL.mul_comm(vb, fa)))),
89        Equal.trans(Nat, Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(Nat.mul(g, va), fb)), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(g, Nat.mul(va, fb))), Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))),
90          Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(g, Nat.mul(fa, vb)), k), Nat.mul(Nat.mul(g, va), fb), Nat.mul(g, Nat.mul(va, fb)),
91            Equal.sym(Nat, Nat.mul(g, Nat.mul(va, fb)), Nat.mul(Nat.mul(g, va), fb), NL.mul_assoc(g, va, fb))),
92          Equal.sym(Nat, Nat.mul(g, Nat.add(Nat.mul(fa, vb), Nat.mul(va, fb))), Nat.add(Nat.mul(g, Nat.mul(fa, vb)), Nat.mul(g, Nat.mul(va, fb))), ndist_l(g, Nat.mul(fa, vb), Nat.mul(va, fb)))))))
93
94# The reverse mode returns g times the forward-mode derivative.
95def nbwd_ok(e: NE, +x: Nat, +g: Nat) -> {nbwd(e, x, g) == Nat.mul(g, nfwd(e, x)) : Nat}:
96  match e:
97    case NCst{n}:
98      Equal.sym(Nat, Nat.mul(g, 0n), 0n, NL.mul_zero(g))
99    case NX{}:
100      Equal.sym(Nat, Nat.mul(g, 1n), g, NL.mul_one_r(g))
101    case NAdd{+a, +b}:
102      Equal.trans(Nat, Nat.add(nbwd(a, x, g), nbwd(b, x, g)), Nat.add(Nat.mul(g, nfwd(a, x)), nbwd(b, x, g)), Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))),
103        Equal.cong(Nat, Nat, k => Nat.add(k, nbwd(b, x, g)), nbwd(a, x, g), Nat.mul(g, nfwd(a, x)), nbwd_ok(a, x, g)),
104        Equal.trans(Nat, Nat.add(Nat.mul(g, nfwd(a, x)), nbwd(b, x, g)), Nat.add(Nat.mul(g, nfwd(a, x)), Nat.mul(g, nfwd(b, x))), Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))),
105          Equal.cong(Nat, Nat, k => Nat.add(Nat.mul(g, nfwd(a, x)), k), nbwd(b, x, g), Nat.mul(g, nfwd(b, x)), nbwd_ok(b, x, g)),
106          Equal.sym(Nat, Nat.mul(g, Nat.add(nfwd(a, x), nfwd(b, x))), Nat.add(Nat.mul(g, nfwd(a, x)), Nat.mul(g, nfwd(b, x))), ndist_l(g, nfwd(a, x), nfwd(b, x)))))
107    case NMul{+a, +b}:
108      nmul_case(g, nfwd(a, x), nfwd(b, x), nval(a, x), nval(b, x), nbwd(a, x, Nat.mul(g, nval(b, x))), nbwd(b, x, Nat.mul(g, nval(a, x))), nbwd_ok(a, x, Nat.mul(g, nval(b, x))), nbwd_ok(b, x, Nat.mul(g, nval(a, x))))
109
110# LAW (reverse_eq_forward): for any expression made of constants, X, sums
111# and products, the reverse-mode gradient (starting with gradient 1 at the output)
112# equals the forward-mode derivative. It holds for every x.
113law reverse_eq_forward:
114  for +e: NE
115  for +x: Nat
116  {nbwd(e, x, 1n) == nfwd(e, x) : Nat}
117
118def reverse_eq_forward(e, x):
119  Equal.trans(Nat, nbwd(e, x, 1n), Nat.mul(1n, nfwd(e, x)), nfwd(e, x), nbwd_ok(e, x, 1n), NL.mul_one_l(nfwd(e, x)))
120
121
122# =====================================================================
123# 2) Scalar autograd in F32 (reverse mode, micrograd style)
124# =====================================================================
125
126type G is Data:
127  GCst{c: F32}
128  GVar{i: Nat}
129  GAdd{a: G, b: G}
130  GMul{a: G, b: G}
131  GRelu{a: G}
132  GTanh{a: G}
133  GExp{a: G}
134
135def lookup(env: List<&2, F32>, +i: Nat) -> F32:
136  match env i:
137    case Nil{} _:
138      0.0
139    case Con{h, t} 0n:
140      h
141    case Con{h, t} 1n+p:
142      lookup(t, p)
143
144def gval(e: G, +env: List<&2, F32>) -> F32:
145  match e:
146    case GCst{c}:
147      c
148    case GVar{i}:
149      lookup(env, i)
150    case GAdd{a, b}:
151      (gval(a, env) + gval(b, env) : F32)
152    case GMul{a, b}:
153      (gval(a, env) * gval(b, env) : F32)
154    case GRelu{a}:
155      F32.max(gval(a, env), 0.0)
156    case GTanh{a}:
157      F32.tanh(gval(a, env))
158    case GExp{a}:
159      F32.exp(gval(a, env))
160
161# adds v to the gradient of variable i
162def acc_add(acc: List<&2, F32>, +i: Nat, +v: F32) -> List<&2, F32>:
163  match acc i:
164    case Nil{} _:
165      Nil{}
166    case Con{+h, t} 0n:
167      (h + v : F32) <> t
168    case Con{h, t} 1n+p:
169      h <> acc_add(t, p, v)
170
171def mask2(pos: Bool) -> F32:
172  match pos:
173    case True{}:
174      1.0
175    case False{}:
176      0.0
177
178# local derivative of relu: 1 if the input > 0, else 0
179def relu_d(+x: F32) -> F32:
180  mask2(F32.is_gt(x, 0.0))
181
182def tanh_d(+x: F32) -> F32:
183  (1.0 - (F32.tanh(x) * F32.tanh(x) : F32) : F32)
184
185# g is the gradient arriving at the node; acc accumulates the gradient of each variable
186def gbwd(e: G, +env: List<&2, F32>, +g: F32, acc: List<&2, F32>) -> List<&2, F32>:
187  match e:
188    case GCst{c}:
189      acc
190    case GVar{+i}:
191      acc_add(acc, i, g)
192    case GAdd{a, b}:
193      gbwd(b, env, g, gbwd(a, env, g, acc))
194    case GMul{+a, +b}:
195      gbwd(b, env, (g * gval(a, env) : F32), gbwd(a, env, (g * gval(b, env) : F32), acc))
196    case GRelu{+a}:
197      gbwd(a, env, (g * relu_d(gval(a, env)) : F32), acc)
198    case GTanh{+a}:
199      gbwd(a, env, (g * tanh_d(gval(a, env)) : F32), acc)
200    case GExp{+a}:
201      gbwd(a, env, (g * F32.exp(gval(a, env)) : F32), acc)
202
203# gradient of the expression with respect to each variable of env
204def ggrad(e: G, +env: List<&2, F32>) -> List<&2, F32>:
205  gbwd(e, env, 1.0, T.fill.l(List.length(&2, F32, env), 0.0))
206
207# =====================================================================
208# 3) Layers with tensors: the shape of each gradient is imposed by the TYPE
209# =====================================================================
210
211# y = x·W + b           x: n x i,  W: i x o,  b: o   ->   y: n x o
212def linear(-n: Nat, +i: Nat, +o: Nat, x: T.Mat<n, i>, w: T.Mat<i, o>, b: T.Vec<o>) -> T.Mat<n, o>:
213  T.Mat.add_row(n, o, T.Mat.matmul(n, i, o, x, w), b)
214
215# Given the gradient dy arriving at y, returns (dx, (dW, db)):
216#   dx = dy · Wᵀ   (n x o)·(o x i) = n x i
217#   dW = xᵀ · dy   (i x n)·(n x o) = i x o
218#   db = sum of the rows of dy
219# If any of these products had the wrong dimensions, the program would not compile.
220def linear_bwd(+n: Nat, +i: Nat, +o: Nat, x: T.Mat<n, i>, w: T.Mat<i, o>, +dy: T.Mat<n, o>) -> T.Mat<n, i> & (T.Mat<i, o> & T.Vec<o>):
221  (T.Mat.matmul(n, o, i, dy, T.Mat.transpose(i, o, w)), (T.Mat.matmul(i, n, o, T.Mat.transpose(n, i, x), dy), T.Mat.col_sums(n, o, dy)))
222
223def mask.l(xs: List<&2, F32>) -> List<&2, F32>:
224  match xs:
225    case Nil{}:
226      Nil{}
227    case Con{h, t}:
228      relu_d(h) <> mask.l(t)
229
230def mask.rows(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
231  match xs:
232    case Nil{}:
233      Nil{}
234    case Con{r, t}:
235      mask.l(r) <> mask.rows(t)
236
237# y = relu(x); the gradient passes where x > 0
238def relu_bwd(-r: Nat, -c: Nat, x: T.Mat<r, c>, dy: T.Mat<r, c>) -> T.Mat<r, c>:
239  match x:
240    case T.Mat{rows}:
241      T.Mat.mul(r, c, dy, T.Mat{mask.rows(rows)})
242
243# one-hot: n rows of c numbers, with 1 at the label's position
244def onehot.cell(hit: Bool) -> F32:
245  mask2(hit)
246
247def onehot.row(c: Nat, +k: Nat, +j: Nat) -> List<&2, F32>:
248  match c:
249    case 0n:
250      Nil{}
251    case 1n+p:
252      onehot.cell(Nat.is_eq(k, j)) <> onehot.row(p, k, 1n+j)
253
254def onehot.rows(labels: List<&2, Nat>, +c: Nat) -> List<&2, List<&2, F32>>:
255  match labels:
256    case Nil{}:
257      Nil{}
258    case Con{k, t}:
259      onehot.row(c, k, 0n) <> onehot.rows(t, c)
260
261def log.l(xs: List<&2, F32>) -> List<&2, F32>:
262  match xs:
263    case Nil{}:
264      Nil{}
265    case Con{+h, t}:
266      F32.log(F32.max(h, 0.000000000001)) <> log.l(t)
267
268def log.rows(xs: List<&2, List<&2, F32>>) -> List<&2, List<&2, F32>>:
269  match xs:
270    case Nil{}:
271      Nil{}
272    case Con{r, t}:
273      log.l(r) <> log.rows(t)
274
275def rows.total(xs: List<&2, List<&2, F32>>, acc: F32) -> F32:
276  match xs:
277    case Nil{}:
278      acc
279    case Con{r, t}:
280      rows.total(t, (acc + T.sum.go(r, 0.0) : F32))
281
282def ce_loss2(-n: Nat, -c: Nat, +cc: Nat, +nf: F32, p: T.Mat<n, c>, labels: List<&2, Nat>) -> F32:
283  match p:
284    case T.Mat{rows}:
285      (0.0 - (rows.total(T.rows.zip_mul(log.rows(rows), onehot.rows(labels, cc)), 0.0) / nf : F32) : F32)
286
287# mean cross-entropy loss: -(1/n) Σ log p[i][label_i]
288def ce_loss(+n: Nat, +c: Nat, logits: T.Mat<n, c>, labels: List<&2, Nat>) -> F32:
289  ce_loss2(n, c, c, F32.from_nat(n), T.Mat.softmax(n, c, logits), labels)
290
291# gradient of the loss with respect to the logits: (softmax - one-hot) / n
292def ce_grad(+n: Nat, +c: Nat, logits: T.Mat<n, c>, labels: List<&2, Nat>) -> T.Mat<n, c>:
293  T.Mat.scale(n, c, (1.0 / F32.from_nat(n) : F32), T.Mat.sub(n, c, T.Mat.softmax(n, c, logits), T.Mat{onehot.rows(labels, c)}))
294
295def sgd(-r: Nat, -c: Nat, +lr: F32, w: T.Mat<r, c>, dw: T.Mat<r, c>) -> T.Mat<r, c>:
296  T.Mat.sub(r, c, w, T.Mat.scale(r, c, lr, dw))
297
298def sgd_vec(-n: Nat, +lr: F32, b: T.Vec<n>, db: T.Vec<n>) -> T.Vec<n>:
299  T.Vec.sub(n, b, T.Vec.scale(n, lr, db))