[Paper Review] Label-Assemble: Leveraging Multiple Datasets with Partial Labels

Mintong Kang, Li, Bowen+5 more|arXiv (Cornell University)|Sep 25, 2021

COVID-19 diagnosis using AI4 citations

TL;DR

Label-Assemble proposes a novel framework to leverage partially labeled medical datasets by treating negative and related disease labels as auxiliary signals, using class queries, pseudo-labeling, and consistency regularization. It achieves state-of-the-art performance—improving COVID-19 diagnosis accuracy to 99.3% and PDAC detection sensitivity to 84.0%—by assembling pathologically and spatially related labels from heterogeneous public datasets without requiring full annotations.

ABSTRACT

The success of deep learning relies heavily on large labeled datasets, but we often only have access to several small datasets associated with partial labels. To address this problem, we propose a new initiative, "Label-Assemble", that aims to unleash the full potential of partial labels from an assembly of public datasets. We discovered that learning from negative examples facilitates both computer-aided disease diagnosis and detection. This discovery will be particularly crucial in novel disease diagnosis, where positive examples are hard to collect, yet negative examples are relatively easier to assemble. For example, assembling existing labels from NIH ChestX-ray14 (available since 2017) significantly improves the accuracy of COVID-19 diagnosis from 96.3% to 99.3%. In addition to diagnosis, assembling labels can also improve disease detection, e.g., the detection of pancreatic ductal adenocarcinoma (PDAC) can greatly benefit from leveraging the labels of Cysts and PanNets (two other types of pancreatic abnormalities), increasing sensitivity from 52.1% to 84.0% while maintaining a high specificity of 98.0%.

Motivation & Objective

• Address the challenge of limited labeled data in medical imaging by exploiting partially labeled public datasets.
• Overcome the limitations of existing methods that only use labeled portions of data, ignoring the potential of unlabeled or partially labeled examples.
• Investigate whether learning from negative or related disease examples can improve the performance of models for rare or novel diseases.
• Develop a unified framework capable of dynamically integrating varying, heterogeneous, and potentially conflicting labels across multiple datasets.
• Demonstrate that assembling labels from related diseases (pathologically or spatially) leads to significant performance gains in both diagnosis and detection tasks.

Proposed method

• Introduce a learnable class query mechanism that encodes visual tasks and enables dynamic integration of partial labels across datasets with varying categories.
• Use a shared convolutional neural network (CNN) to extract image features, followed by a prediction head that computes predictions via inner product between features and class queries.
• Apply binary cross-entropy loss (ℒ_bce) when labels are available, and unsupervised losses (ℒ_pseudo and ℒ_consist) for missing labels using pseudo-labeling and consistency regularization.
• Leverage pseudo-labeling to generate predictions for unlabeled samples and enforce consistency across augmentations to mitigate domain shift across datasets.
• Train the model end-to-end using a combination of supervised and self-supervised objectives to exploit both labeled and unlabeled data from heterogeneous sources.
• Adapt the framework to multi-task learning by enabling the model to handle arbitrary, variable numbers of classes through the class query mechanism, unlike fixed-category classification.

Experimental results

Research questions

• RQ1 Can learning from negative examples in partially labeled datasets improve the performance of a model for a target disease, such as COVID-19?
• RQ2 Does assembling labels from pathologically similar diseases (e.g., pneumonia for COVID-19) lead to better diagnostic performance than relying solely on positive examples?
• RQ3 Can leveraging labels from spatially related abnormalities (e.g., cysts and PanNETs for PDAC) enhance disease detection sensitivity?
• RQ4 Is it possible to achieve performance comparable to full-label learning using only partial labels from multiple heterogeneous datasets?
• RQ5 How does the proposed Label-Assemble framework compare to existing methods that only utilize labeled data or require full annotations?

Key findings

• Label-Assemble improves COVID-19 diagnosis accuracy from 96.3% to 99.3% by assembling labels from NIH ChestX-ray14, demonstrating that learning from negative and related disease examples enhances decision boundary learning.
• The method increases sensitivity for pancreatic ductal adenocarcinoma (PDAC) detection from 52.1% to 84.0% by incorporating labels from related pancreatic abnormalities (cysts and PanNETs), while maintaining high specificity (98.0%).
• Performance gains are positively correlated with inter-class similarity: for lung nodule classification, the Pearson correlation between improvement and similarity to assembled disease is r=0.83 (p=4.93e-4).
• Label-Assemble achieves state-of-the-art mAUC of 0.832 on NIH ChestX-ray14, outperforming prior methods across 13 out of 14 thorax diseases, and improves CheXpert performance by 1.8% over the baseline.
• The framework achieves performance on par with full-label learning while reducing annotation cost by 40% (75,310 partial labels vs. 105,434 full labels), showing that complete label completion is not necessary.
• The method outperforms existing partial-label baselines such as DoDNet and Med3D, especially in multi-task and semi-supervised settings, due to its adaptive class query and consistency regularization components.