[Proposal] Support data whitening in SAE training

[chanind]

Data Whitening Improves Sparse Autoencoder Learning shows that whitening the input data during training results in better SAEs. This seems like a pretty intuitive and sort of obvious thing to do, it's surprising nobody has done this so far. We already support normalizing activations using normalize_activations="expected_average_only_in", we could add an additional option "covariance_whitening" implementing the technique from this paper.

[xXCoolinXx]

I'm in favor of testing this approach, but their result needs to be replicated before it should be added to the repo.

I looked more deeply into the paper, as I was thinking about referencing it as future work for my paper, and the graphs and results they report are very strange. On the common evaluation metrics (CE, FVE, k-sparse, etc.) they don't report results for the individual models, but instead average metrics between Pythia 160M and Gemma 2 2B. I feel like that is a completely inappropriate way to report their results, as I'm almost certain that the performance difference for Pythia 160M and Gemma 2 2B with whitening is quite large. They also don't test with modern SAE architectures. Further, they only test interpretability with k-sparse probing (a deeply flawed metric!) and provide neither qualitative nor autointerp evidence that their latents are actually more interpretable.

I'm not familiar enough with the theory to know if data whitening should or should not work in this setting - I would guess it has significantly different characteristics since there isn't a single input manifold due to superposition - but I don't think one could say that the above paper provides sufficient evidence that it does work.