Removing layer privileging in the importance loss

[danbraunai]

Migrated from: goodfire-ai/spd-gf#56
Original author: @leesharkey

---

Removing layer privileging in the importance loss using the learned linear transformation proposed by Lucius

Currently, SPD and APD both privilege layers. We think this has been okay for the small models that we've trained things on, but in general some networks might not privilege layers strongly enough for this to be a valid thing to do. We can't, for instance, rule out the possibility that we struggle to get things working on resid-mlp with more than 3 layers is due to this issue.

Lucius proposed some scheme to avoid this involving learning a linear transformation that would take as input the causal importances and transform them in some way such that the final causal importance loss does not privilege layers but where we still get to penalize causal importances of the subcomponents. I assume this transformation needs to be constrained to have certain properties (otherwise it can just learn to be a matrix of all zeros, which would minimize the causal importance loss).

[danbraunai]

This is something that we're just keeping in mind as we go along. Notably, Lucius thinks it may be possible in SPD to add back an L2 weight norm to potentially achieve this.