PETRA: Parallel End-to-end Training with Reversible Architectures
Authors: Stephane Rivaud
Louis Fournier
Thomas Pumir
Eugene Belilovsky
Michael Eickenberg
Edouard Oyallon
Venue: ICLR 2025 Spotlight (2025)
Keywords: Model parallelism
Delayed gradient
Reversible architectures
Primary area: optimization
TL;DR: We show how combining reversible architectures with delayed gradient approaches for model parallelism allows to achieve computational speedups with drastic memory reduction.

Abstract: 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 standard computer vision benchmarks, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models.

Published as a conference paper at ICLR 2025
PETRA: PARALLEL END-TO-END TRAINING OF RE- VERSIBLE ARCHITECTURES
Stéphane Rivaud 1,5 Louis Fournier' Thomas Pumir
Eugene Belilovsky*, Mickael Eickenberg" Edouard Oyallon
ISIR - -Sorbonne Université, Paris
FFlatiron Institute, New York
Milaa, Montréal *Concordia Université, Montréal
LLSN Université Paris-Sackap CNRS, Inria, Orsay Helmmai,, San Francisco
stephane.c a, rivaud@inria.
ABSTRACT
Reversible architectures have been shown to be capable of performing on par with their non-reversible architectures, being applied in deep learning for memory sav- ings and generative modeling. In this work, show how reversible architectures can solve challenges in parallelizing deep model training. We introduce PETRA, a novel alternative backpropagation for parallelizing gradient computations. PE- TRA facilitates effective model parallelism by enabling stages (i.e., a set of layers) compute independently on different devices, while only needing communi- cate activations and gradients between each other. By decoupling the forward and backward passes and keeping single updated version of the parameters, the need for weight stashing is also removed. We develop custom autograd-liks training framework for PETRA, and we demonstrate its effectiveness on CIFAR- 10, ImageNer32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models.
INTRODUCTION
First-order methods using stochastic gradients computed via backpropagation mini-batches de-facto standard for computing parameter updates in Deep Neural Networks (LeCun et al, 2013). As datasets and models continue to grow (Alabdulmohsin et al. 2022) there is urgent need for memory-efficient and scalable parallelization
deep learning tramming across multiple workers. Data parallelism via mini-batches (LeCun et al.) 2015) has been widely adopted in deep learning frameworks (Li et al., (020). This approach computes gradients across model replicas distributed among workers, yet requires frequent synchronization aggregate gradients, leading high com- munication costs, as well as substantial memory redundancy. Furthermore, with the increasing size and scale of models exceeding that of the growth of on-device memory. the forward and backward passes now often exceed a single device's memory capacity (Ren et al., 202T). To further address these issues, methods have attempted mitigate this memory overhead and to parallelize the sequen- tial backpropagation steps themselves across devices, while computing exact gradients. Techniques like optimizer sharding (Rajbhandari et al., 2020), tensor parallelism (Shoeybi et al. (019), activa- tion checkpointing (Chen et al., 2016), pipelining (Huang et al.) 2019), have been deployed indi- vidually or combined, leading for instance to the development 3D parallelism (Smith et al. 2022), popular methodology which improves the efficiency of the backpropagation implementation. On the other hand, the fundamental inefficiency underlying the parallelization backpropagation has not been addressed by these methods.
However, the use of exact gradient restricts algorithmic choices and parallel implementations, as highlighted by Jaderberg et al. (2017). For instance, backpropagation is backward locked: the in- puts of each layer must be propagated through the network and preserved until an error signal is
During one year leave now back CNRS, Sorbonne University.

---

Published as a conference paper at ICLR 2025
a) Backpropagation
Loss
Stage
Layer
Stage
Stage
Propagation batch
b) PETRA
Stage 2
Stage
Stage
Time
Figure 1: Comparison of PETRA with standard backpropagation. This approach splits the stages model and decouples their forward and backward passes, resulting a sixfold increase in parallelization speed in this example.
retropropagated to the layer of origin. This requirement enforces a synchronous dependency among subsequent layers and requires them to systematically store intermediary activations, potentially im- peding overall resource efficiency workers must wait for each other to continue their computations and release memory used for activations. To unlock the potential of backpropagation, inexact back- propagation procedures have been proposed. These procedures are generally conceptualized within the context of model parallelism, where a neural network is split into stages that can process their activations in parallel, potentially on multiple devices. For example, some methods use outdated pa- rameters or activations, such as double-buffered pipelining (Harlap et al.) 2018) or delayed gradient approaches (Zhuang et al. 2021b). However, these methods introduce significant memory overhead due to the use of ad hoc buffers for activations, parameters, both. Following opposite direction, local learning methods (Nøkland & Eidnes 2019, Belilovsky et al., 2020), which estimate inexact gradients via a local auxiliary neural network, pave the way parallel gradient computations but often lead to unrecoverable performance drops (Fournier et al. (002)). This underscores the need for robust alternative to backpropagation, with limited memory overhead.
In this work, we introduce PETRA (Parallel End-to-End Training with Reversible Architectures), novel method designed parallelize gradient computations within reversible architectures with minimal computational overhead. Reversible architectures are an ideal candidate for this task, as they can significantly reduce memory overhead during standard backpropagation with limited com- munication costs. Furthermore, reversibility is minor requirement,
many studies have demon- strated that standard architectures can be adapted into reversible ones without any performance drops (ooee et al. 2017 Jacobsen et al. 2018b Mangalam et al. 2022; Kitaev et al. 2020). By allow- ing parameters to evolve parallel and by computing an approximate inversion during backward, propose an effective alternative to backpropagation which allows high model parallelism with a constant communication overhead and no additional parameter or activation buffers. fact, for constant increase in communication overhead, PETRA achieves linear speedup compared to standard backpropagation with respect to the number of stages the network is split into. We illus- trate our approach in Fig. [ by contrasting the evolution of PETRA with standard backpropagation pass.
Contributions. Our contributions are follows: (1) We introduce PETRA, a streamlined ap- proach for parallelizing the training of reversible architectures. This method leverages a delayed, approximate inversion of activations during the backward pass, allowing for enhanced computa- tional efficiency. (2) Our technique significantly reduces memory overhead by minimizing the ne- cessity store extensive computational graphs. (3) It enables the parallelization of forward and backward pass computations across multiple devices, effectively distributing the workload and re- ducing training time. (4) We validate the efficacy of PETRA through rigorous testing bench- mark datasets such as CIFAR-10. ImageNet-32, and ImageNet, where demonstrates robust per- formance with minimal impact accuracy. (5) We observe significant empirical throughput
2

---

Published as a conference paper at ICLR 2025
increase when using PETRA. (6) Additionally, we provide flexible reimplementation of the au- tograd system in PyTorch, specifically tailored for our experimental setup, which is available at https://gi http://gttuub... wwtp:a/tpoeeeeppeeeeeeeeeeee
2 RELATED WORK
Reversible architectures. Reversible DNNs are composed of layers that invertible, meaning that the input of layer can be computed from its output. This approach allows to avoid the need to store intermediary activations during the forward pass by reconstructing them progressively during the backward pass (Gomez et al. 2017). at the cost of an extra computation per layer. Invertible net- works further improve this method by removing dimensionality reduction steps such as downsam- plings, making the networks fully invertible (Jacobsen et al., 2018a). Reversibility is not restricted to a type of architecture or tasks and has been extensively used for generative models (Dinh et al.) 2014] for ResNets (Gomez et al. (017), and Transformers (Mangalam et al. 2022). However, as far as we know, reversible architectures have never been used to enhance varallelizate capabilities.
Alternatives to backpropagation. Multiple alternatives to backpropagation have been proposed previously to improve over its computational efficiency. For instance, DNI (Jaderberg et al.] (017) is the first to mention the backpropagation inefficiency and its inherent synchronization locks. How-
ever, they address those locks with a method non-competitive with simple baselines. Local (or
greedy) learning (Nøkland & Eidnes 2019 Belilovsky et al. 2019) propose
layerwise losses
to decouple the training of layers. allowing them to train in parallel (Belilovsky et al.] (202)) Local learning in videos (Malinowski et al.) 2021) notably uses the similarity between successive tempo- ral features to remove buffer memory. However, the difference in training dynamics between local
training and backpropagation still limits such approaches (Fournier et al. 2023
Wang et al. 2021
Pipeline parallclism. Pipelining encompasses a range of model parallel techniques that divide the components of a network into stages that compute in parallel, while avoiding idle workers. Initially popularized by Huang et (2019). a batch of data is divided into micro-batches that are processed independently each stage. Although more efficient pipelining schedules have been proposed (Fan) et al. 2021), notably to mitigate the peak memory overhead, keeping an exact batch gradient compu- tation requires leaving a bubble of idle workers. By alternating one forward and one backward pass for each worker, PipeDrremm (Narayanan et al.) (019) can allow to get rid of idleness bubbles, but at the expense of introducing staleness in the gradients used. Narayanan et al. (202I) mitigates this staleness to only one optimization step by accumulating gradients, thus also reducing the parameter memory overhead to only two versions of the parameters. Nevertheless, these approaches still suffer from quadratic activation memory overhead with regard to the number of stages, micro-batcl activations pile up in buffers, especially for early layers. Some implementations propose to limit this overhead by combining activation checkpointing (Chen et al., (0166) with pipelining (Kim al., 2020
Liu et al.; 2023), although the memory overhead still scales with the number of stages.
Delayed gradient. By allowing stale gradients in the update process, these previous methods provide the context for our approach. Delayed gradient optimization methods are model parallel techniques that aim to decouple and process layers in parallel during backpropagation. In these ap- proaches, delays occur stage-wise: the backward pass may be computed with outdated parameters or activations compared to the forward pass. For instance, [Huo al. (2018a) proposes a feature replay approach, where a forward pass first stores intermediary activations, which are then "replayed" to compute the backward pass in parallel. This method still requires heavy synchronization between layers, yielding lock on computations. In Zhuang et al. (2020) and Zhuang et al.] (2021a), stale gradients are computed from older parameter versions differing from the parameters used during the update. This staleness can be mitigated: [Zhuang et al., (2021a) 'shrinks' the gradient by the delay value, but more advanced techniques also exist (Yang et al., 2021 Kosson et al., 202I). I). Still, these methods are limited like previous pipelining methods by their memory overhead as the computa- tional graph is fully stored. A first step to reduce this, as proposed in Diversely Stale Parameters (DSP) (Xu et al. 2019), PipeMar Yanng et al. (02I) and (Kosson et al. 2021). 1) is to keep single set of parameters and approximate the gradients computed during the backward pass with the updated parameters, which differ from the ones used in the forward pass. This requires, like in activation
3

---

Published as a conference paper at ICLR 2025
F
F
2
(a) ResNet
(b) RevNet Forward
(c) RevNet Reverse
Figure 2: Differences between the residual block of a ResNet and its reversible counterpart. (a) Forward of residual block. (b) Forward and (c) Reverse forward reversible residual block. For reversible blocks, in Gomez et al. (2017), the input æi is doubled in size and split equally into (x].. along its channels. The function Fj includes skip-connection while F, does not.
checkpointing, an additional reconstruction of the computational graph. Furthermore, the quadratic activation memory overhead still limits the scalability of these methods for a large number of stages.
3 METHOD
3.1 STANDARD BACKPROPAGATION
We consider DNN composed of J stages (e.g., a layer or a set of layers) An input zo is propagated through the network, recursively defined by
I. (o
where Fj is the j-th stage parameterized by fj. The backpropagation algorithm is the ubiquitous algorithm to compute parameter gradients. First, an input is propagated through the network with forward pass, while storing its intermediate activations. A scalar loss L is then deduced from the corresponding output z.. Parameter gradients are then computed during the backward pass by taking advantage of the chain rule: starting from the last stage with 8, VayLe C, the gradients with regard to the activations are given by
69 A V
8,+1.
(2)
and the gradients with regard to the parameters are defined as
(3)
Note that these computations follow a synchronous and sequential order. The parameters 0, can then be updated given their gradient estimate Aj. using any optimizer.
3.2 REVERSIBLE ARCHITECTURES
We focus on the reversible neural networks presented in Gomez et al. (0117), although our method dependent on this architecture. Note that this is a weak restriction many architectures adaptable to reversible ones Mangalam et al. 20022) In practice, only few stages which do not preserve feature dimensionality are not reversible and correspond to the downsampling blocks in the ResNet. Fig. highlighss how reversible residual blocks F; differ from their standard counter- part. The input is split into two equal-size inputs, along the channel dimension, that are propagated forward according to Fig. 2b using an ad-hoe operator F,- It can be reconstructed by reverse prop- agating the output according to Fig. 2c by subtracting the output of F, rather than adding like in the previous forward.
4

---

Published as a conference paper at ICLR 2025
Table 1: Comparisons with other methods in ideal setting for one stage. We compare several methods to compute gradient estimate in a model parallel setting: classical backpropagation, its reversible counterpart (Gomez a 2017). the Delayed gradients approach offZuuang et al. (2020) and its improvements using checkpointing by Xu al. (019), and our proposed approach. Here, J the total number of stages while j is the stage index. For the sake of simplicity, we assume that backward pass requires approximately 2 times more FLOPSs than forward pass. Full Gruph indi- cates that it is required to store the full computational graph of a local forward pass. With a limited increase in communication volume and FLOPs, PETRA requires the least storage of all methods while being linearly faster than backpropagation. We assume that the forward and backward passes be executed in parallel for PETRA delayed gradients, making the backward pass responsible for most of the computation time in parallelizable approaches.
Storage
Comm. Volume
FLOPs
Mean time per batch
Methods
Activations
Params.
Backpropagation
Full Graph (FG)
1
1
3J
3/
Reversible backprop.
0
1
4
4.
4J
Delayed gradients
2(J - j)x FG
2(J-))
1
3.
2
+ Checkpointing
2(J -j)
1
1
4./
3
PETRA (ours)
0
1
4
4.J
3
Reversible stages. In order to compute the exact gradients during the backpropagation phase, each reversible stage needs retrieve its output from the stage above. We note F; the reverse stage function, which reconstructs the input from the output. We recursively apply the reconstruction to the final activation zj, such that
(1 I (FF,
8z-1 6j
Note that reconstructing the input in our procedure is computationally equivalent to recomputing the activations in activation checkpointing, meaning is equivalent single forward pass. Thus, this augmented backward procedure is equivalent to one regular forward call and backward call. However, one should observe that since the input I--1 must be sent to the reversible stages, this doubles the cost of backward communications.
Nee-revee stages. In practice, a reversible architecture includes layers that reduce dimen- sionality for computational efficiency, which thus correspond to non-invertible functions. For those very few stages, we employ buffer mechanism store activations and, like activation checkpoint- ing, recompute the computational graph with forward pass during the backward pass. Note that this would not the case when using invertible (i.e., bijective) architectures (Jacobsen et al., 2018a) 8a), which use invertible downsampling.
3.3 A PARALLELIZABLE APPROACH: PETRA
As with any model parallel training technique, PETRA requires to partition the network architecture into stages F,, that are distributed across distinct devices. Each device needs only to communicate with its neighboring devices 1 and j + 1. The pseudo-code in Alg. [ details the operations performed by each device, and the whole algorithm execution can be summarized as follows. The first device sequentially accesses mini-batcher ceiving its input - from the previous stage, each stage processes it in forward mode and passes it to the next stage, until the final stage is reached. The final stage evaluates the loss and computes the gradients with regard to its input and parameters, thus initiating the backward process, which is performed in parallel of the forward process. In it, each stage processes the input and its associated gradient from the next stage. This means first reconstructing the computational graph, either while reconstructing the input 2 for reversible stages or with forward pass as in activation check-
initiating the data propagation process. When re-
5

---

Published as a conference paper at ICLR 2025
Loss
Loss
Buffer size 2
Forward connection (activations)
Stage 2
Backward connection (gradients)
Stage
Buffer size
Backward connection (reconstructed activations)
Buffer size
Stage
Gradient computation
a) Delayed Gradients
b) PETRA
Figure 3: Comparison of memory use between PETRA and a standard Delayed Gradient method (Zuuang et al., 2020). By avoiding weight stashing and reversing the output into the in- put during the backward phase, we are able to fully decouple the forward and backward phases in all reversible stages, with no memory overhead, compared standard delayed gradient approaches.
pointing otherwise. Then, the parameter gradient approximation A and the input gradient computed before passing the latter to the previous stage. For intermediary reversible stages, this translates into the following equations, where corresponds the current time step of the training,
zj F5(2
= (
0et1 0,,F, 1'
0')
(5)
A 141 ( T
0C+1 Optimizer (0), A
Note that this complete set of equations effectively decouples communications, computations, and parameter updates between independent devices. Indeed, reversible stages are able to operate with- out maintaining any state between the forward and corresponding backward phase by simply avoid- ing weight stashing, similarly to Xu et al. (2019), and by reversing the output into the input during the backward phase, removing the need for an input buffer. As parameters are updated between the forward and backward phases, the reversible stage produces an approximate input reconstruction, thus evaluating gradients with an approximate set of inputs and parameters during the backward phase. We illustrate in Fig. the mechanism of PETRA compared to standard delayed gradient approaches that rely on additional buffers (Zhuang al. 2021b 2020
Complexity analysis. We now discuss the benefits of our method, which are summarized in Tab. 1 In this discussion, we assume a homogeneous setting in which almost identical stages are dis- tributed across J devices uniformly. First, we consider the backpropagation setting, assuming model parallelism strategy: a standard backpropagation pass requires storing locally both the pa- rameters and the computational graph and due the update lock of backpropagation (Jaderberg et al. 2017). requires synchronization between subsequent layers which impede the speed of com- putations. Standard Delayed Gradients strategies as implemented in Zhhang et al. (2021b) 2020) allow to unlock this barrier, but they require buffers for storing both the computational graph and parameters which can become impractical when using large models. In In|Xu al. 2019), an activation checkpointing strategy removes the need for storing parameters, yet it requires small computational overhead of 33% (assuming a backward pass is approximatively two times slower than forward pass, Fig. 6 of Huo et al. (2018b) and Mizutani & Dreyfus (200T)). To avoid storing activations, we rely on reversible architectures (Gomez et al., 2017) which increases the amount of forward communications by factor of 2 and backward communication by factor of 4 activations sizes

---

Published as a conference paper at ICLR 2025
Algorithm Worker perspective for training in parallel with PETRA, on stage j. assuming initial- ized parameters 0g and time step I, as well as an accumulation factor 1.
1: In parallel on the j-th stage, S < J, perform:
2: Forward Communications and Computations:
3: If j = 1 then
4: Zo Readiaaast
5: Else
6:
25-1 f Wait and Receive f-omj-1
7:
If stage] is not reversible:
8: Buffer; - 3ij
9:
2j (z),,.)))
10: Send to (+1(2,)))))))))
11:
Backward Communications and Computations:
12:
(ă f1+1) < Wait and Receive fomm++1
13:
If stage is reversible:
14:
ăj-1 t ( and keep computational graph in memory
(īj,0z)
Else
16:
8 + Buffer,
2j - (. 8,)
recompute the computational graph
8 + 1- 00j+1
A + A, + OgF, (3j-( -,,,,, 0 85+1
20:
If mod k = then:
Update parameters 9 with A,
22:
A 0
1+1+1
24:
Send (o5-1(
25:
26:
In parallel the final stage J. perform:
27:
ÆJ-1 <- Wait and Receive from
28:
L F,
29:
8, &
30:
AJ + C
31:
If mod k = then:
Update parameters 0, with A
33:
A
t +++
35:
Send to J (Z
double and one has pass both activations and gradients at the same time during backward. None of the aforementioned methods scale with the depth J: PETRA combines all the advantages of the previous methods, allowing efficient parallelization scaling linearly with memory overhead, while leading to a limited and constant overhead in computations and communications.
4 NUMERICAL EXPERIMENTS
4.1
CLASSIFICATION ACCURACY
We now describe experimental setup CIFAR-10
(Krizhevsky
2009
ImageNet-
32 Chrabaszez et al. 2017), and ImageNet (Deng et al. 2009).
Experimental setup. All our experiments use standard SGD optimizer with a Nesterov momen- tum factor of 0.9. We train all models for 300 epochs on CIFAR-10 and 90 epochs on ImageNet32 and ImageNet. We apply standard data augmentation, including horizontal flip. random cropping, and standard normalization but do not follow the more involved training settings of Wightman jet al. (202I), which potentially leads to higher accuracy. We perform a warm-up of 5 epochs where the learning rate linearly increases from 0.1, following Goyal al. (2017). Then, the learning

---

Published as a conference paper at ICLR 2025
Table Classification accuracies using our PETRA method with RevNets, compared stan-
dard backpropagation on ResNets and RevNets on CIFAR-10, R-00 ImageNet32, and ImageNet. Our method delivers competitive results with backpropagation, even on ImageNet.
Method
Model
Param. count
CIFAR-10
ImNet32
ImNet
Backprop
ResNet18 (PyTorch)
11.7M
69.8
Backprop
ResNet18 (Ours)
11.7M
95.0
54.0
70.8
Backprop
RevNet18 (Ours)
12.2M
94.9
54.6
70.8
PETRA
RevNet18 (Ours)
12.2M
94.9
54.6
71.0
Backprop
ResNet34 (PyTorch)
21.8M
73.3
Backprop
ResNet34 (Ours)
21.8M
95.5
56.5
74.0
Backprop
RevNet34 (Ours)
22.3M
95.3
56.4
73.2
PETRA
RevNet34 (Ours)
22.3M
94.8
56.1
73.5
Backprop
ResNet50 (Py (PyTorch)
25.6M
76.1
Backprop
ResNet50 (Ours)
25.6M
94.8
58.8
75.6
Backprop
RevNet50 (Ours)
30.4M
95.2
59.7
75.4
PETRA
RevNet50 (Ours)
30.4M
94.5
59.6
74.8
rate decayed by a factor of 0.1 at epochs 30, 60, and 80 for ImageNet32 and ImageNet is decayed at epochs 150 and 225 for CIFAR-10. We use weight decay of 5e-4 for CIFAR-10 and le-4 for ImageNet32 and ImageNet. As suggested in Goyal et al. (2017), we do not apply weight decay the batch norm learnable parameters and biases of affine and convolutional layers. For our standard backpropagation experiments, we follow the standard practice and use a batch size of 128 on ImageNet32 and CIFAR-10, 10, and 256 ImageNet32. However,
made few adaptations to train our models with PETRA. As suggested by Zhuang et al. (2020 202Ia). we employ an accumu- lation factor & and batch size of 64, which allows to reduce the effective staleness during training: in this case, k batches of data must he successively processed before updating the parameters of a stage (see Alg. [].. Such gradient accumulation however also increases the effective batch size, and we apply the training recipe used in Goyal et al. (2017) to adjust the learning rate; note that we use the average of the accumulated gradients instead of the sum. The base learning rate is thus given by 6ik with k the accumulation factor. the formula 1r = 0.1 256
Model adaptations. For designing our RevNet architectures, we adopt methodology similar to Gomez et al.] (2017): the number of channels in each stage is multiplied by 2 to account for the second data stream according to Fig. 2 However, as the stage function F,, operates only on one of the two streams, the number of parameters stays almost the same between residual block and its revertible counterpart. Consequently, the DNNs are split to preserve each residual block, resulting in 10 stages for RevNet18, and 18 stages for RevNet34 and RevNet50: thus varying the level of staleness between configurations. On CIFAR-10, the input layer uses 3x3 convolutions instead of 7x7 convolutions and does not perform max-pooling. The running statistics of batch normalization layers are updated when recomputing the activations during the backward pass and are then used during model evaluation the running statisties not updated during the forward pass.
Performance comparison. Tab. reports our numerical accuracy several vision datasets, comparing a backpropagation performance from an official PyTorch implementation of ResNets (the numbers can be found as of https: : http//tttccc.hcco org/ hub wttgghubbpyyoortttt httuub/ytoort__iiiiiioooo vision. resnet/). for our own implementation of ResNets and RevNets in our custom computational framework, and our proposed method, PETRA. For PETRA, we report the best classification aceu- racy after the last learning rate drop, using the best value (picked on the training set) of accumulation steps within {1, 2, 4, 8, 16, 32). Our CIFAR- accuracies are averaged over runs, with vari- ance smaller than 0.1. We observe that while our reversible models have about the same parameter count, they all perform in the same range of accuracy as their non-reversible counterparts. Only the RevNet-50 leads to small drop in accuracy pling layers removes this gap at the expense of a substantial increase the parameter count (30.4M to 50M). For the stake of comparison with the original ResNets, did not include this result.
ImageNet of about 0.6%: using different downsam-
8

---

Published as a conference paper at ICLR 2025
80.0
Backprop
Validation accuract
77.5
PETRA
75.0
70.0
67.5
65.0
62.5
60.0
10
15
25 30
Accumulation steps
Figure 4: Validation accuracy of PETRA and backpropagation for a various number of accu- mulation steps, for a RevNet18 trained on ImageNet with k E {1, 4, 8, 16, 32). The validation accuracies are averaged over the last 10 epochs. the number of accumulation steps increases, the effective staleness in PETRA decreases, closing the gap with standard backpropagation.
Table 3:
Memory savings for RevNet50 on ImageNet with our method for different configura- tions. We indicate the of memory buffers for inputs parameters. The savings computed with respect to the first configuration, where inputs and buffers are stored. Our method achieves 54.3% memory reduction over the base configuration of Delayed Gradients.
Buffer Params.
Memory (GB)
Saving (%)
Input
44.5
0.0
X
43.6
2.0
V 21.2
52.3
X
20.3
54.3
Impact the accumulation k. We the impact of the accumulation on a RevNet- 18 trained via PETRA for various values of accumulations & E {1, 2, 4, 8, 16, 32} on the ImageNet dataset. Fig. 4 indicates that our method can benefit from large accumulation factors, with the well-known trade-off of large batches mentioned in Goyal et al. (217)) Increasing the accumulation factor reduces the effective staleness during training. and closes the performance gap with standard backpropagation with perfect matching for k 32. This confirms that this large-batch training recipe derived for synchronous data parallelism is also particularly suited for our model parallel approach.
4.2 TECHNICAL DETAILS
A note on the implementation. We shortly describe our implementation details. We base our method PyTorch (Ansel et al., 2024), although we require significant modifications the Au- tograd framework in order to manage delayed first-order quantities consistently with PETRA. We rely heavily on the Vector Jacobian Product of PyTorch to compute gradients during the backward pass of each stage, but other backends could be used. The backward pass for reversible stages only necessitates a reconstruction step and a backward step - a naive implementation would use recon- struction step, followed by forward and a backward step. This is because we only need the output gradient as well as the computational graph of F3 compute the input and parameter gradients at line 12 and 13 of Alg. [] which can be obtained during the input reconstruction phase. For non- reversible stages,
reconstruct the computational graph with forward pass the input retrieved from the buffer during the backward pass. Our models can run on a single A100, 80GB to casily compare training dynamics, distributed over 10 or 18 GPUs when training with RevNet- 18 or RevNer-34 or 50.
Memory benefits and training time. Here, we discuss the practical memory savings and through- put speedup. To better understand the advantage of our method compared other delayed gradient

---

Published as a conference paper at ICLR 2025
Table Performance on CIFAR-100 with our method for different configurations, for different RevNets. Here, we use no accumulation (k 1) to better pinpoint the effect of the staleness. We indicate the use of memory buffers for inputs or parameters and whether using exact reversible backpropagation or using delayed gradients. We note that using the loss of both of the buffers does not have a particular effect on performance compared to the delay, further justifying our approach.
Delayed gradients
Buffer Params.
Accuracy RevNet-34
Input
RevNet-18
RevNet-50
77.5
78.1
78.5
VV
V
76.5
76.3
77.7
76.4
75.5
76.7
X
V
75.4
76.4
78.3
X
75.9
75.2
77.2
Table Mean iteration time of one micro-batch of size 256 of CIFAR-10 data, when training with PETRA or Reversible backpropagation. Training is distributed in both cases on 10 (for the RevNet- 18) or 18 GPUs (for the 3ev5e).34/50). We observe significant speed-up for the models.
Model
Rev. backprop.
PETRA
Speed-up
ReyNet-18
160.3ms
53.4ms
3.0x
RevNet-34
235.6ms
97.1ms
2.4x
approaches (Harlap et al. 2018 Xu et al. 2019 Kosson et al. 2021). we emphasize the practical memory savings associated with different methods in Tab. 3 We estimate the memory needed in gigabytes, as the sum of the size of the model, the input and parameter buffers (excluding the input buffer of the first stage, which uses retrievable dataset inputs). Here we do not include the effect of gradient accumulation, which depends on h and would only affect the parameter buffer size, which small in our case. Note that the batch size also affects the memory savings, and we set it to 64 for consistency with Tab. [2] Storing both inputs and parameters into buffer corresponds to the PipeDream approach (Harlap et al., 2018), while only storing inputs would correspond to the ap- proach in Xu et al. (2019); Kosson al. (202I). Not storing inputs (lines 3 and 4) is only applicable to reversible architectures. The input buffer has the biggest impact on the total memory, being re- sponsible for 52.3% of the memory footprint. Dropping the parameter buffer with PETRA pushes the savings further up to 54.3% for a ReyNet50 on ImageNet. We report on-device stage memory in Tab. 6 where note that non-reversible stages account for the majority of total memory use, indicating that savings would be much higher when using fully invertible architectures.
We also measure the effective throughput of training with PETRA ResNet-18 on 10 GPUs and a ResNer-34 on 18 GPUs, and compare it with basic model parallelism, where batch computations not overlapped between stages. After letting a warm-up period of 500 iterations, we report the effective throughput by measuring the processing time of 50 mini-batches and averaging it. As can be seen in Tab. our method achieves wall-clock time speedup compared basic model paral- lelism. While our main objective is to show the empirical effectiveness of our training procedure for convergence, we also observe significant speed-up, despite using relatively unbalanced stages. A more efficient implementation in the future will allow further training speed improvements.
Impact of the buffers on gradient estimation Three different approximations in gradient esti- mation are used PETRA which may affect our performances, and we investigate in Tab. their impact on CIFAR-100. The use of a gradient computed on previous iterations of the parameters is necessary to speed-up computations. This has the most impact on performances (line 2), but note that using accumulation vanishes this performance drop as shown in Fig. 4 Line 4 should corre- sponds to reversible backpropagation with delays, which is equivalent to line 2, but we choose to reconstruct the input using the latest parameters in the memory buffer, to better characterize the im- pact of reconstructing inputs with the latest parameters, as done PETRA. Lines and indicate that approximating the input affects the learning dynamics more than approximating the weights.
10

---

Published as a conference paper at ICLR 2025
5 CONCLUSION
In this work, we introduce PETRA, novel model parallel training technique for reversible ar- chitectures which is a novel promising alternative backpropagation. It achieves a significant parallelization with a limited overhead compared to standard backpropagation or other competitive alternatives to end-to-end training, like delayed gradients approaches. Our method has the potential to achieve linear speedup compared to standard backpropagation and allows reversible layers to op- crate without any parameter or activation buffers, effectively decoupling the forward and backward phases. Despite using approximate delayed gradient estimate, method delivers competitive performances compared to standard backpropagation on standard computer vision datasets.
In future work, aim implement and optimize PETRA for Large Language Models (LLMs), with a first baseline being Reformers (Kitaev et al.] 2020), invertible transformers that have been shown scale. This will validate further PETRA's effectiveness and robustness on felly-reversible architectures, solidifying its potential as a cutting-edge training technique.
ACKNOWLEDGMENTS
This work was supported by Project ANR-21-CE23-0NE ADONIS and PEPR IA on grant SHARP ANR-23-PER 23-PEIA-0008 and Sorbonne Center for Artificial Intelligence (SCAI) of Sorbonne Univer- sity (IDEX SUPER 1IIDDEX-0004). EB acknowledges funding from FRQNT New Scholar and support for computation from CFL SR acknowledges funding from European Union under GA no. 135782 (MANOLO project). This work was granted access to the AI resources of IDRIS under the allocations 2023-A0151014526 made by GENCI.
REFERENCES
Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300-22312, 2022.
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky.
Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary Vito,, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael La- zos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk. Shunting Zhang, Michael Suo, Phil Tillet, Xu Zhao, Fikan Wang, Keren Zhou, Richard Zou, Xiaodong Wang, Ajit Mathews, William Wen, Gregory Chanan, Peng Wu, and Soumith Chintala. Py- torch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Sup- port for Programming Languages and Operating Systems, Volume 2, ASPLOS '24, pp. 929-947, New York, NY, USA, 2024. Association for Computing Machinery. ISBN 9798400703850, doi: 10.1145/36206553664666 36446666.... URL //t::///oiii org/10 1145 620655666666666 3640366
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to imagenet. In International conference on machine learning, pp. 583-593. PMLR, 2019.
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Decoupled greedy learning of cnns. In International Conference on Machine Learning, pp. 745- PMLR, 2020.
Eugene Belilovsky, Louis Leconte, Lucas Caccia, Michael Eickenberg, and Edouard Oyallon. De- coupled greedy learning of cnns for synchronous and asynchronous distributed learning. arXiv preprint arXiv:2 106.006666 106.06401, 2021.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016.
Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. downsampled variant of imagenet as an alternative to the cifar datasets, 2017.
11

---

Published as a conference paper at ICLR 2025
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale scale hi- erarchical image database. In 2009 JEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components esti- mation. arXiv preprini arXiv: 140..511, 2014.
Shiqing Fan, Yi Rong, Chen Meng. Zongyan Cao, Siyu Wang, Zhen Zheng, Chuan Wu, Guoping Long, Jun Yang. Lixue Xia, et al. Dapple: A pipelined data parallel approach for training large models. In Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, pp. 431 145, 2021.
Louis Fournier, Stéphane Rivaud, Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Can forward gradient match backpropagation? Fortieth International Conference on Machine Learming. 2023.
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger Grosse. The reversible residual net- work: Backpropagation without storing activations. Advances in neural information processing systems, 30, 2017.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint (xXiv:/706.06677 2017.
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint ar arXiv:1806.0333 2018.
Yanping Huang, Youlong Cheng. Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee. Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Zhouyuan Huo, Bin Gu, and Heng Huang. Training neural networks using features replay, Advances in Neural Information Processing Systems, 31, 201 Sa.
Zhouyuan Huo, Bin Gu, Heng Huang, et al. Decoupled parallel backpropagation with convergence guarantee. International Conference on Machine Learning, pp. 2098-2106. PMLR, 2018b.
Jörn-Henrik Jacobsen, Aroold Smeulders, and Edouard Oyallon. i-revnet: Deep invertible networks. arXiv preprint arxiv:1802.00000000 2018a.
Jörn-Henrik Jacobsen, Amoll W. M. Smeulders, and Edouard Oyallon. i-revnet: Deep invertible networks. ArXiv, abs/ 180.00888,, 2018b. URL ott.aars../pppp..eetaaaaaaaaaaa.
CorpusID:34333
Max Jaderberg, Wojciech Marian Czarnecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled neural interfaces using synthetic gradients. In Inter- national conference machine learning, pp. 1627-1635. PMLR, 2017.
Chiheon Kim, Heungsub Lee, Myungryong Jeong, Woonhyuk Back, Boogeon Yoon, Ildoo Kim, Sungbin Lim, and Sungwoong Kim. torchgpipe: On-the-fly pipeline parallelism for training giant models, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
Atli Kosson, Vitaliy Chiley, Abhinav Venigalla, Joel Hestness, and Urs Koster. Pipelined backprop- agation at scale: training large models without batches. Proceedings of Machine Learning and Systems, 3:479-501, 2021.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. 7/api semant i cacholas wttpg//oopss..11020000
URL https
12

---

Published as a conference paper at ICLR 2025
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, al. Pytorch distributed: Experiences accelerating data parallel training. arXiv preprint arXiv:2006. 15704, 2020.
Yuliang Liu, Shenggui Li, Jiarui Fang, Yanjun Shao, Boyuan Yao, and Yang You. Colossal-auto: Unified automation of parallelization and activation checkpoint for large-scale models, 2023.
Mateusz Malinowski, Dimitrios Vytiniotis, Grzegorz Swirszcz, Viorica Patraucean, and Joao Car- reira. Gradient forward-propagatin propagation for large-scale temporal video modelling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9249-9259, 2021.
Karttikeya Mangalam, Haoqi Fan, Yanghao Li, Chao-Yuan Wu, Bo Xiong, Christoph Feichtenhofer, and Jitendra Malik. Reversible vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10830-10840, 2022.
E. Mizutani and S.E. Dreyfus. On complexity analysis of supervised mip-learning for algorithmic comparisons. In IJCNNO1. International Joint Conference Neural Networks. Proceedings (Cat. No.01CH37222), volume 1, pp. 347-352 vol.1, 2001. doi: 1091 9999044.
Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gre- gory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline par- allelism for dnn training. In Proceedings of the 27th ACM Symposium on Operating Systems Principlex, pp. 1-15, 2019.
Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline-parallel dan training. In International Conference on Machine Learning, pp. 7937-7947. PMLR, 2021.
Arild Nøkland and Lars Hiller Eidnes. Training neural networks with local error signals. In Inter- national conference on machine learning, pp. 4839 4850. PMLR, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models, 2020.
Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Min- jia Zhang. Dong Li, and Yuxiong He. {Zero-offload]: Democratizing (billion-scale) model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pp. 551-564, 2021.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv: (ate::999000000 0000000033,..... 2019.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, al. Using deep- speed and megatron to train megatron-turing nlg 530b, large-scale generative language model. arXiv preprint arXiv:220..11991 2022.
Yulin Wang. Zanlin Ni, Shiji Song, Le Yang, and Gao Huang. Revisiting locally supervised learning: an alternative to end-to-end training. arXiv preprint 1rXiv:2101.10832. 2021.
Ross Wightman, Hugo Touvron, and Hervé Jégou. Resnet strikes back: An improved training procedure timm, 2021.
Xu, Zhouyuan Huo, and Heng Huang. On the acceleration of deep learning model paral- lelism with staleness. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recog- nition (CVPR), pp. 2085-2094, 2019. URL https://appi www.sticscooleee.....o.. org/
CorpusID:219633AM
Bowen Yang, Jian Zhang, Jonathan Li, Christopher Ré, Christopher Aberger, and Christopher De Sa. Pipemare: Asynchronous pipeline parallel dnn training. Proceedings of Machine Learning and Systems, 3:269-296 2021.
13

---

Published as a conference paper at ICLR 2025
Huiping Zhuang, Zhiping Lin, and Kar-Ann Toh. Accumulated decoupled learning: Mitigating gradient staleness later- model parallelization. arXiv preprint arXiv:2012.03747, 2020.
Huiping Zhuang, Yi Wang, Qinglai Liu, and Zhiping Lin. Fully decoupled neural network learning using delayed gradients. IEEE transactions on neural networks and learning systems, 33(10); 6013-6020, 2021a.
Huiping Zhuang. Zhenyu Weng, Fulin Luo, Toh Kar-Ann, Haizhou Li, and Zhiping Lin. Accumu- lated decoupled learning with gradient staleness mitigation for convolutional neural networks, In International Conference Machine Learning, pp. 12935-12944. PMLR, 2021b.
14

---

Published as a conference paper at ICLR 2025
Table 6: Memory usage of each stage when training with PETRA on CIFAR-IO with a batch size of 256. Stages with the same memory usage are grouped together for ease of read.
RevNet-18
Stage(s) index(es)
0
1-2
5
7
Memory (GB)
1,15
1,01
2,67
0,52
1,13
0,31
0,57
0,25
0,11
RevNet-34
Stage(s) index(es)
0
1-4
6-8
10-12
13
14-16
17
Memory (GB)
1,15
1,01
4,2
0,52
1,66
0,31
0,71
0,25
0,11
A MEMORY USAGE BY STAGE
In Tab. 6 we report the memory usage with PETRA of each stage of the networks trained. Note that some stages have a much higher memory footprint than others, which is due the presence of non-reversible stages in RevNets, which requires activation buffers.
B QUALITY OF GRADIENT APPROXIMATION AND DELAYS IN DEPTH
PETRA gradient approximations The PETRA optimization procedure estimates gradients with two approximations. First, PETRA estimates delayed gradients, with a delay Tj 2(J- for each block j of a network partitioned into J blocks. According to Eq. 5 for the j-th layer, this would mean that A (z) 33) The implicit underlying hypothesis is that the delayed gradients can serve as an approximation the end-to-end gradient, or, more largely, as a descent direction. Note that this is the gradient computed by standard approaches like Zhuang et al. (2020). Then, PETRA makes further approximations on this delayed gradient. First, having no parameter buffers, each layer will only the latest available in memory weights fj, which will differ from
those used in the forward pass o¹ T. Second, the input used for gradient computations is not stored buffer for reversible layers, but reconstructed. For both the input reconstruction and the Jacobian computations, the layers use the "up-to-date" parameters 0j.
Framework To investigate the quality of these approximations empirically, we trained a Re on CIFAR-10, divided into 10 stages, while tracking approximation quality metrics for each layer throughout training. These metrics are computed 15 times per epoch and averaged at the end of the epoch. The optimization procedure is the same the used to obtain Tab. without gradient accumulation to emphasize the impact of delay. Note that the network has non-reversible stages, and thus input buffers, at stages {3, 7): the first stage is not reversible but can retrieve its input from the dataset. Our model is trained via PETRA, and we compare, at given snapshots throughout the training. gradients following PETRA, standard delayed gradient approaches Zhuang al. (2020), and standard backprop gradients at various depths.
To compare two different gradients, the two metries we record are their cosine similarity and their norm ratio:
Cos-Sim(x,,)) = (x, v)
Norm-Ratio(z, y) ||||| |g>||
||-||| ||vll
Approximation with the standard delayed gradient We report these metrics throughout training in Fig. [5] We observe several tendencies. First, in Fig. [a we observe that the gradient computed by PETRA is indeed a good approximation of the standard delayed gradient, as predicted. The alignment improves during training, with particular jump at the last learning rate drop. This
expected as the discrepancy between 0, and 8²-7 becomes negligible as the model converges. Similarly, the alignment improves for later layers, where fewer delayed parameters are used, and with a smaller delay. The reconstruction error also lower, although this error "resets" after each input buffer. Note also that the ratio of the norms of the gradients, in Fig. 5c follows a similar trend, and stays consistently close to 1.
Approximation with the end-to-end gradient Then, in 5b also compare the alignment of the PETRA gradient with the end-to-end one and the alignment of the delayed gradient with the end-to-end one. These values are much lower for early layers, where the gradient computed
15

---

Published as a conference paper at ICLR 2025
Nocta delayed grudkonts
gradicnt
Stape
15.0
Alignment of PITTA with the condent dwayed gradint
1.0
Grealee
Epoch
100
1.0
203
5.00
300
(b) Cosine similaritics between PETRA and end- to-end gradients (full line), and between the stan- dard delayed gradients (dashed line) and end-to- end
Epech
(a) Cosine similarity between PETRA and stan- dard Delayed Gradients.
17.3
Norm motee FLTIM delayed gradienss wita the exc enc-br-end gracient
Batto the some betivenn PETRA and the calayed gradient
15.3
1.0
12.5
50.3
Grublet
3.2
2.3
Upart
1:0
1:5
200
2.30
300
(d) Norm ratios between PETRA and end-to-end gradients (full line) and the standard delayed gra- dient and end-to-end gradient (dashed line).
(c) Norm ratio between PETRA (numerator) and standard delayed gradients (denominator).
Figure 5: Cosine similarities and norm ratios between gradients throughout training. Each point rep- resents the average of 15 measurements during 1 epoch. Values smoothed with rolling window of size 10. Color corresponds to the stage index. The approximation is noticeably better after the last learning rate drop, and for later stages. Although PETRA approximates well the standard delay gradient, it also approximates better the end-to-end gradient compared to standard delay gradient approaches.
16

---

Published as a conference paper at ICLR 2025
by both methods depends much more the delay, and is thus less aligned with the end-to-end gradient. However, we surprisingly observe that the gradient computed by PETRA shows a better alignment with the end-to-end gradient compared the delayed one. Although surprising, this can be explained by the fact that PETRA does not use delayed parameters during the backward pass, for both the input reconstruction and the Jacobian computations. Although reducing the alignment with the delayed gradient, this increases the alignment with the end-to-end one. The norm ratio between the delayed and PETRA gradients and the end-to-end ones is also quite high early in training, before coming close to at convergence. Here again, the ratio is smaller for the PETRA gradient.
(a) Cosine similarities at epoch 25, i.e. 10 epoch after warm-up phase.
(b) Cosine similarities at epoch 125, i.e. 25 epoch before the first learning rate drop.
Costine Depch (spoch 17g)
Coalne Dopth
Drail pare PCTM in Der
P97TI tn Diiend
Lager
(c) Cosine similaritics epoch 175, i.e. 25 epoch after the first learning rate drop.
(d) Cosine similarities epoch 125, i.e. 25 epoch after the second learning rate drop.
Figure 6: Cosine similarities against stage index of RevNet18 on CIFAR-10 at epochs 25, 125, 175 and 250. Approximations degrades between epoch 25 and 125, consistently with Fig. but improves noticeably after each learning rate drop.
For a clearer comparison, we also provide in Fig. 6 cosine similarities values between the three gradients depending on the layer index, at different training epochs. We once again observe that the similarities improve as layer indexes increase and that the discrepancies between the gradients decrease across stages as the learning rate becomes smaller. We also observe that the similarity between PETRA and end-to-end is sometimes higher than for the delayed gradient, suggesting that using the up-to-date weights for the Jacobian computations and input reconstructions might help to mitigate the staleness.
17


Reviews:
[{"review_id": "T5hGsFAPG5", "reviewer": "Reviewer_A9EF", "summary": "The authors propose a new algorithm for training reversible models. Compared to backpropagation, it can be run on each layer in parallel and with a reduced memory cost. They show empirically the advantages of their algorithm on RevNet models for image classification.", "questions": "- How is the approximated gradient influenced by the depth of the model? I would expect the error to increase as the model gets deeper.\n\nI find the paper very interesting and am ready to increase my grade should my remarks be addressed by the authors.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is well written, clear, and has helpful illustrations.\n- The algorithm seems simple, natural and intuitive.\n- While the algorithm relies on reversible layers, it can still be mixed with standard non-reversible layers, for which a standard backpropagation is performed.\n- The authors validate their algorithm with thorough experiments and analyses.", "weaknesses": "1. Invertible networks are currently not very used. This limits the direct applications of the algorithm. However I am aware that PETRA could motivate such the use of such architectures.\n2. The experiments are only performed on RevNet models for image classification. As mentioned in the conclusion, it would be very nice to see experiments on more tasks and models. Indeed, as PETRA is applicable to only a subset of models (reversible models), it is frustrating to only see experiments on a single architecture.\n3. Lines 509-510: I think you meant RevNet instead of ResNet.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8RNNiqYmW6", "reviewer": "Reviewer_hEA1", "summary": "PETRA is a model-parallel training method for reversible neural networks that decouples forward and backward passes, eliminating the need for activation or parameter buffers. This enables efficient parallel computation across devices with reduced memory overhead. PETRA matches backpropagation in accuracy on datasets like CIFAR-10 and ImageNet, while also achieving notable speed and memory savings, making it a potential alternative for large-model training.", "questions": "How PETRA perform on large model and more complex task, such as pretraining language model? The experiment in the paper is weak. The scalability of PETRA can not be verified by the current empirical results. Experiments on distributed pretraining for llm is necessary to validate the efficiency of PETRA, for example: experiments on Pile dataset with varying model size.\n\nIs the reversible architecture necessary for PETRA? For models that integrate both reversible and non-reversible layers, how does PETRA manage memory savings and efficiency, and could these hybrid architectures affect its scalability benefits?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The PETRA paper presents a new alternative for large-scale neural network training, offering efficient parallelization by decoupling forward and backward passes, which enables stages to compute independently across devices. Utilizing reversible architectures, PETRA removes the need for activation and parameter storage, achieving up to 54.3% memory savings, making it especially valuable for training large models. It demonstrates accuracy comparable with backpropagation on datasets like CIFAR-10 and ImageNet.", "weaknesses": "Dependency on Reversible Architectures: The approach is designed specifically for reversible architectures, which may limit its application to models that can be easily adapted to this structure. Non-reversible architectures, such as standard ResNets or some types of transformers, may not benefit as fully from PETRA’s memory and efficiency gains.\nIncreased Communication Overhead: While PETRA reduces memory usage, its reversible stages require additional communication overhead during the backward pass, which could affect scalability on very large, distributed systems. And the PETRA propose dividing a model into some\nScalability Constraints with Non-Reversible Layers: Although PETRA performs well on reversible architectures, any non-reversible stages still require stored activations, potentially increasing memory use and complicating scalability for models that include such layers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MNCS4Dyc7L", "reviewer": "Reviewer_md3p", "summary": "The paper proposes to perform model parallel training using reversible architectures. Compared to delayed gradient, the proposed method is more memory efficient since it does not need to stash weights. It is shown that on shallower architecture the performance is slightly better than regular backprop and on deeper architecture such as ResNet-50, there is a slight drop but not significant. Overall, the work is likely to have a big impact as a way to scale up model parallel training.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper demonstrated that activation reconstruction can work well with out-of-sync backward weights, and the reconstructed activations can be used to update weights.\n- The paper has shown real computation and memory savings.", "weaknesses": "- It would be nice to see at what scale the method starts to break down (say when there is more and more delay in reconstruction). And show a plot on reconstruction error and final performance as a function of the number of delay steps. The model depth can be another variable to explore, aside from the few standard model architectures, perhaps sweeping a wider range of depths.\n- Algorithm 1 is a little hard to process.\n- The method relies on gradient accumulation to fully match with the It is unclear to me how gradient accumulation would have any impact when a large batch / data parallel is employed. This may not be a concern for LLMs, but for ImageNet and SSL training, many use very large batch sizes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qCibJSzdOP", "reviewer": "Reviewer_ggqx", "summary": "This paper proposes to utilize the concept of reversible architectures to improve parallelization in DNN training. A model is split into multiple stages that are trained asynchronously; i.e. in a model parallel fashion. Leveraging reversibility, the training of the different stages is effectively decoupled. This scheme offers a linear speedup in the number of stages relative to end-to-end backprop, while reducing the memory footprint. The method is evaluated using ResNets/RevNets with three different image classification benchmarks.", "questions": "How did you partition the architectures for your experiments? How many layers/blocks in each stage? Were they all the same size? And if so, would that not bring them out of sync during training such that top layers/stage were idle a lot of the time? The size of the feature maps is decreasing in the layer index, no? Thus, the lower layers/stages would consume more memory and compute than the top ones?\n\nPerhaps you could add some information about this in the appendix :-)", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well-written and easy to follow. The idea of utilizing reversibility for parallelization is a nice, simple, and novel idea! Consequently, I find myself sufficiently convinced that the method works --- albeit, that the empirical evaluation is somewhat limited. The novelty and applicability of the method mostly outweighs my concerns about the evaluation.", "weaknesses": "My only objection to this work is the limited number of experiments. They are limited to ResNet/Revnet 18/34/50 and CIFAR10, ImageNet-32, and ImageNet. It would definitely improve the paper to have at least a few more architectures included.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)