FastFormers: 233x Faster Transformers inference on CPU
By: Parth Chokhra
Tags: Data Science, Technology, AI, Machine Learning, Deep Learning

Since the birth of BERT followed by that of Transformers have dominated NLP in nearly every language-related tasks whether it is Question-Answering, Sentiment Analysis, Text classification or Text Generation. Transformers enjoys much better accuracy on all these tasks unlike RNN and LSTM the problem of vanishing gradients, which hampers learning of long data sequences. Also unlike Transformers; RNN and LSTM are not scalable as they have to take into account the output of the previous neuron.

Now the main problem with Transformers is they are highly computed intensive in both training and inference. While the training part can be solved by using pretrained language models (Open-Sourced by Large Corporates like Google, Facebook and OpenAI 😏) and fine-tuning them on our dataset. Now the latter problem is addressed by FastFormers, a set of recipes to achieve efficient inference-time performance for transformer-based models on various NLU tasks.

“Applying these proposed recipes to the SuperGLUE benchmark, authors were able to achieve from 9.8x up to 233.9x speed-up compared to out-of-the-box models on CPU. On GPU, we also achieve up to 12.4x speed-up with the presented methods.” - FastFormers

The paper FastFormers: Highly Efficient Transformer Models for Natural Language Understanding mainly focuses on providing highly efficient inference for Transformer models which enables deployment in large scale production scenarios. The authors specifically focus on inference time efficiency since it mostly dominates the cost of production deployment. In this blog, we gonna walk through all the problems and challenges this paper addresses.

So how did they address the problem of high inefficient inference time of the Transformers?

They mainly utilize three methods i.e. Knowledge Distillation, Structured Pruning and Model Quantization.

The first step is Knowledge Distillation which deals with reducing the size of model wrt depth and hidden states without compromising with accuracy.

which deals with reducing the size of model wrt depth and hidden states without compromising with accuracy. Second, Structured Pruning that reduces the size of the models by reducing the number of self-attention heads while trying to preserve the accuracy as well.

that reduces the size of the models by reducing the number of self-attention heads while trying to preserve the accuracy as well. Finally, Model Quantization which enables faster model executions by optimally utilizing hardware acceleration capabilities. On CPU, 8-bit integer quantization method is applied while on GPU, all the model parameters are converted into 16-bit floating-point data type to maximally utilize efficient Tensor Cores.

In-depth Walkthrough

Knowledge Distillation: Knowledge distillation refers to the idea of model compression by teaching a smaller network, step by step, exactly what to do using a bigger already-trained network. While large models have higher knowledge capacity than small models, this capacity might not be fully utilized. It can be computationally just as expensive to evaluate a model even if it utilizes little of its knowledge capacity. Knowledge distillation transfers knowledge from a large model to a smaller model without loss of validity. As smaller models are less expensive to evaluate, they can be deployed on less powerful hardware like a smartphone.

Knowledge distillation methods: Two approaches are particularly used namely task-specific and task-agnostic distillation.

In the taskspecific distillation, authors distill fine-tuned teacher models into smaller student architectures following the procedure proposed by TinyBERT.In the task-agnostic distillation approach, authors directly apply fine-tuning on general distilled models to tune for a specific task.

Summary of the workflow for knowledge distillation. Courtesy: Floydhub

Knowledge distillation results: In the experiments, authors have observed that distilled models do not work well when distilled to a different model type. Therefore, authors restricted our setup to avoid distilling Roberta model to BERT or vice versa. The results of knowledge distillation on the tasks are summarized with the teacher models on validation dataset in the below table. (Student referred to Distilled Model)

Accuracy of teacher and student models on the validation data set for each task of SuperGLUE benchmark with knowledge distillation. Courtesy: FastFormers

Neural Network Pruning: Neural network pruning is a method of compression that involves removing weights from a trained model. In agriculture, pruning is cutting off unnecessary branches or stems of a plant. In machine learning, pruning is removing unnecessary neurons or weights. Neural network pruning techniques can reduce the parameter counts of trained networks by over 90%, decreasing storage requirements and improving the computational performance of inference without compromising accuracy. This helps in decreasing the size or energy consumption of the trained neural network and helps to make inference more efficient. Pruning makes the network more efficient and lighter.

Synapses and neurons before and after pruning. Courtesy: Link

Structured pruning methods: The first step of our structured pruning method is to identify the least important heads in Multi-Head Attention and the least important hidden states in the feed-forward layers.

First-order method for computing the importance score, which utilizes the first-order gradient information instead of using magnitude-based pruning.

Before doing the importance score computation, the authors add a mask variable to each attention head for the gradient computation of the heads. Then the authors run forward and backward passes of the model on the entire validation data set, then the absolute values of the gradients are accumulated. These accumulated values are then used as importance scores which we use to sort the importance of the heads and the intermediate hidden states.

Based on the target model size, authors select a given number of top heads and top hidden states from the network. Once the sorting and selection steps are done, authors re-group and reconnect the remaining heads and hidden states which result in a smaller sized model. When heads and hidden states are pruned, authors use the same pruning ratio across different layers. This enables further optimizations to work seamlessly with the pruned models.

In the experiments, authors found out that the pruned model can get better accuracy when it goes through another round of knowledge distillation. So knowledge Distillation is again applied to the model.

Model Quantization: Quantization refers to techniques for performing computations and storing tensors at lower bit widths than floating-point precision. A quantized model executes some or all of the operations on tensors with integers rather than floating-point values. This allows for a more compact model representation and the use of high performance vectorized operations on many hardware platforms.

8-bit quantized matrix multiplications on the CPU:8-bit quantized matrix multiplication brings a significant amount of speed-up compared to 32-bit floating-point arithmetic, thanks to the reduced number of CPU instructions.

16-bit model conversion for the GPU: V100 GPU supports full 16-bit operations for the Transformer architecture. Also, 16-bit floating-point operations do not require special handling of inputs and outputs except for having smaller value ranges. This 16-bit model conversion brings quite significant speed gain since the Transformer models are memory bandwidth bound workload. About 3.53x speed-up depending on the model settings was observed.

On top of the structural and numerical optimizations applied, authors also utilize various ways to further optimize the computations especially Multi-processing optimization and Computational graph optimizations.

Combined results

The table below speaks all how effective the results are given below