Positive-envelope stability for complex state updates / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json, math, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2784
  6rng = np.random.default_rng(SEED)
  7
  8
  9def envelope(A_factors):
 10    A = np.eye(A_factors[0].shape[0], dtype=np.complex128)
 11    P = np.eye(A.shape[0], dtype=np.float64)
 12    for F in A_factors:
 13        A = F @ A
 14        P = np.abs(F) @ P
 15    return A, P
 16
 17
 18def power_radius(P, iters=100):
 19    v = np.ones(P.shape[0]) / math.sqrt(P.shape[0])
 20    for _ in range(iters):
 21        z = P @ v
 22        n = np.linalg.norm(z)
 23        if n == 0: return 0.0
 24        v = z / n
 25    return float(np.linalg.norm(P @ v) / (np.linalg.norm(v) + 1e-15))
 26
 27
 28def verify():
 29    # Prediction 1: |F_L...F_1| <= product |F_k| entrywise for every draw.
 30    violations, ratios = [], []
 31    for _ in range(5000):
 32        L = int(rng.integers(1, 5))
 33        fs = [(rng.normal(size=(4,4)) + 1j*rng.normal(size=(4,4))) / math.sqrt(8)
 34              for _ in range(L)]
 35        A, P = envelope(fs)
 36        violations.append(float(np.max(np.abs(A)-P)))
 37        ratios.append(float(np.max(np.abs(A)/(P+1e-15))))
 38    max_violation = max(violations)
 39
 40    # Prediction 2: two equal paths have cancellation ratio |cos(theta/2)|.
 41    thetas = np.linspace(0, math.pi, 9)
 42    cancel_obs = []
 43    cancel_pred = []
 44    for th in thetas:
 45        z = 1.0 + np.exp(1j*th)
 46        cancel_obs.append(abs(z)/2.0)
 47        cancel_pred.append(abs(math.cos(th/2)))
 48
 49    # Prediction 3: scalar repeated dynamics. Envelope threshold is gamma=1/2,
 50    # while actual threshold is gamma=1/(2|cos(theta/2)|).
 51    gamma_grid = np.linspace(.1, 4.0, 391)
 52    threshold_rows = []
 53    for th in [0.0, math.pi/2, math.pi*0.9, math.pi]:
 54        c = abs(math.cos(th/2))
 55        env_factor = 2.0 * gamma_grid
 56        actual_factor = 2.0 * c * gamma_grid
 57        env_cross = gamma_grid[np.argmax(env_factor >= 1.0)]
 58        act_cross = gamma_grid[np.argmax(actual_factor >= 1.0)] if c > 1e-12 and actual_factor.max() >= 1.0 else None
 59        threshold_rows.append({
 60            'theta': float(th), 'predicted_envelope_gamma': .5,
 61            'observed_envelope_gamma': float(env_cross),
 62            'predicted_actual_gamma': None if c == 0 else float(1/(2*c)),
 63            'observed_actual_gamma': None if act_cross is None else float(act_cross),
 64            'cancellation_factor': float(c),
 65        })
 66    return {
 67        'entrywise_bound': {'trials': 5000, 'max_violation': max_violation,
 68                            'max_absA_over_P': max(ratios)},
 69        'phase_cancellation': [{'theta': float(t), 'predicted_ratio': float(p),
 70                                'observed_ratio': float(o)}
 71                               for t,p,o in zip(thetas,cancel_pred,cancel_obs)],
 72        'growth_thresholds': threshold_rows,
 73    }
 74
 75
 76def torch_benchmark():
 77    # Optional torch benchmark: learn a stable scalar-output complex recurrence.
 78    # Both models have identical factorized transition parameterization and steps.
 79    try:
 80        import torch
 81        torch.manual_seed(SEED)
 82        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 83        try:
 84            torch.zeros(1, device=device).sum().item()
 85        except Exception:
 86            device = 'cpu'
 87        dtype = torch.complex64
 88        n, T, B, steps = 4, 35, 64, 350
 89        g = torch.Generator(device=device).manual_seed(SEED)
 90        # Stable teacher with long memory; inputs are real but state is complex.
 91        F1t = (0.34*torch.eye(n, device=device, dtype=dtype) +
 92               0.025*(torch.randn(n,n, generator=g, device=device)+1j*torch.randn(n,n,generator=g,device=device)))
 93        F2t = (0.90*torch.eye(n, device=device, dtype=dtype) +
 94               0.025*(torch.randn(n,n, generator=g, device=device)+1j*torch.randn(n,n,generator=g,device=device)))
 95        At = F2t @ F1t
 96        x = torch.randn(B,T,1, generator=g, device=device, dtype=torch.float32).to(dtype)
 97        with torch.no_grad():
 98            h = torch.zeros(B,n,device=device,dtype=dtype); ys=[]
 99            for t in range(T):
100                h = h @ At.T + x[:,t] @ torch.ones(1,n,device=device,dtype=dtype)*.12
101                ys.append(h[:,0].real)
102            y = torch.stack(ys,1)
103
104        def run(kind):
105            torch.manual_seed(SEED+ (0 if kind=='baseline' else 1))
106            F1 = torch.nn.Parameter((.72*torch.eye(n,dtype=dtype,device=device) + .11*(torch.randn(n,n,device=device)+1j*torch.randn(n,n,device=device))) )
107            F2 = torch.nn.Parameter((.72*torch.eye(n,dtype=dtype,device=device) + .11*(torch.randn(n,n,device=device)+1j*torch.randn(n,n,device=device))) )
108            read = torch.nn.Parameter(torch.randn(n,device=device,dtype=dtype)*.1)
109            opt = torch.optim.Adam([F1,F2,read], lr=.012)
110            t0=time.perf_counter(); unstable=0
111            for _ in range(steps):
112                opt.zero_grad(); A=F2@F1; P=torch.abs(F2)@torch.abs(F1)
113                h=torch.zeros(B,n,device=device,dtype=dtype); pred=[]
114                for t in range(T):
115                    h=h@A.T + x[:,t]@torch.ones(1,n,device=device,dtype=dtype)*.12
116                    pred.append((h*read).sum(1).real)
117                loss=torch.stack(pred,1).sub(y).pow(2).mean()
118                if kind=='envelope':
119                    # Smooth differentiable conservative penalty: row-sum envelope above target.
120                    env=P.sum(1).amax()
121                    loss=loss + .08*torch.relu(env-.96).pow(2) + .01*(P.sum()-torch.abs(A).sum())/(P.sum()+1e-6)
122                loss.backward()
123                gn=torch.nn.utils.clip_grad_norm_([F1,F2,read], 10.0)
124                if not torch.isfinite(loss) or not torch.isfinite(gn): unstable += 1
125                opt.step()
126            with torch.no_grad():
127                A=F2@F1; P=torch.abs(F2)@torch.abs(F1)
128                # long rollout against teacher from fresh sequence
129                xx=torch.randn(B,120,1,generator=g,device=device,dtype=torch.float32).to(dtype)
130                hp=torch.zeros(B,n,device=device,dtype=dtype); ht=torch.zeros_like(hp)
131                err=[]
132                for t in range(120):
133                    hp=hp@A.T+xx[:,t]@torch.ones(1,n,device=device,dtype=dtype)*.12
134                    ht=ht@At.T+xx[:,t]@torch.ones(1,n,device=device,dtype=dtype)*.12
135                    err.append(((hp*read).sum(1).real-(ht[:,0].real))**2)
136                longerr=torch.stack(err,1).mean().sqrt().item()
137                rhoP=power_radius(P.detach().cpu().numpy())
138                rhoA=max(abs(np.linalg.eigvals(A.detach().cpu().numpy())))
139            return {'final_train_loss':float(loss.item()),'long_rollout_rmse':longerr,
140                    'unstable_steps':unstable,'rho_A':float(rhoA),'rho_P':float(rhoP),
141                    'seconds':time.perf_counter()-t0,'device':device}
142        return {'baseline':run('baseline'),'envelope':run('envelope')}
143    except Exception as e:
144        return {'error':repr(e), 'note':'torch benchmark failed; algebraic verification remains valid'}
145
146
147if __name__ == '__main__':
148    out={'seed':SEED,'verification':verify(),'benchmark':torch_benchmark()}
149    Path('results.json').write_text(json.dumps(out, indent=2))
150    print(json.dumps(out, indent=2))