main.bend
mlbend-generative@0.1.0.0 · 260 lines · 0 definitions
1# mlbend-generative: the generative skeleton, natively in Bend 2.
2#
3# import mlbend-generative@0.1.0.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# The builders recurse on the count left, the current index riding after
17# it, and the law's variables are ordered the same way — the structural
18# variable before the riding one — which is what lets the checker follow
19# the induction (NOTES.md, 2026-10-25). What is tested: every F32 number,
20# against the NumPy twin (reference/test_generative.py), and the
21# bijection of mix32: the inverse is closed form (the half-width xor-shift
22# is an involution, the 13-shift inverts by the triple xor, the multiplies
23# by their modular inverses mod 2^32) and generative/tests/ok.bend checks
24# it on a grid — a law would need bit-level equational reasoning the
25# kernel has not been asked to do, so the bijection is the debt of issue
26# #10, not a claim.
27
28import Base
29
30# ---- the counter RNG ----
31
32# murmur3's fmix32 finaliser, over U32 arithmetic that wraps
33def mix32(+z: U32) -> U32:
34 +a = U32.xor(z, U32.shrn(z, 16n))
35 +b = (a * 2246822519 : U32)
36 +c = U32.xor(b, U32.shrn(b, 13n))
37 +e = (c * 3266489917 : U32)
38 U32.xor(e, U32.shrn(e, 16n))
39
40# the 13-shift undone: xor the shifted value twice more
41def unxor13(+y: U32) -> U32:
42 a = U32.xor(y, U32.shrn(y, 13n))
43 U32.xor(a, U32.shrn(y, 26n))
44
45# the inverse of mix32, in closed form: the 16-shift is an involution on
46# 32 bits; the 13-shift inverts by the triple xor; the multiplies invert
47# by their modular inverses mod 2^32
48def mix32_inv(+z: U32) -> U32:
49 +d = U32.xor(z, U32.shrn(z, 16n))
50 +c = (d * 2828982549 : U32)
51 +b = unxor13(c)
52 +a = (b * 3066638151 : U32)
53 U32.xor(a, U32.shrn(a, 16n))
54
55# the hash of one counter: the step salt times the golden ratio, the
56# index added beside it
57def hash_of(+s: Nat, +i: Nat) -> U32:
58 +k = (U32.from_nat(s) * 2654435769 : U32)
59 (k + (U32.from_nat(i) + 305419898 : U32) : U32)
60
61# a uniform in (0, 1]: the low 24 bits of the hash plus one, exactly
62# representable, so log(u) is never log(0)
63def uniform(+s: Nat, +i: Nat) -> F32:
64 h = mix32(hash_of(s, i))
65 m = ((U32.and(h, U32.from_nat(16777215n)) + 1 : U32) : U32)
66 (F32.from_nat(U32.to_nat(m)) / 16777216.0 : F32)
67
68# one Gaussian of a counter, by parity: the even index of the pair brings
69# the radius, the odd one the angle
70def gauss_body(is_even: Bool, +r: F32, +ang: F32) -> F32:
71 match is_even:
72 case True{}:
73 (r * F32.cos(ang) : F32)
74 case False{}:
75 (r * F32.sin(ang) : F32)
76
77def gauss_at(+s: Nat, +j: Nat) -> F32:
78 +pair = Nat.sub(j, Nat.mod(j, 2n))
79 +ua = uniform(s, pair)
80 +ub = uniform(s, (1n + pair : Nat))
81 +r = F32.sqrt((0.0 - (2.0 * F32.log(ua) : F32) : F32))
82 gauss_body(Nat.is_eq(Nat.mod(j, 2n), 0n), r, (6.2831855 * ub : F32))
83
84# The noise of counter s, left numbers from index `from` on. The count
85# left is the recursion; the index rides after it.
86def noise_from(+s: Nat, +left: Nat, +from: Nat) -> List<&2, F32>:
87 match left:
88 case 0n:
89 Nil{}
90 case 1n+rest:
91 Con{gauss_at(s, from), noise_from(s, rest, (1n + from : Nat))}
92
93# the noise of counter s: m numbers, the index 0 first
94def noise_of(+s: Nat, +m: Nat) -> List<&2, F32>:
95 noise_from(s, m, 0n)
96
97# LAW 1 (noise_length): m numbers asked for are m numbers built.
98law noise_length:
99 for +s: Nat
100 for +m: Nat
101 for +from: Nat
102 {List.length(&2, F32, noise_from(s, m, from)) == m : Nat}
103
104# m numbers asked for are m numbers built
105def noise_length(s, m, from):
106 match m:
107 case 0n:
108 {==}
109 case 1n+rest:
110 %noise_length(s, rest, (1n + from : Nat)) :
111 {1n+List.length(&2, F32, noise_from(s, rest, (1n + from : Nat))) == 1n+_ : Nat}
112 {==}
113
114# ---- the cosine schedule ----
115#
116# ᾱ_t = f(t)/f(0) with f(t) = cos²(((t/T) + 0.008)/1.008 · π/2): a closed
117# form per step, so the table needs no running product and every entry is
118# reproducible on its own.
119
120def cos_f(+t: Nat, +steps: Nat) -> F32:
121 +x = (((F32.from_nat(t) / F32.from_nat(steps) : F32) + 0.008 : F32) / 1.008 : F32)
122 +c = F32.cos((1.5707964 * x : F32))
123 (c * c : F32)
124
125# ᾱ_t of the cosine schedule over `steps` steps (t = 1 .. steps)
126def alpha_at(+t: Nat, +steps: Nat) -> F32:
127 (cos_f(t, steps) / cos_f(0n, steps) : F32)
128
129# β_t, the ratio the reverse step removes: 1 - ᾱ_t/ᾱ_{t-1}, clamped at
130# 0.99 — the cosine schedule's last ratio is 1 in F32 (ᾱ_T ~ 1e-15), and
131# the reverse step's 1/sqrt(1-β) would amplify any prediction error
132def beta_at(+t: Nat, +steps: Nat) -> F32:
133 F32.min((1.0 - (alpha_at(t, steps) / alpha_at(Nat.sub(t, 1n), steps) : F32) : F32), 0.99)
134
135# The schedule from step `from` on, `left` entries. The count left is the
136# recursion; the step rides after it.
137def alphas_from(+steps: Nat, +left: Nat, +from: Nat) -> List<&2, F32>:
138 match left:
139 case 0n:
140 Nil{}
141 case 1n+rest:
142 Con{alpha_at(from, steps), alphas_from(steps, rest, (1n + from : Nat))}
143
144# ᾱ_1 .. ᾱ_steps, in order
145def alphas(+steps: Nat) -> List<&2, F32>:
146 alphas_from(steps, steps, 1n)
147
148# LAW 2 (alphas_length): a schedule of t entries is t entries.
149law alphas_length:
150 for +t: Nat
151 for +left: Nat
152 for +from: Nat
153 {List.length(&2, F32, alphas_from(t, left, from)) == left : Nat}
154
155# a schedule of t entries is t entries
156def alphas_length(t, left, from):
157 match left:
158 case 0n:
159 {==}
160 case 1n+rest:
161 %alphas_length(t, rest, (1n + from : Nat)) :
162 {1n+List.length(&2, F32, alphas_from(t, rest, (1n + from : Nat))) == 1n+_ : Nat}
163 {==}
164
165# ---- the linear schedule ----
166#
167# For few-step samplers (tens of steps, not thousands): a linear β is
168# gentle at the top (the cosine schedule's last ratios approach 1, and a
169# reverse step with β near 1 multiplies any prediction error by
170# 1/sqrt(1-β), so few-step cosine diverges). β runs from 0.0001 to 0.28,
171# which lands ᾱ_T near 1e-4 — pure noise — over 64 steps with every
172# reverse scale at most 1.18.
173
174# β_t of the linear schedule over `steps` steps (t = 1 .. steps)
175def lin_beta_at(+t: Nat, +steps: Nat) -> F32:
176 (0.0001 + ((0.2799 * (F32.from_nat(t) / F32.from_nat(steps) : F32)) : F32) : F32)
177
178# The schedule ᾱ_1 .. from `from` on, `left` entries, the running product
179# riding beside. The count left is the recursion; the step and the
180# product ride after it.
181def lin_alphas_from(+steps: Nat, +left: Nat, +from: Nat, +acc: F32) -> List<&2, F32>:
182 match left:
183 case 0n:
184 Nil{}
185 case 1n+rest:
186 +acc2 = (acc * (1.0 - lin_beta_at(from, steps) : F32) : F32)
187 Con{acc2, lin_alphas_from(steps, rest, (1n + from : Nat), acc2)}
188
189# ᾱ_1 .. ᾱ_steps of the linear schedule, in order
190def lin_alphas(+steps: Nat) -> List<&2, F32>:
191 lin_alphas_from(steps, steps, 1n, 1.0)
192
193# LAW 4 (lin_alphas_length): a linear schedule of t entries is t entries.
194law lin_alphas_length:
195 for +t: Nat
196 for +left: Nat
197 for +from: Nat
198 for +acc: F32
199 {List.length(&2, F32, lin_alphas_from(t, left, from, acc)) == left : Nat}
200
201# a linear schedule of t entries is t entries
202def lin_alphas_length(t, left, from, acc):
203 match left:
204 case 0n:
205 {==}
206 case 1n+rest:
207 %lin_alphas_length(t, rest, (1n + from : Nat),
208 (acc * (1.0 - lin_beta_at(from, t) : F32) : F32)) :
209 {1n+List.length(&2, F32, lin_alphas_from(t, rest, (1n + from : Nat),
210 (acc * (1.0 - lin_beta_at(from, t) : F32) : F32))) == 1n+_ : Nat}
211 {==}
212
213# ---- the sinusoidal time embedding ----
214
215# the frequency of rung p: a decade ladder 0.001 .. 1.0
216def freq_of(+p: Nat) -> F32:
217 match p:
218 case 0n:
219 0.001
220 case 1n:
221 0.01
222 case 2n:
223 0.1
224 case _:
225 1.0
226
227# The embedding of step t over `steps`, from rung `p` on, `left` rungs —
228# each rung a sine and a cosine. The rungs left are the recursion; the
229# rung index rides after them.
230def temb_from(+t: Nat, +steps: Nat, +left: Nat, +p: Nat) -> List<&2, F32>:
231 match left:
232 case 0n:
233 Nil{}
234 case 1n+rest:
235 +tau = (F32.from_nat(t) / F32.from_nat(steps) : F32)
236 +f = freq_of(p)
237 Con{F32.sin((tau * f : F32)), Con{F32.cos((tau * f : F32)),
238 temb_from(t, steps, rest, (1n + p : Nat))}}
239
240# the embedding of step t: k rungs, the lowest frequency first
241def temb_of(+t: Nat, +steps: Nat, +k: Nat) -> List<&2, F32>:
242 temb_from(t, steps, k, 0n)
243
244# LAW 3 (temb_length): an embedding of k rungs is 2k numbers.
245law temb_length:
246 for +t: Nat
247 for +steps: Nat
248 for +k: Nat
249 for +p: Nat
250 {List.length(&2, F32, temb_from(t, steps, k, p)) == Nat.mul(k, 2n) : Nat}
251
252# an embedding of k rungs is 2k numbers
253def temb_length(t, steps, k, p):
254 match k:
255 case 0n:
256 {==}
257 case 1n+rest:
258 %temb_length(t, steps, rest, (1n + p : Nat)) :
259 {1n+(1n+List.length(&2, F32, temb_from(t, steps, rest, (1n + p : Nat)))) == 1n+(1n+_) : Nat}
260 {==}