[Paper Review] A Multi-Task Self-Normalizing 3D-CNN to Infer Tuberculosis Radiological Manifestations

Pedro M. Gordaliza, J.J. Vaquero+3 more|arXiv (Cornell University)|Jul 29, 2019

COVID-19 diagnosis using AI9 references3 citations

TL;DR

This paper proposes a multi-task self-normalizing 3D-CNN using SELU activations and uncertainty-weighted loss to infer tuberculosis radiological manifestations from CT scans, achieving F₁-scores above 0.9 and a root mean square error of 1.16 for nodule counting, outperforming baseline models with reduced need for batch normalization and hyperparameter tuning.

ABSTRACT

We propose a learning method well-suited to infer the presence of Tuberculosis (TB) manifestations on Computer Tomography (CT) scans mimicking the radiologist reports. Latent features are extracted from the CT volumes employing the V-Net encoder and those are the input to a Feed-Forward Neural Network (FNN) for multi-class classification. To overtake the issues (e.g., exploding/vanishing gradients, lack of sensibility) that normally appear when training deep 3D models with datasets of limited size and composed of large volumes, our proposal employs: 1) At the network architecture level, the scaled exponential linear unit (SELU) activation which allows the network self-normalization, and 2) at the learning phase, multi-task learning with a loss function weighted by the task homoscedastic uncertainty. The results achieve F1-scores close to or above 0.9 for the detection of TB lesions and a Root Mean Square Error of 1.16 for the number of nodules.

Motivation & Objective

• To develop a deep learning model that mimics radiologist reports for TB on CT scans without relying on lesion segmentation masks.
• To address challenges in training deep 3D-CNNs on small, high-dimensional CT datasets, such as vanishing/exploding gradients and poor sensitivity.
• To reduce dependency on manual hyperparameter tuning of loss weights in multi-task learning.
• To improve model generalization and convergence by integrating self-normalizing properties and uncertainty-based loss weighting.
• To enable automated, robust extraction of TB radiological biomarkers for clinical and public health use.

Proposed method

• The model uses a V-Net encoder to extract 3D volumetric features from CT scans at 1×1×2 mm isotropic resolution.
• Latent features are fed into a shared feed-forward neural network (FNN₁), followed by two task-specific FNNs for regression (nodule count) and binary classification (TB manifestations).
• Scaled Exponential Linear Units (SELU) are used as activation functions to enable self-normalizing properties, eliminating the need for batch normalization.
• Alpha-dropout is applied instead of standard dropout to maintain the self-normalizing behavior under stochastic regularization.
• A homoscedastic uncertainty-weighted loss function dynamically adjusts the contribution of each task (regression and binary classification) during training, improving loss stability and convergence.
• Data augmentation via elastic transformations and Gaussian noise is applied during training to improve robustness on limited data.

Experimental results

Research questions

• RQ1 Can a self-normalizing 3D-CNN effectively infer TB radiological manifestations from CT scans without requiring lesion segmentation masks?
• RQ2 Does the use of SELU activations and uncertainty-weighted multi-task learning improve model convergence and performance on small, high-dimensional medical imaging datasets?
• RQ3 How does the proposed model compare to conventional batch normalization and PReLU-based architectures in terms of F₁-score and regression error for TB lesion detection?
• RQ4 To what extent does uncertainty-weighted loss reduce the need for manual hyperparameter tuning of task weights in multi-task learning?
• RQ5 Can the model generalize well across folds in 5-fold cross-validation while maintaining high performance on both regression and binary classification tasks?

Key findings

• The proposed SELU-based model achieved a mean F₁-score of 0.90 for binary TB manifestation detection across all lobes, with individual task F₁-scores ranging from 0.88 to 0.94.
• The model achieved a root mean square error (RMSE) of 0.68 for nodule count prediction across all folds, significantly outperforming the BN+PReLU baseline (RMSE = 0.78).
• The SELU model showed improved convergence in validation loss across all folds, as evidenced by the loss curve in Figure 2, with more stable and lower loss values.
• No statistically significant differences were found between models in paired t-tests (p > 0.05), indicating consistent performance gains across folds.
• The model achieved an F₁-score of 0.92 for 'tree in bud' and 0.91 for 'consolidation' across all folds, demonstrating strong performance on rare or subtle TB signs.
• The use of SELU and uncertainty-weighted loss reduced the need for batch normalization and manual loss weight tuning, simplifying training and reducing computational complexity.