86 lines
5.4 KiB
Python
86 lines
5.4 KiB
Python
"""Frozen-model head interventions and continuous influence concentration."""
|
|
import numpy as np
|
|
from functional.model import StreamingLM,PREFIX,rope,causal_conv,attention
|
|
|
|
|
|
def concentration(d):
|
|
"""Last axis donor head; zero total is undefined, represented by NaNs."""
|
|
d=np.asarray(d,dtype=np.float64)
|
|
if not np.isfinite(d).all() or np.any(d<0):raise ValueError('Finite nonnegative influences required')
|
|
h=d.shape[-1];total=d.sum(-1);valid=total>0
|
|
p=np.divide(d,total[...,None],out=np.zeros_like(d),where=valid[...,None]);square=np.sum(p*p,-1)
|
|
neff=np.divide(1,square,out=np.full_like(total,np.nan),where=valid)
|
|
ordered=np.sort(p,axis=-1)[...,::-1]
|
|
return {'total':total,'max_absolute':d.max(-1),'valid':valid,'effective_heads':neff,'effective_fraction':neff/h,
|
|
'concentration':np.where(valid,(square-1/h)/(1-1/h),np.nan),
|
|
'top1':np.where(valid,ordered[...,0],np.nan),'top25':np.where(valid,ordered[...,:max(1,h//4)].sum(-1),np.nan),
|
|
'top50':np.where(valid,ordered[...,:max(1,h//2)].sum(-1),np.nan),
|
|
'bottom25':np.where(valid,ordered[...,-max(1,h//4):].sum(-1),np.nan),
|
|
'heads90':np.where(valid,1+np.sum(np.cumsum(ordered,axis=-1)<.9,axis=-1),np.nan)}
|
|
|
|
|
|
def masks_from_influences(d,weak=True):
|
|
"""d [donor,T,recipient]; output masks [recipient,T,donor]."""
|
|
ranked=np.argsort(d.transpose(2,1,0),axis=-1,kind='stable');n=d.shape[0]//2
|
|
selected=ranked[...,:n] if weak else ranked[...,-n:]
|
|
mask=np.zeros_like(ranked,dtype=np.float32);np.put_along_axis(mask,selected,1,axis=-1)
|
|
return mask
|
|
|
|
|
|
def js_rows(a,b):
|
|
a=np.asarray(a,dtype=np.float64);b=np.asarray(b,dtype=np.float64);m=(a+b)/2
|
|
def term(x):return np.where(x>0,x*np.log(np.maximum(x,1e-300)/np.maximum(m,1e-300)),0)
|
|
return np.maximum(0,.5*np.sum(term(a)+term(b),axis=-1))
|
|
|
|
|
|
class InfluenceLM(StreamingLM):
|
|
capturing=False
|
|
def linear(self,x,name,rows=256):
|
|
y=super().linear(x,name,rows)
|
|
if self.capturing and (name.endswith('self_attn.o_proj') or name.endswith('linear_attn.out_proj')):
|
|
layer=int(name.split('.')[1]);nh=self.config['num_attention_heads'] if 'self_attn' in name else self.config['linear_num_value_heads']
|
|
xx=x.reshape(*x.shape[:2],nh,-1);dim=xx.shape[-1];t=self.tensors[PREFIX+name+'.weight'];od,idim=t.shape
|
|
out=np.empty((*x.shape[:2],nh,od),dtype=np.float32)
|
|
for start in range(0,od,256):
|
|
self.guard();stop=min(od,start+256);w=t.read(slice(start*idim,stop*idim)).reshape(stop-start,idim)
|
|
for head in range(nh):out[:,:,head,start:stop]=xx[:,:,head]@w[:,head*dim:(head+1)*dim].T
|
|
del w;t.evict()
|
|
err=float(np.linalg.norm(out.sum(2).astype(np.float64)-y)/max(np.linalg.norm(y.astype(np.float64)),1e-30))
|
|
if err>2e-5:raise AssertionError(('head sum mismatch',layer,err))
|
|
self.head_outputs=out;self.head_sum_errors.append({'layer':layer,'relative_l2':err})
|
|
return y
|
|
|
|
def baseline(self,ids,root):
|
|
h=self.embedding(ids);self.head_sum_errors=[]
|
|
for layer,kind in enumerate(self.config['layer_types']):
|
|
raw=h.copy();n=self.normalization(h,f'layers.{layer}.input_layernorm');self.capturing=True
|
|
out,_=self.full_attention(n,layer) if kind=='full_attention' else (self.linear_attention(n,layer),None)
|
|
self.capturing=False;heads=self.head_outputs;self.head_outputs=None
|
|
mid=raw+out;h=mid+self.mlp(self.normalization(mid,f'layers.{layer}.post_attention_layernorm'),layer)
|
|
np.savez(root/f'layer-{layer}.npz',raw=raw,out=out,heads=heads,next_raw=h)
|
|
print('baseline layer',layer,flush=True)
|
|
return h
|
|
|
|
def recipient(self,raw,layer):
|
|
h=self.normalization(raw,f'layers.{layer}.input_layernorm');c=self.config;b,t=h.shape[:2]
|
|
if c['layer_types'][layer]=='full_attention':
|
|
p=f'layers.{layer}.self_attn.';nh=c['num_attention_heads'];nk=c['num_key_value_heads'];d=c['head_dim']
|
|
q=self.linear(h,p+'q_proj').reshape(b,t,nh,2*d)[...,:d];k=self.linear(h,p+'k_proj').reshape(b,t,nk,d);v=self.linear(h,p+'v_proj').reshape(b,t,nk,d)
|
|
q=self.normalization(q,p+'q_norm');k=self.normalization(k,p+'k_norm');r=c['rope_parameters'];q,k=(rope(z,r['rope_theta'],r['partial_rotary_factor']) for z in (q,k))
|
|
k=np.repeat(k,nh//nk,axis=2);v=np.repeat(v,nh//nk,axis=2)
|
|
_,prob=attention(q,k,v)
|
|
return {'Q':q,'K':k,'V':v,'P':prob}
|
|
p=f'layers.{layer}.linear_attn.';kh=c['linear_num_key_heads'];vh=c['linear_num_value_heads'];kd=c['linear_key_head_dim'];vd=c['linear_value_head_dim']
|
|
x=causal_conv(self.linear(h,p+'in_proj_qkv'),self.small(p+'conv1d.weight')[:,0,:]);q,k,v=np.split(x,[kh*kd,2*kh*kd],axis=-1)
|
|
q=q.reshape(b,t,kh,kd);k=k.reshape(b,t,kh,kd);v=v.reshape(b,t,vh,vd)
|
|
if vh!=kh:q,k=np.repeat(q,vh//kh,axis=2),np.repeat(k,vh//kh,axis=2)
|
|
q=q/np.sqrt(np.sum(q*q,axis=-1,keepdims=True)+1e-6)/np.float32(np.sqrt(kd));k=k/np.sqrt(np.sum(k*k,axis=-1,keepdims=True)+1e-6)
|
|
return {'Q':q,'K':k,'V':v}
|
|
|
|
def intervene(self,cache,layer,masks):
|
|
# Each scenario has a full coherent token-dependent deletion mask.
|
|
removed=np.einsum('bth,thd->btd',masks,cache['heads'][0],optimize=True)
|
|
mid=cache['raw']+cache['out']-removed
|
|
nxt=mid+self.mlp(self.normalization(mid,f'layers.{layer}.post_attention_layernorm'),layer)
|
|
return self.recipient(nxt,layer+1)
|