Training a neural network using stochastic whitening batch normalization

A method performed by one or more computers implementing a stochastic whitening batch normalization (SWBN) layer that is between a first neural network layer and a second neural network layer in a neural network system, wherein the first neural network generates first layer outputs having a plurality of components, and wherein the method comprises:

• during training of the neural network system on a plurality of batches of training data and for each of the plurality of batches:
• obtaining a respective first layer output for each of the plurality of training data in a batch;
• determining a plurality of normalization statistics for the batch from the first layer outputs;
• generating a respective normalized output for each training datum in the batch by normalizing each component of each first layer output using the determined plurality of normalization statistics;
• updating a respective whitening matrix of the SWBN layer based on a whitening loss function using a covariance matrix determined based at least in part on the normalized output;
• performing stochastic whitening on the normalized components of each first layer output to whiten data for each training datum in the batch using the updated respective whitening matrix of the SWBN layer;
• transforming the whitened data for each training datum in the batch by using a first set of task parameters for rescaling the whitened data and a second set of task parameters for shifting the whitened data;
• generating a respective SWBN layer output for each of the plurality of training data from the transformed whitened data for each training datum in the batch;
• and
• providing the SWBN layer output as an input to the second neural network layer,
• wherein the first and second sets of task parameters are decoupled from the whitening matrix such that:
• the first and second sets of task parameters are updated in a backpropagation phase of the neural network based on a task loss function, and
• the respective whitening matrix is only updated in a forward propagation phase based on the whitening loss function separate from the task loss function, and
• wherein the SWBN layer is one of a plurality of SWBN layers of the neural network, and each SWBN layer of the plurality of SWBN layers has its own respective whitening matrix for performing stochastic whitening.

A non-transitory memory storing one or more programs, which, when executed by one or more processors of a device, cause the device to be configured to perform operations to implement a neural network, wherein the neural network comprises:

• a stochastic whitening batch normalization (SWBN) layer between a first neural network layer and a second neural network layer, wherein the first neural network layer generates first layer outputs having a plurality of components, and wherein the SWBN layer is configured to:
• during training of a neural network system on a plurality of batches of training data and for each of the plurality of batches:
• obtaining a respective first layer output for each of the plurality of training data in a batch;
• determining a plurality of normalization statistics for the batch from the first layer outputs;
• generating a respective normalized output for each training datum in the batch by normalizing each component of each first layer output using the determined plurality of normalization statistics;
• updating a respective whitening matrix of the SWBN layer based on a whitening loss function using a covariance matrix determined based at least in part on the normalized output;
• performing stochastic whitening on the normalized components of each first layer output to whiten data for each training datum in the batch using the updated respective whitening matrix of the SWBN layer;
• transforming the whitened data for each training datum in the batch by using a first set of task parameters for rescaling the whitened data and a second set of task parameters for shifting the whitened data;
• generating a respective SWBN layer output for each of the plurality of training data from the transformed whitened data for each training datum in the batch;
• and
• providing the SWBN layer output as an input to the second neural network layer,
• wherein the first and second sets of task parameters are decoupled from the whitening matrix such that:
• the first and second sets of task parameters are updated in a backpropagation phase of the neural network based on a task loss function, and
• the respective whitening matrix is only updated in a forward propagation phase based on the whitening loss function separate from the task loss function, and
• wherein the SWBN layer is one of a plurality of SWBN layers of the neural network, and each SWBN layer of the plurality of SWBN layers has its own respective whitening matrix for performing stochastic whitening.