RadEdit: testando a robustez de modelos de visão biomédica por meio de edição de imagens com difusão

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

December 20, 2023

Autores: Fernando Pérez-García, 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

cs.AI

Resumo

Conjuntos de dados de imagens biomédicas são frequentemente pequenos e tendenciosos, o que significa que o desempenho de modelos preditivos no mundo real pode ser substancialmente menor do que o esperado em testes internos. Este trabalho propõe o uso de edição generativa de imagens para simular mudanças nos conjuntos de dados e diagnosticar modos de falha de modelos de visão biomédica; isso pode ser usado antes da implantação para avaliar a prontidão, potencialmente reduzindo custos e danos aos pacientes. Métodos de edição existentes podem produzir alterações indesejáveis, com correlações espúrias aprendidas devido à co-ocorrência de doenças e intervenções de tratamento, limitando a aplicabilidade prática. Para resolver isso, treinamos um modelo de difusão de texto para imagem em múltiplos conjuntos de dados de raios-X de tórax e introduzimos um novo método de edição, o RadEdit, que utiliza múltiplas máscaras, quando presentes, para restringir alterações e garantir consistência nas imagens editadas. Consideramos três tipos de mudanças nos conjuntos de dados: mudança de aquisição, mudança de manifestação e mudança populacional, e demonstramos que nossa abordagem pode diagnosticar falhas e quantificar a robustez do modelo sem a necessidade de coleta adicional de dados, complementando ferramentas mais qualitativas para IA explicável.

English

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.