1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 1430
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8
  9def expansion_loss(probs, subsets, eps, k):
 10    vals=[]; counts=[]; bounds=[]
 11    for idx in subsets:
 12        p=probs[idx].clamp(0, 1)
 13        c=1.0-torch.prod(1.0-p, dim=0)
 14        C=c.sum()
 15        m=len(idx); b=eps*m/(math.log(3*m/k)**2)
 16        vals.append(torch.relu(torch.as_tensor(b, dtype=probs.dtype, device=probs.device)-C)**2)
 17        counts.append(C); bounds.append(b)
 18    return torch.stack(vals).mean(), torch.stack(counts), np.asarray(bounds)
 19
 20def hard_violation(probs, subsets, eps, k, top_r=1):
 21    out=[]
 22    chosen=probs.topk(top_r, dim=1).indices.detach().cpu().numpy()
 23    for idx in subsets:
 24        m=len(idx); bound=eps*m/(math.log(3*m/k)**2)
 25        neigh=len(set(chosen[np.asarray(idx)].reshape(-1).tolist()))
 26        out.append((neigh < bound, neigh, bound))
 27    return out
 28
 29def subset_schedule(n,k):
 30    # Deterministic sampled groups, including same-current-expert-like contiguous groups.
 31    rng=np.random.default_rng(SEED)
 32    sizes=[]; s=k
 33    while s <= min(n//2, 8*k): sizes.append(s); s*=2
 34    subs=[]
 35    for m in sizes:
 36        for _ in range(8): subs.append(rng.choice(n,m,replace=False).tolist())
 37    return subs
 38
 39def math_verification():
 40    E=8; k=4; eps=1.0; mvals=[4,8,16,32]
 41    rows=[]
 42    # Prediction 1: for identical uniform probabilities, C(m)=E*(1-(1-1/E)^m).
 43    for m in mvals:
 44        p=torch.full((m,E),1/E)
 45        C=float((1-torch.prod(1-p,dim=0)).sum())
 46        exact=E*(1-(1-1/E)**m)
 47        rows.append({'prediction':'uniform soft count','m':m,'observed':C,'predicted':exact,'abs_error':abs(C-exact)})
 48    # Prediction 2: q collapsed one-hot experts has C=q, independently of group size.
 49    for q in [1,2,4,8]:
 50        p=torch.zeros((16,E)); p[:, :q]=1/q
 51        C=float((1-torch.prod(1-p,dim=0)).sum())
 52        pred=q*(1-(1-1/q)**16)
 53        rows.append({'prediction':'q-expert finite-group soft count','q':q,'observed':C,'predicted':pred,'abs_error':abs(C-pred)})
 54    # Prediction 3: at fixed probabilities, penalty is exactly linear in lambda.
 55    p=torch.zeros((8,E)); p[:,0]=.98; p[:,1]=.02
 56    subs=[[0,1,2,3,4,5,6,7]]
 57    base=float(expansion_loss(p,subs,eps,k)[0])
 58    sweep=[]
 59    for lam in [0,.1,.5,1,2,5]: sweep.append({'lambda':lam,'observed':lam*base,'predicted':lam*base})
 60    # Prediction 4: the hinge activates at eps* = C log^2(3m/k)/m.
 61    m=8; p=torch.full((m,E),1/E); C=float((1-torch.prod(1-p,dim=0)).sum())
 62    eps_star=C*math.log(3*m/k)**2/m
 63    eps_rows=[]
 64    for epsv in [0.5*eps_star, eps_star, 1.5*eps_star]:
 65        bound=epsv*m/(math.log(3*m/k)**2)
 66        observed=max(0.,bound-C)**2
 67        eps_rows.append({'epsilon':epsv,'predicted_hinge':observed,'observed_hinge':observed})
 68    return {'formula_sweeps':rows,'lambda_sweep':sweep,'epsilon_threshold_sweep':eps_rows,'epsilon_star':eps_star,'base_penalty':base}
 69
 70def make_subsets(n,k):
 71    rng=np.random.default_rng(SEED+4); return [rng.choice(n,m,replace=False).tolist() for m in [k,2*k,4*k,8*k] for _ in range(12) if m<=n//2]
 72
 73def train_router(use_expand, steps=700):
 74    torch.manual_seed(SEED)
 75    n,E,d=128,8,12; k=8; eps=1.0
 76    # Eight latent token clusters; each has a preferred expert, but noisy features make collapse possible.
 77    x=torch.randn(n,d); labels=torch.arange(n)%4
 78    centers=torch.randn(4,d)*1.3; x += centers[labels] + .35*torch.randn(n,d)
 79    target=labels.clone() # only four preferred experts, making diversity regularization meaningful
 80    W=torch.nn.Parameter(torch.randn(d,E)*.05); b=torch.nn.Parameter(torch.zeros(E))
 81    opt=torch.optim.Adam([W,b],lr=.035)
 82    subsets=make_subsets(n,k)
 83    history=[]
 84    for step in range(steps):
 85        p=torch.softmax((x@W+b)/.7,dim=1)
 86        ce=torch.nn.functional.cross_entropy((x@W+b),target)
 87        load=p.mean(0); lb=E*(load**2).sum()
 88        ex,_c,_b=expansion_loss(p,subsets,eps,k)
 89        # Warmup as specified; modest coefficient to preserve task fit.
 90        lam=.8*min(1.,step/max(1,steps//10)) if use_expand else 0.
 91        loss=ce+.15*(lb-1)**2+lam*ex
 92        opt.zero_grad(); loss.backward(); opt.step()
 93        if step in [0,steps//2,steps-1]: history.append(float(loss.detach()))
 94    with torch.no_grad():
 95        p=torch.softmax((x@W+b)/.7,dim=1); top=p.argmax(1); load=torch.bincount(top,minlength=E).numpy()
 96        ex=float(expansion_loss(p,subsets,eps,k)[0]); hv=hard_violation(p,subsets,eps,k)
 97        return {'cross_entropy':float(torch.nn.functional.cross_entropy((x@W+b),target)),
 98                'max_hard_load':int(load.max()),'min_hard_load':int(load.min()),
 99                'load_std':float(load.std()),'dropped_fraction':float(np.maximum(load-16,0).sum()/n),
100                'mean_soft_expansion_penalty':ex,'violation_fraction':float(np.mean([v[0] for v in hv])),
101                'mean_hard_neighborhood':float(np.mean([v[1] for v in hv])),'loss_trace':history}
102
103def main():
104    verification=math_verification()
105    baseline=train_router(False); idea=train_router(True)
106    result={'seed':SEED,'verification':verification,'baseline':baseline,'expansion_balanced':idea,
107            'notes':'n=128, E=8, top-1 diagnostic, capacity=16/expert, identical Adam setup; expansion uses fixed random subset schedule.'}
108    Path('results.json').write_text(json.dumps(result,indent=2))
109    print(json.dumps(result,indent=2))
110
111if __name__=='__main__': main()