text:
\section{Local Average Window}

\section{Sparsity Ratio Schedule}
\label{app:schedule}

The cubic schedule of sparsity ratio is widely applied by many existing methods \citep{zhu2017prune,sanh2020movement,zafrir2021prune}, which includes initial and final warmups:
\begin{equation}\label{eq:cubic_schedule}
\rr^{(t)} = \begin{cases}
r^{(0)} & 0 \leq t<t_{i} \\
r^{(T)}+ \left(r^{(0)}-r^{(T)}\right)\left(1-\frac{t-t_{i}-t_{f}}{T-t_i - t_f}\right)^{3} & t_{i} \leq t<T-t_{f} \\
r^{(T)} &\text { o.w. }
\end{cases}
\end{equation}
Same as \citet{sanh2020movement}, we applies this cubic schedule to all baseline methods and our algorithm.

\section{GLUE Dataset Statistics}
\label{app:dataset-glue}

We present the dataset statistics of GLUE \citep{wang2018glue} in the following table.
\begin{table*}[htb!]
	\begin{center}
		\begin{tabular}{l|l|c|c|c|c|c}
			\toprule
			\bf Corpus &Task& \#Train & \#Dev & \#Test   & \#Label &Metrics\\ \midrule
			\multicolumn{6}{@{\hskip1pt}r@{\hskip1pt}}{Single-Sentence Classification (GLUE)} \\ \hline
			CoLA & Acceptability&8.5k & 1k & 1k & 2 & Matthews corr\\ \hline
			SST & Sentiment&67k & 872 & 1.8k & 2 & Accuracy\\ \midrule
			\multicolumn{6}{@{\hskip1pt}r@{\hskip1pt}}{Pairwise Text Classification (GLUE)} \\ \hline
			MNLI & NLI& 393k& 20k & 20k& 3 & Accuracy\\ \hline
			RTE & NLI &2.5k & 276 & 3k & 2 & Accuracy \\ \hline

			QQP & Paraphrase&364k & 40k & 391k& 2 & Accuracy/F1\\ \hline
			MRPC & Paraphrase &3.7k & 408 & 1.7k& 2&Accuracy/F1\\ \hline
			QNLI & QA/NLI& 108k &5.7k&5.7k&2& Accuracy\\ \midrule
			\multicolumn{5}{@{\hskip1pt}r@{\hskip1pt}}{Text Similarity (GLUE)} \\ \hline
			STS-B & Similarity &7k &1.5k& 1.4k &1 & Pearson/Spearman corr\\ \bottomrule

		\end{tabular}
	\end{center}
	\vskip -0.05in
	\caption{Summary of the GLUE benchmark.}
	\label{tab:glue}
\end{table*}

\section{Natural Language Understanding}\label{sec:app_nlu}

\subsection{Training Details}

{\bf Implementation Details.} The implementation of {PLATON} on BERT-base is based on publicly available \textit{MT-DNN} \citep{liu2019multi,mtdnn2020demo}\footnote{\url{https://github.com/microsoft/MT-DNN}} code-base. The implementation of DeBertaV3-base \citep{he2021debertav3} on GLUE is based on \textit{Huggingface Transformers}\footnote{\url{https://github.com/huggingface/transformers}} \citep{wolf2019huggingface} code-base.

{\bf Hyper-parameter Details.} We select $ \beta_{1} $ in range of $ \{0.75, 0.80, 0.85, 0.90\} $, find $ 0.85 $ generally perform best and fix it as $ 0.85 $ for all experiments. We select $ \beta_{2} $ from $ \{0.850, 0.900, 0.950, 0.975\} $. We choose AdamW as optimizer and select learning rate from $\{3\times 10^{-5}, 5\times 10^{-5}, 8\times 10^{-5}, 1\times 10^{-4}  \}$ and batch size from $\{ 8, 16, 32 \}$ and fix them for each dataset. Table~\ref{tab:app_glue_setup} lists the detailed setup of each dataset. For baseline methods, we set the hyper-parameters all as reported by \citep{sanh2020movement}.

We found that {PLATON} is not sensitive to its hyper-parameters $ \beta_{1} $ and $ \beta_{2} $. The performance of {PLATON} does not alter dramatically among different hyper-parameter setup.

\begin{table*}[h!]
\vspace{-1mm}
\caption{Hyper-parameter setup of {PLATON} for GLUE benchmark.}
\label{tab:app_glue_setup}
\begin{center}
\begin{small}
\begin{tabular}{l|c|cccccccc}
\toprule
\multirow{1}*{\bf Ratio} & {\bf Hyper-parameter} & {\bf MNLI} & {\bf RTE} & {\bf QNLI}  & {\bf MRPC} & {\bf QQP } & {\bf SST-2} & {\bf CoLA} & {\bf STS-B} \\
\midrule
~ & \# epochs &   8  &  20 & 10 & 10 & 10 & 6 & 15 & 15 \\
~ & Batch size &   32 &  16 & 32 & 8 & 32 & 32 & 32 & 16 \\
~ & Learning rate & $ 8\times 10^{-5}$ & $1\times 10^{-4}$ & $ 1\times 10^{-4} $ & $ 1\times 10^{-4} $  & $ 1\times 10^{-4} $ & $ 8\times 10^{-5}$ & $ 1\times 10^{-4} $ & $1\times 10^{-4}$ \\
\midrule
\multirow{5}*{20\%} & $ t_i $ & 5400 & 200 & 2000 & 300 & 5400 & 1000 & 500 & 500 \\
~ & $ t_f $ &  22000 & 1200 & 12000 & 900 & 22000 & 5000 & 1500 & 2500 \\
~ & $ \beta_{1} $ & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 \\
~ & $ \beta_{2} $ &  0.85 & 0.99 & 0.95 & 0.95 & 0.85 & 0.85 & 0.95 & 0.85  \\
\midrule
\multirow{5}*{15\%} & $ t_i $ & 5400 & 200 & 2000 & 300 & 5400 & 1000 & 500 & 500  \\
~ & $ t_f $ &  22000 & 1200 & 12000 & 900 & 22000 & 5000 & 1500 & 2500 \\
~ & $ \beta_{1} $ &  0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 \\
~ & $ \beta_{2} $ &  0.90 & 0.50 & 0.90 &0.95 & 0.90 & 0.90 & 0.975 & 0.90 \\
\midrule
\multirow{5}*{10\%} & $ t_i $ & 5400 & 200 & 2000 & 300 & 5400 & 1000 & 500 & 500  \\
~ & $ t_f $ &  22000 & 1200 & 12000 & 900 & 22000 & 5000 & 1500 & 2500 \\
~ & $ \beta_{1} $ & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 & 0.85 \\
~ & $ \beta_{2} $ & 0.85 & 0.95 & 0.95 & 0.95 & 0.90 & 0.85 & 0.95 & 0.95 \\
\bottomrule
\end{tabular}
\end{small}
\end{center}
\end{table*}

\section{Question Answering}\label{sec:app_squad}

\subsection{Dataset}
Following \citet{sanh2020movement}, we also choose SQuAD v1.1 \citep{rajpurkar2016squad} to evaluate the performance of {PLATON} on question answering task.

\subsection{Training Details}
We set the batch size as 16, the number of epochs for fine-tuning as 10, the optimizer as AdamW and the learning rate as $ 3\times 10^{-5} $ for all algorithms. We select the number of epochs as 10.
For {PLATON}, we set its cubic sparsity schedule as $ t_i = 5400 $ and $ t_f = 22000 $. The baselines are all configured as reported by \citet{sanh2020movement}. The other hyper-parameters are reported in Table~\ref{tab:app_squad_setup}.

\begin{table*}[htb!]
\vspace{-2mm}
\caption{Hyper-parameter setup of {PLATON} on question answering tasks (SQuAD v1.1, \citet{rajpurkar2016squad}).}
\vspace{3mm}
\label{tab:app_squad_setup}
\begin{center}
\begin{small}
\begin{tabular}{l|cccccc}
\toprule
\multirow{1}*{\bf Ratio} & {10\%}  & {15\%}  & {20\%}  & {30\%} & {40\%} & {50\%}  \\
\midrule
$ \beta_{1} $ & 0.85 &  0.85  & 0.85  & 0.85  & 0.85  & 0.85 \\
$ \beta_{2} $ & 0.975 & 0.975 & 0.975 & 0.975 & 0.900 & 0.975 \\
\bottomrule
\end{tabular}
\end{small}
\end{center}
\end{table*}

\section{Image Classification}\label{sec:app_vit}
For all image classification tasks we use the Vision Transformer model (ViT) \citep{dosovitskiy2020image}. We use the base model with an input patch size of $16 \times 16$. The ViT model is pre-trained on the Imagenet-21K dataset.

\subsection{Datasets}
We evaluate the {PLATON} method on two datasets, CIFAR-100 and the ILSVRC-2012 ImageNet dataset. The CIFAR-100 has 100 classes containing 600 images each. We use a image resolution of $224 \time 224$ for CIFAR-100. There are 500 training images and 100 testing images per class. The ILSVRC-2012 ImageNet dataset contains 1.2 million images spanning 1000 categories. We use a image resolution of $384 \times 384$ for ImageNet.

\subsection{Training Details}
{\bf Implementation Details.} The implementation of ViT follows the codebase of its pytorch version\footnote{\url{https://github.com/jeonsworld/ViT-pytorch}}. We use SGD with momentum \citep{qian1999momentum} as optimizer. For CIFAR100, we set batch size as 512 and learning rate as $0.03$. For ImageNet, we set batch size as 150 and learning rate as 0.03.

\textbf{Finetuning Details.} We fine-tune the ViT model with a learning rate 0f 0.03. We use a batch size of 512 for CIFAR-100 and 150 for ImageNet due to memory constraints. In addition, we use gradient clipping at norm 1. We also use a cosine decay for the learning rate with a warmup of 1000 steps.

\textbf{{PLATON} Details.}
We use the following hyperparameters for the {PLATON} algorithm. For the ImageNet dataset we use a $t_i = 7000$ step warmup period, a $t_f = 10000$ step final warmup period, and finetune for 60000 steps total. We set the exponential moving average parameter to 0.85.

For the CIFAR-100 dataset we use a 2000 step warmup period, a 8000 step final warmup period, and finetune for 20000 steps total. We set the exponential moving average parameter to 0.85.

\section{Background}\label{sec:preliminary}

We briefly review background and related works.

\subsection{Pre-trained Transformer-based Models}\label{sec:pertrained_models}
Pre-trained Transformer-based models \citep{devlin2018bert,liu2019roberta,brown2020language,dosovitskiy2020image,he2021deberta,he2021debertav3} have manifested superior performance in various natural language processing and computer vision tasks. These models are often pre-trained on enormous amounts of unlabeled data in a unsupervised/self-supervised manner such that they can learn rich semantic knowledge. By further fine-tuning these pre-trained models, we can effectively transfer such knowledge to benefit downstream tasks.

\subsection{Importance Scores for Pruning}\label{sec:weight_importance}

We denote $ \bm{\theta} = [\theta_1, \theta_2, \dots, \theta_d] \in \mathbb{R}^{d}$ parameters of a model, and further define $ \bTheta_{j,-j} = [0, \dots, 0, \theta_j , 0, \\ \dots, 0] \in \mathbb{R}^{d}$. The importance score of $\bm{\theta}$ is denoted as $ S(\bm{\theta}) \in \mathbb{R}^d$, where its $j$-th index $S_j(\bm{\theta})$ is the importance score for $\theta_j$. When the context is clear, we simply write $ S $.

{\bf Parameter magnitude} is an effective importance metric for model pruning \citep{han2015learning,han2015deep,paganini2020iterative,zhu2017prune,renda2020comparing,zafrir2021prune}. It defines the importance of a weight as its magnitude, i.e., $ S_{j} = |\theta_{j}|$. Parameters with small magnitude are pruned. Such a simple metric, however, cannot properly quantify a weight's contribution to the model output. This is because a small weight can yield a huge influence to the model output due to the complex compositional structure of neural networks, and vice versa.

{\bf Sensitivity of parameters} is another importance metric \citep{molchanov2019importance,sanh2020movement,liang2021super}. It essentially approximates the change in loss when a parameter is zeroed out. If the removal of a parameter causes a large influence on the loss, then the model is sensitive to it.
More specifically, the sensitivity of the $j$-th parameter in $\bm{\theta}$ is defined as the magnitude of the gradient-weight product:
\begin{equation}\label{eq:sensitivity}
	I_j = |\bTheta_{j,-j}^{\top} \nabla\cL(\bm{\theta})|.
\end{equation}
This definition is derived from the first-order Taylor expansion of $ \cL(\cdot) $ with respect to $ \theta_j $: $ I_j $ approximates the absolute change of the loss given the removal of $ \theta_j $,
\begin{align}\label{eq:sensitivity_taylor}
	\bTheta_{j,-j}^{\top} \nabla\cL(\bm{\theta}) \approxeq \cL(\bm{\theta}) - \cL(\bm{\theta} - \bTheta_{j,-j}).
\end{align}

\subsection{Iterative Pruning}

Iterative pruning methods \citep{han2015learning,louizos2017learning,sanh2020movement} have demonstrated to be effective in natural language processing and computer vision.
These methods prune model weights based on the ranking of their importance scores, where weights with small scores are pruned.
Specifically, at the $t$-th iteration, we first take a stochastic gradient descent step,
\begin{align*}
\widetilde{\bm{\theta}}^{(t)} = \bTheta^{(t)} - \alpha \nabla \cL(\bTheta^{(t)}),
\end{align*}
where $\alpha>0$ is the learning rate. Then, given the importance score $ \Sc^{(t)} $ (e.g., magnitude or sensitivity) for $\bm{\theta}$, the weights are pruned following
\begin{align*}
\bm{\theta}^{(t+1)} = \mathcal{T}(\widetilde{\bm{\theta}}^{(t)}, \Sc^{(t)}),
\end{align*}
where the $j$-th entry of $\mathcal{T}(\widetilde{\bm{\theta}}, \Sc^{(t)})$ is defined as
\begin{align} \label{eq:proj_function}
[\mathcal{T}(\widetilde{\bm{\theta}},S)]_j = \left\{
\begin{array}{lc}
\widetilde{\bm{\theta}}_j &\textrm{if}~S_j~\textrm{is~in~the~top~$r\%$ of $S$}, \\
0 &\textrm{otherwise}.
\end{array}
\right.
\end{align}
Here $r$ is a pre-defined parameter that determines the percentage of nonzero weights in the pruned model.

\section{Conclusion}
\label{sec:conclusion}

We propose {PLATON}, which is a model pruning algorithm that considers both sensitivity of model weights and uncertainty of importance estimation.
In {PLATON}, we use exponential moving average to smooth the sensitivity computed on mini-batches. Moreover, we quantify uncertainty of the importance estimation via the local temporal variation. These approaches effectively reduce variability of the importance score (smoothed sensitivity times uncertainty) and stabilize training.
We conduct extensive experiments on natural language understanding, question answering and image classification. Results show that {PLATON} significantly outperforms existing approaches. Our approach is particularly effective in the high sparsity regime.
For example, when pruning DeBERTaV3-base on the SQuAD v1.1 dataset, our method achieves a more than $5.0\%$ improvement when $90\%$ of the weights are pruned.
Moreover, training of {PLATON} is more stable and less sensitive to the hyper-parameters.

\section{Discussion}\label{sec:discussion}
Movement pruning \citep{sanh2020movement} is one of the most popular model pruning methods.
However, empirically we observe that it suffers from training instability and even divergence.
This might be due to the following reasons:

First, as mentioned earlier, the importance scores (sensitivity) estimated on mini-batches have high variability in movement pruning. This causes training instability as a weight may frequently alternates between being retained and being pruned. By considering uncertainty of the importance estimation, variability of importance scores in {PLATON} is drastically smaller (Figure~\ref{fig:sensitivity_distribution}, bottom), yielding a more stable training process and better model performance.

Second, a weight is masked instead of zeroed-out when it is deemed unimportant in movement pruning. Subsequently, if the rank of the weight's importance score raises in later iterations, it is unmasked and its value is restored.
However, the may be a number of iterations between a weight being masked and being reactivated, such that the weight's value becomes stale. In this case, such a stale weight can have a huge influence to the model, rendering training unstable.
In contrast, in {PLATON}, we let a weight starts from zero when reactivated.
This is natural because a pruned weight is still included in the training dynamics as a zero value.

\section{Experiments}\label{sec:experiments}
We implement {PLATON} for pruning pre-trained BERT-base \cite{devlin2018bert}, DeBERTaV3-base \cite{he2021deberta} and ViT-B16 \cite{dosovitskiy2020image} during fine-tuning.
We evaluate effectiveness of the proposed algorithm on natural language understanding (GLUE, \citealt{wang2018glue}), question answer (SQuAD v1.1, \citealt{rajpurkar2016squad}) and image classification (CIFAR100, \citealt{krizhevsky2009learning}; ImageNet, \citealt{deng2009imagenet}) tasks.
All the gains have passed significance tests with $p<0.05$.

\textbf{Implementation details.}
We use \textit{PyTorch} \citep{paszke2019pytorch} to implement all the algorithms.
Our implementation is based on the publicly available \textit{MT-DNN} \citep{liu2019multi,mtdnn2020demo}\footnote{\url{https://github.com/microsoft/MT-DNN}} and \textit{Huggingface Transformers}\footnote{\url{https://github.com/huggingface/transformers}} \citep{wolf2019huggingface} code-base.
All the experiments are conducted on NVIDIA V100 GPUs.

In Algorithm~\ref{alg:our_algorithm}, a weight may be zeroed-out (pruned) and reactivated in later iterations.
In this case, we let the weight starts from zero instead of from the value before zeroing-out.
This is because a pruned weight is still included in the training dynamics as a zero value, such that its is natural to restart it from zero.

We use a cubic schedule \cite{zhu2017prune,sanh2020movement,zafrir2021prune} to gradually increase the sparsity level during pruning. Please refer to Appendix~\ref{app:schedule} for details.

\vspace{0.05in} \noindent
{\bf Baselines.}
We compare {PLATON} with the following methods:

$\bullet$
\textit{$\ell_0$ regularization} \citep{louizos2017learning} is an effective modeling pruning method. The method adds a penalty to the proportion of remaining weights.

$\bullet$
\textit{Magnitude pruning} \citep{zhu2017prune} propose an automated gradual pruning method, which is a magnitude-based pruning method. The method enables masked weights to be updated, and has achieved state-of-the-art results among magnitude-based approaches.

$\bullet$
\textit{Movement pruning} \citep{sanh2020movement} is a state-of-the-art method for model pruning. The method applies an iterative pruning strategy, where sensitivity is used as the importance metric. The approach considers the changes in weights during training, and has achieved superior performance.

$\bullet$
\textit{Soft movement pruning} \citep{sanh2020movement} is a soft version of movement pruning based on the binary mask function in \citet{mallya2018piggyback}.

We use publicly available implementation\footnote{\url{https://github.com/huggingface/transformers/tree/master/examples/research_projects/movement-pruning}} to run all the baselines. Please refer to \citet{sanh2020movement} and references therein for more details.

\setlength{\tabcolsep}{0.38em}
\begin{table*}[t!]

	\caption{Results with BERT-base on GLUE development set. Here \textit{Ratio} is the proportion of remaining weights. Results with \emph{N.A.} indicate the model does not converge. The best results on each dataset are shown in \textbf{bold}. We report mean of $5$ runs using different random seeds.}
	\vspace{-3mm}
	\label{tab:glue_datasets}
	\begin{center}
		\begin{small}
			\begin{tabular}{l|l|cccccccc}
				\toprule
				\multirow{2}*{\bf Ratio} & \multirow{2}*{\bf Method} & {\bf MNLI} & {\bf RTE} & {\bf QNLI}  & {\bf MRPC} & {\bf QQP } & {\bf SST-2} & {\bf CoLA} & {\bf STS-B} \\
				~ & ~ & {m / mm} & {Acc} & {Acc} & {Acc / F1} & {Acc / F1} & {Acc} & {Mcc} & { P/S Corr} \\
				\midrule
				{\bf 100\%} & {$\text{BERT}_{\text{base}}$} & {84.6 / 83.4} & {69.3} & {91.3} & {86.4 /  90.3} & {91.5 / 88.5} & {92.7} & {58.3} & {90.2 / 89.7}
				\\
				\midrule
				\multirow{5}*{\bf 20\%} & {$ \ell_{0} $ Regularization} & {80.5 / 81.1} & {63.2} & {85.0} & {75.7 / 80.2} & {88.5 / 83.3} & {85.0} & \emph{N.A.} & {82.8 / 84.7} \\
				~ & {Magnitude} & {81.5 / 82.9} & {65.7} & {89.2} & {79.9 / 86.2}  & {86.0 / 83.8} & {84.3} & {42.5} & {86.8 / 86.6} \\
				~ & {Movement} & {80.6 / 80.8} & \emph{N.A.} & {81.7} & {68.4 / 81.1} & {89.2 / 85.7} & {82.3} & \emph{N.A.} & \emph{N.A.}
				\\
				~ & {Soft-Movement} & {81.6 / 82.1} & {62.8} & {88.3} & {80.9 / 86.7} & {90.6 / \textbf{87.5}} & {89.0} & {48.5} & {87.8 / 87.5}  \\
				~ & {PLATON} & {\bf83.1 / 83.4} & {\bf68.6} & {\bf90.1} & {\bf85.5 / 89.8} & {\bf90.7 / 87.5} & {\bf91.3} & {\bf54.5} & {\bf89.0 / 88.5}
				\\
				\midrule
				\multirow{5}*{\bf 15\%} & {$ \ell_{0} $ Regularization} & {79.1 / 79.8} & {62.5} & {84.0} & {74.8 / 79.8} & {87.9 / 82.3} & {82.8} & \emph{N.A.} & {81.8 / 84.2} \\
				~ & {Magnitude} & {80.1 / 80.7} & {64.6} & {88.0} & {69.6 / 79.4} & {83.6 / 79.2} & {82.8} & \emph{N.A.} & {85.4 / 85.0}
				\\
				~ & {Movement} & {80.1 / 80.3} & \emph{N.A.} & {81.2} & {68.4 / 81.0} & {89.6 / 86.1} & {81.8} & \emph{N.A.} & \emph{N.A.}
				\\
				~ & {Soft-Movement} & {81.2 / 81.7} & {60.2} & {87.2} & {81.1 / 87.0} & {90.4 / 87.1} & {88.4} & {40.8} & {86.9 / 86.6}  \\
				~ & {PLATON} & {\bf82.7 / 83.0} & {\bf65.7} & {\bf89.9} & {\bf85.3 / 89.5} & {\bf90.5 / 87.3} & {\bf91.1} & {\bf52.5} & {\bf88.4 / 87.9}
				\\
				\midrule
				\multirow{5}*{\bf 10\%} & {$ \ell_{0} $ Regularization} & {78.0 / 78.7} & {59.9} & {82.8} & {73.8 / 79.5} & {87.6 / 82.0} & {82.5} & \emph{N.A.} & {82.7 / 83.9} \\
				~ & {Magnitude} & {78.8 / 79.0} & {57.4} & {86.6} & {70.3 / 80.3}  & {78.8 / 77.0} & {80.7} & \emph{N.A.} & {83.4 / 83.3} \\
				~ & {Movement} & {79.3 / 79.5} & \emph{N.A.} & {79.2} & {68.4 / 81.2} & {89.1 / 85.4} & {80.2} & \emph{N.A.} & \emph{N.A.} \\
				~ & {Soft-Movement} & {80.7 / 81.1} & {58.8} & {86.6} & {79.7 / 85.9} & {\textbf{90.2} / 86.7} & {87.4} & \emph{N.A.} & {86.5 / 86.3}  \\
				~ & {PLATON} & {\bf82.0 / 82.2} & {\bf65.3} & {\bf88.9} & {\bf84.3 / 88.8} & {\bf90.2 / 86.8} & {\bf90.5} & {\bf44.3} & {\bf87.4 / 87.1} \\
				\bottomrule
			\end{tabular}
		\end{small}
	\end{center}

\end{table*}

\subsection{Natural Language Understanding}\label{sec:glue_expriments}
{\bf Models and Datasets.} We evaluate the pruning performance of BERT-base \citep{devlin2018bert} and DeBERTaV3-base \citep{he2021deberta} using the proposed algorithm.
We conduct experiments on the General Language Understanding Evaluation (GLUE, \citealt{wang2018glue}) benchmark.
The benchmark includes two single-sentence classification tasks: SST-2 \citep{sst2013} and CoLA \citep{cola2018}.
GLUE also contains three similarity and paraphrase tasks: MRPC \citep{mrpc2005}, STS-B \citep{sts-b2017} and QQP.
There are also four natural language inference tasks in GLUE: MNLI \citep{mnli2018}; QNLI \citep{squad1}; RTE \citep{rte1, rte2, rte3, rte5}; and WNLI \citep{wnli2012}.
Following previous works, we exclude WNLI in the experiments.
Dataset details are summarized in Appendix~\ref{app:dataset-glue}.

\vspace{0.05in} \noindent
{\bf Implementation Details.}
We select the exponential moving average parameters $ \beta_{1} $ from $\{0.75, 0.80, \\ 0.85, 0.90\}$ and $ \beta_{2} $ from $ \{0.850, 0.900, 0.950, 0.975\} $. We select the learning rate from $\{3\times 10^{-5}, 5\times 10^{-5}, 8\times 10^{-5}, 1\times 10^{-4}  \}$ and the batch size from $\{ 8, 16, 32 \}$. More details are presented in Appendix~\ref{sec:app_nlu}.

\vspace{0.05in} \noindent
{\bf Main results. }
We compare {PLATON} with the baseline methods under different sparsity levels: $90\%$, $85\%$ and $80\%$.
Table~\ref{tab:glue_datasets} shows experimental results on the GLUE development set. We see that {PLATON} achieves better or on par performance compared with existing approaches on all the datasets under all the sparsity levels.
For example, when the sparsity level is $80\%$, {PLATON} achieves a $83.1\%$ accuracy on the MNLI-m dataset, which is $1.5\%$ higher than the best-performing baseline (soft movement pruning).
The performance gain of our method is even more significant on small datasets. For example, {PLATON} outperforms existing approaches by more than $5\%$ on RTE when the sparsity level is $80\%$.
We remark that while movement pruning and its soft version perform well on large datasets (e.g., MNLI and QQP), it behaves poorly or even cannot converge on small datasets (e.g., RTE, CoLA and STS-B). We provide plausible explanations in Section~\ref{sec:discussion}.

Table~\ref{tab:deberta_experiemnts} summarizes experimental results on MNLI and SST-2 for pruning DeBERTaV3-base. Similar to Table~\ref{tab:glue_datasets}, {PLATON} significantly outperforms the baseline method on all the datasets under all the sparsity levels. From the results, we see that for the DeBERTa model, our method obtains more performance gain when the sparsity level is high. For example, in the $70\%$ sparse case, {PLATON} outperforms the baseline by $1.9\%$ ($94.6$ v.s. $92.7$) on the SST-2 dataset; while in the $90\%$ sparse case, our method achieves a $6.6\%$ improvement ($90.0$ v.s. $83.4$).

\setlength{\tabcolsep}{0.38em}
\begin{table}[htb!]

	\caption{Results with DeBERTaV3-base on SQuAD v1.1, MNLI and SST-2. Here \textit{Ratio} is the proportion of remaining weights. The best results on each dataset are shown in \textbf{bold}.}
	\vspace{-3mm}
	\label{tab:deberta_experiemnts}
	\begin{center}
		\begin{small}
			\begin{tabular}{c|c|ccc}
				\toprule
				\multirow{2}*{\bf Ratio} & \multirow{2}*{\bf Method} & {\bf MNLI} & {\bf SST-2} & {\bf SQuAD} \\
				~ & ~ & {m / mm} & {Acc}  & { EM / F1} \\
				\midrule
				{\bf 100\%} & {$\text{BERT}_{\text{base}}$} & {89.9 / 90.1}  & {95.6} & {84.6 / 92.0} \\
				\midrule
				\multirow{2}*{\bf 30\%} & {Magnitude} & {87.6 / 87.3} & {92.7} & {82.2 / 89.9} \\
				~ & {{PLATON}} & {\bf 88.9 / 88.8} & {\bf 94.6} & {\bf 83.1 / 90.9} \\
				\midrule
				\multirow{2}*{\bf 20\%} & {Magnitude} & {82.3 / 83.8} & {90.8} & {78.8 / 86.7} \\
				~ & {{PLATON}} & {\bf 87.2 / 87.0} & {\bf93.1} & {\bf81.9 / 89.8} \\
				\midrule
				\multirow{2}*{\bf 15\%} & {Magnitude} & {80.7 / 81.0} & {88.5} & {75.8 / 84.6} \\
				~ & {{PLATON}} & {\bf 85.8 / 85.9} & {\bf 92.3} & {\bf 81.2 / 89.0}  \\
				\midrule
				\multirow{2}*{\bf 10\%} & {Magnitude} & {77.1 / 78.0} & {83.4} & {70.5 / 80.5} \\
				~ & {{PLATON}} & {\bf 83.4 / 83.5} & {\bf 90.0} & {\bf 79.0 / 87.1} \\
				\bottomrule
			\end{tabular}
		\end{small}
	\end{center}
	\vspace{-3mm}
\end{table}

\subsection{Question Answering}

{\bf Models and Datasets.}
We evaluate performance of the proposed algorithm on a question answering dataset (SQuAD v1.1, \citealt{squad1}), where we use {PLATON} to prune BERT-base and DeBERTaV3-base.
Question answering is treated as a sequence labeling problem, where we predict the probability of each token being the start and end of the answer span. The dataset contains $88k$ training and $10k$ validation samples.

\vspace{0.05in} \noindent
{\bf Implementation Details.}  We set the batch size as $16$, and the number of epochs for fine-tuning as $10$.
We use AdamW \citep{loshchilov2017decoupled} as the optimizer and we set the learning rate as $ 3\times 10^{-5} $. Please refer to Appendix~\ref{sec:app_squad} for more details.

\vspace{0.05in} \noindent
{\bf Main Results. }
Table~\ref{tab:squad_experiemnts} summarizes experimental results, where we prune a pre-trained BERT-base model under $6$ different sparsity settings. From the results, we see that {PLATON} consistently outperforms existing approaches under all sparsity levels in terms of the two evaluation metrics: exact match (EM) and F1.
Notice that movement pruning and soft movement pruning are more efficient in the high-sparsity regime (e.g., $90\%$ sparse); while in the low-sparsity regime (e.g., $50\%$ sparse), these methods behave on par or worse than magnitude pruning and $\ell_0$ regularization. Our method, on the other hand, is effective under all sparsity levels.
Also note that {PLATON} is more effective in the high-sparsity regime. For example, in the $50\%$ sparse case, our method outperforms the best-performing baseline (magnitude pruning) by $0.2\%$ in terms of F1; while in the $85\%$ sparse case, {PLATON} achieves a $2.5\%$ gain.

Table~\ref{tab:deberta_experiemnts} summarizes results of pruning DeBERTaV3-base on the SQuAD v1.1 dataset. Similar to the findings in Table~\ref{tab:squad_experiemnts}, {PLATON} significantly outperforms the baseline method. Additionally, our method is also more effective in the high-sparsity regime when pruning DeBERTa models.

\begin{table*}[htb!]
	\caption{Results with $\text{BERT}_{\text{base}}$ on SQuAD v1.1. We report EM/F1. Here \textit{Ratio} is the proportion of remaining weights. The best results in each setting are shown in \textbf{bold}.}
	\vspace{-3mm}
	\label{tab:squad_experiemnts}
	\begin{center}
		\begin{tabular}{l|cccccc}
			\toprule
			{\bf Ratio} & {\bf 10\%} & {\bf 15\%} & {\bf 20\%} & {\bf 30\%} & {\bf 40\%} & {\bf 50\%} \\
			\midrule
			$\text{BERT}_{\text{base}}$ & \multicolumn{6}{|c}{80.4 / 88.1} \\ \midrule
			{$ \ell_{0} $ Regularization} & {68.9 / 80.0} & {70.7 / 80.9} & {72.0 / 81.9} & {73.1 / 82.8} & {74.3 / 83.9} & {75.1 / 84.6}  \\
			{Magnitude} & 67.7 / 78.2 & 73.4 / 82.9 & 75.9 / 84.8 & 77.4 / 86.5 & 77.9 / 86.7 & \textbf{78.5} / 87.0 \\
			{Movement}  & 71.9 / 81.9 & 72.1 / 81.8 & 72.3 / 82.0 & 71.6 / 81.9 & 72.7 / 82.8 & 73.4 / 83.0 \\
			{Soft-Movement}   & 71.0 / 81.0 & 72.2 / 82.2 & --    & 75.3 / 84.6 & --   & 77.0 / 85.8 \\
			\midrule
			{PLATON} & {\bf 74.2 / 83.8} & {\bf 75.9 / 85.4} & {\bf 76.8 / 86.1} & {\bf 77.5 / 86.7} & {\bf 78.0 / 86.9 } & {\bf 78.5 / 87.2}
			\\
			\bottomrule
		\end{tabular}
	\end{center}
	\vspace{-5mm}
\end{table*}

\subsection{Image Classification}
{\bf Models and Datasets.}
We apply {PLATON} to prune a ViT-B16 model \citep{dosovitskiy2020image}, and we evaluate model performance on two image classification datasets: CIFAR100 \cite{krizhevsky2009learning} and ImageNet \cite{deng2009imagenet}.

\vspace{0.05in} \noindent
{\bf Implementation Details.} We implement ViT using the following codebase\footnote{\url{https://github.com/jeonsworld/ViT-pytorch}}.
We use SGD with momentum \cite{qian1999momentum} as the optimizer. For CIFAR100, we set the batch size as $512$ and the learning rate as $0.03$. For ImageNet, we set the batch size as $150$ and the learning rate as $0.03$. Detailed settings are deferred to Appendix~\ref{sec:app_vit}.

\vspace{0.05in} \noindent
{\bf Main Results.}
Experimental results are summarized in Table~\ref{tab:vit_experiemnts}. We see that our method significantly outperforms existing methods on both the datasets (CIFAR100 and ImageNet) under all the sparsity levels. For example, the accuracy of {PLATON} is $76.3\%$ on ImageNet when the sparsity level is $80\%$, whereas the accuracy of magnitude pruning and movement pruning is only $66.5\%$ and $66.6\%$, respectively. Note that performance gain of our method is more significant in the high sparsity regime.

	\begin{table}[htb!]
		\caption{Results with ViT-B16 on CIFAR100 and ImageNet. Here \textit{Ratio} is the proportion of remaining weights. The best results on each dataset are shown in \textbf{bold}.}
		\vspace{-2mm}
		\label{tab:vit_experiemnts}
		\begin{center}
			\begin{tabular}{c|c|@{\hskip4pt}c@{\hskip4pt}@{\hskip4pt}c@{\hskip4pt}@{\hskip4pt}c@{\hskip4pt}|@{\hskip4pt}c@{\hskip4pt}}
				\toprule
				~ & {\bf Ratio} & {Magnitude} & {Movement} & {{PLATON}} & {\small ViT-B16} \\
				\midrule
				\multirow{3}*{\rotatebox{90}{\small CIFAR100}} & {30\%} & {88.9} & {78.5} & {\bf90.1} & \multirow{3}*{92.3} \\
				~ & {20\%} & {84.6} & {78.6} & {\bf87.3} & ~ \\
				~ & {10\%} & {64.3} & {77.5} & {\bf81.2} \\
				\midrule
				\multirow{3}*{\rotatebox{90}{\small ImageNet}} & {40\%} & {81.5} & {69.3} & {\bf82.6} & \multirow{3}*{83.5} \\
				~ & {30\%} & {78.9} & {68.6} & {\bf80.5} \\
				~ & {20\%} & {66.5} & {66.6} & {\bf76.3} \\
				\bottomrule
			\end{tabular}
		\end{center}
		\vspace{-5mm}
	\end{table}

\subsection{Analysis}

\vspace{-2mm}
\begin{figure*}[htb!]
	\centering
	\begin{subfigure}{0.32\textwidth}
		\centering
		\includegraphics[width=1.0\textwidth]{ratio-qqp.png}
		\vspace{-5mm}
		\caption{QQP}
	\end{subfigure}
	\begin{subfigure}{0.32\textwidth}
		\centering
		\includegraphics[width=1.0\textwidth]{ratio-mnli.png}
		\vspace{-5mm}
		\caption{MNLI-m}
	\end{subfigure}
	\begin{subfigure}{0.32\textwidth}
		\centering
		\includegraphics[width=1.0\textwidth]{ratio-squad.png}
		\vspace{-5mm}
		\caption{SQuAD v1.1}
	\end{subfigure}
	\vspace{-3mm}
	\caption{Model performance under different pruning ratio. We prune a BERT-base model during fine-tuning. Here the $x$-axis is the proportion of remaining weights.}
	\label{fig:pruning-ratio}
	\vspace{-2mm}
\end{figure*}

\textbf{Different levels of pruning ratio.}
Figure~\ref{fig:pruning-ratio} illustrates experimental results of pruning a BERT-base model during fine-tuning under different sparsity levels. We see that on all the three datasets (QQP, MNLI-m and SQuAD v1.1), {PLATON} achieves consistent performance improvement under all the sparsity levels compared with the baseline. Note that the performance gain is more significant when the sparsity level is high ($>50\%$). For example, {PLATON} outperforms the baseline by $0.4\%$ on MNLI-m when the sparsity level is $10\%$; while our method achieves a more than $3.0\%$ gain in accuracy when the sparsity level is $70\%$.

\vspace{0.1in} \noindent
\textbf{Variants of the importance score.}
Recall that in {PLATON}, the importance score is the product of sensitivity and uncertainty ($S=\overline{\I} \odot \overline{\U}$). In Table~\ref{tab:importance}, we examine variants of the importance score. From the results, we see that using only the sensitivity ($S=\overline{\I}$) or the uncertainty ($S=\overline{\U}$) deteriorates model performance by over $1.0\%$ on all the three datasets.
We also examine the case where $S=\overline{\I}/\overline{\U}$, i.e., we prune the weights with large uncertainty. In this case, the model fails to converge on CoLA, and the model performance drastically drops on the other datasets (i.e., $8.6\%$ on RTE and $9.6\%$ on SST-2). This indicates that weights with high uncertainty should be retained and further explored.

\begin{table*}[thb!]
	\centering
	\hspace{-4.5mm}
	\begin{minipage}{0.6\textwidth}
		\includegraphics[width=0.45\textwidth]{ablation-beta1.png} \hfill
		\includegraphics[width=0.45\textwidth]{ablation-beta2.png}
		\vspace{-0.1in}
		\captionof{figure}{Results on MNLI-m under different exponential moving average parameters. Left: $\beta_{1}$; Right: $\beta_{2}$. By default, we set $\beta_{1}=0.85$ and $\beta_{2}=0.85$.}
		\label{tab:ablation-beta}
	\end{minipage} \hspace{0.1in}
	\begin{minipage}{0.35\textwidth}
		\vspace{-5mm}
		\centering
		\captionof{table}{\small Variants of the importance score. We prune a BERT-base model.}
		\vspace{-2mm}
		\label{tab:importance}
		\begin{tabular}{l|ccc}
			\toprule
			& {\bf SST-2} & {\bf RTE} & {\bf CoLA} \\ \midrule
			{PLATON} & 90.5 & 65.3 & 44.3 \\ \midrule
			$S=\overline{\I}$ & 89.4 & 63.6 & 42.8 \\
			$S=\overline{\U}$ & 89.3 & 64.3 & 40.6 \\
			$S=\overline{\I} / \overline{\U}$ & 80.9 &  56.7 & \emph{N.A.}  \\
			\bottomrule
		\end{tabular}
	\end{minipage}

\end{table*}

\vspace{0.1in} \noindent
\textbf{Sensitivity to exponential moving average.}
Table~\ref{tab:ablation-beta} summarizes experimental results when we change the exponential moving average parameters $\beta_{1}$ and $\beta_{2}$.
In these experiments, we prune a pre-trained BERT-base model on the MNLI-m dataset.
From the results, we see that our method is robust to these two hyper-parameters.
In practice we fix $\beta_{1}=\beta_{2}=0.85$. Even when these selected hyper-parameters are not optimal, performance of {PLATON} is still better than the baseline methods.

\section{Extension to Structured Pruning}
\label{sec:structured}

Algorithm~\ref{alg:our_algorithm} targets on unstructured pruning \citep{lecun1990optimal,han2015learning,frankle2018lottery,sanh2020movement}, which prunes model weights entry-wise.
However, it is notoriously difficult to speedup inference of the pruned model, which requires low-level implementation that is not supported by existing deep learning libraries such as PyTorch \citep{paszke2019pytorch} and Tensorflow \citep{abadi2016tensorflow}.
In contrast, structured pruning \citep{mccarley2019structured,fan2019reducing,sajjad2020poor,lagunas2021block} methods remove columns or blocks of the weights, such that the pruned model can be easily compressed. As such, inference speedup via structured pruning is more implementation-friendly.

\subsection{Uncertainty Adjusted Structured Pruning}
Our method can be extended to structured pruning. In this case, the proposed importance score is computed group-wise instead of entry-wise (c.f. \eqref{eq:importance_score}).
Specifically, we divided the parameter $\bm{\theta}$ into $p$ disjoint groups denoted by
\begin{align*}
    \bm{\theta} = \{\bm{\theta}_{\mathcal{G}_1}, \bm{\theta}_{\mathcal{G}_2}, \dots, \bm{\theta}_{\mathcal{G}_p}\}.
\end{align*}
For example, for a weight matrix, $\mathcal{G}_j$ can denote its $j$-th column.
Then the sensitivity of $ \bTheta_{\mathcal{G}_{j}} $ is defined as
\begin{align}
    I_{\GG_{j}} = \left| \bTheta_{\mathcal{G}_{j}}^{\top} \nabla_{\GG_{j}} \cL(\bm{\theta}) \right|.
\end{align}
We apply exponential moving average to smooth the sensitivity of each group:
\begin{align*}
    \overline{\I}^{(t+1)}_{\GG_{j}} = \beta_{1}\Ibar^{(t)}_{\GG_{j}} + (1-\beta_{1}) \I^{(t)}_{\GG_{j}}.
\end{align*}
Similar to \eqref{eq:uncertainty} and \eqref{eq:ema_uncertainty}, we quantify the uncertainty of sensitivity of each group and then apply exponential moving average as the following:
\begin{align*}
    & \U^{(t)}_{\GG_{j}} = \left| \I^{(t)}_{\GG_{j}} - \Ibar^{(t)}_{\GG_{j}} \right|, \quad \overline{\U}^{(t+1)}_{\GG_{j}} = \beta_{2} \Ubar^{(t)}_{\GG_{j}} + (1-\beta_{2}) \U^{(t)}_{\GG_{j}}.
\end{align*}
The importance score is defined as (c.f. \eqref{eq:importance_score})
\begin{align*}
    &\Sc^{(t)}_{\GG_{j}} = \Ibar^{(t)}_{\GG_{j}} \cdot \Ubar^{(t)}_{\GG_{j}}.
\end{align*}
We prune the model using \eqref{eq:proj_function} in a group fashion, i.e., we zero-out the entire group $\bTheta_{\mathcal{G}_{j}}$ if $\Sc^{(t)}_{\GG_{j}}$ is not in top-$r$\%.

\begin{table}[h!]
	\vspace{-1mm}
	\caption{Results with $\text{BERT}_{\text{base}}$ on SQuAD v1.1. We report EM/F1. Here \textit{Ratio} is the proportion of remaining weights.}
	\vspace{-3mm}
	\label{tab:sturctured_experiemnts}
	\begin{center}
		\begin{tabular}{c|cc}
			\toprule
			{\bf Ratio} & {{PLATON}} & {{PLATON}\textsubscript{structure}}  \\
			\midrule
			{40\%} & 78.0 / 86.9 & 75.6 / 84.9 \\
			{50\%} & 78.5 / 87.2 & 77.0 / 85.9 \\
			\bottomrule
		\end{tabular}
	\end{center}
	\vspace{-5mm}
\end{table}

\subsection{Experimental Results}
We evaluate performance of the proposed structured pruning method by pruning BERT-base on SQuAD v1.1.
Table~\ref{tab:sturctured_experiemnts} summarizes experimental results.
It is expected that performance of structured pruning is lower than the unstructured alternative, because we have less control over individual weights. For example, important weights and unimportant ones may locate in the same column and are pruned together.
From the results, we see that performance of {PLATON}\textsubscript{structure} drops by about $1.5$ and $2.0$ when the sparsity levels are $50\%$ and $60\%$, respectively.

\section{Introduction}\label{sec:introduction}
Large Transformer-based models have exhibited superior performance in various tasks, such as natural language processing \citep{devlin2018bert,liu2019roberta,he2021deberta,radford2019language,brown2020language} and computer vision \citep{dosovitskiy2020image}.
These models advance start-of-the-art results in natural language understanding (e.g, GLUE, \citealt{wang2018glue}), question answering (e.g. SQuAD, \citealt{squad1}) and image classification (e.g., pImageNet, \citealt{deng2009imagenet}).
However, they typically incur massive memory footprint, e.g., BERT models \citep{devlin2018bert} contain up to 345 million parameters; Vision transformer models (ViT, \citealt{dosovitskiy2020image}) consist of up to 300 million parameters and GPT-3 models \citep{brown2020language} comprise up to 175 billion parameters.
Consequently, it becomes very challenging and expensive to deploy these models to real-world applications due to memory and energy consumption requirements. Moreover, these models' significant inference latency raises a barrier for their practicality, especially for high throughput environments such as e-commerce search platforms.

Confronting the abovementioned challenges, pruning methods are widely applied
at only a small expense of model performance. These methods prune the parameters based on certain importance scores to significantly reduce model sizes \citep{han2015learning,han2016eie,zhu2017prune}. Popular importance scores are based on the parameters' magnitude  \citep{han2015learning,han2015deep,paganini2020iterative} or sensitivity \citep{molchanov2016pruning,molchanov2019importance,ding2019global,sanh2020movement,liang2021super}. Parameters with low importance scores are considered redundant and expected to be pruned.

Existing pruning methods mainly fall into two categories: \emph{one-shot pruning} \citep{lee2018snip,frankle2018lottery,chen2020lottery,liang2021super,zafrir2021prune} and \emph{iterative pruning} \citep{han2015learning,zhu2017prune,paganini2020iterative,louizos2017learning,sanh2020movement}.
One-shot pruning specifies the sparsity pattern of a fully-trained dense model based on the weights' importance scores, and then trains a sparse model via ``rewinding'' (i.e., prune the fully-trained model and then re-train the pruned model). Such an approach has been widely utilized in exploring the Lottery Ticket Hypothesis \citep{frankle2018lottery,chen2020lottery,liang2021super}. However, determining the sparsity pattern based on fully-trained models fails to take the complicated training dynamics into consideration.
Some weights that are important to early stage of training do not necessarily have large importance scores in the fully-trained model, and therefore get pruned. Consequently, the rewinding process only yields models with worse generalization performance, especially in the high sparsity regime \citep{liang2021super}.

By contrast, iterative pruning conducts training and pruning simultaneously \citep{zhu2017prune}. Hence, the sparsity pattern can adapt to the parameter values and their dynamics throughout the entire training process. As such, iterative pruning attains significant improvement in the high sparsity regime over one-shot pruning \citep{sanh2020movement}.
Popular iterative pruning methods consider two approaches: sparsity regularization \citep{louizos2017learning,behnke2021pruning} and cardinality constraint \citep{han2015learning,sanh2020movement}. The former applies sparsity-inducing regularization (e.g., $\ell_{0}$ or $\ell_{1}$ regularization) to shrink certain weights toward zero values. The later truncates undesired weights according to the rank of their importance scores.

Unfortunately, iterative pruning methods \citep{han2015learning,louizos2017learning,sanh2020movement} suffer from a significant drawback: a weight's importance score cannot accurately reflect its contribution to model performance. This is mainly due to two types of uncertainty during training: (i) As the amount of data is quite large, the importance scores at each iteration are computed based on a randomly selected small batch of the training data, and therefore are subject to large variance due to stochastic sampling; (ii) The weights' importance scores can drastically vary due to the complicated training dynamics and optimization choices (e.g., dropout). Therefore, some weights can frequently alternate between being pruned and being activated, which causes training instability or even divergence.

\begin{figure*}[t!]
	\centering
	\includegraphics[width=0.98\textwidth]{our_importance_vs_sensitivity_distribution5.pdf}
	\vspace{-3mm}
	\caption{Variability of importance score when pruning BERT-base on RTE using iterative pruning. We compare two importance metrics: sensitivity in \eqref{eq:sensitivity} (top subfigure) and our proposed importance score in \eqref{eq:importance_score} (bottom subfigure). We calculate the importance scores every 10 steps and randomly sample 12 weights (one from each layer) for illustration purposes. Each violin plot corresponds to the importance scores' distribution of one sampled weight. Note that the plot is in \emph{log} scale. We remark the remaining weights behave similarly.
	}
	\label{fig:sensitivity_distribution}
	 \vspace{-2mm}
\end{figure*}

As an example, Figure~\ref{fig:sensitivity_distribution} (top) illustrates variability of the weights' importance scores (referred to as sensitivity in the figure). We see that during training, importance scores of the sampled weights vary by order of magnitudes. This indicates that some weights indeed alternate between being important (retained) and unimportant (pruned).

To resolve this issue, we propose \emph{{PLATON}} (\underline{P}runing \underline{LA}rge \underline{T}ransf\underline{O}rmer with u\underline{N}certainty), which prunes the weights by considering both their importance and uncertainty (explained below). For the weights with low importance but high uncertainty, our method tends to retain them. More specifically, {PLATON} contains two important components: (i) We consider exponential moving average of the importance scores, which is a smoothed version of the metric. This encourages the model to retain weight whose importance scores abruptly drop due to training instability; (ii) We quantify uncertainty of importance estimation by its local temporal variation, which is defined as the absolute difference between the importance score at the current iteration and the score's exponential moving average of the previous iterations.
Such quantification can be regarded as upper confidence bound (UCB) of estimated importance.
A large local temporal variation implies a high uncertainty in the importance score at the current iteration. Therefore, it is not a reliable indicator of importance, and we should retain the weight even though its score is relatively small.
In {PLATON}, we choose the product between the smoothed importance score and the local temporal variation as a new importance metric for the weights.
From Figure~\ref{fig:sensitivity_distribution} (bottom), we see that variability of the proposed importance metric is drastically smaller.

We conduct extensive experiments on a wide range of tasks and models to demonstrate the effectiveness of {PLATON}.
We evaluate the performance of our method using BERT-base \citep{devlin2018bert} and DeBERTaV3-base \citep{he2021debertav3} on natural language understanding (GLUE, \citealt{wang2018glue}) and question answering (SQuAD, \citealt{rajpurkar2016squad}) datasets.
We also apply our method to ViT-B16 \citep{dosovitskiy2020image} and evaluate the performance on CIFAR100 \citep{Krizhevsky09learningmultiple} and ImageNet \citep{deng2009imagenet}.
Experimental results show that {PLATON} significantly outperforms the baselines, especially under high sparsity settings (e.g., retaining only $10\%$ of the weights). For example, under $90\%$ sparsity on the MNLI dataset, our method achieves a $1.2\%$ accuracy improvement compared with state-of-the-art approaches.

\section{Method}\label{sec:method}
We present our algorithm -- {PLATON} (\underline{P}runing \underline{LA}rge \underline{T}ransf\underline{O}rmer with u\underline{N}certainty) to overcome the drawbacks of existing pruning methods. Our proposed algorithm contains two important components: (i) sensitivity smoothing using exponential moving average, which reduces the non-negligible variability due to training dynamics and mini-batch sampling. (ii) uncertainty qualification using local temporal variation, which drives the algorithm to explore the potentially important weights.

\paragraph{Sensitivity Smoothing}
In practice, the sensitivity in \eqref{eq:sensitivity} exhibits large variability because of two reasons.
First, $I_j$ in \eqref{eq:sensitivity} is computed batch-wise, i.e., in each iteration we randomly sample a mini-batch and compute sensitivity of the weights using this batch.
Second, the complicated training dynamics and optimization settings (e.g., dropout) further amplify such variability.
To mitigate this issue, we propose to smooth $ I_{j} $ by exponential moving average.
Concretely, at the $t$-th iteration, we have the smoothed sensitivity
\begin{equation}\label{eq:ema_sensitivity}
\Ibart_{j} = \beta_{1} \Ibar_{j}^{(t-1)} + (1-\beta_{1})\It_{j},
\end{equation}
where $ \beta_{1} \in (0, 1) $ is a hyper-parameter and  $\Ibar_{j}^{(0)} = 0~\textrm{for~all}~j = 1 \dots d.$

The smoothed sensitivity $ \Ibar_{j} $ can effectively reduce variability.
That is, for a weight $\theta_j$, an abrupt drop in its sensitivity $I_j$ causes limited impact on $\Ibar_{j}$.
Therefore, training using the smoothed sensitivity is stable because the pruning algorithm outputs a stable sparsity pattern.
Furthermore, the exponential moving average emphasizes recent sensitivity scores and drops the stale information, which further benefits training.

\paragraph{Uncertainty Quantification}
Besides sensitivity smoothing, we also directly consider the uncertainty of importance estimation to reduce variability.
Specifically, we quantify the estimation uncertainty by the sensitivity's local temporal variation, defined as
\begin{equation}\label{eq:uncertainty}
\Ut_{j} = |\It_{j} - \Ibart_{j}|.
\end{equation}
Similar to \eqref{eq:ema_sensitivity}, we further apply exponential moving averaging to $ \Ut_{j} $:
\begin{equation}\label{eq:ema_uncertainty}
\Ubart_{j} = \beta_{2} \Ubar_{j}^{(t-1)} + (1-\beta_{2}) \Ut_{j},
\end{equation}
where we set $\Ubar_{j}^{(0)}=0$ for all $j=1\cdots d$.
The uncertainty quantifies the variability by considering the difference between the current sensitivity and its historical average.
A large $ \Ubart_{j} $ indicates that the sensitivity $I_j^{(t)}$ computed using the current batch significantly deviates from the weight's historical sensitivity $\Ibar_{j}^{(t)}$.
This implies that there exists high uncertainty in $ \Ibart_{j} $, and hence, $\Ibar_{j}^{(t)}$ is not yet a reliable indicator of the importance of $\theta_j$.
In this sense, $\Ubart_{j}$ can be regarded as an upper confidence bound of estimated importance $\Ibart_{j}$.

\begin{algorithm}[htb!]
	\caption{{PLATON}}
	\label{alg:our_algorithm}
	\begin{algorithmic}[1]
		\STATE {\bfseries Input:} Dataset $\DD$; total training iterations $T$; exponential moving average parameters $\beta_{1}$ and $\beta_{2}$; remaining weights ratio $r$; learning rate $\alpha$.
		\STATE \textbf{Initialize:} $\overline{\U}^{(0)}=0$, $\overline{\I}^{(0)}=0$;
		\FOR{$ t = 0, \dots, T-1 $}
		\STATE Sample a mini-batch from $\mathcal{D}$;
		\STATE Compute the gradient $ \nabla \cL(\bTheta^{(t)}) $;
		\STATE Compute $ \It_{j} = |{\bm{\theta}}^{(t) \top}_{j,-j} \nabla \cL(\bTheta^{(t)})| $;
		\STATE Compute $ \Ibart_{j} = \beta_{1} \Ibar_{j}^{(t-1)} + (1-\beta_{1})\It_{j} $;
		\STATE Compute $ \Ut_{j} = |\It_{j} - \Ibart_{j}| $;
		\STATE Compute $ \Ubart_{j} = \beta_{2} \Ubar_{j}^{(t-1)} + (1-\beta_{2}) \Ut_{j}  $;
		\STATE Compute $ S^{(t)} = \Ibar^{(t)} \odot \Ubar^{(t)}$;
		\STATE Update $ \bm{\theta}^{(t+1)} = \mathcal{T} ( \bTheta^{(t)} - \alpha \nabla \cL(\bTheta^{(t)}) , S^{(t)})$;
		\ENDFOR
		\STATE \textbf{Output:} Pruned model $\bm{\theta}^{(T)}$.
	\end{algorithmic}
\end{algorithm}

\paragraph{Algorithm}

The most important difference between our proposed algorithm and existing ones is the importance score, which is defined as:
\begin{equation}\label{eq:importance_score}
\Sc^{(t)} = \Ibar^{(t)} \odot \Ubar^{(t)},
\end{equation}
where $ \odot $ is Hadamard product.
As can be seen from \eqref{eq:importance_score}, $\Ibar^{(t)}$ measures the sensitivity of weights while $ \Ubar^{(t)} $ quantities the uncertainty of sensitivity estimation.
The product yields a trade-off between $\Ibar^{(t)}$ and $\Ubar^{(t)}$. Specifically, when a weight $\theta_j$ has a high uncertainty $\Ubar^{(t)}$, even though its sensitivity $\Ibar^{(t)}$ at the current iteration is low, it may still significantly increase due to the high variability introduced by the mini-batch sampling and complicated training dynamics.
Therefore, we make a conservative choice and our proposed algorithm tends to retain it and explore it for a longer time.

\begin{wrapfigure}{r}{65mm}
    \centering
    \vspace{-8mm}
   	\includegraphics[width=0.42\textwidth]{sensitivity_uncertainty_cdf.pdf}
   	\vspace{-10mm}
    \caption{\small The cumulative distribution function (CDF) of $\Ibart_{j}$ and $\Ubart_{j}$ over $j$ at $t=500$ when pruning BERT-base on RTE. }
    \label{fig:cdf}
\end{wrapfigure}

We remark \eqref{eq:importance_score} shares the same spirit with UCB methods for bandit problems \citep{lai1985asymptotically,zhang2021biased}. Each parameter is considered as an arm, $\Ibart_{j}$ is regarded as estimated rewards and $\Ubart_{j}$ controls the level of uncertainty.
Since $ \Ibart_{j} $ and $ \Ubart_{j} $ are highly skewed to zero as shown by Figure~\ref{fig:cdf}, \eqref{eq:importance_score} applies a logarithmic transformation to make them distribute more evenly:
\begin{align*}
    \Sc^{(t)}_{j} = \exp\Big( \log(\Ibart_{j}) + c \log(\Ubart_{j}) \Big),
\end{align*}
where we let $c=1$. The above equation aligns to the policy of UCB. The upper confidence bound $\Ubart_{j}$ quantifies the uncertainty of importance estimation and promotes the exploration.

With our defined importance score, our algorithm prunes the model weights after each gradient decent step, i.e.,
\begin{equation}\label{eq:pruning_alg}
\bm{\theta}^{(t+1)} = \mathcal{T} ( \bTheta^{(t)} - \alpha \nabla \cL(\bTheta^{(t)}) , S^{(t)}),
\end{equation}
where $\mathcal{T}$ is defined in \eqref{eq:proj_function}.
We summarize our proposed algorithm \emph{{PLATON}} in Algorithm~\ref{alg:our_algorithm}.

redpajama_set_name: RedPajamaArXiv