\section{Introduction}\label{Introduction}


Transformer-based language models have achieved promising results on natural language understanding tasks, \emph{e.g.,}\xspace Q\&A~\citep{qu2019bert,yang2020bert}, relation extraction~\citep{xue2020embarrassingly, zhou2020document} and dialogue system~\citep{ni2021recent}. Recently, on vision tasks, transformers~\citep{dosovitskiy2020image,zhou2021deepvit, xue2021go, xue2022one} also outperform convolution-based models by a large margin. With sufficient training data,  transformer-based models can be scaled to trillions of trainable parameters~\citep{fedus2021switch, du2021glam}. Through scaling along width (\emph{i.e.,}\xspace hidden dimension) and depth (\emph{i.e.,}\xspace number of transformer blocks), these huge transformers show effectiveness across various tasks and even areas.





When we use transformer, we typically follow existing work to set width and depth. For instance, we usually set the width of transformer-base model as 768 and the depth as 12. An interesting question here is: \textit{Why do we select these fixed hyper-parameters, even for problems in different areas?} To answer this question, we revisit the conventional configurations from representative studies. For vision transformer~\citep{dosovitskiy2020image}, authors set the base ViT configuration according to those used in BERT~\citep{devlin2018bert}.  BERT selects such configuration following OpenAI GPT~\citep{radford2018improving}. OpenAI also follows the original transformer paper~\citep{vaswani2017attention}. In the original transformer paper, \citet{vaswani2017attention} conduct a set of ablation studies on machine translation task to find the optimal configurations. That is, for a good range of tasks, we have largely follow the transformer configuration based on an ablation study on machine translation task, \emph{i.e.,}\xspace a sequence-to-sequence task.


Nowadays, transformer-based models can be trained with various different training objectives or strategies. Taking vision transformer~\citep{dosovitskiy2020image} as an example, we can train transformer from scratch with a supervised learning setting for image classification.
In this straightforward image classification task, each image is modeled as a sequence of multiple tokens, and each token corresponds to a patch in the image. We use the  global information (from all tokens/patches of the image) to predict a single label, the category of the image. Here, as the training objective is to capture the global information of an image, the differences between token representations would not be considered directly. This image classification task is quite different from  machine translation task, which requests for strong understanding of each input token and generating another sequence. Hence, it is natural to assume that different optimal transformer configurations exist for these two different tasks. 

Previous work has tried to train a deeper transformer from scratch. However, as reported in~\citep{zhou2021deepvit,gong2021vision}, training by classification task (\emph{i.e.,}\xspace using the global signal of the input sequence) has the over-smoothing problem. That means, at the deeper transformer layers, all token representations tend to be identical. Such an issue harms the scalability of training vision transformer, especially scaling along depth. When we scale to a larger model, we only get a slight improvement or even worse accuracy. Recently, \citet{zhou2021deepvit, gong2021vision} show that, when we add special-designed regularization to avoid the ``uniform tokens'' (\emph{i.e.,}\xspace the over-smoothing problem), it is possible to train deeper transformer on the sequence (image) classification setting. 


Different from training from scratch, the masked autoencoder is a two-stage training framework, including pre-training and fine-tuning. Given partially masked input sequence, the pre-training stage aims to recover the original unmasked sequence. The fine-tuning stage is similar to the aforementioned training from scratch.
With masked autoencoder, recent studies~\citep{bao2021beit,he2021masked} successfully train large scale transformers,  even without using additional training data compared with supervised learning. This is counterintuitive because we usually assume more training data is the key in self-supervised pre-training to improve effectiveness. However, masked autoencoder is able to scale transformer to larger or deeper without the need of bigger dataset. This result motivates us to rethink the reason behind it. 

We hypothesize that masked autoencoder training can alleviate the over-smoothing issue, which is an important reason why masked autoencoder can help to scale transformer up. The main reason behind this is the different training objectives.
In masked autoencoder frameworks (\emph{e.g.,}\xspace BERT, BEiT), the target is to recover the masked tokens according to the unmasked tokens. Compared with training transformer from scratch under supervised learning, whose target is a simple classification task, masked autoencoder framework is using a sequence labeling target. The sequence labeling task requires the model to learn semantic information from neighboring unmasked tokens. Since different masked tokens have different unmasked neighboring tokens, the unmasked token representations must carry sufficient semantics for accurate prediction of the masked tokens, which in turn prevents the token representations to become identical. In a word, we can inference the masked autoencoder's training objective helps to alleviate over-smoothing problem by its regularization on token differences. To justify the reasoning above, we conduct both theoretical analysis and experimental verification. The results show that the over-smoothing issue is indeed alleviated in masked autoencoder. Compared with training under masked autoencder, training transformer by a simple classification task (\emph{e.g.,}\xspace training vision transformer from scratch) does not have such strength.



If the masked autoencoder well alleviates the over-smoothing issue, which is a challenge for scaling transformer along depth, does this mean the masked autoencoder prefers deeper models than wider ones for some tasks? We argue that we need to re-visit the configurations for different training objectives, especially for masked autoencoder. To this end, we conduct experiments to investigate the masked autoencoder configurations and propose our idea, Bamboo. Bamboo is, when we train transformer with masked autoencoder, we suggest to use deeper and narrower configurations to achieve better effectiveness with comparable computation cost as a typical setting. To evaluate our new configurations, we conduct comprehensive experiments on computer vision and natural language processing tasks. On vision tasks, we evaluate our configuration on large-scale vision transformer training. With Bamboo configurations, the masked autoencoders outperforms baselines by a large margin. For instance, on ImageNet, with comparable number of trainable parameters and computation cost, our narrower and deeper base-scale masked autoencoder, Bamboo-B, outperform MAE-B by 0.9\% in terms of top-1 accuracy. On natural language processing tasks, we conduct experiments on BERT. Results show that our configurations can improve BERT-L by 1.1 points on GLUE datasets.


In summary, our main contributions are three folds:
\begin{itemize}
\item  We propose the insight that masked autoencoder is better at scaling along depth. Training objective of masked autoencoder helps to handle transformer over-smoothing. We justify this reasoning by both theoretical analysis and experimental verification.

\item  We argue that the existing transformer configurations cannot fully use the strength of masked autoencoder. To this end, we propose Bamboo, an idea to scale transformer along depth with training with masked autoencoder. We show that narrower and deeper version show better performance than existing configurations, in a plug-and-play manner.

\item  We further verify our Bamboo configurations on larger scale vision transformer pre-training and language tasks. Results show that our Bamboo achieves state-of-the-art top-1 accuracy on image classification, and outperforms original BERT configurations by 1.1 points on GLUE.

\end{itemize}

The following sections are organized based on the analysis process of this work. In Section~\ref{Over-smoothing in Masked Autoencoder}, we briefly review the over-smoothing problem in transformer and show the strength of masked autoencoder in handling this issue theoretically and experimentally. We then conduct experiments to investigate scaling masked autoencoder along depth in Section~\ref{Bamboo}. Based on the consist sweet depth across scales, we suggest the Bamboo idea, 
using narrower and deeper configurations in masked autoencoder training. Then, we adapt the new configurations to larger scale, and conduct evaluations across different areas, vision tasks in Section~\ref{Evaluation Vision} and NLP tasks in Section~\ref{Evaluation Language}. Finally, we discuss the difference between this work and existing ones in Section~\ref{Discussion}.













\section{Over-smoothing in Masked Autoencoder}\label{Over-smoothing in Masked Autoencoder}

The over-smoothing issue is common in graph neural networks. When we stack many graph convolution networks, the node representations tend to be identical~\citep{chen2020measuring}. Recent studies show that transformer has a similar problem~\citep{shi2021revisiting}, known as ``uniform tokens'' in transformer models. In deep transformer, each token representation can be seen as a node in the graph, and each attention score is an edge. Therefore, different token representations  tend to be identical when we scale transformer to deeper. In this work, we mainly focus on the over-smoothing problem due to different training settings. 

\subsection{Theoretical Analysis}

We do theoretical analysis to compare the over-smoothing issue on sequence-level supervised learning, and masked antoencoder-based self-supervised learning, respectively. We use supervised vision transformer (ViT)~\citep{dosovitskiy2020image} and vision masked autoencoder (MAE)~\citep{he2021masked} as two representative platforms to show our insights in both theoretical and experimental analysis. 

Formally, transformer block's output (\emph{i.e.,}\xspace token representations) can be written as $h=\{h_0, h_1, ..., h_{T-1}\}$, where $h_t \in \mathbb{R}^{d}$, $d$ is the hidden dimension, $T$ is the number of tokens at different transformer blocks. Usually, $T$ is fixed across transformer blocks. For the token representations after $l^{\mathrm{th}}$ transformer block, we denote the token representations as $h^l$.

ViT is trained from scratch. The input is $x=\{x_0, x_1, ..., x_{T-1}\}$, and the label is $y \in \mathbb{R}^C$, where $C$ is the number of classes. We analysis the over-smoothing issue when we use either the [CLS] token or average token representations as the prediction head.


If we use [CLS] token to classify an input image, the training target is to minimize the distance between $f(h^L_0)$ and $y$, where $f$ is a linear projection. Observe that the final classification only relays on one token representation from the [CLS] token. More importantly, when over-smoothing happens, \emph{i.e.,}\xspace $h^L_0=h^L_1=...=h^L_{T-1}$, the prediction head remains meaningful for the task. If we use the average pooling prediction head, we reach the same conclusion. In this case, we are to minimize the distance between $f(\bar{h}^L)$ and $y$, where $\bar{h}^L$ is the average pooling vector from $L^\mathrm{th}$ layer token representations. If $h^L_0=h^L_1=...=h^L_{T-1}=\bar{h}^L$, the classification output remains the same. However, model effectiveness will be limited since less attention layers are effectively used.

Then, what would happen if we use masked autoencoder to pre-train? For masked autoencoder, the input is $x=\{x_0, x_1, ..., x_{T-1}\}$. If $x_t$ token is masked, we denote this masked token as $x^m_t$. The set of the masked tokens is $x^m$. The corresponding original token set before masking is $\tilde{x}^m$. During masked autoencoder pre-training, we are to minimize the distance between $h^m$ and $\tilde{x}^m$. We now state a key lemma to show that the masked autoencoder training target regularizes token representations. For simplicity, we treat each token representation as a scalar to analyze the variance $\mathrm{Var(\cdot)}$. The following results can be extended to the vector case directly.


\begin{lemma}[Variance of final token representations is bounded in masked autoencoder.]\label{lemma: variance_1}
{In masked autoencoder, given the variance of masked token representation $\sigma^2$, the variance of final token representations is bounded by the masked tokens:
\begin{equation}
    \mathrm{Var}({h^L}^m) \ge  \sigma^2
\end{equation}
}
\end{lemma}
The proof of the Lemma above is in Appendix~\ref{proof_1}. As the transformer can also be seen as a weighted multi-view graph, given the Lemma above, following the over-smoothing issue in graph neural network, we can also infer that the intermediate token representations are also bounded:

\begin{lemma}[Variance of intermediate token representation is bounded in masked autoencoder.]\label{lemma: variance_2}
{In masked autoencoder, given the  variance of intermediate token representations from $l^{\mathrm{th}}$ layer is bounded by the variance of  the token representation from ${l+1}^{\mathrm{th}}$ layer:
\begin{equation}
    \mathrm{Var}({h^l}^m) > \mathrm{Var}({h^{l+1}}^m) 
\end{equation}
}
\end{lemma}
The proof of the Lemma above is in Appendix~\ref{proof_2}. Given the theoretical analysis above, we can find both the intermediate and final token representations are bounded by the variance of the masked tokens, \emph{i.e.,}\xspace $\mathrm{Var}({h^l}^m)>\mathrm{Var}({h^{l+1}}^m)\ge  \sigma^2$. Therefore, the over-smoothing issue is alleviated in masked autoencoder.

In fact, revisiting the training of ViT from scratch again, Lemma~\ref{lemma: variance_2} also holds for the image classification setting. However, the difference is, Lemma~\ref{lemma: variance_1} cannot work on training ViT from scratch. Without the masked token modeling target in MAE, there is no bound to push last layer's token representations to be different. Since the Lemma~\ref{lemma: variance_2} still holds, the variance of token representations decreases layer-by-layer. When we scale transformer to deeper, over-smoothing happens.


We do not take residual connection into consideration in the two lemmas above. In reality, the residual connection is a key module in transformer training to alleviate over-smoothing. The residual connection would keep variance increasing along depth, although over-smoothing may still be observed~\citep{dong2021attention}. We thus, based on the two lemmas above, state the following theorem to show the variance in masked autoencoder increases faster than the variance in training transformer by simple classification setting. 

\begin{theorem}[Variance in masked autoencoder increases faster than the variance in training transformer by simple classification.]\label{lemma: variance_3}
{Given a converaged masked autoencoder $g_1(\cdot)$ and transformer trained by simple classification $g_2(\cdot)$ with residual connection, the increase in variance $\Delta \mathrm{MAE}=\mathrm{Var}(x+g_1(x))-\mathrm{Var}(x)$ is larger than, $\Delta \mathrm{ViT}=\mathrm{Var}(x+g_2(x))-\mathrm{Var}(x)$:

\begin{equation}
    \Delta \mathrm{MAE} > \Delta \mathrm{ViT}
\end{equation}
}
\end{theorem}
The proof of the theorem above is in Appendix~\ref{proof_3}.

In summary, when we train transformer on the supervised learning setting directly and treat the task as a simple classification task, the transformer does not need to consider the differences across token representations. Such behavior means there is no regularization to avoid over-smoothing. For masked autoencoder training framework, the training target is to recover the masked tokens base on the neighboring tokens.
Since the different masked tokens have different semantic-related neighbors, the different unmasked tokens must cover different semantics in deeper layers. Therefore, although we are not regularizing the transformer model in a straightforward way, the masked autoencoder training framework still pushes the token representations to be different.  


\subsection{Experimental Verification}


\begin{figure}[t]
\begin{minipage}[b]{0.4\linewidth}
\centering
\includegraphics[width=1.0\textwidth]{standard_deviation.pdf}
\subcaption{Average standard deviation}
\label{fig:over-smoothing-std}
\end{minipage}\hspace{5mm}
\begin{minipage}[b]{0.4\linewidth}
\centering
\includegraphics[width=1.0\textwidth]{cosine_similarity.pdf}
\subcaption{Average patch-pair cosine similarity}
\label{fig:over-smoothing-cos}
\end{minipage}
\centering
\vspace{-0.1cm}
\caption{Over-smoothing analysis of ViT and MAE, at different transformer layers.}
\label{fig:over-smoothing}
\vspace{-0.8cm}
\end{figure}

\textbf{Standard deviation.} To verify our reasoning above, we conduct two sets of quantitative analysis. For the first set of quantitative analysis, we focus on the standard deviation of token representations at different transformer layers. To compare the over-smoothing issue between supervised learning and masked autoencoder, we calculate the mean standard deviation of the token representations at different transformer layers:
\begin{equation}
\small
    {ms}^{i} =  \sqrt{\frac{1}{T-1} \sum\limits_{t=0}^{T-1}  (h_t^i - \bar{h^i})^2 }
\end{equation}
where $\bar{h^i}=\frac{1}{T-1}{\sum\limits_{t=0}^{T-1} h_t^i}$ is the mean vector of token representations at $i^{\mathrm{th}}$ transformer block. 
We use a deeper version of ViT-L and MAE-L as platforms to validate our analysis above.  Note that, the ViT-L is trained on the supervised learning setting and treats image classification as a simple classification task, by average pooling prediction head. The MAE-L is trained on the masked autoencoder setting, which means the pre-training target is closed to sequence labeling task. We then fine-tune the pre-trained model for image classification. We set the width as 768 and the depth as 48. The total number of trainable parameters and computation cost is comparable with the original configuration (\emph{i.e.,}\xspace width is 1024, depth is 24). The results are shown in Figure~\ref{fig:over-smoothing-std}. 

Observe that mean standard deviations of ViT-L and MAE-L both increase along depth. Such a result matches well with our theoretical analysis. The residual connection within transformer indeed alleviates over-smoothing (to some extent)~\citep{dong2021attention}. Nevertheless, the over-smoothing issue still exists even if we have residual connection as stated in~\citep{dong2021attention,shi2021revisiting}. We can observe that the increasing on MAE-L is much faster than that on ViT-L, which matches well with our Theorem~\ref{lemma: variance_3}. That means the deeper transformer blocks can learn different semantics for different token representations in MAE. In other words, the over-smoothing issue is alleviated in the transformer model pre-trained with the masked autoencoder setting. 

\textbf{Patch-pair cosine similarity.} To further verify that the over-smoothing issue can be alleviated in masked autoencoder, following \citet{gong2021vision}, we compare the patch-pair cosine similarity between ViT-L and MAE-L with deeper and narrower configuration. If there is an over-smoothing issue in the model, we should observe that the patch-pair cosine similarity increases along depth. The more serious the over-smoothing issue is, the faster cosine similarity increases. To remove the impact from input representations (residual connection), we use the zero-centered token representations $\tilde{h^i}= h^i - \bar{h^i}$ for evaluation. The results are shown in Figure~\ref{fig:over-smoothing-cos}. The cosine similarity of ViT increases along depth due to over-smoothing. However, for the model pre-trained by masked autoencoder framework, cosine similarity keeps constant along depth. This comparison is interesting, as we can barely observe over-smoothing issue on the model pre-trained by the masked autoencoder, even if we are using a deeper model than usual. 







\section{Bamboo}\label{Bamboo}

Our analysis above shows that the masked antoencoder  helps to alleviate the over-smoothing issue, which is a main challenge in scaling transformer along depth. Based on this finding, we propose Bamboo, an idea of using deeper and narrower transformer configurations, for masked autoencoder pre-training. Certainly,  deeper and narrower configurations do not always improve performance, as there are a few other reasons hindering scaling transformer along depth. In this section, we devote to find a sweet point following Bamboo idea, to achieve effective masked autoencoder by experiments.





\begin{table}[t]
\small
\caption{The configurations we used to scale transformer along depth. For a fair comparison, we keep a comparable computation cost with the original transformer configurations (\emph{i.e.,}\xspace depth=12 for Base, depth=24 for Large).}
\vspace{5pt}
\centering
\begin{tabular}{l llll llll}
\toprule 
Scale                  & \multicolumn{4}{l}{Base} & \multicolumn{4}{l}{Large} \\  \cmidrule(lr){2-5}\cmidrule(lr){6-9}
Depth             & 12   & 24   & 48   & 96  & 24    & 48   & 60   & 96  \\ 
Width             & 768  & 512  & 384  & 256 & 1024  & 768  & 640  & 512 \\
\#Attention Heads & 12   & 8    & 6    & 4   & 16    & 12   & 10   & 8   \\
Computation Cost  & 1$\times$    & 0.9$\times$  & 1$\times$    & 0.9$\times$ & 1$\times$    & 1.1$\times$  & 1$\times$    & 1$\times$  \\ \bottomrule
\end{tabular}
\vspace{-0.3cm}
\label{tbl:config-explore}
\end{table}






To validate our reasoning above, we conduct experiments on training transformer on ILSVRC-2012 ImageNet~\citep{deng2009imagenet} (ImageNet-1K), and report the top-1 accuracy. We use the original viT trained by supervised learning and MAE pre-trained masked image modeling as the test platform. For a fair comparison with the original transformer configurations, when we scale transformer to deeper, we also reduce the width to keep comparable number of trainable parameters and computation cost. The configurations we used are summarized in Table~\ref{tbl:config-explore}.


For supervised learning models (\emph{i.e.,}\xspace ViT), we train base scale models for 300 epochs and large scale models for 200 epochs following \citet{he2021masked}. We use RandAugment~\citep{cubuk2020randaugment}, drop path~\citep{huang2016deep}, mixup, cutmix~\citep{yun2019cutmix}, label smoothing~\citep{szegedy2016rethinking} for data augmentation.  Detailed hyper-parameters are summarized in Appendix~\ref{appendix:hyper-parameter}. For the masked autoencoder, we pre-train the base scale models for 1600 epochs and fine-tune for 100 epochs. For large models, we pre-train for 800 epochs and fine-tune for 50 epochs. We use the per-patch normalization following \citet{he2021masked} for better representations. 



\begin{figure}[t]
\begin{minipage}[b]{0.4\linewidth}
\centering
\includegraphics[width=1.0\textwidth]{base-sweet.pdf}
\subcaption{Scaling supervised ViT along depth}
\label{fig:sweet-base}
\end{minipage}\hspace{5mm}
\begin{minipage}[b]{0.4\linewidth}
\centering
\includegraphics[width=1.0\textwidth]{large-sweet.pdf}
\subcaption{Scaling MAE along depth}
\label{fig:sweet-large}
\end{minipage}
\centering
\caption{Scaling ViT and MAE with comparable number of parameters and computation cost.}
\label{fig:sweet}
\vspace{-0.6cm}
\end{figure}




The results of scaling to deeper transformers are summarized in Figure~\ref{fig:sweet}. For both the base-scale and large-scale models, we report the top-1 accuracy on ImageNet-1K dataset. We can observe the models trained by supervised image classification directly (\emph{i.e.,}\xspace ViT-B and ViT-L) cannot improve accuracy when we use the deeper architectures. For ViT-B, the top-1 accuracy on ImageNet-1K keeps comparable when the depth is smaller than 50, but after that, there is a significant accuracy drop. For ViT-L, the accuracy decreases even faster than the ViT-B. There is a significant accuracy drop when we use the narrower and deeper models. However, for masked autoencoders, we can observe totally different patterns. Even if we keep the comparable trainable parameters and computation cost, with such a simple modification, the masked autoencoders gain significant improvements. More importantly, when we scale to 48 layers, both 
MAE-B and MAE-L reach sweet points. When scaling to 96 layers, we observe the training is unstable compared with shallower models. This result matches well with our observation in Figure~\ref{fig:over-smoothing-cos}. At around 40th layer, the over-smoothing issue starts to happen slightly in masked autoencoder, which is much later than the ViT model. We suggest another reason of the unstable deep model training is the too large model updates~\citep{wang2022deepnet}. In this work, we focus on the over-smoothing issue of deep transformer training instead. The large model updates in deeper layers are out-of-scope. We leave that as our future work.


\begin{table}[t]
\small
\caption{Re-designed configurations under Bamboo idea.}
\vspace{5pt}
\centering
\begin{tabular}{l ll ll ll}
\toprule 
Scale                  & \multicolumn{2}{l}{Base} & \multicolumn{2}{l}{Large} & \multicolumn{2}{l}{Huge} \\  \cmidrule(lr){2-3} \cmidrule(lr){4-5}\cmidrule(lr){6-7}
                 & Original & Bamboo & Original & Bamboo & Original & Bamboo  \\ \midrule
Depth             & 12   & 48   & 24   & 48    & 32   & 64    \\ 
Width             & 768  & 384  & 1024  & 768 & 1280  & 896 \\
\#Attention Heads & 12   & 6    & 16    & 12   & 16    & 14    \\
Computation Cost  & 1$\times$    & 1$\times$  & 1$\times$    & 1.1$\times$ & 1$\times$    & 1$\times$   \\ \bottomrule
\end{tabular}
\label{tbl:bamboo}
\end{table}

According to the experimental results above, we find that masked autoencoder can indeed scale transformer well along depth. Even if we keep comparable trainable parameters and computation cost with the original transformer, the model achieves better accuracy. Another observation is, both transformer-base and transformer-large reach their sweet points at around 50 layers. We thus recommend a new set of transformer configurations in Table~\ref{tbl:bamboo} following our Bamboo idea, which are deeper and narrower than the original transformer configurations.





\section{Evaluation on Vision Tasks}\label{Evaluation Vision}


\subsection{Settings} 

We further evaluate our deeper and narrower transformer configurations on vision task. We train different models with more training epochs and compare with state-of-the-art vision models. We conduct experiments on ImageNet-1K and compare with recent supervised vision models \emph{e.g.,}\xspace DeepViT~\citep{zhou2021deepvit} and DeiT~\citep{touvron2021training},  and self-supervised vision models \emph{e.g.,}\xspace DINO~\citep{caron2021emerging}, MoCo v3~\citep{chen2021empirical}, BEiT~\citep{bao2021beit} and MAE~\citep{he2021masked}. Compared with the MAE, \textit{the only difference is the configurations}. That is, MAE uses original transformer configurations and we use our Bamboo configurations. The experiments are to verify that such a simple modification can improve model effectiveness. 

We evaluate the models on three different scales, \emph{i.e.,}\xspace base, large, and huge. The data augmentation setting is exactly the same as MAE for a fair comparison. Again, the only difference is that we use the Bamboo configurations instead of the original transformer configurations. During fine-tuning, we use the same script with training ViT from scratch. We fine-tune base models for 100 epochs. For large and huge models, we fine-tune them for 50 epochs following existing work~\citep{bao2021beit, he2021masked}. 


\begin{table}[t]
\vspace{-0.5cm}
\small
\caption{Top-1 accuracy on ImageNet-1K. We report two versions of ViT training from scratch. The first one is from original ViT paper~\citep{dosovitskiy2020image}, and the second one is from \citet{he2021masked}'s re-implementation with strong data augmentation. For MAE-B, we reproduce the results by running the official code and obtain a slightly different result (denoted by 83.3*). The original result is 83.6.}
\vspace{5pt}
\centering
\begin{tabular}{l llll}

\toprule 
Method             & Pre-train Data    & Base & Large & Huge \\ \midrule
ViT from scratch~\citep{dosovitskiy2020image}            & - & 77.9 & 76.5 & - \\
DeepViT~\citep{zhou2021deepvit}             & - & 80.9 & - & - \\ 
DeiT~\citep{touvron2021training}             & - & 81.8 & - & - \\
ViT from scratch~\citep{he2021masked}            & - & 82.1 & 81.5 & 80.9  \\ \midrule
DINO~\citep{caron2021emerging}             & IN1K & 82.8 & - & -  \\
MoCo v3~\citep{chen2021empirical}             & IN1K & 83.2 & 84.1 & -  \\ \midrule
BEiT~\citep{bao2021beit}             & IN1K + DALL-E & 83.2 & 85.2 & -  \\
MaskFeat~\citep{wei2021masked}             & IN1K & 84.0 & 85.7 & -  \\
IBOT~\citep{zhou2021ibot}             & IN1K & 84.0 & 84.8 & -  \\ \midrule
MAE~\citep{he2021masked} (Direct baseline)           & IN1K & 83.3* & 85.9 & 86.9 \\
Bamboo (Ours)             & IN1K & \textbf{84.2} & \textbf{86.3} & \textbf{87.1} \\\bottomrule
\end{tabular}
\label{tbl:exp-cv-sota}
\end{table}

\subsection{Results} 


Results on ImageNet-1K are reported in Table~\ref{tbl:exp-cv-sota}. For training ViT from scratch, we report the original results~\citep{dosovitskiy2020image} and the results with strong data augmentation~\citep{he2021masked}. A few recent work~\citep{zhou2021deepvit,touvron2021training} focusing on training ViT from scratch on ImageNet-1K are also included. In general, we can find the models pre-trained by self-supervised learning (\emph{e.g.,}\xspace DINO, MoCO v3, MAE) perform much better than training from scratch. If we only consider the self-supervised learning approaches, masked image modeling-based methods (\emph{e.g.,}\xspace BEiT, MaskFeat, IBOT, MAE) outperform the contrastive learning-based methods (\emph{e.g.,}\xspace DINO, MoCo v3) significantly, especially on larger scale. 

Since we are focusing on the scalability of training transformer with masked autoencoder, we choose MAE as our direct baseline, recent work with better scalability. We train MAE with Bamboo configurations and report the results in Table~\ref{tbl:exp-cv-sota}. Observe that our Bamboo achieves the best top-1 accuracy on all scales. On the base scale, Bamboo achieves state-of-the-art performance, 84.2 top-1 accuracy, which is 0.9 (0.6) points higher than MAE-B. When we scale the model up to large scale and huge scale, Bamboo remains the best performer and achieves 86.3 and 87.1 top-1 accuracy respectively. Note that, compared to other scales, the improvement is not so significant on the huge scale. One reason is, the original huge configuration has been deep (\emph{i.e.,}\xspace 32 layers), which is closer to the sweet point. Similarly, this can also explain why the configurations designed under Bamboo can improve MAE-B significantly.


\section{Evaluation on Language Tasks}\label{Evaluation Language}

\subsection{Settings} 

We further evaluate our Bamboo configurations on language tasks. We select BERT~\citep{devlin2018bert} as the platform to evaluate our Bamboo configuration because it is widely used on many language tasks. Note that BERT is a post-layer normalization transformer model. If our configuration shows better performance on BERT, that means it has good generalization performance.  We follow the BERT paper to use Wikipedia and bookscorpus to pre-train. The masked ratio is 15\% as by default. During pre-training, we use LAMB optimizer and set batch size and learning rate as 4096 and 1.76e-3, respectively, following~\citet{you2019large}.


During fine-tuning, we conduct experiments on General Language Understanding Evaluation (GLUE) benchmark.
The GLUE benchmark~\citep{wang-etal-2018-glue} is widely used in natural language understanding tasks, which include 8 tasks, \emph{i.e.,}\xspace CoLA, MNLI, MRPC, QNLI, QQP, RTE, SST-2 and STS-B. We set the learning rate as 1e-5 or 2e-5. Batch size is fixed as 32. For small datasets, \emph{i.e.,}\xspace CoLA, MRPC, RTE and STS-B, we fine-tune for 10 epochs. For larger datasets, \emph{i.e.,}\xspace MNLI, QNLI, QQP and SST-2, we fine-tune for 3 epochs. Matthew’s correlation is used as metric for CoLA. For MNLI, we report the average accuracy on MNLI-m and MNLI-mm. QNLI and RTE also adapt the accuracy as metric. The results on MRPC and QQP are reported with the average of F1 and accuracy. We use Spearman correlation on STS-B. For all fine-tuning results, we run the code 5 times and report the median.

\begin{table*}[t]
\small
\caption{Results of fine-tuning on GLUE benchmark. }
\vspace{-0.3cm}
\label{nlp:table:main}
\begin{center}
\begin{tabular}{l llllllll l}

\toprule
Method                   & CoLA & MNLI & MRPC & QNLI & QQP & RTE & SST-2  & STS-B & Avg \\ \midrule
BERT-B   & 59.6 & 83.7 & 88 & 90.4 & 88.8 &  68.6  & 91.5 & 89.4 & 82.5 \\ 
BERT-L   & 60.9 & 86.2 & 89.3 & 92.3 & \textbf{89.6} &  73.1  & 92.5 & 90.4 & 84.3 \\ \midrule
Bamboo-B   & 60.5 & 84.3 & 88.4 & 90.5 & 89.0 & 70.0 &  92.2  & 89.5 & 83.1 \\ 
Bamboo-L   & \textbf{62.9} & \textbf{87.1} & \textbf{89.8} & \textbf{92.4} & 89.4 &  \textbf{77.3}  & \textbf{93.8} & \textbf{90.6} & \textbf{85.4} \\ 
\bottomrule
\end{tabular}
\end{center}
\vspace{-0.6cm}
\end{table*}

\subsection{Results}

The results on language tasks are reported in Table~\ref{nlp:table:main}. Under the same pre-training and fine-tuning settings, models with Bamboo configurations outperform BERT by a large margin. Bamboo-L achieves the best performance in Table~\ref{nlp:table:main}. Compared with BERT-L, our Bamboo-L wins on 7 out of 8 datasets, and surpassed BERT-L by 1.1 points on average. It is also notable Bamboo configurations can outperform baselines on both large datasets (\emph{e.g.,}\xspace MNLI) and small datasets (\emph{e.g.,}\xspace CoLA).


\section{Discussion}\label{Discussion}

\textbf{Compared with simply scaling along depth}, this work keeps comparable computation cost and number of trainable parameters. When scaling along depth, we also make the deep transformer narrower. Under such setting, the deeper and narrower configurations re-designed under Bamboo idea can still outperform the baseline configurations, suggesting that we should consider narrower and deeper transformer when training by masked autoencoder. One related work is~\citep{tay2021scale}, an empirical study of practical scaling of transformer, which has a similar observation, deeper and narrower can improve accuracy. On this basis, our insight is, the insight in~\citep{tay2021scale} depends on different training objectives. 

\textbf{Compared with brute-force hyper-parameter tuning}, we provide both theoretical analysis and experimental verification to justify our insight, \emph{i.e.,}\xspace masked autoencoder alleviates over-smoothing issue in transformer. Motivated by this, we propose Bamboo by experimental investigation. There is no guarantee to ensure the configurations are always the optimal ones as different tasks may require different configurations. We believe it is impossible in deep learning to know the optimal configurations before experiments. However, we justify the tendency that masked autoencoder gets more benefits from deeper configurations. Our insight may instruct future works to consider configurations according to the training objectives.


\textbf{Instead of proposing a new approach or a new set of configuration,} this work focuses on an existing but neglected problem. After revisiting, we find configurations should be re-designed for different training objectives. We then re-design a set of configurations for masked autoencoder, which can improve transformer on both vision and language tasks. On vision tasks, we even achieve SoTA top-1 accuracy on ImageNet. However, we highlight our main contribution is an insight instead of a SoTA model, and this insight is orthogonal to the advanced masked antoencoder modifications in the future. 





\section{Conclusion}\label{Conclusion}

We first rethink the masked autoencoder's scalability in this work. Through theoretical analysis, we show that training transformer with masked autoencoder can alleviate the over-smoothing issue when scaling transformer along depth. We then verify this finding through experiments. The results show that, compared with supervised learning, model pre-trained by masked autoencoder can handle over-smoothing (\emph{i.e.,}\xspace uniform tokens) better. Under this insight, we rethink the widely used configurations in vision and language transformer, and suggest deeper transformer when training with masked autoencoder. We then explore this strategy on base and large scale vision transformer and re-design a set of vision transformer configurations for masked autoencoder empirically. To further verify the effectiveness of our configuration, we conduct experiments on larger scale vision transformer training and achieved SoTA top-1 accuracy with comparable computation cost on ImageNet-1K. To show our configurations are better than original transformer configurations, we further conduct experiments on language tasks. Compared with widely used language masked autoencoder, BERT, our model surpass it by 1.1 points on average on GLUE benchmark. In summary, our work provides an insight to explain why masked autoencoder can scale transformer along depth, and we propose Bamboo, a set of configurations to fully utilize such strength. 


In the future, we plan to integrate the existing methods focusing on scaling along depth, \emph{e.g.,}\xspace DeepNorm~\citep{wang2022deepnet} to further scale transformer. Another future direction is to verify our Bamboo configurations on other masked autoencoder-based models (\emph{e.g.,}\xspace BEiT~\citep{bao2021beit}, RoBERTa~\citep{liu2019roberta}).










   

   
















\printbibliography

\newpage

\section*{Checklist}




\begin{enumerate}


\item For all authors...
\begin{enumerate}
  \item Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
   
    \answerYes{}
  \item Did you describe the limitations of your work?
   
    \answerYes{}
  \item Did you discuss any potential negative societal impacts of your work?
   
    \answerNA{}
  \item Have you read the ethics review guidelines and ensured that your paper conforms to them?
   
    \answerYes{}
\end{enumerate}


\item If you are including theoretical results...
\begin{enumerate}
  \item Did you state the full set of assumptions of all theoretical results?
    \answerYes{}
  \item Did you include complete proofs of all theoretical results?
    \answerYes{}
\end{enumerate}


\item If you ran experiments...
\begin{enumerate}
  \item Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
   
    \answerYes{}
  \item Did you specify all the training details (\emph{e.g.,}\xspace, data splits, hyperparameters, how they were chosen)?
    \answerYes{}
    \item Did you report error bars (\emph{e.g.,}\xspace, with respect to the random seed after running experiments multiple times)?
    \answerYes{}
    \item Did you include the total amount of compute and the type of resources used (\emph{e.g.,}\xspace, type of GPUs, internal cluster, or cloud provider)?
    \answerYes{}
\end{enumerate}


\item If you are using existing assets (\emph{e.g.,}\xspace, code, data, models) or curating/releasing new assets...
\begin{enumerate}
  \item If your work uses existing assets, did you cite the creators?
    \answerYes{}
  \item Did you mention the license of the assets?
    \answerYes{}
  \item Did you include any new assets either in the supplemental material or as a URL?
    \answerYes{}
  \item Did you discuss whether and how consent was obtained from people whose data you're using/curating?
    \answerNA{}
  \item Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
    \answerYes{}
\end{enumerate}


\item If you used crowdsourcing or conducted research with human subjects...
\begin{enumerate}
  \item Did you include the full text of instructions given to participants and screenshots, if applicable?
    \answerNA{}
  \item Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
    \answerNA{}
  \item Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
    \answerNA{}
\end{enumerate}


\end{enumerate}





\newpage

