LayerPipe2: Multistage Pipelining and Weight Recompute via Improved Exponential Moving Average for Training Neural Networks

TL;DR

LayerPipe2 provides a principled, retiming-based theory for pipelined backpropagation, deriving explicit delay assignments ($Delay(l)=2S(l)$) and extending to arbitrary multistage partitions. It introduces a storage-efficient weight reconstruction via a pipeline-aware EMA, reducing memory from a full history to a compact moving-average while preserving convergence. Experimental results on CIFAR-100 with ResNet-18 demonstrate that the proposed EMA recovers performance close to exact weight stashing, validating the approach and its scalability. The work connects delayed-gradient theory with practical pipeline design, enabling deeper, more flexible training pipelines with controlled compute-communication tradeoffs.

Abstract

In our prior work, LayerPipe, we had introduced an approach to accelerate training of convolutional, fully connected, and spiking neural networks by overlapping forward and backward computation. However, despite empirical success, a principled understanding of how much gradient delay needs to be introduced at each layer to achieve desired level of pipelining was not addressed. This paper, LayerPipe2, fills that gap by formally deriving LayerPipe using variable delayed gradient adaptation and retiming. We identify where delays may be legally inserted and show that the required amount of delay follows directly from the network structure where inner layers require fewer delays and outer layers require longer delays. When pipelining is applied at every layer, the amount of delay depends only on the number of remaining downstream stages. When layers are pipelined in groups, all layers in the group share the same assignment of delays. These insights not only explain previously observed scheduling patterns but also expose an often overlooked challenge that pipelining implicitly requires storage of historical weights. We overcome this storage bottleneck by developing a pipeline--aware moving average that reconstructs the required past states rather than storing them explicitly. This reduces memory cost without sacrificing the accuracy guarantees that makes pipelined learning viable. The result is a principled framework that illustrates how to construct LayerPipe architectures, predicts their delay requirements, and mitigates their storage burden, thereby enabling scalable pipelined training with controlled communication computation tradeoffs.

Figures (5)

• Figure 1: Dataflow graph of a modern pipeline parallel systems with stage-based pipelines and stashing for weights and activations.
• Figure 2: Conceptual delayed least–mean–square (DLMS) adaptation, in which an $M$-sample delay is introduced in the coefficient update path. This abstraction models the effect of stale gradients within iterative learning systems and provides a historical analogue to delayed gradient application in pipelined training.
• Figure 3: Retiming-based derivation of LayerPipe pipelining. Delays are first inserted at feedforward cutsets and weight-update edges, then progressively retimed across cutsets so that delay elements accumulate at layer boundaries. The resulting temporal separation forms explicit pipeline stages and reveals activation and weight stashing as structural consequences of delay motion. The delay placements match the analytical rule $\mathrm{Delay}(l)=2S(l)$ derived in the text.
• Figure 4: Illustration of retiming applied over a grouped two-layer stage. Delays are inserted at feedforward cutsets and gradient feedback edges, then redistributed using backward and forward retiming cutsets. One delay element is left at each downstream stage boundary, yielding a structure in which both layers share identical delay requirements. This demonstrates that delay placement is determined by the number of pipeline stages that follow the grouped region, not by the number of layers it contains.
• Figure 5: Test accuracy over fifty epochs on CIFAR--100 for four weight–handling strategies under pipelined training. Explicit weight stashing (baseline) produces stable convergence. Using the latest weights degrades performance, and a fixed EMA only partially recovers accuracy. The proposed pipeline-aware EMA reconstructs historical weights sufficiently accurately to match the baseline, but without storing multiple weight versions.