Lung Infection Severity Prediction Using Transformers with Conditional TransMix Augmentation and Cross-Attention

2025-10-08

AI Summary

To address class imbalance and insufficient feature representation in medical imaging data for severity prediction of pulmonary infections—particularly pneumonia—this paper proposes the QCross-Att-PVT network and Conditional Online TransMix augmentation. QCross-Att-PVT introduces a parallel encoder architecture with cross-gated attention to enable effective multi-scale feature aggregation, while Conditional Online TransMix dynamically adjusts sample mixing strategies based on prediction confidence to mitigate label skew. Evaluated on two benchmark datasets—RALO CXR and Per-COVID-19 CT—the method achieves state-of-the-art performance in accuracy, AUC, and robustness compared to leading models. It delivers clinically interpretable, generalizable AI support for risk stratification and personalized intervention in pulmonary infection management.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Multi-modal Vision
• Knowledge Representation and Reasoning: Diagnosis and Abductive Reasoning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Responsible Web: Algorithmic accountability and transparency on the web

Abstract

Lung infections, particularly pneumonia, pose serious health risks that can escalate rapidly, especially during pandemics. Accurate AI-based severity prediction from medical imaging is essential to support timely clinical decisions and optimize patient outcomes. In this work, we present a novel method applicable to both CT scans and chest X-rays for assessing lung infection severity. Our contributions are twofold: (i) QCross-Att-PVT, a Transformer-based architecture that integrates parallel encoders, a cross-gated attention mechanism, and a feature aggregator to capture rich multi-scale features; and (ii) Conditional Online TransMix, a custom data augmentation strategy designed to address dataset imbalance by generating mixed-label image patches during training. Evaluated on two benchmark datasets, RALO CXR and Per-COVID-19 CT, our method consistently outperforms several state-of-the-art deep learning models. The results emphasize the critical role of data augmentation and gated attention in improving both robustness and predictive accuracy. This approach offers a reliable, adaptable tool to support clinical diagnosis, disease monitoring, and personalized treatment planning. The source code of this work is available at https://github.com/bouthainas/QCross-Att-PVT.

Problem

• Predicting lung infection severity from medical imaging using AI
• Addressing dataset imbalance with conditional augmentation techniques
• Improving diagnostic accuracy for pneumonia and COVID-19 cases

Innovation

• Transformer-based architecture with cross-gated attention
• Conditional Online TransMix for data augmentation
• Multi-scale feature integration for severity prediction