Diffusion Model Driven Test-Time Image Adaptation for Robust Skin Lesion Classification

2405.11289

YC

0

Reddit

0

Published 5/21/2024 by Ming Hu, Siyuan Yan, Peng Xia, Feilong Tang, Wenxue Li, Peibo Duan, Lin Zhang, Zongyuan Ge
Diffusion Model Driven Test-Time Image Adaptation for Robust Skin Lesion Classification

Abstract

Deep learning-based diagnostic systems have demonstrated potential in skin disease diagnosis. However, their performance can easily degrade on test domains due to distribution shifts caused by input-level corruptions, such as imaging equipment variability, brightness changes, and image blur. This will reduce the reliability of model deployment in real-world scenarios. Most existing solutions focus on adapting the source model through retraining on different target domains. Although effective, this retraining process is sensitive to the amount of data and the hyperparameter configuration for optimization. In this paper, we propose a test-time image adaptation method to enhance the accuracy of the model on test data by simultaneously updating and predicting test images. We modify the target test images by projecting them back to the source domain using a diffusion model. Specifically, we design a structure guidance module that adds refinement operations through low-pass filtering during reverse sampling, regularizing the diffusion to preserve structural information. Additionally, we introduce a self-ensembling scheme automatically adjusts the reliance on adapted and unadapted inputs, enhancing adaptation robustness by rejecting inappropriate generative modeling results. To facilitate this study, we constructed the ISIC2019-C and Dermnet-C corruption robustness evaluation benchmarks. Extensive experiments on the proposed benchmarks demonstrate that our method makes the classifier more robust across various corruptions, architectures, and data regimes. Our datasets and code will be available at url{https://github.com/minghu0830/Skin-TTA_Diffusion}.

Create account to get full access

or

If you already have an account, we'll log you in

Overview

  • This paper introduces a novel approach for improving the robustness of skin lesion classification models by adapting the input images at test time using a diffusion model.
  • The method aims to make skin lesion classification models more robust to distribution shift and domain mismatch between the training and test data.
  • The authors demonstrate the effectiveness of their approach on several public skin lesion datasets, achieving state-of-the-art performance compared to existing test-time adaptation techniques.

Plain English Explanation

Skin cancer is a serious health concern, and accurately identifying different types of skin lesions is crucial for early diagnosis and treatment. However, machine learning models for skin lesion classification can struggle when faced with images that are different from the ones they were trained on. This is known as the "distribution shift" problem.

The researchers in this paper developed a new technique to address this challenge. They used a type of machine learning model called a "diffusion model" to automatically adapt the input images at test time, before they are fed into the skin lesion classification model. The diffusion model learns to transform the input images in a way that makes them more similar to the training data, thereby improving the classification model's performance.

By applying this test-time adaptation approach, the researchers were able to significantly boost the accuracy of skin lesion classification, even on datasets that were very different from the training data. This is an important step towards building more robust and reliable AI systems for medical image analysis, which could have a real impact on patient outcomes.

The key innovation in this work is the use of a diffusion model for test-time adaptation, which is a powerful but often overlooked technique in machine learning. The researchers' success in applying it to the skin lesion classification problem demonstrates its potential to be a valuable tool for improving the reliability of AI systems in a wide range of domains.

Technical Explanation

The researchers propose a novel approach called "Diffusion Model Driven Test-Time Image Adaptation" (DMTA) to improve the robustness of skin lesion classification models. DMTA works by using a pre-trained diffusion model to transform the input images at test time, making them more similar to the training data distribution.

The diffusion model is trained in an unsupervised manner on a large, diverse dataset of skin lesion images. During test time, the diffusion model is used to generate a sequence of intermediate images that gradually transform the input image to be more like the training distribution. The final adapted image is then fed into the skin lesion classification model, which is expected to perform better due to the reduced distribution shift.

The researchers evaluate DMTA on several public skin lesion datasets, including ISIC 2018, Derm7pt, and MClass-1000. They compare DMTA to existing test-time adaptation methods, such as MetaDerm and DG-TTA, and demonstrate state-of-the-art performance on all datasets.

The key insights from this work are:

  1. Diffusion models are effective for test-time adaptation: The researchers show that diffusion models can learn rich representations of skin lesion images, which can then be leveraged to adapt the input images at test time.
  2. Test-time adaptation can significantly improve skin lesion classification: By applying DMTA, the researchers are able to achieve substantial performance gains over baseline models, highlighting the importance of addressing distribution shift in this domain.
  3. DMTA is a versatile and generalizable approach: The researchers demonstrate the effectiveness of DMTA on multiple skin lesion datasets, suggesting that it could be a valuable tool for improving the robustness of medical image classification models in general.

Critical Analysis

The researchers have presented a compelling approach for addressing the problem of distribution shift in skin lesion classification. By using a diffusion model for test-time adaptation, they have shown that it is possible to significantly improve the performance of skin lesion classification models, even when faced with data that is quite different from the training distribution.

One potential limitation of the work is the reliance on a large, diverse dataset of skin lesion images for training the diffusion model. In practice, such datasets may not always be readily available, and the performance of DMTA could be sensitive to the quality and coverage of the training data. The researchers could have explored ways to mitigate this, such as using few-shot or self-supervised learning techniques to train the diffusion model.

Additionally, the paper does not provide a detailed analysis of the types of distribution shifts that DMTA is able to handle. It would be helpful to understand the specific image transformations that the diffusion model is able to learn, and how they relate to the common sources of distribution shift in skin lesion datasets (e.g., differences in patient demographics, imaging modalities, or skin types).

Despite these minor limitations, the researchers have made a valuable contribution to the field of medical image analysis. The DMTA approach represents an important step towards building more robust and reliable AI systems for skin cancer diagnosis, and the insights gained from this work could be relevant to other medical imaging tasks as well.

Conclusion

This paper introduces a novel test-time adaptation approach called "Diffusion Model Driven Test-Time Image Adaptation" (DMTA) that leverages a diffusion model to improve the robustness of skin lesion classification models. The key innovation is the use of a diffusion model to transform the input images at test time, making them more similar to the training data distribution and thereby boosting the classification performance.

The researchers demonstrate the effectiveness of DMTA on several public skin lesion datasets, achieving state-of-the-art results compared to existing test-time adaptation techniques. This work represents an important step towards building more reliable and trustworthy AI systems for medical image analysis, which could have a significant impact on early detection and treatment of skin cancer.

The insights gained from this research could also be applicable to other medical imaging tasks where distribution shift is a challenge, and the use of diffusion models for test-time adaptation could emerge as a valuable tool in the broader field of machine learning.



This summary was produced with help from an AI and may contain inaccuracies - check out the links to read the original source documents!

Related Papers

šŸ‘€

RadEdit: stress-testing biomedical vision models via diffusion image editing

Fernando P'erez-Garc'ia, Sam Bond-Taylor, Pedro P. Sanchez, Boris van Breugel, Daniel C. Castro, Harshita Sharma, Valentina Salvatelli, Maria T. A. Wetscherek, Hannah Richardson, Matthew P. Lungren, Aditya Nori, Javier Alvarez-Valle, Ozan Oktay, Maximilian Ilse

YC

0

Reddit

0

Biomedical imaging datasets are often small and biased, meaning that real-world performance of predictive models can be substantially lower than expected from internal testing. This work proposes using generative image editing to simulate dataset shifts and diagnose failure modes of biomedical vision models; this can be used in advance of deployment to assess readiness, potentially reducing cost and patient harm. Existing editing methods can produce undesirable changes, with spurious correlations learned due to the co-occurrence of disease and treatment interventions, limiting practical applicability. To address this, we train a text-to-image diffusion model on multiple chest X-ray datasets and introduce a new editing method RadEdit that uses multiple masks, if present, to constrain changes and ensure consistency in the edited images. We consider three types of dataset shifts: acquisition shift, manifestation shift, and population shift, and demonstrate that our approach can diagnose failures and quantify model robustness without additional data collection, complementing more qualitative tools for explainable AI.

Read more

4/4/2024

šŸ›ø

Subject-Diffusion:Open Domain Personalized Text-to-Image Generation without Test-time Fine-tuning

Jian Ma, Junhao Liang, Chen Chen, Haonan Lu

YC

0

Reddit

0

Recent progress in personalized image generation using diffusion models has been significant. However, development in the area of open-domain and non-fine-tuning personalized image generation is proceeding rather slowly. In this paper, we propose Subject-Diffusion, a novel open-domain personalized image generation model that, in addition to not requiring test-time fine-tuning, also only requires a single reference image to support personalized generation of single- or multi-subject in any domain. Firstly, we construct an automatic data labeling tool and use the LAION-Aesthetics dataset to construct a large-scale dataset consisting of 76M images and their corresponding subject detection bounding boxes, segmentation masks and text descriptions. Secondly, we design a new unified framework that combines text and image semantics by incorporating coarse location and fine-grained reference image control to maximize subject fidelity and generalization. Furthermore, we also adopt an attention control mechanism to support multi-subject generation. Extensive qualitative and quantitative results demonstrate that our method outperforms other SOTA frameworks in single, multiple, and human customized image generation. Please refer to our href{https://oppo-mente-lab.github.io/subject_diffusion/}{project page}

Read more

5/21/2024

DG-TTA: Out-of-domain medical image segmentation through Domain Generalization and Test-Time Adaptation

DG-TTA: Out-of-domain medical image segmentation through Domain Generalization and Test-Time Adaptation

Christian Weihsbach, Christian N. Kruse, Alexander Bigalke, Mattias P. Heinrich

YC

0

Reddit

0

Applying pre-trained medical segmentation models on out-of-domain images often yields predictions of insufficient quality. Several strategies have been proposed to maintain model performance, such as finetuning or unsupervised- and source-free domain adaptation. These strategies set restrictive requirements for data availability. In this study, we propose to combine domain generalization and test-time adaptation to create a highly effective approach for reusing pre-trained models in unseen target domains. Domain-generalized pre-training on source data is used to obtain the best initial performance in the target domain. We introduce the MIND descriptor previously used in image registration tasks as a further technique to achieve generalization and present superior performance for small-scale datasets compared to existing approaches. At test-time, high-quality segmentation for every single unseen scan is ensured by optimizing the model weights for consistency given different image augmentations. That way, our method enables separate use of source and target data and thus removes current data availability barriers. Moreover, the presented method is highly modular as it does not require specific model architectures or prior knowledge of involved domains and labels. We demonstrate this by integrating it into the nnUNet, which is currently the most popular and accurate framework for medical image segmentation. We employ multiple datasets covering abdominal, cardiac, and lumbar spine scans and compose several out-of-domain scenarios in this study. We demonstrate that our method, combined with pre-trained whole-body CT models, can effectively segment MR images with high accuracy in all of the aforementioned scenarios. Open-source code can be found here: https://github.com/multimodallearning/DG-TTA

Read more

4/11/2024

GDA: Generalized Diffusion for Robust Test-time Adaptation

GDA: Generalized Diffusion for Robust Test-time Adaptation

Yun-Yun Tsai, Fu-Chen Chen, Albert Y. C. Chen, Junfeng Yang, Che-Chun Su, Min Sun, Cheng-Hao Kuo

YC

0

Reddit

0

Machine learning models struggle with generalization when encountering out-of-distribution (OOD) samples with unexpected distribution shifts. For vision tasks, recent studies have shown that test-time adaptation employing diffusion models can achieve state-of-the-art accuracy improvements on OOD samples by generating new samples that align with the model's domain without the need to modify the model's weights. Unfortunately, those studies have primarily focused on pixel-level corruptions, thereby lacking the generalization to adapt to a broader range of OOD types. We introduce Generalized Diffusion Adaptation (GDA), a novel diffusion-based test-time adaptation method robust against diverse OOD types. Specifically, GDA iteratively guides the diffusion by applying a marginal entropy loss derived from the model, in conjunction with style and content preservation losses during the reverse sampling process. In other words, GDA considers the model's output behavior with the semantic information of the samples as a whole, which can reduce ambiguity in downstream tasks during the generation process. Evaluation across various popular model architectures and OOD benchmarks shows that GDA consistently outperforms prior work on diffusion-driven adaptation. Notably, it achieves the highest classification accuracy improvements, ranging from 4.4% to 5.02% on ImageNet-C and 2.5% to 7.4% on Rendition, Sketch, and Stylized benchmarks. This performance highlights GDA's generalization to a broader range of OOD benchmarks.

Read more

4/3/2024