Re-using Adversarial Mask Discriminators for Test-time Training under Distribution Shifts
arXiv:2108.11926 · doi:10.59275/j.melba.2022-bd5e
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 \textit{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.
arXiv commentAccepted for publication at the Journal of Machine Learning for Biomedical Imaging (MELBA) https://www.melba-journal.org/papers/2022:014.html
References in corpus (40)
- Adam: A Method for Stochastic Optimization
- Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
- Overcoming catastrophic forgetting in neural networks
- Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks
- Spectral Normalization for Generative Adversarial Networks
- Conditional Random Fields as Recurrent Neural Networks
- Generative Adversarial Network in Medical Imaging: A Review
- On Calibration of Modern Neural Networks
- A continual learning survey: Defying forgetting in classification tasks
- Improved Training of Wasserstein GANs
- Domain Generalization: A Survey
- When Does Label Smoothing Help?
- CHAOS Challenge -- Combined (CT-MR) Healthy Abdominal Organ Segmentation
- Anatomically Constrained Neural Networks (ACNN): Application to Cardiac Image Enhancement and Segmentation
- On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima
- Domain Generalization via Model-Agnostic Learning of Semantic Features
- Embracing Imperfect Datasets: A Review of Deep Learning Solutions for Medical Image Segmentation
- Constrained-CNN losses for weakly supervised segmentation
- Tent: Fully Test-time Adaptation by Entropy Minimization
- A Topological Loss Function for Deep-Learning based Image Segmentation using Persistent Homology
- Disentangled Representation Learning in Cardiac Image Analysis
- Continual Learning with Deep Generative Replay
- Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation
- Amortised MAP Inference for Image Super-resolution
- Cardiac Segmentation with Strong Anatomical Guarantees
- Do CIFAR-10 Classifiers Generalize to CIFAR-10?
- A Meta-Transfer Objective for Learning to Disentangle Causal Mechanisms
- Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation
- MEMO: Test Time Robustness via Adaptation and Augmentation
- Test-Time Training with Self-Supervision for Generalization under Distribution Shifts
- Cardiac MRI Segmentation with Strong Anatomical Guarantees
- Incorporating prior knowledge in medical image segmentation: a survey
- A Large-Scale Study on Regularization and Normalization in GANs
- Memorization Precedes Generation: Learning Unsupervised GANs with Memory Networks
- Adversarially Learned Anomaly Detection
- On the Effectiveness of Least Squares Generative Adversarial Networks
- MT3: Meta Test-Time Training for Self-Supervised Test-Time Adaption
- Is Discriminator a Good Feature Extractor?
- Fence GAN: Towards Better Anomaly Detection
- Collaborative Sampling in Generative Adversarial Networks