Van der Corput progressive expert scheduler / vdc_moe_experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  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()