ICXML : An In-Context Learning Framework for Zero-Shot Extreme Multi-Label Classification (2024)

Abstract: This paper focuses on the task of Extreme Multi-Label Classification (XMC) whose goal is to predict multiple labels for each instance from an extremely large label space. While existing research has primarily focused on fully supervised XMC, real-world scenarios often lack supervision signals, highlighting the importance of zero-shot settings. Given the large label space, utilizing in-context learning approaches is not trivial. We address this issue by introducing In-Context Extreme Multi-label Learning (ICXML), a two-stage framework that cuts down the search space by generating a set of candidate labels through in-context learning and then reranks them. Extensive experiments suggest that ICXML advances the state of the art on two diverse public benchmarks.

## introduction extreme multi-label classification (xmc) deals with the classification of instances into a set of relevant labels from a large label set @xcite @xcite . it finds applications in various domains, including text categorization @xcite , recommendation systems @xcite , image tagging @xcite , and so on. unlike conventional multi-label classification, where the number of labels is relatively small, xmc involves an exponentially larger label space, e.g., in the 10 6 magnitude. this poses significant computational and modeling challenges. while existing research has primarily focused on supervised xmc, real-world applications often encounter challenges in obtaining complete supervision signals. scenarios arise during test sessions when new labels emerge without any assigned input instances @xcite , or when both instances and labels are available, but the corresponding relations between them are unknown @xcite . this task, which is the focus of this paper, is called zero-shot xmc. zero-shot xmc can be seen as a retrieval problem, where the test instance is considered as the query and candidate labels are retrieved in response to the given input. methods based on lexical matching, such as tf-idf @xcite and @xcite , and semantic matching, such as dense retrieval models @xcite , can be adopted for this task. state-of-the-art approaches for zero-shot xmc, such as rts @xcite and maclr @xcite , also belong to this category. a major shortcoming with these approaches is that there is little lexical or semantic overlap between the test instance (i.e., queries) and the label space (i.e., documents). one may argue that large language models (llms) can be used to generate labels for each test input. however, the labels that llms generate may not be in the acceptable label set and unlike conventional classification tasks, the label set is too large to be given to the llms in their prompt. therefore, using llms for this task is either impractical or extremely expensive. in this work, we put together the benefits of both retrieval-and generation-based approaches by introducing icxml-a two-stage framework designed for zero-shot xmc. in the first stagegenerate, icxml enriches the intrinsic capabilities of large language models by generating and/or retrieving demonstrations in a zero-shot manner. the obtained outputs are generated by prompting the model with the help of a support set of generated demonstrations. subsequently, the generated outputs are adjusted to align with the label space, resulting in a condensed shortlist of candidate labels. in the second stage -rerank, we leverage the capabilities of llms to perform multi-label classification by reintroducing the refined candidate label shortlist along with the test instance as input. this approach capitalizes on the language model's inherent potential to handle multiple labels concurrently, thereby augmenting its performance in the context of extreme multi-label classification tasks. in summary, our main contributions are three fold: our implementation scripts and codes are publicly available for research purposes at https: //github.com/yaxinzhuars/icxml . ## problem formulation let x and y respectively denote the input and output spaces. in this work, we focus on text data, thus each @xmath0. given the focus on xmc, the output space is extremely large, e.g., |y| ∼ 10 6 . the goal is to map each input x to a small subset of labels y ⊂ y. label reranking content-based label centric retrieval-based icl-based inference label space input instances demonstration generation label space mapping candidate label shortlisting figure 1: an illustration of the proposed generate-rerank framework. for a given test input xi, we generate demonstrations di to facilitate icl-based shortlisting. subsequently, this shortlisted set ȳi is provided to llm for listwise re-ranking, culminating in the final results y * i . following xiong et al. (2021); zhang et al. (2022a), we also consider a scenario where some input instances, called the input corpus, are available, however the mapping {(x i , y i )|x @xmath1, y i ⊆ y} is not available for training. ## the icxml framework benefiting from the high performance and rich knowledge encoded within llm parameters, we can easily generate a set of labels by describing classification tasks and inputs using prompts. however, for a classification task with a predetermined label set, this approach results in uncertain outcomes in the absence of guidance from fewshot example pairs, referred to as demonstrations @xcite . while common approaches involve incorporating the label candidate list into the prompt, this approach becomes impractical or extremely expensive when faced with an extreme label space, e.g., a label space of 10 6 magnitude. to address this issue, we propose a two-stage framework illustrated by 1. the first stage is generate, including demonstration generation (section 4.1) and candidate shortlisting (section 4.2) in the figure. we perform in-context learning using llm ϕ to generate labels, and the pseudo demonstrations for this stage is generated by ϕ using a prompt-guided approach. the second stage is rerank, where we utilize another prompt-guided method and ϕ for selecting top labels from candidate labels as decribed in section 4.3. ## expanding icxml by utilizing an input corpus our methodology adopts a novel approach that avoids the use of training samples, setting it apart from conventional models. this design choice aligns with zero-shot learning paradigms where, although a corpus of training instances x train and their corresponding labels may be available, they are not utilized in a paired manner. our framework maintains a degree of adaptability and can be readily extended to a more flexible setting by substituting demonstration generation with demonstration retrieval, when x train is available. for each test input x @xmath2, we select the top m neighbor instances from x train as follows: subsequently, for each z j @xmath3, we proceed with the methodology delineated in section 4.1, titled "content-based demonstration generation", to construct d i . ## experiments 6.1 data we evaluate the effectiveness of our approach on the following large-scale datasets: lf-amazon-131k in item recommendation domain and lf-wikiseealso in wikipedia articles title tagging domain, where 131k and 320k denote the size of label space @xcite . the dataset statistics are presented in