Training neural networks using distributed batch normalization

A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform operations for distributing the training of a neural network across a plurality of devices, wherein the neural network has a plurality of network parameters, wherein the neural network has one or more batch normalization layers, wherein each batch normalization layer applies batch normalization to outputs of a corresponding other layer in the neural network, wherein each device maintains a respective replica of the neural network, and wherein the operations comprise:

• maintaining data that partitions the plurality of devices into a plurality of sub-groups that each comprise two or more of the plurality of the devices, wherein each of the plurality of devices comprises one or more hardware accelerators for training neural networks;
• receiving, at each of the plurality of devices, a respective batch of training examples;
• training the neural network on the respective batches of training examples received by the devices, the training comprising:
• performing, by each device, a forward pass through the replica maintained by the device to determine a respective network output for each training example in the batch received by the device, comprising, for each batch normalization layer in the neural network:
• generating, by each of the plurality of devices and using the replica maintained by the device, a respective output of the corresponding other layer for each training example in the batch received by the device,
• determining, by each of the plurality of devices, a per-replica mean and a per-replica variance of the respective outputs of the corresponding other layer for the training examples in the batch received by the device;
• determining, for each sub-group, a distributed mean and a distributed variance from the per-replica means and the per-replica variances for the two or more devices in the sub-group and determined from the respective batches received by the two or more devices in the sub-group by performing, over a network that connects the two or more devices in the group using one or more communication links, a distributed reduction operation on the per-replica means and the per-replica variances for the two or more devices in the sub-group and determined from the respective batches received by the two or more devices in the sub-group;
• and
• applying, by each device, batch normalization to the respective outputs of the corresponding other layer generated by the device using the distributed mean and the distributed variance for the sub-group to which the device belongs;
• and
• performing, by each device, a backward pass through the replica maintained by the device to determine a respective update to the parameters of the neural network for each training example in the batch received by the device.

One or more non-transitory computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations for distributing the training of a neural network across a plurality of devices, wherein the neural network has a plurality of network parameters, wherein the neural network has one or more batch normalization layers, wherein each batch normalization layer applies batch normalization to outputs of a corresponding other layer in the neural network, wherein each device maintains a respective replica of the neural network, and wherein the operations comprise:

• maintaining data that partitions the plurality of devices into a plurality of sub-groups that each comprise two or more of the plurality of the devices, wherein each of the plurality of devices comprises one or more hardware accelerators for training neural networks;
• receiving, at each of the plurality of devices, a respective batch of training examples;
• training the neural network on the respective batches of training examples received by the devices, the training comprising:
• performing, by each device, a forward pass through the replica maintained by the device to determine a respective network output for each training example in the batch received by the device, comprising, for each batch normalization layer in the neural network:
• generating, by each of the plurality of devices and using the replica maintained by the device, a respective output of the corresponding other layer for each training example in the batch received by the device,
• determining, by each of the plurality of devices, a per-replica mean and a per-replica variance of the respective outputs of the corresponding other layer for the training examples in the batch received by the device;
• determining, for each sub-group, a distributed mean and a distributed variance from the per-replica means and the per-replica variances for the two or more devices in the sub-group and determined from the respective batches received by the two or more devices in the sub-group by performing, over a network that connects the two or more devices in the group using one or more communication links, a distributed reduction operation on the per-replica means and the per-replica variances for the two or more devices in the sub-group and determined from the respective batches received by the two or more devices in the sub-group;
• and
• applying, by each device, batch normalization to the respective outputs of the corresponding other layer generated by the device using the distributed mean and the distributed variance for the sub-group to which the device belongs;
• and
• performing, by each device, a backward pass through the replica maintained by the device to determine a respective update to the parameters of the neural network for each training example in the batch received by the device.