Original: An In-depth Study of Stochastic Backpropagation
Revised: An In-depth Study of Stochastic Backpropagation

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: In this paper, we provide an in-depth study of Stochastic Backpropagation (SBP)when training deep neural networks for standard image classification and objectdetection tasks.
Revised: In this paper, we provide an in-depth study of Stochastic Backpropagation (SBP)when training deep neural networks for standard image classification and objectdetection tasks.

Original: During backward propagation, SBP calculates the gradients byonly using a subset of the feature maps to save memory and the computation cost.
Revised: During backward propagation, SBP calculates the gradients byonly using a subset of the feature maps to save memory and the computation cost.

Original: We interpret SBP as an efficient way to implement stochastic gradient decent byperforming backpropagation dropout, which leads to considerable memory savingand training process speedup, with a minimal impact on the overall model accuracy.
Revised: We interpret SBP as an efficient way to implement stochastic gradient decent byperforming backpropagation dropout, which leads to considerable memory savingand training process speedup, with a minimal impact on the overall model accuracy.

Original: We offer a good practice to apply SBP in training deep image classificationmodels, which can be adopted in learning a wide range of deep neural networks.
Revised: We offer a good practice to apply SBP in training deep image classificationmodels, which can be adopted in learning a wide range of deep neural networks.

Original: Experiments on image classification and object detection show that SBP can saveup to 40% of GPU memory with less than 1% accuracy degradation
Revised: Experiments on image classification and object detection show that SBP can saveup to 40% of GPU memory with less than 1% accuracy degradation

Original:  Introduction
Revised:  Introduction

Original: A common practice to improve accuracy when training deep networks is to increase the inputresolution [20], model depth [8], and width [23].
Revised: A common practice to improve accuracy when training deep networks is to increase the inputresolution [20], model depth [8], and width [23].

Original: This can significantly increase the GPU memoryusage, making training on devices with limited resources difficult.
Revised: This can significantly increase the GPU memoryusage, making training on devices with limited resources difficult.

Original: For example, when training anobject detector using ConvNeXt-Base [14] as the backbone with batch size 2, over 17 GB of memoryis required, which will not be able to fit on most modern GPUs.
Revised: For example, when training anobject detector using ConvNeXt-Base [14] as the backbone with batch size 2, over 17 GB of memoryis required, which will not be able to fit on most modern GPUs.

Original: Many prior studies have focused on designing memory-efficient training method, especially forlarge model training on high resolution images or videos.
Revised: Many prior studies have focused on designing memory-efficient training method, especially forlarge model training on high resolution images or videos.

Original: Mixed-precision training [17] uses lowerprecision to represent the network weights and activations for certain layers of the network.
Revised: Mixed-precision training [17] uses lowerprecision to represent the network weights and activations for certain layers of the network.

Original: Gradientcheckpointing [3] restores intermediate feature maps by recomputing nodes in the computationgraph during the backward pass, which requires extra computations.
Revised: Gradientcheckpointing [3] restores intermediate feature maps by recomputing nodes in the computationgraph during the backward pass, which requires extra computations.

Original: Gradient accumulation splits amini-batch into smaller chunks to calculate gradients iteratively, which slows down training process.
Revised: Gradient accumulation splits amini-batch into smaller chunks to calculate gradients iteratively, which slows down training process.

Original: A recently developed method, Stochastic Backpropagation (SBP) [4], offers a new angle on memoryefficient training.
Revised: A recently developed method, Stochastic Backpropagation (SBP) [4], offers a new angle on memoryefficient training.

Original: It performs partial execution of gradients during backpropagation on selectedframes when training video models [1, 19].
Revised: It performs partial execution of gradients during backpropagation on selectedframes when training video models [1, 19].

Original: It achieves considerable memory saving by only cachinga subset of frame feature maps without significant accuracy loss.
Revised: It achieves considerable memory saving by only cachinga subset of frame feature maps without significant accuracy loss.

Original: However, the prior work [4] onlyapplies SBP to video-based models of two architectures and does not thoroughly explore the designspace of the method.
Revised: However, the prior work [4] onlyapplies SBP to video-based models of two architectures and does not thoroughly explore the designspace of the method.

Original: The explanation provided in the paper for the effectiveness of SBP is simple butnarrow – it attributes the validity of the approximation to the high redundancy of video.
Revised: The explanation provided in the paper for the effectiveness of SBP is simple butnarrow – it attributes the validity of the approximation to the high redundancy of video.

Original: This leaves many questions around SBP unanswered.
Revised: This leaves many questions around SBP unanswered.

Original: Here we ask the questions: (1) Is SBP’seffectiveness limited to operating on the redundancy of video frames or can it effectively leveragespatial redundancy?
Revised: Here we ask the questions: (1) Is SBP’seffectiveness limited to operating on the redundancy of video frames or can it effectively leveragespatial redundancy?

Original: (2) If so, can SBP be generalized to image-based tasks, with more limitedredundancy?
Revised: (2) If so, can SBP be generalized to image-based tasks, with more limitedredundancy?

Original: (3) Is SBP generalizable to different networks (e.g., MLP, CNN, Transformer) andoperators?
Revised: (3) Is SBP generalizable to different networks (e.g., MLP, CNN, Transformer) andoperators?

Original: (4) What are the key design choices when implementing SBP?
Revised: (4) What are the key design choices when implementing SBP?

Original: To answer these questions, we first formalize the SBP process in a general form and analyze theeffect of SBP on the calculation of gradients.
Revised: To answer these questions, we first formalize the SBP process in a general form and analyze theeffect of SBP on the calculation of gradients.

Original: With this formulation, we show how to generalize the
Revised: With this formulation, we show how to generalize the

Original: SBP idea to image models and present an in-depth study of the design choices available when using
Revised: SBP idea to image models and present an in-depth study of the design choices available when using

Original: SBP.
Revised: SBP.

Original: During backward propagation, SBP calculates gradients from a subset of spatial feature mapsto save memory and computation cost.
Revised: During backward propagation, SBP calculates gradients from a subset of spatial feature mapsto save memory and computation cost.

Original: We observe that the calculated gradients by SBP are highlycorrelated with the standard SGD gradients, indicating SBP is a reasonable approximation.
Revised: We observe that the calculated gradients by SBP are highlycorrelated with the standard SGD gradients, indicating SBP is a reasonable approximation.

Original: We alsoprovide a new, straightforward implementation of the idea that requires only a few lines of code.
Revised: We alsoprovide a new, straightforward implementation of the idea that requires only a few lines of code.

Original: Weexplore several important design strategies which have a noticeable impact on the performance.
Revised: Weexplore several important design strategies which have a noticeable impact on the performance.

Original: Wevalidate the generalizability of SBP on two common tasks (image classificaton and detection) withtwo popular network architectures: ViT [7] and ConvNeXt [14].
Revised: Wevalidate the generalizability of SBP on two common tasks (image classificaton and detection) withtwo popular network architectures: ViT [7] and ConvNeXt [14].

Original: We show that SBP can effectivelytrain image tasks with less than 1% loss of accuracy in a memory-efficient fashion.
Revised: We show that SBP can effectivelytrain image tasks with less than 1% loss of accuracy in a memory-efficient fashion.

Original:  Related Work
Revised:  Related Work

Original: Since deep neural networks are highly extended in both depth [8] and breadth [23] in recent years,designing memory efficient training methods has attracted a fair amount of attention.
Revised: Since deep neural networks are highly extended in both depth [8] and breadth [23] in recent years,designing memory efficient training methods has attracted a fair amount of attention.

Original: By saving andrecomputing nodes in the computation graph during backpropagation, Gradient Checkpoint [3] cansave a large amount of memory (up to about 50%).
Revised: By saving andrecomputing nodes in the computation graph during backpropagation, Gradient Checkpoint [3] cansave a large amount of memory (up to about 50%).

Original: Besides, gradient accumulation splits the largesample batch into several mini batches and runs them iteratively without updating model parameters.
Revised: Besides, gradient accumulation splits the largesample batch into several mini batches and runs them iteratively without updating model parameters.

Original: This saves the GPU memory proportionally without affecting the accuracy, because accumulatingthe gradients of these sub-iteration and then updating the model parameters is identical to directlyoptimizing the model global batch size.
Revised: This saves the GPU memory proportionally without affecting the accuracy, because accumulatingthe gradients of these sub-iteration and then updating the model parameters is identical to directlyoptimizing the model global batch size.

Original: However, these two methods slow down the training process.
Revised: However, these two methods slow down the training process.

Original: Multigrid [22] proposes to use different mini-batch shapes to speedup the model training.
Revised: Multigrid [22] proposes to use different mini-batch shapes to speedup the model training.

Original: Sparse
Revised: Sparse

Original: Network [6] particularly targets the recognition tasks but can only save the memory theoretically.
Revised: Network [6] particularly targets the recognition tasks but can only save the memory theoretically.

Original: Sideways [15] and its follow-up [16] reduce the memory cost by overwriting activations whenevernew ones become available, but they are limited to causal models.
Revised: Sideways [15] and its follow-up [16] reduce the memory cost by overwriting activations whenevernew ones become available, but they are limited to causal models.

Original: SBP [4] reduces a large portion of
Revised: SBP [4] reduces a large portion of

Original: GPU memory by only backpropagating gradients from incomplete execution.
Revised: GPU memory by only backpropagating gradients from incomplete execution.

Original:  Understanding Stochastic Backpropagation (SBP)
Revised:  Understanding Stochastic Backpropagation (SBP)

Original: In this section, we first formalize the process of Stochastic Backpropagation (SBP) (sec. 3.1) andshow how to compute the gradient for two common operations (sec. 3.2).
Revised: In this section, we first formalize the process of Stochastic Backpropagation (SBP) (sec. 3.1) andshow how to compute the gradient for two common operations (sec. 3.2).

Original: With these formulationsestablished, we analyze how to perform the chain rule of gradient calculations for SBP (sec. 3.3).
Revised: With these formulationsestablished, we analyze how to perform the chain rule of gradient calculations for SBP (sec. 3.3).

Original: Finally, we present a simple implementation of SBP in practice with only a few lines of code (sec. 3.4).
Revised: Finally, we present a simple implementation of SBP in practice with only a few lines of code (sec. 3.4).

Original: 3.1 Formulation of SBP
Revised: 3.1 Formulation of SBP

Original: In this section, we present a general formulation of SBP [4].
Revised: In this section, we present a general formulation of SBP [4].

Original: We denote a loss function L (Θ) tolearn a model on a dataset D
Revised: We denote a loss function L (Θ) tolearn a model on a dataset D

Original: = { s i } Ni =1 , where Θ is the collection (vector, matrix, or tensor) of allparameters of the model, N is the total number of training samples, s i ∈ R T × H × W × C is an inputdata sample, where H and W are spatial resolution height and width respectively, C is the input 70 channel size, T = 1 represents an image and T > 1 represents a video.
Revised: = { s i } Ni =1 , where Θ is the collection (vector, matrix, or tensor) of allparameters of the model, N is the total number of training samples, s i ∈ R T × H × W × C is an inputdata sample, where H and W are spatial resolution height and width respectively, C is the input 70 channel size, T = 1 represents an image and T > 1 represents a video.

Original: In each training iteration,the model with SBP processes a mini-batch of data samples S ∈ R B × T × H × W × C ⊆ D to computethe loss in the forward pass, where B is the batch size.
Revised: In each training iteration,the model with SBP processes a mini-batch of data samples S ∈ R B × T × H × W × C ⊆ D to computethe loss in the forward pass, where B is the batch size.

Original: During the backward pass, different fromthe traditional backpropagation that relies on the full feature maps X i ∈ R B × T i × H
Revised: During the backward pass, different fromthe traditional backpropagation that relies on the full feature maps X i ∈ R B × T i × H

Original: i × W i ×
Revised: i × W i ×

Original: C i ( i is theindex of the layers), an incomplete execution for backpropagation (i.e., SBP) only utilizes a subset ofthe feature maps X
Revised: C i ( i is theindex of the layers), an incomplete execution for backpropagation (i.e., SBP) only utilizes a subset ofthe feature maps X

Original: subi
Revised: subi

Original: [Equation]
Revised: [Equation]

Original: subi
Revised: subi

Original: [Equation]
Revised: [Equation]

Original: H subi × W subi × C i to approximate the gradient calculation.
Revised: H subi × W subi × C i to approximate the gradient calculation.

Original: That is, itupdates the model weights as
Revised: That is, itupdates the model weights as

Original: [Equation]
Revised: [Equation]

Original: where η is the learning rate.
Revised: where η is the learning rate.

Original: SBP saves memory by only caching a subset of feature maps in thebackward pass, but it keeps all the full forward paths.
Revised: SBP saves memory by only caching a subset of feature maps in thebackward pass, but it keeps all the full forward paths.

Original: Note that the method in [4] is a special case of(1) as it performs on the frame level, i.e., using a subset of video frames X T subi i ∈ R T subi × H
Revised: Note that the method in [4] is a special case of(1) as it performs on the frame level, i.e., using a subset of video frames X T subi i ∈ R T subi × H

Original: i × W i × C i .
Revised: i × W i × C i .

Original: 3.2 Backward Gradient Calculation
Revised: 3.2 Backward Gradient Calculation

Original: Next, we derive equations to compute the gradient during the backwards phase of SBP.
Revised: Next, we derive equations to compute the gradient during the backwards phase of SBP.

Original: We denotethat at layer
Revised: We denotethat at layer

Original: i of the model, an operator f i with learnable weights
Revised: i of the model, an operator f i with learnable weights

Original: W f i to processes the feature maps
Revised: W f i to processes the feature maps

Original: X i ∈ R B × T i × H
Revised: X i ∈ R B × T i × H

Original: i × W i × C ini , where C ini refers to the input channel size.
Revised: i × W i × C ini , where C ini refers to the input channel size.

Original: 3.2.1 Calculation of Linear Layers
Revised: 3.2.1 Calculation of Linear Layers

Original: We start with exploring f i as a point-wise convolutional layer (i.e., PW-Conv).
Revised: We start with exploring f i as a point-wise convolutional layer (i.e., PW-Conv).

Original: PW-Conv is widely 86 used in the state-of-the-art computer vision models such as ViT [7], ConvNeXt [14], and Video
Revised: PW-Conv is widely 86 used in the state-of-the-art computer vision models such as ViT [7], ConvNeXt [14], and Video

Original: Swin [13].
Revised: Swin [13].

Original: It is equivalent to the linear layers in Transformer MLP blocks [14], and is also known as 88 channel-wise fully connected layer.
Revised: It is equivalent to the linear layers in Transformer MLP blocks [14], and is also known as 88 channel-wise fully connected layer.

Original: It can be either implemented through the linear layer, or the 1 × 1convolution.
Revised: It can be either implemented through the linear layer, or the 1 × 1convolution.

Original: It is mainly used to enrich the channel information with weights W f i ∈ R C ini × C outi , andthe analysis on PW-Conv can be easily extended to other layers.
Revised: It is mainly used to enrich the channel information with weights W f i ∈ R C ini × C outi , andthe analysis on PW-Conv can be easily extended to other layers.

Original: The forward pass of this layer is (the 91 bias term is omitted):
Revised: The forward pass of this layer is (the 91 bias term is omitted):

Original: [Equation]
Revised: [Equation]

Original: 2) where X i +1 ∈ R B × T
Revised: 2) where X i +1 ∈ R B × T

Original: i × H
Revised: i × H

Original: i
Revised: i

Original: × W i × C outi .
Revised: × W i × C outi .

Original: During the traditional backward pass with full backpropagation, the gradients of both activations and weights are calculated by:
Revised: During the traditional backward pass with full backpropagation, the gradients of both activations and weights are calculated by:

Original: [Equation]
Revised: [Equation]

Original: It requires storing the weights W f i ∈ R C ini × C outi and the full feature maps X i ∈ R B × T
Revised: It requires storing the weights W f i ∈ R C ini × C outi and the full feature maps X i ∈ R B × T

Original: i × H
Revised: i × H

Original: i
Revised: i

Original: × W i × C inito compute the gradients.
Revised: × W i × C inito compute the gradients.

Original: In many cases, B × T i × H
Revised: In many cases, B × T i × H

Original: i × W i
Revised: i × W i

Original: ≫ C outi , i.e., caching of feature mapsdominates the memory usage.
Revised: ≫ C outi , i.e., caching of feature mapsdominates the memory usage.

Original: In order to effectively save GPU memory during training, SBP onlycaches a subset of the feature map for backpropagation.
Revised: In order to effectively save GPU memory during training, SBP onlycaches a subset of the feature map for backpropagation.

Original: In the backward pass of SBP case, we split the full feature maps X i into two subsets X keepi and
Revised: In the backward pass of SBP case, we split the full feature maps X i into two subsets X keepi and

Original: X dropi , which denote the corresponding indices of the sampled feature maps that are kept anddropped, respectively.
Revised: X dropi , which denote the corresponding indices of the sampled feature maps that are kept anddropped, respectively.

Original: An example of X keepi can be X i [: , even index on T, : , : , :] for video, or X i [: , :, even index on H, odd index on W, :] for image (e.g., Fig.
Revised: An example of X keepi can be X i [: , even index on T, : , : , :] for video, or X i [: , :, even index on H, odd index on W, :] for image (e.g., Fig.

Original: 3 left), and X dropi is the complementarypart of the feature maps.
Revised: 3 left), and X dropi is the complementarypart of the feature maps.

Original: We have X keepi ∩ X dropi = ∅ and X keepi ∪ X dropi =
Revised: We have X keepi ∩ X dropi = ∅ and X keepi ∪ X dropi =

Original: X i .
Revised: X i .

Original: We denotekeep-ratio r as the number of gradients kept indices j ∈ Z keepi over the number of all indices.
Revised: We denotekeep-ratio r as the number of gradients kept indices j ∈ Z keepi over the number of all indices.

Original: One advantage of simplifying the SBP problem by using PW-Conv operator is that the nodes in 105 spatial or temporal dimensions can be calculated independently through the forward and backwardpasses.
Revised: One advantage of simplifying the SBP problem by using PW-Conv operator is that the nodes in 105 spatial or temporal dimensions can be calculated independently through the forward and backwardpasses.

Original: Therefore, we can update the forward pass separately on the kept and dropped indices:
Revised: Therefore, we can update the forward pass separately on the kept and dropped indices:

Original: [Equation]
Revised: [Equation]

Original: as well as the backward:
Revised: as well as the backward:

Original: [Equation]
Revised: [Equation]

Original: SBP drops all the gradients calculation with superscripts drop and keeps all the gradients calculationwith superscripts keep , thus it saves memory and computation on the dependencies of droppedindices.
Revised: SBP drops all the gradients calculation with superscripts drop and keeps all the gradients calculationwith superscripts keep , thus it saves memory and computation on the dependencies of droppedindices.

Original: In other words, the memory usage and the backward computation cost is proportional tothe keep-ratio of the original case.
Revised: In other words, the memory usage and the backward computation cost is proportional tothe keep-ratio of the original case.

Original: Mathematically, it sets dX dropi +1 = 0 (thus dX dropi = 0 ) and usesdW keepf i =
Revised: Mathematically, it sets dX dropi +1 = 0 (thus dX dropi = 0 ) and usesdW keepf i =

Original: X keepiT dX keepi +1 as a stochastic approximation of the original weight gradients dW f i toupdate the model weights.
Revised: X keepiT dX keepi +1 as a stochastic approximation of the original weight gradients dW f i toupdate the model weights.

Original: 3.2.2 Calculation of Convolutional Layers
Revised: 3.2.2 Calculation of Convolutional Layers

Original: We then derive the calculation of a general convolutional layer (ConvLayer) as f i .
Revised: We then derive the calculation of a general convolutional layer (ConvLayer) as f i .

Original: Unlike the PW- 116 Conv, nodes of general ConvLayer in spatial or temporal dimensions are no longer independent inforward and backward calculations.
Revised: Unlike the PW- 116 Conv, nodes of general ConvLayer in spatial or temporal dimensions are no longer independent inforward and backward calculations.

Original: Assume that we apply a convolutional kernel k × k with stride s .
Revised: Assume that we apply a convolutional kernel k × k with stride s .

Original: The regular back-propagation of the ConvLayer is calculated as follows:
Revised: The regular back-propagation of the ConvLayer is calculated as follows:

Original: [Equation]
Revised: [Equation]

Original: 6) where ∗ is the convolution operation, zeropad is to pad the edge of the matrix with zeros, and W ′ f irepresents the 180 ◦ rotation of W f i .
Revised: 6) where ∗ is the convolution operation, zeropad is to pad the edge of the matrix with zeros, and W ′ f irepresents the 180 ◦ rotation of W f i .

Original: In the gradient calculation of ConvLayer at a given location j on the input feature map as x ji ∈ X
Revised: In the gradient calculation of ConvLayer at a given location j on the input feature map as x ji ∈ X

Original: i ,all the elements from dX i +1 inside the region of the kernel size k will contribute to the gradientcalculation of element x ji .
Revised: i ,all the elements from dX i +1 inside the region of the kernel size k will contribute to the gradientcalculation of element x ji .

Original: Therefore, in SBP of ConvLayer, when we set dX dropi +1 = 0 for certainareas, in general, two things will happen: (a) even if j ∈ Z drop and dx ji +1 = 0 , dx ji will not be 0 dueto the contribution of the neighbors of j on dX i +1 ; (b) dx ji will be an approximated gradient unlessall dx pi +1 with p in the neighbor areas of j in terms of the kernel size k are kept and there is no chainrule effect on dx pi +1 .
Revised: Therefore, in SBP of ConvLayer, when we set dX dropi +1 = 0 for certainareas, in general, two things will happen: (a) even if j ∈ Z drop and dx ji +1 = 0 , dx ji will not be 0 dueto the contribution of the neighbors of j on dX i +1 ; (b) dx ji will be an approximated gradient unlessall dx pi +1 with p in the neighbor areas of j in terms of the kernel size k are kept and there is no chainrule effect on dx pi +1 .

Original: However, there is a special scenario: when s ≥ k , the neighbor effects disappeardue to the fact that dX i +1 needs to be zero interweaved in gradient calculation.
Revised: However, there is a special scenario: when s ≥ k , the neighbor effects disappeardue to the fact that dX i +1 needs to be zero interweaved in gradient calculation.

Original: As a result, nodes inspatial or temporal dimensions are getting more independent, which is similar to that of PW-Conv.
Revised: As a result, nodes inspatial or temporal dimensions are getting more independent, which is similar to that of PW-Conv.

Original: 3.2.3 Discussion
Revised: 3.2.3 Discussion

Original: From the derivation above, we see that the SBP approximation performs differently when the operatorchanges.
Revised: From the derivation above, we see that the SBP approximation performs differently when the operatorchanges.

Original: For layers such as linear layer or PW-Conv, the gradients (i.e., dx i ) are either 0 or exact 1 .
Revised: For layers such as linear layer or PW-Conv, the gradients (i.e., dx i ) are either 0 or exact 1 .

Original: While for convolutional layers with kernel size larger than 1, most of the gradients on the activationsare neither 0 nor exact.
Revised: While for convolutional layers with kernel size larger than 1, most of the gradients on the activationsare neither 0 nor exact.

Original: Due to space limitation, we only derived SBP on two most popular operators.
Revised: Due to space limitation, we only derived SBP on two most popular operators.

Original: For derivation of other operators, please refer to the supplemental material.
Revised: For derivation of other operators, please refer to the supplemental material.

Original: 3.3 Chain Rule Effect
Revised: 3.3 Chain Rule Effect

Original: We have derived the calculation of SBP on a single layer in the previous section.
Revised: We have derived the calculation of SBP on a single layer in the previous section.

Original: In this section, wederive the calculation of SBP on a stack of multiple layers.
Revised: In this section, wederive the calculation of SBP on a stack of multiple layers.

Original: For simplicity, we take two layers as anexample, which can be easily generalized to more layers.
Revised: For simplicity, we take two layers as anexample, which can be easily generalized to more layers.

Original: 3.3.1 Two PW-Conv Layers
Revised: 3.3.1 Two PW-Conv Layers

Original: We first consider a two PW-Conv layer scenario.
Revised: We first consider a two PW-Conv layer scenario.

Original: We assume a layer before f i , i.e., f i − 1 , is another 142 PW-Conv layer with weights W f
Revised: We assume a layer before f i , i.e., f i − 1 , is another 142 PW-Conv layer with weights W f

Original: [Equation]
Revised: [Equation]

Original: The forward pass of this layer is (the bias term is also omitted):
Revised: The forward pass of this layer is (the bias term is also omitted):

Original: [Equation]
Revised: [Equation]

Original: Then we can calculate the gradients of weights and activations in layer
Revised: Then we can calculate the gradients of weights and activations in layer

Original: f
Revised: f

Original: i − 1 similar to Eq.
Revised: i − 1 similar to Eq.

Original: [Equation]
Revised: [Equation]

Original: We first consider the case when there is no additional gradient dropout at layer f
Revised: We first consider the case when there is no additional gradient dropout at layer f

Original: i − 1 .
Revised: i − 1 .

Original: Since dX i +1has gradients at kept indices (i.e., dX keepi +1 ) and 0 at dropped indices (i.e., dX dropi +1 ), Eq.
Revised: Since dX i +1has gradients at kept indices (i.e., dX keepi +1 ) and 0 at dropped indices (i.e., dX dropi +1 ), Eq.

Original: (8) can beupdated to:
Revised: (8) can beupdated to:

Original: [Equation]
Revised: [Equation]

Original: From the equation, we can see that if there is no additional gradient dropout applied on layer f
Revised: From the equation, we can see that if there is no additional gradient dropout applied on layer f

Original: i − 1 ,through the chain rule, the keep and drop set will transfer from X i +1 to X i − 1 identically.
Revised: i − 1 ,through the chain rule, the keep and drop set will transfer from X i +1 to X i − 1 identically.

Original: If we apply gradient dropout on layer f i − 1 as well, i.e., on layer
Revised: If we apply gradient dropout on layer f i − 1 as well, i.e., on layer

Original: f i − 1 , we split dX i into dX keepi anddX dropi .
Revised: f i − 1 , we split dX i into dX keepi anddX dropi .

Original: Assume that the keep / drop indices set at layer i + 1 is Z keepi +

Revised: Assume that the keep / drop indices set at layer i + 1 is Z keepi +

Original: and Z dropi +1 , and at layer i is
Revised: and Z dropi +1 , and at layer i is

Original: Z keepi and Z dropi , respectively.
Revised: Z keepi and Z dropi , respectively.

Original: Then through Eq.
Revised: Then through Eq.

Original: (8) and Eq.
Revised: (8) and Eq.

Original: (9), the keep / drop indices set on layeri − 1 will be:
Revised: (9), the keep / drop indices set on layeri − 1 will be:

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (9) becomes:
Revised: (9) becomes:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where keep ′
Revised: where keep ′

Original: =
Revised: =

Original: Z keepi − 1 and drop ′ =
Revised: Z keepi − 1 and drop ′ =

Original: Z dropi − 1 .
Revised: Z dropi − 1 .

Original: # f: an arbitrary operation # grad_keep_idx: sampled indices where gradients are kept # grad_drop_idx: sampled indices where gradients are dropped def sbp_f(f, inputs, grad_keep_idx, grad_drop_idx): # initiate outputs outputs = torch.zeros(output_shape, device=inputs.device)
Revised: # f: an arbitrary operation # grad_keep_idx: sampled indices where gradients are kept # grad_drop_idx: sampled indices where gradients are dropped def sbp_f(f, inputs, grad_keep_idx, grad_drop_idx): # initiate outputs outputs = torch.zeros(output_shape, device=inputs.device)

Original: # forward with gradient calculation, gradients will be calculated with torch.autograd with torch.enable_grad(): outputs[grad_keep_idx] = f(inputs[grad_keep_idx]) # forward without gradient calculation with torch.no_grad(): outputs[grad_drop_idx] = f(inputs[grad_drop_idx]) return outputs
Revised: # forward with gradient calculation, gradients will be calculated with torch.autograd with torch.enable_grad(): outputs[grad_keep_idx] = f(inputs[grad_keep_idx]) # forward without gradient calculation with torch.no_grad(): outputs[grad_drop_idx] = f(inputs[grad_drop_idx]) return outputs

Original: 3.3.2 PW-Conv + ConvLayers
Revised: 3.3.2 PW-Conv + ConvLayers

Original: Now let’s look at the chain-rule effect of the case when f i − 1 is a convolutional layer with kernel size 158 k > 1.
Revised: Now let’s look at the chain-rule effect of the case when f i − 1 is a convolutional layer with kernel size 158 k > 1.

Original: Assume that we apply a convolutional kernel k × k with stride s .
Revised: Assume that we apply a convolutional kernel k × k with stride s .

Original: Using the chain rule, we canupdate the Eq.
Revised: Using the chain rule, we canupdate the Eq.

Original: (6) as follows:
Revised: (6) as follows:

Original: [Equation]
Revised: [Equation]

Original: We can see from the equation that each dx i − 1 receives the backward gradients from the outputs ofits neighbors, which might be dropped out, thus the gradient at dx i − 1 is no longer exact.
Revised: We can see from the equation that each dx i − 1 receives the backward gradients from the outputs ofits neighbors, which might be dropped out, thus the gradient at dx i − 1 is no longer exact.

Original: From the Eq.
Revised: From the Eq.

Original: The dropout at layer i makes the matrix dX i sparse (only indices at X keepi +1 arenon-zero).
Revised: The dropout at layer i makes the matrix dX i sparse (only indices at X keepi +1 arenon-zero).

Original: If we conduct dropout at the current layer f i − 1 (i.e., we have dX keepi and dX dropi ), thematrix dX i will become even more sparse unless the dropout is happening on the same location.
Revised: If we conduct dropout at the current layer f i − 1 (i.e., we have dX keepi and dX dropi ), thematrix dX i will become even more sparse unless the dropout is happening on the same location.

Original: Ifwe stack layers further and further in this way, the chain rule effect may lead to a very sparse outputderivative matrix (e.g., dX i ), in which the nodes are getting more spatial or temporal independentlyin gradient calculation like what is happening in the PW-Conv operator.
Revised: Ifwe stack layers further and further in this way, the chain rule effect may lead to a very sparse outputderivative matrix (e.g., dX i ), in which the nodes are getting more spatial or temporal independentlyin gradient calculation like what is happening in the PW-Conv operator.

Original: However, stacking ConvLayers (due to the space limitation, we skip the math derivation) will nothave such a trend, as the gradient dropout in ConvLayer only creates approximated gradient but notmany zeros, because of the neighbor effects of gradient calculation.
Revised: However, stacking ConvLayers (due to the space limitation, we skip the math derivation) will nothave such a trend, as the gradient dropout in ConvLayer only creates approximated gradient but notmany zeros, because of the neighbor effects of gradient calculation.

Original: 3.3.3 Discussion
Revised: 3.3.3 Discussion

Original: Overall, in terms of the chain rule effect, we can see from Eq.
Revised: Overall, in terms of the chain rule effect, we can see from Eq.

Original: (11) that stacking PW-Conv layers,especially in an interweaving way may lead to a sparse keep set.
Revised: (11) that stacking PW-Conv layers,especially in an interweaving way may lead to a sparse keep set.

Original: However, if we always drop gradientat the same positions across layers, this won’t happen.
Revised: However, if we always drop gradientat the same positions across layers, this won’t happen.

Original: Stacking PW-Conv + ConvLayers (in Eq.
Revised: Stacking PW-Conv + ConvLayers (in Eq.

Original: (12))also tends to make the keep set sparse.
Revised: (12))also tends to make the keep set sparse.

Original: While stacking ConvLayers doesn’t have such effect, thoughit stacks the non-exact gradient at almost every position and tends to make the gradients graduallymove away from the original ones.
Revised: While stacking ConvLayers doesn’t have such effect, thoughit stacks the non-exact gradient at almost every position and tends to make the gradients graduallymove away from the original ones.

Original: 3.4 Efficient Implementation
Revised: 3.4 Efficient Implementation

Original: With a deeper understanding of how SBP works, we provide a simple and efficient implementation ofthe SBP technique in Alg.
Revised: With a deeper understanding of how SBP works, we provide a simple and efficient implementation ofthe SBP technique in Alg.

Original: We point out that this is more efficient than the prior work [4] as there is(1) no need to do the re-forward in the backward, thus it is computationally more efficient, and (2) noneed to manually cache the intermediate feature maps and fill zero gradients in the backward, whichis much simpler.
Revised: We point out that this is more efficient than the prior work [4] as there is(1) no need to do the re-forward in the backward, thus it is computationally more efficient, and (2) noneed to manually cache the intermediate feature maps and fill zero gradients in the backward, whichis much simpler.

Original: We emphasize that it is very simple and easy to implement, it only needs 3 linesof code to do the forward, and the backward (no need to implement manually) can be seamlesslyhandled by the ease of “autograd engine” in the deep learning frameworks.
Revised: We emphasize that it is very simple and easy to implement, it only needs 3 linesof code to do the forward, and the backward (no need to implement manually) can be seamlesslyhandled by the ease of “autograd engine” in the deep learning frameworks.

Original: It supports point-wiseconvolutional layers, the dot-product attention layers, and multilayer perception (MLP) in popularnetworks including ViT [7] and its variants
Revised: It supports point-wiseconvolutional layers, the dot-product attention layers, and multilayer perception (MLP) in popularnetworks including ViT [7] and its variants

Original:  Design Strategies
Revised:  Design Strategies

Original: In this section, we investigate the effectiveness of different design strategies of SBP on image tasks.
Revised: In this section, we investigate the effectiveness of different design strategies of SBP on image tasks.

Original: We use ViT-Tiny [7] model to evaluate the top-1 accuracy on ImageNet [18] validation dataset.
Revised: We use ViT-Tiny [7] model to evaluate the top-1 accuracy on ImageNet [18] validation dataset.

Original: The
Revised: The

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: ViT-Tiny is a smaller version of the ViT [7] variants, it uses patch size 16 with 12 transformer blocks(more than 80 layers), the embedding dimension is 192 and number of heads is 3, i.e., each head hasdimension 64.
Revised: ViT-Tiny is a smaller version of the ViT [7] variants, it uses patch size 16 with 12 transformer blocks(more than 80 layers), the embedding dimension is 192 and number of heads is 3, i.e., each head hasdimension 64.

Original: The experiments are trained for 100 epochs on the ImageNet training dataset.
Revised: The experiments are trained for 100 epochs on the ImageNet training dataset.

Original: Theconclusions remain for larger models and higher training epochs.
Revised: Theconclusions remain for larger models and higher training epochs.

Original: How many layers should be applied with SBP?
Revised: How many layers should be applied with SBP?

Original: To answer this question, we progressively apply
Revised: To answer this question, we progressively apply

Original: SBP to transformer layers from the last layers to early layers and present their results in Fig.
Revised: SBP to transformer layers from the last layers to early layers and present their results in Fig.

Original: 1 Left.
Revised: 1 Left.

Original: As shown, we observe a sweet spot when applying SBP to 8 out of 12 block layers, which saves 33%of GPU memory with an accuracy drop within 1.5%.
Revised: As shown, we observe a sweet spot when applying SBP to 8 out of 12 block layers, which saves 33%of GPU memory with an accuracy drop within 1.5%.

Original: Although applying SBP on all the transformerlayers can save the most GPU memory, it causes a large accuracy drop.
Revised: Although applying SBP on all the transformerlayers can save the most GPU memory, it causes a large accuracy drop.

Original: Therefore, unless statedotherwise, we apply SBP on 2/3 of all layers of ViT models for the rest of this work.
Revised: Therefore, unless statedotherwise, we apply SBP on 2/3 of all layers of ViT models for the rest of this work.

Original: How much gradients should be kept in SBP?
Revised: How much gradients should be kept in SBP?

Original: In this section, we discuss the trade-off between thegradients keep-ratio and accuracy in SBP.
Revised: In this section, we discuss the trade-off between thegradients keep-ratio and accuracy in SBP.

Original: As a reminder, keep-ratio r is the number of gradients keptindices over the number of all indices, and it represents the percentage of spatial feature map valuesused to calculate the gradients.
Revised: As a reminder, keep-ratio r is the number of gradients keptindices over the number of all indices, and it represents the percentage of spatial feature map valuesused to calculate the gradients.

Original: In general, the higher the keep-ratio is, the more gradient informationis preserved, and the higher the accuracy is, but the less memory is saved.
Revised: In general, the higher the keep-ratio is, the more gradient informationis preserved, and the higher the accuracy is, but the less memory is saved.

Original: A natural question arises:what is a good keep-ratio to balance the accuracy drop and memory saving?
Revised: A natural question arises:what is a good keep-ratio to balance the accuracy drop and memory saving?

Original: In Fig.
Revised: In Fig.

Original: 1 Middle, we show the trade-off between model accuracy and the keep-ratio of gradients.
Revised: 1 Middle, we show the trade-off between model accuracy and the keep-ratio of gradients.

Original: Inthis case, we apply the same keep-ratio to all the SBP layers at the same gradients kept indices and wecall it uniform keep-ratios method.
Revised: Inthis case, we apply the same keep-ratio to all the SBP layers at the same gradients kept indices and wecall it uniform keep-ratios method.

Original: We observe that setting the keep-ratio to be 0.5 is a golden rule,and having a keep-ratio of 0.25 is too aggressive, which results in a dramatic accuracy drop (3.04%)of the model.
Revised: We observe that setting the keep-ratio to be 0.5 is a golden rule,and having a keep-ratio of 0.25 is too aggressive, which results in a dramatic accuracy drop (3.04%)of the model.

Original: We also explore different keep-ratio methods by linearly increasing or decreasing thekeep-ratios from early to late layers while keeping their average keep-ratio to be 0.5.
Revised: We also explore different keep-ratio methods by linearly increasing or decreasing thekeep-ratios from early to late layers while keeping their average keep-ratio to be 0.5.

Original: Specifically, thekeep-ratios are [0 . 25 , 0 .
Revised: Specifically, thekeep-ratios are [0 . 25 , 0 .

Original: 75] (increasing) or its inverse (decreasing)on 8 transformer blocks.
Revised: 75] (increasing) or its inverse (decreasing)on 8 transformer blocks.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 1 Right, both models perform significantly worse.
Revised: 1 Right, both models perform significantly worse.

Original: We simulate the gradient calculation on ViT-Tiny with 2000 data batches randomly sampled from the
Revised: We simulate the gradient calculation on ViT-Tiny with 2000 data batches randomly sampled from the

Original: ImageNet training dataset.
Revised: ImageNet training dataset.

Original: We apply different keep-ratio methods on the model with the same fixedweights, and calculate the correlation (measured by cosine similarity) between the weights gradientof applying SBP to the original exact gradient without applying SBP.
Revised: We apply different keep-ratio methods on the model with the same fixedweights, and calculate the correlation (measured by cosine similarity) between the weights gradientof applying SBP to the original exact gradient without applying SBP.

Original: From Fig.
Revised: From Fig.

Original: 2, we observe that theweights gradient has a stronger correlation by using uniform keep-ratios compared to increasing or
Revised: 2, we observe that theweights gradient has a stronger correlation by using uniform keep-ratios compared to increasing or

Original: [Figure]
Revised: [Figure]

Original: decreasing keep-ratios.
Revised: decreasing keep-ratios.

Original: We believe uniform keep-ratios can keep the consistency of gradient droppinglocations between different layers and can preserve gradient information at these spatial locations.
Revised: We believe uniform keep-ratios can keep the consistency of gradient droppinglocations between different layers and can preserve gradient information at these spatial locations.

Original: However, non-uniform keep-ratios will have different dropping locations between different layers.
Revised: However, non-uniform keep-ratios will have different dropping locations between different layers.

Original: As an overall result, it may loss more gradient information on more spatial locations and produce aless accurate estimation of the original gradients and may have a lower accuracy performance.
Revised: As an overall result, it may lose more gradient information on more spatial locations and produce aless accurate estimation of the original gradients and may have a lower accuracy performance.

Original: How to sample the gradient keep mask?
Revised: How to sample the gradient keep mask?

Original: SBP calculates the gradients by only using a subset offeature maps.
Revised: SBP calculates the gradients by only using a subset offeature maps.

Original: A key question is how to sample the subset of feature maps to mask the kept gradients.
Revised: A key question is how to sample the subset of feature maps to mask the kept gradients.

Original: If we fix the spatial locations of the gradient kept indices for every training step, the backward onlypropagates the valid gradients (non-zeros) on these fixed locations, which forces the model to onlylearn from fixed partial spatial information and therefore severely hurts the model performance.
Revised: If we fix the spatial locations of the gradient kept indices for every training step, the backward onlypropagates the valid gradients (non-zeros) on these fixed locations, which forces the model to onlylearn from fixed partial spatial information and therefore severely hurts the model performance.

Original: Abetter strategy is to randomly sample the gradient kept indices in each training step, so that it canstatistically visit every spatial location with equal importance.
Revised: Abetter strategy is to randomly sample the gradient kept indices in each training step, so that it canstatistically visit every spatial location with equal importance.

Original: We examine two mask sampling strategies: grid-wise mask sampling and random mask sampling.
Revised: We examine two mask sampling strategies: grid-wise mask sampling and random mask sampling.

Original: See Fig.
Revised: See Fig.

Original: 3 left as a particular example of keep-ratio 0.5.
Revised: 3 left as a particular example of keep-ratio 0.5.

Original: We report the accuracy in Fig.
Revised: We report the accuracy in Fig.

Original: 4 left.
Revised: 4 left.

Original: Overall,grid-wise mask sampling achieves 0 .
Revised: Overall,grid-wise mask sampling achieves 0 .

Original: 5% ∼ 1 .
Revised: 5% ∼ 1 .

Original: 7% higher accuracy over random mask sampling.
Revised: 7% higher accuracy over random mask sampling.

Original: In
Revised: In

Original: Fig.
Revised: Fig.

Original: 3, A stronger correlation is also observed by using the grid-wise sampling.
Revised: 3, A stronger correlation is also observed by using the grid-wise sampling.

Original: One explanationis that grid-wise sampling is a more structured sampling method and it provides a more accurateestimation of gradients calculation.
Revised: One explanationis that grid-wise sampling is a more structured sampling method and it provides a more accurateestimation of gradients calculation.

Original: Unless stated otherwise, we use grid-wise sampling in this paper.
Revised: Unless stated otherwise, we use grid-wise sampling in this paper.

Original: How to apply SBP on MHSA?
Revised: How to apply SBP on MHSA?

Original: In general, the transformer block consists of two sub-blocks: MultiHead Self-Attention (MHSA) and Multi-layer perception (MLP).
Revised: In general, the transformer block consists of two sub-blocks: MultiHead Self-Attention (MHSA) and Multi-layer perception (MLP).

Original: Theoretically, the memory usage ofapplying SBP on MLP (two linear layers) is proportional to the keep-ratio r .
Revised: Theoretically, the memory usage ofapplying SBP on MLP (two linear layers) is proportional to the keep-ratio r .

Original: However, the memorysaving of applying SBP on MHSA depends on the sampling method on the query, key, value (QKV),and the ratio of head dimension d over the number of tokens n .
Revised: However, the memorysaving of applying SBP on MHSA depends on the sampling method on the query, key, value (QKV),and the ratio of head dimension d over the number of tokens n .

Original: From [4], activation maps neededfor gradients calculation in MHSA include the inputs and outputs ( 2 hdn ), QKV vectors (each withhdn ), the attention weight maps ( 2 hnn before and after the softmax), thus requiring the memory of3 hdn + 2 hnn in total, where h is the number of heads.
Revised: From [4], activation maps neededfor gradients calculation in MHSA include the inputs and outputs ( 2 hdn ), QKV vectors (each withhdn ), the attention weight maps ( 2 hnn before and after the softmax), thus requiring the memory of3 hdn + 2 hnn in total, where h is the number of heads.

Original: Dropping gradients on query only (Drop
Revised: Dropping gradients on query only (Drop

Original: Query) or on all QKV (Drop QKV) enjoys the following memory usage ratio:
Revised: Query) or on all QKV (Drop QKV) enjoys the following memory usage ratio:

Original: Query only:
Revised: Query only:

Original: [Equation]
Revised: [Equation]

Original: In video transformers, the number of tokens n is much larger than the head dimension d (i.e., invideo Swin Transformer, n = 392 and d = 32 , dn = 0 . 082 ) and thus the attention weight maps 2 hnn
Revised: In video transformers, the number of tokens n is much larger than the head dimension d (i.e., invideo Swin Transformer, n = 392 and d = 32 , dn = 0 . 082 ) and thus the attention weight maps 2 hnn

Original: [Figure]
Revised: [Figure]

Original: comprise most of the memory.
Revised: comprise most of the memory.

Original: Therefore, dropping gradients on query only in the video model isgood enough to save memory so that its memory is linearly proportional to the keep ratio r [4].
Revised: Therefore, dropping gradients on query only in the video model isgood enough to save memory so that its memory is linearly proportional to the keep ratio r [4].

Original: However, in image task models such as ViT with a typical input size 224 and patch size 16, n =14 × 14 = 196 and d = 64 , dn = 0 . 327 .
Revised: However, in image task models such as ViT with a typical input size 224 and patch size 16, n =14 × 14 = 196 and d = 64 , dn = 0 . 327 .

Original: With keep-ratio r = 0 .
Revised: With keep-ratio r = 0 .

Original: 5 (or r = 0 . 25 ), dropping gradientson query alone uses 0 .
Revised: 5 (or r = 0 . 25 ), dropping gradientson query alone uses 0 .

Original: 61 × (or 0 .
Revised: 61 × (or 0 .

Original: 42 × ) memory while dropping all QKV reduces its memory usageratio to 0 .
Revised: 42 × ) memory while dropping all QKV reduces its memory usageratio to 0 .

Original: 46 × (or 0 . 22 × ).
Revised: 46 × (or 0 . 22 × ).

Original: From Fig.
Revised: From Fig.

Original: 4, dropping gradients on all QKV also gains better accuracy.
Revised: 4, dropping gradients on all QKV also gains better accuracy.

Original: Wealso report the performance of dropping gradients along the attention heads, but we did not observeaccuracy improvement.
Revised: Wealso report the performance of dropping gradients along the attention heads, but we did not observeaccuracy improvement.

Original: Hence, we apply SBP on all QKV for image-based transformers.
Revised: Hence, we apply SBP on all QKV for image-based transformers.

Original: We plot the cosine similarity of the weights gradients of applying SBP and the original exact gradientsin Fig.
Revised: We plot the cosine similarity of the weights gradients of applying SBP and the original exact gradientsin Fig.

Original: One interesting observation is that Drop QKV has a slightly better correlation to Drop Queryand a much higher correlation to Drop Head.
Revised: One interesting observation is that Drop QKV has a slightly better correlation to Drop Queryand a much higher correlation to Drop Head.

Original: We believe an accurate estimation of the gradientsmight be an important factor to the model performance trained with SBP, which provides insight forfurther investigation.
Revised: We believe an accurate estimation of the gradientsmight be an important factor to the model performance trained with SBP, which provides insight forfurther investigation.

Original:  Generalizability of SBP
Revised:  Generalizability of SBP

Original: We evaluate the generalizability of our proposed SBP on two computer vision benchmarks: imageclassification on ImageNet [18] and object detection on COCO [11].
Revised: We evaluate the generalizability of our proposed SBP on two computer vision benchmarks: imageclassification on ImageNet [18] and object detection on COCO [11].

Original: In order to do a fair comparison ofthe same network, we keep the training hyper-parameters of optimizer, augmentation, regularization,batch size and learning rate the same for a given model, and only adopt different stochastic depthaugmentation [10] for different model sizes.
Revised: In order to do a fair comparison ofthe same network, we keep the training hyper-parameters of optimizer, augmentation, regularization,batch size and learning rate the same for a given model, and only adopt different stochastic depthaugmentation [10] for different model sizes.

Original: The Mixed Precision Training [17] method is enabledfor faster training.
Revised: The Mixed Precision Training [17] method is enabledfor faster training.

Original: All experiments are conducted on 8 × Tesla 16GB V100.
Revised: All experiments are conducted on 8 × Tesla 16GB V100.

Original: 5.1 Classification on ImageNet
Revised: 5.1 Classification on ImageNet

Original: We report the memory and accuracy trade-off in Tab. 1 for two state-of-the-art networks: transformerbased ViT [7] and convolutional-based ConvNeXt [14].
Revised: We report the memory and accuracy trade-off in Tab. 1 for two state-of-the-art networks: transformerbased ViT [7] and convolutional-based ConvNeXt [14].

Original: We train 300 epochs for both these twonetworks.
Revised: We train 300 epochs for both these twonetworks.

Original: We use DeiT’s [21] training recipe for ViT and use the official one of ConvNeXt.
Revised: We use DeiT’s [21] training recipe for ViT and use the official one of ConvNeXt.

Original: For ViT models, we apply SBP on the last 8 transformer blocks including MHSA and MLP layerswith the keep-ratio of 0.5, reducing GPU memory usage by approximately 30%.
Revised: For ViT models, we apply SBP on the last 8 transformer blocks including MHSA and MLP layerswith the keep-ratio of 0.5, reducing GPU memory usage by approximately 30%.

Original: Note that thememory saving ratio of the network is not proportional the the drop-ratio ( 1 − keep-ratio) because weonly apply SBP on 23 of all the layers.
Revised: Note that thememory saving ratio of the network is not proportional the the drop-ratio ( 1 − keep-ratio) because weonly apply SBP on 23 of all the layers.

Original: With keep ratio of 0.5, the SBP technique can still effectivelylearn with an accuracy drop of only 0.04% and 0.43% for ViT-Tiny and ViT-Base, respectively.
Revised: With keep ratio of 0.5, the SBP technique can still effectivelylearn with an accuracy drop of only 0.04% and 0.43% for ViT-Tiny and ViT-Base, respectively.

Original: For ConvNeXt models, it builds the basic block with a depth-wise convolutional (DW-Conv) layerwith kernel size 7 × 7 followed by two point-wise convolutional (PW-Conv) layers 2 and the DW-Convlayer only consumes a small portion of GPU memory (less than 10%).
Revised: For ConvNeXt models, it builds the basic block with a depth-wise convolutional (DW-Conv) layerwith kernel size 7 × 7 followed by two point-wise convolutional (PW-Conv) layers 2 and the DW-Convlayer only consumes a small portion of GPU memory (less than 10%).

Original: It has multi-scale stages whereearly stages has higher redundancy and occupies most of the GPU memory.
Revised: It has multi-scale stages whereearly stages has higher redundancy and occupies most of the GPU memory.

Original: Therefore, we apply
Revised: Therefore, we apply

Original: SBP on the PW-Conv and down-sampling layers in the first three stages, and do not apply it on thelast stage.
Revised: SBP on the PW-Conv and down-sampling layers in the first three stages, and do not apply it on thelast stage.

Original: With a keep ratio of 0.5, ConvNeXt-Tiny and ConvNeXt-Base can efficiently learn themodel with a GPU memory saving of above 40% and an accuracy drop within 0.5%.
Revised: With a keep ratio of 0.5, ConvNeXt-Tiny and ConvNeXt-Base can efficiently learn themodel with a GPU memory saving of above 40% and an accuracy drop within 0.5%.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: 5.2 Object Detection and Segmentation on COCO
Revised: 5.2 Object Detection and Segmentation on COCO

Original: We fine-tune Mask R-CNN [9] and Cascade Mask R-CNN
Revised: We fine-tune Mask R-CNN [9] and Cascade Mask R-CNN

Original: [2] on the COCO dataset [11] with
Revised: [2] on the COCO dataset [11] with

Original: ConvNeXt-Tiny and ConvNeXt-Base backbones pretrained on ImageNet-1K, respectively.
Revised: ConvNeXt-Tiny and ConvNeXt-Base backbones pretrained on ImageNet-1K, respectively.

Original: Followingthe same hyper-parameter settings of [14], we apply SBP on the ConvNeXt backbones to savetraining memory.
Revised: Followingthe same hyper-parameter settings of [14], we apply SBP on the ConvNeXt backbones to savetraining memory.

Original: Results in Tab.
Revised: Results in Tab.

Original: 2 show that SBP with keep-ratio of 0.5 can still learn the detectiontask reasonably well with 0.2% and 0.7% loss on the box and mask average precision (AP) for
Revised: 2 show that SBP with keep-ratio of 0.5 can still learn the detectiontask reasonably well with 0.2% and 0.7% loss on the box and mask average precision (AP) for

Original: ConvNeXt-Base and ConvNeXt-Tiny backbones, respectively.
Revised: ConvNeXt-Base and ConvNeXt-Tiny backbones, respectively.

Original: However, it only consumes about0.7 × of the GPU memory.
Revised: However, it only consumes about0.7 × of the GPU memory.

Original: Note that training detection models on high resolution images (up to 800 × 1333) is memory intensiveeven for a very small batch size, e.g., with backbone of ConvNeXt-Base and batch size of 2, thetraining requires 17.4GB of GPU memory which may lead to out of memory error for commondevices such as Tesla 16GB V100.
Revised: Note that training detection models on high resolution images (up to 800 × 1333) is memory intensiveeven for a very small batch size, e.g., with backbone of ConvNeXt-Base and batch size of 2, thetraining requires 17.4GB of GPU memory which may lead to out of memory error for commondevices such as Tesla 16GB V100.

Original: While with the help of SBP, the model can still be effectivelytrained on GPUs with 12.5GB memory, which is much more feasible for many research groups.
Revised: While with the help of SBP, the model can still be effectivelytrained on GPUs with 12.5GB memory, which is much more feasible for many research groups.

Original: Herewe apply SBP only in the backbone, we believe it has a potential to apply SBP on more other layersin the detection model to achieve more promising results in terms of efficiency.
Revised: Herewe apply SBP only in the backbone, we believe it has a potential to apply SBP on more other layersin the detection model to achieve more promising results in terms of efficiency.

Original:  Discussion
Revised:  Discussion

Original: In this work, we present a comprehensive study of the Stochastic Back-Propagation (SBP) mechanismwith a generalized implementation.
Revised: In this work, we present a comprehensive study of the Stochastic Back-Propagation (SBP) mechanismwith a generalized implementation.

Original: We analyze the effect of different design strategies to optimizethe trade-off between accuracy and memory.
Revised: We analyze the effect of different design strategies to optimizethe trade-off between accuracy and memory.

Original: We show that our approach can reduce 40% of the GPUmemory when training image classification models under various deep learning backbones.
Revised: We show that our approach can reduce 40% of the GPUmemory when training image classification models under various deep learning backbones.

Original: Limitations:
Revised: Limitations:

Original: In general, SBP is a memory efficient training method, but as byproduct, it stillcauses slight loss of accuracy.
Revised: In general, SBP is a memory efficient training method, but as byproduct, it stillcauses slight loss of accuracy.

Original: In our current results, SBP produces only a small amount of trainingspeedup( ∼ 1 . 1 × ).
Revised: In our current results, SBP produces only a small amount of trainingspeedup( ∼ 1 . 1 × ).

Original: How to further speed up the training process is part of our future work.
Revised: How to further speed up the training process is part of our future work.

Original: Societal impact: Our work studies the stochastic backpropagation mechanism, which can be usedfor training general deep learning models.
Revised: Societal impact: Our work studies the stochastic backpropagation mechanism, which can be usedfor training general deep learning models.

Original: Misuse can potentially cause societal harm.
Revised: Misuse can potentially cause societal harm.

Original: However, webelieve our work is a general approach and is consist with the common practice in machine learning.
Revised: However, webelieve our work is a general approach and is consist with the common practice in machine learning.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main
Revised: (a) Do the main

Original: claims made
Revised: claims made

Original: in
Revised: in

Original: the abstract and introduction accurately
Revised: the abstract and introduction accurately

Original: reflect
Revised: reflect

Original: the paper’scontributions and scope
Revised: the paper’scontributions and scope

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] See Section 6.
Revised: [Yes] See Section 6.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See
Revised: [Yes] See

Original: Section 6.
Revised: Section 6.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]
Revised: [N/A]

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes] We will add
Revised: [Yes] We will add

Original: instructions in the supplemental material.
Revised: instructions in the supplemental material.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes] See Section 5’s first paragraph and Section 5.1’s first paragraph.
Revised: [Yes] See Section 5’s first paragraph and Section 5.1’s first paragraph.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running exper-
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running exper-

Original: iments multiple times)?
Revised: iments multiple times)?

Original: [No] We will include these error bars in the supplemental
Revised: [No] We will include these error bars in the supplemental

Original: material.
Revised: material.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Section 5’s first paragraph.
Revised: [Yes] See Section 5’s first paragraph.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We cited the all
Revised: [Yes] We cited the all

Original: datasets we used, such as ImageNet [18] ad COCO [11].
Revised: datasets we used, such as ImageNet [18] ad COCO [11].

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [No] ImageNet is under BSD 3-ClauseLicense and COCO is under Creative Commons Attribution 4.0 License.
Revised: [No] ImageNet is under BSD 3-ClauseLicense and COCO is under Creative Commons Attribution 4.0 License.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No]
Revised: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [No]
Revised: [No]

Original: (e) Did you discuss whether the data you are using/curating contains personally identi-
Revised: (e) Did you discuss whether the data you are using/curating contains personally identi-

Original: fiable information or offensive content?
Revised: fiable information or offensive content?

Original: [No] We do not have sufficient resource to
Revised: [No] We do not have sufficient resource to

Original: exhaustively screen offensive content from the datasets we used.
Revised: exhaustively screen offensive content from the datasets we used.

Original: We expect, because
Revised: We expect, because

Original: the datasets only contain images of people with daily activities and common objects,
Revised: the datasets only contain images of people with daily activities and common objects,

Original: they are unlikely to contain offensive content.
Revised: they are unlikely to contain offensive content.

Original: They are both very widely used.
Revised: They are both very widely used.

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: [N/A]
Revised: [N/A]