QUEST : Efficient Extreme Multi-Label Text Classification with Large Language Models on Commodity Hardware (2024)

Abstract: Extreme multi-label text classification (EMTC) involves predicting multiple labels from a vast pool of candidates based on a user’s textual query. While traditional BERT-based methods have shown limited success, large language models (LLMs) have brought new possibilities. It is promising to leverage their remarkable comprehension ability to understand textual queries. However, implementing LLMs is non-trivial for two main reasons. Firstly, real-world EMTC datasets can be extremely large, with candidate product pairs reaching up to ten million in real-world scenarios, which poses significant challenges in data ingestion. Secondly, the large size of LLMs makes computation and memory demands prohibitive for EMTC applications. To this end, we propose QUEST, a Quantized and Efficient Learning with Sampling Technique. QUEST includes a tailored hash sampling module that reduces the data volume to one-fourth of its original size. Additionally, we perform compressive fine-tuning LLMs with only twenty thousand trainable parameters, largely reducing computational requirements. Extensive experiments demonstrate that QUEST outperforms existing methods while requiring fewer computational resources, unlocking efficient EMTC on commodity hardware such as a single Nvidia RTX 3090 GPU with 24 GB of memory.

## introduction extreme multi-label text classification (emtc) has been widely applied in real-world scenarios, e.g., recommendation systems and social networks @xcite , where each input query needs to be assigned multiple output labels from a very large set of candidate labels. as shown in figure 1 , the input consists of text descriptions of purchased items, and the model is required to infer potential products from a vast list. due to * xiao huang is the corresponding author. the complexity brought by its scale, emtc has been extensively studied in natural language processing from various perspectives @xcite @xcite . the existing standard method of emtc is training a deep model to learn representations for input queries and labels. this process transforms tasks into a similarity search problem, where for each query, we search for its nearest neighbors in the label space. early studies design linear models defined on tf-idf feature vectors @xcite . recently, plenty of research has been dedicated to sentence embeddings powered by bert-based models with million-scale parameters @xcite . however, performance remains unsatisfactory due to its limited semantic understanding ability. with the emergence of large language models (llms), which have demonstrated impressive inference capabilities @xcite @xcite , we are motivated to incorporate llms into emtc to better understand the semantics in both queries and labels. however, this task remains challenging for two reasons. directly using llms to process extremely large training data is time-consuming, making it difficult to apply in real-world scenarios. second, the extensive size of llms requires significant computational resources @xcite , with a full llama-7b model consuming 28gb of memory. moreover, the pace of our hardware enhancements is decelerating @xcite . current gpu resources can handle lighter frameworks like bag-of-words and bert-based encoders, but finetuning a full llm on a single machine is impossible. furthermore, the average number of labels per query can reach up to 22.2 @xcite , further increasing the training pairs. processing these extensive texts poses a heavy burden, necessitating the employment of multiple gpus for over a week. to this end, we present a novel quantized and efficient learning framework with sampling technique, i.e., quest, for extreme multi-label text classification, taking into account both memory usage and computation resources. we significantly reduce the search space of a large number of labels and facilitate efficient model training tailored for llms. (i) we introduce a novel hashing-based sampling module that selects more representative and diverse texts as training pairs for contrastive learning. this is inspired by the observation that textual descriptions of labels are highly clustered. we achieve promising performance by using only a small proportion of the original data. (ii) the model is optimized through a prompt-based pairwise contrastive training algorithm on the quantized llm. this effectively reduces the number of trainable parameters and memory usage. contributions: • we formally define the paradigm of efficient learning for emtc with llms, addressing both memory usage and computation resources. ## preliminaries notations. the dataset includes queries and labels, with each consisting of the title and content of the item description. in this paper, @xmath0, c 1 ), ..., (t |x| , c |x| )} is denoted as a collection of input documents, where each pair (t i , c i ) represents the title and content respectively. we denote @xmath1, ..., y |y| } as the set of labels, each of which is characterized by a piece of sentence. the original documents are split into the training dataset x train and the validation dataset x test . different from ordinary classification tasks that only consider a limited number of labels @xcite , the presence of millions of distinct choices makes using softmax classifiers difficult. therefore, these issues are transformed into a nearest-neighbor search problem, where the model seeks to match the representations of queries with relevant tags. another important aspect is its zeroshot capability. in reality, numerous newly released items are introduced to customers, making it impractical to have access to all potential labels during the training process. providing accurate recommendations for unseen products holds great realistic significance. ## overall framework as mentioned in the introduction, this paper aims to achieve a more efficient and faster learning process. the approaches are mainly divided into two lines: increasing batch size and reducing parameters to improve model efficiency; and using diverse sampling to increase data efficiency. in this section, we will offer an overview of the motivation and intuitive explanations for each module. ## experiments we conduct extensive experiments on four real-life datasets. our study aims to address the following research questions: q1: can quest achieve superior prediction performance than current stateof-the-art frameworks? q2: how effective is hash sampling in reducing dataset size without sacrificing performance and consuming excessive time, compared to random sampling and other clustering techniques? q3: how quickly can quest address the emtc issue in comparison to existing baselines, and to what extent can quest be compressed by utilizing various quantization methods? ## conclusions in this paper, we validated the effectiveness of large generative language models in emtc compared to traditional methods like bag-of-words vectors and bert-based models. given the fact that many users cannot afford expensive computational resources, all of our study revolves around this motivation. we employ prompt-tuning, quantization, and hash sampling to mitigate the side effects of the extremely large model size and vast training pairs, enabling training on a single machine within a limited time. the quest model manages to handle large classification tasks and comprehensive experiments confirm the efficacy of each module.