mlsanityData to 7 Oct 2026

PETRA: Parallel End-to-end Training with Reversible Architectures

arXiv 2406.02052 ↗First posted 4 Jun 2024cs.LG stat.MLcited by 6

Not on the map. Why?

Abstract (arXiv)

Reversible architectures have been shown to be capable of performing on par with their non-reversible architectures, being applied in deep learning for memory savings and generative modeling. In this work, we show how reversible architectures can solve challenges in parallelizing deep model training. We introduce PETRA, a novel alternative to backpropagation for parallelizing gradient computations. PETRA facilitates effective model parallelism by enabling stages (i.e., a set of layers) to compute independently on different devices, while only needing to communicate activations and gradients between each other. By decoupling the forward and backward passes and keeping a single updated version of the parameters, the need for weight stashing is also removed. We develop a custom autograd-like training framework for PETRA, and we demonstrate its effectiveness on CIFAR-10, ImageNet32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models.

Key facts. Cited by 6 papers in the map (data through 7 Oct 2026). Builds on 34 of its 47 references that are in the map. First posted on arXiv 4 Jun 2024.

Builds on

34 of 47 references.

References · 34

  1. 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
  2. 2020Reformer: The Efficient Transformercited by 1,658
  3. 2017Accurate, Large Minibatch SGD: Training ImageNet in 1 Hourcited by 1,479
  4. 2019ZeRO: Memory Optimizations Toward Training Trillion Parameter Modelscited by 1,278
  5. 2014NICE: Non-linear Independent Components Estimationcited by 1,052
  6. 2016Training Deep Nets with Sublinear Memory Costcited by 825
  7. 2022Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Modelcited by 557
  8. 2017A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasetscited by 374
  9. 2021ResNet strikes back: An improved training procedure in timmcited by 365
  10. 2021ZeRO-Offload: Democratizing Billion-Scale Model Trainingcited by 339
  11. 2017The Reversible Residual Network: Backpropagation Without Storing Activationscited by 215
  12. 2018i-RevNet: Deep Invertible Networkscited by 137
22 more
  1. 2016Decoupled Neural Interfaces using Synthetic Gradientscited by 135
  2. 2019Training Neural Networks with Local Error Signalscited by 135
  3. 2020Memory-Efficient Pipeline-Parallel DNN Trainingcited by 133
  4. 2022Revisiting Neural Scaling Laws in Language and Visioncited by 129
  5. 2018PipeDream: Fast and Efficient Pipeline Parallel DNN Trainingcited by 116
  6. 2018Greedy Layerwise Learning Can Scale to ImageNetcited by 113
  7. 2020DAPPLE: A Pipelined Data Parallel Approach for Training Large Modelscited by 100
  8. 2019Decoupled Greedy Learning of CNNscited by 76
  9. 2021Revisiting Locally Supervised Learning: an Alternative to End-to-end Trainingcited by 71
  10. 2019PipeMare: Asynchronous Pipeline Parallel DNN Trainingcited by 54
  11. 2018Decoupled Parallel Backpropagation with Convergence Guaranteecited by 39
  12. 2023Reversible Vision Transformerscited by 33
  13. 2020torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Modelscited by 26
  14. 2023Can Forward Gradient Match Backpropagation?cited by 26
  15. 2018Training Neural Networks Using Features Replaycited by 23
  16. 2019On the Acceleration of Deep Learning Model Parallelism with Stalenesscited by 17
  17. 2019Fully Decoupled Neural Network Learning Using Delayed Gradientscited by 12
  18. 2020Pipelined Backpropagation at Scale: Training Large Models without Batchescited by 11
  19. 2021Decoupled Greedy Learning of CNNs for Synchronous and Asynchronous Distributed Learningcited by 6
  20. 2021Gradient Forward-Propagation for Large-Scale Temporal Video Modellingcited by 5
  21. 2023Colossal-Auto: Unified Automation of Parallelization and Activation Checkpoint for Large-scale Modelscited by 5
  22. 2020Accumulated Decoupled Learning: Mitigating Gradient Staleness in Inter-Layer Model Parallelizationcited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 3; 2026: 3.

6 papers.

  1. 2025PosterGen: Aesthetic-Aware Multi-Modal Paper-to-Poster Generation via Multi-Agent LLMscited by 8
  2. 2026Stabilizing Native Low-Rank LLM Pretrainingcited by 6
  3. 2025Model Parallelism With Subnetwork Data Parallelismcited by 3
  4. 2025Learning without Global Backpropagation via Synergistic Information Distillationcited by 1
  5. 2026Unbiased Approximate Vector-Jacobian Products for Efficient Backpropagationcited by 0
  6. 2026Learned Subspace Compression for Communication-Efficient Pipeline Parallelismcited by 0