135 lines
9.9 KiB
Python
135 lines
9.9 KiB
Python
"""Find repeated learned scalar MLP functions on real hidden inputs, no training.
|
|
|
|
Four fixed layers; calibration-only selection; one frozen pair; independent heldout
|
|
functional test and two actual network substitutions. No scale fitting or R.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import resource
|
|
import time
|
|
import numpy as np
|
|
from .model import StreamingLM,silu
|
|
from .run import write_json,save_csv
|
|
|
|
LAYERS=[5,6,20,21]
|
|
|
|
|
|
def pair_metrics(x,y):
|
|
x,y=x.astype(np.float64),y.astype(np.float64)
|
|
xx,yy=float(x@x),float(y@y);delta=float((x-y)@(x-y))
|
|
xc,yc=x-x.mean(),y-y.mean()
|
|
return {'relative_l2_a_to_b':float(np.sqrt(delta/max(xx,1e-30))),
|
|
'relative_l2_b_to_a':float(np.sqrt(delta/max(yy,1e-30))),
|
|
'cosine':float(x@y/max(np.sqrt(xx*yy),1e-30)),
|
|
'pearson':float(xc@yc/max(np.linalg.norm(xc)*np.linalg.norm(yc),1e-30)),
|
|
'rms_a':float(np.sqrt(xx/x.size)),'rms_b':float(np.sqrt(yy/y.size)),
|
|
'std_over_rms_a':float(np.std(x)/max(np.sqrt(xx/x.size),1e-30)),
|
|
'std_over_rms_b':float(np.std(y)/max(np.sqrt(yy/y.size),1e-30)),
|
|
'effective_energy_tokens_a':float(xx*xx/max(float((x**4).sum()),1e-30)),
|
|
'effective_energy_tokens_b':float(yy*yy/max(float((y**4).sum()),1e-30))}
|
|
|
|
|
|
class NeuronReuseLM(StreamingLM):
|
|
def configure(self,layer,target,donor):self.reused_layer,self.target,self.donor=layer,target,donor
|
|
def mlp(self,h,layer):
|
|
p=f'layers.{layer}.mlp.'
|
|
z=silu(self.linear(h,p+'gate_proj'))*self.linear(h,p+'up_proj')
|
|
if layer==self.reused_layer:z[...,self.target]=z[...,self.donor]
|
|
return self.linear(z,p+'down_proj')
|
|
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--root',type=Path,required=True);a=p.parse_args();root=a.root
|
|
protocol={'layers':LAYERS,'unit':'post_silu_gate_times_up_scalar_function; keep each output column',
|
|
'thresholds':[.05,.10,.20],'maximum_directional_relative_l2':.20,'cosine_min':.98,'pearson_min':.95,
|
|
'rms_floor':'at least within-layer median calibration neuron RMS','std_over_rms_min':.25,
|
|
'effective_energy_tokens_min':16,'max_selected_pairs':1,'selection':'smallest worst directional calibration error; then layer/id',
|
|
'heldout_per_text_error_max':.35,'heldout_quality_mean_nll_delta_max':.01,'heldout_quality_sequence_p95_delta_max':.05,
|
|
'training':False,'scale_fitting':False,'cost_required':False}
|
|
protocol_path=root/'neuron-reuse-protocol.json'
|
|
if protocol_path.exists():
|
|
old=json.loads(protocol_path.read_text());assert all(old[k]==v for k,v in protocol.items())
|
|
else:write_json(protocol_path,{**protocol,'registered_unix':time.time()})
|
|
ledger_path=root/'compute-budget-ledger.json';ledger=json.loads(ledger_path.read_text());remaining=int(7200-ledger['spent_seconds']-30)
|
|
if remaining<120:raise TimeoutError('Remaining cumulative budget insufficient')
|
|
started=time.monotonic();result={'state':'running','selected':None,'function_pass':False,'interventions':[]}
|
|
try:
|
|
cfg=json.loads((root/'model_config.json').read_text());lm=NeuronReuseLM(root/'weights',cfg,seconds=remaining)
|
|
data=np.load(root/'baseline-activations.npz');samples=json.loads((root/'samples-tokenized.json').read_text())['samples']
|
|
cal=np.array([i for i,s in enumerate(samples) if s['split']=='calibration']);held=np.array([i for i,s in enumerate(samples) if s['split']=='heldout'])
|
|
activations={};summaries=[];best=None
|
|
for layer in LAYERS:
|
|
h=data[f'mlp_{layer}'];pfx=f'layers.{layer}.mlp.'
|
|
z=silu(lm.linear(h,pfx+'gate_proj'))*lm.linear(h,pfx+'up_proj');activations[layer]=z
|
|
# Profiles use all calibration tokens; never use heldout in selection.
|
|
x=np.ascontiguousarray(z[cal].reshape(-1,z.shape[-1]).T,dtype=np.float64)
|
|
n=x.shape[1];normsq=(x*x).sum(1);means=x.mean(1);variance=np.maximum(normsq-n*means**2,0)
|
|
rms=np.sqrt(normsq/n);floor=float(np.median(rms));effective=normsq**2/np.maximum((x**4).sum(1),1e-30)
|
|
active=(rms>=floor)&(np.sqrt(variance/np.maximum(normsq,1e-30))>=.25)&(effective>=16)
|
|
counts={str(t):0 for t in protocol['thresholds']};tested=0;layerbest=None
|
|
for start in range(0,x.shape[0],256):
|
|
lm.guard();end=min(start+256,x.shape[0]);dot=x[start:end]@x.T
|
|
delta=np.maximum(normsq[start:end,None]+normsq[None,:]-2*dot,0)
|
|
error=np.sqrt(delta/np.maximum(np.minimum(normsq[start:end,None],normsq[None,:]),1e-30))
|
|
cosine=dot/np.maximum(np.sqrt(normsq[start:end,None]*normsq[None,:]),1e-30)
|
|
pearson=(dot-n*means[start:end,None]*means[None,:])/np.maximum(np.sqrt(variance[start:end,None]*variance[None,:]),1e-30)
|
|
unique=np.arange(start,end)[:,None]<np.arange(x.shape[0])[None,:]
|
|
allowed=unique&active[start:end,None]&active[None,:]&(cosine>=.98)&(pearson>=.95)
|
|
tested+=int((unique&active[start:end,None]&active[None,:]).sum())
|
|
for threshold in protocol['thresholds']:counts[str(threshold)]+=int((allowed&(error<=threshold)).sum())
|
|
mask=allowed&(error<=.20)
|
|
if mask.any():
|
|
score=np.where(mask,error,np.inf);row,col=np.unravel_index(np.argmin(score),score.shape)
|
|
candidate={'layer':layer,'a':int(start+row),'b':int(col),'calibration_worst_relative_l2':float(score[row,col])}
|
|
key=lambda c:(c['calibration_worst_relative_l2'],c['layer'],c['a'],c['b'])
|
|
if layerbest is None or key(candidate)<key(layerbest):layerbest=candidate
|
|
if best is None or key(candidate)<key(best):best=candidate
|
|
summaries.append({'layer':layer,'neurons':int(x.shape[0]),'eligible_neurons':int(active.sum()),'eligible_unique_pairs_checked':tested,'rms_floor':floor,'calibration_pair_counts':counts,'best':layerbest})
|
|
write_json(root/'neuron-reuse-search-summary.json',summaries)
|
|
print(json.dumps(summaries[-1]),flush=True)
|
|
np.savez(root/'neuron-reuse-activations.npz',**{f'layer{k}':v for k,v in activations.items()})
|
|
write_json(root/'neuron-reuse-selection.json',{'selected':best,'selection_uses_heldout':False})
|
|
result['selected']=best
|
|
if best is not None:
|
|
layer,i,j=best['layer'],best['a'],best['b'];z=activations[layer]
|
|
metrics=[]
|
|
for split,index in [('calibration',cal),('heldout',held)]:
|
|
metrics.append({'split':split,'scope':'pooled',**pair_metrics(z[index,:,i].reshape(-1),z[index,:,j].reshape(-1))})
|
|
for idx in index:metrics.append({'split':split,'scope':samples[idx]['id'],**pair_metrics(z[idx,:,i],z[idx,:,j])})
|
|
save_csv(root/'neuron-reuse-function-metrics.csv',metrics)
|
|
pooled=next(r for r in metrics if r['split']=='heldout' and r['scope']=='pooled')
|
|
floor=next(r['rms_floor'] for r in summaries if r['layer']==layer)
|
|
function_pass=(max(pooled['relative_l2_a_to_b'],pooled['relative_l2_b_to_a'])<=.20 and pooled['cosine']>=.98 and pooled['pearson']>=.95 and
|
|
min(pooled['std_over_rms_a'],pooled['std_over_rms_b'])>=.25 and min(pooled['effective_energy_tokens_a'],pooled['effective_energy_tokens_b'])>=16 and
|
|
min(pooled['rms_a'],pooled['rms_b'])>=floor and
|
|
all(max(r['relative_l2_a_to_b'],r['relative_l2_b_to_a'])<=.35 for r in metrics if r['split']=='heldout' and r['scope']!='pooled'))
|
|
result.update(function_pass=bool(function_pass),heldout_function_metrics=pooled)
|
|
# Show parameter similarity separately from actual functional similarity.
|
|
wd=[]
|
|
for role in ('gate_proj','up_proj'):
|
|
t=lm.tensors[f'model.language_model.layers.{layer}.mlp.{role}.weight'];d=t.shape[1]
|
|
wa,wb=t.read(slice(i*d,(i+1)*d)),t.read(slice(j*d,(j+1)*d))
|
|
wd.append({'role':role,'weight_cosine':float(wa.astype(np.float64)@wb/max(np.linalg.norm(wa.astype(np.float64))*np.linalg.norm(wb.astype(np.float64)),1e-30))});t.evict()
|
|
result['weight_diagnostics']=wd
|
|
if function_pass:
|
|
ids=np.array([s['token_ids'] for s in samples],dtype=np.int64);base=np.load(root/'baseline-heldout-token-nll.npy').mean(1)
|
|
for target,donor in ((i,j),(j,i)):
|
|
lm.configure(layer,target,donor);h,_=lm.forward(ids[held]);nll=lm.nll(h,ids[held])
|
|
np.save(root/f'neuron-reuse-{target}-uses-{donor}-token-nll.npy',nll)
|
|
delta=nll.mean(1)-base
|
|
result['interventions'].append({'target':target,'donor':donor,'sequence_nll_deltas':delta.tolist(),'mean_nll_delta':float(delta.mean()),'sequence_p95_delta':float(np.quantile(delta,.95)),'passed':bool(delta.mean()<=.01 and np.quantile(delta,.95)<=.05)})
|
|
write_json(root/'neuron-reuse-summary.json',result)
|
|
# No fallback after looking at heldout.
|
|
result.update(state='completed',empirical_repeated_function_positive=bool(result['function_pass'] and len(result['interventions'])==2 and all(r['passed'] for r in result['interventions'])),
|
|
claim_limit='Distinct learned scalar MLP subfunctions on the specified short-text distribution; not whole-layer duplication, global equivalence or broad task-capability proof. No performance requirement.')
|
|
except Exception as exc:result.update(state='stopped',error_type=type(exc).__name__,error=str(exc));raise
|
|
finally:
|
|
elapsed=time.monotonic()-started;result.update(elapsed_seconds=elapsed,peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)
|
|
ledger=json.loads(ledger_path.read_text());ledger['spent_seconds']+=elapsed;ledger['calls'].append({'module':'functional.neuron_reuse','seconds':elapsed});write_json(ledger_path,ledger)
|
|
write_json(root/'neuron-reuse-summary.json',result)
|
|
|
|
|
|
if __name__=='__main__':main()
|