Enabling Model Parallelism for Neural Networks Based on Decoupled Supervised Contrastive Learning

Ming-Yao Ho, Cheng-Kai Wang, Hung-Hsuan Chen

Keywords: model parallelism, contrastive learning, supervised contrastive learning, bubbles

TL;DR: We enable simultaneous parameter gradient computation on different layers, thus achieving better model parallelism.

Abstract: End-to-end backpropagation (BP) is the current standard for training deep neural networks. However, as networks become deeper, BP becomes inefficient for various reasons. This paper introduces a new methodology that decouples BP, transforming a long gradient flow into multiple short ones. This design enables the \emph{simultaneous} computation of parameter gradients in different layers so as to realize better model parallelism. Thorough experiments are presented to demonstrate the efficiency and effectiveness of our model compared to BP, Early Exit, GPipe, and associated learning (AL), a state-of-the-art methodology for backpropagation decoupling. The experimental code is released for reproducibility at \url{https://anonymous.4open.science/r/SCPL-802C/}

Primary Area: general machine learning (i.e., none of the above)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 2161