Whitening Update

[HarrisonSantiago]

• Whiten.py works with data on [N, D] level
• Images.py is a high level wrapper that contains additional support for whitening images
• Patches.py -> Images.py
• Updated paths in misc example notebooks
• Created new example notebook showing how to use whitening

[belsten]

I've noticed an issue with how the ZCA/PCA whitening is done.

As of right now, compute_whitening_stats masks out the eigenvalues/vectors based on n_components. This masking should occur after the inversion of the eigenvalues (i.e. in whiten) and that mask only needs to be applied to the eigenvalues. Right now, zeros are being inverted in whiten resulting in nans.

One way to solve this is pass n_components to whiten, which would be better because then it would be with epsilon, a parameter that has a similar function. Then whiten would perform the logic of keeping however many PCs via masking.

Here is a minimum working example that demonstrates the issue:

A = torch.from_numpy(np.asarray([
    [1,0.2,0.4],
    [0.2,1,0.2],
    [0.4,0.2,1],
])).float()
data = torch.randn(10000,3)@A.T

# set n_components less than 3 and epsilon=0
whitened_stats = compute_whitening_stats(data, n_components=2)
whiten(data,algorithm="zca",stats=whitened_stats)

Output:

tensor([[nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan],
        ...,
        [nan, nan, nan],
        [nan, nan, nan],
        [nan, nan, nan]])

[belsten]

See latest comment regarding PCA/ZCA whitening.

[HarrisonSantiago]

Thanks for catching that! Moved it into whiten and now it only masks eigenvalues. Made some changes in the accompanying notebook and images.py to reflect

[9q9q]

I think this could be unstable if the covariance matrix is ill-conditioned? Since L is triangular, could you use https://pytorch.org/docs/stable/generated/torch.linalg.solve_triangular.html?

[HarrisonSantiago]

Is switching to cholesky-stable what you were thinking? https://colab.research.google.com/drive/1W8FdF-K6hYE3W0wOmRSGVQv9kPLaRH18?usp=sharing

[9q9q]

Thanks for making the notebook, I should've been a good reviewer like Alex and given you an example to start with 😅.

What I had in mind initially was
X_whitened = torch.linalg.solve_triangular(L, X.T, upper=False).T
to solve for the whitened data directly, but I think what you have works as well, plus it's consistent with the other implementations. And maybe it's nice to compute W explicitly if we want to return it in the future (e.g. to whiten other similar data).

Looks like the errors are actually not too far off the for the two methods, but using solve_triangular is ~2x faster (added some timing code): https://colab.research.google.com/drive/1CIfNaWEbOfU9R8WeZYEQ4XepJU6ZWIRy?usp=sharing

And here's a nice writeup I found about matrix inversions in practice: https://gregorygundersen.com/blog/2020/12/09/matrix-inversion/

[belsten]

I also agree that an option to return the whitening matrix would be useful. It's helpful if you're converting between whitened/unwhitened data or want to visualize the whitening transform.

[HarrisonSantiago]

Thanks for sharing the write up! I went with using solve_triangular but not solving for X_whitened directly, I like having the functionality of returning the whitening matrix for all of the methods

[9q9q]

nit

Suggested change

	n_components=None,
	n_components = None,

[HarrisonSantiago]

Done!

[9q9q]

I think it would be helpful for documentation to

• add a few words (super short) for why someone would want to choose PCA or Cholesky over the the default ZCA, and then direct them to the paper or stackexchange for more details.
• add short example usage in a code block that calls compute_whitening_stats and then whiten (not sure how to format this for Sphinx @belsten). Maybe this is excessive, but I think we should be clear because these are separate functions. Or could mention first calling compute_whitening_stats in the stats Parameter field.

What do you think?

[HarrisonSantiago]

I cover all of that in the whitening notebook that'll be added to the documentation. What about adding that notebook to the references to encourage people to check it out?

[9q9q]

A pointer to the notebook would work instead IMO! And it avoids bloat in the docstring.

[HarrisonSantiago]

Done!

[9q9q]

Please also add tests for compute_whitening_stats and whitening using PCA and Cholesky.

[HarrisonSantiago]

Done!