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

## Abstract

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 CIFAR-10, ImageNet32, and ImageNet, achieving competitive accuracies comparable to backpropagation using ResNet-18, ResNet-34, and ResNet-50 models.

## 1 Introduction

First-order methods using stochastic gradients computed via backpropagation on mini-batches are the de-facto standard for computing parameter updates in Deep Neural Networks []. As datasets and models continue to grow [] there is an urgent need for memory-efficient and scalable parallelization of deep learning training across multiple workers. Data parallelism via mini-batches [] has been widely adopted in deep learning frameworks []. This approach computes gradients across model replicas distributed among workers, yet it requires frequent synchronization to aggregate gradients, leading to high communication 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 []. To further address these issues, methods have attempted to mitigate this memory overhead and to parallelize the sequential backpropagation steps themselves across devices, while computing exact gradients. Techniques like optimizer sharding [], tensor parallelism [], activation checkpointing [], or pipelining [], have been deployed individually or combined, leading for instance to the development of 3D parallelism [], a popular methodology which improves the efficiency of the backpropagation implementation. On the other hand, the fundamental inefficiency underlying the parallelization of backpropagation has not been addressed by these methods.

However, the use of exact gradient restricts algorithmic choices and parallel implementations, as highlighted by []. For instance, backpropagation is backward locked: the inputs of each layer must be propagated through the network and preserved until an error signal is retropropagated to the layer of origin. This requirement enforces a synchronous dependency among subsequent layers and requires them to systematically store intermediary activations, potentially impeding overall resource efficiency as workers must wait for each other to continue their computations and release memory used for activations. To unlock the potential of backpropagation, inexact backpropagation 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 parameters or activations, such as double-buffered pipelining [] or delayed gradient approaches []. However, these methods introduce significant memory overhead due to the use of ad hoc buffers for activations, parameters, or both. Following an opposite direction, local learning methods [, ], which estimate inexact gradients via a local auxiliary neural network, pave the way to parallel gradient computations but often lead to unrecoverable performance drops []. This underscores the need for a robust alternative to backpropagation, with limited memory overhead.

Figure: Figure 1: Comparison of PETRA with standard backpropagation. This approach splits the stages of a model and decouples their forward and backward passes, resulting in a sixfold increase in parallelization speed in this example.
Refer to caption: /html/2406.02052/assets/x1.png

In this work, we introduce PETRA (Parallel End-to-End Training with Reversible Architectures), a novel method designed to 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 communication costs. Furthermore, reversibility is a minor requirement, as many studies have demonstrated that standard architectures can be adapted into reversible ones without any performance drops [, , , ]. By allowing parameters to evolve in parallel and by computing an approximate inversion during backward, we propose an effective alternative to backpropagation which allows high model parallelism with a constant communication overhead and no additional buffers. In fact, for a constant increase in communication overhead, PETRA achieves a linear speedup compared to standard backpropagation with respect to the number $J$ of stages the network is split into. We illustrate our approach in Fig. [1](#S1.F1), by contrasting the evolution of PETRA with a standard backpropagation pass.

#### Contributions.

Our contributions are as follows: (1) We introduce PETRA, a streamlined approach for parallelizing the training of reversible architectures. This method leverages a delayed, approximate inversion of activations during the backward pass, allowing for enhanced computational efficiency. (2) Our technique significantly reduces memory overhead by minimizing the necessity to store extensive computational graphs. (3) It enables the parallelization of forward and backward pass computations across multiple devices, effectively distributing the workload and reducing training time. (4) We validate the efficacy of PETRA through rigorous testing on benchmark datasets such as CIFAR-10, ImageNet-32, and ImageNet, where it demonstrates robust performance with minimal impact on accuracy. (5) Additionally, we provide a flexible reimplementation of the autograd system in PyTorch, specifically tailored for our experimental setup, which is available at [https://github.com/streethagore/PETRA](https://github.com/streethagore/PETRA).

## 2 Related work

#### Reversible architectures.

Reversible DNNs are composed of layers that are invertible, meaning that the input of a 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 [], at the cost of an extra computation per layer. Invertible networks further improve this method by removing dimensionality reduction steps such as downsamplings, making the networks fully invertible []. Reversibility is not restricted to a type of architecture or tasks and has been extensively used for generative models [], for ResNets [], and Transformers []. However, as far as we know, reversible architectures have never been used to enhance parallelization capabilities.

#### Alternatives to backpropagation.

Multiple alternatives to backpropagation have been proposed previously to improve over its computational efficiency. For instance, DNI [] is the first to mention the backpropagation inefficiency and its inherent synchronization locks. However, they address those locks with a method non-competitive with simple baselines. Local (or greedy) learning [, ] propose to use layerwise losses to decouple the training of layers, allowing them to train in parallel []. Local learning in videos [] notably uses the similarity between successive temporal features to remove buffer memory. However, the difference in training dynamics between local training and backpropagation still limits such approaches [, ].

#### Pipeline parallelism.

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 [], a batch of data is divided into micro-batches that are processed independently at each stage. Although more efficient pipelining schedules have been proposed [], notably to mitigate the peak memory overhead, keeping an exact batch gradient computation requires leaving a bubble of idle workers. By alternating one forward and one backward pass for each worker, PipeDream [] can allow to get rid of idleness bubbles, but at the expense of introducing staleness in the gradients used. [] 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 a quadratic activation memory overhead with regard to the number of stages, as micro-batch activations pile up in buffers, especially for early layers. Some implementations propose to limit this overhead by combining activation checkpointing [] with pipelining [, ], 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 approaches, delays occur stage-wise: the backward pass may be computed with outdated parameters or activations compared to the forward pass. For instance, [] 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 a lock on computations. In [] and [], stale gradients are computed from older parameter versions differing from the parameters used during the update. This staleness can be mitigated: [] ’shrinks’ the gradient by the delay value, but more advanced techniques also exist [, ]. Still, these methods are limited like previous pipelining methods by their memory overhead as the computational graph is fully stored. A first step to reduce this, as proposed in Diversely Stale Parameters (DSP) [], PipeMare [] and [], is to keep a 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 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.

Figure: Figure 2: Differences between the residual block of a ResNet and its reversible counterpart. (a) Forward of a residual block. (b) Forward and (c) Reverse forward of a reversible residual block. For reversible blocks, similarly to [], the input $x_{j}$ is doubled in size and split equally into $\{x_{j}^{1},x_{j}^{2}\}$ along the channel dimension. The function $\mathcal{F}_{j}$ includes a skip-connection while $\tilde{\mathcal{F}}_{j}$ does not.
Refer to caption: /html/2406.02052/assets/x2.png

## 3 Method

### 3.1 Standard backpropagation

We consider a DNN composed of $J$ stages (e.g., a layer or a set of layers). An input $x_{0}$ is propagated through the network, recursively defined by

$$ $x_{j}\triangleq F_{j}(x_{j-1},\theta_{j})\,,$ (1) $$

where $F_{j}$ is the $j$-th stage parameterized by $\theta_{j}$. The backpropagation algorithm is the ubiquitous algorithm to compute parameter gradients. First, an input is propagated through the network with a forward pass, while storing its intermediate activations. A scalar loss $\mathcal{L}$ is then deduced from the corresponding output $x_{J}$. Parameter gradients are then computed during the backward pass by taking advantage of the chain rule: starting from the last stage with $\delta_{J}=\nabla_{x_{J}}\mathcal{L}$, the gradients with regard to the activations are given by

$$ $\delta_{j}\triangleq\nabla_{x_{j-1}}\mathcal{L}=\partial_{x}F_{j}(x_{j-1},\theta_{j})^{\mathrm{T}}\delta_{j+1}\,,$ (2) $$

and the gradients with regard to the parameters are defined as

$$ $\Delta_{j}\triangleq\nabla_{\theta_{j}}\mathcal{L}=\partial_{\theta}F_{j}(x_{j-1},\theta_{j})^{\mathrm{T}}\delta_{j+1}\,.$ (3) $$

Note that these computations follow a synchronous and sequential order.
The parameters $\theta_{j}$ can then be updated given their gradient estimate $\Delta_{j}$, using any optimizer.

### 3.2 Reversible architectures

We focus on the reversible neural networks presented in [], although our method is not dependent on this architecture. In practice, only a few stages which do not preserve feature dimensionality are not reversible and correspond to the downsampling blocks in the ResNet. Fig. [2](#S2.F2) highlights how reversible residual blocks $F_{j}$ differ from their standard counterpart. The input is split into two equal-size inputs, along the channel dimension, that are propagated forward according to Fig. [2](#S2.F2) using an ad-hoc operator $\tilde{\mathcal{F}}_{j}$. It can be reconstructed by reverse propagating the output according to Fig. [2](#S2.F2), by subtracting the output of $\tilde{\mathcal{F}}_{j}$ rather than adding it like in the previous forward.

#### Reversible stages.

In order to compute the exact gradients during the backpropagation phase, each reversible stage needs to retrieve its output from the stage above. We note $F_{j}^{-1}$ the reverse stage function, which reconstructs the input from the output. We recursively apply the reconstruction to the final activation $x_{J}$, such that

$$ $\displaystyle\begin{bmatrix}x_{j-1}\\ \delta_{j}\end{bmatrix}=\begin{bmatrix}F_{j}^{-1}(x_{j},\theta_{j})\\ \partial_{x}F_{j}(F_{j}^{-1}(x_{j},\theta_{j}),\theta_{j})^{\mathrm{T}}\delta_{j+1}\end{bmatrix}\,.$ (4) $$

Note that reconstructing the input in our procedure is computationally equivalent to recomputing the activations in activation checkpointing, meaning it is equivalent to a 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 $x_{j-1}$ must be sent to the reversible stages, this doubles the cost of backward communications.

#### Non-reversible stages.

In practice, a reversible architecture includes layers that reduce dimensionality for computational efficiency, which thus correspond to non-invertible functions. For those very few stages, we employ a buffer mechanism to store activations and, like activation checkpointing, we recompute the computational graph with a forward pass during the backward pass. Note that this would not be the case for invertible (i.e., bijective) architectures [], which use an invertible downsampling.

### 3.3 A parallelizable approach: PETRA

**Table 1: Comparisons with other methods in an ideal setting for one stage. We compare several methods to compute a gradient estimate in a model parallel setting. Here, $J$ is the total number of stages while $j$ is the stage index. For the sake of simplicity, we assume that a backward pass requires approximately 2 times more FLOPs than a forward pass. Full Graph indicates 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 can be executed in parallel for PETRA or delayed gradients, making the backward pass responsible for most of the computation time in parallelizable approaches.**
|  | Storage | Comm. | FLOPs | Mean time |  |
| --- | --- | --- | --- | --- | --- |
| Methods | Activations | Params. | Volume |  | per batch |
| Backpropagation | Full Graph (FG) | $1$ | $1$ | $3J$ | $3J$ |
| Reversible backprop. [] | 𝟎 0 | $1$ | $4$ | $4J$ | $4J$ |
| Delayed gradients [] | $2(J-j)\times$ FG | $\frac{2(J-j)}{k}$ | $1$ | $3J$ | $2$ |
| + Checkpointing [] | $2(J-j)$ | $1$ | $1$ | $4J$ | $3$ |
| PETRA (ours) | 𝟎 0 | $1$ | 4 | $4J$ | $3$ |

As with any model parallel training technique, PETRA requires to partition the network architecture into stages $F_{j}$ that are distributed across distinct devices. Each device $j$ needs only to communicate with its neighboring devices $j-1$ and $j+1$. The pseudo-code in Alg. [1](#alg1) details the operations performed by each device, and the whole algorithm execution can be summarized as follows. The first device sequentially accesses mini-batches, initiating the data propagation process. When receiving its input $x^{t}_{j-1}$ 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 $\tilde{x}^{t}_{j-1}$ for reversible stages or with a forward pass as in activation checkpointing otherwise. Then, the parameter gradient approximation $\Delta^{t+1}_{j}$ and the input gradient are computed before passing the latter to the previous stage. For intermediary reversible stages, this translates into the following equations, where $t$ corresponds to the current time step of the training,

$$ $\displaystyle\left\{\begin{aligned} x_{j}^{t+1}&=F_{j}(x_{j-1}^{t},\theta_{j}^{t})\\ \tilde{x}^{t+1}_{j-1}&=F_{j}^{-1}(\tilde{x}^{t}_{j},\theta_{j}^{t})\\ \delta_{j}^{t+1}&=\partial_{x}F_{j}(\tilde{x}_{j-1}^{t+1},\theta_{j}^{t})^{\mathrm{T}}\delta^{t}_{j+1}\\ \Delta^{t+1}_{j}&=\partial_{\theta}F_{j}(\tilde{x}_{j-1}^{t+1},\theta_{j}^{t})^{\mathrm{T}}\delta^{t}_{j+1}\\ \theta^{t+1}_{j}&=\text{Optimizer}^{t}_{j}(\theta^{t}_{j},\Delta_{j}^{t+1})\,.\end{aligned}\right.$ (5) $$

Note that this complete set of equations effectively decouples communications, computations, and parameter updates between independent devices. Indeed, reversible stages are able to operate without maintaining any state between the forward and corresponding backward phase by simply avoiding weight stashing, similarly to [], 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. [3](#S3.F3) the mechanism of PETRA compared to standard delayed gradient approaches that rely on additional buffers [, ].

Figure: Figure 3: Comparison of our PETRA method to a standard Delayed Gradient method []. By avoiding weight stashing and reversing the output into the input during the backward phase, we are able to fully decouple the forward and backward phases in all reversible stages, with no memory overhead, compared to standard delayed gradient approaches.
Refer to caption: /html/2406.02052/assets/x3.png

#### Complexity analysis.

We now discuss the benefits of our method, which are summarized in Tab. [1](#S3.T1). In this discussion, we assume a homogeneous setting in which almost identical stages are distributed across $J$ devices uniformly. First, we consider the backpropagation setting, assuming a model parallelism strategy: a standard backpropagation pass requires storing locally both the parameters and the computational graph and due to the update lock of backpropagation [], requires synchronization between subsequent layers which impede the speed of computations. Standard Delayed Gradients strategies as implemented in [, ] 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 [], an activation checkpointing strategy removes the need for storing parameters, yet it requires a small computational overhead of $33\%$ (assuming a backward pass is approximatively two times slower than a forward pass, see Fig. 6 of [] and []). To avoid storing activations, we rely on reversible architectures [] which increases the amount of forward communications by a factor of 2 and backward communication by a factor of 4 – activations sizes double and one has to 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 an efficient parallelization, while leading to a limited overhead in computations and communications.

Figure: Algorithm 1 Worker perspective for training in parallel with PETRA, on a stage $j$, assuming initialized parameters $\theta_{j}$ and time step $t$, as well as an accumulation factor $k>1$.

## 4 Numerical experiments

### 4.1 Classification accuracy

We now describe our experimental setup on CIFAR-10 [], ImageNet-32 [], and ImageNet [].

#### Experimental setup.

All our experiments use a standard SGD optimizer with a Nesterov momentum 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 we do not follow the more involved training settings of [], which potentially leads to higher accuracy. We perform a warm-up of 5 epochs where the learning rate linearly increases from 0 0 to $0.1$, following  []. Then, the learning rate is decayed by a factor of $0.1$ at epochs 30, 60, and 80 for ImageNet32 and ImageNet – it is decayed at epochs 150 and 225 for CIFAR-10. We use a weight decay of 5e-4 for CIFAR-10 and 1e-4 for ImageNet32 and ImageNet. As suggested in [], we do not apply weight decay on 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, and 256 on ImageNet32. However, we made a few adaptations to train our models with PETRA. As suggested by [, ], we employ an accumulation factor $k$ and a batch size of 64, which allows to reduce the effective staleness during training: in this case, $k$ batches of data must be successively processed before updating the parameters of a stage (see Alg. [1](#alg1)). Such gradient accumulation however also increases the effective batch size, and we apply the training recipe used in [] 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 the formula $\texttt{lr}=0.1\frac{64k}{256}$, with $k$ the accumulation factor.

#### Model adaptations.

For designing our RevNet architectures, we adopt a methodology similar to []: the number of channels in each stage is multiplied by 2 to account for the second data stream according to Fig. [2](#S2.F2). However, as the stage function $\tilde{\mathcal{F}}_{j}$ operates only on one of the two streams, the number of parameters stays almost the same between a 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 statistics are not updated during the forward pass.

#### Performance comparison.

Tab. [2](#S4.T2) reports our numerical accuracy on several vision datasets, comparing a backpropagation performance from an official PyTorch implementation of ResNets (the numbers can be found as v1 of [https://pytorch.org/hub/pytorch_vision_resnet/](https://pytorch.org/hub/pytorch_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 accuracy 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-10 accuracies are averaged over 3 runs, with a variance 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 a small drop in accuracy on ImageNet of about 0.6%: using different downsampling layers removes this gap at the expense of a substantial increase in the parameter count ($30.4$M to $50$M). However, we decided not to include this result for the sake of comparison with respect to the original ResNets.

**Table 2: Classification accuracies using our PETRA method with RevNets, compared to standard backpropagation on ResNets and RevNets on CIFAR-10, 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 (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 |

Figure: Figure 4: Validation accuracy of PETRA and backpropagation for a various number of accumulation steps, for a RevNet18 trained on ImageNet with $k\in\{1,2,4,8,16,32\}$. The validation accuracies are averaged over the last 10 epochs. As the number of accumulation steps increases, the effective staleness in PETRA decreases, closing the gap with standard backpropagation.
Refer to caption: /html/2406.02052/assets/x4.png

#### Impact of the accumulation k 𝑘 k .

We test the impact of the accumulation on a RevNet-18 trained via PETRA for various values of accumulations with $k$ spanning $\{1,2,4,8,16,32\}$ on the ImageNet dataset. Fig. [4](#S4.F4) indicates that our method can benefit from large accumulation factors, with the well-known trade-off of large batches mentioned in []. 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 on PyTorch [], although we require significant modifications to the Autograd 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 a reconstruction step, followed by a forward and a backward step. This is because we only need the output gradient as well as the computational graph of $\tilde{\mathcal{F}}_{j}$ to compute the input and parameter gradients at line 12 and 13 of Alg. [1](#alg1), which can be obtained during the input reconstruction phase. For non-reversible stages, we reconstruct the computational graph with a forward pass on the input retrieved from the buffer during the backward pass. Our models can run on a single A100, 80GB.

**Table 3: Memory savings for RevNet50 on ImageNet with our method for different configurations. We indicate the use of memory buffers for inputs or parameters. The savings are 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 | Memory (GB) | Saving (%) |  |
| --- | --- | --- | --- |
| Input | Params. |  |  |
| $\surd$ | $\surd$ | 44.5 | 0.0 |
| $\surd$ | $\times$ | 43.6 | 2.0 |
| $\times$ | $\surd$ | 21.2 | 52.3 |
| $\times$ | $\times$ | 20.3 | 54.3 |

#### Memory benefits and training time.

To better understand the advantage of our method compared to other delayed gradient approaches [, , ], we emphasize the practical memory savings associated with different methods in Tab. [3](#S4.T3). We estimate the memory needed in gigabytes, as the sum of the model size, the input buffer size, and the parameter buffer size, while excluding the input buffer size of the first stage, which corresponds to retrievable dataset inputs.
We do not include the effect of gradient accumulation since it depends on the value of $k$ and only affects the length of the parameter buffer, which is small in our case, i.e., we use $k=1$. Note that the batch size also affects the memory savings, and we set it to 64 for consistency with Tab. [2](#S4.T2). Storing both inputs and parameters into a buffer corresponds to the PipeDream approach []. Only storing inputs into buffers would correspond to the approach in [, ]. The third and fourth lines are only applicable to reversible architectures as they do not store the input into buffers. As can be seen, the input buffer has the biggest impact on the total memory needed, being responsible for 52.3% of the memory footprint. Dropping the parameter buffer in PETRA pushes the memory savings further to 54.3% for a RevNet50 on ImageNet. Note that non-reversible stages account for the majority of total memory use, meaning that savings would be much higher for fully invertible architectures.

## 5 Conclusion

In this work, we introduce PETRA, a novel model parallel training technique for reversible architectures which is a novel promising alternative to 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 operate without any parameter or activation buffers, effectively decoupling the forward and backward phases. Despite using an approximate delayed gradient estimate, our method delivers competitive performances compared to standard backpropagation on standard computer vision datasets.

In future work, we aim to implement and optimize PETRA for Large Language Models (LLMs), with a first baseline being Reformers [], invertible transformers that have been shown to scale. This will validate PETRA’s effectiveness and robustness, solidifying its potential as a cutting-edge training technique.

## Acknowledgments and Disclosure of Funding

This work was supported by Project ANR-21-CE23-0030 ADONIS, EMERG-ADONIS from Alliance SU, and Sorbonne Center for Artificial Intelligence (SCAI) of Sorbonne University (IDEX SUPER 11-IDEX-0004). This work was granted access to the AI resources of IDRIS under the allocations 2023-A0151014526 made by GENCI. We thank polymathic-ai for the relevant discussions.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
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.

I find the paper very interesting and am ready to increase my grade should my remarks be addressed by the authors.
soundness: 3
strengths: - The paper is well written, clear, and has helpful illustrations.
- The algorithm seems simple, natural and intuitive.
- While the algorithm relies on reversible layers, it can still be mixed with standard non-reversible layers, for which a standard backpropagation is performed.
- The authors validate their algorithm with thorough experiments and analyses.
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.
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.
2. 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.
3. Lines 509-510: I think you meant RevNet instead of ResNet.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.

Is 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?
soundness: 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.
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.
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.
Increased 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
Scalability 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: N/A
soundness: 3
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.
- The paper has shown real computation and memory savings.
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.
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.
- Algorithm 1 is a little hard to process.
- 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.

---

confidence: 5
contribution: 3
initial_rating: 8
presentation: 4
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?

Perhaps you could add some information about this in the appendix :-)
soundness: 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.
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.
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.

Meta-review: 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.

Decision: Accept (Spotlight)