A Debiased Nearest Neighbors Framework for Multi-Label Text Classification

TL;DR

This work tackles MLTC with a kNN retrieval backbone by identifying embedding alignment and confidence estimation biases that hinder performance. It introduces DENN, a debiased framework combining Debiased Contrastive Learning to align embeddings with label co-occurrence and Debiased Confidence Estimation to adaptively fuse kNN and inductive predictions without adding parameters. Across four benchmarks (AAPD, RCV1-V2, Amazon-531, EUR-LEX57K), DENN achieves state-of-the-art results, driven by improved neighbor consistency and per-sample confidence. The approach yields robust gains across label frequencies and maintains practical efficiency, underscoring the value of debiasing in retrieval-augmented MLTC.

Abstract

Multi-Label Text Classification (MLTC) is a practical yet challenging task that involves assigning multiple non-exclusive labels to each document. Previous studies primarily focus on capturing label correlations to assist label prediction by introducing special labeling schemes, designing specific model structures, or adding auxiliary tasks. Recently, the $k$ Nearest Neighbor ($k$NN) framework has shown promise by retrieving labeled samples as references to mine label co-occurrence information in the embedding space. However, two critical biases, namely embedding alignment bias and confidence estimation bias, are often overlooked, adversely affecting prediction performance. In this paper, we introduce a DEbiased Nearest Neighbors (DENN) framework for MLTC, specifically designed to mitigate these biases. To address embedding alignment bias, we propose a debiased contrastive learning strategy, enhancing neighbor consistency on label co-occurrence. For confidence estimation bias, we present a debiased confidence estimation strategy, improving the adaptive combination of predictions from $k$NN and inductive binary classifications. Extensive experiments conducted on four public benchmark datasets (i.e., AAPD, RCV1-V2, Amazon-531, and EUR-LEX57K) showcase the effectiveness of our proposed method. Besides, our method does not introduce any extra parameters.

Figures (7)

• Figure 1: Illustration of two bias problems and corresponding solutions. The degree of greenness of each sample is proportional to the label similarity to the anchor. The yellow samples indicate that they do not have the same label as the anchor.
• Figure 2: Illustration of the Debiased Nearest Neighbors (DENN) Framework. The upper part denotes the training phase, where binary cross-entropy loss and debiased contrastive loss are used to train the model. The lower part represents the inference phase, which consists of three steps, building datastore, getting the two predictions (i.e., $\mathbf{\hat{y}_{knn}}$ and $\mathbf{\hat{y}_{clf}})$, and adaptive combining them with debiased confidence.
• Figure 3: Effects of the number of retrieved nearest neighbors $k$ on two datasets.
• Figure 4: Effects of $\gamma$ on two datasets.
• Figure 5: Effects of the size of datastore on two datasets.