Layer‐parallel training of residual networks with auxiliary variable networks

Abstract

AbstractGradient‐based methods for training residual networks (ResNets) typically require a forward pass of input data, followed by back‐propagating the error gradient to update model parameters, which becomes time‐consuming as the network structure goes deeper. To break the algorithmic locking and exploit synchronous module parallelism in both forward and backward modes, auxiliary‐variable methods have emerged but suffer from communication overhead and a lack of data augmentation. By trading off the recomputation and storage of auxiliary variables, a joint learning framework is proposed in this work for training realistic ResNets across multiple compute devices. Specifically, the input data of each processor is generated from its low‐capacity auxiliary network (AuxNet), which permits the use of data augmentation and realizes forward unlocking. Backward passes are then executed in parallel, each with a local loss function derived from the penalty or augmented Lagrangian (AL) method. Finally, the AuxNet is adjusted to reproduce updated auxiliary variables through an end‐to‐end training process. We demonstrate the effectiveness of our method on ResNets and WideResNets across CIFAR‐10, CIFAR‐100, and ImageNet datasets, achieving speedup over the traditional layer‐serial training approach while maintaining comparable testing accuracy.