main.bend

mlbend-generative@0.1.1.0 · 422 lines · 36 definitions

1# mlbend-generative: the generative skeleton, natively in Bend 2.
2#
3#   import mlbend-generative@0.1.1.0/main.bend as G
4#
5# What a generative model needs besides the network: randomness that a
6# verifier can draw again, a noise schedule, and a time embedding. The
7# randomness is a pure function of counters — murmur3's fmix32 finaliser
8# over (step salt · φ, index), Box-Muller over uniforms on a 24-bit grid —
9# so a NumPy twin draws the same numbers bit for bit. The schedule is the
10# cosine one (Nichol & Dhariwal): ᾱ_t = f(t)/f(0) in closed form, so every
11# entry is computed on its own and the table needs no running product.
12#
13# What is proved: the shape — the noise of any counter is exactly m
14# numbers (noise_length), a schedule over t steps is exactly t
15# (alphas_length), an embedding of k rungs is exactly 2k (temb_length),
16# a strided walk over n entries emits exactly the entries its count says
17# (strided_length), and a DDIM step over n elements is n elements
18# (ddim_length).
19# The builders recurse on the count left, the current index riding after
20# it, and the law's variables are ordered the same way — the structural
21# variable before the riding one — which is what lets the checker follow
22# the induction (NOTES.md, 2026-10-25). What is tested: every F32 number,
23# against the NumPy twin (reference/test_generative.py), and the
24# bijection of mix32: the inverse is closed form (the half-width xor-shift
25# is an involution, the 13-shift inverts by the triple xor, the multiplies
26# by their modular inverses mod 2^32) and generative/tests/ok.bend checks
27# it on a grid — a law would need bit-level equational reasoning the
28# kernel has not been asked to do, so the bijection is the debt of issue
29# #10, not a claim.
30
31import Base
32
33# ---- the counter RNG ----
34
35# murmur3's fmix32 finaliser, over U32 arithmetic that wraps
36def mix32(+z: U32) -> U32:
37  +a = U32.xor(z, U32.shrn(z, 16n))
38  +b = (a * 2246822519 : U32)
39  +c = U32.xor(b, U32.shrn(b, 13n))
40  +e = (c * 3266489917 : U32)
41  U32.xor(e, U32.shrn(e, 16n))
42
43# the 13-shift undone: xor the shifted value twice more
44def unxor13(+y: U32) -> U32:
45  a = U32.xor(y, U32.shrn(y, 13n))
46  U32.xor(a, U32.shrn(y, 26n))
47
48# the inverse of mix32, in closed form: the 16-shift is an involution on
49# 32 bits; the 13-shift inverts by the triple xor; the multiplies invert
50# by their modular inverses mod 2^32
51def mix32_inv(+z: U32) -> U32:
52  +d = U32.xor(z, U32.shrn(z, 16n))
53  +c = (d * 2828982549 : U32)
54  +b = unxor13(c)
55  +a = (b * 3066638151 : U32)
56  U32.xor(a, U32.shrn(a, 16n))
57
58# the hash of one counter: the step salt times the golden ratio, the
59# index added beside it
60def hash_of(+s: Nat, +i: Nat) -> U32:
61  +k = (U32.from_nat(s) * 2654435769 : U32)
62  (k + (U32.from_nat(i) + 305419898 : U32) : U32)
63
64# a uniform in (0, 1]: the low 24 bits of the hash plus one, exactly
65# representable, so log(u) is never log(0)
66def uniform(+s: Nat, +i: Nat) -> F32:
67  h = mix32(hash_of(s, i))
68  m = ((U32.and(h, U32.from_nat(16777215n)) + 1 : U32) : U32)
69  (F32.from_nat(U32.to_nat(m)) / 16777216.0 : F32)
70
71# one Gaussian of a counter, by parity: the even index of the pair brings
72# the radius, the odd one the angle
73def gauss_body(is_even: Bool, +r: F32, +ang: F32) -> F32:
74  match is_even:
75    case True{}:
76      (r * F32.cos(ang) : F32)
77    case False{}:
78      (r * F32.sin(ang) : F32)
79
80def gauss_at(+s: Nat, +j: Nat) -> F32:
81  +pair = Nat.sub(j, Nat.mod(j, 2n))
82  +ua = uniform(s, pair)
83  +ub = uniform(s, (1n + pair : Nat))
84  +r = F32.sqrt((0.0 - (2.0 * F32.log(ua) : F32) : F32))
85  gauss_body(Nat.is_eq(Nat.mod(j, 2n), 0n), r, (6.2831855 * ub : F32))
86
87# The noise of counter s, left numbers from index `from` on. The count
88# left is the recursion; the index rides after it.
89def noise_from(+s: Nat, +left: Nat, +from: Nat) -> List<&2, F32>:
90  match left:
91    case 0n:
92      Nil{}
93    case 1n+rest:
94      Con{gauss_at(s, from), noise_from(s, rest, (1n + from : Nat))}
95
96# the noise of counter s: m numbers, the index 0 first
97def noise_of(+s: Nat, +m: Nat) -> List<&2, F32>:
98  noise_from(s, m, 0n)
99
100# LAW 1 (noise_length): m numbers asked for are m numbers built.
101law noise_length:
102  for +s: Nat
103  for +m: Nat
104  for +from: Nat
105  {List.length(&2, F32, noise_from(s, m, from)) == m : Nat}
106
107# m numbers asked for are m numbers built
108def noise_length(s, m, from):
109  match m:
110    case 0n:
111      {==}
112    case 1n+rest:
113      %noise_length(s, rest, (1n + from : Nat)) :
114        {1n+List.length(&2, F32, noise_from(s, rest, (1n + from : Nat))) == 1n+_ : Nat}
115      {==}
116
117# ---- the cosine schedule ----
118#
119# ᾱ_t = f(t)/f(0) with f(t) = cos²(((t/T) + 0.008)/1.008 · π/2): a closed
120# form per step, so the table needs no running product and every entry is
121# reproducible on its own.
122
123def cos_f(+t: Nat, +steps: Nat) -> F32:
124  +x = (((F32.from_nat(t) / F32.from_nat(steps) : F32) + 0.008 : F32) / 1.008 : F32)
125  +c = F32.cos((1.5707964 * x : F32))
126  (c * c : F32)
127
128# ᾱ_t of the cosine schedule over `steps` steps (t = 1 .. steps)
129def alpha_at(+t: Nat, +steps: Nat) -> F32:
130  (cos_f(t, steps) / cos_f(0n, steps) : F32)
131
132# β_t, the ratio the reverse step removes: 1 - ᾱ_t/ᾱ_{t-1}, clamped at
133# 0.99 — the cosine schedule's last ratio is 1 in F32 (ᾱ_T ~ 1e-15), and
134# the reverse step's 1/sqrt(1-β) would amplify any prediction error
135def beta_at(+t: Nat, +steps: Nat) -> F32:
136  F32.min((1.0 - (alpha_at(t, steps) / alpha_at(Nat.sub(t, 1n), steps) : F32) : F32), 0.99)
137
138# The schedule from step `from` on, `left` entries. The count left is the
139# recursion; the step rides after it.
140def alphas_from(+steps: Nat, +left: Nat, +from: Nat) -> List<&2, F32>:
141  match left:
142    case 0n:
143      Nil{}
144    case 1n+rest:
145      Con{alpha_at(from, steps), alphas_from(steps, rest, (1n + from : Nat))}
146
147# ᾱ_1 .. ᾱ_steps, in order
148def alphas(+steps: Nat) -> List<&2, F32>:
149  alphas_from(steps, steps, 1n)
150
151# LAW 2 (alphas_length): a schedule of t entries is t entries.
152law alphas_length:
153  for +t: Nat
154  for +left: Nat
155  for +from: Nat
156  {List.length(&2, F32, alphas_from(t, left, from)) == left : Nat}
157
158# a schedule of t entries is t entries
159def alphas_length(t, left, from):
160  match left:
161    case 0n:
162      {==}
163    case 1n+rest:
164      %alphas_length(t, rest, (1n + from : Nat)) :
165        {1n+List.length(&2, F32, alphas_from(t, rest, (1n + from : Nat))) == 1n+_ : Nat}
166      {==}
167
168# ---- the linear schedule ----
169#
170# For few-step samplers (tens of steps, not thousands): a linear β is
171# gentle at the top (the cosine schedule's last ratios approach 1, and a
172# reverse step with β near 1 multiplies any prediction error by
173# 1/sqrt(1-β), so few-step cosine diverges). β runs from 0.0001 to 0.28,
174# which lands ᾱ_T near 1e-4 — pure noise — over 64 steps with every
175# reverse scale at most 1.18.
176
177# β_t of the linear schedule over `steps` steps (t = 1 .. steps)
178def lin_beta_at(+t: Nat, +steps: Nat) -> F32:
179  (0.0001 + ((0.2799 * (F32.from_nat(t) / F32.from_nat(steps) : F32)) : F32) : F32)
180
181# The schedule ᾱ_1 .. from `from` on, `left` entries, the running product
182# riding beside. The count left is the recursion; the step and the
183# product ride after it.
184def lin_alphas_from(+steps: Nat, +left: Nat, +from: Nat, +acc: F32) -> List<&2, F32>:
185  match left:
186    case 0n:
187      Nil{}
188    case 1n+rest:
189      +acc2 = (acc * (1.0 - lin_beta_at(from, steps) : F32) : F32)
190      Con{acc2, lin_alphas_from(steps, rest, (1n + from : Nat), acc2)}
191
192# ᾱ_1 .. ᾱ_steps of the linear schedule, in order
193def lin_alphas(+steps: Nat) -> List<&2, F32>:
194  lin_alphas_from(steps, steps, 1n, 1.0)
195
196# LAW 4 (lin_alphas_length): a linear schedule of t entries is t entries.
197law lin_alphas_length:
198  for +t: Nat
199  for +left: Nat
200  for +from: Nat
201  for +acc: F32
202  {List.length(&2, F32, lin_alphas_from(t, left, from, acc)) == left : Nat}
203
204# a linear schedule of t entries is t entries
205def lin_alphas_length(t, left, from, acc):
206  match left:
207    case 0n:
208      {==}
209    case 1n+rest:
210      %lin_alphas_length(t, rest, (1n + from : Nat),
211        (acc * (1.0 - lin_beta_at(from, t) : F32) : F32)) :
212        {1n+List.length(&2, F32, lin_alphas_from(t, rest, (1n + from : Nat),
213          (acc * (1.0 - lin_beta_at(from, t) : F32) : F32))) == 1n+_ : Nat}
214      {==}
215
216# ---- the sinusoidal time embedding ----
217
218# the frequency of rung p: a decade ladder 0.001 .. 1.0
219def freq_of(+p: Nat) -> F32:
220  match p:
221    case 0n:
222      0.001
223    case 1n:
224      0.01
225    case 2n:
226      0.1
227    case _:
228      1.0
229
230# The embedding of step t over `steps`, from rung `p` on, `left` rungs —
231# each rung a sine and a cosine. The rungs left are the recursion; the
232# rung index rides after them.
233def temb_from(+t: Nat, +steps: Nat, +left: Nat, +p: Nat) -> List<&2, F32>:
234  match left:
235    case 0n:
236      Nil{}
237    case 1n+rest:
238      +tau = (F32.from_nat(t) / F32.from_nat(steps) : F32)
239      +f = freq_of(p)
240      Con{F32.sin((tau * f : F32)), Con{F32.cos((tau * f : F32)),
241        temb_from(t, steps, rest, (1n + p : Nat))}}
242
243# the embedding of step t: k rungs, the lowest frequency first
244def temb_of(+t: Nat, +steps: Nat, +k: Nat) -> List<&2, F32>:
245  temb_from(t, steps, k, 0n)
246
247# LAW 3 (temb_length): an embedding of k rungs is 2k numbers.
248law temb_length:
249  for +t: Nat
250  for +steps: Nat
251  for +k: Nat
252  for +p: Nat
253  {List.length(&2, F32, temb_from(t, steps, k, p)) == Nat.mul(k, 2n) : Nat}
254
255# an embedding of k rungs is 2k numbers
256def temb_length(t, steps, k, p):
257  match k:
258    case 0n:
259      {==}
260    case 1n+rest:
261      %temb_length(t, steps, rest, (1n + p : Nat)) :
262        {1n+(1n+List.length(&2, F32, temb_from(t, steps, rest, (1n + p : Nat)))) == 1n+(1n+_) : Nat}
263      {==}
264
265# ---- the strided subsequence ----
266#
267# A few-step sampler (DDIM) visits every k-th rung of the schedule
268# instead of every rung. The walk is over the list itself: it can only
269# emit entries the list really has, so no index can fall off the end and
270# answer 0.0 — the silent out-of-range of `at`. `skip` says how many
271# entries drop before the first emission; a sampler anchors the top rung
272# with skip = (T-1) mod k, which lands the last emission on the last
273# entry whatever T and k are.
274
275# How many entries the walk emits over n entries. The count recurses on
276# the count left, exactly as the walk recurses on the list left, and the
277# two matches (on the count, on the skip) are in the same order — that
278# is what lets the law below follow both by induction.
279def groups_of(+n: Nat, +k: Nat, +skip: Nat) -> Nat:
280  match n:
281    case 0n:
282      0n
283    case 1n+m:
284      match skip:
285        case 0n:
286          (1n + groups_of(m, k, Nat.sub(k, 1n)) : Nat)
287        case 1n+r:
288          groups_of(m, k, r)
289
290# Every k-th entry of xs, the (skip+1)-th first: drop `skip`, emit one,
291# drop k-1, emit one, and so on to the end of the list. The list is
292# matched first (binders in order), the skip beside it.
293def strided_from(-A: Data, xs: List<&2, A>, +k: Nat, +skip: Nat) -> List<&2, A>:
294  match xs:
295    case Nil{}:
296      Nil{}
297    case Con{h, t}:
298      match skip:
299        case 0n:
300          Con{h, strided_from(A, t, k, Nat.sub(k, 1n))}
301        case 1n+r:
302          strided_from(A, t, k, r)
303
304# every k-th entry of xs, the k-th first
305def strided(-A: Data, xs: List<&2, A>, +k: Nat) -> List<&2, A>:
306  strided_from(A, xs, k, Nat.sub(k, 1n))
307
308# LAW 5 (strided_length): every k-th entry of n entries is exactly the
309# number the count says.
310law strided_length:
311  for -A: Data
312  for +xs: List<&2, A>
313  for +k: Nat
314  for +skip: Nat
315  {List.length(&2, A, strided_from(A, xs, k, skip)) == groups_of(List.length(&2, A, xs), k, skip) : Nat}
316
317# every k-th entry of n entries is exactly the number the count says
318def strided_length(A, xs, k, skip):
319  match xs:
320    case Nil{}:
321      {==}
322    case Con{h, t}:
323      match skip:
324        case 0n:
325          %strided_length(A, t, k, Nat.sub(k, 1n)) :
326            {1n+List.length(&2, A, strided_from(A, t, k, Nat.sub(k, 1n))) == 1n+_ : Nat}
327          {==}
328        case 1n+r:
329          %strided_length(A, t, k, r) :
330            {List.length(&2, A, strided_from(A, t, k, r)) == _ : Nat}
331          {==}
332
333# ---- the DDIM step (Song et al., 2020) ----
334#
335# The same forward model, a different sampler. From abar at the current
336# rung, abar_prev at the rung below and the noise level eta:
337#
338#   sigma  = eta·sqrt((1-abar_prev)/(1-abar))·sqrt(1-abar/abar_prev)
339#   x0     = (x - sqrt(1-abar)·eps)/sqrt(abar)
340#   x_prev = sqrt(abar_prev)·x0 + sqrt(1-abar_prev-sigma²)·eps + sigma·z
341#
342# eta = 0 never reads z, so the step is a pure function of the model and
343# the rungs; eta = 1 over every rung reproduces the ancestral DDPM step,
344# because that sigma² is exactly the posterior beta-tilde_t (the algebra
345# is checked against the ancestral mean in reference/test_ddim.py to
346# float32 rounding).
347
348def sigma_of(+abar: F32, +abar_prev: F32, +eta: F32) -> F32:
349  +above = F32.max(0.0, (1.0 - abar_prev : F32))
350  +below = F32.max(0.0, (1.0 - (abar / abar_prev : F32) : F32))
351  (eta * (F32.sqrt((above / (1.0 - abar : F32) : F32)) * F32.sqrt(below) : F32) : F32)
352
353# The five coefficients of one step: the four roots and the noise scale.
354# They are computed once per step, not once per element — the CIFAR
355# probe measures 786k elements in a step, and five sqrts each would be
356# paid 786k times. Every radicand is clamped at 0, so float rounding
357# cannot hand sqrt a negative and turn a whole grid into NaN.
358type Step is Data:
359  Step{c0: F32, c1: F32, c2: F32, c3: F32, sig: F32}
360
361def step_of(+abar: F32, +abar_prev: F32, +eta: F32) -> Step:
362  +sig = sigma_of(abar, abar_prev, eta)
363  +c0 = (1.0 / F32.sqrt(abar) : F32)
364  +c1 = F32.sqrt(F32.max(0.0, (1.0 - abar : F32)))
365  +c2 = F32.sqrt(F32.max(0.0, abar_prev))
366  +c3 = F32.sqrt(F32.max(0.0, (1.0 - (abar_prev + (sig * sig : F32) : F32) : F32)))
367  Step{c0, c1, c2, c3, sig}
368
369# one element of the step: x0 folded into x_prev
370def ddim_elem(+co: Step, +x: F32, +eps: F32, +z: F32) -> F32:
371  match co:
372    case Step{c0, c1, c2, c3, sig}:
373      +x0 = ((x - (c1 * eps : F32)) * c0 : F32)
374      (((x0 * c2 : F32) + (c3 * eps : F32) : F32) + (sig * z : F32) : F32)
375
376# the head of a list, or 0.0 when it is empty: the missing eps or z of a
377# step whose z is not drawn at all (eta = 0 passes Nil, and sigma·0.0 is
378# 0.0)
379def head_or_zero(xs: List<&2, F32>) -> F32:
380  match xs:
381    case Nil{}:
382      0.0
383    case Con{h, t}:
384      h
385
386# the rest of a list, or empty: the tail beside head_or_zero above
387def tail_or_nil(xs: List<&2, F32>) -> List<&2, F32>:
388  match xs:
389    case Nil{}:
390      Nil{}
391    case Con{h, t}:
392      t
393
394# One step over the whole sample: elementwise, and the shape follows x —
395# n elements in, n elements out (LAW 6). eps and z ride beside it: a
396# shorter list pads with 0.0, which is exactly right for the z of an
397# eta = 0 step and is caught by the twin's numbers otherwise.
398def ddim(+co: Step, x: List<&2, F32>, +eps: List<&2, F32>, +z: List<&2, F32>) -> List<&2, F32>:
399  match x:
400    case Nil{}:
401      Nil{}
402    case Con{h, t}:
403      Con{ddim_elem(co, h, head_or_zero(eps), head_or_zero(z)),
404        ddim(co, t, tail_or_nil(eps), tail_or_nil(z))}
405
406# LAW 6 (ddim_length): a step over n elements is n elements.
407law ddim_length:
408  for +co: Step
409  for +x: List<&2, F32>
410  for +eps: List<&2, F32>
411  for +z: List<&2, F32>
412  {List.length(&2, F32, ddim(co, x, eps, z)) == List.length(&2, F32, x) : Nat}
413
414# a step over n elements is n elements
415def ddim_length(co, x, eps, z):
416  match x:
417    case Nil{}:
418      {==}
419    case Con{h, t}:
420      %ddim_length(co, t, tail_or_nil(eps), tail_or_nil(z)) :
421        {1n+List.length(&2, F32, ddim(co, t, tail_or_nil(eps), tail_or_nil(z))) == 1n+_ : Nat}
422      {==}