"""Audit the saved SMS model with exact traces and interventions. From the repository root: .venv/bin/python blog/src/content/drafts/feature-detectors-how-neurons-read-and-write/experiment/analyze_interpretability.py Requires the training example's PyTorch environment plus matplotlib. Writes a JSON audit, per-message activation archive, and PNG figure alongside this script; downloads data into /tmp. The activation archive covers both train and test messages at the saved final checkpoint, not intermediate training epochs. The initialization sweep fixes both data split and minibatch order to seed 42. """ import argparse from collections import Counter import contextlib import hashlib import io import itertools import json from pathlib import Path import torch from torch.nn import functional as F from torch.utils.data import DataLoader import train_sms as sms HERE = Path(__file__).resolve().parent def metrics(logits, labels): pred = logits > 0 tn = int((~pred & (labels == 0)).sum()) fp = int((pred & (labels == 0)).sum()) fn = int((~pred & (labels == 1)).sum()) tp = int((pred & (labels == 1)).sum()) return dict(accuracy=(tn + tp) / len(labels), cross_entropy=F.binary_cross_entropy_with_logits(logits, labels).item(), spam_precision=tp / (tp + fp) if tp + fp else 0, spam_recall=tp / (tp + fn), confusion=[[tn, fp], [fn, tp]]) def trace(model, x): z = model.W(x) h = z.relu() contributions = h * model.V.weight[0] s = contributions.sum(-1) + model.V.bias[0] torch.testing.assert_close(s, model(x)) return z, h, contributions, s def trace_rows(model, x): z, h, c, s = trace(model, x) return [dict(counts=xx.tolist(), preactivations=zz.tolist(), activations=hh.tolist(), contributions=cc.tolist(), logit=ss.item(), probability=ss.sigmoid().item()) for xx, zz, hh, cc, ss in zip(x, z, h, c, s)] def pattern_table(x, y): counts = Counter((tuple(map(int, xx.tolist())), int(yy)) for xx, yy in zip(x, y)) return [dict(counts=list(p), nonspam=counts[p, 0], spam=counts[p, 1]) for p in sorted({key[0] for key in counts})] def export_activations(model, splits, saved, output): """Save final-checkpoint traces with the inputs needed to interpret them.""" records = [] for split, rows, data in splits: x, labels = data.tensors traces = trace_rows(model, x) for index, ((message, label), row) in enumerate(zip(rows, traces)): assert label == int(labels[index]) identity = " ".join(message.lower().split()) row.update(split=split, split_index=index, message=message, message_sha256=hashlib.sha256(identity.encode()).hexdigest(), label=label, prediction=int(row['logit'] > 0)) records.append(row) archive = dict( schema_version=1, checkpoint='final', seed=saved['seed'], epochs=saved['epochs'], torch_version=torch.__version__, dataset_sha256=saved['dataset_sha256'], vocabulary=list(sms.VOCABULARY), parameters=saved['parameters'], output_bias=model.V.bias.item(), description='Final saved model evaluated without updates on train and held-out messages. ' 'Preactivations precede ReLU; activations follow ReLU. ' 'Contributions are activations times outgoing weights; ' 'their sum plus output_bias equals logit. Labels/predictions: 0=non-spam, 1=spam.', records=records) # One complete record per line keeps the archive easy to search and compact. header = json.dumps({k: v for k, v in archive.items() if k != 'records'}, indent=2) output.write_text(header[:-2] + ',\n "records": [\n' + ',\n'.join(' ' + json.dumps(row, ensure_ascii=False) for row in records) + '\n ]\n}\n') restored = json.loads(output.read_text()) assert restored == archive for split, _, data in splits: subset = [row for row in restored['records'] if row['split'] == split] restored_h = torch.tensor([row['activations'] for row in subset]) restored_logits = torch.tensor([row['logit'] for row in subset]) torch.testing.assert_close(model.V(restored_h).squeeze(-1), model(data.tensors[0])) torch.testing.assert_close(restored_logits, model(data.tensors[0])) return dict(file=output.name, count=len(records), checkpoint='final', split_counts={split: len(rows) for split, rows, _ in splits}, sha256=hashlib.sha256(output.read_bytes()).hexdigest(), round_trip_validated=True) def exact_path_attribution(model, x): """Integrate the input gradient exactly along t*x, t in [0,1]. Each ReLU preactivation is beta_i + t*(w_i dot x). Its derivative contributes for exactly the fraction of this segment where it is positive. These are contributions to s(x)-s(0), not to the absolute score. """ beta = model.W.bias delta = x @ model.W.weight.T end = beta + delta active_fraction = torch.zeros_like(delta) active_fraction[(beta > 0) & (end > 0)] = 1 enters = (beta <= 0) & (end > 0) exits = (beta > 0) & (end <= 0) active_fraction[enters] = (end / delta)[enters] active_fraction[exits] = (-beta / delta)[exits] per_neuron_feature = (x[:, None, :] * model.W.weight[None, :, :] * model.V.weight[0][None, :, None] * active_fraction[:, :, None]) attribution = per_neuron_feature.sum(1) actual_delta = model(x) - model(torch.zeros_like(x)) torch.testing.assert_close(attribution.sum(-1), actual_delta, atol=1e-5, rtol=1e-5) return attribution, per_neuron_feature def plot_audit(model, probe_rows, output): import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import numpy as np fig, axes = plt.subplots(1, 2, figsize=(12, 4.4), layout='constrained') weights = model.W.weight.detach().numpy() im = axes[0].imshow(weights, cmap='RdBu_r', vmin=-3, vmax=3, aspect='auto') axes[0].set_xticks(range(3), sms.VOCABULARY) neuron_labels = [f'Neuron {i}\nβ = {model.W.bias[i].item():+.3f}; ' f'v = {model.V.weight[0, i].item():+.3f}' for i in range(2)] axes[0].set_yticks(range(2), neuron_labels) axes[0].set_title('Learned incoming weights (seed 42)') for i in range(2): for j in range(3): axes[0].text(j, i, f'{weights[i,j]:+.3f}', ha='center', va='center') fig.colorbar(im, ax=axes[0], shrink=.75, label='Incoming weight') labels = ['No tracked words', 'claim', 'prize', 'meeting', 'claim + prize + meeting*'] ys = np.arange(len(labels)) for offset, key, color, label in [(-.22, 0, '#d55e00', 'Neuron 0: h₀v₀'), (0, 1, '#0072b2', 'Neuron 1: h₁v₁')]: axes[1].barh(ys + offset, [r['contributions'][key] for r in probe_rows], height=.2, color=color, label=label) axes[1].barh(ys + .22, [model.V.bias.item()] * len(ys), height=.2, color='#666666', label='Output bias') axes[1].scatter([r['logit'] for r in probe_rows], ys, color='black', marker='x', label='Total score', zorder=4) axes[1].set_yticks(ys, labels) axes[1].invert_yaxis() axes[1].axvline(0, color='black', linewidth=.6) axes[1].set_xlabel('Contribution to spam logit (positive favors spam)') axes[1].set_title('Exact writing contributions\n*Combination absent from train and test') axes[1].legend(fontsize=8, loc='upper center', bbox_to_anchor=(.5, -.14), ncol=2) fig.savefig(output, dpi=160) plt.close(fig) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--data-dir', type=Path, default=Path('/tmp/read-write-sms-data')) parser.add_argument('--output', type=Path, default=HERE / 'interpretability-results.json') parser.add_argument('--seeds', type=int, nargs='+', default=list(range(10)) + [42]) args = parser.parse_args() torch.set_num_threads(1) saved_path = HERE / 'example-results.json' saved = json.loads(saved_path.read_text()) rows, original_count, digest = sms.load_messages(args.data_dir) assert digest == saved['dataset_sha256'] train_rows, test_rows = sms.split_messages(rows, saved['seed']) train_data, test_data = sms.dataset(train_rows), sms.dataset(test_rows) x, y = test_data.tensors train_x, train_y = train_data.tensors model = sms.ReadWriteClassifier() model.load_state_dict({k: torch.tensor(v) for k, v in saved['parameters'].items()}) model.eval() with torch.no_grad(): assert sms.evaluate(model, test_data) == saved['held_out'] z, h, c, logits = trace(model, x) ablations = [] for removed in ([0], [1], [0, 1]): intervened_h = h.clone() intervened_h[:, removed] = 0 intervened = model.V(intervened_h).squeeze(-1) expected_delta = -c[:, removed].sum(-1) torch.testing.assert_close(intervened - logits, expected_delta, atol=1e-5, rtol=1e-5) ablations.append(dict(zeroed_neurons=removed, metrics=metrics(intervened, y), changed_predictions=int(((intervened > 0) != (logits > 0)).sum()), max_identity_error=float((intervened-logits-expected_delta).abs().max()))) interventions = [] for j, word in enumerate(sms.VOCABULARY): for intervention in ('remove_all', 'add_one'): changed = x.clone() if intervention == 'remove_all': changed[:, j] = 0 else: changed[:, j] += 1 new_logits = model(changed) delta = new_logits - logits interventions.append(dict(word=word, intervention=intervention, changed_inputs=int((changed != x).any(-1).sum()), changed_predictions=int(((new_logits > 0) != (logits > 0)).sum()), score_delta_min=float(delta.min()), score_delta_max=float(delta.max()), score_delta_mean=float(delta.mean()), counterfactual_metrics_with_original_labels=metrics(new_logits, y))) profiles = [] for neuron in range(2): classes = [] for label in (0, 1): mask = y == label classes.append(dict(label=label, count=int(mask.sum()), active_count=int((h[mask, neuron] > 0).sum()), mean_activation=float(h[mask, neuron].mean()), mean_contribution=float(c[mask, neuron].mean()))) indices = torch.argsort(h[:, neuron], descending=True, stable=True)[:5].tolist() examples = [] for i in indices: example = trace_rows(model, x[i:i+1])[0] example.update(test_index=i, label=int(y[i]), message_excerpt=test_rows[i][0][:160]) examples.append(example) profiles.append(dict(neuron=neuron, classes=classes, top_held_out=examples)) probe_x = torch.tensor([[0,0,0], [1,0,0], [0,1,0], [0,0,1], [1,1,1], [2,1,0], [0,1,1]], dtype=torch.float32) probes = trace_rows(model, probe_x) attribution, per_neuron = exact_path_attribution(model, probe_x) exact_path_attribution(model, x) for i, p in enumerate(probes): p.update(train_support=int((train_x == probe_x[i]).all(-1).sum()), test_support=int((x == probe_x[i]).all(-1).sum()), zero_baseline_feature_attributions=attribution[i].tolist(), zero_baseline_attributions_by_neuron=per_neuron[i].tolist()) grid = torch.tensor(list(itertools.product(range(4), repeat=3)), dtype=torch.float32) exact_path_attribution(model, grid) grid_traces = trace_rows(model, grid) groups = pattern_table(x, y) empirical_ceiling = sum(max(g['spam'], g['nonspam']) for g in groups) / len(y) equivalent_rule = (x[:, 0] + x[:, 1]) > 0 report = dict( provenance=dict(torch_version=torch.__version__, dataset_url=sms.DATA_URL, dataset_sha256=digest, source_rows=original_count, unique_messages=len(rows), split_seed=saved['seed'], train_count=len(train_rows), test_count=len(test_rows), saved_report_sha256=hashlib.sha256(saved_path.read_bytes()).hexdigest(), source_sha256=hashlib.sha256(Path(sms.__file__).read_bytes()).hexdigest(), analysis_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest()), parameters=saved['parameters'], baseline=metrics(logits,y), saved_metrics_reproduced_exactly=True, held_out_patterns=groups, train_patterns=pattern_table(train_x, train_y), empirical_pattern_majority_accuracy_ceiling=empirical_ceiling, claim_or_prize_rule_disagreements=int((equivalent_rule != (logits > 0)).sum()), activation_profiles=profiles, neuron_ablations=ablations, word_count_interventions=interventions, probes=probes, grid_0_to_3=grid_traces, validation=dict(forward_decomposition=True, ablation_delta_identity=True, path_attribution_completeness=True)) report['activation_archive'] = export_activations( model, [('train', train_rows, train_data), ('test', test_rows, test_data)], saved, args.output.with_name(args.output.stem + '-activations.json')) plot_audit(model, probes[:5], args.output.with_suffix('.png')) sweep = [] for seed in args.seeds: torch.manual_seed(seed) fresh = sms.ReadWriteClassifier() initial = {k: v.detach().tolist() for k,v in fresh.named_parameters()} loader = DataLoader(train_data, batch_size=64, shuffle=True, generator=torch.Generator().manual_seed(saved['seed'])) with contextlib.redirect_stdout(io.StringIO()): history = sms.train(fresh, loader, saved['epochs'], saved['learning_rate']) result = dict(initialization_seed=seed, minibatch_seed=saved['seed'], split_seed=saved['seed'], held_out=sms.evaluate(fresh,test_data), initial_parameters=initial, parameters={k: v.detach().tolist() for k,v in fresh.named_parameters()}, final_train_loss=history[-1]) with torch.no_grad(): fresh_h = fresh.W(x).relu() result['held_out_active_counts'] = (fresh_h > 0).sum(0).tolist() result['baseline_zero_input_activations'] = fresh.W(torch.zeros(1,3)).relu()[0].tolist() result['predictions_disagree_with_saved'] = int(((fresh(x)>0)!=(logits>0)).sum()) if seed == saved['seed']: assert result['parameters'] == saved['parameters'] assert history == saved['train_loss'] result['exact_saved_training_reproduction'] = True sweep.append(result) print(f"Initialization {seed}: accuracy={result['held_out']['accuracy']:.4f}; " f"recall={result['held_out']['spam_recall']:.4f}; " f"active counts={result['held_out_active_counts']}", flush=True) report['fixed_split_initialization_sweep'] = sweep args.output.write_text(json.dumps(report, indent=2)+'\n') print(f'Wrote {args.output} and {args.output.with_suffix(".png")}') if __name__ == '__main__': main()