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

arXiv: https://arxiv.org/abs/2512.08160 · HTML: https://mlsanity.com/paper/2512.08160/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 30 references that are in the map.
- First posted on arXiv 9 Dec 2025.

## Builds on (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [Energy and Policy Considerations for Deep Learning in NLP](https://mlsanity.com/paper/1906.02243/) (2019)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Deep Learning for Medical Image Analysis](https://mlsanity.com/paper/1708.08987/) (2017)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Beyond Data and Model Parallelism for Deep Neural Networks](https://mlsanity.com/paper/1807.05358/) (2018)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Exponential Moving Average of Weights in Deep Learning: Dynamics and Benefits](https://mlsanity.com/paper/2411.18704/) (2024)
- [On the Acceleration of Deep Learning Model Parallelism with Staleness](https://mlsanity.com/paper/1909.02625/) (2019)
- [A Gradient-Interleaved Scheduler for Energy-Efficient Backpropagation for Training Neural Networks](https://mlsanity.com/paper/2002.05529/) (2020)
- [LayerPipe: Accelerating Deep Neural Network Training by Intra-Layer and Inter-Layer Gradient Pipelining and Multiprocessor Scheduling](https://mlsanity.com/paper/2108.06629/) (2021)
- [SCV-GNN: Sparse Compressed Vector-based Graph Neural Network Aggregation](https://mlsanity.com/paper/2304.13532/) (2023)
- [SpikePipe: Accelerated Training of Spiking Neural Networks via Inter-Layer Pipelining and Multiprocessor Scheduling](https://mlsanity.com/paper/2406.06879/) (2024)

## Abstract (arXiv)

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.