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

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 object detection 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 object detection tasks.

Original: During backward propagation, SBP calculatesthe gradients byonly using  a subset of feature maps to savethe GPU memory and computational cost.
Revised: During backward propagation, SBP calculates gradients by using only a subset of feature maps to save GPU memory and computational cost.

Original: We interpret SBP as an efficient way to implement stochastic gradient decent by performing backpropagation dropout, which leads to considerable memory saving and 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 by performing backpropagation dropout, which leads to significant memory saving and training run-time reduction, with a minimal impact on the overall model accuracy.

Original: We offer some good practices to apply SBP in training image recognition models, which can be adopted in learning a wide range of deep neural networks.
Revised: We offer best practices to apply SBP for training image recognition models, 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 save up to 40% of GPU memory with less than 1%of accuracy degradation.
Revised: Experiments on image classification and object detection show that SBP can save up to 40% of GPU memory with less than 1% accuracy degradation.

Original: Code is available at: https://github.com/amazon-research/stochastic-backpropagation
Revised: Code is available at: https://github.com/amazon-research/stochastic-backpropagation

Original:  Introduction
Revised:  Introduction

Original: A common practice to improve accuracy when training deep neural networks is to increase the input resolution [19], model depth [7] and model width [21].
Revised: A common practice to improve accuracy when training deep neural networks is to increase the input resolution [19], model depth [7] and model width [21].

Original: However, this can significantly increase the GPU memory usage, making training on devices with limited resources difficult.
Revised: However, this can significantly increase the GPU memory usage, making training on devices with limited resources difficult.

Original: For example,when training an object detector using ConvNeXt-Base [14] as the backbone with batch size 2, it requires over 17 GB of memory, which cannot fit on most modern GPUs.
Revised: For example, training an object detector using ConvNeXt-Base [14] as the backbone with batch size 2 requires over 17 GB of memory, which cannot fit on many modern GPUs.

Original: Many prior studies have focused on designing memory-efficient training method, especially for large model training on high resolution images or videos.
Revised: A number of prior works have explored memory-efficient training designs, especially for large model training on high resolution images or videos.

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

Original: Gradient checkpointing [2] restores intermediate feature maps by recomputing nodes in the computation graph during the backward pass, which requires extra computations.
Revised: Gradient checkpointing [2] restores intermediate feature maps by recomputing nodes in the computation graph during the backward pass which requires extra computations.

Original: Gradient accumulation [12] splits a mini-batch into smaller chunks to calculate gradients iteratively, which slows down training process.
Revised: Gradient accumulation [12] splits a mini-batch into smaller chunks to calculate gradients iteratively, which slows down training process.

Original: A recent method, Stochastic Backpropagation (SBP) [3], offers a new perspective for memoryefficient model training.
Revised: A recent method, Stochastic Backpropagation (SBP) [3], offers a new perspective for memoryefficient model training.

Original: It performs partial execution of gradients during backpropagation on randomly selected frames when training video models [13, 24, 23, 22].
Revised: It performs partial execution of gradients during backpropagation on randomly selected frames when training video models [13, 24, 23, 22].

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

Original: However, SBP in [3] is only applied to video-based models using two architectures and its design space has not be thoroughly explored to wider tasks.
Revised: However, SBP in [3] is only applied to video-based models using two architectures and its design space has not be thoroughly explored on wider tasks.

Original: The explanation provided in [3] for the effectiveness of SBP is simple but narrow — it attributes the validity of the method to the high redundancy of video.
Revised: The explanation provided in [3] for the effectiveness of SBP is simple but narrow — it attributes the validity of the method to the high redundancy of video frames.

Original: This leaves many questions around SBP unanswered: (1) Is SBP’s effectiveness limited tooperating on the  redundancy of videoframes or can it effectively leverage spatial redundancy?
Revised: This leaves a number of questions around SBP unanswered: (1) Is SBP ’s effectiveness limited to the temporal redundancy of video or can it effectively leverage spatial redundancy?

Original: (2) If so, can
Revised: (2) If so, can

Original: SBP be generalized to image-based tasks, with more limited redundancy?
Revised: SBP be generalized to image-based tasks, with more limited redundancy?

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

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 processin a general form and analyze the effect of SBP on the calculation of gradients.
Revised: To answer these questions, we first formalize the SBP process and analyze the effect of SBP on the calculation of gradients.

Original: With this formulation, we show how to generalize the idea of SBPto image models and present an in-depth study of the design choices when using SBP.
Revised: With this formulation, we show how to generalize the idea of SBPto image models and present an in-depth study of the design choices when using SBP.

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

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

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

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

Original: We validate the generalizability of SBP on two common tasks (image classification and object detection) with two popular network architectures: ViT [6] and ConvNeXt [14].
Revised: We validate the generalizability of SBP on two common tasks (image classification and object detection) with two popular network architectures (ViT [6] and ConvNeXt [14]).

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

Original:  Related Work
Revised:  Related Work

Original: Since deep neural networks are highly extended in both depth [7] and breadth [21] in recent years,  designing memory-efficient training methods has attracted a fair amount of attention.
Revised: Since deep neural networks have expanded in both depth [7] and breadth [21] in recent years, the area of designing memory-efficient training methods has attracted a fair amount of attention.

Original: By saving and recomputing nodes in the computation graph during backpropagation, gradient checkpointing [2] can save a large amount of memory.
Revised: By saving and recomputing nodes in the computation graph during backpropagation, gradient checkpointing [2] can save a large amount of memory.

Original: Besides, gradient accumulation [12] splits the large sample batch into several mini batches and runs them iteratively without updating model parameters.
Revised: Gradient accumulation [12] splits the large sample 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 accumulating the gradients of these sub-iterations and then updating the model parameters are identical to directly optimizing the model with a global batch size.
Revised: This saves the GPU memory proportionally without affecting the accuracy, because accumulating the gradients of these sub-iterations and then updating the model parameters are identical to directly optimizing the model with a 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 [20] proposes to use different mini-batch shapes to speedup the model training.
Revised: Multigrid [20] proposes to use different mini-batch shapes to speedup the model training.

Original: Sparse network [5] particularly targets the recognition tasks but can only save the memory theoretically.
Revised: Sparse network [5] 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 whenever new 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 whenever new ones become available, but they are limited to causal models.

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

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

Original: In this paper, we generalize SBP to more standard computer vision tasks including image recognition.
Revised: In this paper, we generalize SBP to more standard computer vision tasks including image recognition.

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) and show 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) and show how to compute the gradient for two common operations (sec. 3.2).

Original: With these formulations established, we analyze how to perform the chain rule of gradient calculations for SBP (sec. 3.3).
Revised: With these formulations established, 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: We first present a general formulation of SBP [3].
Revised: We first present a general formulation of SBP [3].

Original: We denote a loss function L (Θ) to learn a model on a dataset D = { d j } Nj =1 , where Θ is the collection (vector, matrix, or tensor) of all parameters of the model, N is the total number of training samples, d j ∈ R T × H × W × C is an input data sample, where H and W are spatial resolution height and width respectively, C is the input channel size,
Revised: We denote a loss function L (Θ) to learn a model on a dataset D = { d j } Nj =1 , where Θ is the collection (vector, matrix, or tensor) of all parameters of the model, N is the total number of training samples, d j ∈ R T × H × W × C is an input data sample, where H and W are spatial resolution height and width respectively, C is the input channel size,

Original: T = 1 represents an image and T > 1 represents a video.
Revised: 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 compute the loss after 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 compute the loss after the forward pass, where B is the batch size.

Original: During the backward pass, different from the traditional backpropagation that relies on the full feature maps X i ∈ R B × T i × H
Revised: During the backward pass, different from the 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 the index of the layers), an incomplete execution for backpropagation ( i.e., SBP) only utilizes a subset of the feature maps X subi ∈ R B × T subi × H subi × W subi ×
Revised: C i ( i is the index of the layers), an incomplete execution for backpropagation ( i.e., SBP) only utilizes a subset of the feature maps X subi ∈ R B × T subi × H subi × W subi ×

Original: C i to approximate the gradient calculation.
Revised: C i to approximate the gradient calculation.

Original: Specifically, it updates the model weights as
Revised: Specifically, it updates the model weights as

Original: [Equation]
Revised: [Equation]

Original:  where η is the learning rate, ∇ Θ is the collection of model weight gradients, Θ iter and Θ iter +1 are model weights at the current and next iteration, respectively.
Revised: ) where η is the learning rate, ∇ Θ is the collection of model weight gradients, Θ iter and Θ iter +1 are model weights at the current and next iteration, respectively.

Original: SBP keeps allthe full forward paths, butit can save memory by only caching a subset of feature maps in the backward pass.
Revised: SBP computes all forward paths, but only caches a subset of feature maps in the backward pass.

Original: Note that the method in [3] is a special case of (1) as it performs on the temporal dimension and uses a subset of video frames X T subi
Revised: Note that the method in [3] is a special case of (1) as it operates on the temporal dimension by using a subset of video frames X T subi

Original: i
Revised: i

Original: [Equation]
Revised: [Equation]

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

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

Original: We denote that at layer i of the model, an operator f i with learnable weights W f i to processes the feature maps
Revised: We denote that at layer i of the model, an operator f i with learnable weights 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 (PW-Conv) layer.
Revised: We start by exploring f i as a point-wise convolutional (PW-Conv) layer.

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

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

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

Original: It is mainly used to enrich the channel information with weights W f i ∈ R C ini × C outi , and the analysis on
Revised: It is mainly used to enrich the channel information with weights W f i ∈ R C ini × C outi , and the analysis on

Original: PW-Conv can be easily extended to other layers.
Revised: PW-Conv can be easily extended to other layers.

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

Original: [Equation]
Revised: [Equation]

Original: where X i +1 ∈ R ( B × T i × H
Revised: where X i +1 ∈ R ( B × T i × H

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

Original: × C outi
Revised: × C outi

Original: ( X i and X
Revised: ( X i and X

Original: i +1 are reshaped as a matrix format).
Revised: i +1 are reshaped as a matrix format).

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 full feature maps X i ∈
Revised: It requires storing the weights W f i ∈ R C ini × C outi and full feature maps X i ∈

Original: R ( B × T
Revised: R ( B × T

Original: i × H
Revised: i × H

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

Original: × C ini to compute the gradients.
Revised: × C ini to 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 maps dominates the memory usage.
Revised: ≫ C outi , i.e., caching of feature maps dominates the memory usage.

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

Original: In the backward pass of SBP process, we split the full feature maps X i into two subsets X keepi and X dropi , which denote the corresponding indices of the sampled feature maps that are kept and dropped , respectively.
Revised: In the backward pass of SBP process, we split the full feature maps X i into two subsets X keepi and X dropi , which denote the corresponding indices of the sampled feature maps that are kept and dropped , 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 input, or X i [: , : , even index on H, odd index on W, :] for image input (e.g., Fig.

Original: 3 left), and X dropi is the complementary part of the feature maps.
Revised: 3 left), and X dropi is the complementary part 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 denote keep-ratio r as the number of gradient kept indices j ∈ Z keepi over the number of all indices.
Revised: We denote keep-ratio r as the number of gradient 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 spatial or temporal dimensions can be calculated independently through the forward and backward passes.
Revised: One advantage of simplifying the SBP problem by using PW-Conv operator is that the nodes in spatial or temporal dimensions can be calculated independently through the forward and backward passes.

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 pass:
Revised: (4) as well as the backward pass:

Original: [Equation]
Revised: [Equation]

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

Original: In other words, the memory usage and the backward computation cost are proportional to the keepratio of the original case.
Revised: In other words, the memory usage and the backward computation cost are proportional to the keepratio of the original case.

Original: Mathematically, it is equivalent to set dX dropi +

Revised: Mathematically, it is equivalent to set dX dropi +

Original: = 0 (thus dX dropi = 0 ) and uses dW keepf i = X keepiT dX keepi
Revised: = 0 (thus dX dropi = 0 ) and uses dW keepf i = X keepiT dX keepi

Original: +1 as a stochastic approximation of the original weight gradients dW f i to update the model weights.
Revised: +1 as a stochastic approximation of the original weight gradients dW f i to update 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- Conv, nodes of general ConvLayer with kernel size > 1 in spatial or temporal dimensions are no longer independent in forward and backward calculations.
Revised: Unlike the PW- Conv, nodes of general ConvLayer with kernel size > 1 in spatial or temporal dimensions are no longer independent in forward 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 backpropagation of the ConvLayer is calculated as follows:
Revised: The regular backpropagation of the ConvLayer is calculated as follows:

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

Original: ) where ∗ is the convolution operation, zeropad is to pad the edge of the matrix with zeros, and W ′ f
Revised: ) where ∗ is the convolution operation, zeropad is to pad the edge of the matrix with zeros, and W ′ f

Original: i represents the 180 ◦ rotation of W f i .
Revised: i represents the 180 ◦ rotation of W f i .

Original: From Eq.
Revised: From Eq.

Original: (6), in the gradient calculation of ConvLayer at a given location j on the input feature map as x ji ∈ X
Revised: (6), 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 gradient calculation of dx ji .
Revised: i , all the elements from dX i +1 inside the region of the kernel size k will contribute to the gradient calculation of dx ji .

Original: Therefore, in SBP of ConvLayer, when we set dX dropi +1 = 0 for certain areas, in general, two things will happen: (a) even if j ∈ Z drop and dx ji +1 = 0 , dx ji will not be 0 due to the non-zero gradient contribution of the neighbors of j on dX i +1 ; (b) dx j i will be an approximated gradient unless all dx pi +1 with p in the neighbor areas of j in terms of the kernel size k ( i.e., indices p are within the convolutional kernel of index j ) are kept and there is no chain rule effect on dx pi +1 .
Revised: Therefore, in SBP of ConvLayer, when we set dX dropi +1 = 0 for certain areas, in general, two things will happen: (a) even if j ∈ Z drop and dx ji +1 = 0 , dx ji will not be 0 due to the non-zero gradient contribution of the neighbors of j on dX i +1 ; (b) dx j i will be an approximated gradient unless all dx pi +1 with p in the neighbor areas of j in terms of the kernel size k ( i.e., indices p are within the convolutional kernel of index j ) are kept and there is no chain rule effect on dx pi +1 .

Original: However, there is a special scenario: when s ≥ k , the neighbor effects disappear due 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 disappear due to the fact that dX i +1 needs to be zero interweaved in gradient calculation.

Original: As a result, nodes in spatial or temporal dimensions are getting more independent, which is similar to the case of PW-Conv.
Revised: As a result, nodes in spatial or temporal dimensions are patch-wise independent, which is similar to the case of PW-Conv.

Original: 3.2.3 Discussion
Revised: 3.2.3 Discussion

Original: From the above derivation, we see that the approximation effect of SBP performs differently when the operator changes.
Revised: From the above derivation, we see that the approximation effect of SBP performs differently when the operator changes.

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

Original: While for convolutional layers with kernel size larger than 1, most of the gradients on the activations are neither 0 nor exact.
Revised: While for convolutional layers with kernel size larger than 1, most of the gradients on the activations are 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 supplementary material.
Revised: For derivation of other operators, please refer to the supplementary 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, we derive the calculation of SBP on a stack of multiple layers.
Revised: In this section, we derive the calculation of SBP on a stack of multiple layers.

Original: For simplicity, we take two layers as an example, which can be easily generalized to more layers.
Revised: For simplicity, we take two layers as an example, 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 scenario with two PW-Conv layers.
Revised: We first consider a scenario with two PW-Conv layers.

Original: We assume a layer before f i , i.e., f i − 1 , is another PW-Conv layer with weights W f i − 1 .
Revised: We assume a layer before f i , i.e., f i − 1 , is another PW-Conv layer with weights W f i [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 + has non-zero gradients at kept indices ( i.e., dX keepi +1 ) and zero gradients at dropped indices ( i.e., dX dropi +1 ), Eq.
Revised: Since dX i + has non-zero gradients at kept indices ( i.e., dX keepi +1 ) and zero gradients at dropped indices ( i.e., dX dropi +1 ), Eq.

Original: (8) can be updated to:
Revised: (8) can be updated 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 and dX dropi .
Revised: f i − 1 , we split dX i into dX keepi and dX dropi .

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

Revised: Assume that the keep and 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 and drop indices set on layer i − 1 will be:
Revised: (9), the keep and drop indices set on layer i − 1 will be:

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

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

Original: [Equation]
Revised: [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: 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 ConvLayer with kernel size k > 1 and stride s .
Revised: Now let’s look at the chain-rule effect of the case when f i − 1 is a ConvLayer with kernel size k > 1 and stride s .

Original: Using the chain rule, we can update the Eq.
Revised: Using the chain rule, we can update the Eq.

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

Original: [Equation]

Original: [Equation]
Revised: [Equation]

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

Original: From Eq.
Revised: From Eq.

Original: (12), the dropout at layer i makes the matrix dX
Revised: (12), the dropout at layer i makes the matrix dX

Original: i sparse (only indices at Z keepi +1 are non-zero).
Revised: i sparse (only indices at Z keepi +1 are non-zero).

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

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

Original: =
Revised: =

Original: Z dropi .
Revised: Z dropi .

Original: If we stack PW-Conv + ConvLayer further and further in an interweaving way ( i.e., Z dropi +

Revised: If we stack PW-Conv + ConvLayer further and further in an interweaving way ( i.e., Z dropi +

Original: ̸ =
Revised: ̸ =

Original: Z dropi ), the chain rule effect may lead to a very sparse output derivative matrix ( e.g., dX i − 1 ).
Revised: Z dropi ), the chain rule effect may lead to a very sparse output derivative matrix ( e.g., dX i − 1 ).

Original: However, stacking ConvLayers (due to the space limitation, we skip the math derivation) will not have such a trend, as the gradient dropout in ConvLayer only creates approximated gradient but not many zeros, because of the neighbor effect of gradient calculation.
Revised: However, stacking ConvLayers (due to the space limitation, we skip the math derivation) will not have such a trend, as the gradient dropout in ConvLayer only creates approximated gradient but not many zeros, because of the neighbor effect 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 Z keepi − 1 and a vanishing gradient of dW f i − 1 .
Revised: (11) that stacking PW-Conv layers, especially in an interweaving way, may lead to a sparse keep set Z keepi − 1 and a vanishing gradient of dW f i − 1 .

Original: However, in practice, we always drop gradients at the same positions across layers to avoid the gradient vanishingproblem and to achieve good accuracy performance.
Revised: Thus, in practice, we recommend to drop gradients at the same positions across layers to avoid the gradient vanishing and maintain high accuracy.

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 and thus should be avoided.

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

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 of the SBP technique in Alg.
Revised: With a deeper understanding of how SBP works, we provide a simple and efficient implementation of the SBP technique in Alg.

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

Original: We emphasize that it is very simple and easy to implement, it only needs 3 lines of code to do the forward, and the backward (no need to implement manually) can be seamlessly handled 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 lines of code to do the forward, and the backward (no need to implement manually) can be seamlessly handled by the ease of “autograd engine” in the deep learning frameworks.

Original: It supports point-wise convolutional layers, the dot-product attention layers, and multilayer perception (MLP) in popular networks including ViT [6] and its variants [6, 11, 13], ConvNeXt [14], and hybrid architectures [4].
Revised: It supports point-wise convolutional layers, the dot-product attention layers, and multilayer perception (MLP) in popular networks including ViT [6] and its variants [6, 11, 13], ConvNeXt [14], and hybrid architectures [4].

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():
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():

Original: [Figure]
Revised: [Figure]

Original: [grad_
Revised: [grad_

Original: [Figure]
Revised: [Figure]

Original: _idx]
Revised: _idx]

Original: [Figure]
Revised: [Figure]

Original: f(inputs[grad_drop_idx]
Revised: f(inputs[grad_drop_idx]

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

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 [6] model to evaluate the top-1 accuracy on ImageNet [18] validation dataset.
Revised: We use ViT-Tiny [6] model to evaluate the top-1 accuracy on ImageNet [18] validation dataset.

Original: The
Revised: The

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

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

Original: The conclusions remain  for larger models and more training epochs.
Revised: The conclusions remain consistent for larger models and longer training schedules.

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 reduces GPU memory consumption by 33% with an accuracy drop within 1.5%.

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

Original: Therefore, unless stated otherwise, we apply SBP on 2/3 of all layers of ViT models for the rest of this work.
Revised: Therefore, unless stated otherwise, 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 the gradients keep-ratio and accuracy in SBP.
Revised: In this section, we discuss the trade-off between the gradient keep-ratio and accuracy in SBP.

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

Original: In general, the higher the keep-ratio is, the more gradient information is preserved, and the higher the accuracy is, the less memory is saved.
Revised: In general, the higher the keep-ratio is, the more gradient information is preserved, and the higher the accuracy is, 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: In this case, we apply the same keep-ratio to all the SBP layers at the same gradient kept indices and we call it uniform keep-ratios method.
Revised: In this case, we apply the same keep-ratio to all the SBP layers at the same gradient kept indices. We name this strategy "uniform keep-ratio".

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 achieves a reasonable trade off while 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 the keep-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 the keep-ratios from early to late layers while keeping their average keepratio to be 0.5.

Original: Specifically, the keep-ratios are [0 . 25 , 0 .
Revised: Specifically, we use keep-ratios of [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 ImageNet training dataset.
Revised: We simulate the gradient calculation on ViT-Tiny with 2000 data batches randomly sampled from the ImageNet training dataset.

Original: We apply different keep-ratio methods on the same model with the same fixed weights, and calculate the correlation (measured by cosine similarity) between the weights gradient of applying SBP and the original exact weights gradient without applying SBP, i.e., Cosine ( dW SBP , dW no − SBP ) .
Revised: We apply different keep-ratio methods on the same model with the same fixed weights, and calculate the correlation (measured by cosine similarity) between the weight gradients of applying SBP and the original exact weights gradients, i.e., Cosine ( dW SBP , dW no − SBP ) .

Original: From Fig.
Revised: From Fig.

Original: 2, we observe that the weights gradient has a stronger correlation by using uniform keep-ratios compared to the case of increasing or decreasing keepratios.
Revised: 2, we observe that the weight gradients have a stronger correlation when using uniform keep-ratios compared to the case of increasing or decreasing keep-ratios.

Original: We believe that  uniform keep-ratioscan keep the consistency of gradient dropping locations between different layers and can preserve gradient information at these spatial locations.
Revised: We believe that the fact that uniform keep-ratios keep the location of the dropped gradients consistency between different layers helps preserve gradient information at these spatial locations.

Original: However, non-uniform keep-ratios will have differentdropping locations between different layers. As an overall result, it may lose more gradient information on more spatial locations and produce a less accurate estimation of the original gradients, and hence it may have a lower accuracy performance.
Revised: However, non-uniform keep-ratios will drop at different locations between different layers, which adds noise to produce a less accurate estimation of the original gradients, and hence lowers accuracy.

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 of feature maps.
Revised: SBP calculates the gradients by only using a subset of feature 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 only propagates the valid gradients (non-zeros) on these fixed locations, which forces the model to only learn 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 only propagates the valid gradients (non-zeros) on these fixed locations, which forces the model to only learn from fixed partial spatial information and therefore severely hurts the model performance.

Original: A better strategy is to randomly sample the gradient kept indices in each training step, so that it can statistically visit every spatial location with equal importance.
Revised: A better strategy is to randomly sample the gradient kept indices in each training step, so that it can statistically 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 explanation is that grid-wise samplingis a more structured sampling method and it provides a more accurate estimation of gradients calculation.
Revised: Grid-wise sampling more structured sampling approach seems to provides a more accurate estimation of gradients.

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 of applying SBP on MLP (two linear layers) is proportional to the keep-ratio r .
Revised: Theoretically, the memory usage of applying SBP on MLP (two linear layers) is proportional to the keep-ratio r .

Original: However, the memory saving 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 memory saving 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 [3], activation maps needed
Revised: From [3], activation maps needed

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

Original: for gradients calculation in MHSA include the inputs and outputs ( 2 hdn ), QKV vectors (each with hdn ), the attention weight maps ( 2 hnn before and after the softmax), thus requiring the memory of 3 hdn + 2 hnn in total, where h is the number of heads.
Revised: for gradients calculation in MHSA include the inputs and outputs ( 2 hdn ), QKV vectors (each with hdn ), the attention weight maps ( 2 hnn before and after the softmax), thus requiring the memory of 3 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 only) or on all QKV (Drop QKV) enjoys the following memory usage ratio:
Revised: Query only) or on all QKV (Drop QKV) enjoys the following memory usage ratio:

Original: [Equation]
Revised: [Equation]

Original: In video transformers, the number of tokens n is much larger than the head dimension d (i.e., in video
Revised: In video transformers, the number of tokens n is much larger than the head dimension d (i.e., in video

Original: Swin Transformer [13] , n = 392 and d = 32 , dn = 0 . 082 ) and thus the attention weight maps 2 hnn comprise most of the memory.
Revised: Swin Transformer [13] , n = 392 and d = 32 , dn = 0 . 082 ) and thus the attention weight maps 2 hnn comprise most of the memory.

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

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 , d n = 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 , d n = 0 . 327 .

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

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

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

Original: 42 × ) memory while dropping all QKV reduces its memory usage ratio to 0 .
Revised: 42 × ) memory while dropping all QKV reduces its memory usage ratio 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: We also report the performance of dropping gradients along the attention heads, but we did not observe accuracy improvement.
Revised: We also report the performance of dropping gradients along the attention heads, but we did not observe accuracy 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 gradients in Fig.
Revised: We plot the cosine similarity of the weights gradients of applying SBP and the original exact gradients in Fig.

Original: One interesting observation is that Drop QKV has a slightly better correlation to Drop
Revised: One interesting observation is that Drop QKV has a slightly better correlation to Drop

Original: Query only and a much higher correlation to Drop Head.
Revised: Query only and a much higher correlation to Drop Head.

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

Original:  Generalizability of SBP
Revised:  Generalizability of SBP

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

Original: In order to do a fair comparison of the 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 depth
Revised: In order to do a fair comparison of the 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 depth

Original: [Table]
Revised: [Table]

Original: augmentation [9] for different model sizes.
Revised: augmentation [9] for different model sizes.

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

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

Original: Each accuracy is reported with an average result of three runs with different random seeds.
Revised: Each accuracy is reported with an average result of three runs with different random seeds.

Original: For more details of experimental settings, please refer to the supplementary material.
Revised: For more details of experimental settings, please refer to the supplementary material.

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 [6] 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 [6] and convolutional-based ConvNeXt [14].

Original: We train 300 epochs for both these two networks by following the training recipe of ConvNeXt [14]
Revised: We train 300 epochs for both these two networks by following the training recipe of ConvNeXt [14]

Original: [Table]
Revised: [Table]

Original: For ViT models, we apply SBP on the last 8 transformer blocks including MHSA and MLP layers with 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 layers with the keep-ratio of 0.5, reducing GPU memory usage by approximately 30%.

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

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

Original: For ConvNeXt models, it builds the basic block with a depth-wise convolutional (DW-Conv) layer with kernel size 7 × 7 followed by two point-wise convolutional (PW-Conv) layers 3 and the DW-Conv layer 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) layer with kernel size 7 × 7 followed by two point-wise convolutional (PW-Conv) layers 3 and the DW-Conv layer only consumes a small portion of GPU memory (less than 10%).

Original: It has multi-scale stages where early stages have higher redundancy and occupy most of the GPU memory.
Revised: It has multi-scale stages where early stages have higher redundancy and occupy 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 the last stage.
Revised: SBP on the PW-Conv and down-sampling layers in the first three stages, and do not apply it on the last stage.

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

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

Original: We fine-tune the vanilla Mask R-CNN [8] and Cascade Mask R-CNN [1] on the COCO dataset [10] with ConvNeXt-Tiny and ConvNeXt-Base backbones pretrained on ImageNet-1K, respectively.
Revised: We fine-tune the vanilla Mask R-CNN [8] and Cascade Mask R-CNN [1] on the COCO dataset [10] with ConvNeXt-Tiny and ConvNeXt-Base backbones pretrained on ImageNet-1K, respectively.

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

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

Original: 2 show that SBP with keep-ratio of 0.5 can still learn the detection task reasonably well with 0.2% and 0.7% loss on the box and mask average precision (AP) for ConvNeXt-Base and ConvNeXt-Tiny backbones, respectively.
Revised: 2 show that SBP with keep-ratio of 0.5 can still learn the detection task reasonably well with 0.2% and 0.7% loss on the box and mask average precision (AP) for ConvNeXt-Base and ConvNeXt-Tiny backbones, respectively.

Original: However, it only consumes about 0.7 × of the GPU memory.
Revised: However, it only consumes about 0.7 × of the GPU memory.

Original: Note that training detection models on high resolution images (up to 800 × 1333) is memory intensive even for a very small batch size.
Revised: Note that training detection models on high resolution images (up to 800 × 1333) is memory intensive even for a very small batch size.

Original: For example, with backbone of ConvNeXt-Base and batch size of 2, the model training requires 17.4GB of GPU memory which may lead to out of memory error for common devices such as Tesla 16GB V100.
Revised: For example, with backbone of ConvNeXt-Base and batch size of 2, the model training requires 17.4GB of GPU memory which may lead to out of memory error for common devices such as Tesla 16GB V100.

Original: While with the help of SBP, the model can still be effectively trained 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 effectively trained on GPUs with 12.5GB memory, which is much more feasible for many research groups.

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

Original: [Table]
Revised: [Table]

Original:  Discussion
Revised:  Discussion

Original: In this work, we present a comprehensive study of the Stochastic Backpropagation (SBP) mechanism with a generalized implementation.
Revised: In this work, we present a comprehensive study of the Stochastic Backpropagation (SBP) mechanism with a generalized implementation.

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

Original: We show that our approach can reduce up to 40% of the
Revised: We show that our approach can reduce up to 40% of the

Original: GPU memory when training image recognition models under various deep learning backbones.
Revised: GPU memory when training image recognition models under various deep learning backbones.

Original: Limitations:
Revised: Limitations:

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

Original: In our current results, SBP produces only a small amount of training speedup( ∼ 1 . 1 × ).
Revised: In our current results, SBP produces only a small amount of training speedup( ∼ 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 used for training general deep learning models.
Revised: Societal impact: Our work studies the stochastic backpropagation mechanism, which can be used for training general deep learning models.

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

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

Original:  Acknowledgments and Disclosure of Funding
Revised:  Acknowledgments and Disclosure of Funding

Original: We thank the anonymous reviewers for their helpful suggestions.
Revised: We thank the anonymous reviewers for their helpful suggestions.

Original: This work was funded by Amazon.
Revised: This work was funded by Amazon.

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (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 Section 6.
Revised: [Yes] See Section 6.

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

Original: [Yes] 2.
Revised: [Yes] 2.

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] (b) Did you include complete proofs of all theoretical results?
Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

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

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See Section 5’s first paragraph and Section 8.3 in the supplementary material.
Revised: [Yes] See Section 5’s first paragraph and Section 8.3 in the supplementary material.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they 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 experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] Each accuracy report is an average result of running experiments three times with different random seeds.
Revised: [Yes] Each accuracy report is an average result of running experiments three times with different random seeds.

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

Original: [Yes] See Section 5’s first paragraph and Section 8.3 in the supplementary material.
Revised: [Yes] See Section 5’s first paragraph and Section 8.3 in the supplementary material.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We cited the all datasets we used, such as ImageNet [18] ad COCO [10].
Revised: [Yes] We cited the all datasets we used, such as ImageNet [18] ad COCO [10].

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-Clause License and COCO is under Creative Commons Attribution 4.0 License.
Revised: [No] ImageNet is under BSD 3-Clause License 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] (d) Did you discuss whether
Revised: [No] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: [No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [No] We do not have sufficient resource to exhaustively screen offensive content from the datasets we used.
Revised: [No] We do not have sufficient resource to exhaustively screen offensive content from the datasets we used.

Original: We expect, because the datasets only contain images of people with daily activities and common objects, they are unlikely to contain offensive content.
Revised: We expect, because the datasets only contain images of people with daily activities and common objects, 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... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

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