Bayesian Meta-Learning with Expert Feedback for Task-Shift Adaptation through Causal Embeddings
Abstract
Meta-learning methods perform well on new within-distribution tasks but often fail when adapting to out-of-distribution target tasks, where transfer from source tasks can induce negative transfer. We propose a causally-aware Bayesian meta-learning method, by conditioning task-specific priors on precomputed latent causal task embeddings, enabling transfer based on mechanistic similarity rather than spurious correlations. Our approach explicitly considers realistic deployment settings where access to target-task data is limited, and adaptation relies on noisy (expert-provided) pairwise judgments of causal similarity between source and target tasks. We provide a theoretical analysis showing that conditioning on causal embeddings controls prior mismatch and mitigates negative transfer under task shift. Empirically, we demonstrate reductions in negative transfer and improved out-of-distribution adaptation in controlled simulations and a real-world clinical prediction setting for cross-disease transfer, where causal embeddings align with underlying clinical mechanisms; we include the judgments from a medical expert in the clinical prediction task and obtain improved performance in predictions of unseen diseases.
Lay Summary
Machine learning systems often struggle when applied to new situations that differ from their training data. In healthcare, for example, diseases can look similar on the surface while arising from very different biological mechanisms. This can cause models to rely on misleading patterns and make unreliable predictions for new patient groups or diseases. We developed a method that transfers knowledge between related prediction problems by focusing on the underlying mechanisms that connect them. Our approach can extend to new prediction problems, by asking a domain expert for simple comparisons ("is the new problem more similar to A or B?"). The model uses this information to guide how knowledge is transferred from previous (related) problems. In healthcare, this means that a medical expert can indicate which diseases share similar biological causes, helping the model adapt to new diseases with only limited data. Our method improves prediction performance in challenging settings where standard approaches are less reliable. This helps build more dependable machine learning systems, especially for rare and emerging diseases and changing patient populations.