172 lines
9.2 KiB
Python
172 lines
9.2 KiB
Python
"""Attention functions and discrete diagnostics, no continuous alignment/fitting."""
|
|
import numpy as np
|
|
from functional.model import StreamingLM,PREFIX,attention,rope,sigmoid
|
|
|
|
|
|
def assignment(cost):
|
|
"""Exact minimum-cost bijection via O(n^2 2^n) dynamic programming."""
|
|
cost=np.asarray(cost,dtype=np.float64);n=cost.shape[0]
|
|
assert cost.shape==(n,n) and n<=12 and np.isfinite(cost).all()
|
|
states={0:(0.,())}
|
|
for i in range(n):
|
|
nxt={}
|
|
for mask,(value,perm) in states.items():
|
|
for j in range(n):
|
|
if mask&(1<<j):continue
|
|
key=mask|(1<<j);candidate=(value+float(cost[i,j]),perm+(j,))
|
|
if key not in nxt or candidate<nxt[key]:nxt[key]=candidate
|
|
states=nxt
|
|
return list(states[(1<<n)-1][1])
|
|
|
|
|
|
def inverse(perm):
|
|
out=[0]*len(perm)
|
|
for i,j in enumerate(perm):out[j]=i
|
|
return out
|
|
|
|
|
|
def output_metrics(a,b):
|
|
"""Per-text directional error, cosine and output RMS. No fitted scaling."""
|
|
x=a.reshape(a.shape[0],-1).astype(np.float64)
|
|
y=b.reshape(b.shape[0],-1).astype(np.float64)
|
|
xx=np.sum(x*x,-1);yy=np.sum(y*y,-1);diff=np.sum((x-y)**2,-1)
|
|
return {'error_a':np.sqrt(diff/np.maximum(xx,1e-30)),
|
|
'error_b':np.sqrt(diff/np.maximum(yy,1e-30)),
|
|
'cosine':np.sum(x*y,-1)/np.maximum(np.sqrt(xx*yy),1e-30),
|
|
'rms_a':np.sqrt(xx/x.shape[1]),'rms_b':np.sqrt(yy/y.shape[1])}
|
|
|
|
|
|
def distribution_metrics(a,b):
|
|
"""Aligned heads; primary queries have at least 8 legal keys (t>=7)."""
|
|
a,b=a.astype(np.float64),b.astype(np.float64)
|
|
assert a.shape==b.shape and a.shape[-1]>=8
|
|
m=(a+b)*.5
|
|
def term(x):return np.where(x>0,x*np.log(np.maximum(x,1e-300)/np.maximum(m,1e-300)),0)
|
|
js=np.maximum(.5*np.sum(term(a)+term(b),axis=-1),0)
|
|
eligible=js[...,7:]
|
|
overlaps=[]
|
|
for t in range(7,a.shape[-2]):
|
|
aa=np.argsort(a[:,:,t,:t+1],axis=-1)[...,-5:]
|
|
bb=np.argsort(b[:,:,t,:t+1],axis=-1)[...,-5:]
|
|
overlaps.append(np.mean((aa[...,None]==bb[...,None,:]).any(-1),-1))
|
|
# First-key ablation is diagnostic, not an unreported filter of positional repetitions.
|
|
ar,br=a[...,1:],b[...,1:];ma,mb=ar.sum(-1),br.sum(-1)
|
|
valid=(ma>1e-8)&(mb>1e-8)
|
|
ar=ar/np.maximum(ma[...,None],1e-300);br=br/np.maximum(mb[...,None],1e-300)
|
|
mr=(ar+br)*.5
|
|
cond=.5*np.sum(np.where(ar>0,ar*np.log(np.maximum(ar,1e-300)/np.maximum(mr,1e-300)),0)+np.where(br>0,br*np.log(np.maximum(br,1e-300)/np.maximum(mr,1e-300)),0),-1)
|
|
sink=[]
|
|
for bi in range(a.shape[0]):
|
|
which=valid[bi,:,7:];values=cond[bi,:,7:]
|
|
sink.append(float(np.maximum(values[which],0).mean()) if which.any() else None)
|
|
diag=np.stack([a[:,:,t,t] for t in range(7,a.shape[-2])],-1).mean((1,2))
|
|
local=np.stack([a[:,:,t,max(0,t-2):t+1].sum(-1) for t in range(7,a.shape[-2])],-1).mean((1,2))
|
|
return {'js':eligible.mean((1,2)),'js_p95':np.quantile(eligible.reshape(a.shape[0],-1),.95,axis=1),
|
|
'top5':np.stack(overlaps,-1).mean((1,2)),'js_all_queries':js.mean((1,2)),
|
|
'first_key_mass_a':a[:,:,7:,0].mean((1,2)),'first_key_mass_b':b[:,:,7:,0].mean((1,2)),
|
|
'self_mass_a':diag,'recent3_mass_a':local,'js_without_first_key':sink}
|
|
|
|
|
|
def cross_head_output_metrics(a,b):
|
|
"""All head pairs, per text. a/b: [B,T,H,Dhidden]."""
|
|
aa=a.transpose(0,2,1,3).reshape(a.shape[0],a.shape[2],-1).astype(np.float64)
|
|
bb=b.transpose(0,2,1,3).reshape(b.shape[0],b.shape[2],-1).astype(np.float64)
|
|
xx=np.sum(aa*aa,-1);yy=np.sum(bb*bb,-1)
|
|
dot=np.einsum('bif,bjf->bij',aa,bb,optimize=True)
|
|
diff=np.maximum(xx[:,:,None]+yy[:,None,:]-2*dot,0)
|
|
return {'error_a':np.sqrt(diff/np.maximum(xx[:,:,None],1e-30)),
|
|
'error_b':np.sqrt(diff/np.maximum(yy[:,None,:],1e-30)),
|
|
'cosine':dot/np.maximum(np.sqrt(xx[:,:,None]*yy[:,None,:]),1e-30),
|
|
'rms_a':np.broadcast_to(np.sqrt(xx/aa.shape[-1])[:,:,None],dot.shape),
|
|
'rms_b':np.broadcast_to(np.sqrt(yy/bb.shape[-1])[:,None,:],dot.shape)}
|
|
|
|
|
|
def value_gram_cka(a,b):
|
|
"""Centered token Gram agreement, diagnostic only; not interchangeability."""
|
|
a=a.astype(np.float64)-a.mean(axis=1,keepdims=True,dtype=np.float64)
|
|
b=b.astype(np.float64)-b.mean(axis=1,keepdims=True,dtype=np.float64)
|
|
ga=np.einsum('btd,bsd->bts',a,a);gb=np.einsum('btd,bsd->bts',b,b)
|
|
return np.sum(ga*gb,(1,2))/np.maximum(np.sqrt(np.sum(ga*ga,(1,2))*np.sum(gb*gb,(1,2))),1e-30)
|
|
|
|
|
|
class QKVModel(StreamingLM):
|
|
def project(self,context,gate,layer):
|
|
return self.linear((context*sigmoid(gate)).reshape(*context.shape[:2],-1),f'layers.{layer}.self_attn.o_proj')
|
|
|
|
def project_heads(self,context,gate,layer):
|
|
"""Each head's contribution in the common residual output coordinates."""
|
|
c=context*sigmoid(gate);b,t,h,d=c.shape
|
|
weight=self.tensors[PREFIX+f'layers.{layer}.self_attn.o_proj.weight']
|
|
outdim,indim=weight.shape;assert indim==h*d
|
|
out=np.empty((b,t,h,outdim),dtype=np.float32)
|
|
for start in range(0,outdim,256):
|
|
self.guard();end=min(start+256,outdim)
|
|
w=weight.read(slice(start*indim,end*indim)).reshape(end-start,indim)
|
|
for head in range(h):out[:,:,head,start:end]=c[:,:,head,:]@w[:,head*d:(head+1)*d].T
|
|
del w;weight.evict()
|
|
return out
|
|
|
|
def parts(self,h,layer,heads=True):
|
|
p=f'layers.{layer}.self_attn.';c=self.config;d=c['head_dim'];nh=c['num_attention_heads'];nk=c['num_key_value_heads']
|
|
mixed=self.linear(h,p+'q_proj').reshape(*h.shape[:2],nh,2*d)
|
|
q,gate=mixed[...,:d],mixed[...,d:].copy()
|
|
k=self.linear(h,p+'k_proj').reshape(*h.shape[:2],nk,d)
|
|
v=self.linear(h,p+'v_proj').reshape(*h.shape[:2],nk,d)
|
|
q=self.normalization(q,p+'q_norm');k=self.normalization(k,p+'k_norm')
|
|
r=c['rope_parameters'];q,k=(rope(x,r['rope_theta'],r['partial_rotary_factor']) for x in (q,k))
|
|
context,probs=attention(q,k,v)
|
|
out={'q':q,'k':k,'p':probs,'v':v,'gate':gate,'context':context,'output':self.project(context,gate,layer)}
|
|
if heads:out['heads']=self.project_heads(context,gate,layer)
|
|
return out
|
|
|
|
def transplant(self,own,donor,layer,kind,target_head=None,donor_head=None,perm=None):
|
|
repeats=self.config['num_attention_heads']//self.config['num_key_value_heads']
|
|
if kind=='full_attention':return donor['output']
|
|
if kind=='qkv_head':return own['output']-own['heads'][:,:,target_head]+donor['heads'][:,:,donor_head]
|
|
probs=own['p'].copy();v=own['v'].copy()
|
|
if kind=='qk_map':probs=donor['p'][:,perm]
|
|
elif kind=='qk_head':probs[:,target_head]=donor['p'][:,donor_head]
|
|
elif kind=='v_group':v[:,:,target_head]=donor['v'][:,:,donor_head]
|
|
else:raise ValueError(kind)
|
|
context=np.einsum('bhts,bshd->bthd',probs,np.repeat(v,repeats,axis=2))
|
|
return self.project(context,own['gate'],layer)
|
|
|
|
def head_transplant_output(self,own,donor,layer,target_head,donor_head):
|
|
"""QK-only transplanted target head contribution; all other heads omitted."""
|
|
repeats=self.config['num_attention_heads']//self.config['num_key_value_heads']
|
|
c=own['context'].copy()
|
|
c[:,:,target_head]=np.einsum('bts,bsd->btd',donor['p'][:,donor_head],own['v'][:,:,target_head//repeats])
|
|
# Zero remaining head contexts before projection to isolate this learned subfunction.
|
|
for head in range(c.shape[2]):
|
|
if head!=target_head:c[:,:,head]=0
|
|
return self.project(c,own['gate'],layer)
|
|
|
|
def value_group_output(self,own,donor,layer,target_group,donor_group):
|
|
repeats=self.config['num_attention_heads']//self.config['num_key_value_heads'];c=np.zeros_like(own['context'])
|
|
for head in range(target_group*repeats,(target_group+1)*repeats):
|
|
c[:,:,head]=np.einsum('bts,bsd->btd',own['p'][:,head],donor['v'][:,:,donor_group])
|
|
return self.project(c,own['gate'],layer)
|
|
|
|
|
|
class InterventionLM(QKVModel):
|
|
def configure(self,candidate,reverse=False,cached=False):
|
|
self.candidate=candidate;self.reverse=reverse;self.use_cache=cached;self.donor_cache=None
|
|
self.target_layer=candidate['j'] if reverse else candidate['i']
|
|
self.donor_layer=candidate['i'] if reverse else candidate['j']
|
|
self.target_head=candidate.get('b') if reverse else candidate.get('a')
|
|
self.donor_head=candidate.get('a') if reverse else candidate.get('b')
|
|
self.perm=inverse(candidate['perm']) if reverse and candidate.get('perm') else candidate.get('perm')
|
|
if cached:assert self.donor_layer<self.target_layer,'Cannot reuse future layer state'
|
|
|
|
def full_attention(self,h,layer,shared_p=None):
|
|
if not hasattr(self,'candidate'):return super().full_attention(h,layer,shared_p)
|
|
if self.use_cache and layer==self.donor_layer:
|
|
self.donor_cache=self.parts(h,layer,heads=self.candidate['kind']=='qkv_head')
|
|
return self.donor_cache['output'],self.donor_cache['p']
|
|
if layer!=self.target_layer:return super().full_attention(h,layer,shared_p)
|
|
own=self.parts(h,layer,heads=self.candidate['kind']=='qkv_head')
|
|
donor=self.donor_cache if self.use_cache else self.parts(h,self.donor_layer,heads=self.candidate['kind']=='qkv_head')
|
|
assert donor is not None
|
|
out=self.transplant(own,donor,layer,self.candidate['kind'],self.target_head,self.donor_head,self.perm)
|
|
return out,own['p']
|