import json, math, random from pathlib import Path import numpy as np import torch from torch import nn SEED=17 random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) try: device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device.type=='cuda': torch.zeros(1,device=device) except Exception: device=torch.device('cpu') torch.set_default_dtype(torch.float32) class Coupling(nn.Module): def __init__(self,d,mask,cond=2,hidden=32,clamp=1.5): super().__init__(); self.d=d; self.register_buffer('mask',mask.float()); self.clamp=clamp inp=d+cond+1 self.net=nn.Sequential(nn.Linear(inp,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,2*d)) # near identity initialization, but trainable nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias) def params_at(self,x,c,t): u=torch.cat([x*self.mask, c, t],-1); out=self.net(u) raw_s,b=out[...,:self.d],out[...,self.d:] s=self.clamp*torch.tanh(raw_s)* (1-self.mask) b=b*(1-self.mask) return s,b def forward(self,x,c,t): s,b=self.params_at(x,c,t); y=x*self.mask+(1-self.mask)*(x*torch.exp(s)+b) return y, s.sum(-1) def inverse(self,y,c,t): s,b=self.params_at(y,c,t); x=y*self.mask+(1-self.mask)*(y-b)*torch.exp(-s) return x, -s.sum(-1) class Flow(nn.Module): def __init__(self,d,K=4,cond=2): super().__init__(); masks=[] for k in range(K): masks.append(torch.tensor([(i+k)%2 for i in range(d)])) self.layers=nn.ModuleList([Coupling(d,m,cond) for m in masks]); self.d=d def forward(self,z,c,t): h=z; ld=torch.zeros(z.shape[0],device=z.device) for layer in self.layers: h,a=layer(h,c,t); ld=ld+a return h,ld def inverse(self,x,c,t): h=x; ld=torch.zeros(x.shape[0],device=x.device) for layer in reversed(self.layers): h,a=layer.inverse(h,c,t); ld=ld+a return h,ld def log_prob(self,x,c,t): z, invld=self.inverse(x,c,t) logq=-0.5*(z*z+math.log(2*math.pi)).sum(-1) return logq+invld def score(self,x,c,t): xx=x.detach().requires_grad_(True); lp=self.log_prob(xx,c,t).sum() return torch.autograd.grad(lp,xx)[0] class DiagMLP(nn.Module): def __init__(self,d,cond=2): super().__init__(); self.d=d; self.net=nn.Sequential(nn.Linear(cond+1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2*d)) def log_prob(self,x,c,t): o=self.net(torch.cat([c,t],-1)); mu,raw=o[...,:self.d],o[...,self.d:] logstd=2*torch.tanh(raw); return (-0.5*((x-mu)/logstd.exp())**2-logstd-0.5*math.log(2*math.pi)).sum(-1) def score(self,x,c,t): xx=x.detach().requires_grad_(True); lp=self.log_prob(xx,c,t).sum(); return torch.autograd.grad(lp,xx)[0] def sample_data(n,d,device): # task c=(separation, parity feature); conditional symmetric two-component mixture sep=torch.rand(n,1,device=device)*1.5+1.5; aux=torch.sin(sep) c=torch.cat([sep,aux],-1); t=torch.zeros(n,1,device=device) sign=torch.where(torch.rand(n,1,device=device)<.5,-1.,1.) direction=torch.ones(1,d,device=device); direction[:,1::2]=-1 means=sign*sep*direction x=means+0.35*torch.randn(n,d,device=device) return x,c,t def analytic_score(x,c): sep=c[:,0:1]; direction=torch.ones(1,x.shape[1],device=x.device); direction[:,1::2]=-1 m=sep*direction; var=.35**2 lp=torch.stack([-0.5*((x-mm)**2).sum(-1)/var for mm in [m,-m]],-1) w=torch.softmax(lp,-1) scores=torch.stack([-(x-mm)/var for mm in [m,-m]],1) return (w.unsqueeze(-1)*scores).sum(1) def train(model,d,steps=700): model.to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3) for i in range(steps): x,c,t=sample_data(96,d,device); loss=-model.log_prob(x,c,t).mean(); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): x,c,t=sample_data(512,d,device); nll=-model.log_prob(x,c,t).mean().item() # score needs graph, use fixed batch x,c,t=sample_data(256,d,device); pred=model.score(x,c,t); truth=analytic_score(x,c) return nll, (pred-truth).pow(2).mean().sqrt().item() def math_checks(): d=8; x=torch.randn(64,d,device=device); c=torch.randn(64,2,device=device); t=torch.randn(64,1,device=device) rows=[] # constant alpha is imposed through a tiny custom coupling: test the exact formulas directly for gamma in [0.,.25,.5,1.,1.5]: mask=torch.tensor([1,0,1,0,1,0,1,0],device=device).float(); nB=int((1-mask).sum()) alpha=gamma*torch.ones_like(x)*(1-mask); beta=.13*torch.ones_like(x)*(1-mask) y=x*mask+(1-mask)*(x*torch.exp(alpha)+beta) xr=y*mask+(1-mask)*(y-beta)*torch.exp(-alpha) recon=(xr-x).abs().max().item(); observed=y.shape[0] # logdet is nB*gamma; additive K=3 prediction is 3*nB*gamma one=alpha.sum(-1).mean().item(); comp=(3*alpha.sum(-1)).mean().item() rows.append({'gamma':gamma,'pred_logdet':nB*gamma,'obs_logdet':one,'pred_3x':3*nB*gamma,'obs_3x':comp,'recon_max':recon}) # Direct Jacobian check on the implemented nonlinear coupling (single sample). layer=Coupling(d, torch.tensor([1,0,1,0,1,0,1,0],device=device).float(), cond=2).to(device) with torch.no_grad(): layer.net[-1].bias[:d].fill_(0.35) layer.net[-1].bias[d:].fill_(0.11) xx=torch.randn(d,device=device,requires_grad=True); cc=torch.randn(1,2,device=device); tt=torch.randn(1,1,device=device) yy,ld=layer(xx[None],cc,tt); xr,ild=layer.inverse(yy,cc,tt) J=torch.autograd.functional.jacobian(lambda u: layer(u[None],cc,tt)[0][0],xx) sign,actual=torch.linalg.slogdet(J) rows.append({'nonlinear_jacobian_logdet':actual.item(),'analytic_logdet':ld.item(), 'jacobian_abs_error':abs(actual.item()-ld.item()),'jacobian_sign':sign.item(), 'nonlinear_inverse_max_error':(xr[0]-xx).abs().max().item()}) return rows def main(): checks=math_checks(); results={'device':str(device),'math_checks':checks,'models':{}} for d in [2,8]: # reset reproducibly per dimension torch.manual_seed(SEED+d) results['models'][str(d)]={} results['models'][str(d)]['flow']=train(Flow(d),d) torch.manual_seed(SEED+d) results['models'][str(d)]['diag_mlp']=train(DiagMLP(d),d) Path('results.json').write_text(json.dumps(results,indent=2)) print(json.dumps(results,indent=2)) if __name__=='__main__': main()