Large Language Model as a Teacher for Zero-shot Tagging at Extreme Scales

TL;DR

This work tackles extreme zero-shot multi-label text classification by combining a strong but costly LLM-based labeling phase with a lightweight bi-encoder for inference. The LMTX framework uses an LLM as a teacher to curate high-quality pseudo labels from a label shortlist produced by ANNS, and then trains a DistilBERT-based bi-encoder with a triplet loss to align document and label embeddings in a shared space. It achieves state-of-the-art results on EZ-XMC benchmarks while maintaining practical inference efficiency, since the LLM is not required during prediction. The study further analyzes the impact of teacher models, training data size, and sampling strategies, highlighting curriculum learning as a key factor in improving pseudo-label quality and model performance.

Abstract

Extreme Multi-label Text Classification (XMC) entails selecting the most relevant labels for an instance from a vast label set. Extreme Zero-shot XMC (EZ-XMC) extends this challenge by operating without annotated data, relying only on raw text instances and a predefined label set, making it particularly critical for addressing cold-start problems in large-scale recommendation and categorization systems. State-of-the-art methods, such as MACLR and RTS, leverage lightweight bi-encoders but rely on suboptimal pseudo labels for training, such as document titles (MACLR) or document segments (RTS), which may not align well with the intended tagging or categorization tasks. On the other hand, LLM-based approaches, like ICXML, achieve better label-instance alignment but are computationally expensive and impractical for real-world EZ-XMC applications due to their heavy inference costs. In this paper, we introduce LMTX (Large language Model as Teacher for eXtreme classification), a novel framework that bridges the gap between these two approaches. LMTX utilizes an LLM to identify high-quality pseudo labels during training, while employing a lightweight bi-encoder for efficient inference. This design eliminates the need for LLMs at inference time, offering the benefits of improved label alignment without sacrificing computational efficiency. Our approach achieves superior performance and efficiency over both LLM and non-LLM based approaches, establishing a new state-of-the-art in EZ-XMC.

Figures (6)

• Figure 1: An example of how to construct training pairs using state-of-the-art methods MACLR xiong-etal-2022-extreme and RTS zhang-etal-2022-structural-contrastive. MACLR utilizes the 'Title' of a document to generate pseudo labels, while the 'Description' serves as the training document. Conversely, RTS forms its training pairs by selecting two random segments from the 'Description'. Differently, our proposed model, LMTX, adopts a more refined approach. It selects 'computer accessories' as a pseudo positive label from a predefined set, a choice validated by the LLM model.
• Figure 2: The process of getting feedback from LLM model for training the bi-encoder. First, for a given document, the (pre) trained bi-encoder and ANNS are employed to create a short-list of potential labels. Next, the LLM assesses the relevance between the labels in this shortlist and the document. Finally, the selected labels are utilized to further train the bi-encoder.
• Figure 3: Effect of training sample size on LMTX performance and training time.
• Figure 4: Comparative impact of negative sampling strategies on precision and recall performance.
• Figure 5: Impact of shortlist size on performance metrics for EURLex-4k dataset.