"""Experimental paper-derived DiffusionBlocks/AR/SAT/NAT CPU adapters. Not official Sakana AR code. Equations/noise-band design follow paper2506.14202v3 and SakanaAI/DiffusionBlocks@ade0b08. Upstream Apache2license is retained separately. All dense projection 'binary' operations emulate STE signs in FP32, not packed. Generation takes prefixes only, never reference continuations. """ import math from statistics import NormalDist import torch from torch import nn from torch.nn import functional as F N=NormalDist();PM=-1.2;PS=1.2;SMIN=.002;SMAX=80. def cdf(s):return N.cdf((math.log(s)-PM)/PS) def sigma(p):return math.exp(PM+PS*N.inv_cdf(p)) def boundaries(n):return [sigma(cdf(SMIN)+(cdf(SMAX)-cdf(SMIN))*i/n) for i in range(n+1)] class Linear(nn.Linear): def __init__(self,a,b,binary=False):super().__init__(a,b,bias=False);self.binary=binary def forward(self,x): w=self.weight if self.binary: quant=torch.where(w>=0,1.,-1.)*w.detach().abs().mean(1,keepdim=True);w=w+(quant-w).detach() return F.linear(x,w) class RMS(nn.Module): def __init__(self,d):super().__init__();self.weight=nn.Parameter(torch.ones(d)) def forward(self,x):return x*torch.rsqrt(x.square().mean(-1,keepdim=True)+1e-5)*self.weight class Block(nn.Module): def __init__(self,d,h,binary): super().__init__();self.h=h;self.d=d;self.n1=RMS(d);self.n2=RMS(d) self.qkv=Linear(d,3*d,binary);self.out=Linear(d,d,binary);self.up=Linear(d,2*d,binary);self.gate=Linear(d,2*d,binary);self.down=Linear(2*d,d,binary) self.time=nn.Linear(1,d,bias=False) def rope(self,x,pos): # Head widths2,3,4,...: rotate the even prefix; retain any odd final dim. k=(x.shape[-1]//2)*2 f=10000.**(-torch.arange(0,k,2,device=x.device,dtype=x.dtype)/max(1,k));a=pos[:,None]*f u,v=x[...,:k:2],x[...,1:k:2];r=torch.stack((u*a.cos()-v*a.sin(),u*a.sin()+v*a.cos()),-1).flatten(-2) return torch.cat([r,x[...,k:]],-1) def forward(self,x,mask,pos,noise=None): if noise is not None:x=x+self.time(noise) b,t,d=x.shape;q,k,v=self.qkv(self.n1(x)).chunk(3,-1);q,k,v=[z.reshape(b,t,self.h,d//self.h).transpose(1,2) for z in (q,k,v)] q,k=self.rope(q,pos),self.rope(k,pos);a=F.scaled_dot_product_attention(q,k,v,attn_mask=mask).transpose(1,2).reshape(b,t,d) x=x+self.out(a);z=self.n2(x);return x+self.down(F.silu(self.gate(z))*self.up(z)) class Architecture(nn.Module): def __init__(self,spec): super().__init__();self.spec=spec;d=spec['width'];self.emb=nn.Embedding(256,d);self.norm=RMS(d) self.blocks=nn.ModuleList([Block(d,spec['attention_heads'],spec['binary_transformer_projections']) for _ in range(spec['depth'])]);self.diff=spec['family'].startswith('diffusionblocks');self.bands=boundaries(spec['depth']) def logits(self,x):return F.linear(self.norm(x),self.emb.weight)/math.sqrt(self.spec['width']) def forward(self,tokens): assert not self.diff,'Diffusion model has no ordinary AR likelihood forward' x=self.emb(tokens);t=tokens.shape[1];i=torch.arange(t);mask=i[:,None]>=i[None,:] if self.spec['family']=='ar_local16':mask&=(i[:,None]-i[None,:]<16) for block in self.blocks:x=block(x,mask,i) return self.logits(x) def training_mask(self,t): # x_i is the previous clean token for target y_i. Clean keys <=groupstart # are permitted; noisy keys only from the current generated group. k=self.spec['generation_block_size'];i=torch.arange(t);group=(i//k)*k mask=torch.zeros((2*t,2*t),dtype=torch.bool);mask[:t,:t]=i[:,None]>=i[None,:] mask[t:,:t]=i[None,:]<=group[:,None];mask[t:,t:]=(i[None,:]//k)==(i[:,None]//k) return mask def denoise(self,clean,z,sig,block,training=True): b,t,d=z.shape;past=self.emb(clean);sd=.5 cin=torch.rsqrt(sig.square()+sd**2);cout=sig*sd*cin;cskip=sd**2/(sig.square()+sd**2) x=torch.cat([past,z*cin[:,None,None]],1);n=past.shape[1] cond=x.new_zeros((b,n+t,1));cond[:,n:]=(.25*sig.log())[:,None,None] if training:mask=self.training_mask(t);pos=torch.cat([torch.arange(n),torch.arange(t)+1]) else: mask=torch.zeros((n+t,n+t),dtype=torch.bool);mask[:n,:n]=torch.ones((n,n),dtype=torch.bool).tril();mask[n:,:]=True;pos=torch.arange(n+t) hidden=self.blocks[block](x,mask,pos,cond)[:,n:] denoised=hidden*cout[:,None,None]+z*cskip[:,None,None] return self.logits(denoised) def loss(self,batch,g): x,y=batch[:,:-1],batch[:,1:] if not self.diff:return F.cross_entropy(self(x).flatten(0,1),y.flatten()),{'active_blocks':len(self.blocks)} interval=int(torch.randint(len(self.blocks),(1,),generator=g));lo,hi=self.bands[interval:interval+2] delta=math.log(hi)-math.log(lo);lo=max(SMIN,math.exp(math.log(lo)-.1*delta));hi=min(SMAX,math.exp(math.log(hi)+.1*delta)) u=torch.rand(batch.shape[0],generator=g).tolist();sig=torch.tensor([sigma(cdf(lo)+(cdf(hi)-cdf(lo))*v) for v in u]) target=F.normalize(self.emb(y),dim=-1);z=target+sig[:,None,None]*torch.randn(target.shape,generator=g) block=len(self.blocks)-1-interval;logits=self.denoise(x,z,sig,block,True) ce=F.cross_entropy(logits.flatten(0,1),y.flatten(),reduction='none').reshape(y.shape) weight=(sig.square()+.25)/(sig*.5).square() return (ce*weight[:,None]).mean(),{'active_blocks':1,'block':block,'objective':'weighted denoising CE, not AR NLL'} @torch.inference_mode() def generate(self,prefix,new_tokens,seed=1,passes=4): self.eval();g=torch.Generator().manual_seed(seed);tokens=prefix.clone();calls=0;generated=0 while generated=v for v in self.bands)-1));block=len(self.blocks)-1-interval logits=self.denoise(tokens,z,torch.full((b,),s),block,False);calls+=1 if i+1