Machine learning (ML) has been instrumental in advancing healthcare, especially in the realm of medical imaging. However, current models often fall short in explaining how visual changes impact ML decisions, creating a need for transparent models that not only classify medical imagery accurately but also elucidate the signals and patterns they learn. Google’s new framework, StylEx, utilises generative AI to tackle the lack of explainability in AI models applied to medical imaging.
Usually, AI models in computer vision, such as those used in medical imaging, generate heatmaps that indicate the significance of different pixels in an image. These tools can be limited in their capacity to adequately explain the ‘what’ and ‘why’ behind important image features like texture, shape, or size. Contrary to the norm, StylEx, driven by a StyleGAN-based image generator and shaped by a classifier, can generate hypotheses by identifying and visualizing visual signals connected with a classifier’s predictions.
The StylEx methodology consists of four stages. Firstly, it trains a classifier on a medical imaging dataset for a designated task. After accomplishing high performance, it verifies the existence of pertinent signals within the imageries. Secondly, it trains a StyleGAN2 generator to create realistic images while preserving the decision-making of the classifier. In the third stage, it identifies the top attributes in the generator that influence the classifier’s predictions, through manipulation of each co-ordinate in the StyleSpace. Consequently, it produces counterfactual visualizations that display the effect of each attribute independently. Lastly, an interdisciplinary panel of experts reviews these findings to formulate hypotheses for future research, considering both biological and socio-cultural health determinants.
This innovative framework, therefore, enhances the comprehensibility of AI models in medical imaging. It provides an in-depth understanding of the ‘what’ behind model’s decisions by creating counterfactual images and visualising the attributes affecting classifier predictions. Furthermore, the contributions from an interdisciplinary panel of professionals ensure that the analysis encompasses a broad spectrum of perspectives, ranging from clinical to socio-cultural, thereby reducing potential biases and opening new frontiers for scientific exploration.