Accelerating ML Model Inferencing

This is not a unique problem in the travel industry. It can be generalized to verify a whole class of financial and legal agreements on an ongoing basis. Speed is of the essence; quickly detecting violations protects the platform’s viability and stops possible fraud.

We require real-time decisions from unstructured data but would a large language model be up to the task? We started with BERT, the open-source transformer-based machine learning model for Natural Language Processing(NLP). After running sample data, we discovered it could give us inferences reliably — but was too slow.

In fact, there was no way to get results within the constraints of our service-level agreement (SLA). Simulating each call with listing descriptions of 500 words or more, the model took up to two seconds. The speed was too slow for upstream services and didn’t meet the 250 milliseconds SLA we promised.

This was when we embarked on a journey to improve speed while maintaining inference reliability. We went straight to standard model compression techniques which are readily available and with which we were already familiar:

1. Knowledge Distillation (KD) Teacher-to-student transfer reduces model size.
2. Graph compression using Open Neural Network Exchange(ONNX) allows for data transfer to many different environments. It also has training capabilities.
3. Using specialized hardware like GPU/Xeon processors through AWS exchange

We used these techniques to improve speed without significantly impacting the model’s accuracy.

A Bit About Knowledge Distillation (KD)

In machine learning, Knowledge Distillation refers to transferring knowledge from a large model to a smaller one. The end goal is to develop accurate models that are more computationally efficient.

In a nutshell, we want to transfer knowledge from larger models (e.g., deep neural networks) to smaller models (e.g., shallow neural networks).

One solution is to use a technique called Knowledge Distillation. This involves “distilling” the knowledge from your large model into a new one that’s smaller, more manageable, and easier to work with.

The general idea of performing KD is:

• Use a pre-trained teacher model.
• Then, take some training data on which you want to train the Student model.
• Pass data through the Teacher model to generate soft labels (actual labels are “hard” labels).
• Train Student model on data + soft labels.

There are several ways to distill the knowledge of the Teacher model into a simpler Student model. The main techniques are Multi-teacher, Offline, and Online.

Multi-teacher approaches benefit from more diverse data sources and, in some cases, less biased data. For example, for offline KD, the Teacher is trained, and then the model is transferred directly to the Student. Conversely, online KD uses the ensemble result from multiple soft target Student models to generate a Student model.

Technique 1: Use DistilBERT Instead of BERT

• DistilBERT has the same general architecture as BERT but only half the layers of BERT.
• DistilBERT was trained on eight 16GB V100 GPUs for around 90 hours.
• KD is done during the pre-training phase, reducing the size of BERT by 40% while retaining 97% accuracy and becoming 60% faster.

Technique 2: Convert the Model to ONNX

• Convert the DistilBERT model to ONNX format.
• It is effectively a serialized format to represent the model and additionally functions as a model compression technique.

Technique 3: Use a Suitable AWS Instance Type

• Use inf1.6xlarge instance type to bring down latency further (published improvement of 2.3 X)
• 24 vCPU, 48 GiB memory
• 4 AWS Inferentia chips — dedicated to inference with a large onboard memory cache and speeds of 128 TOPS.
• Much better than g4dn.2xlarge(GPU) and c5–9xlarge(compute intensive) instance types.

Our Request/Response

Request

{
  "transactionContext": {
    "transactionId": "123456"
  },
  "attributes": {
    "listingDescription": "Near Xel Casino and the Cosmo Casino, the Tenebian Resort Las Vegas provides a casino, a nightclub, and a grocery/convenience store. For those looking to try their luck, this resort boasts a sportsbook, 137 casino gaming tables, and 7 casino VIP rooms. Indulge in a hot stone massage, a manicure/pedicure, and hydrotherapy at City Ranch Spa + Fitness, the onsite spa. Be sure to enjoy a meal at any of the 21 onsite restaurants, which feature Asian cuisine and happy hour. Yoga classes and aerobics classes are offered at the gym."
  }
}

Response

{
  "modelCode": "DistilBertNLPSupplyListingDescHA",
  "modelVersion": "1.0",
  "refreshVersion": "v20211220",
  "probabilityScore": "1.02e-10",
  "modelSegment": "SUPPLY",
  "topContributor": "listingDescription",
  "scoreDetails": {
    "label": "NonFraud"
  }
}

Numbers

|     Label    | Num requests |  Avg  | tp90 (ms) |  tp99  |  Min  |   Max  | throughput |
|--------------|--------------|-------|-----------|--------|-------|--------|------------|
| HTTP Request |      60      | 77.25 |   96.75   | 110.67 | 68.25 | 126.75 |   12.71    |

Notes:

• Averaged over 12 different runs using Apache Jmeter, free Java open-source software for load testing.
• Predictor service running on AWS ECS with 1 Fargate instance.

More Numbers

| Model           | Size   | Latency tp99 (top 99 %-ile) | Accuracy           |
|-----------------|--------|-----------------------------|--------------------|
| BERT            | 712 MB | >500 ms                     | 91.14              |
| DistilBERT      | 528 MB | 164 ms                      | 90.48              |
| DistilBERT+ONNX | 516 MB | 116 ms                      | same as DistilBERT |

• Instance type used: inf1.6xlarge
• Each input request: 512 tokens

The bottom line was that our work was successful. We met both our accuracy and our latency goals. We were able to improve latency by 95%(two sec to 110 milliseconds) and reduce the model size by 27% while impacting accuracy by only 1%.

Following this general procedure should work for a lot of common NLP challenges. As for us, we plan to use the three-step technique outlined above for various use cases we have today. Then it’s off to the next challenge.