From 79483389d9c13d9fedaca3fe790d98ef862702e2 Mon Sep 17 00:00:00 2001 From: Sohrab Tawana Date: Sun, 7 Jun 2026 21:14:09 +0200 Subject: [PATCH] fix: backprop combined loss (reconstruction + AuxK), not reconstruction only _calculate_loss_and_log builds loss = reconstruction_loss + lambda_aux * aux_loss but returns reconstruction_loss for backprop in both branches, so the AuxK dead-latent revival term never contributes a gradient. Return the combined loss so the auxiliary loss is actually optimized. Co-Authored-By: Claude Opus 4.8 --- crosscode/trainers/topk_crosscoder/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crosscode/trainers/topk_crosscoder/trainer.py b/crosscode/trainers/topk_crosscoder/trainer.py index b5824b7..336f66f 100644 --- a/crosscode/trainers/topk_crosscoder/trainer.py +++ b/crosscode/trainers/topk_crosscoder/trainer.py @@ -36,9 +36,9 @@ def _calculate_loss_and_log( **self._get_fvu_dict(batch_BMPD, train_res.recon_acts_BMPD), } - return reconstruction_loss, log_dict + return loss, log_dict - return reconstruction_loss, None + return loss, None def aux_loss( self, batch_BMPD: torch.Tensor, train_res: ModelHookpointAcausalCrosscoder.ForwardResult