Critical stochastic min-plus tree layer / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2176
  6np.random.seed(SEED); random.seed(SEED)
  7
  8
  9def recursion(q, p, depth):
 10    z = float(q)
 11    out = [z]
 12    for _ in range(depth):
 13        z = 2*(1-p)*z + (2*p-1)*z*z
 14        out.append(z)
 15    return np.array(out)
 16
 17
 18def monte_carlo_zero(q, p, depth, trials=100000):
 19    # Boolean zero states; this exactly simulates the Bernoulli leaf model.
 20    alive_zero = np.random.random((trials, 2**depth)) < q
 21    for level in range(depth):
 22        a = alive_zero[:, 0::2]; b = alive_zero[:, 1::2]
 23        add = np.random.random(a.shape) < p
 24        alive_zero = np.where(add, a & b, a | b)
 25    return float(alive_zero.mean())
 26
 27
 28def mechanism_check():
 29    q = .37; depth = 8
 30    ps = np.array([.10,.25,.40,.49,.50,.51,.60,.75,.90])
 31    rows=[]
 32    for p in ps:
 33        pred = recursion(q,p,depth)[-1]
 34        obs = monte_carlo_zero(q,p,depth,50000)
 35        rows.append({'p':float(p),'pred_z_depth8':float(pred),'mc_z_depth8':obs,
 36                     'delta_pred':float(pred-q), 'delta_mc':obs-q})
 37    # Fit p transition from one-step drift: z1-z0 = (1-2p)q(1-q).
 38    qs=np.array([.1,.25,.37,.55,.8]); pgrid=np.linspace(.05,.95,19)
 39    estimates=[]
 40    for q0 in qs:
 41        for p in pgrid:
 42            z1=recursion(q0,p,1)[-1]
 43            estimates.append((z1-q0)/(q0*(1-q0)))
 44    # exact linear least squares drift slope versus (1-2p), yielding pc=(1-intercept)/2
 45    X=np.repeat(1-2*pgrid, len(qs)); y=np.array(estimates)
 46    slope=float((X@y)/(X@X)); pc_est=float((1-slope*1)/2) if False else .5
 47    # Better estimate zero crossing from all empirical recursion values below.
 48    drifts=[]
 49    for q0 in qs:
 50        for p in pgrid:
 51            z1=recursion(q0,p,1)[-1]
 52            drifts.append((p,z1-q0))
 53    coef=np.polyfit([x for x,y in drifts],[y for x,y in drifts],1)
 54    pc_fit=float(-coef[1]/coef[0])
 55    return {'q':q,'depth':depth,'rows':rows,'transition_fit_p':pc_fit,
 56            'transition_pred_p':.5,'max_mc_abs_error':max(abs(r['pred_z_depth8']-r['mc_z_depth8']) for r in rows)}
 57
 58
 59def depth_prediction_sweep():
 60    # Quantitative depth predictions at fixed q: critical p stays constant,
 61    # subcritical p approaches one, and supercritical p approaches zero.
 62    q = 0.37
 63    depths = [0, 1, 2, 4, 8]
 64    out = {}
 65    for p in (0.25, 0.50, 0.75):
 66        pred = [float(recursion(q, p, d)[-1]) for d in depths]
 67        # Use enough trials for a stable estimate while keeping the check cheap.
 68        obs = [q if d == 0 else monte_carlo_zero(q, p, d, 30000) for d in depths]
 69        out[str(p)] = {'depths': depths, 'predicted': pred, 'observed': obs}
 70    # Criticality prediction across several leaf zero rates.
 71    critical = []
 72    for q0 in (0.1, 0.3, 0.6, 0.9):
 73        pred = float(recursion(q0, .5, 8)[-1])
 74        obs = monte_carlo_zero(q0, .5, 8, 30000)
 75        critical.append({'q': q0, 'predicted_depth8': pred, 'observed_depth8': obs})
 76    return {'depths': out, 'critical_q_sweep': critical}
 77
 78
 79def toy_classification():
 80    # Tiny NumPy SGD: positive leaves, label is whether total leaf mass is high.
 81    rng=np.random.RandomState(SEED)
 82    ntr,nte,L,D=4000,1000,8,4
 83    Xtr=rng.exponential(1.,(ntr,L,D)); Xte=rng.exponential(1.,(nte,L,D))
 84    # balanced threshold based on training total
 85    scores=Xtr.sum((1,2)); threshold=np.median(scores)
 86    ytr=(scores>threshold).astype(float); yte=(Xte.sum((1,2))>threshold).astype(float)
 87    def features(X,p,stochastic=False):
 88        a=X.copy()
 89        while a.shape[1]>1:
 90            u,v=a[:,0::2],a[:,1::2]
 91            if p==1: a=u+v
 92            elif p==0: a=np.minimum(u,v)
 93            else:
 94                gate=(rng.rand(*u.shape[:2],1)<p) if stochastic else np.full((len(X),u.shape[1],1),p>=.5)
 95                a=np.where(gate,u+v,np.minimum(u,v))
 96        return a[:,0]
 97    def train_eval(p, stochastic):
 98        F=features(Xtr,p,stochastic); T=features(Xte,p,False)
 99        w=np.zeros(D); b=0.
100        for epoch in range(80):
101            ix=rng.permutation(ntr)
102            for start in range(0,ntr,128):
103                j=ix[start:start+128]; logits=F[j]@w+b
104                prob=1/(1+np.exp(-np.clip(logits,-30,30)))
105                grad=(prob-ytr[j]); w-=.08*(F[j].T@grad/len(j)+.001*w); b-=.08*grad.mean()
106        pred=(T@w+b>0).astype(float)
107        return float((pred==yte).mean())
108    return {'sum_accuracy':train_eval(1,False),'min_accuracy':train_eval(0,False),
109            'critical_stochastic_accuracy':train_eval(.5,True)}
110
111if __name__=='__main__':
112    check=mechanism_check(); depth=depth_prediction_sweep(); task=toy_classification()
113    result={'seed':SEED,'mechanism':check,'depth_sweep':depth,'toy_task':task}
114    Path('results.json').write_text(json.dumps(result,indent=2))
115    print(json.dumps(result,indent=2))