Bayesian Meta-Learning with Expert Feedback for Task-Shift Adaptation through Causal Embeddings

Lotta Mäkinen; Jorge Loría; Samuel Kaski

cs.LGFeb 23, 2026

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

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 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

Panel of Figure 1 — Figure 1: Illustration of the causal embedding space ( left ) and the mapping to the parameter space ( right ). Each task $t$ has an embedding $z_{t}$ , encoding its causal structure. The linear map $z\mapsto\theta+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

\varepsilon_{\mathrm{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

\varepsilon_{\mathrm{OOD}}

levels for our method against baselines, error bars denote SD. b) Effect of expert query budget on prediction performance, for a single task (

{\varepsilon=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 .

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

\xi_{b}=\arg\!\max_{(i,j)}H[y_{ij}\mid\mathcal{C}]-\mathbb{E}_{z_{t^{\prime}}\sim q_{\varphi}}[H[y_{ij}\mid z_{t^{\prime}}]].

Appendix

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

\tau_{\text{expert}}

for multiple target tasks ordered with increasing distribution shift. Lower

\tau_{\text{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,

\tau_{\text{expert}}=1.0

. The final panel aggregates results across all target tasks.

Method	Hyperparameter	Value
BNN	Inner learning rate	$3\times 10^{-3}$
	Inner temperature	$10^{-1}$
	Global prior std. $\sigma$	$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\!\times\!10^{-5},\,3\!\times\!10^{-3}]$
	Inner temperature	log-uniform $[10^{-4},\,3\!\times\!10^{1}]$
	Global prior std. $\sigma$	log-uniform $[0.02,\,0.2]$
	Prior scaling	log-uniform $[10^{2},\,10^{4}]$
	Model init. log std.	$\{-4,\,-3,\,-2\}$
MAML	Inner learning rate	log-uniform $[10^{-4},\,5\!\times\!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

\chi^{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 (

\bar{z}

) and to the nearest source task. CHI2 distances are reported in units of

\times 10^{3}

. Most OOD task for each metric and method in bold.

Meta-Learning

Meta-learning trains models across tasks so they can adapt quickly to new problems with few examples, limited feedback, or no parameter updates.

See all 164 papers

Demistifying Data and Simulator Assumptions in Supervised Causal Discovery
Sep 29, 2026Pingchuan Ma, Rui Ding, Bojun Huang +1Causal Discovery Causal Structure Learning
Zhejiang University of Technology, China · Microsoft Research Asia · Sony Research +1
Transfer learning for causal forest
Jun 5, 2026Bérénice-Alexia Jocteur, Véronique Maume-Deschamps, Pierre RibereauCausal Effect Estimation Conditional Average Treatment Effect Estimation
ICJ, PSPM · Universit´e Lyon 1, Centrale Lyon, INSA Lyon · Universit´e Jean Monnet, CNRS, ICJ UMR5208, 69622 Villeurbanne, France. +1
MetaCaDI: A Meta-Learning Framework for Causal Discovery from Multiple Environments with Unknown Interventions
Oct 25, 2025Hans Jarett Ong, Yoichi Chikahara, Tomoharu IwataMeta-Learning Causal Discovery
Nara Institute of Science and Technology, Nara, Japan · NTT Communication Science Laboratories, Kyoto, Japan

Sep 29, 2026J/K move · Enter open · S save

Demistifying Data and Simulator Assumptions in Supervised Causal Discovery

Pingchuan Ma, Rui Ding, Bojun Huang, Shuai Wang

Zhejiang University of Technology, China · Microsoft Research Asia · Sony Research+1

Causal Discovery Causal Structure Learning Causal Identifiability Cross-Dataset Generalization

Open paper AskPDF arXiv

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.

Abstract

Meta-Learning

Explore similar work

Demistifying Data and Simulator Assumptions in Supervised Causal Discovery

Transfer learning for causal forest

MetaCaDI: A Meta-Learning Framework for Causal Discovery from Multiple Environments with Unknown Interventions