Is merging worth it? Evaluating the information gain for causal dataset acquisition whilst preserving privacy
Merging datasets across institutions is a lengthy and costly procedure, especially when it involves private information. Data hosts may therefore want to prospectively gauge which datasets are most beneficial to merge with, without revealing sensitive information. For causal estimation this is particularly challenging as the value of a merge will depend not only on the reduction in epistemic uncertainty but also the improvement in overlap. To address this challenge, we introduce the first cryptographically secure information-theoretic approach for quantifying the value of a merge in the context of heterogeneous treatment effect estimation. To the best of our knowledge, this work presents the first privacy-preserving method for dataset acquisition tailored to causal estimation. Our code is available here.
1. Motivation
Merging datasets is often a time-consuming and resource-intensive task. Practitioners thus crucially need methodology for determining in advance if a merge is worth it, whilst also complying with privacy requirements [1]. In this work, we focus on data fusion in the context of heterogeneous treatment effect estimation.
Illustration example: Imagine a hospital wants to assess the effect of a medical intervention on its patients but finds its own data is insufficient to get accurate estimates. Therefore, it plans to select one of K local hospitals in which to conduct a follow-up study. As such a study is expensive and time consuming, the hospital would like to identify in advance which potential dataset would provide the most information upon merging whilst complying with privacy regulations.
→ here, we propose a solution for this problem, allowing for patient outcomes at the other hospitals to be unobserved or simply masked.
To do this we apply techniques from Bayesian experimental design [2–4] to quantify the value of a merge in a principled, information theoretic way.Our solution has similarities with standard Bayesian dataset acquisition [5], however in causal contexts data fusion serves an additional, distinct purpose. Specifically, acquired data should not only enhance our understanding of the outcome function, but also assist in combating the selection bias that is inherent to causal effect estimation [6] by balancing treatment. Put differently, generic data fusion only aims at reducing the epistemic uncertainty stemming from incomplete knowledge of the outcome, whilst causal estimation also seeks to improve treatment overlap [7], i.e. reducing the epistemic uncertainty for counterfactual outcomes.
For privacy, we make use of Secure Multi-Party Computation (MPC) [8, 9]. In our context it allows the different candidate sites to compute their information gains relative to the initial sites’ data, without anyone having to reveal the content of their dataset.
This work was done in collaboration with Jake Fawkes, Desi Ivanova, Uri Shalit, and Chris Holmes. The complete paper can be found on arXiv.
Problem statement:
- We consider an initial host dataset D_0.
- A number of potential datasets are available for merging, which are referred to as the candidate sites and denoted D_e, where e ∈ {0, · · · , K } indexes the datasets.
- We let X , T , and Y represent the covariates, treatment, and outcomes in each dataset.
- We further assume that the outcomes can be unmeasured or masked in the candidate datasets.
Our goal is to estimate which of the candidate datasets De, would reduce the uncertainty over the Conditional Average Treatment Effect (CATE) if we were to observe Y and merge with the host dataset D_0.
We assume that datasets are generated according to the Directed Acyclic Graph (DAG) in Figure 1. This implies:
- P(y|x,t,e) is fixed across environments,
- both covariate distributions and treatment allocation schemes are free to vary i.e. P (x, t|e) can depend on e.
We also assume positivity the whole population, so that ∀x,0 < P(T = 1|X = x) < 1, and consistency. Altogether, this induces that CATE is identified and, constant across environments e.
2. Method
2.1 Quantifying Data Merge Utility through Expected Information Gain
We take a Bayesian approach to modelling the CATE τ(x):
- we define f_θ : X × {0, 1} → Y to be a function representing the expected outcome conditional on covariates and treatment, i.e., fθ (x, t) seeks to approximate E [Y |X = x, T = t]
- we assign a prior distribution P (θ), to the parameters θ and denote the conditional likelihood of the outcome given parameters, covariates, and treatment as P (Y | x, t, θ)
Here, we focus on continuous Y and use a normal likelihood with fixed variance σ.
To measure the value of a merge with a candidate dataset, we measure the expected information gain (EIG), which measures the expected reduction in uncertainty in the parameters, as measured by Shannon entropy, when moving from the post host posterior, P (θ|D0) to the post merge posterior, P (θ|D0, De).
Mathematically, the EIG is expressed as:
which can be rewritten as the expected log ratio of the likelihood to the marginal likelihood:
and is also known as Bayesian Active Learning by Disagreement (BALD) in the active learning literature [10]. This quantity can be approximated using nested Monte Carlo [11, 12] as follows:
Using these estimators we can gauge the value of a merge before observing outcomes in the dataset De. However, whilst this does provide us with information gain on the parameters for the conditional outcome it may not necessarily lead to an improvement in predicting CATE. This is because EIGθ|D0 encourages uniform entropy reduction across all dimensions of θ, not just the ones relevant for CATE estimation. For instance, a dataset of untreated individuals would still provide information about the conditional outcome but much less for the contrast of outcomes required for CATE. This problem motivates our improved methodology given in the next subsection.
2.2 Targeted EIG in CATE Parameters: An Improved Approach to Merge Evaluation
We split the parameter set θ as θ = θc ∪ θnc where θc parameterises the CATE model and θnc is a set of nuisance parameters. Leveraging such a parameterisation, we can prioritise uncertainty reduction in θc and therefore CATE by targeting the following:
where the numerator is the Bayesian distribution of the outcomes conditional on θc only
As P(ye|θc,Xe,te,D0) is given by integrating out θnc, it isn’t available in closed form. Therefore, we use a second nested Monte Carlo estimator to approximate P (ye |θc , Xe , te , D0 ) in the numerator of the previous equation. This comes in addition to the previous nested Monte Carlo estimator of the posterior predictive likelihood P (ye |Xe , te , D0 ) used in Equation (4). Ultimately, we estimated the EIG in CATE parameters:
2.3 Procedure and Model Classes
We apply both procedures to three popular Bayesian causal inference methods: Bayesian Polynomial Regression, Bayesian Causal Forests [13], and Causal Multi-task Gaussian Processes. Further details on the parametrisation of each model can be found in section 3.3 of our paper.
4. Privacy
For privacy, we make use of Multi-Party Computation (MPC) and adopt a semi-honest setting, in which all parties follow the specified protocol, but some corrupt parties will try to learn as much about their peers inputs in the process. The goal is to devise a protocol which will preserve the privacy of the non-corrupt party’s inputs, up to a given computational budget by the adversary. Here, this means that any collection of corrupt sites should be unable to learn anything about the other sites data whilst performing the EIG calculation.
Get Lucile Ter-Minassian’s stories in your inbox
Join Medium for free to get updates from this writer.
For implementing multi-party computation, we employ the open source library CrypTen [14]. CrypTen builds upon PyTorch allowing for standard tensor operations to be performed in an MPC protocol.
4. Experiments and results
4.1 Procedure
We experimentally validate our approach on a setting where the host has to rank a number of candidate datasets based on the estimated gain from merging with each site. We evaluate on a selection of synthetic and semi-synthetic benchmark datasets as this allows us to form a ground truth ranking using the provided CATE values. We ultimately compare the implied rankings with the ground truth ranking given by the Precision in Estimation of Heterogeneous Effects (PEHE) performance of models trained on the merged datasets D0 ∪ De where the PEHE is estimated with hold out samples from P0(x).
This procedure is repeated for each model class, to make sure the same predictor is used to gauge the information gain and give the ground truth ranking, leading to one true ranking per model class.
4.2 Experimental setting
Datasets We validate our framework in a setting where the host must choose between numerous potential candidates, each with different distributions. To do so we begin with the IHDP dataset, and form the host and candidates datasets using subsampling functions, Se(x,t) as detailed above where subsampling function is a logistic regression with random parameters. This ensures that each site has a different covariate distribution. We apply the three methods described in 2.3 to estimate both the two expected information gains (targeted or not) for each candidate site, De.
We ultimately compare the implied rankings with the ground truth ranking given by the PEHE by Spearman ρ and precision at k.
Baselines We provide a number of simple comparison methods as baselines for our task. Specifically, a) ranking by sample size, b) ranking by similarity of covariate distribution measured by a multivariate Gaussian fit to the host, and c) ranking by dissimilarity of treatment allocation measured by the error of a propensity model fit on the host.
Results Table 1 shows the results of our experiment on the IHDP dataset where the host has to rank the best datasets out of 10 candidates. We report the average performance of the rankings across 20 repeated experiments, and according to Spearman ρ [61] and precision at k. We include the best baseline by Spearman ρ performance. Additional metrics and the performance of all baselines can be found in the complete paper.
These results demonstrate that across all models, our EIGθ based approach which focuses on causal parameters outperforms the standard approach which seeks expected entropy reduction in all parameters. Results also show that our method works significantly better than all baselines.
4.3 Other experiments
Similar results can be found in the complete paper.
Furthermore, the document features:
- an illustrative experiment that “forces” a merge to be prioritized with a “natural” candidate for CATE estimation (section 5.1 of the complete paper)
- an MPC experiment. We repeat a similar experiment on the IHDP dataset, this time choosing between twenty datasets. (section 5.3 of the complete paper). It shows that MPC vastly outperforms differential privacy, both in terms of MSE and ranking with DP.
5. Discussion
Limitations
- Computationally costly in high dimensions due to the cost of computing high dimensional determinants and performing multiple rounds of conditional sampling. However, both of these remain negligible compared to the data engineering and expenses of data fusion [16, 17].
- It might be beneficial to consider introducing a problem-specific threshold to determine when a proposed merger is worthwhile.
Conclusion
We introduce an information-theoretic, cryptograpically secure framework for evaluating the utility of potential data merges for causal estimation. To the best of our knowledge, this is the first work addressing this relevant challenge.
Through empirical evaluation, we demonstrate that our framework can reliably rank datasets according to their ability to improve CATE estimation. We show that entropy reduction in the CATE parameters alone gives an improvement when compared to a more standard approach to Bayesian dataset acquisition. Finally, we demonstrate that our cryptographic procedure can be applied with minimal loss of accuracy, vastly outperforming deferentially privacy.
References
[1] K. Abouelmehdi, A. Beni-Hessane, and H. Khaloufi. Big healthcare data: preserving security and privacy. Journal of big data, 5(1):1–18, 2018. 1
[2] K. Chaloner and I. Verdinelli. Bayesian experimental design: A review. Statistical science, pages 273–304, 1995. 2, 4
[3] D. V. Lindley. On a measure of the information provided by an experiment. The Annals of Mathematical Statistics, 27(4):986–1005, 1956. 2, 4
[4] T. Rainforth, A. Foster, D. R. Ivanova, and F. Bickford Smith. Modern bayesian experimental design. Statistical Science, 39(1):100–114, 2024. 2, 4
[5] A. Kirsch, J. Van Amersfoort, and Y. Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. Advances in neural information processing systems, 32, 2019. 2
[6] P. W. Holland. Statistics and causal inference. Journal of the American statistical Association, 81(396):945–960, 1986. 2
[7] D. B. Rubin. Estimating causal effects from large data sets using propensity scores. Annals of internal medicine, 127(8_Part_2):757–763, 1997.
[8] D. Evans, V. Kolesnikov, M. Rosulek, et al. A pragmatic introduction to secure multi-party computation. Foundations and Trends® in Privacy and Security, 2(2-3):70–246, 2018. 2, 6
[9] B. Knott, S. Venkataraman, A. Hannun, S. Sengupta, M. Ibrahim, and L. van der Maaten. Crypten: Secure multi-party computation meets machine learning. Advances in Neural Informa- tion Processing Systems, 34:4961–4973, 2021. 2, 6
[10] N. Houlsby, F. Huszár, Z. Ghahramani, and M. Lengyel. Bayesian active learning for classifica- tion and preference learning. arXiv preprint arXiv:1112.5745, 2011. 4, 6, 9, 22
[11] A. E. Foster. Variational, Monte Carlo and policy-based approaches to Bayesian experimental design. PhD thesis, University of Oxford, 2021. 4
[12] T. Rainforth, R. Cornish, H. Yang, A. Warrington, and F. Wood. On nesting monte carlo estimators. In International Conference on Machine Learning, pages 4267–4276. PMLR, 2018. 4
[13] P. R. Hahn, J. S. Murray, and C. M. Carvalho. Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion). Bayesian Analysis, 15(3):965–1056, 2020. 2, 5, 6, 17
[14] B. Knott, S. Venkataraman, A. Hannun, S. Sengupta, M. Ibrahim, and L. van der Maaten. Crypten: Secure multi-party computation meets machine learning. Advances in Neural Informa- tion Processing Systems, 34:4961–4973, 2021. 2, 6
[15] C.Louizos, U.Shalit, J .M.Mooij, D.Sontag, R.Zemel, and M.Welling. Causaleffectinference with deep latent-variable models. Advances in neural information processing systems, 30, 2017. 9
[16] M. L. Brodie. Data integration at scale: From relational data integration to information ecosystems. In 2010 24th IEEE International Conference on Advanced Information Networking and Applications, pages 2–3. IEEE, 2010. 10
[17] A. Kadadi, R. Agrawal, C. Nyamful, and R. Atiq. Challenges of data integration and interoper- ability in big data. In 2014 IEEE international conference on big data (big data), pages 38–40. IEEE, 2014. 10
