Back Razor: Memory-Efficient Transfer Learning by Self-Sparsified Backpropogation
Abstract
Transfer learning from the model trained on large datasets to customized downstream tasks has been widely used as the pre-trained model can greatly boost the generalizability.
However, the increasing sizes of pre-trained models also lead to a prohibitively large memory footprints for downstream transferring, making them unaffordable for personal devices.
Previous work recognizes the bottleneck of the footprint to be the activation, and hence proposes various solutions such as injecting specific lite modules.
In this work, we present a novel memory-efficient transfer framework called Back Razor , that can be plug-and-play applied to any pre-trained network without changing its architecture.
The key idea of Back Razor is asymmetric sparsifying : pruning the activation stored for back-propagation, while keeping the forward activation dense.
It is based on the observation that the stored activation, that dominates the memory footprint, is only needed for backpropagation.
Such asymmetric pruning avoids affecting the precision of forward computation, thus making more aggressive pruning possible.
Furthermore, we conduct the theoretical analysis for the convergence rate of Back Razor, showing that under mild conditions, our method retains the similar convergence rate as vanilla SGD.
Extensive transfer learning experiments on both Convolutional
Neural Networks and Vision Transformers show that Back Razor could yield up to 97% sparsity , saving 9.2x memory usage, without losing accuracy.
The code is available at: https://github.com/VITA-Group/BackRazor_Neurips22 .
Introduction
Edge devices equipped with deep learning applications are becoming ubiquitous.
However, only deploying pre-trained models is not sufficient: these edge devices may keep collecting novel data that general pre-trained models have never seen before and fine-tuning on them would notably improve the model’s performance.
A typical approach is to upload these newly collected data to cloud servers for fine-tuning, and download the trained model from the cloud to local devices.
However, such a way has potential privacy issues due to data sensitivity and also imposes undesirable pressures on the Internet bandwidth, brought by transmitting both the collected data and trained models.
These challenges motivate researchers to explore on-device learning methods, i.e., how to fine-tune pre-trained models on memory-limited devices.
One of the main challenges of fine-tuning on edge devices is the memory constraints.
Figureexemplifies the memory usage of training some popular networks: ResNet50 with a small batch size of 16 can easily exceed the memory limitation of onboard devices.
The newly emerged Vision
Transformer (ViT) would incur even more memory consumption.
Recent studies show that the activations stored for back-propagation take up a large proportion of memory cost during training.
Therefore, popular model compression techniques such as weight pruning and quantization are not highly efficient in reducing training memory footprints.
Parameter efficient training method like Bitfit also fails to compress the activation to memory limitation (see Figure 1).
Gradient checkpointing can reduce the training memory by storing only a subset of activations, but requires more FLOPs at back-propagation due to the re-computation of discarded activations.
Using a smaller batch size or half-precision for training can also lower the memory cost, but they incur a decrease in the model’s accuracy and training speed.
[1] successfully avoided saving activation via freezing the weights for convolutional neural networks and only learning a newly added the lite residual module, thus no need to store the intermediate activations.
However, it comes at a certain price of the transfer performance.
[Figure]
A new stream of works studied how to directly compress the activations.
They mainly leverage the quantization technique that discretizes the activation tensors so these tensors can be represented by efficient data formats and require less memory [2, 3, 4].
However, the activation pruning techniques have been less explored.
[5] applied random activation masks on each layer in both forward and backward passes to reduce latency.
[6] integrated weight pruning and activation pruning by learning importance scores on neuron’s outputs.
Moreover, they adopt the symmetric scheme that prunes the activation at both the forward and backward processes.
Such a pruning scheme undermines the precision of models’ outputs, leading to significantly worse performance when the pruning ratio is high.
In this work, we present a general activation compression method called Back Razor .
Inspired by the observations that the activations are only stored for back-propagation, we adopt the asymmetric pruning scheme that only prunes the activation during the backward process.
Specifically, after activation is used by a subsequent layer, our method replaces it with a pruned activation which carries a lower memory cost.
During the backward phase, the gradients are calculated with the approximated activation, but only with a small error between those with the exact activation.
Our contributions are summarized as following:
• We propose a memory-efficient transfer learning method called Back Razor that prunes the activation for back-propagation.
Our proposed method can be easily applied to various backbones, including convolutional neural networks (CNNs) and vision transformers (ViTs).
• We provide a theoretical analysis of our self-sparsified back-propagation method for CNNs.
Specifically, we discover that under mild assumptions, our method can share a similar convergence rate as vanilla SGD.
• We extensively verify the proposed Back Razor on both CNNs and ViTs for multiple datasets and show that the proposed method can achieve up to 97% sparsity with saving 9 . 2 × memory saving without losing accuracy.
By further verifying on GPU, the proposed method achieves 2 .
7 × on-device memory efficiency for ViTs.
Related Works 2.1 Pruning and Sparse Training
Pruning can be operated at different levels, e.g. , pruning on weight, weight gradients, activations, and activation gradients.
Pruning on weights are most widely studied since it can accelerate the inference phase [7, 8, 9]; pruning on weight gradients can help reduce communication cost especially in
[Figure]
distributed training [10, 11]; pruning on activation gradients [12, 13, 14, 15] or activation itself have been less addressed since they are mostly dynamic during training.
[5] randomly generated static sparse masks and applied them on each layer in both forward and backward passes.
[6] proposed integral pruning that prunes both activation and weights by consecutively learning and applying two masks on weights and activations.
However these works are adopting the symmetric pruning scheme, i.e. , pruning the activation at both the forward and the backward process.
Pruning the activations at each layer brings errors, and they will accumulate as the forwarding proceeds.
SWAT [16] empirically explores sparsifying both weights and activations for training CNNs from scratch, and the authors also discovered that pruning activations only for backprogration can effectively preserve accuracy better. While our idea shares similarity to SWAT, Back Razor is dedicated to transfer learning and is applicable to general backbones (both CNNs and ViTs), and we also supply the theoretical analysis for such self-sparsified backpropogation for the first time.
2.2 Memory Efficient Transfer Learning
Transfer Learning from a model pre-trained on large datasets to small datasets has been widely used for many applications [17, 18, 19, 20].
In contrast to the pre-training that is always conducted on large clusters, transfer learning may happen on personal computing devices with limited memory constrains.
To fit the large model to small devices, a common choice is to fine-tune the last several layers of the model [21, 17, 22, 23].
However, it can be a sub-optimal choice, especially for the downstream tasks with large distribution differences.
In contrast, fine-tuning the whole can yield better accuracy in most cases [1, 24].
However, it can easily lead to memory outages.
On the other hand, recent parameter efficient transfer learning algorithms [25, 26, 27, 28] only require tuning less than 1% of the parameters.
But this parameter efficiency only brings a mild memory efficiency.
The most relevant prior work is TinyTL [1], it recognizes the main bottleneck of memory is in activation and introduces a new trainable module that only requires down-sampled activation.
However, it still focuses on exact gradient calculation while we employ approximate gradient calculation, which enables compressing activation without changing the model architecture.
A few other works have also explored the approximate activation techniques.
Most of them focus on activation quantization [2, 3, 29, 30].
Recently, activation compression is extended to ViTs [31] and Graph Neural Networks [32].
Method 3.1 Activation is the bottleneck of memory
The forward of neural networks can be express as following
[Equation]
where x and y is the input and label, respectively.
L is the loss between prediction f
( x ; θ ) and label y calculated by loss function L .
f i and θ i denote the function and parameters of i th layer, respectively.
The network is composed of l layers.
We further denote the output of i th layer (a.k.a Sample x , y from D Initialize z {+
+}
= x for i = 1 to l do ▷ Forward pass.
z i
=
f i ( z i − 1 , θ i ) ˜ z i − 1 ← TopK_Prune ( z i − 1 ) ▷ Sparsify z i − 1 and store it.
end for
Calculate the training loss L = L ( z l , y )
Initialize s l ← ∂L∂ z l for i =
l − 1 to 1 do ▷ Backward pass.
Calculate ∂L∂ z i : s i
← s i +1 h
( z ) i
(˜ z i , θ i ) ▷ Calculate the gradient for the next activation Calculate ∂L∂ θ i : s i +1 h
( θ ) i
(˜ z i , θ i ) and update θ i .
end for activation) as z i = f i ( z i − 1 , θ i ) and z 1 = f 1 ( x ; θ 1 ) . When conducting back-propagation, ∂L∂ z i and ∂L∂ θ i are required to be calculated for each activation.
Taking linear layer with z i
=
z i − [-
-]1 W i
+ b i ( θ i =
[ W i , b i ] ) as an example, the gradients required to be calculated are [1]:
[Equation]
As shown in Equation 2, the activation z
i is employed for computing one of the gradients.
This creates the need of saving the activation of the forward process.
Otherwise, extra FLOPs are required for re-computing the activation.
The size of activation is always larger than the weight as weight is always shared across different patches or tokens.
Also, the size of activation can increase linearly with the batch size.
Moreover, in the transformer, the length of input tokens can also lead to a quadratic increase of the memory.
Discussion about freezing weight.
TinyTL [1] pointed out that we can avoid saving activation via freezing the weight W i as other terms do not require the participation of z i .
This is true for most architectures of convolutional neural networks.
However, the emergence of ViT breaks this assumption.
Many operations (e.g. self-attention, Softmax, GeLU [33]) in transformer would involve the activation for computing ∂L∂ z i , which cannot be avoided if previous layers requires updating.
Therefore we adopt the more general method: directly sparsifying the activation for backpropagation.
3.2 Back-propagation activation self-sparsification
In contrast to the activation sparsification [5, 6] that prunes the activation of both forward and backward, our proposed Back Razor leverages the back-propagation activation sparsification: it only prunes the activation that is used for backward while keeping the forward activation dense as shown in Figure 2.
The normal backward process can be formally defined as
[Equation]
3) where h ( z ) i ( z i , θ i ) :
= ∂
z i +1 ∂
z i and h ( θ ) i
( z i , θ i ) :
= ∂
z i +
∂
θ
i , representing the derivatives of z i +1 with respect to z i
and θ i , respectively.
For Back Razor, the backward process can be expressed as
[Equation]
[Equation]
where ˜ z i denotes the pruned activation z i .
We inherit the mathematical forms behind h ( z ) i and h ( θ ) i from the normal backward to preserve the gradient flow, but calculate the gradients on weights and activations with the pruned layerwise activation ˜ z i .
With the above backward, there is no need for saving the dense activation z i .
Instead, we store a sparse version ˜
z
i that is more memory efficient.
The detailed optimization algorithm with the proposed Back Razor can be found in Algorithm 1.
In the forward pass, we would compute with dense activation while only saving its sparse version.
In the backward, we employ the sparse activation for computing the gradient.
We use a simple pruning method for backward activation pruning: prune the smallest magnitude activations.
The smallest magnitude values are usually playing less significant roles than those with larger magnitude.
This pruning method is also widely used for many pruning works [34].
In practice, we would prune all the activations to a fixed sparsity of λ by setting all the values under k th smallest value as zero, where k = λn , n is the total number of values for the activation.
As the ranking can be slow for large activations, instead of ranking the whole activation, we rank the activation of each sample and use the sample-wise threshold for pruning.
This largely shortens the ranking time.
Also, we noted that the thresholds across different samples are very close, which means it is functionally similar to global ranking.
We save the sparse matrix into two parts: a bitmap that indicates the position of the non-zero elements and a smaller dense matrix that contains the values of the non-zero elements.
The bitmap has the same shape as the origin sparse tensor, but it is much smaller since it is in bool type.
The memory cost of this format is n 8 + λnc type byte, where c type is the count of bytes for the employed data type.
For the previous activation pruning method, the error that comes from pruning could gradually accumulate layer by layer in the forward, causing the deviation from the pre-trained model and leading to performance degradation.
In contrast, the proposed Back Razor ensures the correctness of the forward.
Moreover, in the backward progress, the pruning error of many layers would not accumulate: the activations of those layers are only used for computing the gradient of itself (e.g. linear layer, batch/layer normalization layer).
It’s also worth noting that the proposed Back Razor saves the same amount of memory as activation pruning under the same pruning ratio.
Comparison with recent parameter-efficient transfer methods.
Recent studies reveal that finetuning can be parameter efficient.
For example, diffPrune [25] points out that the fine-tuning performance can match fully fine-tuning via only modifying 0.5% of the weights.
However, it employs Gumbel-Softmax [35] to make the masked parameters differentiable, which means the dense activations are required for updating the weight.
Therefore, it cannot yield memory efficiency.
On the other hand, other works also reveal that fine-tuning part of the model can yield competitive performance [26, 27, 28, 36].
Though it can lead to memory saving, we empirically verified that
Back Razor can yield a higher saving ratio.
3.3 Theoretical Results
Now we study the convergence of Back Razor.
We simplify the notion and define L to be the training loss of a convolutional neural network that takes the model’s weights as the input.
θ t = { W i } Li =1 as the flattened vector of the model’s parameters at step t , and g t be the gradient calculated with all the samples in the training dataset.
Furthermore, we denote the stochastic gradient by ˜ g t , which is the gradient calculated with one sample from the training set.
We make several classical assumptions [37]:
Assumption 1.
Let g ( θ ) be the gradient of the objective L at point θ , then for all x
and y we assume there exists a non-negative constant vector β that
[Equation]
Assumption 2.
For each iteration t , the stochastic gradient ˜ g t satisfies the following conditions:
[Equation]
where n is the number of parameters of models and σ 2 j is a constant representing the variance for coordinate j.
Assumption 3.
The objective L is bounded below by L ∗ .
These assumptions are discussed in Appendix A. We first prove the convergence for linear networks.
Lemma 1.
For linear neural networks, the gradient of parameters can decomposed by ˜ g t = ˜
z t ˜ A ⊤ t , and approximated gradient after activation pruning can decomposed by ˜ g ′ t = ˜
z ′ t ˜ A ⊤ t , where ˜ z t is the exact activation and ˜ z ′ t is the pruned activations.
In linear networks, ∂L∂ z i only depends on the model’s weights after the i th layer, so pruning the activation stored for backward does not change the its gradient.
Therefore in Lemma 1, ˜
A t is exclusively constructed by activations’ gradients, and not changed after activation pruning.
We call ˜
A t the transformation matrix for activations at step t .
Theorem 1 (Convergence) .
If Assumption 1-3 hold, ( p + αβ ) K a < 1 where p is the pruning ratio and K a is the squared condition number of the transformation matrix, then
[Equation]
where the T is the number of iterations, β = ∥ β ∥ ∞ , θ 0 is the initial weights and σ 2 = (cid:80) nj =1 σ 2 j .
The Eqn. 5 consists of two terms: the first term will converge to zero as the number of iterations T goes to infinity.
The second term will not vanish, but its value can be controlled by using larger batch size.
Intuitively, it proves that our algorithm reach to the neighborhood of a stationary point, and the radius of the neighborhood is bounded by the gradient variance.
Remark.
The key factor for theoretical convergence is ( p + αβ ) K a .
αβ can be reduced by using smaller learning rate.
The theoretical convergence will not be guaranteed when the transformation matrix is ill-conditioned.
We next discuss the convergence for non-linear networks, and focus on analyzing convolutional neural networks with a common CONV - BN - ReLU structure [15].
The activation pruning happens after the convolutional layer, where the gradient on weights and activations can be calculated as [15]:
[Equation]
(6) which suggests the gradient on z i will not be affected if we use the pruned activation ˜ z
i to obtain the derivatives on θ i when back-warding through the convolutional layer.
The activations are neither stored nor pruned for BN and ReLU operations, so the Lemma 1 holds for convolutional neural networks.
So under the same assumptions, we achieve similar convergence situation for CNNs.
We present more details in Appendix.
We lastly discuss the Transformer structures.
Unfortunately, the self-attention module does not meet the Lipchitz condition [38], therefore the above theoretical analysis will not apply out-of-the-box to them.
However we empirically observe that Back Razor is also applicable on ViTs [39] with decent performance (refer to Section 4.3).
A more rigorous theoretical framework customized for transformers is left as our future work.
Experiment 4.1 Settings
Following the common practice [25, 40, 41], we employ supervised pre-trained models on ImageNet as the starting point of transfer learning.
Specifically, we employ ImageNet-1K and Imagenet-22K for CNNs and ViTs, respectively.
For downstream fine-tuning, we consider eight datasets: Pets [42],
Aircraft [43], CIFAR10, CIFAR100 [44], Flowers [45], Cars [46], CUB [47], and Food [48].
The default image resolution is 224 × 224 following [1].
Our experiments are implemented with Pytorch [49] and conducted on 1080 Ti or [-V
-]V100 GPUs.
To measure the training memory, by default we report the theoretical memory usage following [1].
We also report the actual memory usage measured on the GPU at Section 4.3.
4.2 Back Razor with Convolutional Neural Networks
We choose ProxylessNAS-Mobile [51] as the CNN backbone following TinyTL [1].
It is worth noting that we keep the origin setting of ProxylessNAS-Mobile instead of modifying the architecture as in TinyTL given Back Razor is applicable for any architecture.
We also follow the most of training settings of TinyTL for fair comparison: The fine-tuning epochs and batch size is set as 50 epochs and 8, respectively.
The model is optimized with adam [52] optimizer and cosine learning learning rate schedule [53].
The initial learning rate is tuned for each dataset.
We freeze all the parameters of Batch Normalization layers [54] as they occupy more memory than convolutional layers while having much fewer parameters to update.
By freezing them, we could free the occupied memory with [Table] less influence on the fine-tuning performance.
We also use the origin activation for ReLU6, which is memory efficient (equal to the memory cost of a bitmap with the same shape).
The Back Razor is then applied to convolutional layers and the fully-connected classification head.
Comparison Results: In Table 1, we compare Back Razor with previous transfer learning methods on CNNs.
Specifically, it includes i)
FT-Last: fine-tuning the last fully connected layer of the network, which is also referred as linear classification head.
ii) FT-Norm+Last: fine-tuning the batch normalization layers plus the classification head iii) FT-Full: fine-tuning the full model.
Also, we include two different variants of the previous state-of-the-art (SOTA) method [1]: TinyTL and TinyTL@320 (TinyTL@320 employs a larger resolution image of 320).
Compared with FT-Full, Back Razor@ 90% use 8.7x times smaller memory while achieving comparable performance.
The variance for BackRazor@90% with respect to different random seeds is also comparable with FT-Full and they are both small.
In contrast, FT-Norm+Last takes 4.6x time more memory while being [0.7%, 19.5%, 2.1%, 3.1%, 2.6%, 12.2%, 4.7%, 6.2%] worse than
Back Razor@ 90% in [Pets, Aircraft, CIFAR10, CIFAR100, Flowers, Cars, CUB, Food], respectively.
Though FT-Last does yield slightly less training memory (31MB v.s. 42MB), the performance largely degrades: the accuracy is lower by [1.6%, 42.7%, 11.0%, 14.5%, 6.8%, 39.2%, 7.7%, 14.5%] than
Back Razor@ 90% in [Pets, Aircraft, CIFAR10, CIFAR100, Flowers, Cars, CUB, Food], respectively.
Compared with the previous State-of-The-Art method TinyTL, Back Razor@ 90% is comparable in terms of the training memory usage (37MB v.s 42MB).
Meanwhile, it yields higher accuracy of [1.1%, 12.2%, 1.0%, 1.9%, 1.4%, 5.1%, 3.9%, 3.5%] at [Pets, Aircraft, CIFAR10, CIFAR100, Flowers, Cars, CUB, Food], respectively.
When compared with TinyTL@320, which employs a larger resolution, Back Razor@ 90% saves 1.5 times of the memory while still yielding an improvement of [0.0%, 5.3%, 0.8%, 1.8%, 0.1%, 1.3%, 0.0%, 0.3%] at [Pets, Aircraft, CIFAR10, CIFAR100, Flowers, Cars, CUB, Food], respectively.
It’s also worth noting that Back Razor method can even achieve more aggressive sparsity.
For instance, the performance of Back Razor is [ 92 . 9% , 93 .
5% ] for pruning ratio of [ 90% , 95% , 97% , 99% ] at Pets, respectively.
The performance maintains the same level even with the pruning ratio of 97% and only marginally drops at the pruning ratio of 99% .
A similar trend is also observed in other datasets.
However, the high pruning ratio does not bring further memory reduction since the main bottleneck of memory is no longer the accumulated activations.
Instead, it is the model parameters (of 11.6MB) and one large forward activation (of 24.5MB).
Can we prune at forward pass?
We further compare the performance difference between activation pruning (prune both forward and backward) with Back Razor (only prune backward) in Figure 3.
We implement the activation pruning via directly applying the mask of the Back Razor to the forward pass.
It can be observed that the model can easily collapse (the model can hardly converge and only have accuracy below 5%) for activation pruning even with a mild pruning ratio of 70%.
In contrast, for Back Razor, the performance can be at a high level even with an aggressive pruning ratio of 99% ( the performance is even higher than activation prune at a sparsity of 50%).
This demonstrates that the backward activation can be more sparse than that of the forward. Pruning Ratio A cc u r acy M e m o r y
U sa g e // 0 . [Figure] [Figure] Aircraft [43] 0.
Pruning Ratio A cc u r acy M e m o r y U sa g e // 0 .
[Figure] [Table] Convergence speed in practice: As shown in Figure 4, the convergence speed in transfer experiments is comparable with each other. In the early stage of the training, the convergence speed of Back Razor is even faster.
Even the last stage, FT-Full is only marginally better than Back Razor.
Comparison with SWAT: We further compare with SWAT [16] which leverages the similar insight on the sparsity of backward activation.
Although [16] was originally developed for training from scratch, here we apply it to the identical fair setting of transfer learning .
As shown in Table 2, while similar memory footprint is employed for Back Razor@ 90% and SWAT@ 80% , the proposed BackRazor can yield a higher performance by 0 .
6% and 2 .
6% on CIFAR10 and CIFAR100, respectively.
[Figure] [Table] 4.3 Back Razor with Vision Transformer
In this section [Table] we study Back Razor on ViT [39].
The patch size is by default set as 16.
We employ the standard SGD optimizer with cosine learning rate decay for finetuning.
The training steps are fixed as 20k and the initial learning rate is tuned for each dataset.
We employ a larger batch size of 128 following the common practice for accelerating training [39, 55, 31].
Comparison Results: In Table 3, we compare Back Razor with previous transfer learning methods on ViT, including FT-Last and FT-Full same as defined in Section 4.3.
We further compare two parameter efficient transfer learning method from the literature: Bitfit [26], and the SOTA ViT transfer leanring method, Mesa [26] that is based on backpropagation quantization.
[Table]
Compared with FT-full, Back Razor@ 80% yields comparable performance across the datasets while achieving 3 . 9 × memory saving.
Though FT-Last employs smaller train memory than Back Razor, the accuracy drops by [2.2%, 32.9%, 2.1%, 6.4%, 0.1%, 26.6%, 0.9%, 4.2%] in [Pets, Aircraft, CIFAR10, CIFAR100, Flowers, Cars, CUB, Food], respectively, compared to Back Razor@ 80% .
Bitfit achieves competitive performance via tuning less than 0.1% of the parameters.
However, it can only reduce the memory saving of 2 .
× memory given it requires saving some large activations in full (e.g. attention map).
In contrast, Back Razor can reduce the size of any activations.
While Mesa has similar memory usage and performance compared with Back Razor, we note that its methodology is orthogonal with Back Razor. By combining Meta and Back Razor together, we can achieve 7 . 5 × memory saving with comparable performance to FT-Full.
When employing a larger pruning ratio for BackRazor in ViTs, more memory footprint savings can be achieved at little performance degradation.
Especially, for CIFAR10, the performance of BackRazor only drops by 0.2% compared to FT-Full at a pruning ratio of 95%.
4.4 Back Razor with More Tasks
In this section, we further exploring applying the proposed BackRazor to more downstream tasks.
Semantic Segmentation: As shown in Table 5, when applying to semantic segmentation tasks, BackRazor can yield 5 .
3 × memory efficiency at a pruning ratio of 90% with only a marginal accuracy drop of 0.6% compared to the fully fine-tuning baseline.
Language modeling: As demonstrated in Table 6, when applying Back Razor on the language model,
Back Razor can achieve a significant memory efficiency improvement at a pruning ratio of 90% with only a marginal accuracy drop of 0.6%.
[Table]
[Table]
In this work
[Table]
we propose Back Razor
[Table]
a memory-efficient transfer learning method, which is also the first asymmetric sparse transfer method that only sparsifies the back-propagation activations while keeping the forward.
Through extensive experiments in both Convolutional Neural Networks and
Vision Transformer across multiple datasets
[Table]
we demonstrated the proposed Back Razor can achieve state-of-the-art memory efficiency with an aggressive pruning ratio.
Moreover, it can be combined with the propagation quantization method and yield better memory efficiency.
As limitations, Back
Razor has only been tested on a limited range of tasks
[Table]
and its theoretical underpinning is so far restricted to CNN backbones despite its empirical success in ViTs, calling for continued efforts
[Table]
Acknowledgment
Z. Wang is in part supported by the National Science Foundation under Grant IIS-2212176, and a
Google TensorFlow Model Garden Award.