On the Convergence Theory of Pipeline Gradient-based Analog In-memory Training

TL;DR

This paper analyzes the convergence of asynchronous pipeline gradient-based training on analog in-memory computing hardware, demonstrating that it converges efficiently despite hardware imperfections and stale weights.

Contribution

It provides the first theoretical convergence analysis of asynchronous pipeline SGD on AIMC hardware, showing near-equivalent efficiency to digital SGD.

Findings

Analog-SGD-AP converges with complexity $O(rac{1}{ ext{ε}^2} + rac{1}{ ext{ε}})$

Asynchronous pipelining overlaps computation, nearly matching synchronous pipeline efficiency

Hardware imperfections and stale weights do not significantly hinder convergence

Abstract

Aiming to accelerate the training of large deep neural networks (DNN) in an energy-efficient way, analog in-memory computing (AIMC) emerges as a solution with immense potential. AIMC accelerator keeps model weights in memory without moving them from memory to processors during training, reducing overhead dramatically. Despite its efficiency, scaling up AIMC systems presents significant challenges. Since weight copying is expensive and inaccurate, data parallelism is less efficient on AIMC accelerators. It necessitates the exploration of pipeline parallelism, particularly asynchronous pipeline parallelism, which utilizes all available accelerators during the training process. This paper examines the convergence theory of stochastic gradient descent on AIMC hardware with an asynchronous pipeline (Analog-SGD-AP). Although there is empirical exploration of AIMC accelerators, the theoretical understanding of how analog hardware imperfections in weight updates affect the training of multi-layer DNN models remains underexplored. Furthermore, the asynchronous pipeline parallelism results in stale weights issues, which render the update signals no longer valid gradients. To close the gap, this paper investigates the convergence properties of Analog-SGD-AP on multi-layer DNN training. We show that the Analog-SGD-AP converges with iteration complexity $O(\varepsilon^{-2}+\varepsilon^{-1})$ despite the aforementioned issues, which matches the complexities of digital SGD and Analog SGD with synchronous pipeline, except the non-dominant term $O(\varepsilon^{-1})$. It implies that AIMC training benefits from asynchronous pipelining almost for free compared with the synchronous pipeline by overlapping computation.