Knowledge Distillation for Enhancing Walmart E-commerce Search Relevance Using Large Language Models

TL;DR

This paper introduces a knowledge distillation framework that transfers the high ranking capabilities of large language models into a smaller, efficient model for e-commerce search relevance, achieving improved performance and low latency.

Contribution

The paper presents a novel distillation method that expands training data with unlabeled data and trains a student model to outperform the teacher LLM in search relevance tasks.

Findings

Student model performance improves with more augmented data.

The student model can outperform the teacher model with sufficient data.

Successfully deployed in Walmart's production system with positive metrics.

Abstract

Ensuring the products displayed in e-commerce search results are relevant to users queries is crucial for improving the user experience. With their advanced semantic understanding, deep learning models have been widely used for relevance matching in search tasks. While large language models (LLMs) offer superior ranking capabilities, it is challenging to deploy LLMs in real-time systems due to the high-latency requirements. To leverage the ranking power of LLMs while meeting the low-latency demands of production systems, we propose a novel framework that distills a high performing LLM into a more efficient, low-latency student model. To help the student model learn more effectively from the teacher model, we first train the teacher LLM as a classification model with soft targets. Then, we train the student model to capture the relevance margin between pairs of products for a given query using mean squared error loss. Instead of using the same training data as the teacher model, we significantly expand the student model dataset by generating unlabeled data and labeling it with the teacher model predictions. Experimental results show that the student model performance continues to improve as the size of the augmented training data increases. In fact, with enough augmented data, the student model can outperform the teacher model. The student model has been successfully deployed in production at Walmart.com with significantly positive metrics.