Learning Low-Rank Feature for Thorax Disease Classification

TL;DR

This paper introduces a Low-Rank Feature Learning (LRFL) method that enhances thorax disease classification by reducing noise and background effects, applicable across neural network architectures, and supported by empirical and theoretical insights.

Contribution

The paper proposes a novel LRFL method for neural networks that leverages low-rank features, improving medical image classification performance.

Findings

LRFL improves classification accuracy and mAUC on chest X-ray datasets.

Empirical results show LRFL's effectiveness across CNN and ViT architectures.

Theoretically, LRFL is supported by a generalization bound for low-rank features.

Abstract

Deep neural networks, including Convolutional Neural Networks (CNNs) and Visual Transformers (ViT), have achieved stunning success in medical image domain. We study thorax disease classification in this paper. Effective extraction of features for the disease areas is crucial for disease classification on radiographic images. While various neural architectures and training techniques, such as self-supervised learning with contrastive/restorative learning, have been employed for disease classification on radiographic images, there are no principled methods which can effectively reduce the adverse effect of noise and background, or non-disease areas, on the radiographic images for disease classification. To address this challenge, we propose a novel Low-Rank Feature Learning (LRFL) method in this paper, which is universally applicable to the training of all neural networks. The LRFL method is both empirically motivated by the low frequency property observed on all the medical datasets in this paper, and theoretically motivated by our sharp generalization bound for neural networks with low-rank features. In the empirical study, using a neural network such as a ViT or a CNN pre-trained on unlabeled chest X-rays by Masked Autoencoders (MAE), our novel LRFL method is applied on the pre-trained neural network and demonstrate better classification results in terms of both multiclass area under the receiver operating curve (mAUC) and classification accuracy.