Should topk aux loss detach e before loss calculation?

[chanind]

In dictionary_learning's topk aux loss, the SAE error e is detached before the aux loss is calculated, but in sparsify it's not detached.

Intuitively, it feels more correct to detach since this means the aux loss can only be reduced by pulling dead latents towards the SAE error, not by pulling the SAE error closer to the dead latents.

Should Sparsify also detach the error before calculating aux loss, or is there a reason why it's not ideal to do this?

[norabelrose]

I will look at this later today, I'm not sure, but FWIW I haven't used AuxK loss in a very long time so it is sort of a neglected feature. I don't think it's really ever necessary. MultiTopK would be a better way to get rid of dead latents. Also, our new Signum optimizer seems to reduce dead latents significantly over Adam.

[chanind]

Is the signum optimizer whats in sign_sgd.py? Is there a paper describing the multi topk technique? It looks like it adds a loss that's equivalent to increasing k by 4, sort of like a matryoshka SAE?

[norabelrose]

Yeah. MultiTopK was introduced in the original TopK paper from Gao et al.

[chanind]

Ah you are correct, I clearly didn't read the topk paper closely enough 😅. Thanks for the response and the great library!