Patent ID: 11907825
Assignee: GOOGLE LLC
Field: Computer technology (Electrical engineering)
Classification: CPC G | IPC G

Claim 0:
1. A method of 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 method comprises:
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.