[flyingdoubleG]

Problem

With skip_connection=True, SparseCoder.forward adds the skip term to the main TopK reconstruction:

sae_out = self.decode(top_acts, top_indices)
if self.W_skip is not None:
    sae_out += x.to(self.dtype) @ self.W_skip.mT

The Multi-TopK branch re-decodes with 4 * k latents but does not add it:

top_acts, top_indices = pre_acts.topk(4 * self.cfg.k, sorted=False)
sae_out = self.decode(top_acts, top_indices)
multi_topk_fvu = (sae_out - y).pow(2).sum() / total_variance

The trainer optimizes fvu + auxk_alpha * auxk_loss + multi_topk_fvu / 8. So for a skip-connected coder with multi_topk=True, typically a skip transcoder, the Multi-TopK term trains the 4k latents to reconstruct the part of the target that W_skip already predicts. That conflicts with the main objective, and the logged multi_topk_fvu is measured against the wrong reconstruction.

Repro on CPU: set W_skip = I and W_dec = 0, so that the skip connection alone reconstructs the input exactly:

fvu=0.0000  multi_topk_fvu=1.0370

Fix

Add the skip term to the Multi-TopK reconstruction, as in the main branch.

Test

tests/test_multi_topk.py::test_multi_topk_fvu_includes_skip_connection checks two things:

• in the exact-skip case, both FVUs are 0;
• for random weights, multi_topk_fvu equals the FVU of the 4k decode plus the skip term.

It fails on main and passes with the fix (CPU, eager decoder). tests/test_auxk_loss.py still passes.

🤖 Generated with Claude Code

[CLAassistant]

All committers have signed the CLA.