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

We thank the reviewer for his enthusiastic review and would like to address his questions. ### **How did you partition the architectures for your experiments? How many layers/blocks in each stage?** We partitioned the architectures at the residual block level, as standardly done for RevNets to obtain a fine-grained partitioning. The first convolutional block and the final projection each count as a separate stage; each residual block then also counts as a 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?** We did not fully investigate idling time caused by the non-uniform distribution of FLOPS across stages. In this paper, we wanted to focus on the optimization properties of PETRA and the effect of crucial parameters. Therefore, we left the topic of fully optimizing resource efficiency for future work, which requires additional development efforts for automation. We also found it difficult to obtain a reasonable balance by hand, each configuration being hardware configuration specific. Ideally, we’d like every worker to have the same flops to fully exploit the parallelization potential. We plan to investigate this load balancing problem on transformers which maintains constant dimensionality and FLOPs across depth in future work. ### **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?** In order to properly answer this question, we need to separate compute from memory. - As the size of the feature maps is decreasing with respect to layer index, early layers consume more memory than the top ones. - The ResNet architectures are designed to approximately keep a constant number of FLOPs across layers. The non-reversible layers have more FLOPs than the other ones due to the extra downsampling operation, and the FLOPs of the different downsampling operations are not equal. The residual blocks that do not employ downsampling keep approximately the same number of flops, which is around 28 Giga-FLOPS for a ResNet-50.

We thank the reviewer for his effort in this reviewing process, and would like to comment on the weaknesses reported. ### **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.** As the reviewer correctly noticed, the delay $\tau$ is the main source of numerical instabilities. Therefore, we chose to explore the largest setting up to a RevNet50, by partitioning the model at the residual block level. A coarser partitioning would decrease the delay and make the problem easier in our case. To address the concern of the reviewer, we designed a simple fully-reversible convolutional architecture by stacking reversible residual blocks. We train such architectures with PETRA while increasing the depth of the network and splitting the architecture at the residual block level, thus effectively increasing the delay during training. We will report the result in the rebuttal if they arrive in the following day, and include an appendix with the tables and plots suggested by the reviewer. ### **Algorithm 1 is a little hard to process.** We apologize for this. We deliberately chose a precise formulation of the algorithm to avoid any misinterpretation of the training 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.** We did investigate the effect of large batch sizes on ImageNet training. The available training recipes do not scale to effective batch sizes larger than 2048 without performance degradation; deterioration starts at 4096 in our experiments. Since LLMs are trained with significantly larger batch sizes, gradient accumulation would not be an issue in this scenario. It is indeed often used to handle significant model sizes when the maximum supported batch size on a given hardware configuration is too restricted. We however note that the training recipes allowing us to scale synchronous SGD to large batch sizes are also effective with PETRA according to our experiments.

We thank the reviewer for his dedicated time and feedback in this review process. We would like to comment on the reported weaknesses first: - **Dependency on reversible architectures:** the authors acknowledge that many existing architectures may not easily be made fully revertible. However, we believe that the favorable memory consumption of PETRA should motivate further research in developing fully reversible architectures. - **Increased Communication Overhead:** The communication overhead stems from the transition from a non-reversible to a reversible counterpart in our paper. While this is a valid concern when dealing with very large models, we want to emphasize that the increase is by a fixed factor, and does not depend on the network architecture. Therefore, we do not believe this would be a significant limiting factor of PETRA. - **Scalability Constraints with Non-Reversible Layers:** While this is true that non-reversible layers induced a memory overhead in PETRA, all other model parallel training techniques in the literature providing acceleration do require activation buffers and suffer from the same memory constraint. To the best of our knowledge, PETRA offers the best compromise in terms of memory consumption and acceleration. The authors will now attempt to answer each of the reviewer’s question: ### **How PETRA perform on large model and more complex task, such as pretraining language model?** The implementation used to compute the accuracy figures reported in the experiment section used a simulated environment to assess the numerical stability of PETRA, and did not allow us to perform efficient distributed training. Only recently did we succeed in producing an efficient distributed implementation of PETRA, from which we were able to derive runtime estimates to assess speedups for RevNets. While LLMs are not yet fully integrated into our framework for complete benchmarking, we have promising developments towards this use case with LLama2 on OpenWebText. Still, LLM pre-training requires substantial engineering to be performed efficiently. Thus, we consider that such large-scale studies are out of the scope of this paper. With the resources at our disposal, we focused our experience toward an academic-friendly setup as we personally think that the optimization properties must be studied in detail as we scale up the complexity of the task. [Edit]: One related concern, that another reviewer has raised explicitly, is the extent to which the depth, or more precisely the number of stages since it is related to the induced delay, affects the scalability potential of PETRA. To this end, we have currently performed experiments in Appendix B to ablate the effect of depth on gradient quality compared to backpropagation, and we have included it in the appendix of our current paper. ### **Is the reversible architecture necessary for PETRA?** PETRA does necessitate a reversible architecture. However, reversible architectures can be derived from many existing architectures, with similar optimization properties as shown in the literature; see [1, 2]. PETRA aims to democratize the development of reversible architectures by providing acceleration with limited memory overhead when scaling model size. ### **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?** The presented paper proposes an implementation of PETRA able to handle non-reversible layers by using an activation buffer between the forward and backward pass. This induces a memory overhead similar to any other delayed gradient approach in the literature, which is quantified theoretically in the storage column of table 1. Resorting to buffers is necessary for our experiments which use the reversible counterpart of ResNets that include non shape-preserving layers which cannot be made invertible canonically as presented in the paper. We tried to specifically deal with the memory issue of non-reversible layers by deferring the buffer to the CPU memory. However, we did not succeed in obtaining an efficient implementation of this mechanism, and left this aspect for further optimization. This would require advanced profiling tools that we are trying to integrate within our code base. As a workaround, we also experimented with quantization as a way to decrease buffer size with promising results, but did not fully investigate the impact of such numerical approximation. We will add an experiment in the appendix quantifying the impact of buffer quantization to a lower precision on final model performances and report corresponding memory usage. [1] - Jacobsen, J. H., Smeulders, A., & Oyallon, E. (2018). i-revnet: Deep invertible networks. ICLR 2018. [2] - Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). Reformer: The efficient transformer. ICLR 2020.

We would like to thank the reviewer once again for their suggestion and take this opportunity to address a few points. ## Invertible networks are currently not very used. This limits the direct applications of the algorithm. Invertible networks are currently not widely used but they have been shown to be extendable to all major architectures without loss of accuracy. For example such as [1], [2], [3], and [4]. Each of these studies highlights that invertibility is an intrinsic property of the model that does not degrade performance. On the contrary, it enables reversibility through simple adaptations of the model, often resulting in improvements, computational savings, or enhanced interpretability. ## The experiments are only performed on RevNet models for image classification. The experiments in our work are conducted solely on RevNet models. We understand the reviewers' concerns; this paper serves primarily as a proof of concept, demonstrating that PETRA represents a promising research direction. In Appendix B, we have included additional experiments analyzing gradient approximation quality, emphasizing the potential for this work to extend beyond its initial scope, which we discuss below. ## How is the approximated gradient influenced by the depth of the model? We have proposed a detailed analysis of this in Appendix B. It demonstrates that, when comparing PETRA during training with standard backpropagation and delayed gradients (as used in [5]), PETRA produces consistent gradients. Notably, deeper layers exhibit gradients that are closer to the true end-to-end gradients, whereas delays introduce discrepancies. However, our experiments in Section 4 show that these discrepancies do not negatively impact training. **References** [1] - Kitaev, N., Kaiser, Ł., & Levskaya, A. (2020). Reformer: The efficient transformer. ICLR 2020. [2] - Gomez, A. N., Ren, M., Urtasun, R., & Grosse, R. B. (2017). The reversible residual network: Backpropagation without storing activations. NeurIPS 2017. [3] - Behrmann, J., Grathwohl, W., Chen, R. T., Duvenaud, D., & Jacobsen, J. H. (2019, May). Invertible residual networks. ICML 2019. [4] - Jacobsen, J. H., Smeulders, A., & Oyallon, E. (2018). i-revnet: Deep invertible networks. ICLR 2018. [5] Zhuang, H., Wang, Y., Liu, Q., Zhang, S., & Lin, Z. (2021). Fully Decoupled Neural Network Learning Using Delayed Gradients. IEEE transactions on neural networks and learning systems.

I thank the authors for their rebuttal, which addressed my concerns. I am updating my score.