Model Distillation for Query-Document Matching: Techniques and Optimizations

1. Introduction

Knowledge Distillation (KD) was introduced by Hinton et al. (NIPS 2014) to transfer knowledge from one or more teacher models to a lightweight student model. This article describes how we applied KD to a video query‑doc matching BERT model, achieving a 1‑layer lightweight BERT that is production‑ready.

2. Existing Solutions

Current KD methods such as TinyBERT and DistillBERT only compress the original 12‑layer BERT to at most 4 layers; further compression leads to severe AUC loss. To meet online latency requirements, we propose a series of optimizations.

3. Matching Model Details

The original model encodes a query‑doc pair with a BERT encoder, extracts the CLS token, max‑pooled and average‑pooled hidden vectors, concatenates them, and passes them through two linear layers with Tanh activation to produce a 1‑dimensional matching score.

During training, each triple (query, positive doc, negative doc) is converted into two pairs (query, positive) and (query, negative). Their scores are used to compute a hinge loss.

4. Distillation Framework

We fix the 4‑layer BERT teacher and distill its knowledge into a student model using a combined loss:

Soft loss (MSE) between student and teacher logits.

Hard loss (hinge) between student predictions and ground‑truth labels.

Relevance loss from a high‑performance GBDT teacher.

The overall distill loss is a weighted sum of soft and hard components (weights α and β), with AutoML searching for optimal values.

5. Loss Calculations

Soft loss uses MSE between student and teacher logits. Hard loss is a hinge loss with a threshold of 0.7. Relevance loss is an MSE between student logits and GBDT relevance scores.

Distill loss = α·soft_loss + β·hard_loss, with exponential scaling applied to accelerate convergence.

6. Student Model Optimization

We explored two lightweight student architectures:

ALBERT : Reduced to 1‑layer (1L‑ALBERT) with shared parameters, achieving lower latency and a 1.7% AUC gain over the original 4‑layer BERT.

TextCNN : A tiny CNN with Word2Vec embeddings and QQSeg tokenization, yielding 3.55 ms latency and slightly higher AUC than the 4‑layer teacher.

7. Better Teacher Guidance

In addition to the BERT teacher, we used a high‑performance GBDT ranking model as an auxiliary teacher. Its relevance scores are incorporated via the relevance loss, further improving student performance.

8. AutoML Hyper‑parameter Search

We employed AutoML on the Venus platform to search optimal hyper‑parameters (learning rate, loss weights, etc.) using a 6% data sample for 24 h. The best configuration was fine‑tuned on the full dataset, yielding an additional 0.6% AUC improvement.

9. Experimental Results

Both the 1L‑ALBERT and CNN‑attention students achieve latency below 5 ms and surpass the manually tuned 4‑layer BERT. The 1L‑ALBERT reaches 2.99 ms latency with a 2.4% AUC gain; the CNN model attains 3.55 ms latency with comparable AUC.

10. References

1. Hinton et al., “Distilling the Knowledge in a Neural Network”, NIPS 2014. 2. “Distilling Task‑Specific Knowledge from BERT into Simple Neural Networks”. 3. “ALBERT: A LITE BERT for Self‑Supervised Learning of Language Representations”. 4. “Transformer to CNN: Label‑scarce distillation for efficient text classification”.

Author

Wang Ruichen – Tencent Application Research Engineer, focusing on model compression, AutoML, and KD for CV/NLP.