Multi-Label Text Classification with Label Attention Aware and Correlation Aware Contrastive Learning

1 Introduction

Multi-label text classification (MLTC) is a critical task in natural language processing (NLP), where each document can be associated with multiple relevant labels simultaneously. This task underpins numerous real-world applications, such as document categorization [Rubin et al., 2012], news tagging [Zhang et al., 2019], and research topic classification [Yang et al., 2016]. Compared to single-label classification, MLTC introduces unique challenges due to the inherent complexities of label relationships and overlaps.

∗Corresponding author

One of the primary challenges in MLTC lies in the complex correlations among labels. Labels in MLTC often exhibit hierarchical, semantic, or co-occurrence dependencies. For example, an article discussing climate change may be tagged with "Environment," "Policy," and "Economics." These labels are semantically related, as climate change often involves policy discussions and economic impacts. Ignoring such intricate correlations risks producing incoherent predictions and undermines the interpretability of the model. Additionally, MLTC often faces the issue of partial label overlaps, where a document aligns strongly with some labels but only partially with others. For instance, an article on renewable energy may be highly relevant to "Environment" and "Energy" while weakly overlapping with "Economics" when financial aspects are discussed. These challenges are illustrated in Table 1.

Although existing methods address parts of these challenges, limitations remain. Classical approaches, such as binary relevance and label powerset, treat labels as independent entities, completely ignoring label correlations and overlaps [Guo et al., 2021; Wang et al., 2022]. More recent methods, such as label embedding and latent space learning, model label correlations to some extent [Ying et al., 2021; Xu et al., 2023], but fail to handle the ambiguity caused by partial overlaps. These approaches often treat overlapping labels as either fully dependent or entirely unrelated, leading to suboptimal representation learning and reduced generalization capabilities.

Motivated by these challenges, we propose a unified framework to explicitly address both complex label correlations and partial label overlaps in MLTC. Our solution is rooted in the observation that these two challenges are inherently interconnected and require complementary solutions. We designed the Unified Contextual and Label-Aware Framework (UCLAF), which integrates two key components: (1) the Label Attention Aware Network(LAN), which explicitly captures label dependencies by embedding text and labels into a shared semantic space; and (2) the Correlation Aware Contrastive Learning (CACL) , which introduces a dynamic and fine-grained mechanism for defining and optimizing relationships between samples based on their label sets. Unlike traditional methods[Zhang and Wu, 2024; Khosla et al., 2020] that rely on strict label matching (e.g., exact matches or any overlap), CACL considers the proportional overlap of labels between samples to determine their similarity and effectively mitigates the ambiguities caused by partial label overlaps and enhances representation learning in multilabel contexts. Two components form a complementary and synergistic design. Label Attention Module explicitly models label correlations by aligning text and label representations in a shared semantic space, CACL refines these representations by focusing on inter-sample relationships. This synergy ensures a comprehensive solution to the challenges of multilabel text classification. In summary, the main contributions of this work are as follows:

• We identify partial label overlaps as a key scientific problem in MLTC and analyze its impact on positivenegative sample distinctions and representation learning.
• We propose UCLAF, a unified framework that integrates label-aware representation learning and supervised contrastive learning to address partial label overlaps and model complex label correlations.
• Extensive experiments on benchmark datasets, including AAPD, EURLex demonstrate that UCLAF significantly outperforms state-of-the-art methods in multiple evaluation metrics.

2 Related work

2.1 Multi-Label Learning

Multi-label learning deals with instances associated with multiple interrelated labels and finds applications in domains such as object detection and anomaly detection [Ge et al., 2018; Cheng et al., 2025], image classification [Zhang and Wu, 2024], and text classification [Xiao et al., 2019]. For example, an image may contain several objects, or a text may cover multiple topics [Bogatinovski et al., 2022; Liu et al., 2017a; Liu et al., 2021; Wu et al., 2020]. Early methods approached this by treating multi-label problems as independent single-label tasks, which failed to model label correlations effectively [Boutell et al., 2004; Read et al., 2011]. Recent advances utilize statistical analysis [Wang et al., 2018; Yeh et al., 2017], label embedding techniques like RNNs and GNNs [Yazici et al., 2020; Wang et al., 2016; Durand et al., 2019], or auto-encoders to jointly encode label and instance correlations [Zhao et al., 2021; Bai et al., 2020; Bai et al., 2022]. Despite progress, most work focuses on non-text domains, leaving textual multi-label learning relatively underexplored.

2.2 Multi-label Text Classification

Existing multi-label text classification (MLTC) methods primarily focus on learning text representations and modeling label correlations. Early approaches used CNNs[Moschitti et al., 2014; Kurata et al., 2016] and RNNs[Liu et al., 2016] to capture long-range dependencies in text. We propose that focusing on each label's unique feature representation is a promising strategy for obtaining richer text representations for each label. For example, recent methods[Wang et al., 2023] have incorporated Transformer-based modules to exploit label semantic for capturing high-quality document representations, while long-range word dependencies are modeled through encoders, and multi-label attention mechanisms identify[Xu et al., 2023] the most relevant parts of the text for each label. In terms of label correlation modeling, previous approaches have employed sequence generation models[Nam et al., 2017; Yang et al., 2018; Xiao et al., 2021] and iterative inference mechanisms[Wang et al., 2021]. Additionally, Ma et al.[Ma et al., 2021] applied graph neural networks based on labeled graphs. However, these methods have yet to fully explore the latent semantic space between labels and texts, and the challenge of obtaining more effective feature representations remains a significant issue in MLTC.

In light of these challenges, our paper is specifically designed for multi-label text classification (MLTC). This approach aims to address the existing limitations of complex label correlations and partial overlaps through the integration of a unified framework.

3 Task Formulation

Consider a batch of data B = {(doci, yi)}Ni=1 , where N represents the mini-batch size, and yi = {y(i) j }Lj∈{0,1} denotes the multi-label set for sample i. Here, y(i) j indicates the j-th label of sample i, and L is the total number of labels in the dataset. The classifier computes the probability ˆ yi = {ˆ y1, ˆ y2, . . . , ˆ yL}, where each ˆ yj represents the likelihood of the corresponding label being true.

The binary cross-entropy (BCE) loss is used to measure the discrepancy between the predicted probabilities ˆ yi and the ground truth labels yi. It is defined as:

(1)

where yl and ˆ yl correspond to the ground truth and predicted probability for label l, respectively.

For the contrastive learning framework, which is based on the MoCo structure, we introduce additional notations. Let z(q) i and z(k) i denote the L2-normalized outputs of the query model and key model for sample i, respectively. The key model is updated using momentum-based updates. Furthermore, a queue Q is maintained to store z(k) from previous batches, following the MoCo[He et al., 2020] design.

4 Method

4.1 Label Attention Aware Network

To capture the dependencies between label embeddings and text features, the text data is first embedded into a vector space. Let the text embeddings be denoted as F = {f1, . . . , fN}, where fi ∈ Rd represents the embedding of the i-th text instance. The interaction between a label embedding lk and a text feature fi is quantified using an additive attention mechanism, producing a score sik:

where [a; b] denotes the concatenation of vectors a and b, and vb and Wb are trainable parameters of the attention mechanism. The computed scores sik are then normalized across all text instances using the SoftMax function to obtain attention weights βik:

These attention weights are subsequently used to generate a text-aware attention vector ck ∈ Rd by aggregating the text embeddings:

The attention vector ck encapsulates context-specific information from the text embeddings relevant to label lk. To model the joint relationship between text and label, the attention vector ck is concatenated with the label embedding lk as [lk; ck]. This concatenated representation is passed through two fully connected layers, which refine the features and compute the final output[Zhou et al., 2021]. The output scores for each text-label pair are transformed into probabilities using a sigmoid activation function, enabling multi-label predictions. By explicitly modeling text-aware label dependencies, this mechanism enhances the representation learning process, leading to improved performance in multi-label classification tasks.

5 Correlation Aware Contrastive Learning

In Multi-Label Text Classification (MLTC), identifying positive samples for the anchor point is complex due to the presence of multiple labels, unlike single-label classification where positive samples are clearly identified through exact label matching. Inspired by these preliminary insights[Zhang and Wu, 2024; Khosla et al., 2020], we propose the Correlation Aware Contrastive Learning (CACL), which optimizes the contrastive learning process for both "ANY" (partial label overlap) and "ALL" (exact label match) scenarios. The two scenarios for defining positive samples are: For the general case, we denote A(i) as the index of all samples involved in the loss calculation from both the batch and the queue. The supervised contrastive loss is defined as:

where s(i) p = z(i) q · z(p) k , τ is a temperature scaling parameter, and z(i) q and z(p) k represent the L2-normalized query and key embeddings, respectively. By taking the gradient of L(i) supcon

with respect to z(i) q , we calculate:

This gradient can be decomposed into two components:

where:

Here, ¯ z represents the mean of the positive sample embeddings, while ˆ z represents the weighted average of all embeddings, which prevents representation collapse. This ensures the optimization direction moves z(i) q towards a balanced representation of positive samples while considering all contextual dependencies. To further enhance the method, the proposed CACL loss treats the i-th sample as a distinct instance for each label y(i) j it belongs to. A separate positive sample set is constructed for each y(i) j ∈ y(i) as:

The CACL loss function is defined as:

This loss function generalizes the SupCon loss for multi-label tasks, reducing to the traditional SupCon loss in single-label scenarios where |y(i) | = 1. Figure 2 illustrates how CACL effectively balances contributions from partially overlapping ("ANY") and fully matching ("ALL") scenarios, ensuring robust optimization for multi-label classification tasks.

• ALL: Only samples with exactly the same label set as the anchor are considered positive:

• ANY: Samples with any overlapping label class with the anchor are considered positive:

5.1 D. Training Objective

The UCLAF is an end-to-end MLTC model that is made possible by the constructed label attention and CACL. The goal of UCLAF is to minimize the target loss function Lsum, which includes LCACL and LBCE(ˆ yidoc , yi ). The specific formula is defined as:

Here, λ is the tunable coefficient for the document-label contrastive learning loss function and the Correlation Aware Contrastive Learning (CACL) loss function, respectively, controlling the balance between the two losses. ˆ yi text represents the final probability that the semantic information of the i-th document is used as input.

6 IV. Experiments

In this section, we conduct a series of experiments on realworld text datasets to evaluate the performance of UCLAF and answer the following research questions:

Q1: Does UCLAF outperform current state-of-the-art (SOTA) methods? Q2: Does each module in UCLAF contribute to the overall model performance? Q3: How does correlation-aware contrastive learning operate and improve performance? Q4: How does the label attention network capture the relationships between labels and text? Q5: What is the impact of hyperparameters α and λ on UCLAF's performance?

6.1 Datasets

We evaluated our proposed model on two benchmark datasets: AAPD[Yang et al., 2018] and EURLex[Loza Menc´ıa and F¨urnkranz, 2008]. Table 2 presents the statistics for these datasets, where Ntrn represents the number of samples in the training set, L denotes the total number of labels, ¯L is the average number of labels per document, and ¯Wtrn represents the average number of words per training document.

6.2 Baseline Methods and EvaluationMetrics

To rigorously assess the effectiveness of the proposed UCLAF method, we compare it against several state-of-theart multi-label text classification (MLTC) models that have been widely recognized in recent research. The following baseline methods are selected for comparison:

• XML-CNN[Liu et al., 2017b]: A CNN-based model that employs a dynamic max pooling scheme to capture high-level features in the input text.
• AttentionXML[You et al., 2019]: A BiLSTM-based model that integrates a label tree structure, utilizing label attention to capture contextual semantics and generate label-aware document representations.
• CornetAttentionXML[Xun et al., 2020]: An architecture that combines the AttentionXML text encoder with a Cornet module to effectively model label correlations and dependencies in multi-label classification.
• LightXML[Jiang et al., 2021]: A lightweight deep learning framework that incorporates dynamic negative label sampling. For experimental consistency, we reproduce this model using BERT as the underlying architecture.
• GUDN[Wang et al., 2023]: Amodel that utilizes guided networks to enhance classification performance by improving the representation of labels.
• MLGN[Liu et al., 2023]: Amethod that integrates label semantic information to improve document representations but does not fully address complex label dependencies in multi-label tasks, nor does it effectively model label associations.

6.3 EvaluationMetrics

We evaluate model performance using two widely adopted metrics: precision at k (P@k) and normalized discounted cumulative gain at k (N@k)[You et al., 2019; Xiao et al., 2019],with higher values indicating better performance.

• P@k: Measures the proportion of relevant items among the top-k results.
• N@k: Considers both the relevance and ranking of retrieved items, rewarding relevant items ranked higher.

6.4 Implementation Details

All experiments were conducted on an NVIDIA 3090 GPU using the pre-trained bert-base-uncased model for feature extraction. The maximum input text length was set to 512, with a dropout rate of 0.5 for text representation, and a batch size of 16. For the EURLex dataset, the learning rate was set to 5 × 10-5 , with λ = 0.5 and τ = 3. For the AAPD dataset, the learning rate was set to 1×10-5 , with λ = 0.7 and τ = 5.

6.5 Main Results(Q1)

We evaluated the performance of the proposed UCLAF model against state-of-the-art methods using P@K and N@K as evaluation metrics. The results in Table 3 show that UCLAF achieves superior performance compared to existing models across all metrics. Specifically, UCLAF outperforms the second-best model in terms of P@1, P@3, P@5, N@3, and N@5 by significant margins. For the EURLex dataset, UCLAF improves P@1 by 1.01%, P@3 by 2.62%, P@5 by 0.24%, N@3 by 0.69%, and N@5 by 0.40% compared to the second-best model, which is AttentionXML. Similarly, for the AAPD dataset, UCLAF improves P@1 by 0.59%, P@3 by 0.39%, P@5 by 0.28%, N@3 by 0.54%, and N@5 by 0.53% compared to the second-best model, which is MLGN. These results demonstrate that UCLAF consistently leads to better classification performance, confirming its effectiveness in multi-label text classification tasks.

7 ANALYSIS

7.1 Ablation Study (Q2)

We conducted ablation experiments to evaluate the impact of the Label Attention Aware Network (LAN) and CorrelationAware Contrastive Learning (CACL) modules in UCLAF, with results shown in Table 4. Adding the LAN module improves performance by leveraging label semantics. For example, P@1 increases from 85.30 to 85.58 on AAPD and from 84.55 to 85.42 on EURLex. Similarly, incorporating the CACL module enhances performance by capturing label correlations, as seen in the rise of P@5 from 41.18 to 42.24 on AAPD and P@3 from 73.05 to 74.13 on EURLex. The combination of both LAN and CACL achieves the best results across all metrics, with P@1 reaching 86.69 on AAPD and 87.32 on EURLex. These results demonstrate that LAN

and CACL complement each other effectively, jointly contributing to significant improvements in label alignment and inter-label modeling.

7.2 Analysis of Label Attention Aware (Q3)

The heatmap in AAPD dataset demonstrates the Label Attention Module's ability to capture semantic relationships between texts and labels by aligning them in a shared semantic space. For example, Label 1 shows strong attention on Text 3 and Text 5, indicating effective focus on relevant text segments. Overlapping attention patterns, such as for Label 2 and Label 3 on Text 4 and Text 6, highlight the module's ability to model label correlations, capturing co-occurrence and semantic dependencies between labels.

However, limitations exist. For example, Text 7 consistently receives low attention weights, suggesting weak alignment or poor feature representation. Similarly, Label 4 shows a uniform attention distribution, reducing precision in distinguishing relationships. These results confirm the module's effectiveness in addressing complex label correlations while highlighting areas for further improvement.

7.3 Correlation Aware Contrastive Learning Experiments(Q4)

To evaluate the effectiveness of CACL, we conducted a series of experiments. First, we examined the impact of positive and negative sample pair construction on model performance. Second, we visualized the learned sample representations to provide qualitative insights. Finally, we investigated the influence of different data augmentation strategies on the model's overall performance.

Effects of Positive Sample Generation

As shown in Table 5, CACL significantly outperforms the "ANY" and "ALL" match strategies on the AAPD dataset.

It achieves the highest scores across all metrics, with a P@1 of 87.32, compared to 85.46 (ANY) and 84.32 (ALL). Similar improvements are observed in P@3 and P@5. The key strength of CACL lies in its flexible positive sample selection, associating samples with others that share a subset of their labels. This strategy better suits multi-label tasks by capturing label correlations, leading to more effective representation learning and performance gains.

Performance of Different Data Augmentation Techniques

Table 6 presents a comparison of different data augmentation strategies, including random masking, continuous masking, token shuffling, and the dropout strategy. Among these, the dropout strategy achieves the highest performance, indicating its effectiveness in generating positive samples by randomly removing neuron-level information. Random masking and continuous masking yield comparable results, showing no significant differences in their impact on model performance. Conversely, token shuffling results in the lowest scores, likely due to its disruption of text semantics, which hinders the model's ability to recover meaningful information. These findings underscore the critical role of maintaining semantic structure in data augmentation to ensure effective learning.

Visualization of Learned Representation

To further investigate the interpretability of the high-quality document representations in UCLAF, this section provides a visual analysis of the document representations in UCLAF. We employ t-SNE to visualize the category prototypes, as shown in Figure 4. The results demonstrate that the category prototypes capture meaningful semantic information. Specifically, the distance between different categories is maximized, while the distance between pairs of highly cooccurring categories is minimized, and pairs of categories with low co-occurrence are more distantly separated.

(b) Ours(AAPD)

(c) BCE(EURLex)

Table 6: Performance of different positive example generation techniques on the AAPD dataset.

7.4 Hyper-parameter Analysis(Q5)

Sensitivity Analysis of hyper-parameter threshold λ

λ controls the balance between the contrastive learning loss LCACL and the binary cross-entropy loss LBCE. As is shown in Figure 5(a)(b), when λ is too large, the model overly emphasizes label correlations, which may reduce the performance of individual label classification. Conversely, when λ is too small, the model focuses more on independent label classification but may overlook global semantic relationships between labels. Therefore, λ needs to be adjusted experimentally to balance label correlation modeling and text-label matching performance.

Sensitivity analysis of the hyperparameter τ

The parameter τ controls the sharpness of the similarity distribution in contrastive learning. As is shown in Figure 5(c)(d), when τ is too large, the similarity scores become too uniform, reducing the model's ability to distinguish between positive and negative samples, which weakens its performance. Conversely, when τ is too small, the similarity distribution becomes overly sharp, causing the model to overfit a few high-similarity positive samples while ignoring others, leading to poor generalization. Therefore, τ needs to be carefully tuned to balance the focus on high-quality positive samples and the overall similarity distribution, typically through validation experiments.

8 Conclusion

We proposed the Unified Contextual and Label-Aware Framework (UCLAF) to tackle challenges in multi-label text classification, such as label interdependencies and overlaps. By combining label attention and supervised contrastive learning, UCLAF captures semantic relationships and reduces label ambiguity. Experiments on AAPD, EURLex show that UCLAF outperforms state-of-the-art methods . However, UCLAF relies on static label embeddings, struggles with large-scale datasets due to contrastive learning's computational cost, and is limited to text data. Future work could explore dynamic label embeddings, improve efficiency, and extend UCLAF to multi-modal or hierarchical tasks for greater scalability and versatility.