1import json, time
  2import numpy as np
  3
  4
  5def project(Ubar, Vbar, tol=1e-11, max_iter=80):
  6    """Joint KL projection for unit row marginals and shared column marginal."""
  7    Ubar = np.asarray(Ubar, float); Vbar = np.asarray(Vbar, float)
  8    n, r = Ubar.shape; m = Vbar.shape[0]
  9    if np.any(Ubar <= 0) or np.any(Vbar <= 0): raise ValueError('references must be positive')
 10    z = np.zeros(r)
 11    def state(z):
 12        # subtract maxima only inside exponentials; row normalization is invariant
 13        a = np.log(Ubar) + z[None, :]
 14        b = np.log(Vbar) - z[None, :]
 15        a -= a.max(1, keepdims=True); b -= b.max(1, keepdims=True)
 16        pu = np.exp(a); pu /= pu.sum(1, keepdims=True)
 17        pv = np.exp(b); pv /= pv.sum(1, keepdims=True)
 18        U, V = pu, pv
 19        cu, cv = U.sum(0), V.sum(0)
 20        grad = cu-cv
 21        H = np.diag(cu)-U.T@U + np.diag(cv)-V.T@V
 22        return U,V,grad,H
 23    for it in range(max_iter):
 24        U,V,g,H = state(z)
 25        # gauge z[-1]=0; solve reduced Newton system
 26        if np.max(np.abs(g)) < tol: break
 27        Hr = H[:-1,:-1]; gr = g[:-1]
 28        try: step = np.linalg.solve(Hr, -gr)
 29        except np.linalg.LinAlgError: step = np.linalg.lstsq(Hr + 1e-10*np.eye(r-1), -gr, rcond=None)[0]
 30        old = np.sum(g*g)
 31        alpha = 1.0
 32        while alpha > 1e-10:
 33            zn = z.copy(); zn[:-1] += alpha*step; zn[-1] = 0
 34            gn = state(zn)[2]
 35            if np.sum(gn*gn) < old*(1-1e-4*alpha): z=zn; break
 36            alpha *= .5
 37        else: break
 38    U,V,g,H = state(z)
 39    return U,V,z,it+1,np.max(np.abs(g)),H
 40
 41
 42def hvp(U,V,s):
 43    """Covariance-sum HVP for unit row marginals."""
 44    def one(X):
 45        return (X*s).sum(0) * 0 + (X*s).sum(0) # unused, explicit below
 46    # sum_i [diag(p_i)s - p_i(p_i^T s)]
 47    out = (U*s).sum(0) - (U @ s) @ U
 48    out += (V*s).sum(0) - (V @ s) @ V
 49    return out
 50
 51
 52def dense_apply(U,V,X):
 53    g=U.sum(0)
 54    W=(U/g[None,:]) @ V.T
 55    return W@X, W
 56
 57def thin_apply(U,V,X):
 58    g=U.sum(0)
 59    return U @ ((V.T@X)/g[:,None])
 60
 61
 62def run():
 63    rng=np.random.default_rng(7)
 64    rows=[]
 65    # Prediction 1: exact projection residuals remain at numerical tolerance as n,r vary.
 66    for n in [16,32,64,128]:
 67      for r in [2,4,8]:
 68        if r>n: continue
 69        A=rng.normal(size=(n,r)); B=rng.normal(size=(n,r))
 70        Ub=np.exp(A); Vb=np.exp(B)
 71        U,V,z,its,res,H=project(Ub,Vb)
 72        rows.append({'n':n,'r':r,'residual':float(res),'iters':its,
 73          'row_res':float(max(abs(U.sum(1)-1).max(),abs(V.sum(1)-1).max())),
 74          'col_match':float(abs(U.sum(0)-V.sum(0)).max()),
 75          'h_null':float(np.linalg.norm(H@np.ones(r)) )})
 76    # Prediction 2: analytic covariance HVP agrees with finite differences.
 77    hv=[]
 78    for r in [2,4,8,12]:
 79      n=40; U,V,_,_,_,H=project(np.exp(rng.normal(size=(n,r))),np.exp(rng.normal(size=(n,r))))
 80      z=np.zeros(r); s=rng.normal(size=r); eps=1e-5
 81      # H is the exact analytic derivative in z coordinates
 82      rel=np.linalg.norm(hvp(U,V,s)-H@s)/(np.linalg.norm(H@s)+1e-15)
 83      hv.append({'r':r,'hvp_relative_error':float(rel),'min_gauge_eigen':float(np.linalg.eigvalsh(H[:-1,:-1]).min())})
 84    # Prediction 3: thin application agrees with dense W, while storage ratio grows ~n/r.
 85    app=[]
 86    for n,r in [(32,4),(64,4),(128,8),(256,8)]:
 87      U,V,_,_,res,_=project(np.exp(rng.normal(size=(n,r))),np.exp(rng.normal(size=(n,r))))
 88      X=rng.normal(size=(n,16)); Yd,W=dense_apply(U,V,X); Yt=thin_apply(U,V,X)
 89      app.append({'n':n,'r':r,'apply_rel_error':float(np.linalg.norm(Yd-Yt)/(np.linalg.norm(Yd)+1e-15)),
 90                  'dense_entries':n*n,'factor_entries':2*n*r+r,
 91                  'storage_ratio':float(n*n/(2*n*r+r))})
 92    # Small attention comparison: dense softmax baseline versus exact factor mixer.
 93    n,r,d=256,8,32
 94    Q=rng.normal(size=(n,d)); K=rng.normal(size=(n,d)); X=rng.normal(size=(n,d))
 95    t=time.perf_counter(); S=Q@K.T/np.sqrt(d); S-=S.max(1,keepdims=True); P=np.exp(S); P/=P.sum(1,keepdims=True); yd=P@X; dense_ms=1000*(time.perf_counter()-t)
 96    U,V,_,_,res,_=project(np.exp(Q[:,:r]/3),np.exp(K[:,:r]/3))
 97    t=time.perf_counter(); yi=thin_apply(U,V,X); thin_ms=1000*(time.perf_counter()-t)
 98    result={'prediction_1_feasibility':rows,'prediction_2_hvp':hv,'prediction_3_application':app,
 99            'mini_experiment':{'n':n,'rank':r,'dense_ms':dense_ms,'thin_ms':thin_ms,
100              'dense_attention_row_residual':float(abs(P.sum(1)-1).max()),'idea_ds_residual':float(res),
101              'dense_state_entries':n*n,'idea_factor_entries':2*n*r+r,
102              'output_relative_difference':float(np.linalg.norm(yd-yi)/(np.linalg.norm(yd)+1e-15))}}
103    with open('results.json','w') as f: json.dump(result,f,indent=2)
104    print(json.dumps(result,indent=2))
105
106if __name__=='__main__': run()