Intrinsic Schrödinger Bridge Diffusion / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
 1import json, math
 2from pathlib import Path
 3import numpy as np
 4
 5SEED = 2742
 6D = 8
 7BETA = 2.0
 8KAPPA = 1.5
 9MU = np.zeros(D); MU[0] = 1.0
10
11def proj(x, v):
12    return v - np.sum(x*v, axis=-1, keepdims=True)*x
13
14def normalize(x):
15    return x / np.linalg.norm(x, axis=-1, keepdims=True)
16
17def sphere_run(method, h, T=0.25, ntraj=2000, amp=0.0, seed=0):
18    r = np.random.default_rng(seed)
19    n = int(round(T/h)); h = T/n
20    x = np.zeros((ntraj,D)); x[:,0] = 1.0
21    maxviol = np.zeros(ntraj); energy = np.zeros(ntraj)
22    noise_scale = math.sqrt(2.0*h/BETA)
23    for _ in range(n):
24        drift = KAPPA*proj(x, MU[None,:])
25        u = amp*proj(x, MU[None,:])
26        z = r.normal(size=x.shape)
27        if method in ('intrinsic','projected_each'):
28            z = proj(x,z)
29        xnew = x + h*(drift+u) + noise_scale*z
30        energy += np.sum(u*u,axis=1)*h
31        if method in ('intrinsic','projected_each'):
32            x = normalize(xnew)
33            # Constraint is evaluated on the state actually used by the method.
34            state_violation = np.abs(np.linalg.norm(x,axis=1)-1.0)
35        else:
36            x = xnew
37            state_violation = np.abs(np.linalg.norm(x,axis=1)-1.0)
38        maxviol = np.maximum(maxviol, state_violation)
39    if method == 'projected_final':
40        x = normalize(x)
41    return dict(max_violation=float(np.max(maxviol)), median_violation=float(np.median(maxviol)),
42      mean_final_norm_error=float(np.mean(np.abs(np.linalg.norm(x,axis=1)-1))),
43      mean_alignment=float(np.mean(x[:,0])), mean_control_energy=float(np.mean(energy)),
44      q95_violation=float(np.quantile(maxviol,.95)))
45
46def one_step_violation(h, ntraj=100000, seed=0):
47    r=np.random.default_rng(seed)
48    x=np.zeros((ntraj,D)); x[:,0]=1
49    z=r.normal(size=x.shape)
50    y=x+math.sqrt(2*h/BETA)*z
51    e=np.abs(np.linalg.norm(y,axis=1)-1)
52    return float(np.median(e)), float(np.quantile(e,.95))
53
54def main():
55    r=np.random.default_rng(SEED+1)
56    x=normalize(r.normal(size=(100,D))); v=r.normal(size=(100,D)); p=proj(x,v)
57    projector_resid=max(float(np.max(np.abs(np.sum(x*p,axis=1)))),float(np.max(np.abs(proj(x,p)-p))))
58    hs=[1/25,1/50,1/100,1/200]
59    intrinsic=[sphere_run('intrinsic',h,seed=100+i) for i,h in enumerate(hs)]
60    med=[]; q95=[]
61    for i,h in enumerate(hs):
62        a,b=one_step_violation(h,seed=300+i); med.append(a); q95.append(b)
63    slope=float(np.polyfit(np.log(hs),np.log(med),1)[0])
64    x0=normalize(r.normal(size=(10000,D))); tangent=proj(x0,MU[None,:]); base=float(np.mean(np.sum(tangent*tangent,axis=1)))
65    amps=[.25,.5,1.,2.,4.]
66    costs=[BETA/4*a*a*base for a in amps]
67    cost_slope=float(np.polyfit(np.log(amps),np.log(costs),1)[0])
68    endpoint={
69      'intrinsic':sphere_run('intrinsic',1/200,T=.25,amp=1,seed=2024),
70      'projected_each':sphere_run('projected_each',1/200,T=.25,amp=1,seed=2024),
71      'projected_final':sphere_run('projected_final',1/200,T=.25,amp=1,seed=2024),
72      'ambient':sphere_run('ambient',1/200,T=.25,amp=1,seed=2024)}
73    response=[sphere_run('intrinsic',1/200,T=.25,amp=a,seed=900+i)['mean_alignment'] for i,a in enumerate([0,.5,1,2,4])]
74    report={'seed':SEED,'dimension':D,'beta':BETA,'kappa':KAPPA,
75      'algebra_check':{'max_projector_tangent_or_idempotence_residual':projector_resid},
76      'prediction_checks':{
77       'P1_intrinsic_constraint':{'prediction':'post-retraction intermediate |norm-1| stays at floating-point precision, independent of h','h_values':hs,'observed_max_each':[q['max_violation'] for q in intrinsic],'observed_max_over_sweep':max(q['max_violation'] for q in intrinsic)},
78       'P2_ambient_local_scaling':{'prediction':'one-step unconstrained radial norm violation scales as h^0.5','h_values':hs,'median_violation':med,'q95_violation':q95,'observed_loglog_exponent':slope},
79       'P3_fixed_state_control_cost':{'prediction':'beta/4 E||u||^2 scales exactly as amplitude^2','amplitudes':amps,'costs':costs,'observed_loglog_exponent':cost_slope,'intrinsic_endpoint_alignment_by_amplitude':response}},
80      'endpoint_comparison_equal_steps':endpoint}
81    Path('results.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
82if __name__=='__main__': main()