Van der Corput progressive expert scheduler / vdc_moe_experiment.py
Mechanism confirmed, baseline not beaten
1import json, math, random
2from pathlib import Path
3import numpy as np
4
5
6def vdc(m):
7 x, f = 0.0, 0.5
8 while m:
9 x += (m & 1) * f
10 m >>= 1
11 f *= 0.5
12 return x
13
14
15def sequence(D, n):
16 return np.array([min(D-1, int(D*vdc(m))) for m in range(n)], dtype=int)
17
18
19def dyadic_sets(D):
20 # Every node in the fixed balanced binary tree, represented by a leaf interval.
21 out = []
22 width = D
23 while width >= 1:
24 for lo in range(0, D, width):
25 out.append((lo, lo+width))
26 width //= 2
27 return out
28
29
30def endpoint_leaf_error(seq, D):
31 counts = np.bincount(seq, minlength=D)
32 n = len(seq)
33 return float(np.max(np.abs(counts - n / D)))
34
35
36def discrepancy(seq, D):
37 counts = np.zeros(D, dtype=int)
38 max_leaf = 0.0
39 max_nested = 0.0
40 arg = None
41 sets = dyadic_sets(D)
42 for n, x in enumerate(seq, 1):
43 counts[x] += 1
44 max_leaf = max(max_leaf, float(np.max(np.abs(counts - n/D))))
45 for lo, hi in sets:
46 e = abs(int(counts[lo:hi].sum()) - n*(hi-lo)/D)
47 if e > max_nested:
48 max_nested, arg = e, (n, lo, hi)
49 return max_leaf, max_nested, arg
50
51
52def random_discrepancy(D, n, trials=40, seed=7):
53 rng = np.random.default_rng(seed)
54 vals, nested = [], []
55 sets = dyadic_sets(D)
56 for _ in range(trials):
57 s = rng.integers(0, D, size=n)
58 hist = np.zeros((n, D), dtype=np.int16)
59 hist[np.arange(n), s] = 1
60 counts = np.cumsum(hist, axis=0)
61 vals.append(float(np.max(np.abs(counts[-1] - n / D))))
62 mx = 0.0
63 k = np.arange(1, n + 1)
64 for lo, hi in sets:
65 actual = counts[:, lo:hi].sum(axis=1)
66 target = k * (hi - lo) / D
67 mx = max(mx, float(np.max(np.abs(actual - target))))
68 nested.append(mx)
69 return float(np.mean(vals)), float(np.mean(nested))
70
71
72def assignment_trial(D=8, tokens=128, capacity_factor=1.0, seed=11):
73 # Synthetic learned router preferences: 70% prefer expert 0, remaining are uniform.
74 rng = np.random.default_rng(seed)
75 cap = int(math.ceil(tokens / D * capacity_factor))
76 pref = np.where(rng.random(tokens) < .70, 0, rng.integers(0,D,size=tokens))
77 schedules = {
78 'random': rng.integers(0,D,size=tokens),
79 'greedy': pref.copy(),
80 'vdc': sequence(D,tokens),
81 }
82 result = {}
83 for name, proposed in schedules.items():
84 used = np.zeros(D, dtype=int); accepted = 0
85 # Assignment takes proposed candidate; capacity overflow is dropped.
86 for x in proposed:
87 if used[x] < cap:
88 used[x] += 1; accepted += 1
89 result[name] = {'drop_rate': 1-accepted/tokens,
90 'load_var': float(np.var(used)), 'loads': used.tolist()}
91 return result
92
93
94def main():
95 np.set_printoptions(suppress=True)
96 # Prediction P1: first D points are a permutation, hence zero discrepancy at n=D.
97 # P2: nested discrepancy grows no faster than the supplied logarithmic budget B(D).
98 # P3: at every prefix, VDC leaf discrepancy is <= 1 (a stronger empirical prediction
99 # for this dyadic construction); random routing should grow with D and prefix length.
100 sweep=[]
101 for D in [2,4,8,16,32,64]:
102 n=4*D
103 s=sequence(D,n)
104 leaf,nested,arg=discrepancy(s,D)
105 B=math.log(D)/(3*math.log(2))+1
106 rd,rn=random_discrepancy(D,n)
107 sweep.append({'D':D,'vdc_leaf_max':leaf,'vdc_nested_max':nested,
108 'bound_B':B,'random_leaf_mean':rd,'random_nested_mean':rn,
109 'at_n_D_endpoint_error':endpoint_leaf_error(sequence(D,D),D)})
110 # Capacity sweep averages fixed seeds, same proposals for each method per trial.
111 caps=[]
112 for cf in [0.75,1.0,1.25,1.5]:
113 agg={k:[] for k in ['random','greedy','vdc']}
114 for seed in range(20):
115 r=assignment_trial(capacity_factor=cf,seed=seed)
116 for k in agg: agg[k].append(r[k])
117 caps.append({'capacity_factor':cf, **{k:{'drop_rate':float(np.mean([z['drop_rate'] for z in v])),
118 'load_var':float(np.mean([z['load_var'] for z in v]))} for k,v in agg.items()}})
119 out={'predictions':[
120 'P1: at n=D, VDC visits every leaf exactly once (max leaf error 0).',
121 'P2: max nested prefix error is below B(D)=log(D)/(3 log 2)+1.',
122 'P3: dyadic VDC leaf prefix error stays <=1, while random routing has larger mean imbalance.'
123 ],'sweep':sweep,'capacity_sweep':caps}
124 Path('results.json').write_text(json.dumps(out,indent=2))
125 print(json.dumps(out,indent=2))
126
127if __name__=='__main__': main()