Re-using Adversarial Mask Discriminators for Test-time Training under Distribution Shifts

Gabriele Valvano10000-0002-4764-2997, Andrea Leo2, Sotirios A. Tsaftaris3
1: IMT School for Advanced Studies Lucca, 2: University of Pisa, Italy, 3: University of Edinburgh, UK
Publication date: 2022/05/30
https://doi.org/10.59275/j.melba.2022-bd5e
PDF · Code · arXiv

Abstract

Thanks to their ability to learn flexible data-driven losses, Generative Adversarial Networks (GANs) are an integral part of many semi- and weakly-supervised methods for medical image segmentation. GANs jointly optimise a generator and an adversarial discriminator on a set of training data. After training is complete, the discriminator is usually discarded, and only the generator is used for inference. But should we discard discriminators? In this work, we argue that training stable discriminators produces expressive loss functions that we can re-use at inference to detect and correct segmentation mistakes. First, we identify key challenges and suggest possible solutions to make discriminators re-usable at inference. Then, we show that we can combine discriminators with image reconstruction costs (via decoders) to endow a causal perspective to test-time training and further improve the model. Our method is simple and improves the test-time performance of pre-trained GANs. Moreover, we show that it is compatible with standard post-processing techniques and it has the potential to be used for Online Continual Learning. With our work, we open new research avenues for re-using adversarial discriminators at inference. Our code is available at https://vios-s.github.io/adversarial-test-time-training.

Keywords

GAN · Discriminator · Segmentation · Test-time training · Shape prior

Bibtex @article{melba:2022:014:valvano, title = "Re-using Adversarial Mask Discriminators for Test-time Training under Distribution Shifts", author = "Valvano, Gabriele and Leo, Andrea and Tsaftaris, Sotirios A.", journal = "Machine Learning for Biomedical Imaging", volume = "1", issue = "MICCAI 2021 workshop omnibus special issue", year = "2022", pages = "1--27", issn = "2766-905X", doi = "https://doi.org/10.59275/j.melba.2022-bd5e", url = "https://melba-journal.org/2022:014" }
RISTY - JOUR AU - Valvano, Gabriele AU - Leo, Andrea AU - Tsaftaris, Sotirios A. PY - 2022 TI - Re-using Adversarial Mask Discriminators for Test-time Training under Distribution Shifts T2 - Machine Learning for Biomedical Imaging VL - 1 IS - MICCAI 2021 workshop omnibus special issue SP - 1 EP - 27 SN - 2766-905X DO - https://doi.org/10.59275/j.melba.2022-bd5e UR - https://melba-journal.org/2022:014 ER -

2022:014 cover