[flyingdoubleG]

Description

Follow-up to #754. That PR made the SAE replacement in get_recons_loss keep the original activations at excluded tokens (BOS/EOS/PAD, which run_evals excludes by default), and noted that the zero-ablation reference kept its existing behaviour. That reference still zeroes every position, including BOS, so the two runs that ce_loss_score and kl_div_score compare no longer touch the same positions:

ce_loss_score = (CE_ablation - CE_sae) / (CE_ablation - CE_clean)

CE_sae comes from a run where BOS is left alone, while CE_ablation comes from a run where BOS (usually the attention sink) is also zeroed. The ablation loss is inflated, so both scores come out higher than they should.

On GPT-2 small with gpt2-small-res-jb (4 short prompts, BOS excluded, everything else default):

hook	CE clean	CE SAE	CE ablation, all tokens (current)	CE ablation, non-excluded only	ce_loss_score current → fixed	kl_div_score current → fixed
blocks.4.hook_resid_pre	3.49	3.68	17.07	11.16	0.986 → 0.976	0.990 → 0.983
blocks.8.hook_resid_pre	3.49	3.93	11.35	9.39	0.944 → 0.926	0.972 → 0.964

At layer 4 the unexplained part of the CE loss (1 − score) is under-reported by about 40% (1.4% vs 2.4%).

The fix zeroes only the positions the SAE replaces, using the same mask, for the residual, all-head (hook_z) and single-head hooks. When nothing is excluded, the behaviour is unchanged.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

test_zero_ablation_preserves_excluded_tokens_and_other_heads (3 cases, mirroring the #754 test) checks that the ablation run changes exactly the positions the replacement run changes. All 3 fail on main and pass with the fix; tests/test_evals.py passes in full (95 tests). ruff format --check and ruff check pass on the changed files.

You have tested formatting, typing and tests

• I have run make check-ci to check format and linting. (you can run make format to format code if needed.)