main.bend
bend-ml-autograd@0.1.0.0 · 299 lines · 0 definitions
1# bend-ml-autograd: diferenciação automática com uma lei provada.
2#
3# import bend-ml-autograd@0.1.0.0/main.bend as AG
4#
5# 1) MODELO PROVADO (sobre Nat): o modo reverso do autodiff dá o mesmo resultado
6# que o modo direto (números duais). Prova pelo kernel, usando os lemas de
7# soma e produto do bend-ml-nat-lemmas.
8# 2) AUTOGRAD DE ESCALARES em F32 (estilo micrograd), com a mesma estrutura.
9# 3) CAMADAS COM TENSORES (forward e backward) cujo formato é checado pelo tipo.
10#
11# F32 não é um número real, então a correção NUMÉRICA do F32 é validada por
12# testes contra o PyTorch (reference/test_autograd.py), não por prova.
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) Modelo provado
20# =====================================================================
21
22# Expressões com uma variável NX, sobre Nat (um semianel comutativo).
23type NE is Data:
24 NCst{n: Nat}
25 NX{}
26 NAdd{a: NE, b: NE}
27 NMul{a: NE, b: NE}
28
29# valor da expressão em 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# derivada em modo DIRETO (números duais): derivada da soma e regra do produto
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# derivada em modo REVERSO: g é o gradiente que chega da saída; cada nó repassa
54# g (soma) ou g * valor-do-outro-fator (produto) para os filhos e soma o que
55# chega em 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 (distributividade à esquerda; o nat-lemmas só tem à direita)
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# Caso do produto, com tudo abstraído em números:
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# O modo reverso devolve g vezes a derivada do modo direto.
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# LEI (reverse_eq_forward): para qualquer expressão feita de constantes, X, somas
111# e produtos, o gradiente em modo reverso (começando com gradiente 1 na saída)
112# é igual à derivada em modo direto. Vale para todo 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) Autograd de escalares em F32 (reverse-mode, estilo micrograd)
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# soma v ao gradiente da variável 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# derivada local da relu: 1 se a entrada > 0, senão 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 é o gradiente que chega ao nó; acc acumula o gradiente de cada variável
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# gradiente da expressão em relação a cada variável de 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) Camadas com tensores: o formato de cada gradiente é imposto pelo TIPO
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# Dado o gradiente dy que chega em y, devolve (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 = soma das linhas de dy
219# Se alguma conta tivesse as dimensões erradas, o programa não compilaria.
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); o gradiente passa onde 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 linhas de c números, com 1 na posição do rótulo
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# perda de entropia cruzada média: -(1/n) Σ log p[i][rótulo_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# gradiente da perda em relação aos 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))