DiNNO: Distributed Neural Network Optimization for Multi-Robot Collaborative Learning

摘要

We present DiNNO, a distributed algorithm that enables a group of robots to collaboratively optimize a deep neural network model while communicating over a mesh network. Each robot only has access to its own data and maintains its own version of the neural network, but eventually learns a model that is as good as if it had been trained on all the data centrally. No robot sends raw data over the wireless network, preserving data privacy and ensuring efficient use of wireless bandwidth. At each iteration, each robot approximately optimizes an augmented Lagrangian function, then communicates the resulting weights to its neighbors, updates dual variables, and repeats. Eventually, all robots&#x2019; local model weights reach a consensus. For convex objective functions, this consensus is a global optimum. Unlike many existing methods we test our algorithm on robotics-related, deep learning tasks with nontrivial model architectures. We compare DiNNO to two benchmark distributed deep learning algorithms in (i) an MNIST image classification task, (ii) a multi-robot implicit mapping task, and (iii) a multi-robot reinforcement learning task. In these experiments we show that DiNNO performs well when faced with nonconvex deep learning objectives, time-varying communication graphs, and streaming data. In all experiments our method outperforms baselines, and was able to achieve validation loss equivalent to centrally trained models. See <uri>msl.stanford.edu/projects/dist_nn_train</uri> for videos and code.