💻 computer science

Leveraging Pathology Co-occurrence for Test-Time Adaptation in Chest X-Ray Diagnosis

This paper proposes Co-occurrence Weighted Adaptation (CoWA), a test-time adaptation method for chest X-ray diagnosis that leverages structured disease co-occurrence patterns to downweight noisy predictions and improve model robustness under domain shifts.

Original authors: Woojin Jeong, Yujin Choi, Dongbin Kim, Soyeon Park, Jaewook Lee

Published 2026-07-10

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Woojin Jeong, Yujin Choi, Dongbin Kim, Soyeon Park, Jaewook Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a detective trying to solve a mystery in a city you've never visited before. You have a brilliant guidebook (a pre-trained AI model) that worked perfectly in your hometown, but now you're in a new place with different street signs, different weather, and a different crowd. This is exactly what happens when a medical AI trained on one hospital's X-rays tries to diagnose patients at a different hospital. The equipment is different, the patients look different, and the old guidebook starts making mistakes.

Usually, when an AI gets confused, it tries to "learn on the fly" by looking at the new, unlabeled X-rays and adjusting itself. This is called Test-Time Adaptation (TTA). Think of it like a student taking a pop quiz without a teacher; they have to guess the answers and tweak their brain based on what they think is right.

The Problem with the Old Way
The paper argues that the standard way these students learn is flawed. Traditional methods treat every single X-ray as an independent puzzle. They assume that if a machine thinks a patient has "Pneumonia," it doesn't matter if the machine also thinks they have "Broken Leg" or "Sunburn." It just tries to make every single guess as confident as possible.

But here's the catch: Diseases don't work in isolation. In the real world, certain illnesses love to hang out together. If a patient has a specific type of fluid in their lungs, they are very likely to also have an enlarged heart. These are "co-occurrence patterns." The old methods ignore this social life of diseases. They might get a confident answer that says, "This patient has a broken leg and a sunburn," which is mathematically confident but medically nonsense. By trying to learn from these weird, impossible combinations, the AI actually gets worse at adapting to the new hospital.

The New Solution: CoWA (The "Social Butterfly" Detector)
The authors, a team from Seoul National University and others, propose a new method called CoWA (Co-occurrence Weighted Adaptation).

Imagine the AI is now a detective who carries a "Friendship Chart." As the AI looks at the new X-rays, it doesn't just look at one disease at a time. Instead, it asks: "Do these diseases usually hang out together in this new city?"

Building the Chart: The AI starts looking at the new X-rays and builds a map of which diseases appear together. It's like noticing that in this new neighborhood, people who wear red hats almost always carry umbrellas.
The Reality Check: When the AI makes a prediction, it checks its own "Friendship Chart."
- If the AI says, "Red Hat and Umbrella!" (a common pair), it says, "Good job! This prediction makes sense. I'll trust you."
- If the AI says, "Red Hat and Ice Cream Cone!" (a weird pair that never happens in this city), it says, "Wait a minute. This doesn't fit the pattern. You're probably confused. I'm going to ignore this guess for now."
The Result: The AI only learns from the guesses that fit the local social rules. It filters out the "noise" and the weird, impossible combinations.

What the Experiments Showed
The team tested this idea on four major chest X-ray datasets, creating six different scenarios where the AI had to move from one "city" (dataset) to another. They compared CoWA against the old methods (like TENT and AdaBN) and a model that didn't adapt at all.

The results were promising. In most of the test scenarios, CoWA performed better than the other methods. For example, when moving from the CheXpert dataset to the MIMIC dataset, CoWA achieved a mean score of 68.4, beating the next best method (EATA) which scored 67.9. When moving to the VinDr dataset, CoWA hit 85.6, while the runner-up (TENT) only managed 84.9.

Crucially, the paper suggests that CoWA is particularly good at preventing the AI from crashing on rare diseases. While other methods sometimes made huge mistakes on specific conditions (like Pneumothorax), CoWA kept its cool, ensuring that even the tricky cases didn't cause a total failure.

What This Doesn't Mean
The paper is careful not to overhype. It doesn't claim to have solved the problem of medical AI forever. It explicitly notes that the "Friendship Chart" the AI builds is based on the data it sees right now. If the new hospital has a weird bias or a very small number of patients, the chart might be a little off. The authors suggest that future work could look into making these patterns even more robust, perhaps by understanding the causes of why diseases hang out together, rather than just the patterns themselves.

The Bottom Line
This paper suggests that by teaching AI to respect the "social rules" of diseases—knowing that some illnesses are best friends and never show up alone—we can make medical models much more reliable when they travel to new hospitals. It's a simple but powerful tweak: stop trusting every confident guess, and only trust the ones that make sense in the context of the whole picture.

Technical Summary: Leveraging Pathology Co-occurrence for Test-Time Adaptation in Chest X-Ray Diagnosis

Problem Statement

Deep learning models for chest X-ray classification often suffer from significant performance degradation when deployed at new clinical sites due to domain shifts arising from differences in imaging equipment, acquisition protocols, and patient demographics. While Test-Time Adaptation (TTA) offers a solution by updating pretrained models using only unlabeled target data (without access to source data or target annotations), existing TTA methods are ill-suited for multi-label medical imaging.

Standard TTA approaches, originally designed for single-label natural image classification, minimize prediction entropy uniformly across all samples. This implicitly assumes label independence. However, in clinical settings, pathologies exhibit structured co-occurrence patterns (e.g., certain diseases frequently appear together) rather than occurring in isolation. Existing methods fail to account for these dependencies; they treat all confident predictions as equally reliable, even if a prediction violates plausible label relationships. This can lead to the amplification of inconsistent patterns and the introduction of noisy gradients during adaptation, particularly under domain shifts.

Methodology: Co-occurrence Weighted Adaptation (CoWA)

The authors propose CoWA, a source-free TTA method that leverages disease co-occurrence patterns as a reliability signal to guide adaptation. The framework operates without requiring additional supervision or architectural changes, adapting only the Batch Normalization (BN) affine parameters of the pretrained model.

CoWA consists of three core components:

Target Domain Co-occurrence Matrix Estimation:
The method estimates the label co-occurrence structure directly from incoming test predictions. Model outputs (soft probabilities) are binarized via a fixed threshold to suppress low-confidence activations. A pairwise co-occurrence matrix $S$ is accumulated across processed batches, where $S_{jk}$ counts the joint prediction of pathologies $j$ and $k$ . This is normalized to form an empirical co-occurrence matrix $M$ , which evolves as adaptation proceeds, capturing the target domain's specific dependency structure.
Consistency-Aware Sample Weighting:
For each sample, CoWA computes a consistency score measuring how well its predicted label configuration aligns with the estimated target co-occurrence matrix $M$ . A local co-occurrence pattern $m_i$ is constructed from the sample's predictions. The deviation between $m_i$ and $M$ is calculated using the Frobenius norm. This deviation is converted into a reliability weight $w_i$ :
$w_i = \exp\left(-\frac{\|m_i - M\|_F^2}{\tau}\right)$
where $\tau$ is a temperature parameter. Samples with predictions consistent with the emerging target structure receive weights close to 1, while structurally inconsistent (noisy) predictions are downweighted.
Weighted Entropy Minimization:
Instead of minimizing the uniform average entropy, CoWA minimizes a weighted entropy loss:
$L_{CoWA} = \frac{1}{B} \sum_{i=1}^{B} \max(w_i, w_{min}) \cdot H(\hat{y}_i)$
Here, $H(\cdot)$ is the prediction entropy, $B$ is the batch size, and $w_{min}$ prevents weights from collapsing to zero for numerical stability. This ensures that the model adapts primarily based on structurally coherent samples, reducing the impact of noisy gradients.

Key Contributions

Proposal of CoWA: A novel TTA method that estimates pathology co-occurrence patterns from model predictions and utilizes them as a per-sample reliability signal for multi-label medical image classification.
Identification of Limitations in Standard TTA: The authors demonstrate that standard TTA methods, designed for single-label benchmarks, can be ineffective or harmful in multi-label clinical tasks because they ignore label dependencies.
Empirical Validation: Extensive experiments on public chest X-ray datasets (MIMIC-CXR, CheXpert, VinDr-CXR, NIH Chest X-ray) across six domain shift scenarios show consistent performance gains over established baselines.

Experimental Results

The authors evaluated CoWA against several baselines, including No Adaptation, AdaBN, TENT, CoTTA, EATA, and RoTTA, using AUROC as the primary metric.

Overall Performance: CoWA achieved the best mean AUROC in most domain shift scenarios and remained competitive in others, consistently ranking first or second with small performance gaps.
Robustness to Domain Shift: Unlike baselines such as TENT and AdaBN, which showed high variability (performing well in some shifts but degrading in others), CoWA maintained stable performance.
Class-Specific Improvements: Competing methods often suffered marked drops on low-prevalence pathologies (e.g., Pneumothorax). CoWA consistently alleviated these class-specific degradations. Notably, CoWA was the only method that avoided severe performance drops below the unadapted baseline across all shifts, highlighting its worst-case robustness.
Analysis:
- Matrix Evolution: The estimated co-occurrence matrix in CoWA progressively aligned with the true target structure during adaptation, whereas baselines lost clear dependency structures.
- Weight Reliability: A positive correlation was observed between the assigned weights and prediction quality (AUPRC), confirming that the weights effectively serve as an implicit reliability estimator.
- Ablation: The use of the Frobenius norm ( $\ell_F$ ) for consistency scoring yielded higher mean AUROC and lower variance compared to $\ell_1$ , $\ell_2$ , or $\ell_\infty$ norms. Additionally, CoWA demonstrated more stable BN affine gradient updates compared to TENT.

Significance and Claims

The paper claims that CoWA offers a simple, deployable solution for improving diagnostic robustness under domain shift by incorporating target-domain co-occurrence structure into TTA without requiring additional annotations.

The authors emphasize that in clinical settings, where a severe drop in performance on any single pathology can compromise diagnosis, the "worst-case robustness" provided by CoWA is as critical as mean performance. By downweighting predictions that violate plausible label relationships, CoWA enables adaptation to focus on structurally consistent samples, thereby reducing noisy gradients.

The authors conclude modestly, noting that while CoWA is effective, the estimated patterns may be sensitive to dataset-level biases. They suggest that extending CoWA to capture more robust, causally grounded relationships is a promising direction for future work.