almogtavor/amortized-neuron-pruning-effects

Amortized Neuron-Ablation Effects (for pruning)

Exact mean-ablation effects of individual MLP neurons in transformer LMs, paired with cheap forward/backward signals, for the task of amortized causal intervention-effect prediction at neuron granularity (pruning).

The intervention unit is a single MLP neuron (blocks.{l}.mlp.hook_post channel). For every (prompt, neuron) we record the exact effect of replacing that neuron's last-token activation with its dataset mean, m(x, neuron<-mean) - m(x) (IOI logit-diff metric), together with the first-order attribution g*(mean-a), activation/gradient statistics, and static weight norms.

A neuron is prunable iff this effect is small. The dataset lets you study whether the exact ablation effect (hence safe prunability) is predictable from cheap signals, and how learned prediction compares to weight-magnitude / activation / attribution pruning baselines.

This is a sibling of almogtavor/amortized-causal-intervention-effects (heads + path edges); here the unit is the neuron.

Columns

• —attribution = g*(mean-a): first-order estimate of the mean-ablation effect
• —act, grad, act_grad: last-token activation, d(metric)/d(act), and their product
• —w_in_norm, w_out_norm: static neuron weight norms
• —mean_act, delta: the ablation target and applied change
• —layer_frac, clean_metric: structural / context
• —exact_effect: label -- exact mean-ablation effect (IOI logit diff)
• —layer, neuron, prompt_idx: identity

Configs

One per (model, task). Generated with GPT-2 / Qwen3 on the IOI task.