In-Context Learning for Extreme Multi-Label Classification

TL;DR

This paper introduces a novel in-context learning approach for extreme multi-label classification that leverages a multi-step interaction between language models and retrievers, achieving state-of-the-art results without fine-tuning.

Contribution

The paper presents the $ exttt{Infer--Retrieve--Rank}$ program and the $ exttt{DSPy}$ model, enabling effective multi-label classification with minimal examples and no fine-tuning, adaptable to various datasets.

Findings

Achieved state-of-the-art results on three benchmarks (HOUSE, TECH, TECHWOLF).

Attained competitive performance on a diverse benchmark (BioDEX).

Requires only tens of labeled examples and no fine-tuning.

Abstract

Multi-label classification problems with thousands of classes are hard to solve with in-context learning alone, as language models (LMs) might lack prior knowledge about the precise classes or how to assign them, and it is generally infeasible to demonstrate every class in a prompt. We propose a general program, $\texttt{Infer--Retrieve--Rank}$, that defines multi-step interactions between LMs and retrievers to efficiently tackle such problems. We implement this program using the $\texttt{DSPy}$ programming model, which specifies in-context systems in a declarative manner, and use $\texttt{DSPy}$ optimizers to tune it towards specific datasets by bootstrapping only tens of few-shot examples. Our primary extreme classification program, optimized separately for each task, attains state-of-the-art results across three benchmarks (HOUSE, TECH, TECHWOLF). We apply the same program to a benchmark with vastly different characteristics and attain competitive performance as well (BioDEX). Unlike prior work, our proposed solution requires no finetuning, is easily applicable to new tasks, alleviates prompt engineering, and requires only tens of labeled examples. Our code is public at https://github.com/KarelDO/xmc.dspy.