Mutual Information Guided Reinforcement Learning for Ambiguous Label Disambiguation

1 Introduction

Learning with ambiguous labels has become an increasingly important problem in modern machine learning, arising naturally in scenarios where precise annotations are difficult, costly, or even impossible to obtain. Therefore, partial label learning (PLL) [Cour et al., 2011; Fan et al., 2026a] has attracted much attention to address the impact of ambiguous labels. PLL is a weakly supervised learning framework, where each instance is associated with a candidate label set that contains the unknown ground-truth label but also includes multiple irrelevant labels. Although PLL can significantly reduce annotation cost, accurately identifying groundtruth labels from candidate label sets to reduce the impact of ambiguous labels poses a fundamental challenge to model training.

∗The corresponding author.

Intuitively, the basic method for accomplishing the PLL task is to perform disambiguation. To achieve this purpose, label confidence in PLL has been investigated according to two general strategies, i.e., non-deep disambiguation strategy (NDS) [Berg et al., 2004; Luo and Orabona, 2010; Chai et al., 2019; Fuchs et al., 2024] and deep disambiguation strategy (DDS) [Yao et al., 2020; Fan et al., 2026b; Fan et al., 2021; Li et al., 2025]. For NDS, attempt to progressively identify the true label by iteratively updating disambiguation labels. DDS incorporates some training objectives compatible with stochastic optimization to distinguish the true label from candidate labels, such as adding regularization terms, data augmentation, and adjusting loss weight. Although existing PLL algorithms have achieved certain success, they still have limitations in relying on static or heuristic update rules, implicitly assuming that label correction is a deterministic process driven only by local loss minimization. Therefore, these methods often lack the ability to actively explore different possibilities and are prone to accumulating errors, especially in training iterations that reinforce the impact of early error disambiguation decisions.

To tackle this issue, a novel mutual information guided reinforcement learning framework for partial label disambiguation (MR-PLL) is proposed to address PLL from a decision-theoretic perspective. Unlike existing PLL approaches that treat label correction as a deterministic optimization problem, we formulate the disambiguation process as a reinforcement learning problem in a Markov Decision Process (MDP) [Sutton et al., 1998]. Recently, Reinforcement Learning (RL) [Kaelbling et al., 1996] has provided a principled framework for modeling sequential decision problems, where agents interact with the environment and learn optimal policies by maximizing long-term cumulative rewards. Therefore, RL can dynamically adjust its ambiguous labels correction strategy based on its behavioral feedback and optimize performance through a series of decisions. However, its potential for resolving ambiguous labels in PLL remains largely underexplored. A core challenge lies in designing meaningful reward signals that can effectively guide the label disambiguation process without access to groundtruth labels.

To address the above limitations, this article designs an information weighted reward term for MR-PLL as the core component of RL. The reward integrates classification feedback, label consistency, and the mutual information between feature representations. Mutual Information (MI) provides a measure of similarity between representations of different classes, encouraging agents to promote feature alignment between the same labeled classes and avoid unreliable corrections. Meanwhile, MI estimation can be used to estimate the uncertainty between instances, which is used to adjust the value updates of RL agents. Specifically, the higher the uncertainty of the estimation, the smaller the corresponding update, which avoids ambiguous labeling leading to erroneous updates (values or policy functions) during the learning process, thereby making the agent unstable and forgetting all valuable states previously seen (catastrophic forgetting). Meanwhile, our method utilizes an instance selection-based discriminator to dynamically discriminate whether to retain, modify, or abstain from correcting ambiguous labels. The abstention mechanism is introduced the agent to temporarily withhold highly ambiguous instances from classifier updates, thus mitigating bias accumulation from ambiguous labels. To further enhance robustness, we propose a MI guided soft label fusion strategy. It adaptively blends original candidate labels with propagated pseudo-labels based on each instance reliability estimates. This design lowers gradient variance and stabilizes training, especially in early stages when model predictions are uncertain. The resulting framework unifies RL, information theory, and PLL into an end-to-end optimization pipeline. Furthermore, we provide a theoretical analysis of convergence and bias reduction. The contributions of this work are summarized as follows:

• We formulate PLL as a reinforcement learning problem, providing a novel decision-theoretic perspective on ambiguous supervision.
• We develop an information weighted reward term that encourages effective label disambiguation by aligning classification feedback, label consistency, andmutual information.
• A MI guided soft label fusion strategy and an abstention mechanism are proposed to adaptively balance uncertainty and bias control in label correction.

2 RelatedWork

2.1 Partial Label Learning

To alleviate the impact of ambiguous labels on classifier learning, two general strategies based on machine learning have been proposed for PLL, i.e., non-deep disambiguation strategy (NDS) and deep disambiguation strategy (DDS). The NDS operates on the principle that similar instances tend to share the same label. It resolves label ambiguity by leveraging the relationships among the candidate labels of neighboring instances[Chai et al., 2019]. A wide range of machine learning techniques have been employed to tackle PLL challenges, including k-Nearest Neighbors (KNN)[Luo and Orabona, 2010; Zhang et al., 2016], Expectation Maximization (EM) [Liu and Dietterich, 2012], Low- Rank Representation (LRR) [Fan et al., 2022; Chen et al., 2013], and Support Vector Machine (SVM) [Chai et al., 2019; Fan et al., 2023]. The progress in deep neural networks has motivated significant efforts to develop deep learning-based frameworks for PLL. DDS incorporates some training objectives compatible with stochastic optimization to distinguish the true label from candidate labels, such as adding regularization terms [Wang et al., 2023; Feng et al., 2025; Fan et al., 2026c], data augmentation [Chai et al., 2019], and the weight of the loss function [Fan et al., 2026a; Wen et al., 2021; Fan et al., 2021], have been proposed to deal with ambiguous labels. Instance-dependent PLL offers a more realistic framework for ambiguous supervision. To tackle this challenge, methods such as contrastive learning, Bayesian models, and variational inference have been developed [Xu et al., 2021; Xu et al., 2024; Wu et al., 2024].

The drawback of PLL is that it relies on static disambiguation heuristics, which are prone to error propagation when the ambiguity labels are high. Label disambiguation is inherently a dynamic process, early decisions directly affect the feature representations, which in turn constrain future label predictions. This creates a vicious cycle where an initial labeling error can be reinforced by subsequent updates, consequently causing severe and irreversible performance decline.

2.2 Reinforcement Learning

Reinforcement learning (RL) is a learning paradigm in which an autonomous agent acquires decision-making strategies through continuous interaction with an environment, optimizing its behavior over a sequence of actions to achieve longterm objectives [Sutton et al., 1998]. A defining strength of RL lies in its capacity for effective exploration [Amin et al., 2021] and its ability to optimize policies with respect to delayed and cumulative rewards [Xu et al., 2023]. Owing to these advantages, RL has been widely adopted in realworld environments, especially in robotics-relevant applications [Kober et al., 2013]. Preference-based reinforcement learning (PbRL) further extends this paradigm by enabling agents to learn directly from human preference feedback over pairs of trajectory segments, thereby alleviating the need for explicitly designed reward functions [Christiano et al., 2017]. More recently, PbRL has demonstrated remarkable success in the fine-tuning of Large Language Models (LLMs) [Ouyang et al., 2022], where human preferences provide a scalable and flexible alternative to manually specified reward signals. Unlike supervised learning, RL is a key artificial intelligence paradigm that enables an agent to learn optimal behavior through trial and error, guided by rewards. Therefore, the advantage of RL is that it is more suitable for labeling environments where data is scarce or unavailable.

3 MR-PLLMethod

We propose a novel MI guided RL for ambiguous label disambiguation. The PLL is formulated as the action execution in the MDP, where an agent dynamically decides whether to retain, modify, or abstain from correcting ambiguous labels. Designing a suitable reward function plays a crucial role in building reinforcement learning models for real-world applications in the face of ambiguous labels in PLL. It is difficult to design a reward function that produces credible rewards in the presence of ambiguous labels. This is because when the output of the reward function receives ambiguous feedback from candidate labels, it can lead to biased rewards for the classification model. To address this issue, we design an information weighted reward for actions to guide the action decision process. Meanwhile, the MI guided soft label propagation mechanism is adopted to train the disambiguation policy. The overall framework is illustrated in Figure 1, and the detailed components of the proposed method are introduced in the following subsections.

3.1 ProblemFormulation

Let X = Rd denote a d-dimensional feature space and Y = {yc : c = 1, . . . ,Q} denote the output space with Q class labels. The PLL training dataset D = {(xi, Si)|1 ≤ i ≤ N} of N instances contains feature vectors xi ∈ X and a set of candidate labels Si ⊆ Y. |Si| represents the number of ambiguous labels for instance xi. Particularly, instance xi is annotated by a label vector yi = {yc}Qc=1 , where yic ∈ {0, 1} denotes whether the label yc is present in the candidate labels ('1') or the non-candidate labels ('0'). The objective is to remove the influence of false-positive labels and learn an effective prediction model fθ ◦ gψ : X ↦→ Y that can make correct predictions on unseen inputs, where gψ denotes the feature extractor with parameters ψ, and fθ represents the classifier with parameters θ.

3.2 RL for Ambiguous Label Disambiguation

RL problem is often modeled as a MDP, characterized by the tuple (S,A, π,P,R), where S denotes the state space, A represents the action space, P(s′ | s, a) defines the transition dynamics for s, s′ ∈ S and a ∈ A, and R is a reward function. The goal of an RL agent is to learn a policy π that maximizes the expected cumulative reward.

State Representation

For each instance xi, the state si includes the potential information of an instance to be labeled, including both representation-level, prediction-level, and uncertainty-aware information. Specifically, we define the state as,

where zi = gψ(xi) is the feature representation extracted by the backbone network, pi = softmax(fθ(zi)) is the predicted label probability, H(pi) = - ∑Qc=1 pic log pic is the predictive entropy, max(pi) is the maximum confidence score, and ∆(pi) denotes the margin between the top-2 predicted probabilities. This design provides a sufficient statistic capturing discriminability, uncertainty, and learning dynamics, which is critical for reliable policy learning to select action.

Action Space

Due to the PLL setting, our agent is required to determine the label of each class for instance. To avoid the influence of ambiguous labels, we define the action space A corresponding to three possible disambiguation decisions, such that for each action is defined as ai = {0, 1, 2}. ai = 0 represents keep the original label distribution, ai = 1 indicates the current label needs to be corrected, and ai = 2 indicates that uncertain instance need to be postponed for disambiguation to avoid updating the classifier with this instance. It allows the policy to explicitly model uncertainty and prevents premature and potentially erroneous label corrections.

MI Guided Ambiguous Label Disambiguation

1) KNN Attention for Label Refinement: To mitigate the noisy labels influence, we through KNN attention to identify potential ground truth labels by utilizing complementary information between instances and candidate labels, making the model more resistant to ambiguous labels. The underlying assumption is that a label consistent with those of its k-nearest neighbors is more likely to have the same label. To this end, we aggregate the labels of its k-nearest neighbors in the current state s using an attention mechanism to generate a new label prediction ˆyi, which can be computed as follows,

where Nk(xi) denotes the KNN set of instance xi based on the Euclidean distances calculated in the embedding space, calculated by the backbone network gψ, and the attention weights αij is computed using similarities of instance pairs in the embedding space,

where τ is a temperature hyperparameter, and sim(·, ·) denotes the cosine similarity.

2) Mutual Information Estimation: A critical challenge in ambiguous label disambiguation is determining when such propagation is reliable. To this end, we introduce MI between positive and negative feature representations as a principled confidence measure. Then, we use variational f-divergence estimation [Poole et al., 2019] to maximize MI, which can be calculated by,

where Tω(·) is the density estimate, P denotes the joint probability distribution, and Q denotes the marginal probability distribution. z+i and z-i are the positive and negative instance representations of zi, where z+i selects from the knearest neighbors in zi, and z-i selects based on the intersection of the candidate labels sets between the instances being empty. While maximizing MI, by optimizing the relationship between the joint distribution and the marginal distribution, we significantly enhance the similarity between instances and their positive counterparts while reducing the similarity between instances and negative counterparts. A good representation can effectively improve the disambiguation ability of the model. This design aims to achieve tight aggregation of intra-class instances and effective separation of inter-class instances, thereby enhancing the discriminative and generalization performance of the model.

3) MI Guided Soft Label Fusion: A low mutual information between representations indicates weak semantic alignment and unreliable supervisory signals, suggesting that the label is more likely to be noisy and in need of correction. Therefore, we propose an MI guided soft label fusion,

where Normalize(·) is normalization function and η ∈ (0, 1) is the balance parameter. βi is defined as βi = σ(Ii -E[I]), where σ(·) is the sigmoid function. Detailed derivations of MI guided soft label fusion can be found in Appendix A.1.

4) Deterministic Transition with Stochastic Policy: Given state si, we determine the action ai using the stochastic policy function to enhance the exploratory behavior of the learning process, the label distribution evolves as,

where ˜yi denotes the propagated soft label. Notably, for the abstention action (ai = 2), the classifier parameters are not updated using this instance, ensuring that uncertain supervision does not affect representation learning. This mechanism maintains label uncertainty through soft label distributions and updates them via probabilistic corrections guided by noise likelihood. The introduced stochasticity promotes broader exploration of the state-action space, helping the model escape suboptimal solutions. Meanwhile, gradual label updates reduce instability caused by abrupt label changes, allowing the model to adapt smoothly to evolving learning dynamics and ultimately improving robustness under noisy supervision.

Actor-Critic Framework

MR-PLL adopts an actor-critic framework to learn the policy function by maximizing the expected cumulative reward. In this framework, the policy function πψ is considered as an "actor", where policy network is parameterized by ψ and outputs the probability πψ(a | s) for each action condition on the current state. The value function Vϕ(s) and is parameterized by ϕ is introduced as a "critic" to directly estimate the Q-value.

1) Reward Function: In RL, the reward function assesses the quality of an action a taken in a given state s and directs the learning process toward the desired objective. To achieve the goal of eliminating ambiguity labels for the proposed MRPLL, we design the reward function to evaluate the quality of an action through the labels produced from taking the action. The reward is designed to reflect both classification performance, label consistency, and information reliability.

where γ and λ are balance parameter. Lcls is the partial label classification loss, which can be calculated as,

where I(·) represents the indicator function and ℓ(·) is crossentropy loss. This item learns through a progressive guidance model with soft labels after disambiguation, first learning simple instances with high credibility, and then learning difficult instances. Furthermore, we introduce a label consistency reward LKL = DKL (˜yi∥ˆyi), which can encourage retaining confident predictions and revising uncertain ones. Meanwhile, the role of MI loss is to constrain the reward function to encourage actions that reduce classification loss, while also benefiting instances with high MI.

2) Advantage Estimation: The advantage function is defined as Adv i =Ri -Vϕ(si), which is normalized to stabilize training.

Therefore, updating the learning objectives of actors and critics through minimization,

The final classifier objective is,

where λMI is the weighting coefficient. The final prediction model fθ ◦ gψ is used to make correct predictions on unseen inputs. The learning process of the proposed MR-PLL method is summarized in Algorithm 1. Meanwhile, the theoretical analysis of objective function is elaborated in Appendix A.2.

Algorithm 1 MR-PLL Framework

Input: Dataset D = (xi, yi), classifier fθ ◦ gψ, the policy

function πψ, and the value function Vϕ(s)

Output: y∗ = argmax(fθ(gψ(x)))

1: Initialization:

• 2: while t <maxIter do
• 3: Compute state representation according to Eq.(1);
• 4: MI guided soft label fusion computed according to Eq.(2), Eq.(4) and Eq.(6);
• 5: Policy action: Sample action πψ (a | s);
• 6: Label transition according to Eq.(7);
• 7: Reward computation according to Eq.(8);
• 8: Update actor-critic according to Eq.(10);
• 9: Update classifier according to Eq.(11);
• 10: end while

4 Experiments

4.1 Experimental Details

Datasets To assess the effectiveness of the proposed MRPLL algorithm, we conduct experiments on four widely used benchmark datasets, including MNIST [LeCun et al., 1998], Fashion-MNIST [Xiao et al., 2017], Kuzushiji-MNIST [Clanuwat et al., 2018], and CIFAR-10 [Krizhevsky et al., 2009]. For these datasets, partial label annotations are synthesized following the standard procedure introduced in [Yao et al., 2021]. Specifically, the candidate label set for each instance is generated through Q- 1 independent Bernoulli trials, where the parameter q controls the probability of adding an incorrect label [Fan et al., 2024]. A larger value of q results in candidate sets with higher label ambiguity, thereby increasing the difficulty of the disambiguation task. In our experiments, we consider three noise levels, q ∈ {0.1, 0.3, 0.5}, to evaluate the robustness of MR-PLL under varying degrees of label noise across all four benchmark datasets. In addition to synthetic partial label settings, we further validate our method on five real-world partially labeled datasets, including Lost [Cour et al., 2011], BirdSong [Briggs et al., 2012], MSRCv2 [Liu and Dietterich, 2012], Soccer Player [Zeng et al., 2013], and Yahoo!News [Guillaumin et al., 2010]. Detailed dataset statistics are provided in Appendix B.1.

Comparison Methods. To better verify the superiority of the MR-PLL algorithm, we compare the proposed method with the state-of-the-art PLL methods including CC [Feng et al., 2020], PRODEN [Yao et al., 2021], LW [Lv et al., 2020], CAVL [Zhang et al., 2021], DIRK [Wu et al., 2024], VALEN [Xu et al., 2024] and LS-PLL [Gong et al., 2024]. Each comparison method sets parameters according to suggested settings in the original papers, and the details are shown in Appendix B.2.

Implementation Details. The implementation is based on PyTorch [Paszke et al., 2019] training, and all experiments are conducted with an NVIDIA GeForce RTX 4090 GPU. The backbone network uses an MLP and ConvNet to apply to real-world partially labeled datasets and benchmark datasets, which outputs both class logits and feature representations. The RL component is implemented using an Actor-Critic architecture, where both the actor and critic are parameterized by two-layer multilayer perceptrons with ReLU activations. We train models for 200 epochs, 128 batch size and use the SGD optimizer with a momentum of 0.9. Label refinement is performed using a k-nearest neighbor attention mechanism in the normalized feature space. The number of neighbors is fixed to k = 8, and the temperature parameter in the attention softmax is set to τ = 0.5. The balance parameters η, γ, λ and λMI are set to 0.1, 0.2, 0.6 and 0.2, respectively.

4.2 Experimental Results

Comparisons with State-of-the-arts methods. Across 35 experiments (5 real-world datasets × 7 comparison algorithms) in Table 1, MR-PLL significantly outperforms all other methods in 97.1% of the cases. Especially on datasets BirdSong, MSRCv2, and Lost, the MR-PLL achieved the best performance. Specifically, the mean classification accuracy of MR-PLL is higher than that of the best comparative method by 3.05%, 6.01%, 1.59%, 2.19%. As shown in Table 2, compared with the PLL algorithm designed based on the loss function, MR-PLL consistently performs better than PRODEN, LW, CC, and CAVL in the benchmark dataset. Compared with VALEN, DIRK, and LS-PLL, which have more complex structures, MR-PLL still shows better disambiguation performance and is superior to other methods on the MNIST, F-MNIST, and K-MNIST datasets. We attribute this satisfactory result to the design of the MR-PLL introduction of mutual information guided reinforcement learning disambiguation, which adaptively adjusts the confidence of soft label propagation based on the dependency between feature representations, and explores potential disambiguation strategies through reinforcement learning to gradually improve the

label disambiguation ability of the model.

Analysis of ablation study. To evaluate the effectiveness of each component in MR-PLL, we conduct ablation studies on CIFAR-10 (q = 0.3) and BirdSong. Specifically, ablation experiments were performed on four specific parts, i.e., MR-PLL-w/o-I removes the MI term from the reward and loss function, MR-PLL-w/o-LKL drops the label consistency reward, MR-PLL-w/o-Abs removes the abstain action and force the policy to choose between retain and modify, and MR-PLL-w/o-SLF drops MI guided soft label fusion. As reported in Table 3, all variants perform worse than the full MR-PLL model. Removing the MI term leads to a noticeable accuracy drop, demonstrating the importance of MI in exploiting latent semantic information for label disambiguation. Excluding LKL weakens label consistency during policy exploration, reducing robustness. Eliminating the abstain action degrades performance, indicating that postponing decisions on hard instances helps prevent error accumulation and improves stability. In addition, removing soft label fusion results in lower accuracy, verifying its effectiveness in providing reliable supervision signals for disambiguation. Overall, the superior performance of MR-PLL confirms that each component contributes positively to the disambiguation process.

Visualization of MI estimation in MR-PLL. To further illustrate the effectiveness of MR-PLL, we visualize the learned feature representations on CIFAR-10 using T-SNE [Maaten and Hinton, 2008]. As shown in Figure 2(a), MRPLL learns more discriminative features with clearer class boundaries and better clustering structures. By maximizing MI, samples from the same class become more compact, while features from different classes are better separated. This enhanced representation facilitates similarity estimation among instances, thereby improving ambiguous label detection and disambiguation. Moreover, MI estimation provides a reliable signal for assessing the credibility of candidate labels and dynamically refining the disambiguation process. The confusion matrix further confirms the effectiveness of MR-PLL in achieving accurate label disambiguation.

Uncertainty Quantification of MR-PLL. Existing machine learning algorithms not only require high accuracy, but also must reflect the effectiveness and reliability of the model [Guo et al., 2017]. As shown in Figure 3, we analyzed the output uncertainty of various PLL methods. The results demonstrate that MR-PLL not only significantly mitigates the impact of label ambiguity but also enhances model robustness and reliability. This leads to more credible model outputs and yields competitive expected calibration error (ECE) scores. The results indicate that MR-PLL can enable agents to learn the best disambiguation behavior through trial and error under the guidance of rewards through RL. Meanwhile, the uncertainty estimation of MI is used to guide the reward function and adjust the update of candidate labels. This not only effectively reduces the impact of label ambiguity, but also makes the model more robust and reliable, thereby improving the credibility of the model output and achieving a competitive ECE score.

Convergence analysis. To prove the convergence of MRPLL more intuitively, this section analyzes the convergence curves of MR-PLL on the CIFAR-10 and MNIST datasets with q = 0.3. From Figure 2(b), MR-PLL achieves a good disambiguation effect. With the increase of learning iterations, MR-PLL approaches the steady-state value, so the convergence of the model can be proved through experiments. Meanwhile, from the comprehensive experimental test results, it can be seen that MR-PLL not only has excellent disambiguation performance but also has good stability.

Parameter sensitivity analysis. The MR-PLL involves main parameters. i.e. k, η, γ, λ and λMI . Figure 4 illustrates the performance of MR-PLL under various parameter configurations. For clarity, results are shown only on the BirdSong and MSRCv2 datasets. In these tests, one parameter is changed at a time, while the other parameters are fixed at their default values. As shown in Figure 4, MR-PLL achieves optimal performance at k = 8, η = 0.1, γ = 0.2, λ = 0.6 and λMI = 0.2. These results further demonstrate the effectiveness of the MR-PLL algorithm among its various components. Meanwhile, the per-epoch computational complexity of the MR-PLL algorithm can reach O(N · (d+Q+dlogN)). The detailed calculation process is shown in Appendix B.4.

5 Conclusion

In this paper, we proposed a novel mutual information guided reinforcement learning framework for partial label disambiguation (MR-PLL). The highlight of MR-PLL is that it formulated partial label disambiguation as a sequential Markov decision process to enable dynamic correction of ambiguous labels through reinforcement learning. The integration of mutual information guided rewards and an abstention mechanism allows the model to perform reliable disambiguation while effectively mitigating bias accumulation caused by error disambiguation decisions. The method is superior to the state-of-the-art methods. Experiments on both benchmark and real-world datasets demonstrate its effectiveness. This work highlights the potential of combining reinforcement learning and information-theoretic guidance for addressing PLL and opens new directions for decision-theoretic weakly supervised learning.

Acknowledgments

This research was supported by the National Natural Science Foundation of China (Nos. 62401309, 62376130, 62401362), the China Postdoctoral Science Foundation (No. 2025M771692), the Natural Science Foundation of Shandong Province and Qingdao (No. 2R2024QF119, No. 24-4-4-zrjj89-jch), and open project (No. 2024PY030).