sketch_thompson.nx source
↩ module page · 278 lines · 9193 B
1// sketch_thompson.nx -- Thompson Sampling (Bernoulli-Beta bandit).
2//
3// Thompson 1933 / Russo-Van Roy 2018 modern treatment. Bayesian
4// alternative to UCB1 (sketch_ucb1.nx): instead of a confidence-
5// bound upper-bonus, sample from each arm's posterior and pick the
6// argmax. Empirically tighter regret on stochastic Bernoulli arms;
7// matches UCB asymptotic bounds.
8//
9// CORE OPERATION:
10// For arm i with s_i successes, f_i failures:
11// posterior ~ Beta(s_i + 1, f_i + 1) (uniform Beta(1,1) prior)
12// Per decision:
13// for each arm i: sample x_i ~ Beta(s_i+1, f_i+1)
14// pull arg max x_i
15//
16// SAMPLING WITHOUT f64 (the technical chokepoint):
17// Beta(α, β) for integer α, β equals the distribution of the α-th
18// smallest of α + β − 1 uniforms in [0, 1]. (David-Nagaraja
19// "Order Statistics", standard textbook result.) Exact, no log
20// or special functions needed. Cost O(α + β) per sample.
21//
22// HYBRID FOR LARGE (s + f):
23// - exact order statistic when s + f + 1 ≤ NX_THOMPSON_EXACT_MAX
24// - Normal approximation when larger (cost O(1) instead of O(n)).
25// Beta(α, β) ≈ N(μ, σ²) with μ = α/(α+β), σ² = μ(1−μ)/(α+β+1).
26// Normal sampled via Irwin-Hall sum of 12 uniforms (CLT).
27//
28// Both paths produce Q14 fixed-point samples in [0, Q14] = [0, 1].
29//
30// EXTENDS NISHI-SUBSTRATE BEYOND DATASKETCHES (DS is summarization-
31// only; bandits live elsewhere). Composes against sketch_ucb1 +
32// sketch_epsilon_greedy as the bandit-family triple.
33//
34// LOSSLESS-LANGUAGE DISCIPLINE: nx_thompson_query returns
35// ApproxI64 with NX_ENV_REL_STDDEV = 1/sqrt(count_i + 1), matching
36// posterior shrinkage as samples accumulate.
37
38import "syscalls.nx"
39import "sketch_types.nx"
40import "nx_vecmath.nx"
41
42const NX_THOMPSON_ARMS_MIN: i64 = 2
43const NX_THOMPSON_ARMS_MAX: i64 = 256
44const NX_THOMPSON_Q14: i64 = 16384
45const NX_THOMPSON_EXACT_MAX: i64 = 64 // cutoff for order-stat sampler
46
47// LCG (matches KLL / Reservoir family).
48const NX_THOMPSON_LCG_A: i64 = 1103515245
49const NX_THOMPSON_LCG_C: i64 = 12345
50const NX_THOMPSON_LCG_MOD: i64 = 0x7FFFFFFF
51
52struct Thompson {
53 n_arms: i64,
54 successes: *i64,
55 failures: *i64,
56 total_pulls: i64,
57 rng_state: i64,
58 scratch: *i64, // size NX_THOMPSON_EXACT_MAX
59}
60
61// === construction =================================================
62
63func nx_thompson_alloc(n_arms: i64, seed: i64) -> *Thompson {
64 if n_arms < NX_THOMPSON_ARMS_MIN { return 0 as *Thompson }
65 if n_arms > NX_THOMPSON_ARMS_MAX { return 0 as *Thompson }
66 let raw: *u8 = sys_mmap(56)
67 let t: *Thompson = raw as *Thompson
68 let s_raw: *u8 = sys_mmap(n_arms * 8)
69 t.successes = s_raw as *i64
70 let f_raw: *u8 = sys_mmap(n_arms * 8)
71 t.failures = f_raw as *i64
72 var i: i64 = 0
73 while i < n_arms {
74 t.successes[i] = 0
75 t.failures[i] = 0
76 i = i + 1
77 }
78 t.n_arms = n_arms
79 t.total_pulls = 0
80 t.rng_state = seed | 1
81 let scratch_raw: *u8 = sys_mmap(NX_THOMPSON_EXACT_MAX * 8)
82 t.scratch = scratch_raw as *i64
83 return t
84}
85
86// === LCG + uniform Q14 ============================================
87
88func nx_thompson_rng_next(t: *Thompson) -> i64 {
89 let next: i64 = ((t.rng_state * NX_THOMPSON_LCG_A) + NX_THOMPSON_LCG_C) & NX_THOMPSON_LCG_MOD
90 t.rng_state = next
91 return next
92}
93
94// Uniform in [0, Q14) = 14 low bits of LCG output.
95func nx_thompson_uniform_q14(t: *Thompson) -> i64 {
96 let r: i64 = nx_thompson_rng_next(t)
97 return r & (NX_THOMPSON_Q14 - 1)
98}
99
100// === isqrt (shared shape with stream_stats / ucb1) ================
101
102func nx_thompson_isqrt(x: i64) -> i64 { return vm_isqrt(x) }
103
104// === order-statistic Beta sampler =================================
105//
106// Sample Beta(alpha, beta) for integer alpha, beta with
107// (alpha + beta - 1) ≤ NX_THOMPSON_EXACT_MAX. Returns Q14 sample.
108
109func nx_thompson_sort_scratch(t: *Thompson, n: i64) -> i64 {
110 var i: i64 = 1
111 while i < n {
112 let cur: i64 = t.scratch[i]
113 var j: i64 = i - 1
114 var done: i64 = 0
115 while done == 0 {
116 if j < 0 { done = 1 }
117 if done == 0 {
118 if t.scratch[j] <= cur { done = 1 }
119 if done == 0 {
120 t.scratch[j + 1] = t.scratch[j]
121 j = j - 1
122 }
123 }
124 }
125 t.scratch[j + 1] = cur
126 i = i + 1
127 }
128 return 0
129}
130
131func nx_thompson_sample_order_stat(t: *Thompson, alpha: i64, n_uniform: i64) -> i64 {
132 var i: i64 = 0
133 while i < n_uniform {
134 t.scratch[i] = nx_thompson_uniform_q14(t)
135 i = i + 1
136 }
137 nx_thompson_sort_scratch(t, n_uniform)
138 // alpha-th smallest = index (alpha - 1) zero-based
139 return t.scratch[alpha - 1]
140}
141
142// === Irwin-Hall normal approximation ==============================
143//
144// Z ≈ (U_1 + ... + U_12) − 6·Q14 where U_i ~ Uniform[0, Q14)
145// Mean 0, variance ≈ Q14² in the high-resolution limit. Returns
146// scaled to match a standard normal in Q14.
147
148func nx_thompson_clt_normal_q14(t: *Thompson) -> i64 {
149 var sum: i64 = 0
150 var i: i64 = 0
151 while i < 12 {
152 sum = sum + nx_thompson_uniform_q14(t)
153 i = i + 1
154 }
155 return sum - 6 * NX_THOMPSON_Q14
156}
157
158// === main sampler =================================================
159//
160// Returns one sample from Beta(s + 1, f + 1) in Q14.
161
162func nx_thompson_sample_q14(t: *Thompson, arm: i64) -> i64 {
163 let s: i64 = t.successes[arm]
164 let f: i64 = t.failures[arm]
165 let alpha: i64 = s + 1
166 let beta: i64 = f + 1
167 let n_uniform: i64 = alpha + beta - 1
168 if n_uniform <= NX_THOMPSON_EXACT_MAX {
169 return nx_thompson_sample_order_stat(t, alpha, n_uniform)
170 }
171 // Normal approximation path.
172 let total: i64 = alpha + beta
173 let mu_q14: i64 = (alpha * NX_THOMPSON_Q14) / total
174 // var_q28 = mu_q14 * (Q14 - mu_q14) / (total + 1)
175 // Q14 · Q14 = Q28, then divide by integer keeps Q28.
176 let var_num: i64 = mu_q14 * (NX_THOMPSON_Q14 - mu_q14)
177 let var_q28: i64 = var_num / (total + 1)
178 // sigma_q14 = isqrt(var_q28) -- since sqrt(Q28) = Q14
179 let sigma_q14: i64 = nx_thompson_isqrt(var_q28)
180 let z_q14: i64 = nx_thompson_clt_normal_q14(t)
181 var sample_q14: i64 = mu_q14 + (sigma_q14 * z_q14) / NX_THOMPSON_Q14
182 if sample_q14 < 0 { sample_q14 = 0 }
183 if sample_q14 > NX_THOMPSON_Q14 { sample_q14 = NX_THOMPSON_Q14 }
184 return sample_q14
185}
186
187// === select =======================================================
188//
189// Sample each arm's posterior, return argmax.
190
191func nx_thompson_select(t: *Thompson) -> i64 {
192 if t.n_arms == 0 { return -1 }
193 var best: i64 = 0
194 var best_sample: i64 = nx_thompson_sample_q14(t, 0)
195 var i: i64 = 1
196 while i < t.n_arms {
197 let s: i64 = nx_thompson_sample_q14(t, i)
198 if s > best_sample {
199 best = i
200 best_sample = s
201 }
202 i = i + 1
203 }
204 return best
205}
206
207// === update =======================================================
208//
209// Reward in {0, 1}. 0 = failure, 1 = success. (Continuous-reward
210// variants will need a separate primitive -- this one is the
211// canonical Bernoulli case.)
212
213func nx_thompson_update(t: *Thompson, arm: i64, reward: i64) -> i64 {
214 if arm < 0 { return -1 }
215 if arm >= t.n_arms { return -1 }
216 if reward < 0 { return -1 }
217 if reward > 1 { return -1 }
218 if reward == 1 {
219 t.successes[arm] = t.successes[arm] + 1
220 }
221 if reward == 0 {
222 t.failures[arm] = t.failures[arm] + 1
223 }
224 t.total_pulls = t.total_pulls + 1
225 return 0
226}
227
228// === query / introspection ========================================
229
230func nx_thompson_successes_for(t: *Thompson, arm: i64) -> i64 {
231 return t.successes[arm]
232}
233
234func nx_thompson_failures_for(t: *Thompson, arm: i64) -> i64 {
235 return t.failures[arm]
236}
237
238// Posterior mean estimator: (s + 1) / (s + f + 2) in PPM.
239func nx_thompson_posterior_mean_ppm(t: *Thompson, arm: i64) -> i64 {
240 let s: i64 = t.successes[arm]
241 let f: i64 = t.failures[arm]
242 let alpha: i64 = s + 1
243 let total: i64 = s + f + 2
244 return (alpha * 1000000) / total
245}
246
247func nx_thompson_best_arm(t: *Thompson) -> i64 {
248 var best: i64 = -1
249 var best_mean: i64 = -1
250 var i: i64 = 0
251 while i < t.n_arms {
252 if t.successes[i] + t.failures[i] > 0 {
253 let m: i64 = nx_thompson_posterior_mean_ppm(t, i)
254 if m > best_mean {
255 best = i
256 best_mean = m
257 }
258 }
259 i = i + 1
260 }
261 return best
262}
263
264func nx_thompson_query(t: *Thompson, arm: i64) -> *ApproxI64 {
265 let m: i64 = nx_thompson_posterior_mean_ppm(t, arm)
266 let pulls: i64 = t.successes[arm] + t.failures[arm]
267 var stderr_ppb: i64 = 1000000000
268 let isq: i64 = nx_thompson_isqrt(pulls + 1)
269 if isq > 0 { stderr_ppb = 1000000000 / isq }
270 return nx_approx_new(m, NX_ENV_REL_STDDEV, stderr_ppb,
271 682700000,
272 NX_MATURITY_REFERENCE_IMPL,
273 NX_ADV_HONEST)
274}
275
276func nx_thompson_memory_bytes(t: *Thompson) -> i64 {
277 return 56 + t.n_arms * 16 + NX_THOMPSON_EXACT_MAX * 8
278}