Improvements to numerical stability and performance for ZCA whitening

[mihirparadkar]

This PR simplifies ZCA whitening by not forming the covariance matrix or the full diagonal matrix. Instead it uses the SVD of the data itself and broadcasting division, which is mathematically equivalent but faster and more numerically stable.

Tested numerically by the following

import numpy as np
import numpy.random as npr
import numpy.linalg as linalg

def princomps_old(flat_x):
    sigma = np.dot(flat_x.T, flat_x) / flat_x.shape[0]
    u, s, _ = linalg.svd(sigma)
    principal_components = np.dot(np.dot(u, np.diag(1. / np.sqrt(s + 1e-5))), u.T)
    return principal_components
 
def princomps_new(flat_x):
    n_examples = flat_x.shape[0]
    u, s, vt = linalg.svd(flat_x / np.sqrt(n_examples))
    s_expand = np.hstack((s, np.zeros(vt.shape[0] - n_examples, dtype=flat_x.dtype)))
    principal_components = (vt.T / np.sqrt(s_expand**2 + 1e-5)).dot(vt)
    return principal_components

x = npr.randn(64, 784) #A batch of MNIST data
flat_x = x - x.mean(axis=0)
print(np.allclose(princomps_old(flat_x), princomps_new(flat_x))) #True

[fchollet]

LGTM, thanks!

[fchollet]

Style: use num_examples for consistency

[fchollet]

Style: spaces around **