fix: AuxK loss carries no gradient, now backprop combined loss, not reconstruction only

[SohrabTa]

The TopK / BatchTopK / GroupMax trainer computes the combined loss loss = reconstruction_loss + lambda_aux * aux_loss, but _calculate_loss_and_log returns reconstruction_loss (not loss) in both the logging and non-logging branches. The returned tensor is the one that gets .backward() called on it (base_acausal_trainer.run_batch → base_trainer's loss.div(grad_accum).backward()), so the AuxK dead-latent revival term never contributes any gradient. The auxiliary loss is logged to W&B but is effectively inert, since it does nothing to revive dead latents.

reconstruction_loss = calculate_reconstruction_loss_summed_norm_MSEs(...)
aux_loss = self.aux_loss(batch_BMPD, train_res)
loss = reconstruction_loss + self.cfg.lambda_aux * aux_loss   # combined loss built...

if log:
    log_dict = {... "train/aux_loss": aux_loss.item() ...}
    return reconstruction_loss, log_dict   # ...but reconstruction_loss returned for backprop
return reconstruction_loss, None

Fix

Return loss (the combined reconstruction + AuxK loss) in both branches so the auxiliary term is actually optimized. Two-line change and logging is unaffected.

How we found it

We forked this repo to train Anthropic-style sparse crosscoders on the residual-stream activations of ProtT5 (a protein language model) for a research project on disentangling and interpreting its internal representations. When we inspected our dead-latent counts, they stayed stubbornly high throughout training even though train/aux_loss was being logged, but it kept increasing. Tracing the loss that reaches .backward() showed the AuxK term was logged but never backpropped and the returned tensor was reconstruction_loss, so lambda_aux * aux_loss was dropped from the gradient. The symptom (AuxK appears active in logs but has zero effect on dead latents) makes this easy to miss without following the exact tensor that gets differentiated.

Testing

Verified the combined loss now flows to backprop and that train/loss, train/reconstruction_loss, and train/aux_loss logging is unchanged. We are re-running our full crosscoder training pipeline with this fix.