Title: PETRA: Parallel End-to-end Training with Reversible Architectures
Venue: NeurIPS 2024
Accepted: no
Keywords: Model parallelism

delayed gradient

reversible architecture

memory reduction
Abstract: We propose a new training approach for reversible architectures that enhances computational efficiency by enabling parallel gradient computation across layers, named PETRA (Parallel End-to-end Training of Reversible Architecture). This method diverges from conventional back-propagation by employing an approximate inversion of activations that effectively preserves gradient quality. By reducing the reliance on synchronous operations, our approach achieves high parallelization with only a slight increase in communication overhead. We have tested this method on benchmark datasets including CIFAR-10, ImageNet-32, and ImageNet, and multiple revertible architectures, where PETRA achieved competitive performance with minimal accuracy loss compared to traditional non-revertible training methods. Our method offers a reduced memory footprint compared to delayed gradient or checkpointing techniques. Unlike pipelining strategies, it eliminates the occurrence of bubble effects, enhancing operational efficiency, while being more parallelizable.

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 CIFAR10, 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 [25]. As datasets and

models continue to grow [1] there is an urgent need for memory-efficient and scalable parallelization 17

of deep learning training across multiple workers. Data parallelism via mini-batches [25] has been 18

widely adopted in deep learning frameworks [26]. This approach computes gradients across model 19

replicas distributed among workers, yet it requires frequent synchronization to aggregate gradients, 20

leading to high communication costs, as well as substantial memory redundancy. Furthermore, with 21

the increasing size and scale of models exceeding that of the growth of on-device memory, the 22

forward and backward passes now often exceed a single device's memory capacity [35]. To further 23

address these issues, methods have attempted to mitigate this memory overhead and to parallelize 24

the sequential backpropagation steps themselves across devices, while computing exact gradients. 25

Techniques like optimizer sharding [34], tensor parallelism [36], activation checkpointing [6], or 26

27

pipelining [15], have been deployed individually or combined, leading for instance to the development

of 3D parallelism [37], a popular methodology which improves the efficiency of the backpropagation 28

implementation. On the other hand, the fundamental inefficiency underlying the parallelization of 29

backpropagation has not been addressed by these methods. 30

However, the use of exact gradient restricts algorithmic choices and parallel implementations, as 31

highlighted by [20]. For instance, backpropagation is backward locked : the inputs of each layer 32

must be propagated through the network and preserved until an error signal is retropropagated to the 33

layer of origin. This requirement enforces a synchronous dependency among subsequent layers and 34

requires them to systematically store intermediary activations, potentially impeding overall resource 35

efficiency as workers must wait for each other to continue their computations and release memory 36

used for activations. To unlock the potential of backpropagation, inexact backpropagation procedures 37

a) Backpropagation

b) PETRA

Time

Stage 2

Stage 1

Stage 0

Stage 2

Stage 1

Stage 0

Propagation

of one batch

Layer

Loss

have been proposed. These procedures are generally conceptualized within the context of model 38 parallelism, where a neural network is split into stages that can process their activations in parallel, 39 potentially on multiple devices. For example, some methods use outdated parameters or activations, 40 such as double-buffered pipelining [14] or delayed gradient approaches [44]. However, these methods 41 introduce significant memory overhead due to the use of ad hoc buffers for activations, parameters, 42 or both. Following an opposite direction, local learning methods [33, 4], which estimate inexact 43 gradients via a local auxiliary neural network, pave the way to parallel gradient computations but 44 often lead to unrecoverable performance drops [11]. This underscores the need for a robust alternative 45 to backpropagation, with limited memory overhead. 46

In this work, we introduce PETRA (Parallel End-to-End Training with Reversible Architectures), 47 a novel method designed to parallelize gradient computations within reversible architectures with 48 minimal computational overhead. Reversible architectures are an ideal candidate for this task, 49 as they can significantly reduce memory overhead during standard backpropagation with limited 50 communication costs. Furthermore, reversibility is a minor requirement, as many studies have 51 demonstrated that standard architectures can be adapted into reversible ones without any performance 52 drops [12, 19, 29, 22]. By allowing parameters to evolve in parallel and by computing an approximate 53 inversion during backward, we propose an effective alternative to backpropagation which allows 54 high model parallelism with a constant communication overhead and no additional buffers . In fact, 55 for a constant increase in communication overhead, PETRA achieves a linear speedup compared to 56 standard backpropagation with respect to the number J of stages the network is split into. We illustrate 57 our approach in Fig. 1, by contrasting the evolution of PETRA with a standard backpropagation pass. 58

59

60

61

62

63

64

65

66

67

68

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 we make available to the research community.

69

2 Related work

Reversible architectures. Reversible DNNs are composed of layers that are invertible, meaning 70 that the input of a layer can be computed from its output. This approach allows to avoid the need to 71 store intermediary activations during the forward pass by reconstructing them progressively during 72 the backward pass [12], at the cost of an extra computation per layer. Invertible networks further 73 improve this method by removing dimensionality reduction steps such as downsamplings, making 74

the networks fully invertible [18]. Reversibility is not restricted to a type of architecture or tasks 75 and has been extensively used for generative models [9], for ResNets [12], and Transformers [29]. 76 However, as far as we know, reversible architectures have never been used to enhance parallelization 77 capabilities. 78

79

80

81

82

83

84

85

86

Alternatives to backpropagation. Multiple alternatives to backpropagation have been proposed previously to improve over its computational efficiency. For instance, DNI [20] 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 [33, 3] propose to use layerwise losses to decouple the training of layers, allowing them to train in parallel [5]. Local learning in videos [28] 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 [11, 38].

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

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 [15], a batch of data is divided into micro-batches that are processed independently at each stage. Although more efficient pipelining schedules have been proposed [10], 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 [31] can allow to get rid of idleness bubbles, but at the expense of introducing staleness in the gradients used. [32] 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 [6] with pipelining [21, 27], 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, [16] 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 [42] and [43], stale gradients are computed from older parameter versions differing from the parameters used during the update. This staleness can be mitigated: [43] 'shrinks' the gradient by the delay value, but more advanced techniques also exist [41, 23]. 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) [40], PipeMare [41] and [23], 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.

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

where F j is the j -th stage parameterized by θ 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 L is then deduced from the corresponding output x J . Parameter gradients are then computed during the backward pass by taking

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

<latexi sh

1_b

64="Z7SKX3

YfO

Ck mp9oM/

TD

>A

c

P

E

gR

2

L

r

FI

N

B

+

q

H

j

w

n0

Q

W

8

y

J

d

V

z

v

5

G

u

U

<latexi sh

1_b

64="

K

/I

cT

U

P

2zq+

wuFM

>A

CVX

Z

DN

Q

EH

j

J

R

W0r

d

k

f

o

5

v

g

S

Y

G

p

O

9

7y

n

m

B

3

8

L

<latexi sh

1_b

64="P

FuRgBWj

JV

Z5v

E8

0

zmw

>A

CTX

c

DN

dL

+2

U

QOK

n

H

G9

y

p

Y

r

/

fo

q

k

S

M

7

3

I

<latexi sh

1_b

64="

BD

Y

AfVF

kQSRv9I

jo

NqXZ

>

CU3

c

T

E

0

2

p

d

Ly

5

O

WM

8H

P

K

r

7

w

G

/

m

g

u

z

Jn

+

<latexi sh

1_b

64="rw

P

K

U

I

f

EN

V

ZoD

Qk

>A

C

X

c

B

9

2

+Ly0

R

M

T

j

O

J

zd

/

u

q

g

n

7

8

v

3

5

p

F

m

W

H

G

Y

S

<latexi sh

1_b

64="FCG

n

W

9pU2D

8u

3Q

qz

7

>A

YX

cZVBN

M

EH

L

OPR

k

S

I

5

T

yv

o

/

m+

w

f

r

g

d

0

j

K

J

<latexi sh

1_b

64="P

GfVzywNX

r

M

J

oLHu j0k

>A

CU3

cZ

B

T9

EF

2

Ip

d

R5

O

W

q

Y

8

v

K

7

S

/

D

m

g

Q

n

+

<latexi sh

1_b

64="Gu omRqA

FW+

XkUn

QB0

d

>

C

cZV

N

9

E

2

r

Ly

K

M

T

j

O

J

P

z

/

w

g

7

8

v

3

5

p

I

H

f

Y

D

S

<latexi sh

1_b

64="N

PFO0Ky

B

8n

EA

VZ

CXM

>

S3

c

T9

2

QG

v

w

U

R

Yk

d

Lj mg

W7I

u

/+

f

o

qHr

z

5

D

J

p

<latexi sh

1_b

64="

BD

Y

AfVF

kQSRv9I

jo

NqXZ

>

CU3

c

T

E

0

2

p

d

Ly

5

O

WM

8H

P

K

r

7

w

G

/

m

g

u

z

Jn

+

<latexi sh

1_b

64="rw

P

K

U

I

f

EN

V

ZoD

Qk

>A

C

X

c

B

9

2

+Ly0

R

M

T

j

O

J

zd

/

u

q

g

n

7

8

v

3

5

p

F

m

W

H

G

Y

S

<latexi sh

1_b

64="FCG

n

W

9pU2D

8u

3Q

qz

7

>A

YX

cZVBN

M

EH

L

OPR

k

S

I

5

T

yv

o

/

m+

w

f

r

g

d

0

j

K

J

<latexi sh

1_b

64="P

GfVzywNX

r

M

J

oLHu j0k

>A

CU3

cZ

B

T9

EF

2

Ip

d

R5

O

W

q

Y

8

v

K

7

S

/

D

m

g

Q

n

+

<latexi sh

1_b

64="G

LQoZX

2+T

J

P

I

>A

CS

c

VB

w

EHU

YDOj8

y

k

M

p

9

N

3

g

7

r

W

0q

v

Fu

n

/

f

m

K

<latexi sh

1_b

64="Gu omRqA

FW+

XkUn

QB0

d

>

C

cZV

N

9

E

2

r

Ly

K

M

T

j

O

J

P

z

/

w

g

7

8

v

3

5

p

I

H

f

Y

D

S

+

x

1

j

-

1

x

2

j

-

1

x

1

j

x

2

j

(b) RevNet Forward

(c) RevNet Reverse

˜

F

j

x

j

-

1

x

j

(a) ResNet

x

1

j

-

1

x

2

j

-

1

x

1

j

x

2

j

˜

F

j

-

F

j

advantage of the chain rule: starting from the last stage with δ J = ∇ x J L , the gradients with regard to the activations are given by

and the gradients with regard to the parameters are defined as

Note that these computations follow a synchronous and sequential order. The parameters θ j can then be updated given their gradient estimate ∆ j , using any optimizer.

3.2 Reversible architectures

We focus on the reversible neural networks presented in [12], 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 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. 2b using an ad-hoc operator ˜ F j . It can be reconstructed by reverse propagating the output according to Fig. 2c, by subtracting the output of ˜ 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 -1 j the reverse stage function, which reconstructs the input from the output. We recursively apply the reconstruction to the final activation x J , such that

Note that reconstructing the input in our procedure is computationally equivalent to recomputing the 142 activations in activation checkpointing, meaning it is equivalent to a single forward pass. Thus, this 143 augmented backward procedure is equivalent to one regular forward call and backward call. However, 144 one should observe that since the input x j -1 must be sent to the reversible stages, this doubles the 145 cost of backward communications. 146

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

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

3.3 A parallelizable approach: PETRA

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 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 ˜ x t j -1 for reversible stages or with a forward pass as in activation checkpointing otherwise. Then, the parameter gradient approximation ∆ 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,

Note that this complete set of equations effectively decouples communications, computations, and 168 parameter updates between independent devices. Indeed, reversible stages are able to operate without 169 maintaining any state between the forward and corresponding backward phase by simply avoiding 170 weight stashing, similarly to [40], and by reversing the output into the input during the backward 171 phase, removing the need for an input buffer. As parameters are updated between the forward and 172 backward phases, the reversible stage produces an approximate input reconstruction, thus evaluating 173 gradients with an approximate set of inputs and parameters during the backward phase. We illustrate 174 in Fig. 3 the mechanism of PETRA compared to standard delayed gradient approaches that rely on 175 additional buffers [44, 42]. 176

177

178

179

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 distributed across J devices uniformly. First, we consider the backpropagation setting, assuming a model

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

Loss

Buffer

size 2

Stage 0

Stage 1

Stage 2

Buffer

size 4

Buffer

size 6

Loss

Forward connection

(activations)

Backward connection

(reconstructed activations)

Backward connection

(gradients)

Gradient computation

a) Delayed

Gradients

b) PETRA

parallelism strategy: a standard backpropagation pass requires storing locally both the parameters and the computational graph and due to the update lock of backpropagation [20], requires synchronization between subsequent layers which impede the speed of computations. Standard Delayed Gradients strategies as implemented in [44, 42] 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 [40], 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 [17] and [30]). To avoid storing activations, we rely on reversible architectures [12] 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.

4 Numerical experiments

4.1 Classification accuracy

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

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 [39], which potentially leads to higher accuracy. We perform a warm-up of 5 epochs where the learning rate linearly increases from 0 to 0 . 1 , following [13]. 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 [13], 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 [42, 43], 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). Such gradient accumulation however also increases the effective batch size, and we apply the training recipe used in [13] 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 lr = 0 . 1 64 k 256 , with k the accumulation factor.

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

Algorithm 1 Worker perspective for training in parallel with PETRA, on a stage j , assuming initialized parameters θ j and time step t , as well as an accumulation factor k > 1 .

Model adaptations. For designing our RevNet architectures, we adopt a methodology similar to [12]: 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 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 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/ ), 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

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

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 Mto 50 M). However, we decided not to include this result for the sake of comparison with respect to the original ResNets.

Impact of the accumulation 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 indicates that our method can benefit from large accumulation factors, with the well-known trade-off of large batches mentioned in [13]. 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 [2], 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 ˜ F j to compute the input and parameter gradients at line 12 and 13 of Alg. 1, 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.

Memory benefits and training time. To better understand the advantage of our method compared to other delayed gradient approaches [14, 40, 23], 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 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. Storing both inputs and parameters into a buffer corresponds to the PipeDream approach [14]. Only storing inputs into

0

5

10

15

20

25

30

Accumulation steps

60.0

62.5

65.0

67.5

70.0

72.5

75.0

77.5

80.0

Validation accuracy

Backprop

PETRA

×

×

buffers would correspond to the approach in [40, 23]. The third and fourth lines are only applicable 265 to reversible architectures as they do not store the input into buffers. As can be seen, the input buffer 266 has the biggest impact on the total memory needed, being responsible for 52.3% of the memory 267 footprint. Dropping the parameter buffer in PETRA pushes the memory savings further to 54.3% for 268 a RevNet50 on ImageNet. Note that non-reversible stages account for the majority of total memory 269 use, meaning that savings would be much higher for fully invertible architectures. 270

271

272

273

274

275

276

277

278

279

280

281

282

283

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 [22], 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.

284

References

I. M. Alabdulmohsin, B. Neyshabur, and X. Zhai. Revisiting neural scaling laws in language 285 and vision. Advances in Neural Information Processing Systems , 35:22300-22312, 2022. 286

287

[2]

J. Ansel, E. Yang, H. He, N. Gimelshein, A. Jain, M. Voznesensky, B. Bao, P. Bell, D. Berard,

Burovski, G. Chauhan, A. Chourdia, W. Constable, A. Desmaison, Z. DeVito, E. Ellison, 288

Feng, J. Gong, M. Gschwind, B. Hirsh, S. Huang, K. Kalambarkar, L. Kirsch, M. Lazos, 289

Lezcano, Y. Liang, J. Liang, Y. Lu, C. K. Luk, B. Maher, Y. Pan, C. Puhrsch, M. Reso, 290

Saroufim, M. Y. Siraichi, H. Suk, S. Zhang, M. Suo, P. Tillet, X. Zhao, E. Wang, K. Zhou, 291

Zou, X. Wang, A. Mathews, W. Wen, G. Chanan, P. Wu, and S. Chintala. Pytorch 2: Faster 292

293

machine learning through dynamic python bytecode transformation and graph compilation. In

Proceedings of the 29th ACM International Conference on Architectural Support for Program294

295

ming Languages and Operating Systems, Volume 2

, ASPLOS '24, page 929-947, New York,

NY, USA, 2024. Association for Computing Machinery. 296

E. Belilovsky, M. Eickenberg, and E. Oyallon. Greedy layerwise learning can scale to imagenet. 297 In International conference on machine learning , pages 583-593. PMLR, 2019. 298

E. Belilovsky, M. Eickenberg, and E. Oyallon. Decoupled greedy learning of cnns. In Interna299 tional Conference on Machine Learning , pages 736-745. PMLR, 2020. 300

301

302

303

E. Belilovsky, L. Leconte, L. Caccia, M. Eickenberg, and E. Oyallon. Decoupled greedy learning of cnns for synchronous and asynchronous distributed learning. arXiv preprint arXiv:2106.06401 , 2021.

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost, 2016.

P. Chrabaszcz, I. Loshchilov, and F. Hutter. A downsampled variant of imagenet as an alternative to the cifar datasets, 2017.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

L. Dinh, D. Krueger, and Y. Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8516 , 2014.

S. Fan, Y. Rong, C. Meng, Z. Cao, S. Wang, Z. Zheng, C. Wu, G. Long, J. Yang, L. 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 , pages 431-445, 2021.

L. Fournier, S. Rivaud, E. Belilovsky, M. Eickenberg, and E. Oyallon. Can forward gradient match backpropagation? In Fortieth International Conference on Machine Learning , 2023.

A. N. Gomez, M. Ren, R. Urtasun, and R. B. Grosse. The reversible residual network: Backpropagation without storing activations. Advances in neural information processing systems , 30, 2017.

P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677 , 2017.

A. Harlap, D. Narayanan, A. Phanishayee, V. Seshadri, N. Devanur, G. Ganger, and P. Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03377 , 2018.

Y. Huang, Y. Cheng, A. Bapna, O. Firat, D. Chen, M. Chen, H. Lee, J. Ngiam, Q. V. Le, Y. Wu, 328 et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in 329 neural information processing systems , 32, 2019. 330

Z. Huo, B. Gu, and H. Huang. Training neural networks using features replay. Advances in 331 Neural Information Processing Systems , 31, 2018. 332

333

334

Z. Huo, B. Gu, H. Huang, et al. Decoupled parallel backpropagation with convergence guarantee. In International Conference on Machine Learning , pages 2098-2106. PMLR, 2018.

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

J.-H. Jacobsen, A. Smeulders, and E. Oyallon. i-revnet: Deep invertible networks. arXiv preprint arXiv:1802.07088 , 2018.

J.-H. Jacobsen, A. W. M. Smeulders, and E. Oyallon. i-revnet: Deep invertible networks. ArXiv , abs/1802.07088, 2018.

M. Jaderberg, W. M. Czarnecki, S. Osindero, O. Vinyals, A. Graves, D. Silver, and K. Kavukcuoglu. Decoupled neural interfaces using synthetic gradients. In International conference on machine learning , pages 1627-1635. PMLR, 2017.

C. Kim, H. Lee, M. Jeong, W. Baek, B. Yoon, I. Kim, S. Lim, and S. Kim. torchgpipe: On-the-fly pipeline parallelism for training giant models, 2020.

N. Kitaev, Ł. Kaiser, and A. Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020.

A. Kosson, V. Chiley, A. Venigalla, J. Hestness, and U. Koster. Pipelined backpropagation at scale: training large models without batches. Proceedings of Machine Learning and Systems , 3:479-501, 2021.

A. Krizhevsky. Learning multiple layers of features from tiny images. 2009.

Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. nature , 521(7553):436-444, 2015.

S. Li, Y. Zhao, R. Varma, O. Salpekar, P. Noordhuis, T. Li, A. Paszke, J. Smith, B. Vaughan, P. Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006.15704 , 2020.

Y. Liu, S. Li, J. Fang, Y. Shao, B. Yao, and Y. You. Colossal-auto: Unified automation of parallelization and activation checkpoint for large-scale models, 2023.

M. Malinowski, D. Vytiniotis, G. Swirszcz, V. Patraucean, and J. Carreira. Gradient forwardpropagation for large-scale temporal video modelling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 9249-9259, 2021.

K. Mangalam, H. Fan, Y. Li, C.-Y. Wu, B. Xiong, C. Feichtenhofer, and J. Malik. Reversible vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10830-10840, 2022.

E. Mizutani and S. Dreyfus. On complexity analysis of supervised mlp-learning for algorithmic comparisons. In IJCNN'01. International Joint Conference on Neural Networks. Proceedings (Cat. No.01CH37222) , volume 1, pages 347-352 vol.1, 2001.

D. Narayanan, A. Harlap, A. Phanishayee, V. Seshadri, N. R. Devanur, G. R. Ganger, P. B. Gibbons, and M. Zaharia. Pipedream: Generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles , pages 1-15, 2019.

D. Narayanan, A. Phanishayee, K. Shi, X. Chen, and M. Zaharia. Memory-efficient pipelineparallel dnn training. In International Conference on Machine Learning , pages 7937-7947. PMLR, 2021.

A. Nøkland and L. H. Eidnes. Training neural networks with local error signals. In International conference on machine learning , pages 4839-4850. PMLR, 2019.

S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. Zero: Memory optimizations toward training trillion parameter models, 2020.

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

J. Ren, S. Rajbhandari, R. Y. Aminabadi, O. Ruwase, S. Yang, M. Zhang, D. Li, and Y. He. { Zero-offload } : Democratizing { billion-scale } model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21) , pages 551-564, 2021.

M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 , 2019.

S. Smith, M. Patwary, B. Norick, P. LeGresley, S. Rajbhandari, J. Casper, Z. Liu, S. Prabhumoye, G. Zerveas, V. Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990 , 2022.

Y. Wang, Z. Ni, S. Song, L. Yang, and G. Huang. Revisiting locally supervised learning: an alternative to end-to-end training. arXiv preprint arXiv:2101.10832 , 2021.

R. Wightman, H. Touvron, and H. Jégou. Resnet strikes back: An improved training procedure in timm, 2021.

A. Xu, Z. Huo, and H. Huang. On the acceleration of deep learning model parallelism with staleness. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) , pages 2085-2094, 2019.

B. Yang, J. Zhang, J. Li, C. Ré, C. Aberger, and C. De Sa. Pipemare: Asynchronous pipeline parallel dnn training. Proceedings of Machine Learning and Systems , 3:269-296, 2021.

H. Zhuang, Z. Lin, and K.-A. Toh. Accumulated decoupled learning: Mitigating gradient staleness in inter-layer model parallelization. arXiv preprint arXiv:2012.03747 , 2020.

H. Zhuang, Y. Wang, Q. Liu, and Z. Lin. Fully decoupled neural network learning using delayed gradients. IEEE transactions on neural networks and learning systems , 33(10):6013-6020, 2021.

H. Zhuang, Z. Weng, F. Luo, T. Kar-Ann, H. Li, and Z. Lin. Accumulated decoupled learn398 ing with gradient staleness mitigation for convolutional neural networks. In International 399 Conference on Machine Learning , pages 12935-12944. PMLR, 2021. 400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our novel method is presented in Sec. 3 as described. We provide numerical accuracy and ablations in Sec. 4.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss it along with our experimental results in Sec. 4.2.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: We do not provide theoretical results in this paper.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The experimental results are detailed in the article, and can be reproduced with the provided code.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The data is openly accessible and properly referenced in the text. The code is included in the supplementary material.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The hyperparameters are detailed in the article and the code.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report accuracy averaged over 3 runs on CIFAR10 which have low variance, and we follow to the standard practice on ImageNet and ImageNet32.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

483

484

485

486

487

488

489

490

Justification: Details about resources and executions are presented Sec 4.2.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conforms to the Code of Ethics.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: Our article is concerned with allowing general deep learning optimization, and does not have particular societal impacts.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not present such risks.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: The datasets used are cited, and we credit the author of the code used for our implementation.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The code will be provided with a README file describing how to run experiments.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: No crowdsourcing or human subjects.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

491

Justification: No potential risks or study participants.

492

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn pipelined model training, one important issue is to reduce the bubble sizes. \nOne stream of work is to use the staleness, where the weight discrepancy is mitigated using stashed weights.\nThis work tries to reduce the overhead of storing weights with reversible architectures. \nUsing the non-stashed updated weights, but with restored inputs to each stage, \napproximated gradients are obtained and parallel training is performed.\nThis leads to less memory usage on training at the cost of increased communication. \nTraining results on resnet variants seem to maintain accuracy.\n\nSTRENGTHS:\n- This is a nice adaptation of reversible architectures to pipelined training. If it works well, there is a potential for becoming a new popular pipelined training method.\n\n- The idea of using reconstructed input instead of stored weights seem to be novel.\n\nWEAKNESSES:\n- Insufficient experiment size: Only compared on three different sizes of resnet. This is far from sufficient, especially with the largest model being resnet50. \n\n- No comparison on speedup: speedup on the training time is crucial, but the \"memory benefits and training time\" section does not disclose any data. Since the proposed scheme has larger communication, it is crucial to report the number.\n\n- Classification accuracy drop: The final accuracy drops on all three datasets for resnet50. 0.6%p and 0.7%p are huge drops for those models. Given that this is the largest model among the tested ones, it draws a significant concern on whether this technique would work for larger models such as resnet152 or ViTs.\n\n- There is no analysis or proof on why the proposed scheme would work. Why it is a good approximation, or why it is going to converge, etc.\n\nQUESTIONS:\nThe questions essentially come from the weaknesses.\n\n1. The training time seems to be missing. Could the authors fill it in?\n\n2. How would it work for larger models?\n\n3. Is there a guarantee or an analysis on the convergence or the approximation error?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a method that combines reversible neural networks and parallel distributed training to enable learning with minimal memory usage, while incurring only slight communication and computation overhead. In this approach, the need for storing intermediate activations in traditional backpropagation is eliminated, thus reducing memory constraints and allowing for higher parallelism on the same device. This new method facilitates efficient learning by providing an innovative solution.\n\nSTRENGTHS:\nThe problem setup involving reversible architecture and distributed parallel training is intriguing. High memory consumption is a critical issue in learning, and reversible architecture has been proposed to address this problem. It is anticipated that these advantages can be similarly applied to distributed parallel training. Additionally, the paper is very well-written, making the ideas easy to understand. The figures and tables were also judged to be of high quality and well-prepared.\n\nWEAKNESSES:\nThe main drawback of this paper is the insufficient experimentation. Although using reversible architecture in distributed training is a novel concept, it appears to be merely a combination of existing ideas. For this paper to have a significant impact, it must demonstrate the advantages and benefits of the proposed idea in an actual distributed learning environment. However, the experiments were conducted using only a single A100 GPU, and there is no demonstration of the performance improvements or limitations of the proposed idea in a real distributed environment. The values presented in the tables do not clearly differentiate from what can be achieved with existing reversible architectures. To improve the completeness of this paper, it is essential to analyze scenarios that necessitate the use of multiple GPUs, such as video applications, large-resolution diffusion, and large language models. The current data fails to effectively explain the benefits of the proposed idea.\n\nQUESTIONS:\n1. Additional data is needed to determine the advantages in an actual distributed environment and to assess how much the increase in communication and computation affects speed.\n\n2. Additional experiments are essential to understand how the required minimum number of GPUs changes for tasks that require a lot of memory, or how the tasks that can be run on the same cluster are affected.\n\n3. As mentioned in the conclusion, experiments with large language models are essential.\n\n4. Detailed experiments are also needed to see how these experiments impact accuracy and to analyze the effects on complex tasks and hyperscale models.\n\n5. Where does the accuracy difference between RevNet backprop and PETRA in Table 2 arise? Besides the subtle errors that can occur in precision, are there any other potential sources of error?\n\nLIMITATIONS:\nNot relevant.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the author proposes a new alternative algorithm (Parallel End-to-End Training with Reversible Architectures) for regular backpropagation, which significantly enhance the parallelization with a limited overhead compared to regular backpropagation and other alternatives to end-to-end training.\nSpecifically, the network is split into several stages (one layer or a set of layers) distributed across distinct devices, one batch data is split into several mini-batch data. The first device sequentially accesses the mini-batch data and pass them forward to the next stage until the final stage is reached. The backpropagation is initialized from the final stage to the first stage. It enables a significant parallelization of forward and backward computations across multiple devices.\n\nSTRENGTHS:\n* This paper is well-organized and easy to follow.\n\n* The background information is very rich and makes it easy for someone who is not familiar with this field to understand the relevant techniques including the technique proposed by this paper.\n\n* The figures about the core technique proposed by authors are very clear, which can help readers understand the technique at a glance.\n\n* The paper evaluates the proposed techniques on multiple datasets and networks.\n\nWEAKNESSES:\n* From the comparison between the proposed method and other techniques from related work, it showcases that the proposed method does not have an overall crushing lead. There exists the method which can achieve higher speed and less time than proposed method with storage increased. \n\n* The low or even zero storage on proposed method is mainly due to reversible architectures. Maybe authors can extend proposed parallel training method to some non-reversible architectures (need memory storage for intermediate activations), then compare with other SOTA methods.\n\n* It would be great if authors use more distributed devices to get more stages from a network, in this case, the performance of the proposed method is likely to be deeply explored. Because the proposed technique is aimed to deployed on the distributed devices.\n\nQUESTIONS:\nPlease refer to the weakness section.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose fusing delayed gradient pipeline parallelism with reversible models in order to capture the benefits of the former while mitigating the drawbacks with the latter.\n\nSTRENGTHS:\n- The paper sets up a pretty compelling combination of ideas. This is a great example of a paper that clearly understands the strengths and weaknesses of two disparate techniques and fits them together like puzzle pieces.\n\n- The paper is clear and methodical in laying out the motivation for the approach. By the time the method is introduced, its seems like the natural and obvious choice. This is good writing.\n\n- The concept is solid. I really *want* to like this idea, since it seems to fit together so well.\n\nWEAKNESSES:\n- While the idea is presented fairly clearly, a lot of the analysis is estimates (S4.2) and generalizations (Tab 1). It's fine for motivating the idea, but not really good enough for proving it works as projected. I'm left wondering how much of this method will actually translate to a scaled-up implementation. (No question that it *was* implemented, but a pipeline-parallel model that doesn't actually pipeline across devices is...not particularly compelling.)\n\n- The paper is a fusion of two ideas, designed to capture the computational performance benefits of pipeline parallelism while using reversible models to mitigate memory scaling. Some estimated results of memory footprint are presented in Table 3. No measured results are presented related to parallelism (timing, utilization, etc.). From this paper, it is not possible to determine whether it has succeeded. This is confused further by the section 4.2: \"Memory benefits and training time\" which does not discuss training time at all. The lack of computational results is fairly damning.\n\nQUESTIONS:\nQ1. Most distributed parallel solutions emphasize minimizing interconnect traffic, as it often bottlenecks overall runtime. Table 1 suggests PETRA would increase inter-pipeline-stage traffic by 2-4x (forward/backward). This would suggest that a scaled-up implementation might struggle under increased communication. As the presumed goal of PETRA is to scale up (L64, L115), how do the authors mitigate the increased communication?\n\n(minor) The last column in Table 1 is not described in text. Is this just highlighting that these two are pipelined?\n\nLIMITATIONS:\nAs described in weaknesses. Limitations, like computational performance details, are not well described.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}