[flyingdoubleG]

Summary

With sparsity_include_decoder_norm=True (the default), the SAE, Lorsa, MoLT and Crosscoder encode methods apply the JumpReLU threshold to hidden_pre * ||W_d||, but return hidden_pre unscaled. The JumpReLU Lp (pre-act) loss in SparseDictionary.compute_loss then computes

ReLU(threshold - hidden_pre) * ||W_d||

This compares the threshold with a pre-activation on a different scale from the one it actually gates. The inconsistency shows up under the reparametrization W_E, b_E *= s, W_D /= s. That reparametrization leaves the model's function, l_rec and the tanh sparsity loss unchanged, but l_p changes. Training can therefore lower l_p simply by inflating the encoder and shrinking the decoder, which is exactly the dynamic that sparsity_include_decoder_norm is meant to suppress.

Reproduction

A JumpReLU SAE with every latent rescaled by s, run on the same batch:

s=1.0: l_rec=104.341911  l_p=37.782375
s=2.0: l_rec=104.341911  l_p=32.332272
s=5.0: l_rec=104.341911  l_p=29.326506

Change

In the sparsity_include_decoder_norm case, compare the threshold with the gated quantity:

ReLU(threshold - hidden_pre * ||W_d||)

This is the decoder-norm-scaled counterpart of ReLU(exp(t) - f_i(x)) ||W_{d,i}||. The False case is unchanged, since there the threshold already gates hidden_pre and the ||W_d|| factor gives the invariance. CLT overrides compute_loss and is not affected.

Tests

• New test_jumprelu_lp_loss_is_invariant_to_latent_rescaling[True] fails before this change and passes after. The [False] case is a control and passes both before and after.
• tests/unit/test_sae.py, test_crosscoder.py and test_evaluator.py pass.
• ruff check and ruff format --check are clean.