Learning with Auxiliary Activation for Memory-Efficient Training
Authors: Sunghyeon Woo
Dongsuk Jeon
Venue: ICLR 2023 poster (2023)
Keywords: Memory Efficient Training
Auxiliary Activation
Backpropagation
Deep Learning

Abstract: While deep learning has achieved great success in various fields, a large amount of memory is necessary to train deep neural networks, which hinders the development of massive state-of-the-art models. The reason is the conventional learning rule, backpropagation, should temporarily store input activations of all the layers in the network. To overcome this, recent studies suggested various memory-efficient implementations of backpropagation. However, those approaches incur computational overhead due to the recomputation of activations, slowing down neural network training. In this work, we propose a new learning rule which significantly reduces memory requirements while closely matching the performance of backpropagation. The algorithm combines auxiliary activation with output activation during forward propagation, while only auxiliary activation is used during backward propagation instead of actual input activation to reduce the amount of data to be temporarily stored. We mathematically show that our learning rule can reliably train the networks whose loss landscape is convex if the auxiliary activation satisfies certain conditions. Based on this observation, we suggest candidates of auxiliary activation that satisfy those conditions. Experimental results confirm that the proposed learning rule achieves competitive performance compared to backpropagation in various models such as ResNet, Transformer, BERT, ViT, and MLP-Mixer.

Published as a conference paper at ICLR 2023
LEARNING WITH AUXILIARY MEMORY-EFFICIENT TRAINING
ACTIVATION FOR
Sunghyeon Woo, Dongsuk Jeon
Seoul National University, Seoul, Korea (wsh0917, djjonn @nnu ac.kr
ABSTRACT
While deep learning has achieved great success in various fields, a large amount of memory is necessary train deep neural networks, which hinders the develop- ment of massive state-of-the-art models. The reason is the conventional learning rule, backpropagation, should temporarily store input activations of all the lay- ers in the network. To overcome this, recent studies suggested various memory- efficient implementations of backpropagation. However, those approaches incur computational overhead due the recomputation of activations, slowing down neural network training. In this work, we propose a new learning rule which sig- nificantly reduces memory requirements while closely matching the performance of backpropagation. The algorithm combines auxiliary activation with output ac- tivation during forward propagation, while only auxiliary activation is used during backward propagation instead of actual input activation to reduce the amount of data be temporarily stored. We mathematically show that our learning rule can reliably train the networks if the auxiliary activation satisfies certain conditions. Based on this observation, we suggest candidates of auxiliary activation that sat- isfy those conditions. Experimental results confirm that the proposed learning rule achieves competitive performance compared to backpropagation in various models such as ResNet, Transformer, BERT, VIT, and MLP-Mixer.
INTRODUCTION
Backpropagation (l) is an essential learning rule to train deep neural networks and has proven its out- standing performance in diverse models and tasks. general, the wider and deeper the deep learning model, the better the training performance (2) However, increasing model size unavoidably requires larger memory in training hardware such as GPU (3) 4.. This is because backpropagation needs to temporarily store input activations of all the layers in the network generated in forward propagation as they are later used to update weights backward propagation. Consequently, state-of-the-an deep learning models require substantial memory resources due the large amount of input activa- tion to store. In order to train very deep models with limited hardware resources, the batch size may be reduced (3 4) or many GPUs can be used in parallel (5)
G 7 8, 9 TO: 11.. However, reducing the batch size causes a long training time, and the advantage batch normalization (12) disappears. Also, training huge models such as GPT-3 (13) still requires expensive GPU clusters with thousands of GPUs and incurs high I/O costs even with parallelism.
Recently, a wide range of algorithms have been proposed to alleviate this memory requirement. For instance, a new optimizer (5) or neural network architectures (16) 17 18]
19
20
21 have been
suggested to reduce the memory requirements. Gradient checkpointing 22 23 24
25 26) reduces memory space by only storing some of the input activation during forward propagation. Then, it restores the unsaved input activations through recomputation in backward propagation. In-place ac- tivated batch normalization
(27) merges batch normalization layer and leaky ReLU layer and stores the output activation of the merged layer in forward propagation.
backward propagation, the layer input can be reconstructed for training because the leaky ReLU function is reversible. Sim- ilarly, RevNet (28). Momentum ResNet (29), and Reformer (30) employ reversible neural network architectures, which allow for calculating input activation from output activation in backward prop- agation. Gradient checkpointing and reversible network structures reduce training memory space because they partially store input activations (e.g., input activations of selected layers). However,

---

Published as a conference paper at ICLR 2023
these methods incur additional computational overhead because the unstored input activations must be recomputed during backward propagation. Alternatively, algorithms to approximate activation B5
36) (37 (8). but they suffer from performance degrada- tion or slow down training due to additional computations to quantize and dequantize activations. TinyTT. (39) entirely avoids saving activations by updating only bias parameters while fixing weight parameters. However, it is only applicable to fine-tuning of a pre-trained model.
have been suggested (3T 32 33
34
In this study, we propose new learning rule, Auxiliary Activation Learning, which significantly reduce memory requirements for training deep neural networks without sacrificing training speed. We first introduce the concept of auxiliary activation the training process. Auxiliary activations are combined with output activations and become the input activation of the next layer when processing forward propagation, but only the auxiliary activations are temporarily stored instead of actual input activation for updating weights in backward propagation. To justify our algorithm, we prove that an alternate type of input activation could reliably train neural networks if auxiliary activation satisfies certain conditions. Then, propose multiple candidates of auxiliary activations which meet this criterion. Experimental results demonstrate that the proposed algorithm not only succeeds training ResNet models (40) ImageNet (4I) with similar performance to backpropagation, but is also suitable for training other neural network architectures such as Transformer (42), BERT (43), ViT (44), and MLP-Mixee 445)
AUXILIARY ACTIVATION LEARNING
Compression
Reconstruction
a:
Auxiliary Activation
Memory
W143
Memory
Wiss 1
AW t4s
h
AW
1:42
h
8142
0t+2
a1+z
a++2
8
W1+2
Wuz
AW 242
AW 1+2
h
h
8
O
a
a++1
8 [+1
Wi+1
W1+1
aW
AW
h
h
8;
a1
8:
(a) Conventional memory-saving algorithms
(b) Auxiliary Activation Learning
Figure 1: Conventional memory-efficient training algorithms and Auxiliary Activation Learning.
2.1 MEMORY REQUIREMENTS OF BACKPROPAGATION
To understand why backpropagation uses a large memory space for training, we describe how it trains deep neural network. The training process of backpropagation can be expressed by the equations below.
h1+1 = <6 WW /y + b(+1)
(1)
6, W 8141 O o (yI)
(2)
AW = -7::+1 hi
(3)
Here h;, W4+1, and bi+1 denote the input activation, weight, and bias of the hidden layer + 1, respectively. d, 8, and represent the nonlinear function, backpropagated error, and learning rate, respectively. yi is the output of a linear or convolutional layer (i.e., 3/ = W; h + b)), which becomes the input of the nonlinear function ø. In forward propagation, the input propagates through the network following equation (T). During this process, the generated input activations are stored in memory. In backward propagation, the emor is backpropagated through the network based on
2

---

Published as a conference paper at ICLR 2023
equation The weights are updated using the backpropagated errors and the activations stored in memory (equation (). Due to the time gap between forward and backward computations, all activations obtained in forward propagation must be temporarily stored in memory they are later used for backward propagation. Therefore, as the model size increases, the amount of activation to be stored in memory increases accordingly. To alleviate this issue, gradient checkpointing (22) 23 244 25 26) stores only some of the activations in memory during forward propagation, and the rest of the activations are restored during backpropagation by recomputing equation ( Similarly, activation compression methods (31, 32, B3 35 36, 37) quantize activations before saving them in memory reduce memory requirements, and they are dequantized during backward propagation. However, these methods incur computational overhead and increase training time due to additional processes during forward and backward propagation as depicted in Fig. la
2.2 AUXILIARY ACTIVATION LEARNING: STORE AUXILIARY ACTIVATION INSTEAD OF EXACT ACTIVATION
For memory-efficient training with minimal overhead, here we propose a new learning rule, Aux- iliary Activation Learning. First, we introduce the concept of auxiliary activation in the training process. Then, the algorithm trains the model following the equations below.
h = 0: + at
(4)
01+1 oi W 1+1 h + b(+1)
(5)
6, W 8141 @ Ø (3r)
(6)
A AW = 1 af.
(7)
Here aj represents the newly introduced auxiliary activation of layer In forward propagation, each layer takes the sum of auxiliary activation and the previous layer's 's output activation or input activation h, (equation (4). The generated input activation propagates through the layer following
equation ( In backward propagation, the error 5141 is backpropagated by equation (6). Finally, the weights are updated only using auxiliary activation ay instead of the actual input activation hy following equation (7). The learning process of our algorithm is depicted in Fig. [b Now the memory needs to store the auxiliary activations a, not real activations hy. Consequently, auxiliary activation is reused (e.g., a a:+1) or the amount by auxiliary activation itself is small, the required memory space can be reduced accordingly. Furthermore, it does not require additional computations, avoiding training time increases.
the
The intuition behind only using auxiliary activation for weight updates is that the direction of a weight update obtained using auxiliary activation could be similar to that of backpropagation if carefully choose auxiliary activation. More specifically, we use the sum of the auxiliary activation as and the previous layer's output activation o the input activation hy (equation (4). If we take an existing value (e.g., activation from different layer) as the auxiliary activation, this can be seen as making an additional connection the layer. If we update weights using backpropagation as in equation (3), the input activation h, generated in forward propagation will be used for updating weights: AW
(+1 = -178+1 hf = 7 (or + a;) = o0f n1041aaf Consequently, the weight gradient obtained by using our algorithm can be expressed by A WWAT. = A W + '00+107" Therefore, it is expected that the weight update direction of our algorithm can be similar to that of backpropagation if the effect of the second term is ignorable, which will result in proper training of the neural network. In the next section, we prove that this assumption can be realized if we choose the auxiliary activation appropriately.
ANALYSIS OF AUXILIARY ACTIVATION LEARNING
3.1 HOW AND WHEN USING ALTERNATIVE ACTIVATION CAN TRAIN DEEP NEURAL NETWORKS
This section shows how learning with an alternative form of activation can make the loss of nonlinear networks converge minimum value if specific condition is satisfied. For mathematical analysis, we consider a loss function fz) whose gradient is Lipschitz continuous. h and r are column vectors representing real activations and random alternative activations, respectively. 5 is column vector
3

---

Published as a conference paper at ICLR 2023
of backpropagated errors, and W is flattened vector of weights. When weights updated by backpropagation using exact activations h,
AW Vf(W") = V ................................ 5Lh7,-1)
(8)
where t denotes the current epoch. Contrarily, now we assume that the weights are updated using alternative activations r as below, which represents generalization of our algorithm.
AW = Vec(ōirg... S vF(W")
(9)
W W nAW
(10)
Theorem 1. If (2h: - ri) > 0 in all layers, the step size satisfies 0 <nS L. and the gradient of the loss function ((W) is Lipschitz continuous, then J(W) converges to a minimum value.
Theorem can be proven by applying equations (0- to the quadratic upper bound. A detailed proof is provided in Appendix A Because condition 0 S F. can be easily satisfied if we use small step size, rf (2h: rt) 2 is the most important condition for successful learning that alternative activations should satisfy. We name this condition Learning Criterion. Furthermore, the larger the value of r/ (2h: r!). the better the network converges because the amount of loss reduction in each epoch (K R in Appendix(A) increases. One interesting observation is that rf (2hr- r)) 2r: / - | - +| reaches its maximum when F; h4, which translates conventional backpropagation algorithm. For convenience, we normalize the learning criterion inequality with respect to its maximum value ||//|l2 and define a learning indicator as below.
Learning Indicator r (2h, T()
(11)
1h:122
In summary, alternative activation train the network if the learning indicator in equation 11 is
positive. Moreover, the larger the learning indicator, the better the training performance. Finally, the training performance maximized when the learning indicator (r) hr, backpropagation)
3.2
CONSTRUCTING AUXILIARY ACTIVATIONS
Auxiliary Residual Activation
Auxiliary Sign Activation
Memory
Memory
Residual Connection
h
AWaal
h
bit
AW mal
0142 + 0141
8,42
0/+2
8142
$1+2
8
1
AW L12
AW 1+2
hi41 01+1 +
hit1
bit
8
O
+[++11
$111
8
1
AWL-1
AW i
h;
bbttt
h; 1
h
8,
0
+
St
S
8
1
(a) Auxiliary Residual Activation
(b) Auxiliary Sign Activation
Figure Two candidates of auxiliary activation.
Theorem [ suggests that using an alternate type of activation in calculating weight gradient could still reliably train the neural network. However, in reality, it is cumbersome to meet the learning criterion if we use arbitrary value as the alternate activation (see Appendix To address this, we modify the forward propagation path by adding auxiliary activation to the input activation as in equation (4) in the proposed Auxiliary Activation Learning. Then, the input activation hy o; a, and the alternate activation r; = a share a common term aj, making it easier to meet the learning criterion. Also, if employ some values already existing in the network as auxiliary activation as
4

---

Published as a conference paper at ICLR 2023
this modification can be regarded introducing an additional connection in the layer. The learning indicator of Auxiliary Activation Learning is expressed by
(2h1 - ri) |
af (20) + aa)
Learning Indicator
(12)
|
by applying equation 44) to equation [II). For successful learning, the value of the learning indicator
in equation (12) should be positive. To satisfy this learning criterion, we suggest two candidates of
auxiliary activation as below.
Auxiliary Residual Activation First, jary activations and combine them with output activations, which is identical to having residual connection. Then, the learning process is expressed by
use the previous layer's input activations auxil-
0++1 6 (Wi+1 (ot+1 + hi) b(+1)
(13)
W21+2 116
(14)
Using auxiliary residual activations save training memory if we only store h; and use it for updating W1+1 using backpropagation as well as for updating W1+2 and W4+3 using Auxiliary Activation Learning depicted in Fig. [2 and equation (14). Contrarily, backpropagation needs to store h1, h1+1, and h242 for updating these three layers.
Auxiliary Sign Activations Another option using the sign of the output activation as auxiliary activations. Then, the weight update process is represented by the equations below.
0++1 = < Wi+1 (or + esign(o:)) + b(+1)
(15)
A AW = 110 esign(o;))
(16)
The auxiliary sign activation (s; = (sign(o))) can be guaranteed to satisfy the learning criterion.
Because 2a o 2osign(o:)) (o:))o: 2 0, the learning indicator in equation (12 is always positive.
Using auxiliary sign activation can enable memory-efficient training by storing only a 1-bit sign of
activation instead of actual input activation in high precision shown in Fig. (bb In equations
(15) and (16), we multiply the auxiliary activation by hyperparameter to make the two types of activation comparable in magnitude since the output of the sign function -1 or
EXPERIMENTAL RESULTS
In experiments, we trained various deep neural networks using the proposed learning rule and com- pared its training performance to that of backpropagation. We observed if the selected auxiliary activations satisfied the learning criterion during training and also measured the amount of memory savings and training time. Our algorithm was used to train ResNet (40) as well as transformer and its variants (42 45), which are widely used architectures in vision and NLP-related applications.
4.1 TRAINING RESNET
We applied the proposed Auxiliary Activation Learning to ResNet training using two types of auxil- iary activation as discussed above: Auxiliary Residual Activation (ARA) and Auxiliary Sign Activa- tion (ASA). For ARA, replaced the first convolutional layer, which receives a residual connection from the previous block, with an ARA-Conv layer which performs weight updates using auxiliary residual activations instead of actual activations (Fig, Baj. Hence, memory overhead can be reduced for these ARA-Conv layers. The blocks in ResNet (40) can be categorized into different sets based on their dimension. For example, ResNet-50 consists of blocks with an input channel dimension of 256, 512, 1024, or 2048 as displayed in Fig. 36 In experiments with Auxiliary Activation Learn- ing, the blocks are replaced with ARA blocks with different strides ("sharing stride") each set of blocks as shown in Fig. 3c and 3d For instance, ARA(3.4.4.2) has a sharing stride of 3, 4, 4, and 2 in the sets with a dimension of 256, 512, 1024, and 2048, respectively (Fig. 3d]. There- fore, the larger the sharing stride, the more blocks that replaced by ARA blocks and the larger the amount of memory savings. Similarly, for experiments with auxiliary sign activation, the first convolutional layer was replaced by an ASA-Conv layer, which employs auxiliary sign activation
5

---

Published as a conference paper at ICLR 2023
Input Paol
Input Conw; Pool
Input
Cony, Pool
Blocko
Block0
256
Black 1-1
Block 1
Block 1:1
Ral
Block 1-2
AKA Block 1-2
ARA Block 1-2
Replaces actual activations
Comvax
Block 1-3
512
Block 2-1
Block 2-1
Block 2-1
Convlx1
Block Block 2-3 Bluck
BN
Block 2-3
ARA Black
Block
ARA Biock 2-4
X005 ARA-Commlel
1024
1024
Block 3-1
Block 3-1
Block 3-1
Number of channels
D BN Re
Block 3-2
ARA Block
ARA Block 3-7
Block 3-3
ARA Block
Block 3-4
Block
ARA Block 3
ARA Rlock
Block 3-5
D
Block 3-6
ARA Block 3-6
BN
ReLLU Convixt SNN
2048
Alack
Block 4-1
Block 4-1
Block 4-2
ARA Block 4-2
FC
FC Output
FC Output
Output
(a) ARA Block
(b) BP
(c) ARA(2,2,2,2)
(d) ARA(3,4,4,22
Figure 3: Training ResNet-50 using Auxiliary Residual Activation (ARA).
10a
10a
10a
10a
Distribution
Distribution
10a
10a
10a
10a
10a
-0.5
0.5
-0.1
-0.05
0.05
0.1
Learning Indicator
Learning Indicator
(a) ARA(2.2.2.2)
(b) ASA(2,2,2,2)
Figure 4: Distributions of learning indicator values in ResNet- trained on CIFAR-10.
for weight updates. Other details of the experiments including hyperparameters be found in Appendix H
We first trained ResNet-18 CIFAR-10, CIFAR-100 (46), and Tiay-Im (47 by using BP, ARA, and ASA to confirm the effectiveness of Auxiliary Activation Learning and the correlation between training performance and the value of the learning indicator. Fig. 4a shows that ARA not only satisfies the learning criterion (i.e., keeps learning indicator values positive), but also results in relatively large learning indicator values close 1, This consistent with the training performance displayed in Table M ARA achieves high accuracy and even outperforms backpropagation in some cases. Fig. shows that ASA satisfies the learning criterion, but the values of the indicator are much lower than that of ARA. Accordingly, ASA achieves slightly lower accuracy than BP and ARA. More detailed comparisons between ARA and ASA can found Appendix B addition, the experiments in Appendix C show that the performance is significantly degraded when the learning criterion is not satisfied due to using modifications of auxiliary activation. These results suggest that the proposed learning indicator has strong correlation with training performance.
In the next experiment, we applied ARA to training ResNet-50 and ResNet- 152 on ImageNet with various sharing stride values. Experimental results summarized in Table 2 ARA achieved competitive performance close to BP in all cases, and ARA(3,4,2,22 even achieved slightly higher performance than BP in ResNet-152. Since the mechanism of algorithm is entirely different

---

Published as a conference paper at ICLR 2023
Table
Test accuracy of ResNet-18 trained on CIFAR-IO, CIFAR-100 and Tiny ImageNet.
Dataset
ARA(2,2.2.2)
ASA(2.2,2,2)
CIFAR-10 10
94.77
94.81
94.76
Accuracy (%)
CIFAR- -00 Tiny ImageNet
75.81
75.49
75.31
58.43
58.46
57.01
from other memory-saving techniques, our algorithm can be employed in addition to those conven- tional approaches for even larger memory savings. For experiments, we trained the same networks using gradient checkpointing (GCP) (24) and ActNN (34) which is per-group activation compression training. Then, we additionally applied our ARA these learning rules. More specifically. GCP was applied to the layers that did not employ ARA, whereas ActNN was applied to all layers including ARA_Com layers (see Appendices H..T and H.2 for details). Experimental results clearly show that our algorithm allows for larger memory savings with very small increase in training time. This is in contrast with GCP and ActNN, which noticeably slows down training due to additional compu- tations for reconstructing activations during backward propagation. For instance, applying ActNN alone to ResNet-50 reduces memory space by 11.8x while increasing training time from 18m to 49m compared BP. Contrarily, additionally applying ARA(3,4,6,22 increases the compression 15.01x without impacting training speed. We observed a similar trend experiments under iden- tical compression rates (see Appendix Momentum ResNet (29), which is reversible network for saving training memory in Appendix
to
D). Furthermore, additionally compared ARA with
Table Test accuracy, training memory with compression rate (bracketed), and training time (italic) of one epoch in ResNet training on ImageNet with 512 batch size and six RTX-3090 GPUs.
ARA
ARA
ARA
ARA
Models
Baseline
(2,2,2,2)
(3,4,2,2)
(3,4,4,2)
(3,4.6,2)
76.01
75.97
75.89
75.62
75.23
BP
44.6 GB
(1.12x)
(1.2x)
(1,21x)
(1.22x)
17m 35s
17m 59s
17m 54s
17m 55s
18m 35
76.01
75.97
75.89
75.62
75.23
ResNet-50
GCP
(2.2x)
(2.88x)
(3.44x)
(3.54x)
(3.66x)
36m 42s
37m 2x
37m
36m
37m 14s
75.96
75.93
75.67
75.51
75.12
ActNN
(11.8x)
(14x)
(14.6x)
(14.8x)
(15.01x)
48m 48s
48m 48s
48m 10s
47m 58s
47m 08s
77.38
77.14
77.41
76,84
76.64
BP
90.2 GB
(1.16x)
(1.21x)
(1.27x)
(1.29x)
35m 37s
36m 17s
36m 24s
36m 11s
36m 30s
ResNet-152
T1.38
71.14
77.41
76.84
76.64
GCP
(2.1x)
(2.92x)
(3,26x)
(3.75x)
(3.95x)
1h 18m
1h 19m
1h 19m
1h 19m
1h 20m
4.2 TRAINING TRANSFORMER, VIT, AND MLP-MIXER
Transformer (22 achieved state-of-the-arr performance in NLP, and the models inspired by Trans- former such as BERT (3), VIT (44), and MLP-MiX-- (45) are producing promising results in com- puter vision. The structure of these models is quite distinct from ResNet models (40 48). Hence, applied our Auxiliary Activation Learning to Transformer and its variants to verify the scalability of our algorithm. We trained Transformer on IWSLT (9) and BERT-Large on MRPC (50) and MNLI (IT) datasets. For ViT-Large and MLP-Mixer-Lane not used in this experiment since residual connections are connected to layer normalization layers (2) rather than to cully-connected layers, whereas ARA requires residual connections to directly connect to fully-connected layers to realize equation (3). Instead, we applied ASA linear layers in multi-head attention and feedforward network. Fig. depicts how ASA is applied to those lay- ers. Since ASA-Linear layers store only 1-bit auxiliary sign activation instead of exact activation, memory overhead can be mitigated. We trained networks with BP and Mesa (36) which is activa- tion compression training for Transformer like models. Then, we additionally applied ASA to these algorithms (see Appendix H.2 for details). Mesa stores 8-bit quantized activation for linear layers;
we used CIFAR- 100 (46). Note that ARA was

---

Published as a conference paper at ICLR 2023
Add Norm
Add & Norm
Add Norm
Add Norm
Replaced by
ASA-lnnar
ASA-linear
Atdeliear
ASA-Linear
awalllrry rign
BBC
BBC
eB
biq
Drupout
FFN
Diapout
FFN
Drupout
FFN
Diuvuuu
FFN
BellI
BeUU
ReLL
Number of channels
Add Norm
Add & Norm
Add Norm
Add Norm
Linear
+
+
MHA
Salf Atcontonn
MHA
MHA
MHA
LinE
Lioer
AshcLinw
ASAI
(b) ASA2
ASA3
(d) ASA4
Figure 5: Transformer-like networks trained using Auxiliary Sign Activations (ASA). MHA and FFN represent multihead attention and feedforward network, respectively.
hence, applying ASA these linear layers can further reduce memory space ASA-Linear layers only store 1- 1-bit sign activations, Because Mesa is designed for the networks with self-attention, is not applied to MLP-Mixer which does not employ self-attention layers.
Experimental results are displayed in Table 3
When we apply ASA only, it closely matches or even outperforms BP while using smaller memory space. Since ASA needs extract 1-bit sign from actual activation to construct auxiliary activation, training time is slightly increased. However, negligible compared to Mesa which requires quantization and dequantization during training- For example, when training BERT-Large on MNLI, BP and ASA consumed 2h 37m and 2h 41m,
Table Test/validation scores, compression rate (bracketed), and training time (italic) of one epoch of Transformer, BERT, VIT, and MLP-Mixer training with 4096, 32, 512, and 256 batch sizes, respectively.
Models
Datasel
Baseline
ASAI
ASA2
ASA3
ASA4
35.23
35.44
34.87
34.9
35.02
(1.1))
(1.2x)
Trans-
3m 15s
3m 21x
3m 23y
3m 26s
Sm 38v
IWSLT
Mesa
(1.6x)
(1.7x)
(1.7%)
(1.7x)
(1.8k)
Smt 52s
5m
5m 44s
5m 36s
5m &
88.56
88.69
88.23
88.97
88.32
10.9 GB
(1.1x)
(1.2%)
(1.3x)
(1.3x)
Jm 28s
Ian 29c
Im 30c
Im 30s
37x
MRPCC
88.35
85.25
85. 23.
Mesa
(2.1x)
(2.3%)
(2.4x)
(2.4x)
(2.5k)
2m
55x
Large
86.32
86.65
86.49
86.42
$6.39
10.9 GB
(1.1x)
(1.2x)
(1.3x)
(1.3x)
37m
2h 400m
24 40m
2b 3h-4Jm
2h 47m
86,32
86.37
86.29
86.34
86.17
(2.1x)
(2.3x)
(2.4x)
(2.4k)
(2.5k)
3h 57m
3h. 39m
34 35m
3h 72m
3h 24t
92.3
92.81
92.84
92.97
92.65
48 GB
(1.3x)
(1.5x)
(1.5%)
(1.6x)
VIT: Large
2m 46s
2m 48s
2m 50v
2m 51s
2m 52v
CIFAR-100
92.39
92.75
92.72
92.95
92.84
Mesa
(3.0x)
(3.5%)
(3.6x)
(3.7%)
(4.3%)
4m 46s
4m 15s
3m 57g
3m 49s
3m 33s
Mixer-
91.62
91.45
91.72
91.65
91.91
CIFAR-100
GB
(1.3x)
(1.4x)
(1.7x)
(2.0%)
Large
3mz 39s
3m 47s
3m 50s
3m 57s
4m Is
8

---

Published as a conference paper at ICLR 2023
respectively, while Mesa takes 3h 51m. We maximize memory savings by applying Mesa and our algorithm (ASA) simultaneously. If we train ViT-Large on CIFAR-100 using Mesa only, we can obtain a 3.0x compression rate. On the other hand, applying both Mesa and ASA4 simultaneously results in the compression rate of 4.3x. Furthermore, the training time reduced by 1m 13s because linear layers are replaced with ASA-Linear layers, avoiding costly quantization and dequantization processes for these layers. Experimental results under identical compression rates can be found Appendix D which confirm that ASA enables faster training than Mesa with identical memory savings.
4.3 USAGE CASE OF AUXILIARY ACTIVATION LEARNING
The proposed Auxiliary Activation Learning algorithm allows for training deep neural networks on the systems with limited hardware resources by alleviating memory overheads with minimal computational cost. For example, one may increase the width or depth of a neural network to improve its accuracy. In experiments, our algorithm. For ResNet- 152, we fixed other parameters and
scaled ResNet- 152 and BERT-Large using the same batch
size to confirm the effectiveness increased the number of layers or the width of the bottleneck block following the scheme suggested by Chen et al. (34). We also scaled BERT-Large by increasing the number of transformer blocks or the hidden size following the scheme proposed by Liu et al. (38). The results summarized in Table 4
Table 4: Largest models that can be trained using a single GPU with 24GB memory. (ResNet: depth number of layers. width width of the first bottleneck block. BERT-Large: depth number of transformer blocks, width hidden size.)
Models
ResNet
BERT-Large
Learning rule
BP
ARA
ActNN
ARA+ACINN
BP
ASA4
Mesa
AcNNN
Depth
165
622
718
50
60
64
70
Width
62
214
238
1600
1728
1792
1856
Using ARA, we can train 13% deeper or 22% wider models compared to using BP only. Similarly, additionally applying ARA allows for using 15% deeper 1 11% wider models compared using ActNN only. For BERT-Large, we can train 20% deeper or 8% wider models by using ASA. Ad- ditionally applying ASA allows for using 9% deeper 4% wider model compared to using Mesa only. Our algorithm may be used to increase the batch size for faster training using identical GPU resources, and experimental results are provided in Appendix
5 DISCUSSION
In this work, we proposed a new learning rule, Auxiliary Activation Learning, which effectively alleviates memory overhead in deep neural network training. The algorithm employs an alternate form of activations for weight updates, and hence reduces the amount of data to be stored in memory without additional computations. We proved that the algorithm can train neural networks if an alternative form of activations satisfies a certain condition and also provided performance indicator that can anticipate the training performance of the selected auxiliary activation type. Experimental results confirmed that the proposed algorithm successfully trains various deep neural networks for vision and NI.P tasks with smaller memory without training speed reduction. Furthermore, since algorithm is orthogonal existing memory saving algorithms, it is possible to increase the compression rate further without reduction in training speed by applying existing algorithms and algorithm simultaneously.
Since our study aims to save training memory in linear and convolutional layers, it currently cannot be applied to layers that perform other types of operations such normalization and activation functions. In future work, we will try to address those issues to realize a more memory-efficiee learning algorithm and expand its applicability.

---

Published as a conference paper at ICLR 2023
REPRODUCIBILITY
And it was proven through the ex- perimental results in Section We have detailed various points for experimental reproduction in Section E and Appendix H Furthermore, we support code to reproduce our results as follows: http:://hbbbbbooooooooSouuyg//// Actillatiitttiiiiinnneene
We presented the rationale for our algorithm in Section 3
ACKNOWLEDGMENT
This work was supported by the National Research Foundation of Korea (Grant NRF- 0022RICIC100688) and the Institute of Information & Communications Technology Planning Evaluation (Grant 2021-0-01343).
REFERENCES
[1] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning internal represen- tations by propagation. Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
[2] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. CoRR, abs/2001.08.0 2020.
[3] Zifeng Wu, Chunhua Shen, and Anton den Hengel. High-performance semantic segmen- tation using very deep fully convolutional networks. CoRR, abs/1604.04339, 2016.
[4] Jun- Yan Zhu, Taesung Park, Phillip Isola, and Alexei A. Efros. Unpaired image-to toagge-o-immag translation using cycle-consistent adversarial networks. In IEEE, pages 2242-2251, 2017.
[5] Qinfeng Shi, James Petterson, Gideon Dror, John Langford, Alexander J. Smola, Alexander L. Strehl, and Vishy Vishwanathan. Hash kernels. In A/STATS, volume 5 of JMLR Proceedings, pages 496-503, 2009.
[6] Gideon Mann, Ryan T. McDonald, Mehryar Mohri, Nathan Silberman, and Dan Walker. Effi cient large-scale distributed training of conditional maximum entropy models. NIPS, pages 1231-1239, 2009.
[7] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Quoc V. Le, Mark Z. Mao, Mare' Aurelio Ranzato, Andrew W. Senior, Paul A. Tucker, Ke Yang. and Andrew Ng. Large scale distributed deep networks. In NIPS, pages 1232-1240, 2012.
[8] Ryan McDonald, Keith B. Hall, and Gideon Mann. Distributed training strategies for the structured perceptron. In Human Language Technologies: Conference of the North American Chapter of the Association of Computational Linguistics, pages 4564664, 2010.
[9] Martin Zinkevich, Markus Weimer, Alexander Smola, and Lihong Li. Parallelized stochastic gradient descent. In NIPS, pages 2595-2603, 2010.
[10] Alekh Agarwal and John C. Duchi. Distributed delayed stochastic optimization. In NIPS, pages 771- 2011.
[11] Alckh Agarwal, Olivier Chapelle, Miroslav Dudík, and John Langford. A reliable effective terascale linear learning system. J. Mach. Learn. Res., 15(1):1111-1133, 2014.
[12] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, volume 37 of JMLR Workshop and Conference Proceedings, pages 448-456. JMLR.org, 2015.
[13] Tom B. Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh,
10

---

Published as a conference paper at ICLR 2023
Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learn- ers. In Neurips, 2020.
[14] Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive opti- mization, Advances in Neural Information Processing Systems, 32, 2019.
[15] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
[16] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys (CSUR), 2020.
[17] Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lil- licrap. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[18] Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee whye Teh. Set transformer: A framework for attention-based permutation-invarian neural networks. In International conference machine learning, pages
PMLR, 2019.
[19] Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, et al. Swin transformer v2: Scaling up capacity and resolution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12009-12019, 2022.
[20] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document trans- former. arXiv preprint arXiv:2004.05150 2020.
[21] Hang Zhang, Yeyun Gong, Yelong Shen, Weisheng Li, Jiancheng Lv, Nan Duan, and Weizhu Chen. Poolingformer: Long document modeling with pooling attention. In International Conference on Machine Learning, pages 12437-12446. PMLR, 2021.
[22] Andreas Griewank and Andrea Walther. Algorithm 799: revolve: an implementation of check- pointing for the reverse or adjoint mode of computational differentiation. ACM Trans. Math. Softw., 26(1):19-4) 45, 2000.
[23] James Martens and Ilya Sutskever. Training deep and recurrent networks with hessian-free optimization. In Neural Networks: Tricks of the Trade Second Edition, volume 7700 of Lecture Notes in Computer Science, pages 479-535. Springer, 2012.
[24] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/ http//104061444 2016.
[25] Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Mare Lanctot, and Alex Graves. Memory- efficient backpropagation through time. In NIPS, pages 4125 4135-4133.2220 2016.
[26] Jianwei Feng and Dong Huang. Optimal gradient checkpoint search for arbitrary computation graphs. In IEEE, pages 11433-11442, 2021.
[27] Samuel Rota Bulò, Lorenzo Porzi, and Peter Kontschieder. In-place activated batchnorm for memory-optimized training of dnns. In CVPR, pages 5639-56477, 2018.
[28] Aidan N. Gomez, Mengye Ren, Raquel Urtasun, and Roger B. Grosse. The reversible residual network: Backpropagation without storing activations. In NIPS, pages 2214-2224, 2017.
[29] Michael E Sander, Pieme Ablin, Mathieu Blondel, and Gabriel Peyré. Momentum residual neural networks. In International Conference on Machine Learning, pages 9276-9277 9287..MLR, 2021.
11

---

Published as a conference paper at ICLR 2023
[30] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In ICLR, 2020.
[31] Ayan Chakrabarti and Benjamin Moseley. Backprop with approximate activations for memory- efficient network training. In NIPS, pages 2426-2435, 2019.
[32] R. David Evans, Lufei Liu, and Tor M. Aamodt. JPEG-ACT: accelerating deep learning via transform-batt based lossy compression. In 47th ACM/IEEE Annual International Symposium on Computer Architecture, ISCA 2020, Valencia, Spain, May 30 June 3, 2020, pages 860-873. IEEE, 2020.
[33] Fangeheng Fu, Yuzheng Hu, Yihan He, Jiawei Jiang, Yingxia Shao, Ce Zhang, and Bin Cui. Don't t waste your bits! squeeze activations and gradients for deep neural networks via tinyscript. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Re- search, pages 3304-3314. PMIR, 2020.
[34] Jianfei Chen, Lianmin Zheng. Zhewei Yao, Dequan Wang, lon Stoica, Michael W. Mahoney, and Joseph Gonzalez. Actnn: Reducing training memory footprint via 2-bit activation com- pressed training. Marina Meila and Tong Zhang, editors, Proceedings of the 38th Interna- tional Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, volume 139 Proceedings of Machine Learning Research, pages 1803-1813, PMLR, 2021.
[35] R. David Evans and Tor M. Aamodt. AC-GC:
lossy activation compression with guaranteed convergence. In Marc' Aurelio Ranzato. Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, De- cember 6-14, 2021, virtual, pages 27434-27448, 2021.
[36] Zizheng Pan, Peng Chen, Haoyu He, Jing Liu, Jianfei Cai, and Bohan Zhuang. Mesa: A memory-saving training framework for transformers. CoRR, abs/2111.J. 11111111122 124, 2021.
[37] Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu. Exact: Scalable graph neural networks training via extreme activation compression. In International Conference on Learning Representations, 2021.
[38] Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han, Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, Michael Mahoney, and Alvin Cheung. GACT: activa- tion compressed training for generic network architectures. In Kamalika Chaudhuri, Stefanic Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato, editors, International Con- ference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, vol- ume 162 of Proceedings of Machine Learning Research, pages 14139-14552. PMLR, 2022.
[39] Han Cai, Chuang Gan, Ligeng Zhu, and Song Han. Tinytl: Reduce memory, not parameters for efficient on-device learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, -loiiaFlorrna Balcan, and AssuanTier Lin, editors, Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[40] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE, pages 770-778, 2016.
[41] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael S. Bernstein, Alexander C. Berg, and Li Fei- Fei. Imagenet large scale visual recognition challenge. Int. J. Comput. Vis., 111-)22 2015.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, pages 5998-6008, 2017.
12

---

Published as a conference paper at ICLR 2023
[43] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Do- ran, and Thamar Solorio, editors, Proceedings the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACI-HII 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Pa- pers), pages 4171-4186. Association for Computational Linguistics, 2019.
[44] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth I 16x16 words: Transformers for image recognition at scale. In ICLR, 2021.
[45] Ilya O. Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alcxey Dosovitskiy. Mlp-mixer: An all-mlp architecture for vision. In Neurips, pages 24261-24272, 2021.
[46] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009.
[47] Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 231N, 7(7):3, 2015.
[48] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In BMVC, 2016.
[49] Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, Roldano Cattoni, and Mar-
cello Federico. The IWSLT 2016 evaluation campaign. In /WSLT, 2016.
[50] Bill Dolan and Chris Brockett. Automatically constructing a corpus of sentential paraphrases. In Third International Workshop on Paraphrasing (IWP2005), 2005.
[51] Adina Williams, Nikita Nangia, and Samuel R. Bowman, A broad-coverage challenge corpus for sentence understanding through inference. In Marilyn A. Walker, Heng Ji, and Amanda Stent, editors, Proceedings of the 2018 Conference of the North American Chapter of the As- sociation for Computational Linguistics: Human Language Technologies, NAACL-HLT 2018, New Orleans, Louisiana, USA, June 1-6, 2018, Volume I (Long Pupers), pages 1112-1122. Association for Computational Linguistics, 2018
[52] Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. CoRR, abs/1607.064500. 2016.
[53] Sebastian Ruder. An overview of gradient descent optimization algorithms. abs/1609.04747, 2016.
CoRR,
[54] Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In ICLR, 2017.
[55] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
[56] Nitish Srivastava, Geoffrey E. Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhut- dinov. Dropout: a simple way to prevent neural networks from overfitting. Mach. Learn. Res., 15(1):1929-1158, 2014.
[57] Rafael Müller, Simon Kornblith, and Geoffrey E. Hinton. When does label smoothing help? In NeurlPS, pages 4696 2019.
[58] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Tor- ralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
[59] Timothy P Lillicrap, Daniel Cownden, Douglas B Tweed, and Colin Akerman. Random synaptic feedback weights support error backpropagation for deep learning. Nature communi- cations, 7(1):1-10, 2016.
13

---

Published as a conference paper at ICLR 2023
[60] Arild Nøkland. Direct feedback alignment provides learning in deep neural networks. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg. Isabelle Guyon, and Roman Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pages 1037- 1045, 2016.
[61] Qianli Liao, Joel Z. Leibo, and Tomaso A. Poggio. How important is weight symmetry in backpropagation? In Dale Schuurmans and Michael P. Wellman, editors, Proceedings of the Thirtieth AAAI Conference on Artificial Intelligence, February 12-17, 2016, Phoenix, Arizona, USA, pages 1837-1844. AAAI Press, 2016.
14

---

Published as a conference paper at ICLR 2023
APPENDICES
A PROOF OF THEOREMS
Theorem 1. If (2h; T1) > 0 in all layers, the step size satisfies 0 < S E. and the gradient of the loss function f(W) is Lipschitz continuous, then f(W) converges to a minimum value.
Proof. Since gradient of [(x Lipschitz continuous (i.e. [V [/())]2 L.yy x112).
satisfies f(y) S f(x) + V (x))(( - x)) + T - =||2. By substituting W for x and W for y, obtain
( + V/( (W W + | w'
= =S(W") V/(W" + / 1
-S(W') "
=S(W") GK
where
Kt V/(W") A
M V/(W") AW' "AW'
because we assumed that the step size is < S +. By applying equations and
we obtain
Kt VS(W")TAW'
2 IAW"11
Vf(W')TA'' 1 1 (WW"
(I - (AW') AW'
12 Vediri Vec(28 6, 28th
L-
Σ
Vec Vec(51+1121)) - Vee
2
We can show that Vec(AB") ) || B when A, B, and C are column vectors, and B and C' have the same dimension below. Vec(ABT) V Vec(ACT)
Veo((a:cjlú)
larbı, anb,, ambnllurci a2(1, am
(af + + a 1bcc (af +.....4 .......0....... a )220 + (af + + am bo
(a +... + 2 (bicr +... brca)
-IAA|BCTB
Consequently, by applying this equality the inequality above,
K²
Vec(61+1(2h)) - -r) r) (Vc(6(1+17)
2
1=0
| (1
(2hr rz)
2
15

---

Published as a conference paper at ICLR 2023
because (2h; m1) 0 in all layers by the assumption above. Therefore, 0 Σ f(W++!) - f(W')-nk'* - ((W-t))nkkt--pp of f(W) becomes smaller as the number of epochs increases because K 2 O. Consequently, because f(W) decreases and is bounded below by O, f(W) converges to zero during training even when using alternative activations for updating weights by the monotone convergence theorem.
< ((W)) K" which suggests that the value
Furthermore, if the values of rf (2h; - ri) increases, K", which the amount of loss reduction in each epoch, also increases and therefore the network converges well.
B COMPARISONS BETWEEN AUXILIARY RESIDUAL ACTIVATION AND AUXILIARY SIGN ACTIVATION
Irain loss liny imagenet
Validation loss on liny imageniet
ARU
3030
2002
300
(a) Train loss
(b) Validation loss
Figure Train and Validation losses in ResNet- 18 training on Tiny ImageNet
Table T showw that the training memory of ASA is slightly larger than that of ARA since ASA has to obtain and store 1-bit auxiliary sign activation, while ARA does not need to store actual activation and just reuses auxiliary activation. We also measured the mean values of the learning indicator. The mean values of the learning indicator of ARA are significantly higher than those of ASA. Therefore, expect that the loss of the network would converge better by using ARA
from Theorem 11 compared to ASA. This is consistent with the experimental results in Table ARA achieves higher accuracy than ASA in all experiments. In addition, both train and validation losses converge better when using ARA, as shown in Fig. 6 The processing time for one epoch of ASA is slightly larger than that of ARA since ASA needs pack 1-bit auxiliary sign activation, whereas ARA employs auxiliary residual activation as is.
Table 5: Test accuracy, training memory, training time of epoch, and mean values of learning indicator in ResNet-18 training.
Dataset
metric
BP
ARA (2,2.2.2)
ASA (2,2,2,2)
Accuracy
94.77
94.81
94.76
Training memory
605 MB
546 MB
549 MB
CIFAR-10
Training time
20.83s
21.30s
21.46s
Learning Indicator Accuracy
0.69
0.03
75.81
75.49
75.31
Training memory
605 MB
546 MB
549 MB
100
Training time
20.86s
21.32s
21.52s
Learning Indicator
0.71
0.03
Accuracy
58.43
58.46
57.01
Training memory
1211 MB
1091 MB
1097 MB
Tiny ImageNet
Training time
36.82s
37.54s
37.94s
Learning Indicator
0.7
0.03
16

---

Published as a conference paper at ICLR 2023
C DOES ADDING AUXILIARY ACTIVATION REALLY HELP?
Theorem I suggests that an alternative type of activation can replace actual imput activation in back- ward propagation, long it satísfies the learning criterion. However, it is difficult meet this condition using arbitrary values, so we proposed to add the alternate activation to real activation in forward propagation. Then, the input activation h; Ot ag in forward propagation and the alternate activation a; in backward propagation share a common term, which is expected to help sat- isty the learning criterion. In this section, experimentally verify this claim by comparing ARA and ASA to baselines where the alternate activations are not added to real activation in forward propagation.
In baselines, we replace real activation with the previous layer's activation and 1-bit sign of original activation in backward propagation, while they are not added real activation in forward propaga- tion. They are denoted by Residual Activation (RA) and Sign Activation (SA), respectively. SA is identical to simply quantizing activation into bit conventional backpropagation.
10a
10a
Distribution
Distributtin
Distribution
10a
10
107
10
20
10
Learning
Learning
Learning
RA CIFAR-10
(b) RA CIFAR-100
RA Tiny ImageNet
10*
Distribution
10a
Distribution
Distribution
10a
10a
10
10L
10a
10a
10a
0.5
0.5
0.5
Learning Indicator
Learning Indicator
Learning Indicator
(d) ARA CIFAR-10
(e) ARA CIFAR-100
(f) ARA Tiny ImageNet
Figure 7: Distributions learning indicator values RA and ARA in ResNet-18
Figs. and g display learning indicator values of ResNet-18 trained on CIFAR-10, CIFAR- 100, and Tiny ImageNet. These results show that adding auxiliary activation to real activation in for- ward propagation significantly improves the distributions of learning indicator values by positively skewing them toward This consistent with training performance comparisons provided in Table 6.. ARA and ASA noticeably outperform RA and SA in all experiments, and the performance im- provement is more significant for larger datasets. Therefore, we can conclude that adding auxiliary activation to real activation in forward propagation does help neural network training.
Table 6: Test accuracy of ResNet- trained on CIFAR-10, CIFAR-100, and Tiny ImageNet.
Learning Rule
CIFAR-10
CIFAR-100
Tiny ImageNet
BP
94.77
75.81
58.43
RA (2. 2,
94.1
72.23
53.91
Accuracy (%)
ARA (2, 2,
94.81
75.49
58.46
2,
94.36
74.58
55.76
ASA (2, 2,2)
94.76
75.31
57.01
17

---

Published as a conference paper at ICLR 2023
10*
10a-
10a-
10a
Distribution
Distribution
Distribution
.0.5
0.5
-0.5
0.5
-0.5
0.5
Learning Indicator
Learning Indicator
Learning Indicator
(a) SA on CIFAR-10
(b) SA on CIFAR-100
(c) SA on Tiny ImageNet
10"
10-
Distribution
Distribution
10
Distribution
10*
100
10a
102
10a-
104
10-
0.1
0.05
0.05
0.1
0.1
-0.05
0.05
0.1
Learning Indicator
Learning Indicator
Learning Indicator
(d) ASA on CIFAR-10
(e) ASA on CIFAR-100
(f) ASA on Tiny ImageNet
Figure 8: Distributions of learning indicator values of SA and ASA in ResNet-18
D COMPARISONS BETWEEN AUXILIARY ACTIVATION LEARNING AND OTHER MEMORY SAVING ALGORITHMS UNDER IDENTICAL COMPRESSION RATES
In this section, compare Auxiliary Activation Learning with other memory saving algorithms under identical compression rates. We apply GCP to the last Conv-BN layer residual blocks in Fig. 3a because the dimension of input activation in the last batch normalization layer is equal to that of the ARA-Conv layer, which is not stored in memory during ARA. The input activation of the batch normalization layer is not stored Conv-BN layers, but it recomputed from the input activation of the convolutional layer during backward propagation. Through this, an identical amount of memory saving can be achieved by using GCP. For ActNN, we applied it first convolutional layer of residual blocks as we did for ARA in Fig. [a. which resulted in an identical compression rate. While our ASA stores 1-bit auxiliary sign activation, Mesa stores 8-bit compressed activation. Therefore, we cannot achieve the same memory saving by applying Mesa in the same way as ASA. To overcome this issue, we additionally apply Mesa to the batch normalization and layer normalization layers to obtain the same compression rate ASA. Tables and show that ARA and ASA train the networks faster than other semmory-saving algorithms with identical compression rates as they do not require additional costly computations.
Table 7: Test accuracy, training memory, compression rate (bracketed), and training time of one epoch in ResNet-152 training on ImageNet with 512 batch sizes.
ARA (3,4,222)
GCP
ActNN
Accuracy
77.38
77.41
77.38
77.35 74.5 GB (x1.21) 47m 39s
Training memory
90.2 GB
74.5 GB (x1.21)
74.4 GB (x1.21)
Training time
35m 37s
36m 24s
45m 21s
18

---

Published as a conference paper at ICLR 2023
Table 8: Test accuracy, training memory, and training time of one epoch in ViT-Large training on CIFAR-100.
BP
ASA3
Mcsa
Accuracy
92.93
92.97 33.1 (x1.5) 2m 51s
92.91 31.8 GB (x1.5) 3m 12s
Training memory
Training time
2m 468
E COMPARISONS BETWEEN AUXILIARY RESIDUAL ACTIVATION AND MOMENTUM RESNET.
In this section, we compare ARA with Momentum ResNet (MResNet) (29), which is a reversible network that can be applied to diverse networks with residual connections. For comparisons, trained ResNet-50 on CIFAR- with a 5e-2 learning rate for 100 epochs with 10 warm-up steps and cosine annealing. The results are displayed in Table 99 First, ARA closely matches BP in accuracy, but MResNet exhibits a small accuracy degradation. While ARA requires larger training memory than MResNet, ARA is significantly faster since MResNet needs to recalculate the input of residual blocks along with additional forward propagation during backward propagation.
Table 9: Test Accuracy, training memory, and training time of one epoch ResNet-50 with 128 batch sizes.
Dataset
BP
MResNet
ARA (3,4,2,2)
Accuracy
76.83
76.27
76.78
Training memory
3.30 GB
2.54 GB
2.71 GB
Training time
51.54s
2m 21s
53.49s
MAXIMUM BATCH SIZE AND TRAINING TIME WITH AUXILIARY ACTIVATION LEARNING.
We trained various networks using only three RTX-3090 GPUs to confirm the effectiveness of our algorithm. The experimental results for ResNet-152 and ViT-Large are displayed in Tables TT and and I1 Using memory-saving techniques allows for larger batch size for given GPU memory space. Interestingly, although conventional algorithms such as GCP, ActNN, and Mesa enables a larger batch size, it does not translate faster training speed due to additional computations. In contrast, our ARA and ASA algorithms take full advantage of larger batch size and improve training speed.
Table 10: Maximum batch size and training time of ResNet- 152 on ImageNet. The sharing stride of ARA is (3,4,6,2).
BP
GCP
AcINN
ARA
GCP ARA
AcINN ARA
Max batch Training time
350
700
1950
450
1250
2200
49m
1h 4m
56m
47m
57m
52m
Table 11: Maximum batch size and training time of ViT-Large CIFAR- 100.
BP
Mesa
ASA4
Mesa ASA4
Max batch size
600
1600
900
2100
Training time
mm 53s
2m 9s
1m 49h
Im 59s
G COMPARISONS BETWEEN AUXILIARY ACTIVATION LEARNING AND SM3
In this section, compare ARA with SM3 (14,, which is memory-effiiiett optimizer. While our algorithm aims to reduce the amount of input activation to be stored, SM3 aims to reduce the memory required for storing states generated by accumulating previous gradients of parameters.
19

---

Published as a conference paper at ICLR 2023
For fair comparisons, we measured the total memory allocated to GPU instead of the memory required to store activations. Experimental results are summarized in Table 12 The results show that our ASA4 algorithm requires smaller memory space but takes more time process compared to SM3. This is because the SM3 optimizer is faster than the Adam optimizer that we used for all experiments. Since the theoretical backgrounds behind these two methods are quite different, they can be combined together to obtain even larger savings, As displayed in Table [12 if we apply both methods simultaneously, we can obtain the largest memory savings with minimal impact on the accuracy.
epoch
Table 12: Test Accuracy, training memory, and training time of one epoch in Lert-Large on MRPC with 32 batch sizes.
Dataset
BP
SM3
ASA4
SM3 ARA
Accuracy
88.56 22.2 GB
88.23
88.32
88.15 16,55G Im 22s
Total memory
19.9 GB
19.1 GB
Training time
Im 28s
1m 20s
1m 31s
H EXPERIMENTAL DETAILS
H.1 HOW WE APPLY GRADIENT CHECKPOINITING IN EACH EXPERIMENT
Chen al. (24) suggested that when there a BN-ReLU-Cony layer in the network, applying gra- dient checkpointing to those layers can reduce memory requirements without incurring significant computational overhead. More specifically, only the input activation of the BN layer is stored dur- ing forward propagation, and the rest is recalculated in backward propagation. Therefore, while the BN and ReLU layers have to be recomputed during backward propagation to generate the input activation of the Conv layer, there is need to recompute the Conv layer itself which may lead to large computational overhead. We applied this method to BN EAN-ReLU-Convke layers in Fig. Ba to reduce memory overhead without significant amount of recomputation. In addition, gradient checkpointing to BN BN-ReLU ReW---eLUCCn 1x1- -BN layers in Fig. Bal Although we have to recom- pute Conv1x1, its computational overhead is relatively small compared to Conv3x3 layers. We also applied gradient checkpointing to the first Monv-BN-Rel---A U-MMM--PPooo layers in ResNet. Here only the input of the first Conv layer stored during forward propagation. and the input activation of the other layers is recomputed in backward propagation. Although the first Conv layer has to be recom- puted, a large amount of training memory can be saved because the amount of the input activation in the first Conv, BN, and MaxPool layers is large.
apply
H.2
How WE APPLY ACTIVATION COMPRESSION TRAINING IN EACH EXPERIMENT
Quantitation
Doquantizatio
Quantization
D
Dequantiztion
Canv2d
Conv2e
Canw2d
Com22d
ARA ConvZá
ARA Conv2d
Q
Q
Memory
Memory
4++3
q++2
B
Conv2d
Conv2d
Com2d
Conv2d
Comv2d
ARA-
ARA- Conv2d
(a) ActNN
(b) ActNN with ARA
Figure 9: Training process of ActNN with ARA
ActNN 34) stores 2-bit compressed activations instead of actual activations for all layers except for ReLU and pooling which only needs 1-bit sign activations. We applied ActNN (I.3) to all layers,
20

---

Published as a conference paper at ICLR 2023
which performs per-group quantization with fine- ane-igrieed mixed precision. The per-group quantiza- tion partitions the activation into several groups, and then quantization is performed by groups. The fine-grained mixed precision includes two processes: per-sample allocation and per-layer allocation. In per-sample allocation, the sensitivity, which is related to gradient variance, is calculated for each group and then an optimal number of bits for each group is determined using this sensitivity. Then, per-layer allocation calculates the sensitivity of each layer and then an optimal number of bits for each layer is obtained. When ARA and ActNN applied simultaneously, we perform 2-bit group- wise quantization and per-sample allocation on auxiliary residual activation. However, per-layer allocation is not applied to auxiliary residual activation because is not straightforward to calculate the sensitivity of ARA layers. While ActNN needs to quantize and store activations of all layers, we can skip this process for ARA Conv layers when using ActNN with ARA (Fig.
Quantization
Decuantzzatoon
Guaeniization
Dequentiztion
Linerr
*...*.**
Linar
Linar
Linear- ASA
Linear-
Linear- ASA
bit
bit
Memory
Memory
4++1
Kt-1
S1+2
D
Linear
Linear
Linear
Linear-
Linear-
Linear-
(a) Mesa
(b) Mesa with ASA
Figure 10: Training process of Mesa with ASA.
Pan et al. 36 pointed out that conventional activation compression algorithms do not consider heterogeneous activation distributions in a multi-head attention layer and therefore they are not applicable to Transformer-like models. Therefore, they proposed Mesa which performs head-wise quantization instead of per-tensor quantization. When we combine Mesa and ASA, we replace linear layers with ASA-Linear layers as shown in Fig. TO
H.3 DATASETS
CIFAR-10 and CIFAR-100 The CIFAR- dataset (46) consists 50,000 training images and 10,000 test images which are 32x32 RGB images for 10-class image classification. Likewise, the CIFAR-100 dataset (46) includes 50,000 training images and 10,000 test images with the same resolution for 100-class image classification.
Tiny ImageNet The Tiny-Imagell dataset (47) consists of images of 200 classes and each class has 500 images for training. It also contains 10,000 test images. All images included in Tiny ImageNet are selected from ImageNet and downsized to 64x64.
ImageNet The ImageNet dataset (41) is representative large-scale dataset for image classifica- tion. It consists 1,281,167 training images, 50,000 validation images, and 100,000 test images which are all RGB images. The average image size ImageNet is 469x387, but they are randomly cropped to 224x224 for image classification.
IWSLT 2016 The IWSLT 2016 dataset (49) is constructed for spoken language translation task. We choose IWSLT2016 De-En dataset which is used for German-to- Engo-sh translation. has 1,611 talks, 197K sentences, and 3.96M tokens training set while the validation set consists of 13 talks, 1.1K sentences, and 21K tokens.
The Microsoft Research Paraphrase Corpus (MRPC) dataset (50) consists of 5081 sen- tence pairs from newswire articles. this dataset, the training set contains 4076 sentence pairs with 2753 paraphrases and the test set contains 1725 pairs with 1147 paraphrases.
MRPC
21

---

Published as a conference paper at ICLR 2023
MNLI The Multi-Genre Natural Language Inference (MultiNLI) dataset (51) consists of 433k sentence pairs for ten genres (Face-to- tocsee,o-Fae,, Telephone, etc.). The dataset is divided into 392702 train sets, 9815 validation matched sets which subsets of train sets, and 9832 validation mismatched sets which are included in train sets.
H.4 HYPERPARAMETE
ResNet For training ResNet-18 from scratch on CIFAR-10, CIFAR-100, and Tiny ImageNet, we set the batch size and the total number of epochs 128 and 200, respectively. We applied stochastic gradient descent with momentum (53) along with weight decaying (46) to our experiments. The momentum and weight decay rate layers except for ASA layers was scheduled by cosine annealing (54) with 0.1 initial learning rate during 200 epochs. In comparison, we used 100x higher learning rate for ASA layers to make the magnitude of weight updates comparable to those of other layers. We also set in equations (15) and (I6) 0.01. We used nearly identical However, we set the batch size and the initial learning rate to 512 and 5e-2, respectively, and the number of epochs 90 with warm-up steps.
to 0.9 and le-4, respectively. The learning rate of the
hyperparameters for ImageNet experiments.
Transformer In Transformer, the batch size and the number of epochs are 4096 and 20, respectively. We trained Transformer from scratch on IWSLT and used the Adam optimizer (55) by setting 31 to 0.9, 32 0.98. and 1e-9. In addition, the learning rate was updated by the warm-up learning rate scheduler (54), and the number of warm-up steps was set to 4,000. We applied 100x higher learning rate to ASA layers as ResNet experiments. in equations (5) and (16) for ASA was set to 0.01. The model dimension was 512, the hidden layer dimension was 2048, the number of heads was 8, and the dropout rate (56) was 0.2. Furthermore, we applied label smoothing with a 0.1 label smoothing rate (57) to all experiments in Transformer.
BERT-Large We used Bert-Large (uncased) pretrained on BookCorpus (58) and English
Wikipedia for fine-tuning on MRPC and MNLI. The batch size was set to 32, and we trained BERT- Large for three epochs. The Adam optimizer (55) was used with 0.9 81, 0.999 82. 1e-6 and 0.01 weight decay rate. We also employed warm-up the learning rate up to 10% of the total step, and cosine annealing is used for the remaining epochs (54). The learning rate of cosine annealing was set to 1.5e-5. These learning rate was increased by 100x for ASA layers VIT experiments. e in equations [55 and (16) was to 0.01. The other hyperparameters were selected according to the original BERT paper (43).
ViT-Large We fine-tuned a ViT-Large model pretrained on ImageNet-21k: The model was fine- tuned using 32 batch size on CIFAR-100. We set the training batch size to 512 and experimented with different learning rates of 3e-2, le-1, and 3e-1 achieve the highest performance. For training, learning rate was adjusted by stochastic gradient descent with momentum (53) and warm-up cosine scheduler (54) was applied. The model was trained for 1000 steps with 100 warm-up steps. For
Linear-ASA layers, 100x larger learning rate was used. We also set in equations (15 and [6 to
0.01. The remaining hyperparameters were set according to the original VIT paper (44)
Mr.P-Mixer-Large The MIP-Mixer- Large model was fine-tuned from pretrained model on ImageNet-21k with 16 patch size. We trained the model on CIFAR-10 for 2000 steps with 56 batch size. Like ViT, the SGD optimizer (53) was applied to these experiments. We set the weight decay rate and momentum to 0 and 0.9, respectively. During the first 200 steps, the learning rate warmed up and cosine annealing (54) was employed in the remaining epochs. We set the learning rate 3e-2, le-1, and 3e-1. In the ASA layers, however, the learning rate increased by 100x like Transformer and ViT-Large. of ASA layers equations (15) and (16) was set to 0.01.The other hyperparameters followed the original ML.PMMxer paper (45)
H.5 DEVICES AND MEMORY MEASUREMENTS
In all experiments, perform most of the experiments, we used six GPUs to train ResNet on ImageNet to reduce training time. We calculated the amount of training memory using torch mach.cuda.mer memory allocated function
used Nvidia GeForce RTX 3090 GPUs. While GPU utilized to
22

---

Published as a conference paper at ICLR 2023
in PyTorch instead of nvidia-smi because midia smm includes unused memory space handled by the memory allocator in the repor
I MEMORY AND COMPUTE COMPLEXITY OF AUXILIARY ACTIVATION LEARNING
Table 13: The compute and memory complexity analysis.
BP
GCP
AAL only
AAL with GCP
Compute complexity Memory complexity
3N
4N
3N
4NV
N
2/N
N A
2VN A
In this section, we compare the compute and memory complexity of different learning algorithms. For the sake of simplicity, we assume that the computational cost of forward propagation and back- ward propagation of a single layer is 1. The amount of activations of a single layer is also assumed to be 1. Then we calculate the computational cost and memory space required for training a neural network with N identical layers. BP would need N computations for forward propagation, N com- putations for error propagation, and N computations for calculating weight updates, which results in 3NN computations in total. For memory space, BP has to store N activations generated during forward propagation. On the other hand, GCP requires N additional computations for the recom- putation of activations. However, it can reduce the memory complexity to 2vN when the network divided into VN blocks as suggested in (24). In our Auxiliary Activation Learning (AAL), we have to store auxiliary activations instead of actual activations. We assume that algorithm is applied to A layers in the network. Then, we have to store actual activation for the rest of the layers, which translates
memory complexity of N A. For the layers employing our algorithm, only have to store auxiliary activations and its amount is negligible compared to real activations. For example, Auxiliary Residual Activation (ARA) uses the previous layer's activation auxiliary activation, and hence there is no need to store this auxiliary activation since it is already stored in memory during forward propagation of the previous layer. Auxiliary Sign Activation (ASA) has to store 1-bit sign of actual activations, respectively. but they are also significantly smaller than real activations. Therefore, our algorithm has N A memory complexity and 3N compute complexity. To further reduce memory complexity, we can apply gradient checkpointing to the layers where we do not apply our AAL algorithm. Thus, the memory complexity those layers would decrease to 2 A and the compute complexity becomes 4N due recomputation.
J COMPARISONS TO ALGORIHTM TO APPROXIMATE BACKPROPAGATION
Previous studies have suggested various algorithms to approximate backpropagation 59 60 Our algorithm may be considered in line with such approaches as it replaces actual input activation with auxiliary activation. Furthermore, Uniform Sign-Concordant Feedback (61), which uses the sign of forward weights as feedback weights, looks similar to the concept of ASA which uses the sign of activation as auxiliary activation. However, the motivation quite different. In prior works on approximate backpropagation, the authors aim to remove the need for symmetric feedforward and feedback weights ('weight transport problem") by approximating the feedback path for better clopplussibllty. Contrarily, algorithm reduces memory requirements by using alternate form of activations, but it does not necessarily improve bio-plausibility as we still need symmetric forward and backward paths.
23

---

Published as a conference paper at ICLR 2023
K
CONVERGENCE SPEED OF AUXILIARY ACTIVATION LEARNING
ResNet152 ImageNet
Transformer on IWSLT
100
BP
ARA_Z.
ARA
APA (3, 4, 4,
30-
ARA
25-
Accuracy
20
Ulec
15
BP
10
ASA1
ASA3
ASA4
12
16
Epoch
Epoch
(a) ARA ResNet-152
(b) ASA Transformer
Figure 11: Learning curves of BP and Auxiliary Activation Learning.
24


Reviews:
[{"review_id": "p1UB3es_pds", "reviewer": "Reviewer_91fV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Training deep learning models is becoming increasingly challenging: there is simply not enough memory to store the activations generated by the forward pass until the corresponding gradients are computed. The paper proposes to split each activations in two tensors, $o$ and $a$. The $o$ tensor can be discarded quickly, while the $a$ tensor needs to be preserved in memory to compute the gradients.  The paper develops a theoretical criteria, the learning indicator, under which this partitioning scheme is guaranteed to succeed (subject to some reasonable assumptions). Based on this criteria, the authors design two schemes to partition the activations, auxiliary residual activation (ARA) and auxiliary sign activations (ASA), and empirically demonstrate their effectiveness on a comprehensive set of models.\n", "summary_of_the_review": " The memory bottleneck is an acute problem that impacts most of the deep learning community This paper provides a novel solution that alleviates this problem. Furthermore, this new approach can be effectively combined with existing solutions such as rematerialization to save more memory.\n", "strength_and_weaknesses": "Strengths:\n* The approach is compelling and completely novel as far as I know. It is effective on its own and can be combined with other existing approaches to save more memory.\n* The learning indicator criteria is sound under reasonable assumptions, and used effectively to develop two schemes, ARA and ASA.\n* The empirical evaluation of the scheme is comprehensive and demonstrates that the approach is applicable to a wide range of DNN architecture.\n\nWeaknesses:\n* There is a minor error in Theorem 1: $\\eta$ must be strictly positive for the loss function to converge\n* The convexity assumption under which the learning indicator criteria was developed often does not hold in many cases. That said, most of our theoretical understanding of deep learning suffers from the same limitation, and empirical proof is used to prove that this approach works in practice.\n* Table 1 doesn't list the training speed and memory reduction like the other tables do. It would have been great to be able to compare ASA and ARA side by side. I am especially curious to see whether there is a correlation between training speed and mean value of the learning indicator.\n* In table 2, the speed and accuracy of ARA is compared against that of GCP and ActNN. The comparison would be more fair if there was a way to decrease the compression rate achieved by GCP and ActNN to be in line with these of ARA. The same applies to table 3.\n* Table 4 and 5 propose to use the memory savings to increase the batch size, but this isn't very effective with less than 5% savings in the best case. I wonder if the authors could instead try to increase the width/depth the the neural networks to improve the accuracy of the predictions instead, which I suspect is what most people will use the memory savings for.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "N2NvHcA6FF", "reviewer": "Reviewer_zs6z", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a novel memory-efficient training method Auxiliary Activation Learning, which reduces the amount of data to be stored in memory without additional computations. The experimental results illustrate that Auxiliary Activation Learning can significantly reduce the memory cost on computer vision and NLP tasks without training speed reduction.   ", "summary_of_the_review": "NA", "strength_and_weaknesses": "Strength: \nThe paper proposes a novel algorithm to achieve memory-efficient training of neural networks. \nThe experiments on ResNet, ViT, MLP-Mixer illustrates the efficiency of the proposed method Auxiliary Activation Learning, which can significantly reduce the memory cost without sacrificing training speed. \n\nWeakness: \nSome memory-efficient algorithms are proposed in recent years, such as SM3. Some people also try to propose an efficient architecture to achieve memory-efficient training [1]. Maybe you can provide some introduction and experimental results about that.  \nNLP models are usually very large and memory-efficient training is also very important for them. I find you have reported the results about BERT-Large. It will be great if you can provide more results about that. \n\n[1] Yi Tay, Mostafa Dehghani, Dara Bahri, Donald Metzler. Efficient Transformers: A Survey. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "EAGwqLFgQ24", "reviewer": "Reviewer_YGpd", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new learning rule, named Auxiliary Activation Learning to reduce memory requirements of deep neural networks", "summary_of_the_review": "The paper seems to be well written an experimentation is exhaustive. The mathematical proof seems to be flawed as the convexity constraint would practically not apply to any deep neural network. Or maybe I misunderstood the message authors try to convey here.\nExperimental evaluation even is extensive misses the comparison with recent state-of-the-art work which seem to be better at reducing memory requirements and are better in accuracy.", "strength_and_weaknesses": "Strengths:\n- The proposed strategy seems to reduce memory space (11.8x) for ResNet, while not having a big influence in training time.\n- Extensive evaluation on training times and memory savings for Vision (ResNet) and NLP architectures (BERT, Transformers).\n\nWeaknesses:\n- “For mathematical analysis, we consider a loss function f(x) which is CONVEX, differentiable, and Lipschitz continuous.” Good luck with that. This is a very strong constraint and an unfeasible one when talking about Neural Networks.\n- It only compares with GCP (Chen et al., 2016) and ActNN (Chen et al. 2021), but it failed to compare to recent state-of-the-art methods, most notably Momentum ResNets (Sander et al., 2021). And looking at the Momentum ResNets paper it would seem to me that Momentum ResNets provide much less memory requirements and they provide higher accuracy than what is reported in table 1.\n- Memory savings are usually reported in parameter requirements, here it just says \"reduce memory space (11.8x) for ResNet\" and tables 2 and 3 show minimal memory reduction when compared to backpropagation (1.1x to 1.3x), such memory space is when compared to ActNN (much higher overhead).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "4_DW9GdnU-M", "reviewer": "Reviewer_Urez", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a training method that saves the memory that the activations take. Instead of using the originally dense activation, the authors propose to use the auxiliary activations during the backward process. During the forward pass, the originally dense activation will still be used so the precision is still high. The authors discuss two ways to construct the auxiliary activation: the residual and the sign. Experiments on image classification and NLP tasks show that the training time and the memory can be saved, while the performance on downstream task is not sacrificed too much. ", "summary_of_the_review": "This work proposes an effective way to efficiently train models with less memory cost, although the novelty is not outstanding in its current shape. ", "strength_and_weaknesses": "Strengths:\n- Efficient training is an important topic to study, and experiments in this paper have shown the proposed method can indeed achieve practical efficiency on training.\n- The presentation of this paper is overall clear and it is easy to follow and understand. \n- The proposed method is straight-forward and does not bring much overhead. \n\nWeakness:\n- The proposed technique is not essentially novel. As authors have already discussed, there are already many works trying to approximate the activations during training. Activation quantization and pruning, although argued by the authors to be slower which I do not believe, is a very simple way to achieve memory reduction. \n- (Minor) The baseline on ImageNet looks weaker compared to the commonly used one, which is around 76% (1% lower approximately). It might not be negligible at the scale of ImageNet.  \n- (Minor) $f$ is never assumed to be non-negative (although in common practice we use a non-negative one). It is only twice differentiable and convex.  ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster