1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 620
  6random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_num_threads(4)
  8DTYPE = torch.float64
  9
 10
 11def gm_utility(load, C, eps=1e-12):
 12    # Stable log-space implementation of prod_j (c_j dot load)^(1/d).
 13    factors = load @ C.T
 14    return torch.exp(torch.log(factors.clamp_min(eps)).mean())
 15
 16
 17def concavity_check():
 18    E, d = 5, 5
 19    C = torch.full((d, E), 0.05, dtype=DTYPE)
 20    C += torch.eye(E, dtype=DTYPE)
 21    worst = 1.0
 22    jensen_gaps = []
 23    for _ in range(1000):
 24        x = torch.rand(E, dtype=DTYPE) + .02
 25        y = torch.rand(E, dtype=DTYPE) + .02
 26        t = float(torch.rand(1))
 27        lhs = gm_utility(t*x + (1-t)*y, C).item()
 28        rhs = t*gm_utility(x, C).item() + (1-t)*gm_utility(y, C).item()
 29        gap = rhs-lhs
 30        jensen_gaps.append(gap)
 31    # Hessian eigenvalues at an interior point, with unconstrained positive loads.
 32    x = (torch.rand(E, dtype=DTYPE)+.2).requires_grad_(True)
 33    H = torch.autograd.functional.hessian(lambda z: gm_utility(z, C), x).detach().numpy()
 34    eig = np.linalg.eigvalsh((H+H.T)/2)
 35    return {"max_jensen_violation": float(max(jensen_gaps)),
 36            "mean_jensen_gap": float(np.mean(jensen_gaps)),
 37            "max_hessian_eigenvalue": float(eig.max()),
 38            "min_hessian_eigenvalue": float(eig.min())}
 39
 40
 41def optimize(method, seed=SEED, steps=500):
 42    torch.manual_seed(seed)
 43    T, E, K = 256, 8, 2
 44    # Two task groups have distinct preferred experts, creating a real
 45    # specialization-vs-balance tradeoff.
 46    group = torch.arange(T) % K
 47    target = torch.zeros(T, E, dtype=DTYPE)
 48    target[:, :] = 0.025/(E-1)
 49    target[group == 0, 0] = .975
 50    target[group == 1, 1] = .975
 51    # Router logits are per-token parameters: this isolates the auxiliary
 52    # balancing behavior from architecture and optimizer confounders.
 53    logits = (0.1*torch.randn(T, E, dtype=DTYPE)).requires_grad_(True)
 54    opt = torch.optim.Adam([logits], lr=.12)
 55    C = torch.full((E, E), .03, dtype=DTYPE) + torch.eye(E, dtype=DTYPE)
 56    history=[]
 57    for step in range(steps):
 58        opt.zero_grad()
 59        r = torch.softmax(logits, dim=1)
 60        load = r.mean(0)
 61        task = -(target * torch.log(r.clamp_min(1e-12))).sum(1).mean()
 62        entropy = -(r * torch.log(r.clamp_min(1e-12))).sum(1).mean()
 63        switch = E * (load*load).sum() # differentiable proxy for Switch load loss
 64        if method == "entropy":
 65            aux = -0.08*entropy
 66        elif method == "switch":
 67            aux = 0.08*switch
 68        elif method == "gm":
 69            aux = -0.55*torch.log(gm_utility(load, C).clamp_min(1e-12))
 70        elif method == "gm_symmetric":
 71            Cs = torch.ones((2,E), dtype=DTYPE)
 72            aux = -0.55*torch.log(gm_utility(load, Cs))
 73        loss = task + aux
 74        loss.backward(); opt.step()
 75        if step % 50 == 0 or step == steps-1:
 76            with torch.no_grad():
 77                hard = r.argmax(1)
 78                counts = torch.bincount(hard, minlength=E).double()/T
 79                cv = counts.std(unbiased=False)/(counts.mean()+1e-12)
 80                history.append({"step":step,"loss":float(loss),"task":float(task),
 81                                "entropy":float(entropy),"cv_soft":float(load.std(unbiased=False)/load.mean()),
 82                                "cv_hard":float(cv),"drop_rate_capacity_1":float(torch.clamp(counts*E-1,min=0).sum()/E)})
 83    return history[-1], history
 84
 85
 86def main():
 87    conc = concavity_check()
 88    results={}
 89    for method in ["entropy","switch","gm_symmetric","gm"]:
 90        final, hist=optimize(method)
 91        results[method]={"final":final,"history":hist}
 92    # Directly demonstrate the literal c=1 degeneracy.
 93    load=torch.tensor([.01,.05,.1,.12,.15,.17,.19,.21],dtype=DTYPE,requires_grad=True)
 94    C=torch.ones((2,8),dtype=DTYPE)
 95    u=gm_utility(load,C); u.backward()
 96    results["symmetric_gradient"]={"utility":float(u),"gradient_norm":float(load.grad.norm()),"gradient":load.grad.detach().tolist()}
 97    out={"seed":SEED,"concavity":conc,"results":results}
 98    with open("results.json","w") as f: json.dump(out,f,indent=2)
 99    print(json.dumps(out,indent=2))
100
101if __name__ == "__main__": main()