Generation and Evaluation of Synthetic Data Containing Treatments

[1] Seedat, Nabeel, et al. "Curated llm: Synergy of llms and data curation for tabular augmentation in ultra low-data regimes." arXiv preprint arXiv:2312.12112 (2023).

[2] Li, Jiayu, et al. "TAEGAN: Generating Synthetic Tabular Data For Data Augmentation." arXiv preprint arXiv:2410.01933 (2024).

[3] Liu, Tennison, et al. "GOGGLE: Generative modelling for tabular data by learning relational structure." The Eleventh International Conference on Learning Representations. 2023.

[4] Rashidi, Hooman H., et al. "A novel and fully automated platform for synthetic tabular data generation and validation." Scientific Reports 14.1 (2024): 23312.

[5] Ryan McKenna, Brett Mullins, Daniel Sheldon, and Gerome Miklau. Aim: An adaptive and iterative mechanism for differentially private synthetic data, 2024. URL https://arxiv.org/abs/2201.12677.

[6] Terrance Liu, Giuseppe Vietri, and Steven Z Wu. Iterative methods for private synthetic data: Unifying framework and new methods. Advances in Neural Information Processing Systems, 34: 690–702, 2021.

[7] McKenna, Ryan & Miklau, Gerome & Sheldon, Daniel. (2021). Winning the NIST Contest: A scalable and general approach to differentially private synthetic data. Journal of Privacy and Confidentiality. 11. 10.29012/jpc.778.

[8] Sergul Aydore, William Brown, Michael Kearns, Krishnaram Kenthapadi, Luca Melis, Aaron Roth, and Ankit A Siva. Differentially private query release through adaptive projection. In International Conference on Machine Learning, pp. 457–467. PMLR, 2021.

With the current uniform allocation of ($\delta$, $\epsilon$) amongst the STEAM components, these results can be seen as a lower bound on the potential performance of STEAM in DP generation. As such, we are encouraged that STEAM remains competitive against this updated set of baseline models, and we see attempts at sophisticated optimisation of the privacy budget distribution as a fruitful area of future work.

Update: We have extended our related works on DP generation in Appendix B to better represent the recent literature. We have changed Section 7.3 and added new results to compare DP generation performance between AIM and $STEAM_{AIM}$, on the ACTG dataset. We have added more results for other DP baselines tested (GEM, MST, RAP) in Appendix I.

(C) Formal motivation for new metrics

You are correct that in Section 5, in our illustrative experiment on the failure of current metrics to identify differences in $Q_{Y|W,\mathbf{X}}$, we only highlight the improvement of our $U_{PEHE}$ metric. However, as we state in our original manuscript, a similar experiment, which instead shows how current metrics also struggle to identify differences in $Q_{W|\mathbf{X}}$, in comparison to our $JSD_\pi$ metric, is also included in Appendix C.

Nevertheless, we agree that a more formal approach can be made to show that existing, joint-distribution-level metrics, such as Kullback-Leibler, will tend to lose sight of changes in $Q_{Y|W,\mathbf{X}}$ and $Q_{W|\mathbf{X}}$, under certain assumptions. Such an argument will generalize better than the empirical results currently shown, and in conjunction they will both provide strong motivation for establishing our own set of metrics for synthetic data containing treatments. We make this formal argument in our updated Section 5.1, please see Theorem 1 and the related assumptions in the updated manuscript for context.

Theorem 1 shows that KL divergence loses sensitivity to $W|\mathbf{X}$ and $Y|W,\mathbf{X}$ as $d$ grows, suggesting that it will struggle in selecting between any proposal distributions which take the assumed form, since their scores will converge to similar values despite any difference in their modelling of $P_{W|\mathbf{X}}$ or $P_{Y|W,\mathbf{X}}$. Of course, the assumptions we make here will not necessarily hold for either the true distribution, or the proposal distribution, but we see this result as a nice motivational point. Furthermore, our empirical verification of this phenomena, and its apparent extension to other joint-distribution-level metrics in Appendix D, adds to our motivation for establishing a new set of metrics. Also, we see the importance of our desiderata, outlined in Section 4, as strong motivation in itself for establishing metrics to disentangle between each of these three distributions (in order to answer Q1-3 as posed in Section 1), which metrics that do not directly target these distributions (such as all-encompassing, joint-distribution-level metrics) inherently cannot do.

Since $JSD_\pi$ and $U_{PEHE}$ directly measure the preservation of $P_{W|\mathbf{X}}$ and $P_{Y|W,\mathbf{X}}$ without also trying to incorporate some notion of divergence between the covariate distributions $P_{\mathbf{X}}$ and $Q_{\mathbf{X}}$ , these do not exhibit this 'curse of dimensionality'-type behaviour, assuming the machine-learning models used in their calculation are well-specified.

Update: We have added this new formal argument to Section 5.1 on the inadequacy of existing joint-distribution-level metrics. We have moved the illustrative experiments into the Appendix as further evidence that such phenomena are observable empirically.

(D) Minor changes

We agree that noting the DP composition theorem applies to STEAM is a trivial result, and so we have now labeled it as a proposition, rather than a theorem. We include the proof in Appendix G for the benefit of the causal inference community, who may not be so familiar with the relevant DP theorems.

Thank you once again. We hope that we have addressed all your comments, and we greatly appreciate your feedback.

Thank you for your thoughtful comments and suggestions. We give answers to each of the following in turn, as well as pointing out corresponding updates to the revised manuscript:

• (A) How we determined ‘standard’ baselines
• (B) Comparison with more competitive differentially private generative models
• (C) Formal motivation for new metrics
• (D) Minor changes

As a preface, kindly allow us to emphasise that the main contribution of our work is not a new method for incorporating differential privacy (DP) into synthetic data generation. Rather, we focus on improving the potential of synthetic data for downstream causal inference tasks, by proposing both useful metrics for this use case, and a framework for useful generation that does not make the restrictive assumption that we have full knowledge of the causal graph. In doing so, we mainly focus on the non-DP setting (hence the majority of our comparisons are against non-DP tabular data generators as in Section 7.1 and 7.2), in which there is extensive literature. Nevertheless, we do believe that private generation is an important direction to follow to improve the real-world applications of synthetic data, and so we highlight that STEAM is immediately compatible with DP in Section 6.2, and demonstrate good performance in DP generation in Section 7.3. We agree, however, that with respect to DP generation, our original paper was lacking robust comparisons, and we thank you for particularly helping us improve this.

(A) How we determined ‘standard’ baselines

In our empirical analysis discussed in Section 7.1, we compare STEAM models with a broad array of tabular generative models. While we only include the best performing models in the main body of the paper, Table 11 in Appendix I displays the full set of results for each generator we compare to (CTGAN, TVAE , TabDDPM, ARF, NFlow). Our choice of these ‘standard’ models as baselines is guided by examining the tabular synthetic data generation literature, where we see that a similar selection of baseline models is made in a number of papers (e.g [1] include CTGAN, TVAE, NFlow, TabDDPM, [2] include ARF, TVAE, TabDDPM, [3] include CTGAN, TVAE, NFlow, [4] include TVAE, CTGAN), and by wanting to compare to models that span across different families of generative models (VAE-based, GAN-based, flow-based, diffusion-based, forest-based). Furthermore, as we state in Section 7, we primarily use the $`synthcity`$ package, an open source package which implements an array of generative models, to facilitate reproducibility and fair comparison of these different methods. For DP generation, see (B).

(B) Comparison with more competitive differentially private generative models

We thank you for the added literature on DP generative models, we have now added these to our extended related works discussion in Appendix B. For our empirical study in the DP setting, you are correct in saying that our comparison is less robust, and our choice of baseline is lacking. As such, we have we have run a further set of DP generation experiments, comparing STEAM with four baselines - AIM [5], GEM [6], MST [7], and RAP [8] - on the ACTG dataset used in Section 7.1. We now base our updated Section 7.3 on this experiment, rather than comparing with DP-GAN. Of these four baselines, we found that AIM produced the best data across different $\epsilon$, and its STEAM analogue was the best performing STEAM model. As such, we include these results for AIM vs. $STEAM_{AIM}$ in the main body of the updated manuscript, and include the full results using the other baseline models in Appendix I. For convenience, we will include the updated takeaway from this experiment here.

Figure 5, in Section 7.3 shows the results. $STEAM_{AIM}$ models $P_{Y|W,\mathbf{X}}$ better on all tested values of $\epsilon$, as $U_{PEHE}$ is significantly lower than for standard AIM. $P_{W|\mathbf{X}}$ is better modelled by $STEAM_{AIM}$ at small $\epsilon$, with equivalent performance between the methods at less conservative budgets. $P_{\mathbf{X}}$, on the other hand, is better preserved by standard AIM, scoring higher on $P_{\alpha, \mathbf{X}}$ and $R_{\beta, \mathbf{X}}$ at most $\epsilon$. This is likely because assigning $Q_\mathbf{X}$ one third of the budget of the standard AIM model and having it model largely the same distribution, save for the removed $W$ and $Y$, is prohibitively restrictive given the high-dimensionality of $\mathbf{X}$. As such, with uniform distribution of ($\epsilon, \delta$) across each component, we see a trade-off between $STEAM_{AIM}$ and standard AIM, where $STEAM_{AIM}$ better preserves $P_{W|\mathbf{X}}$ and $P_{Y|W,\mathbf{X}}$, while standard AIM preserves $P_{\mathbf{X}}$ better. Distributing ($\epsilon, \delta$) differently amongst $Q_\mathbf{X}$, $Q_{W|\mathbf{X}}$, and $Q_{Y|W,\mathbf{X}}$ could address this trade-off, as we discuss in Appendix M.

Thank you for your continued engagement with our work, we appreciate the effort you have made during this discussion period. We will respond to each of your points in turn.

1. Changes made to the paper

During this discussion period, we have made an effort to engage with all reviewers and their respective comments, leading to additions and revisions for our submission. We believe that this has been well within the spirit of ICLR's discussion period, and we politely disagree that our changes necessitate new reviews. There has not been substantial change to the setting of our paper, nor our proposed metrics/generation method, nor the conclusions drawn from our results. We believe that all additions and revisions we have made fall in line with what is acceptable, and standard, for a conference rebuttal period, in that they have not substantially changed the main contributions of our work. While these changes have improved our work, they are not so fundamental as to oblige new reviews. We give our reasoning now.

As we state in Section 1, the contributions of our work are the following:

1. Establishing a set of desiderata for what a 'good' synthetic dataset containing treatments involves.

2. Defining a set of metrics to measure adherence to these desiderata.

3. Proposing a generation framework that optimises for these desiderata.

4. Empirically demonstrating state-of-the-art performance in generating synthetic data containing treatments.

These were our main contributions before the discussion period began, and they still remain our main contributions. In total, the key changes made to the main body of the paper during the discussion period amass to 74 lines (with figures and equations). These include:

1. Addition of a small formal result that complements our existing empirical demonstration of a phenomenon which forms part of our motivation to propose new metrics (Section 5.1 lines 216-246). While adding credibility and generality to our previous motivational piece, this formal result does not fundamentally change our motivation for proposing our metrics. This motivation comes directly from our desiderata in Section 4, and from the desire to answer the key questions Q1-3 from Section 1, which remain unchanged.

2. Addition of an extended definition and description of an existing metric which we deploy, to improve accessibility for unfamiliar readers (Section 5.2.1 lines 258-272). Importantly, while this has ostensibly changed the definitions for $P_{\alpha, \mathbf{X}}$ and $R_{\beta, \mathbf{X}}$, in reality they remain exactly the same. We simply bring forward the full, mathematical definition from [1], which we previously withheld and referred to only by way of reference.

3. Re-running a previous experiment on a new dataset, with new baseline methods for more robust conclusions (Section 7.3 lines 486-516). Adding baseline methods and investigating performance on different datasets is an update often made during rebuttal periods, and the broad takeaway from this experiment - that STEAM is competitive in differentially-private generation - remains the same.

Beyond these, additions were made to the appendix, which chiefly serve to better contextualise our work amongst related fields, list results for extended sets of baselines, give advice for practical applications of our metrics, and suggest directions for future work. Again, none of these additions change our paper at a fundamental level, they simply better position our work and strengthen the evidence for the efficacy of our method.

[1] Alaa, Ahmed, et al. "How faithful is your synthetic data? sample-level metrics for evaluating and auditing generative models." International Conference on Machine Learning. PMLR, 2022.

2. Theorem 1

Decomposition of KL divergence via the chain rule is a well-known property, and we have no doubt that it has been used in a variety of applications in the statistics/ML literature. We have not seen a theorem, however, that elaborates on this decomposition and how it leads to KL divergence losing sensitivity to particular one-dimensional conditional distributions, such as $W|\mathbf{X}$ and $Y|W,\mathbf{X}$. If you have such an example in mind, we would happily cite it, as we do not suggest that this theorem is a core contribution of our work, nor that it is strictly necessary for motivating our proposed metrics. We will clarify this now.

We have multiple motivational points for metrics of synthetic data containing treatments, at different parts of our paper. Firstly, we intuitively suggest that a set of metrics in this setting should answer Q1-3 (Section 1), because these are important questions for downstream analysts. Next, as a more formal requirement, we suggest that a metric set should measure adherence to our desiderata, which are very related to Q1-3. Finally, in application, a metric set should be able to conduct correct model selection and rank models based on a particular desideratum.

Existing metrics do not satisfy any of these three points. Inherently, they cannot answer Q1-3 because they do not differentiate between $\mathbf{X},\ W$, and $Y$ (Table 1). Similarly, while they may loosely correlate with our desiderata, they do not directly measure the distributions of interest. As a result, in some circumstances their model selection will be poor, particularly if desideratum (ii) or (iii) is the focus. Formally, we show that this is true for KL divergence in theorem 1, and empirically we show this occurs for a wide array of metrics in Appendix D.

Naturally, the inadequacy of existing metrics acts as motivation to propose new metrics. These metrics should directly target our desiderata, be able to answer Q1-3, and conduct generally correct model selection. These are the three underlying motivations for our metric set. As we describe in Section 5.2, we propose metrics that do indeed directly target an individual desideratum, and that can answer our key questions. Furthermore, because the metrics targeting the $W|\mathbf{X}$ and $Y|W,\mathbf{X}$ distributions do not attempt to also incorporate divergence between covariate distributions, they will not lose sensitivity to these distributions as $d \rightarrow \infty$. We show that this results in better model selection than existing metrics through the empirical investigation in Appendix D.

3. Privacy contribution

Thank you for the additional references, they are useful for clarifying the context of our work, and we will cite them accordingly. We do believe it is important to note, however, that STEAM conducts a particular kind of sequential generation, which is tailored to the setting of synthetic data containing treatments, mimicking the typical DGP for this data. Sequential generation in this setting, without assuming full knowledge of the causal graph, is an important contribution of our work. Indeed, as we reference in Section 1, many related works which focus on synthetic data in medicine do not consider sequential generation, despite its benefits in preserving important relationships for tasks such as treatment effect estimation, as demonstrated in our empirical analyses. We believe that such downstream tasks are largely overlooked in the broad synthetic data literature. Bringing the data needs of the treatment effect estimation community to the attention of the synthetic data community is an important step to make.

We agree with your comment that formal reasoning on utility bounds for STEAM is an interesting direction, and would be a great contribution to make. We see this as an extension beyond the scope of this work, however. Indeed, in our contributions in Section 1, the majority are not related to DP. We do mention, in contribution 3, that STEAM satisfies DP if desired indicating that this setting is not the only focus of this work. Indeed, subsequent sections on differential privacy (Section 6.2 and Section 7.3) are relatively small (less than a page). In contrast, our other contributions, and experiments in the non-DP setting, take up substantially more space in the paper. The desiderata discussion spans 1.5 pages, the metrics span 2 pages, and our method and its performance in the non-DP setting spans 2.5 pages. In comparison, the DP aspect is small. Since privacy is inherently relevant to synthetic data, especially in medicine, we see the compatibility of STEAM with DP as useful, and its competitive performance is promising.

Thank you for your thoughtful comments and suggestions. We give answers to each of the following in turn, as well as pointing out corresponding updates to the revised manuscript:

• (A) Accessibility of metrics
• (B) Comparison with causal generative models
• (C) Guidelines for $\mathcal{F}$ when calculating $U_{PEHE}$
• (D) Extension to arbitrary causal graphs

(A) Accessibility of metrics

We hope that our metrics are used beyond our work, as we think they add important information beyond joint-distribution-level metrics when assessing synthetic data for causal inference, so we are thankful for your help in improving their accessibility. We have added reminders of the definitions of Jensen-Shannon Distance, $P_\alpha$, and $R_\beta$ prior to their use in defining our metrics.

Update: We have added definitions for Jensen-Shannon Distance, $P_\alpha$, and $R_\beta$ to Section 5.2.

(B) Comparison with causal generative models

We agree that adding a discussion and experimental comparison with causal generative models is an important addition to our paper. We will discuss the papers you have cited, and give a broader overview of causal generative models, in our related works in Appendix B to better contextualize our work. In doing so, we will highlight some important points of difference between STEAM and existing methods.

Most notably, [1], [2], and [3] assume that the entire causal graph is known to the data holder, which is a more restrictive assumption than we make with STEAM. We assume that our specification of the underlying DGP ($X \sim P_\mathbf{X},\ W \sim P_{W|\mathbf{X}},\ Y \sim P_{Y|W,\mathbf{X}}$) is correct for datasets containing treatments, but we do not require knowledge of the causal graph, as we do not need to know the causal links between individual variables. We do not assume knowledge of the causal relationships amongst the covariates, nor knowledge of which covariates cause treatment assignment or patient outcomes. As such, we make less restrictive assumptions than these works, and we suggest that our approach is more realistic in complex, real-world scenarios, such as those that arise in medicine, where the true causal graph is unlikely to be available. Note that our setting follows established literature in causal inference that does not impose any assumption on the full underlying causal graph.

Furthermore, the motivation for these works differs from ours. We seek to generate useful synthetic data only from the observational distribution, for use by analysts with goals such as treatment effect estimation (e.g., CATE), while [1], [2], and [3] primarily focus on generating data to answer graph-specific interventional and counterfactual queries that require knowledge of the full causal graph. Nevertheless, you are correct in point out that such methods can be used to generate data from the observational distribution, and we agree that experimental comparisons will strengthen the contribution of our paper.

Experiments with new baselines: In order to compare to some baseline causal generative models (specifically the diffusion-based causal model (DCM) from [3] and the additive noise model (ANM) [5] implementation in the $`DoWhy-GCM`$ python [1]) we must first reconcile their differences in assumptions. For the real-world datasets we use in Section 7.1, we do not have access to the true causal graph, we simply know which features are the treatment and outcome. As such, we must first construct a causal graph to supply to these methods. We do so with three methods:

1. Construction of a naive graph $\mathcal{G}_{naive}$, in which each covariate causes $W$ and $Y$, $W$ causes $Y$, and every pair of covariates has a causal relationship between them;
2. Using the constraint-based PC causal discovery algorithm [4] to discover an estimated graph $\mathcal{G}_{discovered}$ from the Markov Equivalence Class for the true causal graph and;
3. Pruning the discovered causal graph by removing any edges which contradict the DGP we assume. As such, any edges from $Y$ to $W$ or $\mathbf{X}$, or from $W$ to $\mathbf{X}$ are removed to form $\mathcal{G}_{pruned}$.

We compare the DCM and ANM methods using each of these causal graph discovery methods with the best performing STEAM models on each of the datasets used in Section 7.1. We include the new results in the Appendix O, and we add the table in the next comment for convenience.

For each dataset, we see that the relevant STEAM model outperforms all of the causal generative models in almost every metric, only being outperformed to a statistically significant level in $R_{\beta, X}$ on the ACIC dataset. These results validate that, when the true causal graph is not known, our less restrictive assumptions enable more useful generation of synthetic data containing treatments. We also see that the differences between the graph discovery methods are relatively small.

[1] Blöbaum, P., Götz, P., Budhathoki, K., Mastakouri, A. A., & Janzing, D. (2024). DoWhy-GCM: An extension of DoWhy for causal inference in graphical causal models. arXiv:2206.06821.

[2] Sanchez-Martin, P., Rateike, M., & Valera, I. (2021). VACA: Design of Variational Graph Autoencoders for Interventional and Counterfactual Queries. arXiv:2110.14690.

[3] Chao, P., Blöbaum, P., Patel, S., & Kasiviswanathan, S. P. (2024). Modeling Causal Mechanisms with Diffusion Models for Interventional and Counterfactual Queries. arXiv:2302.00860.

[4] Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.

[5] Patrik Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. Advances in neural information processing systems, 21, 2008.

[6] Alicia Curth and Mihaela Van Der Schaar. In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation. In International Conference on Machine Learning, pp. 6623–6642. PMLR, 2023.

[7] Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1810–1818. PMLR, 13–15 Apr 2021. URL https://proceedings.mlr.press/v130/curth21a.html.

[8] Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heteroge- neous treatment effects using machine learning. Proceedings of the national academy of sciences, 116(10):4156–4165, 2019.

[9] Edward H Kennedy et al. Optimal doubly robust estimation of heterogeneous causal effects. arXiv preprint arXiv:2004.14497, 5, 2020.

Furthermore, STEAM’s design can, in principle, incorporate any generative model for $Q_\mathbf{X}$, essentially acting as a wrapper around $Q_\mathbf{X}$ to improve its generation quality for causal inference tasks. This allows STEAM to very easily empower existing generative modelling frameworks, without having to incorporate bespoke generators. Also, it allows STEAM models to continuously improve along with the base generative model. Existing causal generative models do not allow such flexibility, and therefore cannot as easily benefit from future improvements.

Finally, the causal generative models we compared against here cannot immediately satisfy differential privacy, and any causal generative model which requires knowledge of the causal graph must ensure that the causal discovery methods used are also privacy-enabled.

Update: We have added a discussion on causal generative models to our related works section. We have added new experimental results using the causal generative models DCM and ANM as baselines to Appendix O.

(C) Guideline for $\mathcal{F}$ when calculating $U_{PEHE}$

In CATE estimation, model validation is a difficult task [6]. As such, it is reasonable to expect that a set of downstream analysts conducting CATE estimation on $D_{synth}$ will use different learners. Therefore, we want $U_{PEHE}$ to reflect the expected difference in downstream performance between $D_{synth}$ and $D_{real}$ across a diverse array of potential learners, such that it is representative for the entire population of analysts, and has limited bias towards any particular learner class. To achieve this, we propose averaging $U_{PEHE}$ across a family of CATE learners $\mathcal{F}$, and we suggest that larger $|\mathcal{F}|$, and diverse selection of the learners within $\mathcal{F}$, is preferable.

Of course, there is a trade off between the size of $\mathcal{F}$, and therefore the stability of $U_\text{PEHE}$, and the computational cost of repeated CATE estimation. With this in mind, to limit the computation involved in calculating $U_\text{PEHE}$, we suggest that users should be selective of the learners included in $\mathcal{F}$ to maximize learner diversity, and minimise $|\mathcal{F}|$. For example, in our experiments we set $|\mathcal{F}|=4$, and we chose learners from both of the high-level CATE learning strategies described in [7] (i.e. one-step plug-in learners, and two-step learners). Specifically, for the one-step learners we use S- and T- learners [8], and for the two-step learners we use RA and DR learners [9]. All four of these learners conduct CATE estimation differently, and encode different inductive biases in their approaches, and thus they form a good diverse base for $\mathcal{F}$. We see across the runs in Section 7.1 that the 95% confidence intervals we report for $U_\text{PEHE}$ using this $\mathcal{F}$ seem reasonable compared to the magnitude of the metric itself, and we suggest that such $\mathcal{F}$ is a good starting point for a user. If this user has requirements on their computational costs, or desired stability of $U_\text{PEHE}$, then $\mathcal{F}$ can be adjusted accordingly.

For our experiments, on each of the real datasets from Section 7.1, the runtime for calculating $U_{PEHE}$ for one run was 26, 60, and 191 seconds on the ACTG, IHDP, ACIC respectively. Note that these are much less than the typical generation times for each dataset, so this step is unlikely to be a large time burden for the data holder. Also note that these calculations can be parallelized across the learner classes, which we did not do, and this can improve the computational feasibility of using a larger $|\mathcal{F}|$.

Update: We have added a discussion on the considerations that should be made when choosing $\mathcal{F}$, and we have specified our experimental settings and the typical run times for calculating $U_{PEHE}$ on each dataset in Appendix F.

(D) Extension to arbitrary causal graphs

As we note in (B), STEAM places minimal assumptions upon the causal graph of $\mathcal{D}_{real}$, and it does not require full knowledge of variable-level causality. We see that our assumed DGP involving $P_{\mathbf{X}}$, $P_{W|\mathbf{X}}$, and $P_{Y|W,\mathbf{X}}$ is applicable across a very wide range of scenarios for datasets containing treatments, particularly in medical settings, as we describe in Section 6. If this DGP is known to not hold for $\mathcal{D}_{real}$, then it would be unadvisable to apply STEAM in its current form, as its inductive biases may not be helpful. However, if a similar setting arises where the overarching DGP is known, but the specific causal graph is not, then altering the generation order between $\mathbf{X}$, $W$, and $Y$ to mimic this alternative DGP, with a similar underlying motivation as in STEAM - that generation mimicking the real DGP is preferable to joint-level generation - can be done.

An example of a more complicated setting, where STEAM is not immediately applicable is when dealing with longitudinal data, which has time-varying covariates, treatments, and outcomes. In this setting, our assumed DGP would not hold in general, as features at a specific time point will likely have temporal causal relationships with earlier features, e.g. at time t, $X_{t}$ could be affected by $W_{t-1}$, which is not modelled in our assumed DGP. Potential alterations to STEAM could adjust for this time-varying DGP, where one could consider a ‘base’ STEAM model, as described in this work, operating at each time step, with additional inputs from the immediately preceding time step to model temporal causal relationships.

Update: We have included this discussion on potential extensions of STEAM to alternative DGPs in our future works Appendix P.

Thank you once again. We hope that we have addressed all your comments, and we greatly appreciate your feedback.

Table 16: $P_{\alpha, \mathbf{X}}$, $R_{\beta, \mathbf{X}}$, $JSD_\pi$, and $U_\text{PEHE}$ values for the best performing STEAM models from Appendix I in comparison to causal generative models. Averaged over 20 runs, with 95% confidence intervals. Bold indicates the statistically significant best performing model on that metric and dataset.

* $\mathcal{G}$ pruned is the same as $\mathcal{G}$ discovered for IHDP

$\dagger$ Excessive runtime caused the exclusion of $\mathcal{G}_\text{naive}$ ACIC results

We are extremely thankful for your helpful feedback and commitment to reviewing our work. We see that, in addressing your points, the manuscript has substantially improved, and we are very glad that you agree and have increased your score from a 5 to an 8!

Thank you for your very thorough comment and revision to the paper. I feel like the issues I pointed at have been largely addressed, and I don't find any major weakness left. I have changed my rating accordingly.

Thank you for your thoughtful comments and suggestions. We give answers to each of the following in turn, as well as pointing out corresponding updates to the revised manuscript:

• (A) Performance at smaller $\delta$
• (B) Distribution of $\epsilon$ in STEAM
• (C) Performance for $\epsilon < 1$
• (D) Minor changes

(A) Performance at smaller $\delta$

You are correct that, for the differential privacy (DP) experiment in Section 7.3, we set $\delta = 1/n$, with $n = 1000$. In settings where more privacy is necessary, then $\delta$ may need to be reduced. We have run new DP experiments with $\mathbf{\delta = 10^{-6}}$ and we include these results in Section 7.3. Please note as well that, in this section, we now use the real world ACTG dataset from Section 7.1, to inspect DP performance in a more realistic setting, and we compare against an updated set of more competitive baseline DP generative models. We discuss these results further in addressing point (C).

Update: We have included new results in Section 7.3 which compare STEAM with standard DP generative models when $\mathbf{\delta = 10^{-6}}$.

(B) Distribution of $\mathbf{\epsilon}$ in STEAM

In our original manuscript (Section 7.3) we explore the performance of STEAM in DP generation with uniform distribution of $\epsilon$ across $Q_\mathbf{X}$, $Q_{W|\mathbf{X}}$, and $Q_{Y|W,\mathbf{X}}$. While this is one potential mechanism to ensure DP, a more general approach is to distribute $\epsilon$ according to some preference function $f:(0,\infty) \times \triangle^2 \rightarrow \epsilon \cdot \triangle^2$ (where $\triangle^2$ is the 2-simplex) which takes input of the budget $\epsilon$ and weights $\mathbf{w}$ for the relative importance of good modelling in $Q_\mathbf{X}$, $Q_{W|\mathbf{X}}$, and $Q_{Y|W,\mathbf{X}}$, and outputs a corresponding $\epsilon$ distribution. For example, a simple preference function definition would be $f(\epsilon, \mathbf{w}) = \epsilon \cdot \mathbf{w}$ where $\mathbf{w}$ could be defined by a data holder with some prior knowledge of the importance level of each component distribution to downstream analysts. Another approach, if it is not necessary to specify the desired $\epsilon$ distribution a priori, is to treat it as a hyperparameter, to be tuned over a series of runs to optimize some metric, such as a combination of $U_{PEHE}$, $P_{\alpha, X}$, $R_{\beta, X}$, and $\text{JSD}_{\pi}$.

Update: We have extended our discussion on the above alternatives to uniform distribution of the privacy budget amongst the STEAM component models in Appendix M.

(C) Performance for $\epsilon < 1$

In Section 7.3, in the new experiment on the ACTG dataset, we added new results for $\mathbf{\epsilon \in [0.25, 0.5, 1, 2, 3, 5, 10, 15]}$ against an updated set of DP baselines. We display results for the AIM [1] DP generative model and $\text{STEAM}_{AIM}$ in the main body of the paper, with extended results using GEM [2], MST [3], and RAP [4] DP generative models and their STEAM analogues in Appendix I.

We generally see that $STEAM_{AIM}$ performs comparatively better in preserving $P_{W|\mathbf{X}}$, and $P_{Y|W,\mathbf{X}}$, and comparatively worse for $P_{\mathbf{X}}$ at the same privacy budget as AIM. We do not suggest that the comparatively worse modelling of $P_{\mathbf{X}}$ is because of any particular hyperparameter settings in the DP methods used within STEAM. Rather, it is a natural consequence of the inverse relationship between $\epsilon$ and the noise required to satisfy DP. Since we allocate $\epsilon$ uniformly across $Q_{\mathbf{X}}$, $Q_{W|\mathbf{X}}$ and $Q_{Y|W,\mathbf{X}}$ in this experiment, each STEAM component model only receives one third of the budget compared to the baseline AIM model, and this seems to become prohibitively restrictive for modelling $P_{\mathbf{X}}$. Intuitively, this makes sense, since $P_{\mathbf{X}}$ is likely the hardest distribution to model given its high dimensionality, and STEAM’s $Q_{\mathbf{X}}$ and the standard AIM method are modelling it almost the same way, save for the removed $W$ and $Y$ variables from STEAM, yet $Q_{\mathbf{X}}$ receives one third of the budget. To mitigate this effect, since STEAM outperforms in the $U_{PEHE}$ and $JSD_{\pi}$ metric, and the performance of STEAM’s $Q_{Y|W,\mathbf{X}}$ and $Q_{W|\mathbf{X}}$ do not appear to be largely affected by reductions in $\epsilon$, one could redistribute some of the budget assigned to $Q_{Y|W,\mathbf{X}}$ and $Q_{W|\mathbf{X}}$ to $Q_{\mathbf{X}}$, for a more even performance across all three distributions.

Update: We have included new results in Section 7.3 which showcase DP generation for $\epsilon < 1$, and we have updated the takeaways to include discussion on why $STEAM_{AIM}$ preserves $Q_{\mathbf{X}}$ worse than AIM.

(D) Minor changes

Thank you for pointing out the mistake on line 225, your interpretation is indeed correct. We have now updated this to say ‘as $\mathbf{X}$ grows in dimensionality’.

Thank you once again. We hope that we have addressed all your comments, and we greatly appreciate your feedback.

References

[1] Ryan McKenna, Brett Mullins, Daniel Sheldon, and Gerome Miklau. Aim: An adaptive and iterative mechanism for differentially private synthetic data, 2024. URL https://arxiv.org/abs/2201.12677.

[2] Terrance Liu, Giuseppe Vietri, and Steven Z Wu. Iterative methods for private synthetic data: Unifying framework and new methods. Advances in Neural Information Processing Systems, 34: 690–702, 2021.

[3] Ryan McKenna, Gerome Miklau, and Daniel Sheldon. Winning the nist contest: A scalable and general approach to differentially private synthetic data. arXiv preprint arXiv:2108.04978, 2021.

[4] Sergul Aydore, William Brown, Michael Kearns, Krishnaram Kenthapadi, Luca Melis, Aaron Roth, and Ankit A Siva. Differentially private query release through adaptive projection. In International Conference on Machine Learning, pp. 457–467. PMLR, 2021.

References

[1] Patrick Chao, Patrick Blöbaum, Sapan Patel, and Shiva Prasad Kasiviswanathan. Modeling causal mechanisms with diffusion models for interventional and counterfactual queries, 2024. URL https://arxiv.org/abs/2302.00860.

[2] Patrick Blöbaum, Peter Götz, Kailash Budhathoki, Atalanti A Mastakouri, and Dominik Janzing. Dowhy-gcm: An extension of dowhy for causal inference in graphical causal models. Journal of Machine Learning Research, 25(147):1–7, 2024.

[3] Patrik Hoyer, Dominik Janzing, Joris M Mooij, Jonas Peters, and Bernhard Schölkopf. Nonlinear causal discovery with additive noise models. Advances in neural information processing systems, 21, 2008.

[4] Peter Spirtes, Clark Glymour, and Richard Scheines. Causation, prediction, and search. MIT press, 2001.

[5] Ahmed Alaa, Boris Van Breugel, Evgeny S. Saveliev, and Mihaela van der Schaar. How faithful is your synthetic data? Sample-level metrics for evaluating and auditing generative models. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 290–306. PMLR, 17–23 Jul 2022. URL https://proceedings.mlr.press/v162/alaa22a.html.

[6] Van Breugel, Boris, et al. "Membership inference attacks against synthetic data through overfitting detection." arXiv preprint arXiv:2302.12580 (2023).

[7] Ryan McKenna, Brett Mullins, Daniel Sheldon, and Gerome Miklau. Aim: An adaptive and iterative mechanism for differentially private synthetic data, 2024. URL https://arxiv.org/abs/2201.12677.

[8] Terrance Liu, Giuseppe Vietri, and Steven Z Wu. Iterative methods for private synthetic data: Unifying framework and new methods. Advances in Neural Information Processing Systems, 34: 690–702, 2021.

[9] Ryan McKenna, Gerome Miklau, and Daniel Sheldon. Winning the nist contest: A scalable and general approach to differentially private synthetic data. arXiv preprint arXiv:2108.04978, 2021.

[10] Sergul Aydore, William Brown, Michael Kearns, Krishnaram Kenthapadi, Luca Melis, Aaron Roth, and Ankit A Siva. Differentially private query release through adaptive projection. In International Conference on Machine Learning, pp. 457–467. PMLR, 2021.

Thank you for your thoughtful comments and suggestions. We give answers to each of the following in turn, as well as pointing out corresponding updates to the revised manuscript:

• (A) Comparison with causal generative models
• (B) Discussion on precision and recall
• (C) Insights on the interaction between differential privacy and treatment effect estimation
• (D) Added differential privacy experiment on real data
• (E) Minor updates

(A) Comparison with causal generative models

You are correct that existing works consider the generation of synthetic data from a known causal graph. Causal generative models (e.g. [1], [2]) typically try to approximate the underlying structural causal model (SCM) of a dataset, to answer observational, interventional, and counterfactual queries. We agree that adding a discussion and experimental comparison with causal generative models is an important addition to our paper, and there are important differences to clarify in the assumptions and motivation between such methods and ours. Most notably, causal generative models typically assume that the entire causal graph is known to the data holder, which is a more restrictive assumption than we make with STEAM. We assume that our specification of the underlying DGP ($X \sim P_\mathbf{X},\ W \sim P_{W|\mathbf{X}},\ Y \sim P_{Y|W,\mathbf{X}}$) is correct for datasets containing treatments, but we do not require knowledge of the causal graph, as we do not need to know the causal links between individual variables. We do not assume knowledge of the causal relationships amongst the covariates, nor knowledge of which covariates cause treatment assignment or patient outcomes. As such, we make less restrictive assumptions than works which require knowledge of a causal graph, and we suggest that our approach is more realistic in complex, real-world scenarios, such as those that arise in medicine, where the true causal graph is unlikely to be available. Note that our setting follows established literature in causal inference that does not impose any assumption on the full underlying causal graph.

Furthermore, the motivation for these works differs from ours. We seek to generate useful synthetic data only from the observational distribution, for use by analysts with goals such as treatment effect estimation (e.g., CATE), while [1], [2], and other causal generative models primarily focus on generating data to answer graph-specific interventional and counterfactual queries that require knowledge of the full causal graph. Nevertheless, these methods can be used to generate data from the observational distribution, and we agree that experimental comparisons will strengthen the contribution of our paper.

Experiments with new baselines: In order to compare to some baseline causal generative models (specifically the diffusion-based causal model (DCM) from [1] and the additive noise model (ANM) [3] implementation in the $`DoWhy-GCM`$ python [2]) we must first reconcile their differences in assumptions. For the real-world datasets we use in Section 7.1, we do not have access to the true causal graph, we simply know which features are the treatment and outcome. As such, we must first construct a causal graph to supply to these methods. We do so with three methods:

1. Construction of a naive graph $\mathcal{G}_{naive}$, in which each covariate causes $W$ and $Y$, $W$ causes $Y$, and every pair of covariates has a causal relationship between them;
2. Using the constraint-based PC causal discovery algorithm [4] to discover an estimated graph $\mathcal{G}_{discovered}$ from the Markov Equivalence Class for the true causal graph and;
3. Pruning the discovered causal graph by removing any edges which contradict the DGP we assume. As such, any edges from $Y$ to $W$ or $\mathbf{X}$, or from $W$ to $\mathbf{X}$ are removed to form $\mathcal{G}_{pruned}$.

We compare the DCM and ANM methods using each of these causal graph discovery methods with the best performing STEAM models on each of the datasets used in Section 7.1. We include the new results in the Appendix O, and we add the table in the next comment for convenience.

For each dataset, we see that the relevant STEAM model outperforms all instantiations of the causal generative models in almost every metric, only being outperformed to a statistically significant level in $R_{\beta, X}$ on the ACIC dataset. These results validate that, when the true causal graph is not known, our less restrictive assumptions enable more useful generation of synthetic data containing treatments. We also see that the differences between the graph discovery methods are relatively small.

Table 16: $P_{\alpha, \mathbf{X}}$, $R_{\beta, \mathbf{X}}$, $JSD_\pi$, and $U_\text{PEHE}$ values for the best performing STEAM models from Appendix I in comparison to causal generative models. Averaged over 20 runs, with 95% confidence intervals. Bold indicates the statistically significant best performing model on that metric and dataset.

* $\mathcal{G}$ pruned is the same as $\mathcal{G}$ discovered for IHDP

$\dagger$ Excessive runtime caused the exclusion of $\mathcal{G}_\text{naive}$ ACIC results

Furthermore, STEAM’s design can, in principle, incorporate any generative model for $Q_\mathbf{X}$, essentially acting as a wrapper around $Q_\mathbf{X}$ to improve its generation quality for causal inference tasks. This allows STEAM to very easily empower existing generative modelling frameworks, without having to incorporate bespoke generators. Also, it allows STEAM models to continuously improve along with the base generative model. Existing causal generative models do not allow such flexibility, and therefore cannot as easily benefit from future improvements.

Finally, the causal generative models we compared against here cannot immediately satisfy differential privacy, and any causal generative model which requires knowledge of the causal graph must ensure that the causal discovery methods used are also privacy-enabled.

Update: We have added a discussion on causal generative models to our related works section. We have added new experimental results using the causal generative models DCM and ANM as baselines to Appendix O.

(B) Discussion on precision and recall

The alpha-precision and beta-recall metrics we use in Section 5.2.1 were introduced in [5] as metrics for evaluating the quality of synthetic data. They are extensions of the conventional precision/recall metrics in the supervised learning setting to the synthetic data generation setting. Intuitively, alpha-precision captures how much of the synthetic data falls within the support of real data, whereas beta-recall reflects how much of the real data is covered by the support of synthetic data. Together, they capture the faithfulness and diversity of the synthetic data. Since their introduction in 2022, they have been used in variety of publications (197 citations on Google Scholar), hence our description of their use as 'widespread'.

Update: We have updated our descriptions of alpha-precision and beta-recall in Section 5, and changed equation (3) and (4) to incorporate the full definitions from [5].

(C) Insights on the interaction between differential privacy and treatment effect estimation

Our consideration of differential privacy (DP) stems from our desire to facilitate data release. In this work, we consider the application of releasing synthetic data for causal inference tasks, such as treatment effect estimation, which means that all of the covariates, the treatment assignment and the response must be released, and are subject to potential privacy attacks. Only ensuring DP for a subset of variables (such as only the covariates) violates the DP of the entire dataset, as DP for a subset of variables does not lead to DP for all variables. Therefore, pursuing this would remove any guarantees on privacy for individuals whose data are being released. This could be ok on a case-by-case basis, but we hope for our method to be as broadly applicable as possible, and so seek to adhere to DP across the entire dataset.

Furthermore, to your question on whether a data holder can release ‘a realistically differentially private dataset where causal effects are still discoverable’, DP is a formal and provable way to establish data privacy - however it is not directly associated with one particular style of privacy attack (such as membership inference or attribute inference). In fact, a criticism of DP is that its practical implication is very often unclear [6]. For instance, DP provides little insight into how many individuals in the data are vulnerable to particular re-identification attacks. While increasing the degree of DP will generally reduce the success of all privacy attacks, it does not establish provable guarantees for any particular methodology. This has triggered new lines of research into privacy enhancement for specific attack modality, such as in [6], which is fundamentally different from DP privacy. Since we use DP in particular in the work, we do not focus on any particular attack method.

Nevertheless, if we assume that a data holder wants to release a DP synthetic dataset for CATE estimation, but they have some threshold on the degree of added error that analysts will encounter using this dataset for this task, we can show that STEAM facilitates the release of more private datasets than standard DP generation. We can see this is true simply by examining the DP generation results in Section 7.3, and the full results on the extended set of baselines in Appendix I. Comparing each STEAM model with its standard analogue, for any threshold a data holder may set on an acceptable $U_{PEHE}$ score (which measures the difference in the CATEs found from the synthetic vs. real data) we see that the STEAM model always allows a synthetic dataset with smaller $\epsilon$ (and therefore greater degree of privacy) before hitting this set threshold.

(D) Added differential privacy experiment on real data

We have run a further set of DP generation experiments, comparing STEAM with four extra baselines - AIM [7], GEM [8], MST [9], and RAP [10] - on the ACTG dataset used in Section 7.1, for a more robust DP comparison on more realistic data. We now base our updated Section 7.3 using this experiment. Of these four baselines, we found that AIM produced the best data across different epsilon, and its STEAM analogue was the best performing STEAM model. As such, we include these results for AIM vs. $STEAM_{AIM}$ in the main body of the updated manuscript, and include the full results with the other baseline models in the Appendix I

We will now discuss our updated takeaways from this DP generation experiment. Figure 5, in Section 7.3 shows the results. $STEAM_{AIM}$ models $P_{Y|W,\mathbf{X}}$ better on all tested values of $\epsilon$, as $U_{PEHE}$ is significantly lower than for standard AIM. $P_{W|\mathbf{X}}$ is better modelled by $STEAM_{AIM}$ at small $\epsilon$, with equivalent performance between the methods at less conservative budgets. $P_{\mathbf{X}}$, on the other hand, is better preserved by standard AIM, scoring higher on $P_{\alpha, \mathbf{X}}$ and $R_{\beta, \mathbf{X}}$ at most $\epsilon$. This is likely because assigning $Q_\mathbf{X}$ one third of the budget of the standard AIM model and having it model largely the same distribution, save for the removed $W$ and $Y$, is prohibitively restrictive given the high-dimensionality of $\mathbf{X}$. As such, with uniform distribution of ($\epsilon, \delta$) across each component, we see a trade-off between $STEAM_{AIM}$ and standard AIM, where $STEAM_{AIM}$ better preserves $P_{W|\mathbf{X}}$ and $P_{Y|W,\mathbf{X}}$, while standard AIM preserves $P_{\mathbf{X}}$ better. Distributing ($\epsilon, \delta$) differently amongst $Q_\mathbf{X}$, $Q_{W|\mathbf{X}}$, and $Q_{Y|W,\mathbf{X}}$ could address this trade-off, as we discuss in Appendix M.

Update: We have changed Section 7.3 so that we compare DP generation on the real world ACTG dataset, between AIM and $STEAM_{AIM}$. We have added more results for other DP baselines tested (GEM, MST, RAP) in Appendix I.

(E) Minor updates

Thank you for helping us improve the writing and flow of our paper. We summarise some such changes we have made that stem from your review here:

• Described our motivation for STEAM earlier, moving a sentiment similar to the sentence ‘Mimicry of the real DGP acts as an inductive bias, pushing Q closer towards the P in structure, and directly targeting each distributions from our desiderata’ in Section 6, into our contributions in Section 1
• Added citation for T-learner in Section 7
• Defined acronym JSD in Section 5
• Clearly explained that Section 7.1 is indeed focussing on the non-DP setting
• Made Section 4 more conclusive by adding pointers to later parts in the paper, where we show the poor performance of both existing metrics and generative models in respecting our desiderata.

Thank you once again. We hope that we have addressed all your comments, and we greatly appreciate your feedback.