Bayesian Meta-Learning with Expert Feedback for Task-Shift Adaptation through Causal Embeddings
Authors: Lotta Mäkinen, Jorge Loría, Samuel Kaski
Organizations: Department of Computer Science, Aalto University, Espoo, Finland · ELLIS Institute Finland, Espoo, Finland · Department of Computer Science, University of Manchester, United Kingdom
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 both controlled simulations and a large-scale real-world clinical prediction setting for cross-disease transfer, where causal embeddings align with underlying clinical mechanisms.
Figures & tables
Figure 1: Illustration of the causal embedding space ( left ) and the mapping to the parameter space ( right ). Each task t has an embedding zt , encoding its causal structure. The linear map z↦θ+Wz transforms embeddings into task-specific priors, enabling tasks with similar causal mechanisms to have similar priors.
Figure 2: Performance of meta-learning models under increasing task-shift for Experiment 1. a) AUROC as a function of the distribution shift εOOD . b) Change in log loss relative to no transfer baseline (BNN), where negative values indicate positive transfer and positive values negative transfer. Error bars denote standard deviation across 30 runs.
Figure 3: Method comparison with expert inferred embeddings for Experiment 2 across 30 runs. a) Average AUROC across increasing distribution shift εOOD levels for our method against baselines, error bars denote SD. b) Effect of expert query budget on prediction performance, for a single task ( ε=4.0 ).
Method
J44
J45
G45
I21
CHI2, exp.
0.730 (.028)
0.571 (.015)
0.631 (.013)
0.690 (.015)
ICP, exp.
0.820 (.007)
0.620 (.007)
0.655 (.007)
0.693 (.008)
MR, exp.
0.805 (.009)
0.618 (.009)
0.656 (.007)
0.677 (.009)
CHI2, oracle
0.815 (.010)
0.597 (.007)
0.671 (.007)
0.705 (.006)
ICP, oracle
0.816 (.006)
0.622 (.004)
0.673 (.007)
0.714 (.007)
MR, oracle
0.766 (.011)
0.602 (.006)
0.651 (.009)
0.711 (.008)
Table 1: Average AUROC (SD) on target tasks from the UKBB, over 10 seeds. CHI2, ICP, and MR correspond todifferent embeddings in our method. Best in bold .
Algorithm B.3 Expert-guided inference of target task embedding
Figure 4: Performance of causally-aware meta-learning models under increasing noise in the causal task embeddings.
Figure 5: Expert-noise sensitivity under active querying. RMSE as a function of the number of expert queries at different noise levels τexpert for multiple target tasks ordered with increasing distribution shift. Lower τexpert corresponds to noisier expert feedback. The final panel aggregates results across all target tasks.
Figure 6: Random vs. active (BALD) querying. RMSE as a function of the number of expert queries for multiple target tasks ordered with increasing distribution shift, τexpert=1.0 . The final panel aggregates results across all target tasks.
Method
Hyperparameter
Value
BNN
Inner learning rate
3×10−3
Inner temperature
10−1
Global prior std. σ
1.0
Model init. log std.
−1
MAML
Inner learning rate
10−2
Outer learning rate
10−3
Appendix
Table 2: Final hyperparameter values used in synthetic experiments.
ICD-10
Category
Disease description
Role
Cases
Prevalence (%)
I21
C
Acute myocardial infarction
Target
3,941
1.79
G45
C
Transient ischaemic attack (TIA)
Target
2,399
1.09
J44
R
Chronic obstructive pulmonary disease
Target
5,850
2.65
J45
R
Asthma
Target
4,318
1.96
E11
C
Type 2 diabetes mellitus
Source
9,571
4.34
I10
C
Essential (primary) hypertension
Source
32,153
14.58
Appendix
Table 4: Overview of disease tasks used in the experiments, including ICD-10 codes, disease category (cardiovascular-metabolic = C or respiratory = R), clinical descriptions, task roles, and prevalence in the UK Biobank cohort (N=220,571).
Method
Hyperparameter
Search Range
BNN
Inner learning rate
log-uniform [3×10−5,3×10−3]
Inner temperature
log-uniform [10−4,3×101]
Global prior std. σ
log-uniform [0.02,0.2]
Prior scaling
log-uniform [102,104]
Model init. log std.
{−4,−3,−2}
MAML
Inner learning rate
log-uniform [10−4,5×10−2]
Appendix
Table 5: Hyperparameter tuning ranges.
AUPRC
MCC
Method
J44
J45
G45
I21
J44
J45
G45
I21
CHI2, exp
0.074 (.013)
0.029 (.002)
0.017 (.001)
0.039 (.004)
0.122 (.020)
0.030 (.007)
0.040 (.005)
0.075 (.008)
ICP, exp.
0.119 (.009)
0.036 (.001)
0.019 (.001)
0.041 (.002)
0.176 (.005)
0.055 (.004)
0.047 (.004)
0.076 (.004)
MR, exp.
0.107 (.006)
0.038 (.002)
0.020 (.001)
0.039 (.002)
0.164 (.009)
0.055 (.004)
0.046 (.003)
0.070 (.005)
CHI2, oracle
0.126 (.009)
0.032 (.001)
0.020 (.001)
0.043 (.003)
0.174 (.013)
0.044 (.005)
0.053 (.003)
0.082 (.004)
ICP, oracle
0.120 (.007)
0.036 (.002)
0.021 (.001)
0.045 (.003)
0.172 (.007)
0.054 (.003)
0.052 (.004)
0.085 (.005)
Appendix
Table 6: Average (SD) AUPRC and MCC on the UK Biobank results (over 10 seeds) on target tasks. CHI2, ICP, and MR correspond to our method with different embeddings. Best in bold .
Method
Elapsed Time
Peak Memory
(hh:mm:ss)
(GiB)
CHI2, exp
00:13:09
7.5
ICP, exp
00:12:13
7.1
MR, exp
00:20:39
7.7
CHI2, oracle
00:08:11
7.1
ICP, oracle
00:08:21
7.0
Appendix
Table 7: Computational cost of different methods on the UK Biobank experiments. Reported values correspond to average wall-clock time and peak memory usage per task.
Figure 7: Principal component projections of causal task embeddings for Invariant Causal Prediction (ICP), Mendelian Randomisation (MR) and χ2 (CHI2). Colors indicate disease category (cardiovascular-metabolic vs. respiratory). Hollow markers indicate target tasks.
Figure 8: Three-nearest neighbor graphs of causal task embeddings. Nodes represent disease prediction tasks and edges connect tasks with similar causal structures. Colors indicate disease category (cardiovascular-metabolic vs. respiratory). Hollow markers indicate target tasks.
ICP
MR
CHI2
Target
Avg. source
Nearest source
Avg. source
Nearest source
Avg. source
Nearest source
J44
3.39 (1)
2.29 (1)
0.69 (4)
0.59 (3)
29.19 (4)
38.63 (4)
J45
2.61 (2)
2.10 (3)
0.94 (2)
0.72 (2)
29.54 (1)
39.39 (1)
G45
1.93 (3)
2.15 (2)
0.79 (3)
0.50 (4)
29.46 (2)
39.24 (2)
I21
1.42 (4)
1.38 (4)
1.02 (1)
1.23 (1)
29.34 (3)
39.12 (3)
Appendix
Table 8: Characterization of the target tasks in the embedding space, reported as the Euclidean distance and rank (in parenthesis, 1 indicates the most OOD, 4 is the least) to average source task ( zˉ ) and to the nearest source task. CHI2 distances are reported in units of ×103 . Most OOD task for each metric and method in bold.
Supervised causal discovery learns to infer causal structure for a new dataset from training datasets paired with structural labels. These training pairs are typically simulated, making the simulator both a source of supervision and a carrier of assumptions about causal graphs, mechanisms, and noise. Understanding the resulting predictions therefore requires examining how these assumptions supplement the information available in observational data, which may be compatible with multiple causal graphs. This paper examines that relationship across representative methods available through June 2026. We organize these methods by prediction target, prediction granularity, encoder, structural decoder, and training regime to relate what each method predicts to how it uses data and simulator-based supervision. Using this framework, we distinguish two questions: whether the target is identifiable under the assumed model class, and whether a trained predictor generalizes beyond its training distribution. Restrictions on mechanisms and noise can make otherwise ambiguous causal directions identifiable, but predictive accuracy under those restrictions does not establish transfer when they change. This distinction motivates evaluation that matches metrics to the identifiable graph target and tests changes in graphs, mechanisms, and noise between training and deployment. Extending such evaluation to real data also requires documenting the external causal evidence and uncertainty behind benchmark reference graphs. Together, these analyses guide method comparison and identify open questions in transfer, test-time adaptation, and uncertainty assessment.
Pingchuan Ma, Rui Ding, Bojun Huang +1
Zhejiang University of Technology, China · Microsoft Research Asia · Sony Research +1
Transfer learning addresses the challenge of transfering knowledge from one domain to another. Traditional transfer learning focuses on adapting models trained on a source domain (with a lot of observations) to improve performance on a target domain (with few observations). In this work we consider the case of a model shift and we focus on the transfer learning applied to a causal forest namely HTERF. This causal forest aims to estimate the Conditional Average Treatment Effect (CATE). The approach considered is the offset method presented by Wang (2016) adapted to a causal context. This method relies on the use of intermediate models in order to estimate the offset between source and target distributions. Our main result is a bound on the CATE error of HTERF on target depending on the error of the intermediate models. Simulation studies show the good performances of this approach in different settings on simulations and on a real-world dataset.
Bérénice-Alexia Jocteur, Véronique Maume-Deschamps, Pierre Ribereau
ICJ, PSPM · Universit´e Lyon 1, Centrale Lyon, INSA Lyon · Universit´e Jean Monnet, CNRS, ICJ UMR5208, 69622 Villeurbanne, France. +1
Uncovering the causal mechanisms of complex real-world systems remains a significant challenge, as these systems often entail high data collection costs and involve unknown interventions. We introduce MetaCaDI, the first framework to cast the identification of unknown interventions as a meta-learning problem, explicitly leveraging a jointly learned causal graph. MetaCaDI is a Bayesian framework that learns a shared causal structure across multiple environments and is optimized to rapidly adapt to new, few-shot intervention target identification tasks. A key innovation is our model's analytical adaptation, which uses a closed-form solution to bypass expensive and potentially unstable gradient-based bilevel optimization. Extensive experiments on synthetic and complex gene expression data demonstrate that MetaCaDI significantly outperforms state-of-the-art methods. It excels at identifying intervention targets from as few as 3 samples - where existing methods collapse to random chance - while robustly recovering the shared causal graph, proving its effectiveness in data-scarce scenarios.
Hans Jarett Ong, Yoichi Chikahara, Tomoharu Iwata
Nara Institute of Science and Technology, Nara, Japan · NTT Communication Science Laboratories, Kyoto, Japan