Drainage: A Unifying Framework for Addressing Class Uncertainty

TL;DR

This paper introduces a 'drainage node' framework for deep learning that reallocates probability mass to handle noisy labels, class ambiguity, and out-of-distribution samples, improving robustness and accuracy.

Contribution

The paper proposes a novel drainage node mechanism integrated into neural networks to better manage uncertainty and noise, outperforming existing methods in noisy label scenarios.

Findings

Achieves up to 9% accuracy improvement in high-noise regimes.

Matches or surpasses state-of-the-art on real-world noisy datasets.

Demonstrates denoising effect and stability in decision boundaries.

Abstract

Modern deep learning faces significant challenges with noisy labels, class ambiguity, as well as the need to robustly reject out-of-distribution or corrupted samples. In this work, we propose a unified framework based on the concept of a "drainage node'' which we add at the output of the network. The node serves to reallocate probability mass toward uncertainty, while preserving desirable properties such as end-to-end training and differentiability. This mechanism provides a natural escape route for highly ambiguous, anomalous, or noisy samples, particularly relevant for instance-dependent and asymmetric label noise. In systematic experiments involving the addition of varying proportions of instance-dependent noise or asymmetric noise to CIFAR-10/100 labels, our drainage formulation achieves an accuracy increase of up to 9\% over existing approaches in the high-noise regime. Our results on real-world datasets, such as mini-WebVision, mini-ImageNet and Clothing-1M, match or surpass existing state-of-the-art methods. Qualitative analysis reveals a denoising effect, where the drainage neuron consistently absorbs corrupt, mislabeled, or outlier data, leading to more stable decision boundaries. Furthermore, our drainage formulation enables applications well beyond classification, with immediate benefits for web-scale, semi-supervised dataset cleaning, and open-set applications.