Benchmarking Our Baseline Model

Now that our PerformanceBenchmark is complete, let’s give it a spin! For the baseline model we just need to pass the pipeline and dataset we wish to perform the benchmark on, and we’ll collect the results in the perf_metrics dictionary to keep track of each model’s performance:

pb = PerformanceBenchmark(pipe, clinc["test"])
perf_metrics = pb.run_benchmark()

Model size (MB) - 418.17
Average latency (ms) - 46.05 +- 10.13
Accuracy on test set - 0.867

Now that we have a reference point, let’s look at our first compression technique: knowledge distillation.

Knowledge Distillation for Fine-tuning

So how is knowledge actually “distilled” or transferred from the teacher to the student during training? For supervised tasks like fine-tuning, the main idea is to augment the ground truth labels with a distribution of “soft probabilities” from the teacher which provide complementary information for the student to learn from. For example, if our BERT-base classifier assigns high probabilities to multiple intents, then this could be a sign that these intents lie close to each other in the feature space. By training the student to mimic these probabilities, the goal is to distill some of this “dark knowledge”⁶ that the teacher has learnt; knowledge which is not available from the labels alone.

Mathematically, the way this works is as follows. Suppose we feed an input sequence $x$ to the teacher to generate a vector of logits $?\left(x\right)$ = [$z_1\left(x\right),...,z_N\left(x\right)$]. We can convert these logits into probabilities by applying a softmax function

but this isn’t quite what we want because in many cases the teacher will assign a high probability to one class, with all other class probabilities close to zero. When that happens, the teacher doesn’t provide much additional information beyond the ground truth labels, so instead we “soften” the probabilities by scaling the logits with a positive temperature hyperparameter $T$ before applying the softmax:

As shown in Figure 3-4, higher values of $T$ produce a softer probability distribution over the classes and reveal much more information about the decision boundary that the teacher has learned for each training example. When $T=1$ we recover the original softmax distribution.

Figure 3-4. Comparison of a hard label which is one-hot encoded (left), softmax probabilities (middle) and softened class probabilities (right).

Since the student also produces softened probabilities $q_i\left(x\right)$ of its own we can use the Kullback-Leibler (KL) divergence

to measure the difference between the two probability distributions and thereby define a knowledge distillation loss:

where $T^2$ is a normalization factor to account for the fact that the magnitude of the gradients produced by soft labels scales as $1/T^2$. For classification tasks, the student loss is then a weighted average of the distillation loss with the usual cross-entropy loss $L_{CE}$ of the ground truth labels:

where $\alpha$ is a hyperparameter that controls the relative strength of each loss. A diagram of the whole process is shown in Figure 3-5 and the temperature is set to 1 at inference time to recover the standard softmax probabilities.

Figure 3-5. Cartoon of the knowledge distillation process.

Knowledge Distillation for Pretraining

Knowledge distillation can also be used during pretraining to create a general-purpose student that can be subsequently fine-tuned on downstream tasks. In this case, the teacher is a pretrained language model like BERT which transfers its knowledge about masked-language-modeling to the student. For example, in the DistilBERT paper,⁷ the masked-language-modeling loss $L_{mlm}$ is augmented with a term from knowledge distillation and a cosine embedding loss $L_{cos}=1-cos(h_s,h_t)$ to align the directions of the hidden state vectors between the teacher and student:

Since we already have a fine-tuned BERT-base model, let’s see how we can use knowledge distillation to fine-tune a smaller and faster model. To do that we’ll need a way to augment the cross-entropy loss with a $L_{KD}$ term; fortunately we can do this by creating our own trainer! Let’s take a look at how to do this in the next section.

Creating a Knowledge Distillation Trainer

To implement knowledge distillation we need to add a few things to the Trainer base class:

• The new hyperparameters $\alpha$ and $T$ which control the relative weight of the distillation loss and how much the probability distribution of the labels should be smoothed.

• The fine-tuned teacher model, which in our case is BERT-base

• A new loss function that includes the cross-entropy loss with the knowledge distillation loss.

Adding the new hyperparameters is quite simple since we just need to subclass TrainingArguments and include them as new attributes:

from transformers import TrainingArguments

class DistillationTrainingArguments(TrainingArguments):
    def __init__(self, *args, alpha=0.5, temperature=2.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.alpha = alpha
        self.temperature = temperature

For the trainer itself, we want a new loss function so the way to implement this is by subclassing Trainer and overriding the compute_loss function to include the knowledge distillation loss term $L_{KD}$:

import torch.nn as nn
import torch.nn.functional as F
from transformers import Trainer

class DistillationTrainer(Trainer):
    def __init__(self, *args, teacher_model=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.teacher_model = teacher_model

    def compute_loss(self, model, inputs):
        outputs_stu = model(**inputs)
        # Extract cross-entropy loss and logits from student
        loss_ce = outputs_stu.loss
        logits_stu = outputs_stu.logits
        # Extract logits from teacher
        with torch.no_grad():
            outputs_tea = self.teacher_model(**inputs)
            logits_tea = outputs_tea.logits
        # Soften probabilities and compute distillation loss
        loss_fct = nn.KLDivLoss(reduction="batchmean")
        loss_kd = self.args.temperature ** 2 * loss_fct(
            F.log_softmax(logits_stu / self.args.temperature, dim=-1),
            F.softmax(logits_tea / self.args.temperature, dim=-1))
        # Return weighted student loss
        return self.args.alpha * loss_ce + (1. - self.args.alpha) * loss_kd

Let’s unpack this code a bit. When we instantiate DistillationTrainer we pass a teacher_model argument with a teacher that has already been fine-tuned on our task. Next, in the compute_loss function we extract the logits from the student and teacher, scale them by the temperature and then normalize them with a softmax before passing them to PyTorch’s nn.KLDivLoss function for computing the KL divergence. Since nn.KLDivLoss expects the inputs in the form of log-probabilities, we’ve used the F.log_softmax function to normalize the student’s logits, while the teacher’s logits are converted to probabilities with a standard softmax. The reduction=batchmean argument in nn.KLDivLoss specifies that we average the losses over the batch dimension.

Choosing a Good Student Initialization

Now that we have our custom trainer, the first question you might have is which pretrained language model should we pick for the student? In general we should pick smaller model for the student to reduce the latency and memory footprint, and a good rule of thumb from the literature⁸ is that knowledge distillation works best when the teacher and student are of the same model type. One possible reason for this is that different model types, say BERT and RoBERTa, can have different output embedding spaces which hinders the ability of the student to mimic the teacher. In our case study, the teacher is BERT-base so DistilBERT is natural candidate to intitialize the student since it has 40% less parameters and has been shown to achieve strong results on downstream tasks.

First we’ll need to tokenize and encode our queries, so let’s instantiate the tokenizer from DistilBERT and create a simple function to take care of the preprocessing:

student_ckpt = "distilbert-base-uncased"
student_tokenizer = AutoTokenizer.from_pretrained(student_ckpt)

def tokenize_text(batch, tokenizer):
    return tokenizer(batch["text"], truncation=True)

clinc_enc = clinc.map(tokenize_text, batched=True, remove_columns=["text"],
                      fn_kwargs={"tokenizer": student_tokenizer})
clinc_enc.rename_column_("intent", "labels")

Here we’ve removed the text column since we no longer need it and we’ve also used the fn_kwargs argument to specify which tokenizer should be used in the tokenize_text function. We’ve also renamed the intent column to labels so it can be automatically detected by the trainer. Now that our texts are processed, the next thing to do is instantiate DistilBERT for fine-tuning. Since we will be doing multiple runs with the trainer, we’ll use a function to initialize the model with each new run:

import torch
from transformers import AutoConfig

num_labels = intents.num_classes
id2label = bert_model.config.id2label
label2id = bert_model.config.label2id
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

student_config = (AutoConfig
                  .from_pretrained(student_ckpt, num_labels=num_labels,
                                   id2label=id2label, label2id=label2id))

def student_init():
    return (AutoModelForSequenceClassification
            .from_pretrained(student_ckpt, config=student_config).to(device))

Here we’ve also specified the number of classes our model should expect, and used the baseline model’s configuration to provide the mappings id2label and label2id between ID and intent name. Next, we need to define the metrics to track during training. As we did in the performance benchmark, we’ll use accuracy as the main metric so we can reuse our accuracy_score function in the compute_metrics function that we’ll include in the trainer:

def compute_metrics(pred):
    predictions, labels = pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy_score.compute(predictions=predictions, references=labels)

In this function, the predictions from the sequence modeling head come in the form of logits, so we use the np.argmax function to find the most confident class prediction and compare that against the ground truth labels.

Finally, we just need to define the training arguments. To warm-up, we’ll set $\alpha =1$ to see how well DistilBERT performs without any signal from the teacher:⁹

batch_size = 48

student_training_args = DistillationTrainingArguments(
    output_dir="checkpoints", evaluation_strategy = "epoch", num_train_epochs=5,
    learning_rate=2e-5, per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size, alpha=1, weight_decay=0.01)

Next we load the teacher model, instantiate the trainer and start fine-tuning:

teacher_checkpoint = "lewtun/bert-base-uncased-finetuned-clinc"
teacher_model = (AutoModelForSequenceClassification
                 .from_pretrained(teacher_checkpoint, num_labels=num_labels)
                 .to(device))

distil_trainer = DistillationTrainer(model_init=student_init,
    teacher_model=teacher_model, args=student_training_args,
    train_dataset=clinc_enc['train'], eval_dataset=clinc_enc['validation'],
    compute_metrics=compute_metrics, tokenizer=student_tokenizer)

distil_trainer.train();

The accuracy on the validation set looks quite good compared to the 94% that BERT-base teacher achieves. Now that we’ve fine-tuned DistilBERT, we can wrap it in a TextClassificationPipeline and run it through our performance benchmark:

pipe = TextClassificationPipeline(
    model=distil_trainer.model.to("cpu"), tokenizer=distil_trainer.tokenizer)
optim_type = "DistilBERT"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

Model size (MB) - 255.89
Average latency (ms) - 24.13 +- 10.06
Accuracy on test set - 0.856

To compare these results against our baseline, let’s create a scatter plot of the accuracy against the latency, with the radius of each point corresponding to the size of the model. The following function does what we need and marks the current optimization type as a dashed circle to aid the comparison to previous results:

import pandas as pd

def plot_metrics(perf_metrics, current_optim_type):
    df = pd.DataFrame.from_dict(perf_metrics, orient='index')

    for idx in df.index:
        df_opt = df.loc[idx]
        if idx == current_optim_type:
            plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,
                        alpha=0.5, s=df_opt["size_mb"], label=idx,
                        marker='$u25CC$')
        else:
            plt.scatter(df_opt["time_avg_ms"], df_opt["accuracy"] * 100,
                        s=df_opt["size_mb"], label=idx, alpha=0.5)

    legend = plt.legend(bbox_to_anchor=(1,1))
    for handle in legend.legendHandles:
        handle.set_sizes([20])

    plt.ylim(80,90)
    plt.xlim(5, 53)
    plt.ylabel("Accuracy (%)")
    plt.xlabel("Average latency (ms)")
    plt.show()

plot_metrics(perf_metrics, optim_type)

From the plot we can see that by using a smaller model we’ve managed to decrease the average latency by almost a factor of two. And all this at the price of just over a 1% reduction in accuracy! Let’s see if we can close that last gap by including the distillation loss the teacher and finding good values for $\alpha$ and $T$.

Finding Good Hyperparameters with Optuna

So what values of $\alpha$ and $T$ should we pick? We could do a grid search over the 2D parameter space but a much better alternative is to use Optuna,¹⁰ which is an optimization framework designed for just this type of task. Optuna formulates the search problem in terms of an objective function that is optimized through multiple trials. For example, suppose we wished to minimize Rosenbrock’s “banana function”

which is a famous test case for optimization frameworks. As shown in Figure 3-6, the function gets its name from the curved contours and has a global minimum at $(x,y)=(1,1)$. Finding the valley is an easy optimization problem, but converging to the global minimum is not.

Figure 3-6. Plot of the Rosenbrock function of two variables

In Optuna, we can find the minimum of $f(x,y)$ by defining an objective function that returns the value of $f(x,y)$:

def objective(trial):
    x = trial.suggest_float("x", -2, 2)
    y = trial.suggest_float("y", -2, 2)
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

The trial.suggest_float object specifies the parameter ranges to sample uniformly from and Optuna also provides suggest_int and suggest_categorical for integer and categorical parameters respectively. Optuna collects multiple trials as a study so to create one we just pass the objective function to study.optimize as follows:

study = optuna.create_study()
study.optimize(objective, n_trials=1000)

Once the study is completed, we can then find the best parameters as follows:

study.best_params

{'x': 1.0294476378522224, 'y': 1.0595653705812769}

We see that with 1,000 trials, Optuna has managed to find values for $x$ and $y$ that are reasonably close to the global minimum. To use Optuna in Transformers, we use a similar logic by first defining the hyperparameter space that we wish to optimize over. In addition to $\alpha$ and $T$, we’ll include the number of training epochs as follows:

def hp_space(trial):
    return {"num_train_epochs": trial.suggest_int("num_train_epochs", 5, 10),
        "alpha": trial.suggest_float("alpha", 0, 1),
        "temperature": trial.suggest_int("temperature", 2, 20)}

Running the hyperparameter search with the Trainer is then quite simple; we just need to specify the number of trials to run and a direction to optimize for. Since we want the best possible accuracy, we pick direction="maximize" in the Trainer.hyperparameter_search function and pass the hyperparameter search space as follows:

best_run = distil_trainer.hyperparameter_search(
    n_trials=20, direction="maximize", hp_space=hp_space)

The hyperparameter_search method returns a BestRun object which contains the value of the objective that was maximized (by default the sum of all metrics) and the hyperparameters it used for that run:

best_run

BestRun(run_id='4', objective=3080.872670967742,
 > hyperparameters={'num_train_epochs': 8, 'alpha': 0.31235083318309453,
 > 'temperature': 16})

This value of $\alpha$ tells us that most of the training signal is coming from the knowledge distillation term. Let’s update our trainer with these values and run the final training run:

for k,v in best_run.hyperparameters.items():
    setattr(distil_trainer.args, k, v)

distil_trainer.train();

Remarkably we’ve been able to train the student to match the accuracy of the teacher, despite having almost half the number of parameters! Let’s save the model for future use:

distil_trainer.save_model("models/distilbert-base-uncased-distilled-clinc")

Benchmarking Our Distilled Model

Now that we have an accurate student, let’s create a pipeline and redo our benchmark to see how we perform on the test set:

pipe = TextClassificationPipeline(
    model=distil_trainer.model.to("cpu"), tokenizer=distil_trainer.tokenizer)
optim_type = "Distillation"
pb = PerformanceBenchmark(pipe, clinc["test"], optim_type=optim_type)
perf_metrics.update(pb.run_benchmark())

Model size (MB) - 255.89
Average latency (ms) - 24.58 +- 7.66
Accuracy on test set - 0.871

To put these results in context, let’s also visualise them with our plot_metrics function:

plot_metrics(perf_metrics, optim_type)

As expected, the model size and latency remain essentially unchanged compared to the DistilBERT benchmark, but the accuracy has improved and even surpassed the performance of the teacher! We can actually compress our distilled model even further using a technique known as quantization. That’s the topic for the next section.

Quantization Strategies

Quantizing Transformers in PyTorch

The main bottleneck for running inference with Transformers is the compute and memory bandwidth associated with the enormous number of weights in these models. For this reason, dynamic quantization is currently the best approach for Transformer-based models in NLP. In smaller computer vision models the limiting factor is the memory bandwidth of the activations which is why static quantization is generally used and quantization aware training in cases where the performance drops are too significant.

Implementing dynamic quantization in PyTorch is quite simple and can be done with a single line of code:

from torch.quantization import quantize_dynamic

model_ckpt = "models/distilbert-base-uncased-distilled-clinc"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = (AutoModelForSequenceClassification
         .from_pretrained(model_ckpt).to("cpu"))

model_quantized = quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)

Here we pass to quantize_dynamic the full-precision model and specify the set of PyTorch layer classes in that model that we want to quantize. The dtype argument specifies the target precision and can be fp16 or qint8.

Optimizing for Transformer Architectures

We’ve just seen that the ONNX Runtime was very good at optimizing our distilled model out of the box. However, the ONNX Runtime library also offers an optimizer module that contains some Transformer-specific optimizations that we can try to see if the model is fully optimized or not. To use the optimizer module we first need to define some optimization options that are specific to our model. In our case, DistilBERT belongs to the bert model type so we need to use the BertOptimizationOptions class from onnxruntime_tools:

from onnxruntime_tools.transformers.onnx_model_bert import (
    BertOptimizationOptions)

model_type = "bert"
opt_options = BertOptimizationOptions(model_type)
opt_options.enable_embed_layer_norm = False

Here we’ve disabled the norm optimization on the embedding layer to get better model size compression. Now that we’ve specified the model options, we can then run optimizer.optimize_model to optimize the ONNX model specifically for BERT-like architectures:

from onnxruntime_tools import optimizer

opt_model = optimizer.optimize_model(
    "onnx/model.onnx", model_type, num_heads=12, hidden_size=768,
    optimization_options=opt_options)
opt_model.save_model_to_file("onnx/model.opt.onnx")

Here we’ve specified the number of heads and hidden size in our DistilBERT model. The last thing to do is create an inference session for our optimized model, wrap it in a pipeline and run it through our benchmark:

onnx_model_opt = create_model_for_provider("onnx/model.opt.onnx")
pipe = OnnxPipeline(onnx_model_opt, tokenizer)
optim_type = "Distillation + ORT (optimized)"
pb = OnnxPerformanceBenchmark(pipe, clinc["test"], optim_type,
                              model_path="onnx/model.opt.onnx")
perf_metrics.update(pb.run_benchmark())

Model size (MB) - 255.86
Average latency (ms) - 11.22 +- 3.52
Accuracy on test set - 0.871

plot_metrics(perf_metrics, optim_type)

Okay, it seems that our original ORT optimization was already close to the optimal one for this architecture. Let’s now see what happens if we add quantization to the mix. Similar to PyTorch, ORT offers three ways to quantize a model: dynamic, static, and quantization aware training. As we did with PyTorch, we’ll apply dynamic quantization to our distilled model. In ORT, the quantization is applied through the quantize_dynamic function which requires a path to the ONNX model to quantize, a target path to save the quantized model to, and the data type to reduce the weights to:

from onnxruntime.quantization import quantize_dynamic, QuantType

model_input = "onnx/model.onnx"
model_output = "onnx/model.quant.onnx"
quantize_dynamic(model_input, model_output, weight_type=QuantType.QInt8)

Now that the model is quantized, let’s run it through our benchmark:

onnx_quantized_model = create_model_for_provider(model_output)
pipe = OnnxPipeline(onnx_quantized_model, tokenizer)
optim_type = "Distillation + ORT (quantized)"
pb = OnnxPerformanceBenchmark(pipe, clinc["test"], optim_type,
                              model_path=model_output)
perf_metrics.update(pb.run_benchmark())

Model size (MB) - 185.71
Average latency (ms) - 6.95 +- 4.75
Accuracy on test set - 0.875

plot_metrics(perf_metrics, optim_type)

Wow, ORT quantization has reduced the model size and latency by around a factor of two compared to the model obtained from PyTorch quantization (the Distillation + quantization blob). One reason for this is that PyTorch only optimizes the nn.Linear modules, while ONNX quantizes the embedding layer as well. From the plot we can also see that applying ORT quantization to our distilled model has provided an almost 7-fold gain compared to our BERT baseline!

This concludes our analysis of techniques to speed-up Transformers for inference. We have seen that methods such as quantization reduce the model size by reducing the precision of the representation. Another strategy to reduce the size is to remove some weights altogether - this technique is called weight pruning and is the focus of the next section.

Sparsity in Deep Neural Networks

As shown in Figure 3-11, the main idea behind pruning is to gradually remove weight connections (and potentially neurons) during training such that the model becomes progressively sparser. The resulting pruned model has a smaller number of non-zero parameters which can then be stored in a compact sparse matrix format. Pruning can be also combined with quantization to obtain further compression.

Figure 3-11. Weights and neurons before and after pruning. Image from Learning both Weights and Connections for Efficient Neural Networks by S. Han et al (2015).

Weight Pruning Methods

Mathematically, the way most weight pruning methods work is to calculate a matrix $?\in {ℝ}^{n\times n}$ of importance scores and then select the top-$k$ percent of weights by importance:

In effect, $k$ acts as a new hyperparameter to control the amount of sparsity in the model, that is the proportion of weights that are zero-valued. Lower values of $k$ correspond to sparser matrices. From these scores we can then define a mask matrix $?\in {\{0,1\}}^{n\times n}$ that masks the weights $W_{ij}$ during the forward pass with some input $x_i$ and effectively creates a sparse network of activations $a_i$:

As discussed in the tongue-in-cheek Optimal Brain Surgeon paper¹⁶, at the heart of each pruning method are a set of questions that need to be considered:

• Which weights should be eliminated?

• How should the remaining weights be adjusted for best performance?

• How can such network pruning be done in a computationally efficient way?

The answers to these questions inform how the score matrix $?$ is computed, so let’s begin by looking at one of the earliest and most popular pruning methods: magnitude pruning.

Magnitude Pruning

As the name suggests, magnitude pruning calculates the scores according to the magnitude of the weights $?={\left(\mid W_{ij}\mid \right)}_{1\le j,j\le n}$ and then derives the masks from $?={\mathrm{Top}}_k\left(?\right)$. In the literature it is common to apply magnitude pruning in an iterative fashion¹⁷ by first training the model to learn which connections are important and pruning the weights of least importance. The sparse model is then re-trained and the process repeated until the desired sparsity is reached.

One drawback with this approach is that it is computationally demanding: at every step of pruning we need to train the model to convergence. For this reason it is generally better to gradually increase the initial sparsity $s_i$ (which is usually zero) to a final value $s_f$ after some number of steps $N$:¹⁸

$$s_t=s_f+(s_i-s_f){\left(1-\frac{t-t_0}{N\Delta t}\right)}^3\mathrm{for}t\in \{t_0,t_0+\Delta t,...,t_0+N\Delta t\}.$$

Here the idea is to update the binary masks $?$ every $\Delta t$ steps to allow masked weights to reactivate during training and recover from any potential loss in accuracy that is induced by the pruning process. As shown in Figure 3-12, the cubic factor implies that the rate of weight pruning is highest in the early phases (when the number of redundant weights is large) and gradually tapers off.

Figure 3-12. The cubic sparsity scheduler used for pruning.

One problem with magnitude pruning is that it is really designed for pure supervised learning, where the importance of each weight is directly related to the task at hand. By contrast, in transfer learning the importance of the weights is primarily determined by the pretraining phase, so magnitude pruning can remove connections that are important for the fine-tuning task. Recently, an adaptive approach¹⁹ called movement pruning has been proposed by the HuggingFace team - let’s take a look.

Movement Pruning

The basic idea behind movement pruning is to gradually remove weights during fine-tuning such that the model becomes progressively sparser. The key novelty is that both the weights and the scores are learned during fine-tuning. So instead of deriving the scores directly from the weights (like magnitude pruning does), the scores in movement pruning are arbitrary, and learned through gradient descent like any other neural network parameter. This implies that in the backward pass, we also track the gradient of the loss $L$ with respect to the scores $S_{ij}$. We can calculate the gradient from the expression of the activations $a_i$ as follows:

$$\frac{\partial L}{\partial S_{ij}}=\frac{\partial L}{\partial a_i}\frac{\partial a_i}{\partial S_{ij}}=\frac{\partial L}{\partial a_i}{W}_{ij}{x}_j.$$

Once the scores are learned, it is then straightforward to generate the binary mask using $?={\mathrm{Top}}_k\left(?\right)$. There is also a “soft” version of movement pruning where instead of picking the top-$k$% of weights, one uses a global threshold $\tau$ to define the binary mask: $?=(?>\tau )$.

The intuition behind movement pruning is that the weights which are “moving” the most from zero are the most important ones to keep. To see this, we first note that the gradient of $L$ with respect to the weights $W_{ij}$ is given by

$$\frac{\partial L}{\partial W_{ij}}=\frac{\partial L}{\partial a_i}{M}_{ij}{x}_j,$$

which can be combined with the expression for $\partial L/\partial S_{ij}$ to yield

$$\frac{\partial L}{\partial S_{ij}}=\frac{\partial L}{\partial W_{ij}}{W}_{ij}{M}_{ij}.$$

Since the scores are increasing when the gradient $\partial L/\partial S_{ij}$ is negative, we see that this occurs under two scenarios (we can drop $?$ since it’s a positive matrix):

$$\begin{array}{cccc}a)\hfill & {\displaystyle \frac{\partial L}{\partial W_{ij}}}<0\hfill & \mathrm{and}\hfill & W_{ij}>0\hfill \\ b)\hfill & {\displaystyle \frac{\partial L}{\partial W_{ij}}}>0\hfill & \mathrm{and}\hfill & W_{ij}<0\hfill \end{array}$$

In other words, the positive weights increase during fine-tuning and vice versa for the negative weights which is equivalent to saying that the scores increase as the weights move away from zero. As shown in Figure 3-13, this behavior differs from magnitude pruning which selects as the most important weights those which a furthest from zero.

Figure 3-13. Comparison of weights removed (in grey) during magnitude pruning (left) and movement pruning (right).

These differences between the two pruning methods are also evident in the distribution of the remaining weights. As shown in Figure 3-14, magnitude pruning produces two clusters of weights, while movement pruning produces a smoother distribution.

Figure 3-14. Distribution of remaining weights for magnitude pruning (MaP) and movement pruning (MvP)

In this chapter we’ll examine how well movement pruning works with a top-$k$ scorer on our intent classifier. As of this book’s writing, Transformers does not support pruning methods “out of the box”, so we’ll have to implement the main classes ourselves. Fortunately, the code used to produce the results from the movement pruning paper is available in the examples/research$\_$projects/movement-pruning folder of the Transformers repository, so we have a great foundation to work from. Let’s get started!

Creating Masked Transformers

To implement movement pruning, we’ll need a few different ingredients:

• A ${\mathrm{Top}}_k$ operator that we can use to binarize the scores by selecting the top-$k$% of weights.

• A way to apply the adaptive masking on-the-fly to our BERT-base model.

• A cubic sparsity scheduler.

Let’s start by implementing the ${\mathrm{Top}}_k$ binarizer. From the definition, we need to calculate a binary mask matrix $?$ from a real-valued matrix $?$ if and only if $S_{ij}$ is among the $k$% highest values of $?$. Since the back-propagated gradients will flow through the binary mask, we’ll use the autograd.Function from PyTorch to compute the mask on the forward pass and automatically calculate the gradients in the backward pass:

from torch.autograd import Function

class TopKBinarizer(Function):
    @staticmethod
    def forward(ctx, inputs, threshold):
        # Get threshold from column in validation set
        if not isinstance(threshold, float):
            threshold = threshold[0]
        # Sort the inputs
        mask = inputs.clone()
        _, idx = inputs.flatten().sort(descending=True)
        # Get number of elements above the threshdold
        j = int(threshold * inputs.numel())
        # Zero-out elements below the threshold
        flat_out = mask.flatten()
        flat_out[idx[j:]] = 0
        flat_out[idx[:j]] = 1
        return mask

    @staticmethod
    def backward(ctx, gradOutput):
        return gradOutput, None

Let’s test this out on a random matrix of scores:

from torch.autograd import Variable

torch.manual_seed(123)
dtype = torch.FloatTensor
scores = Variable(torch.randn(2, 2).type(dtype), requires_grad=True)
scores

tensor([[-0.1115,  0.1204],
        [-0.3696, -0.2404]], requires_grad=True)

Now let’s see what happens if we zero-out half of the scores:

topk = TopKBinarizer()
topk.apply(scores, 0.5)

tensor([[1., 1.],
        [0., 0.]], grad_fn=<TopKBinarizerBackward>)

Great, this makes sense since the entries in the top row of scores are the ones with the highest value. Okay, so now that we have a way to binarize the scores, the next step is to implement a fully connected layer that can calculate the binary mask on-the-fly and multiply this by the weight matrix and inputs. As of this book’s writing, there is no simple way to do this beyond extending the various BERT classes in Transformers. We refer the reader to the implementation in the Transformers repository, but note that the main ingredient is a replacement of all nn.Linear layers with a new layer that computes the mask on the forward pass:

from torch.nn import init

class MaskedLinear(nn.Linear):
    def __init__(self, in_features, out_features, bias=True, mask_scale = 0.0):
        super(MaskedLinear, self).__init__(
            in_features=in_features, out_features=out_features, bias=bias)
        self.mask_scale = mask_scale
        self.mask_init = mask_init
        self.mask_scores = nn.Parameter(torch.Tensor(self.weight.size()))
        self.init_mask()

    def init_mask(self):
        init.constant_(self.mask_scores, val=self.mask_scale)

    def forward(self, input, threshold):
        # Get the mask
        mask = TopKBinarizer.apply(self.mask_scores, threshold)
        # Mask weights with computed mask
        weight_thresholded = mask * self.weight
        # Compute output (linear layer) with masked weights
        return F.linear(input, weight_thresholded, self.bias)

This new linear layer can then be used to build up a custom MaskedBertModel, MaskedBertForSequenceClassification and so on by allowing the threshold parameter to be passed along with the inputs. We won’t show the explicit code here but refer the reader to the examples/research$\_$projects/movement-pruning folder of the Transformers repository. These new masked classes work in the same way the ordinary ones, so let’s load the configuration and model_init so we can do multiple fine-pruning runs:

from pruning import MaskedBertConfig, MaskedBertForSequenceClassification

masked_config = MaskedBertConfig(num_labels=num_labels)

def model_init():
    return (MaskedBertForSequenceClassification
            .from_pretrained(bert_ckpt, config=masked_config).to(device))

Creating a Pruning Trainer

Now that we have our masked model, the next step is to implement a custom trainer that we can use for fine-pruning. Similar to knowledge distillation, we’ll need a few ingredients:

• New hyperparameters like the amount of sparsity to start and end with during the training run. We’ll also need to specify what fraction of steps we use for warmup and cool down which are important.

• A way to optimize the new learning rate for the scores.

• A custom loss that can calculate the threshold at each step and feed that to the model to generate the loss.

The new training arguments are simple to include, and again we just subclass TrainingArguments:

class PruningTrainingArguments(TrainingArguments):
    def __init__(self, *args, initial_threshold=1., final_threshold=0.1,
                 initial_warmup=1, final_warmup=2,
                 mask_scores_learning_rate=1e-2, **kwargs):
        super().__init__(*args, **kwargs)
        self.initial_threshold = initial_threshold
        self.final_threshold = final_threshold
        self.initial_warmup = initial_warmup
        self.final_warmup = final_warmup
        self.mask_scores_learning_rate = mask_scores_learning_rate

For the trainer we’ll have to do a bit more work, so let’s look at a skeleton of what we need to override:

class PruningTrainer(Trainer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.t_total = (len(self.get_train_dataloader()) //
                        self.args.gradient_accumulation_steps
                        * self.args.num_train_epochs)


    def create_optimizer_and_scheduler(self, num_training_steps):
        pass

    def compute_loss(self, model, inputs):
        pass

    def _schedule_threshold(self, step, total_step, warmup_steps,
        initial_threshold, final_threshold, initial_warmup, final_warmup):
        pass

First we need to implement the cubic sparsity scheduler. This is similar to the equation we saw earlier but in movement pruning we also allow for some amount of cool-down steps $t_f$ so the definition is as follows:

The following function implements this logic:

def _schedule_threshold(self, step, total_step, warmup_steps,
        initial_threshold, final_threshold, initial_warmup, final_warmup):
        if step <= initial_warmup * warmup_steps:
            threshold = initial_threshold
        elif step > (total_step - final_warmup * warmup_steps):
            threshold = final_threshold
        else:
            spars_warmup_steps = initial_warmup * warmup_steps
            spars_schedu_steps = ((final_warmup + initial_warmup)
                                  * warmup_steps)
            mul_coeff = 1 - ((step - spars_warmup_steps)
                             / (total_step - spars_schedu_steps))
            threshold = final_threshold + (
                (initial_threshold - final_threshold) * (mul_coeff ** 3))
        return threshold

PruningTrainer._schedule_threshold = _schedule_threshold

In addition to the usual inputs, the masked model expects the sparsity threshold produced from the sparsity scheduler. A simple way to provide this information is to overwrite the compute_loss function of the Trainer and extract the threshold at each training step:

def compute_loss(self, model, inputs):
    threshold = self._schedule_threshold(step=self.state.global_step+1,
        total_step=self.t_total, warmup_steps=self.args.warmup_steps,
        final_threshold=self.args.final_threshold,
        initial_threshold=self.args.initial_threshold,
        final_warmup=self.args.final_warmup,
        initial_warmup=self.args.initial_warmup)
    inputs["threshold"] = threshold
    outputs = model(**inputs)
    loss, _ = outputs
    return loss

PruningTrainer.compute_loss = compute_loss

As noted earlier a key property of movement pruning is that the scores are learned during fine-tuning. This means that we need to instruct the Trainer to optimize for the usual weights and these new score parameters. The way this is done in practice is to overwrite the Trainer.create_optimizer_and_scheduler function and indicate which parameters belong to the optimizer’s grouped parameters variable:

from transformers import AdamW, get_linear_schedule_with_warmup

def create_optimizer_and_scheduler(self, num_training_steps: int):
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [{
        "params": [p for n, p in self.model.named_parameters()
                   if "mask_score" in n and p.requires_grad],
        "lr": self.args.mask_scores_learning_rate},
        {"params": [p for n, p in self.model.named_parameters()
                    if "mask_score" not in n and p.requires_grad
                    and not any(nd in n for nd in no_decay)],
         "lr": self.args.learning_rate,
         "weight_decay": self.args.weight_decay},
        {"params": [p for n, p in self.model.named_parameters()
                    if "mask_score" not in n and p.requires_grad
                    and any(nd in n for nd in no_decay)],
         "lr": self.args.learning_rate,
         "weight_decay": 0.0}]

    self.optimizer = AdamW(optimizer_grouped_parameters,
                           lr=self.args.learning_rate,
                           eps=self.args.adam_epsilon)
    self.lr_scheduler = get_linear_schedule_with_warmup(
        self.optimizer, num_warmup_steps=self.args.warmup_steps,
        num_training_steps=self.t_total)

PruningTrainer.create_optimizer_and_scheduler = create_optimizer_and_scheduler

Now that we’ve created out trainer it’s time to give it a spin!

Fine-Pruning With Increasing Sparsity

To evaluate the effect of pruning we’ll fine-tune BERT-base on our dataset at increasing levels of sparsity. We expect some accuracy drop compared to the 94.3% that the unpruned model achieves, but hopefully it is not too much. First we need to define the base training arguments for our runs:

num_train_epochs = 5
logging_steps = len(clinc_enc['train']) // batch_size
warmup_steps = logging_steps * num_train_epochs * 0.1
mask_scores_learning_rate = 1e-2

pruning_training_args = PruningTrainingArguments(
    output_dir="checkpoints", evaluation_strategy="epoch", learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size, logging_steps=logging_steps,
    warmup_steps=warmup_steps, num_train_epochs=num_train_epochs,
    mask_scores_learning_rate=mask_scores_learning_rate, weight_decay=0.01)

Next, we’ll gradually decrease the threshold parameter in our mask which is equivalent to increasing the sparsity of the weights. Since Transformers are quite robust to sparsity with movement pruning, we’ll start by retaining 30% of the weights and decrease down to 1%. The following loop implements the sparsity increase by updating the training arguments and validation set with the current threshold value, fine-tuning and then saving the models and accuracies for further analysis:

accuracies = {}

for threshold in [0.3, 0.1, 0.05, 0.03, 0.01]:
    model_ckpt = f"prunebert-{int(threshold * 100)}"
    pruning_training_args.final_threshold = threshold
    pruning_training_args.run_name = model_ckpt
    # Include the current sparsity in the validation set
    eval_ds = clinc_enc['validation'].map(lambda x : {'threshold': threshold})

    pruning_trainer = PruningTrainer(model_init=model_init,
    args=pruning_training_args, train_dataset=clinc_enc["train"],
    eval_dataset=eval_ds, tokenizer=bert_tokenizer,
    compute_metrics=compute_metrics)

    pruning_trainer.train()
    pruning_trainer.save_model(f"models/{model_ckpt}")
    preds = pruning_trainer.evaluate()
    accuracies[threshold] = preds["eval_accuracy"]

    wandb.finish()

As shown in Figure 3-15, we can see that pruning only has a small impact on accuracy and only starts to degrade once we start pruning more than 95% of the weights!

Figure 3-15. Effect of removing weights on BERT-base’s accuracy

Since the best performing model appears to have around 5% of the weights, let’s use this one to count the true number of remaining values and convert the model into a format suitable for native model classes in Transformers.

Counting the Number of Pruned Weights

Now that we’ve pruned a model, let’s do a sanity check to count the number of parameters we’ve removed. A simple way to do this is via PyTorch’s state_dict object which we encountered when saving the model to calculate its size. First, let’s load the state_dict associated with our pruned model:

prunebert_model_ckpt = "models/prunebert-5"
prunebert_args = torch.load(
    f"{prunebert_model_ckpt}/training_args.bin", map_location="cpu")
state_dict = torch.load(
    f"{prunebert_model_ckpt}/pytorch_model.bin", map_location="cpu")

To count the number of pruned weights, we’ll iterate throught the dictionary and count the number of elements in every layer that was masked with a layer_name.mask_scores and calculate the sparsity from this relative to the other layers. Since we’ve only pruned the trainable parameters of the model we’ll exclude the embedding parameters from the count, so we arrive at the following code:

# Number of remaining (not pruned) params in the encoder
remaining_count = 0
# Number of params in the encoder
encoder_count = 0
# Fraction of remaining weights
final_threshold = prunebert_args.final_threshold

for name, param in state_dict.items():
    if "encoder" not in name:
        continue

    if "mask_scores" in name:
        mask_ones = TopKBinarizer.apply(param, final_threshold).sum().item()
        remaining_count += mask_ones
    else:
        encoder_count += param.numel()
        if "bias" in name or "LayerNorm" in name:
            remaining_count += param.numel()

print(f"Remaining weights: {100 * remaining_count / encoder_count:.2f}%")

Remaining weights: 5.13%

Pruning Once and For All

Now that we’re confident that we’ve pruned the model as expected, the final step is convert the model back into a form that is suitable for the standard BertForSequenceClassification class. Here we need to collect all the dense tensors and apply the mask to those which have been marked with the mask_scores name. The resulting state_dict can then be saved along with the configuration and tokenizer from our fine-pruned model:

import shutil

model_path = prunebert_model_ckpt
target_path = f"models/bertarized"
model = torch.load(f"{prunebert_model_ckpt}/pytorch_model.bin")
pruned_model = {}

for name, tensor in model.items():
    if "embeddings" in name or "LayerNorm" in name or "pooler" in name:
        pruned_model[name] = tensor
    elif "classifier" in name:
        pruned_model[name] = tensor
    elif "bias" in name:
        pruned_model[name] = tensor
    else:
        if "mask_scores" in name:
            continue
        prefix_ = name[:-6]
        scores = model[f"{prefix_}mask_scores"]
        mask = TopKBinarizer.apply(scores, final_threshold)
        pruned_model[name] = tensor * mask

shutil.copytree(model_path, target_path, dirs_exist_ok=True)
torch.save(pruned_model, f"{target_path}/pytorch_model.bin")

As a sanity check, we can now load our fine-pruned model with the BertConfig and AutoModelForSequenceClassification as follows:

from transformers import BertConfig

config = BertConfig.from_pretrained(target_path, id2label=id2label,
                                    label2id=label2id)
model = (AutoModelForSequenceClassification
                .from_pretrained(target_path, config=config).to('cpu'))

Finally we can run our model through the performance benchmark:

pipe = TextClassificationPipeline(model=model, tokenizer=bert_tokenizer)
PerformanceBenchmark(pipe, clinc["test"]).run_benchmark();

Model size (MB) - 418.13
Average latency (ms) - 90.73 +- 45.57
Accuracy on test set - 0.840

Looking at the at the benchmark you might be surprised: the pruned model is neither smaller nor faster! The reason is that we stored weights as dense matrices which occupy the same space irrespective of how many values are set to zero. Similarly a matrix multiplication does not get faster if more values are zero. Unfortunately, modern frameworks still lack fast sparse operations and hence it is hard to get a speedup from pruning. However, we can store the matrices in a more compact format which we will explore in the next section.

Quantizing and Storing in Sparse Format

As of this book’s writing, pruning in PyTorch or TensorFlow does not lead to improved inferences times or a reduced model size since a dense tensor filled with zeroes is still dense. However, when combined with compression algorithms like gzip, pruning does allow us to reduce the size of the model on disk. To get the most amount of compression, we’ll first apply quantization to our pruned model:

from torch.quantization import quantize_dynamic

quantized_model = quantize_dynamic(model=model, qconfig_spec={torch.nn.Linear},
                                   dtype=torch.qint8)

qtz_st = quantized_model.state_dict()

Next let’s wrap this quantized model in a pipeline so we can get a sense for it’s size:

pipe = TextClassificationPipeline(model=quantized_model,
                                  tokenizer=bert_tokenizer)
PerformanceBenchmark(pipe, clinc["test"]).compute_size();

Model size (MB) - 173.15

X = np.random.uniform(size=(3, 3))
X[X < 0.6] = 0
X

array([[0.        , 0.        , 0.        ],
       [0.60754485, 0.        , 0.        ],
       [0.94888554, 0.96563203, 0.80839735]])

from scipy.sparse import csr_matrix

X_csr = csr_matrix(X)
print(X_csr)

  (1, 0)        0.6075448519014384
  (2, 0)        0.9488855372533332
  (2, 1)        0.9656320330745594
  (2, 2)        0.8083973481164611

for name, param in qtz_st.items():
    if "dtype" not in name and param.is_quantized:
        scale = param.q_scale()
        zero_point = param.q_zero_point()
        print(f"Layer name - {name}")
        print(param)
        break

Layer name - bert.encoder.layer.0.attention.self.query._packed_params.weight
tensor([[0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        ...,
        [0.0000, 0.0535, 0.0576,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000]],
       size=(768, 768), dtype=torch.qint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.004113025031983852,
       zero_point=0)

param.int_repr()

tensor([[ 0,  0,  0,  ...,  0,  0,  0],
        [ 0,  0,  0,  ...,  0,  0,  0],
        [ 0,  0,  0,  ...,  0,  0,  0],
        ...,
        [ 0, 13, 14,  ...,  0,  0,  0],
        [ 0,  0,  0,  ...,  0,  0,  0],
        [ 0,  0,  0,  ...,  0,  0,  0]], dtype=torch.int8)

print(csr_matrix(param.int_repr()[:1]))

  (0, 73)       16
  (0, 91)       -2
  (0, 102)      7
  (0, 249)      -9
  (0, 250)      7
  (0, 277)      20
  (0, 374)      -4
  (0, 395)      5
  (0, 490)      7
  (0, 496)      -16
  (0, 520)      11
  (0, 560)      -13
  (0, 575)      -5
  (0, 581)      1
  (0, 638)      20
  (0, 644)      1
  (0, 655)      5
  (0, 719)      1
  (0, 739)      20

elementary_qtz_st = {}

for name, param in qtz_st.items():
    if "dtype" not in name and param.is_quantized:
        scale = param.q_scale()
        zero_point = param.q_zero_point()
        elementary_qtz_st[f"{name}.scale"] = scale
        elementary_qtz_st[f"{name}.zero_point"] = zero_point
        # Convert sparse quantized tensors into CSR format
        int_repr = param.int_repr()
        int_repr_cs = csr_matrix(int_repr)
        elementary_qtz_st[f"{name}.int_repr.data"] = int_repr_cs.data
        elementary_qtz_st[f"{name}.int_repr.indptr"] = int_repr_cs.indptr
        elementary_qtz_st[
            f"{name}.int_repr.indices"
        ] = np.uint16(int_repr_cs.indices)
        elementary_qtz_st[
            f"{name}.int_repr.shape"
        ] = int_repr_cs.shape
    else:
        elementary_qtz_st[name] = param

torch.save(elementary_qtz_st, "tmp.pt")
!du -h tmp.pt

110M    tmp.pt