| """The complete c8 milestone, isolated from the architecture-search code. |
| |
| Frozen reader weights; online habit and attention schema. No semantic labels, |
| personality target, progress reward, text memory or optimization objective. |
| """ |
| import copy |
| import math |
| import os |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| def unit(x): |
| return x / x.norm(dim=-1, keepdim=True).clamp_min(1e-8) |
|
|
|
|
| def carry(a, b, x): |
| c = (a*b).sum(-1, keepdim=True) |
| w = b-c*a |
| sn = w.norm(dim=-1, keepdim=True) |
| w = w/sn.clamp_min(1e-8) |
| xa, xw = (x*a).sum(-1, keepdim=True), (x*w).sum(-1, keepdim=True) |
| out = x + (c*xa-sn*xw-xa)*a + (sn*xa+c*xw-xw)*w |
| return torch.where(sn > 1e-8, out, x) |
|
|
|
|
| class Reader(nn.Module): |
| def __init__(self, weights): |
| super().__init__() |
| w = torch.load(weights, map_location='cpu', weights_only=True) |
| self.d = w['rln.weight'].numel() |
| self.rln = nn.LayerNorm(self.d) |
| self.rk, self.rv = nn.Linear(self.d,256), nn.Linear(self.d,256) |
| self.rq = nn.Linear(8,256) |
| self.rq_bias = nn.Parameter(torch.zeros(256)) |
| self.rpos = nn.Parameter(torch.zeros(16)) |
| self.to_pic = nn.Sequential(nn.Linear(256,128),nn.GELU(),nn.Linear(128,32)) |
| self.dec = nn.Sequential(nn.Linear(32,64),nn.GELU(),nn.Linear(64,8)) |
| self.load_state_dict({k:v for k,v in w.items() if k in self.state_dict()}, strict=True) |
| self.eval().requires_grad_(False) |
|
|
| def read(self, s, hist, habit, extra): |
| n,l,_ = hist.shape |
| a = self.rln(hist) |
| k,v = self.rk(a).view(n,l,4,64),self.rv(a).view(n,l,4,64) |
| q = (self.rq_bias.expand(n,-1)+self.rq(s)).view(n,4,64) |
| if extra is not None: q = q+extra |
| if habit is not None: q = q+torch.einsum('nhdm,nm->nhd',habit,s) |
| att = torch.softmax(torch.einsum('nhd,nlhd->nhl',q,k)/8 + self.rpos[-l:][None,None],-1) |
| r = torch.einsum('nhl,nlhd->nhd',att,v).reshape(n,-1) |
| raw = self.to_pic(r) |
| picture = math.sqrt(32)*raw/raw.norm(dim=-1,keepdim=True).clamp_min(1e-6) |
| dest = self.dec(picture) |
| dest = dest/dest.norm(dim=-1,keepdim=True).clamp_min(1e-6) |
| return picture,dest,att,k |
|
|
|
|
| class Life: |
| def __init__(self, reader, seed=40, copies=4, schema='own', habit=True, |
| exchange=None, at=1000, until=2000, freeze_at=None, freeze_habit_at=None): |
| self.reader=reader |
| self.dev=next(reader.parameters()).device |
| self.config=dict(seed=seed,copies=copies,schema=schema,habit=habit,exchange=exchange, |
| at=at,until=until,freeze_at=freeze_at,freeze_habit_at=freeze_habit_at) |
| self.t=0 |
| self.s=unit(torch.randn(copies,8,generator=torch.Generator().manual_seed(seed))).to(self.dev)[:1].repeat(copies,1) |
| self.prev_g=None |
| self.hist=torch.zeros(copies,0,reader.d,device=self.dev) |
| self.habit=torch.zeros(copies,4,64,8,device=self.dev) if habit else None |
| self.cap=float(.5*reader.rq.weight.norm()) |
| self.rate=None |
| self.Z=torch.zeros(copies,16,16,device=self.dev) |
| self.C=torch.zeros(copies,256,16,device=self.dev) |
| self.proj=torch.randn(16,256,generator=torch.Generator().manual_seed(1616)).to(self.dev)/16 |
| self.kappa=None |
| self.shift=None |
|
|
| @torch.no_grad() |
| def step(self, h): |
| cfg=self.config |
| self.hist=torch.cat([self.hist,h.to(self.dev)[:,None]],1)[:,-16:] |
| n,l,_=self.hist.shape |
| keys=self.reader.rk(self.reader.rln(self.hist)).view(n,l,4,64) |
| z=unit(torch.einsum('ij,bj->bi',self.proj,keys.mean(1).reshape(n,-1)))*4 |
| A=self.C@torch.linalg.inv(self.Z+1e-3*torch.eye(16,device=self.dev)) |
| used=A |
| in_window=cfg['at'] <= self.t < cfg['until'] |
| if in_window and cfg['exchange']: |
| if cfg['exchange']=='swap': used=A[torch.tensor([1,0]+list(range(2,n)),device=self.dev)] |
| elif cfg['exchange']=='random': |
| if self.shift is None: |
| gen=torch.Generator().manual_seed(4646) |
| size=(A[0]-A[1]).norm() |
| self.shift=torch.zeros_like(A) |
| for b in (0,1): |
| r=torch.randn(A.shape[1:],generator=gen).to(self.dev) |
| self.shift[b]=r*size/r.norm().clamp_min(1e-12) |
| used=A+self.shift |
| pull=torch.einsum('bij,bj->bi',used,z) |
| extra=None |
| if cfg['schema']!='none' and self.t>=300: |
| if self.kappa is None: |
| self.kappa=float(.5*self.cap/pull.norm(dim=-1).mean().clamp_min(1e-12)) |
| extra=(self.kappa*pull).view(n,4,64) |
| picture,g,att,k=self.reader.read(self.s,self.hist,self.habit,extra) |
| kp=(torch.einsum('nhl,nlhd->nhd',att,keys)-keys.mean(1)).reshape(n,-1) |
| pred=torch.einsum('bij,bj->bi',A,z) |
| kc=(keys-keys.mean(1,keepdim=True)).reshape(n,l,-1) |
| expect=torch.einsum('bd,bld->bl',pred,kc) |
| record=dict(s=self.s.clone(),g=g.clone(),picture=picture.clone(),att=att.mean(1).clone(), |
| schema_error=((pred-kp)**2).sum(-1),schema_total=(kp**2).sum(-1), |
| expected=expect,schema=A.clone(),habit=None if self.habit is None else self.habit.clone()) |
| if cfg['freeze_at'] is None or self.t<cfg['freeze_at']: |
| zl,kl=(torch.roll(z,-1,0),torch.roll(kp,-1,0)) if cfg['schema']=='neighbor' else (z,kp) |
| self.Z=.998*self.Z+torch.einsum('bi,bj->bij',zl,zl) |
| self.C=.998*self.C+torch.einsum('bi,bj->bij',kl,zl) |
| new_s=self.s if self.prev_g is None else carry(self.prev_g,g,self.s) |
| if self.habit is not None and self.prev_g is not None and l==16 and (cfg['freeze_habit_at'] is None or self.t<cfg['freeze_habit_at']): |
| kap=torch.einsum('nhl,nlhd->nhd',att,k)-k.mean(1) |
| w=torch.einsum('nhd,nm->nhdm',kap,self.s) |
| if self.rate is None: self.rate=float(self.cap/(1000*float(w.flatten(1).norm(dim=-1).mean().clamp_min(1e-12)))) |
| self.habit=self.habit+self.rate*w |
| self.habit=self.habit*(self.cap/self.habit.flatten(1).norm(dim=-1).clamp_min(1e-12)).clamp(max=1).view(n,1,1,1) |
| self.s,self.prev_g=new_s,g |
| self.t+=1 |
| return record |
|
|
| def snapshot(self): |
| state={k:(v.detach().cpu().clone() if isinstance(v,torch.Tensor) else copy.deepcopy(v)) |
| for k,v in self.__dict__.items() if k not in ('reader','dev')} |
| return dict(format=1,state=state,torch_rng=torch.random.get_rng_state(),numpy_rng=np.random.get_state(), |
| cuda_rng=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []) |
|
|
| def save(self,path): |
| path=Path(path); path.parent.mkdir(parents=True,exist_ok=True) |
| tmp=path.with_suffix(path.suffix+'.partial') |
| torch.save(self.snapshot(),tmp); os.replace(tmp,path) |
|
|
| def restore(self,snapshot): |
| for k,v in snapshot['state'].items(): setattr(self,k,v.to(self.dev) if isinstance(v,torch.Tensor) else copy.deepcopy(v)) |
| torch.random.set_rng_state(snapshot['torch_rng']) |
| np.random.set_state(snapshot['numpy_rng']) |
| if snapshot['cuda_rng'] and torch.cuda.is_available(): torch.cuda.set_rng_state_all(snapshot['cuda_rng']) |
| return self |
|
|
