Intrinsic Schrödinger Bridge Diffusion / experiment.py
Mechanism confirmed, baseline not beaten
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()