Memory efficient scalable deep learning with model parallelization
Methods and systems for training a neural network include sampling multiple local sub-networks from a global neural network. The local sub-networks include a subset of neurons from each layer of the global neural network. The plurality of local sub-networks are trained at respective local processing devices to produce trained local parameters. The trained local parameters from each local sub-network are averaged to produce trained global parameters.
1. A method for training a neural network, comprising:
sampling a plurality of local sub-networks from a global neural network, wherein the local sub-networks comprise a subset of neurons from each layer of the global neural network;
training the plurality of local sub-networks at respective local processing devices to produce trained local parameters; and
averaging the trained local parameters from each local sub-network to produce trained global parameters.
2. The method of claim 1 , wherein sampling the plurality of local sub-networks comprises fixed sampling.
3. The method of claim 1 , wherein sampling the plurality of local sub-networks comprises dynamic sampling.
4. The method of claim 1 , wherein training the plurality of local sub-networks comprises sharing local parameters between local processing devices when parameters overlap in respective sub-networks.
5. The method of claim 1 , wherein the global parameters are stored at a single global parameter server.
6. The method of claim 1 , further comprising performing a classification task using the trained global parameters.
7. The method of claim 6 , wherein performing the classification task comprises performing distributed matrix multiplication.
8. A system for training a neural network, comprising:
a plurality of local sub-network processing devices, each comprising:
a neural network module comprising neurons that represent a subset of neurons from each layer of a global neural network, the neural network module being configured to train a local sub-network to produce trained local parameters; and
a global parameter server configured to average the trained local parameters from each local sub-network to produce trained global parameters.
9. The system of claim 8 , wherein the neurons of each neural network module represent a local sub-network that is sampled from the global neural network based on fixed sampling.
10. The system of claim 8 , wherein the neurons of each neural network module represent a local sub-network that is sampled from the global neural network based on dynamic sampling.
11. The system of claim 8 , wherein the neural network module of each local sub-network processing device is further configured to share local parameters between local processing devices when parameters overlap in respective sub-networks.
12. The system of claim 8 , wherein the plurality of local sub-network processing devices are further configured to perform a classification task using the trained global parameters.
13. The system of claim 12 , wherein the plurality of local sub-network processing devices are further configured to perform distributed matrix multiplication.