Learning Low-Rank Feature for Thorax Disease Classification

Yancheng Wang, Rajeev Goel, Utkarsh Nath, Alvin C Silva, Teresa Wu, Yingzhen Yang

Keywords: Low-Rank Feature Learning, Low Frequency Property, Thorax Disease Classification

TL;DR: We propose a novel we propose a novel Low-Rank Feature Learning (LRFL) method for thorax disease classification, which learns low-rank features of a neural network so as to effectively reduce the adverse effect of background or non-disease areas.

Abstract: Deep neural networks, including Convolutional Neural Networks (CNNs) and Visual Transformers (ViT), have achieved stunning success in the 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 that 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 a Low Frequency Property (LFP) and theoretically motivated by our sharp generalization bound for neural networks with low-rank features. LFP not only widely exists in deep neural networks for generic machine learning but also exists in all the thorax medical datasets studied in this paper. 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 demonstrates better classification results in terms of both multi-class area under the receiver operating curve (mAUC) and classification accuracy than the current state-of-the-art. The code is available at https://github.com/Statistical-Deep-Learning/LRFL.

Primary Area: Machine learning for healthcare

Flagged For Ethics Review: true

Submission Number: 19189