85 lines
7.2 KiB
Python
85 lines
7.2 KiB
Python
"""All adjacent layers, input-dependent influence distributions; CPU only."""
|
|
import argparse,json,time,resource,hashlib,traceback
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from tokenizers import Tokenizer
|
|
from functional.model import StreamingLM
|
|
from .core import InfluenceLM,concentration,masks_from_influences,js_rows
|
|
|
|
|
|
def dump(path,data):path.write_text(json.dumps(data,indent=2,allow_nan=False)+'\n')
|
|
|
|
|
|
def distances(base,changed):
|
|
out={}
|
|
for key in ('Q','K','V'):
|
|
a=base[key].astype(np.float64);b=changed[key].astype(np.float64)
|
|
out[key]=np.linalg.norm(b-a,axis=-1)/np.maximum(np.linalg.norm(a,axis=-1),1e-12)
|
|
out['QKV']=np.sqrt(sum(out[k]**2 for k in ('Q','K','V'))/3)
|
|
if 'P' in base:out['JS']=js_rows(base['P'],changed['P']).transpose(0,2,1)
|
|
return out
|
|
|
|
|
|
def link(lm,cache,layer,path,seed):
|
|
started=time.monotonic();base=lm.recipient(cache['next_raw'],layer+1);hs=cache['heads'].shape[2];t=cache['raw'].shape[1];hr=base['Q'].shape[2]
|
|
# Whole donor head across the full sequence is the single-ablation intervention.
|
|
singles={k:[] for k in ('Q','K','V','QKV')+ (('JS',) if 'P' in base else ())}
|
|
for start in range(0,hs,4):
|
|
index=list(range(start,min(start+4,hs)));masks=np.zeros((len(index),t,hs),dtype=np.float32)
|
|
for n,h in enumerate(index):masks[n,:,h]=1
|
|
scores=distances(base,lm.intervene(cache,layer,masks))
|
|
for k,v in scores.items():singles[k].extend(v)
|
|
singles={k:np.asarray(v,dtype=np.float32) for k,v in singles.items()}
|
|
# Identical cardinality weak/strong/random removals, ranking by balanced QKV influence.
|
|
# One coherent full-sequence mask per recipient head; only that head's output scored.
|
|
rank=singles['QKV'];rng=np.random.default_rng(seed)
|
|
weak=masks_from_influences(rank,True);strong=masks_from_influences(rank,False)
|
|
random=np.zeros_like(weak)
|
|
for r in range(hr):
|
|
for q in range(t):random[r,q,rng.choice(hs,hs//2,replace=False)]=1
|
|
joint={};sum_individual={}
|
|
for label,masks in [('weak_half',weak),('strong_half',strong),('random_half',random),('all_heads',np.ones_like(weak))]:
|
|
values={k:[] for k in singles}
|
|
for start in range(0,hr,4):
|
|
stop=min(start+4,hr);scores=distances(base,lm.intervene(cache,layer,masks[start:stop]))
|
|
for k,v in scores.items():
|
|
for n,r in enumerate(range(start,stop)):values[k].append(v[n,:,r])
|
|
for k,v in values.items():
|
|
joint[label+'_'+k]=np.asarray(v,dtype=np.float32).T
|
|
sum_individual[label+'_'+k]=np.sum(singles[k].transpose(2,1,0)*masks,axis=-1).T
|
|
null=distances(base,lm.intervene(cache,layer,np.zeros((1,t,hs),dtype=np.float32)))
|
|
if max(float(np.max(v)) for v in null.values())>2e-5:raise AssertionError('Null intervention changed recipient')
|
|
payload={'single_'+k:v for k,v in singles.items()};payload.update({'joint_'+k:v for k,v in joint.items()});payload.update({'individual_sum_'+k:v for k,v in sum_individual.items()})
|
|
payload.update({'mask_weak_half':weak,'mask_strong_half':strong,'mask_random_half':random})
|
|
payload.update({'null_'+k:v for k,v in null.items()})
|
|
np.savez_compressed(path,**payload)
|
|
return {'source':layer,'recipient':layer+1,'source_heads':hs,'recipient_query_heads':hr,'full_attention_recipient':'P' in base,'elapsed_seconds':time.monotonic()-started,'null_max':max(float(np.max(v)) for v in null.values())}
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--out',type=Path,required=True);p.add_argument('--seconds',type=int,default=7200);a=p.parse_args()
|
|
a.out.mkdir(parents=True,exist_ok=False);(a.out/'cache').mkdir();(a.out/'raw').mkdir();started=time.monotonic()
|
|
protocol={'scope':'all 23 adjacent layer transitions; source attention head contribution to next effective Q/K/V','samples':'4 previously authored zh/en texts, first32 tokens, descriptive','primary_queries':list(range(7,32)), 'influence':'L2 change / baseline token recipient-head L2; QKV score sqrt(mean of three squared relative changes))','linear_QK':'after causal convolution/SiLU and delta-rule L2 normalization; Q additionally divided by sqrt(key dim)','full_QK':'QK norm and partial RoPE applied; K/V repeated to query heads for indexing (not independent KV observations)','single_ablation':'remove complete source head projected output at all tokens; recompute source postattention normalization, MLP, residual and next input normalization','joint_ablation':'weak/strong/random half and all attention heads; one coherent token-varying removal mask for each recipient query head; ranking uses QKV single effects; no fitting','continuous':'effective heads=1/sum(p^2); fraction=effective/H; concentration=(sum(p^2)-1/H)/(1-1/H)','zero_total':'undefined concentration (NaN arrays), count explicitly','cost_budget_seconds':a.seconds,'no_training':True,'keep_all_weights_and_caches':True,'limitations':['not final LM quality or pruning proof','individual effects not additive; normalized scores not true contribution probabilities','token-varying masks are input-specific diagnostic and not a deployable fixed head subset','linear attention does not have an ordinary softmax probability matrix','Q/K/V scales handled separately; combined QKV is a declared balanced relative-effect score']}
|
|
dump(a.out/'protocol.json',protocol);status={'state':'running','completed_links':[]};dump(a.out/'status.json',status)
|
|
try:
|
|
lm=InfluenceLM(a.source/'weights',json.loads((a.source/'model_config.json').read_text()),seconds=a.seconds)
|
|
samples=json.loads(Path('crossqkv/fresh_samples.json').read_text())['samples'];tok=Tokenizer.from_file(str(a.source/'tokenizer.json'));controls=[]
|
|
for si,s in enumerate(samples):
|
|
lm.guard();s['token_ids']=tok.encode(s['text']).ids[:32];ids=np.array([s['token_ids']],dtype=np.int64);cache=a.out/'cache'/s['id'];cache.mkdir()
|
|
final=lm.baseline(ids,cache);np.save(cache/'final.npy',final)
|
|
for x in lm.head_sum_errors:controls.append({'sample':s['id'],**x})
|
|
if si==0:
|
|
reference,_=StreamingLM.forward(lm,ids);normalized=lm.normalization(final,'norm');err=float(np.linalg.norm(normalized-reference)/max(np.linalg.norm(reference),1e-30));assert err<2e-5;status['canonical_forward_relative_error']=err
|
|
dump(a.out/'head-sum-controls.json',controls);dump(a.out/'samples.json',samples)
|
|
for layer in range(23):
|
|
target=a.out/'raw'/f'{s["id"]}-link-{layer:02d}.npz'
|
|
with np.load(cache/f'layer-{layer}.npz') as data:loaded={k:data[k] for k in data.files}
|
|
info=link(lm,loaded,layer,target,31000+si*100+layer)
|
|
del loaded
|
|
status['completed_links'].append({'sample':s['id'],**info});status['elapsed_seconds']=time.monotonic()-started;dump(a.out/'status.json',status)
|
|
print('completed',s['id'],layer,'seconds',round(info['elapsed_seconds'],2),flush=True)
|
|
status.update(state='completed',elapsed_seconds=time.monotonic()-started,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss);dump(a.out/'status.json',status)
|
|
except BaseException as e:
|
|
status.update(state='failed',error=repr(e),elapsed_seconds=time.monotonic()-started);dump(a.out/'status.json',status);raise
|
|
if __name__=='__main__':main()
|