As artificial intelligence (AI) continues to revolutionize healthcare and medical image analysis, ensuring safe and effective AI deployment in open clinical environments has become paramount. However, many existing medical AI methods prioritize model performance, focusing on achieving higher accuracy rather than clinical applicability. In this talk, I will present our works on building clinically applicable deep learning systems for medical image analysis, emphasizing domain generalization, continual learning, and multi-modality learning. Our aim is to inspire the development of more reliable and effective medical AI systems, ultimately enhancing patient care and outcomes. Additionally, I will discuss up-to-date progress and promising future directions in this critical domain.

Recently, there has been a growing surge of interest in enabling machine learning systems to generalize well to Out-of-Distribution (OOD) data. Most efforts are devoted to advancing optimization objectives that regularize models to capture the underlying invariance; however, there often are compromises in the optimization process of these OOD objectives: i) Many OOD objectives have to be relaxed as penalty terms of Empirical Risk Minimization (ERM) for the ease of optimization, while the relaxed forms can weaken the robustness of the original objective; ii) The penalty terms also require careful tuning of the penalty weights due to the intrinsic conflicts between ERM and OOD objectives. Consequently, these compromises could easily lead to suboptimal performance of either the ERM or OOD objective. To address these issues, we introduce a multi-objective optimization (MOO) perspective to understand the OOD optimization process, and propose a new optimization scheme called PAreto Invariant Risk Minimization (PAIR). PAIR improves the robustness of OOD objectives by cooperatively optimizing with other OOD objectives, thereby bridging the gaps caused by the relaxations. Then PAIR approaches a Pareto optimal solution that trades off the ERM and OOD objectives properly. Extensive experiments on challenging benchmarks, WILDS, show that PAIR alleviates the compromises and yields top OOD performances.

Developing robust approaches for domain generalization is critical for the real world deployment of deep learning models. Here, we address a particular domain generalization challenge: selective classification for automated medical image diagnosis. In this setting, models must learn to abstain from making predictions when label confidence is low, especially when tested with samples that deviate significantly from the training set (covariate shift). Using the example of diabetic retinopathy detection we show that even state-of-the-art deep learning models, including Bayesian networks, fail during selective classification under covariate shift. Bayesian estimates of predictive uncertainty do not generalize well under covariate shift yielding catastrophic performance drops during referral. We identify the source of these failures and propose several post hoc referral solutions that enable reliable selective classification under covariate shift.

Mark has spent the last year developing a state-of-the-art neural network for his favourite domain. But after deployment, everyone starts complaining. But surely they really should stop using it in settings that different from Mark's training scenario? Robustness matters. Any machine learning method needs to be broadly realistically applicable, must specify the domain of application, and must work across that domain. The restriction that that test environment match the training setting is neither well defined nor realistically applicable. Hence it is our responsibility to deal with domain shift, and all that entails. But why is it so hard? Why do our models break when we try to use them in the real world? Why do neural networks seem particularly susceptible to this? How do we understand and mitigate these issues? And what tools are at our disposal to build models that deal better with domain shift from the outset? We take a causal look at domain shift and look at approaches that enable improved capability when training and test domains are different.

We study the relationship between the entropy of intermediate representations and a model’s robustness to distributional shift. We train models consisting of two feed-forward networks end-to-end separated by a discrete n-bit channel on an unsupervised contrastive learning task. Different masking strategies are applied after training that remove a proportion of low-entropy bits, high-entropy bits, or randomly selected bits, and the effects on performance are compared to the baseline accuracy with no mask. We hypothesize that the entropy of a bit serves as a guide to its usefulness out-of-distribution (OOD). Through experiment on three OOD datasets we demonstrate that the removal of low-entropy bits can notably benefit OOD performance. Conversely, we find that top-entropy masking disproportionately harms performance both in-distribution (InD) and OOD.

Deep reinforcement learning repeatedly succeeds in closed, well-defined domains such as games (Chess, Go, StarCraft). The next frontier is real-world scenarios, where setups are numerous and varied. For this, agents need to learn the underlying rules governing the environment, so as to robustly generalize to conditions that differ from those they were trained on. Model-based reinforcement learning algorithms, such as the highly successful MuZero, aim to accomplish this by learning a world model. However, leveraging a world model has not consistently shown greater generalization capabilities compared to model-free alternatives. In this work, we propose improving the data efficiency and generalization capabilities of MuZero by explicitly incorporating the symmetries of the environment in its world-model architecture. We prove that, so long as the neural networks used by MuZero are equivariant to a particular symmetry group acting on the environment, the entirety of MuZero's action-selection algorithm will also be equivariant to that group. We evaluate Equivariant MuZero on procedurally-generated MiniPacman and on Chaser from the ProcGen suite: training on a set of mazes, and then testing on unseen rotated versions, demonstrating the benefits of equivariance. Further, we verify that our performance improvements hold even when only some of the components of Equivariant MuZero obey strict equivariance, which highlights the robustness of our construction.