Learning to Clean: Reinforcement Learning for Noisy Label Correction

AI summary

Overview

• Research area: Machine learning — specifically learning with noisy labels, recast as a reinforcement learning (RL) problem.
• Technical level: Intermediate to Advanced. Readers need some familiarity with Markov decision processes, actor-critic policy gradient methods, and standard noisy-label benchmarks.
• Scope: The paper proposes RLNLC, a framework that treats noisy label correction as a sequential decision-making problem and learns a policy network to decide which training labels to fix, evaluated on four benchmark datasets.

What This Paper Is About

Supervised deep learning assumes labels are correct, but real-world labels are often wrong, which degrades model generalization. Most existing fixes (sample selection, robust losses, label filtering) make static, one-shot decisions and cannot explore alternatives or learn from the long-term consequences of a correction. This paper reframes label correction as a reinforcement learning problem, where an agent repeatedly decides which labels to change, guided by a reward that reflects how well the resulting labels fit the data.

Key Contributions

1. RL formulation of label correction: The paper defines noisy label correction as a Markov Decision Process with a state space (data plus current labels), an action space (per-instance binary correction decisions), and a reward mechanism.
2. A representation-based policy function: The policy is built on top of a deep feature extraction network and decides correction probabilities using k-nearest-neighbor label agreement in the embedding space, computed with an attention-weighted aggregation of neighbor labels.
3. A two-part reward function: The reward combines a Label Consistency Reward (negative KL divergence between current labels and k-NN predicted labels) with a Noisy Label Alignment Reward (negative KL divergence between corrected labels in the noisy subset and their k-nearest neighbors in the clean subset).
4. An efficient critic encoding scheme plus actor-critic optimization: Because states and actions scale with dataset size, the paper replaces the state-action pair with the deterministically derived next state and compresses it into a fixed-length vector of bin proportions, then optimizes the policy with a SARSA-style temporal-difference critic.
5. Benchmark evaluation: Experiments across IDN, symmetric, and real-world noise conditions show RLNLC outperforming prior state-of-the-art methods on the reported tables.

Main Findings

• CIFAR10-IDN and CIFAR100-IDN results: At noise rates 0.20, 0.30, 0.40, 0.45 and 0.50, RLNLC reaches 97.3, 97.1, 96.9, 96.6 and 95.8 percent test accuracy on CIFAR10-IDN, and 80.5, 80.1, 78.5, 77.2 and 74.7 percent on CIFAR100-IDN (standard deviations 0.1 to 0.9 in parentheses). For comparison, the reproduced SSR baseline scores 96.5, 96.5, 96.3, 95.9 and 94.1 on CIFAR10-IDN, and 78.8, 78.6, 77.0, 75.0 and 72.8 on CIFAR100-IDN; DivideMix scores 94.8, 94.6, 94.5, 94.1 and 93.0, and 77.1, 76.3, 70.8, 57.8 and 58.6 respectively.
• Animal-10N results: RLNLC achieves 90.2 percent (standard deviation 0.1), compared with SSR at 87.7, SSR+ at 88.5, SURE at 89.0, Nested-CE at 84.1 (0.1), PLC at 83.4 (0.4) and CE at 79.4 (0.1). The dataset has roughly 8 percent noisy labels.
• Food-101N results: RLNLC reaches 89.2 percent (0.1), ahead of LongReMix at 87.3, PLC at 85.2 (0.0), DeepSelf at 85.1, BARE at 84.1, CleanNet at 83.9 and CE at 81.6.
• Symmetric noise results: With symmetric noise at rates 0.50, 0.80 and 0.90, RLNLC scores 97.4 (0.1), 95.8 (0.2) and 82.1 (0.7) on CIFAR10, and 81.2 (0.3), 70.6 (0.4) and 44.2 (0.8) on CIFAR100. DivideMix scores 94.4, 92.9 and 75.4 on CIFAR10, and 74.2, 59.6 and 31.0 on CIFAR100.
• Consistent margin across noise regimes: The reported gains hold on instance-dependent noise, symmetric noise and naturally noisy web data, which is the paper's central empirical claim.
• Efficiency motivation for the critic: The paper states that binning the next state into a proportion vector is needed because states and actions scale with dataset size; it does not report a wall-clock or memory comparison in the provided content.

Methodology in Plain English

The method treats the whole training dataset as the "state" of an environment. Each training example carries a label vector that can be changed over time. A deep network first learns an embedding space by pretraining with standard cross-entropy on the noisy data.

For each example, the method looks at its k nearest neighbors in that embedding space and builds a new predicted label by taking an attention-weighted average of the neighbors' current labels, where attention weights come from cosine similarity with a temperature parameter.

The policy network then compares this neighbor-based prediction with the example's current label. If classes other than the current one have higher probability under the neighbor prediction, the probability of correcting that example goes up. The formula sums the probabilities of classes more likely than the current predicted class and normalizes by the sum of probabilities of classes at least as likely as it; the probability is zero when there is no disagreement. Each correction decision is then sampled from a Bernoulli distribution, so the process explores different correction patterns rather than always applying the most obvious fix.

After the action is sampled, the transition is deterministic: examples with action 0 keep their label, examples with action 1 take the neighbor-aggregated label.

Two reward signals score the outcome. The Label Consistency Reward measures, using a fixed pretrained backbone, how well the new labels agree with their k nearest neighbors across the whole dataset. The Noisy Label Alignment Reward splits data into a "clean" subset (labels unchanged) and a "noisy" subset (labels corrected), and measures how well the corrected labels align with their k nearest neighbors drawn from the clean subset. Both are negative KL divergences; the composite reward exponentiates the weighted sum so it falls in (0, 1].

The policy is the actor; a Q-network is the critic, updated with SARSA temporal-difference errors. To keep the critic's input small, the method uses the fact that the next state is uniquely determined by the state-action pair, computes a per-instance reward, sorts instances into N_b bins by that reward, and encodes the state as the vector of bin proportions.

For final deployment, the trained policy is run for T' steps starting from the original dataset, and the labels in the final state are treated as cleaned labels. The prediction model is pretrained on the noisy data, the feature extractor is trained as part of policy learning, and the final model is fine-tuned on the cleaned labels with cross-entropy.

Training setup: ResNet-34 for CIFAR10-IDN, ResNet-50 for CIFAR100-IDN and Food-101N, VGG-19 with batch normalization for Animal-10N, and ResNet-18 for the symmetric-noise CIFAR experiments. Optimization uses SGD with momentum 0.9, batch size 128, L2 coefficient 5×10⁻⁴, initial learning rate 0.01 reduced to 0.001 at the halfway point, and a 50-epoch warmup. The policy network is trained for 500 epochs on CIFAR10-IDN, CIFAR100-IDN and Animal-10N with T = 10. The provided content is truncated at this point, so the remaining hyperparameters, ablations and additional analyses are not reported here.

Why This Matters

• Research impact: The work opens a new angle on noisy-label learning by replacing static filtering or reweighting with a sequential, reward-driven correction process. It also connects the noisy-label literature to the actor-critic tooling that has driven recent RL successes in other domains.
• Real-world applications:
  • Web-scraped image datasets, where labels come from search queries and captions, as in the Food-101N setting.
  • Medical imaging archives, where annotations from multiple experts disagree on ambiguous cases.
  • Industrial quality control and defect detection, where operators occasionally mislabel items under time pressure.
  • Crowdsourced and user-generated annotation pipelines, where subjective errors are common on hard or ambiguous examples.
• Industry relevance: Any organization that trains classifiers on cheaply collected labels — rather than expensively curated ones — can in principle use a label-cleaning policy to improve model accuracy without re-annotating the dataset. The critic binning scheme is explicitly designed to keep the approach scalable to large training sets.

Future Directions

• Noise-type generalization: The paper separates class-conditional noise (CCN) from instance-dependent noise (IDN) and focuses on IDN; whether the RL formulation transfers as cleanly to CCN is an open question.
• Scale and compute cost: The method requires nearest-neighbor search, a critic network and repeated policy rollouts over the full dataset. Its behavior on datasets far larger than the benchmarks used here is not established in the reported content.
• Hyperparameter and ablation analysis: The provided content does not report sensitivity to the discount factor, the reward trade-off weight, the number of bins, the trajectory lengths T and T', or the number of neighbors, nor does it include an ablation isolating each reward component in the material available here.
• Integration with modern training paradigms: Combining the learned correction policy with large-scale pretraining, semi-supervised pipelines, or human-feedback-style reward models — the last being an RL direction the related work highlights — is a natural extension.

Target Audience

Researchers and graduate students working on learning with noisy labels, weakly supervised learning, or data-centric AI; practitioners who train classifiers on crowdsourced, web-scraped, or otherwise imperfectly annotated data; and RL researchers interested in applying sequential decision-making to data curation rather than control or generation. Readers without RL background will need to consult the cited policy-gradient and actor-critic references to follow the optimization section.

Authors’ abstract

The challenge of learning with noisy labels is significant in machine learning, as it can severely degrade the performance of prediction models if not addressed properly. This paper introduces a novel framework that conceptualizes noisy label correction as a reinforcement learning (RL) problem. The proposed approach, Reinforcement Learning for Noisy Label Correction (RLNLC), defines a comprehensive state space representing data and their associated labels, an action space that indicates possible label corrections, and a reward mechanism that evaluates the efficacy of label corrections. RLNLC learns a deep feature representation based policy network to perform label correction through reinforcement learning, utilizing an actor-critic method. The learned policy is subsequently deployed to iteratively correct noisy training labels and facilitate the training of the prediction model. The effectiveness of RLNLC is demonstrated through extensive experiments on multiple benchmark datasets, where it consistently outperforms existing state-of-the-art techniques for learning with noisy labels.