Research map / Architectures and efficiency / RNNs and normalization
PETRA: Parallel End-to-end Training with Reversible Architectures
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PETRA, a model-parallel training method for reversible networks that decouples forward and backward passes across stages without weight stashing or activation buffers, and reports accuracy close to backpropagation on CIFAR-10, ImageNet32 and ImageNet with RevNets.
Problem
Backpropagation is backward-locked and sequential, limiting model-parallel parallelization; existing delayed-gradient and pipeline methods need large parameter and activation buffers that scale with the number of stages.
Approach
Uses the invertibility of reversible residual stages to reconstruct inputs during the backward pass, so each stage keeps no input or parameter buffers; stages update with the latest parameters, giving delayed approximate gradients that run in parallel with the forward pass, with gradient accumulation to reduce staleness.
Key results
- PETRA on RevNet-18 reaches 71.0% top-1 on ImageNet vs 70.8% for backpropagation on RevNet-18 (ours), and 94.9% on CIFAR-10.
- PETRA on RevNet-50 reaches 74.8% on ImageNet vs 75.4% for backprop on RevNet-50, a drop of about 0.6 points.
- Mean iteration time speed-up over reversible backprop is 3.0x for RevNet-18 (10 GPUs) and 2.4x for RevNet-34 (18 GPUs).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use PETRA instead of backpropagation | RevNet18 · CIFAR-10 | accuracy: 94.9 vs $95.0 | No clear effect |
| use PETRA instead of backpropagation | RevNet18 · ImageNet32 | accuracy: 54.0 -> $54.6 | Helped |
| use PETRA instead of backpropagation | RevNet18 · ImageNet | accuracy: 70.8 -> $71.0 | Helped |
| use PETRA instead of backpropagation | RevNet34 · CIFAR-10 | accuracy: $95.5 -> 94.8 | Hurt |
| use PETRA instead of backpropagation | RevNet34 · ImageNet32 | accuracy: 56.5 -> 56.1 | Hurt |
| use PETRA instead of backpropagation | RevNet34 · ImageNet | accuracy: $74.0 -> 73.5 | Hurt |
| use PETRA instead of backpropagation | RevNet50 · CIFAR-10 | accuracy: $94.8 -> 94.5 | Hurt |
| use PETRA instead of backpropagation | RevNet50 · ImageNet32 | accuracy: $58.8 -> 59.6 | No clear effect |
| use PETRA instead of backpropagation | RevNet50 · ImageNet | accuracy: 75.6 -> 74.8 | Hurt |
| add delayed gradients (ablation) | RevNet18 with input and parameter buffers · CIFAR-100 | accuracy: 77.5 -> 76.5 | Hurt |
| remove parameter buffer (ablation) | delayed-gradient RevNet with input buffer · CIFAR-100 | accuracy: 76.5 -> 76.4; 76.3 -> 75.5; 77.7 -> 76.7 | Hurt |
| remove input buffer (ablation) | delayed-gradient RevNet with parameter buffer · CIFAR-100 | accuracy: 76.5 -> 75.4; 76.3 -> 76.4; 77.7 -> 78.3 | Mixed |
Surprising findings
- PETRA's gradient approximation aligns better with the end-to-end gradient than the standard delayed gradient does, attributed to using up-to-date parameters in the backward pass.
- Using gradient accumulation with larger k closes the gap to backpropagation, reaching a perfect match at k=32 on ImageNet for RevNet-18.
Limitations
- Non-reversible downsampling stages still require activation buffers and recomputation, so memory savings are smaller for standard RevNets.
- Delayed/stale gradients degrade accuracy without accumulation; RevNet-50 on ImageNet shows a small drop.
- Only evaluated on vision classification with ResNet-style models; LLM application is future work.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet32 ImageNet
Benchmarks: CIFAR-10 CIFAR-100 ImageNet32 ImageNet
Metrics: top-1 accuracy mean iteration time speed-up memory (GB) cosine similarity of gradients norm ratio
Baselines: Backpropagation (ResNet, PyTorch) Reversible backpropagation Delayed gradients PipeDream
Builds on: RevNet delayed gradient methods PipeDream activation checkpointing PyTorch autograd
Closest papers by meaning
- Accelerated training through iterative gradient propagation along the residual path (ICLR 2025 · similarity 0.87)
- Reducing Fine-Tuning Memory Overhead by Approximate and Memory-Sharing Backpropagation (ICML 2024 · similarity 0.85)
- Dr2Net: Dynamic Reversible Dual-Residual Networks for Memory-Efficient Finetuning (CVPR 2024 · similarity 0.84)
- MARS: Unleashing the Power of Variance Reduction for Training Large Models (ICML 2025 · similarity 0.84)
- GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling (NeurIPS 2025 · similarity 0.84)
- Time- Memory- and Parameter-Efficient Visual Adaptation (CVPR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.