1import json, math, os, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1029
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 10    if device.type == "cuda":
 11        torch.cuda.manual_seed_all(SEED)
 12        # A failed CUDA allocation should not invalidate the experiment.
 13        torch.zeros(1, device=device)
 14except Exception:
 15    device = torch.device("cpu")
 16
 17class Coupling(nn.Module):
 18    """Additive coupling C_f,g and its analytically computed inverse."""
 19    def __init__(self, d, scale=0.1):
 20        super().__init__()
 21        h = d // 2
 22        self.f = nn.Linear(h, h, bias=False)
 23        self.g = nn.Linear(h, h, bias=False)
 24        nn.init.normal_(self.f.weight, std=scale / math.sqrt(h))
 25        nn.init.normal_(self.g.weight, std=scale / math.sqrt(h))
 26
 27    def forward(self, x):
 28        a, b = x.chunk(2, dim=-1)
 29        ap = a + self.f(b)
 30        bp = b + self.g(ap)
 31        return torch.cat((ap, bp), dim=-1)
 32
 33    def inverse(self, y):
 34        ya, yb = y.chunk(2, dim=-1)
 35        b = yb - self.g(ya)
 36        a = ya - self.f(b)
 37        return torch.cat((a, b), dim=-1)
 38
 39class InverseCopy(nn.Module):
 40    def __init__(self, source):
 41        super().__init__()
 42        d = source.f.in_features * 2
 43        self.block = Coupling(d, 0.01)
 44        self.block.load_state_dict(source.state_dict())
 45    def forward(self, x):
 46        return self.block.inverse(x)
 47
 48class Pair(nn.Module):
 49    def __init__(self, source):
 50        super().__init__()
 51        self.forward_block = Coupling(source.f.in_features * 2, 0.01)
 52        self.forward_block.load_state_dict(source.state_dict())
 53        self.inverse_block = InverseCopy(source)
 54    def forward(self, x):
 55        return self.inverse_block(self.forward_block(x))
 56
 57class Net(nn.Module):
 58    def __init__(self, d, source=None):
 59        super().__init__()
 60        self.block = Coupling(d, 0.15) if source is None else source
 61        self.out = nn.Linear(d, 1)
 62    def forward(self, x):
 63        return self.out(self.block(x))
 64    def append_pair(self):
 65        old = self.block
 66        self.block = nn.Sequential(old, Pair(old).to(next(old.parameters()).device))
 67
 68
 69def coupling_error(scale, dtype=torch.float32, n=4096):
 70    torch.manual_seed(7)
 71    d = 8
 72    c = Coupling(d, 0.01).to(device=device, dtype=dtype)
 73    with torch.no_grad():
 74        c.f.weight.mul_(scale / 0.01); c.g.weight.mul_(scale / 0.01)
 75        x = torch.randn(n, d, device=device, dtype=dtype)
 76        z = c.inverse(c(x))
 77        err = (z-x).abs().max().item()
 78        rel = (z-x).norm().item() / max(x.norm().item(), 1e-30)
 79    return err, rel
 80
 81def perturbation_sweep():
 82    # At an identity pair, a small untied parameter step produces O(eta) output change.
 83    torch.manual_seed(11)
 84    d=8; x=torch.randn(512,d,device=device)
 85    source=Coupling(d,0.08).to(device)
 86    pair=Pair(source).to(device)
 87    with torch.no_grad(): base=pair(x).clone()
 88    loss=(pair(x)**2).mean(); loss.backward()
 89    grads=[p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p) for p in pair.parameters()]
 90    grad_norm=math.sqrt(sum((g**2).sum().item() for g in grads))
 91    vals=[]
 92    for eta in [1e-5, 3e-5, 1e-4, 3e-4, 1e-3]:
 93        q=Pair(source).to(device); q.load_state_dict(pair.state_dict())
 94        with torch.no_grad():
 95            for p,g in zip(q.parameters(), grads):
 96                p.sub_(eta*g)
 97            delta=(q(x)-base).norm().item()/x.norm().item()
 98        vals.append((eta,delta))
 99    slope=np.polyfit(np.log([a for a,b in vals if b>0]),np.log([b for a,b in vals if b>0]),1)[0]
100    return {"gradient_norm":grad_norm,"sweep":vals,"loglog_slope":float(slope)}
101
102def train(model, x, y, steps, lr=2e-3):
103    opt=torch.optim.Adam(model.parameters(),lr=lr)
104    losses=[]
105    for _ in range(steps):
106        opt.zero_grad(set_to_none=True)
107        loss=((model(x)-y)**2).mean()
108        loss.backward(); opt.step(); losses.append(loss.item())
109    return losses
110
111def mini_experiment():
112    torch.manual_seed(SEED)
113    d=8; n=2048
114    x=torch.randn(n,d,device=device)
115    # A target requiring a nontrivial but small transformation.
116    target=Coupling(d,0.22).to(device)
117    with torch.no_grad(): y=(target(target(x))).sum(dim=1,keepdim=True)/2.0
118    torch.manual_seed(SEED)
119    progressive=Net(d).to(device)
120    t0=time.time(); pre=train(progressive,x,y,250); before=((progressive(x)-y)**2).mean().item()
121    progressive.append_pair()
122    with torch.no_grad(): after=((progressive(x)-y)**2).mean().item()
123    insertion_rel=(after-before)/max(abs(before),1e-30)
124    post=train(progressive,x,y,300)
125    # Standard shallow continuation, and a randomly initialized two-block model.
126    torch.manual_seed(SEED); shallow=Net(d).to(device); train(shallow,x,y,250); shallow_post=train(shallow,x,y,300)
127    torch.manual_seed(SEED+1); random_deep=Net(d).to(device); random_deep.block=nn.Sequential(random_deep.block,Coupling(d,0.15).to(device)); random_post=train(random_deep,x,y,300)
128    return {"device":str(device),"before_insertion_loss":before,"after_insertion_loss":after,"insertion_relative_change":insertion_rel,
129            "paired_loss_after_300":post[-1],"shallow_loss_after_300":shallow_post[-1],"random_deep_loss_after_300":random_post[-1],
130            "paired_best_300":min(post),"wall_seconds":time.time()-t0}
131
132def main():
133    scales=[0.01,0.03,0.1,0.3,1.0]
134    cancel=[]
135    for s in scales:
136        e,r=coupling_error(s); cancel.append({"scale":s,"max_abs_error":e,"relative_error":r,"predicted":"identity; fp error increases with operator products"})
137    pert=perturbation_sweep()
138    mini=mini_experiment()
139    result={"seed":SEED,"math_predictions":[
140        "C^{-1}(C(x))=x, so insertion loss change is zero up to floating point error.",
141        "Composition roundoff should increase with coupling scale (products of f and g operators).",
142        "After untie, a parameter step of size eta changes the pair output proportionally to eta (log-log slope 1)."],
143        "cancellation_sweep":cancel,"untied_step_sweep":pert,"mini_experiment":mini}
144    with open("results.json","w") as f: json.dump(result,f,indent=2)
145    print(json.dumps(result,indent=2))
146
147if __name__ == "__main__": main()