Diffusion Models for Realistic Tabular Simulation

Abstract

Monte Carlo studies in econometrics are often analytically transparent but empirically thin: functional forms, covariate distributions, error distributions, and treatment assignment rules are chosen for tractability rather than for resemblance to the empirical settings where estimators fail. Athey, Imbens, Metzger, and Munro (2024) propose using Wasserstein GANs to learn data-generating processes from real datasets and then use the fitted generator as a simulation design (Athey et al. 2024). This paper develops the same idea with denoising diffusion probabilistic models. The main claim is not that diffusion uniformly dominates GANs. The claim is that finite-sample econometric benchmarking should target the empirical geometry that drives estimator behavior–overlap, leverage, tails, treatment imbalance, and dependence–and that diffusion models are a promising engine for learning and perturbing that geometry because they replace adversarial fitting with supervised denoising. The cost is slower iterative sampling. The evidence below should be read as a runnable practice memo rather than a definitive benchmark.

Introduction

Suppose an applied study observes rows

\[ R_i = (Y_i, W_i, X_i), \qquad i=1,\ldots,n, \]

where \(Y_i\) is an outcome, \(W_i \in \{0,1\}\) is a binary treatment or policy exposure, and \(X_i\) collects pre-treatment covariates. A traditional Monte Carlo design specifies a distribution \(P_0\) for \(R_i\) and then evaluates estimators under repeated draws from \(P_0\). The credibility of the exercise rests on whether \(P_0\) captures the parts of the empirical problem that matter for the estimators being compared: overlap, tail behavior, discreteness, collinearity, heteroskedasticity, treatment imbalance, and nonlinear dependence.

The AIMM paper reframes the problem as a learned simulation design. Let \(P_n\) be the empirical distribution of a real dataset. Fit a generator \(\hat G\) so that samples from \(\hat P_G\) look like draws from \(P_n\). Then use \(\hat P_G\) as the baseline simulation environment. This does not estimate structural primitives; it reduces arbitrary choices in simulation design while preserving the empirical geometry of the application.

This note separates two claims that are easy to conflate. The first is a generative-model comparison: WGANs and DDPMs are alternative ways to approximate a tabular row distribution. The second is a simulation-design claim: the object of an econometric Monte Carlo should be a realistic empirical environment plus controlled perturbations, not only a convenient analytic DGP. The second claim is the more important one. DDPMs matter here only insofar as they make that simulation-design workflow more stable and modular.

Learned Simulation Designs

A GAN has two networks. A generator maps noise \(Z\) into synthetic rows:

\[ \tilde R = G_\theta(Z), \qquad Z \sim P_Z. \]

A discriminator or critic scores whether rows look real. In a Wasserstein GAN, the population objective is

\[ \min_\theta \max_{f \in \mathcal{F}_1} \left\{ \mathbb{E}_{X \sim P_n}[f(X)] - \mathbb{E}_{Z \sim P_Z}[f(G_\theta(Z))] \right\}, \]

where \(\mathcal{F}_1\) is the class of 1-Lipschitz functions (Arjovsky, Chintala, and Bottou 2017). In practice, the Lipschitz constraint is enforced by weight clipping or a gradient penalty (Gulrajani et al. 2017).

For econometric simulation this has several attractions:

• the generator is fast once trained;
• it represents complicated row distributions without a parametric likelihood;
• conditional generation can be implemented by feeding treatment, cohort, or market labels into the generator;
• the Wasserstein critic gives a training signal with a distributional interpretation.

But the adversarial game is also the main cost. The generator improves only through the critic’s current gradients. If the critic is weak, the generator learns the wrong directions. If the critic is too strong or poorly regularized, training can become unstable. Mode collapse is a particularly bad failure for simulation studies because the generated sample can look locally plausible while missing subpopulations that determine estimator risk.

The optimizer choice matters because the problem is a saddle-point problem, not ordinary risk minimization. Simultaneous Adam is convenient, but the game can rotate around equilibria rather than converge. Optimistic Adam, mirror descent, and other game-aware updates are natural modifications because they use information about recent gradients to damp these rotations. That is a repair to the WGAN fitting problem, not a change in the underlying object being learned.

Denoising Diffusion Models

A denoising diffusion probabilistic model starts from the opposite direction. Instead of asking a generator to fool a critic, it defines a known corruption process that gradually turns data into noise. The model then learns to reverse that corruption.

Let \(x_0 \in \mathbb{R}^d\) denote a standardized data row. Choose a noise schedule

\[ 0 < \beta_1 < \cdots < \beta_T < 1, \]

and define

\[ \alpha_t = 1-\beta_t, \qquad \bar\alpha_t = \prod_{s=1}^t \alpha_s. \]

The forward corruption process is

\[ q(x_t \mid x_{t-1}) = \mathcal{N}\left( \sqrt{\alpha_t}x_{t-1}, \beta_t I \right). \]

By Gaussian algebra, one can sample \(x_t\) directly from \(x_0\):

\[ x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\epsilon, \qquad \epsilon \sim \mathcal{N}(0,I). \]

The model is a denoiser \(\epsilon_\theta(x_t,t,c)\), where \(c\) is optional context. In tabular simulations, \(c\) might be treatment status, market, cohort, site, panel time, or a policy environment. The simple DDPM training objective is

\[ \min_\theta \mathbb{E}_{x_0,t,\epsilon} \left[ \left\| \epsilon - \epsilon_\theta(x_t,t,c) \right\|_2^2 \right]. \]

This is ordinary supervised learning. The label is the known noise used to corrupt the clean row. Ho, Jain, and Abbeel show how this objective arises from a variational bound and connects to denoising score matching (Ho, Jain, and Abbeel 2020).

Sampling reverses the chain. Start with

\[ x_T \sim \mathcal{N}(0,I), \]

and recursively apply

\[ \mu_\theta(x_t,t,c) = \frac{1}{\sqrt{\alpha_t}} \left( x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}} \epsilon_\theta(x_t,t,c) \right). \]

Then draw

\[ x_{t-1} = \mu_\theta(x_t,t,c) + \sigma_t z_t, \qquad z_t \sim \mathcal{N}(0,I), \]

with \(\sigma_t\) set by the chosen reverse variance. After reaching \(x_0\), invert the tabular preprocessing map and apply any column constraints: binary rounding, nonnegativity, bounded support, or category decoding.

Density Modeling Interpretation

The DDPM is a density model built from many noisy versions of the data. The forward process defines a family of smoothed densities \(p_t(x)\), where \(p_0\) is the data distribution and \(p_T\) is close to a standard Gaussian. The denoiser learns a vector field that points noisy observations back toward regions of high probability under \(p_t\).

The score of a density is

\[ s_t(x) = \nabla_x \log p_t(x). \]

For the Gaussian corruption step,

\[ \nabla_{x_t}\log q(x_t \mid x_0) = - \frac{x_t-\sqrt{\bar\alpha_t}x_0}{1-\bar\alpha_t} = - \frac{\epsilon}{\sqrt{1-\bar\alpha_t}}. \]

Therefore a noise predictor implies a score estimate:

\[ s_\theta(x_t,t,c) \approx - \frac{\epsilon_\theta(x_t,t,c)} {\sqrt{1-\bar\alpha_t}}. \]

This is the link to score-based generative modeling. Song, Sohl-Dickstein, Kingma, Kumar, Ermon, and Poole formulate the same idea in continuous time: add noise through a stochastic differential equation and sample backward using the estimated score field (Y. Song et al. 2021). DDPMs are one discrete-time implementation of that broader score-based idea.

The econometric analogy is useful. A likelihood model specifies \(p_\theta(x)\) and maximizes \(\sum_i \log p_\theta(X_i)\). A score-based model learns the local derivative \(\nabla_x \log p_t(x)\) along many noisy versions of the distribution. It does not need a normalized likelihood. For simulation, this is enough: if the learned reverse chain approximately transports Gaussian noise into the empirical row distribution, then it supplies draws for a Monte Carlo design.

Tabular Rows

The original DDPM success story is image generation, but the row-distribution problem is different. In image data, all coordinates are intensities with spatial structure. In an econometric row, coordinates may be continuous, binary, categorical, censored, top-coded, or logically constrained.

A practical tabular DDPM therefore needs a representation map

\[ z = S(x), \]

where \(S\) standardizes continuous variables, encodes categorical variables, and records constraints needed to invert the transformation. Continuous variables can be diffused with Gaussian noise. Binary and categorical variables can be handled in several ways:

• keep them as relaxed numeric variables during diffusion and round after sampling;
• use separate output heads and postprocessing rules;
• use discrete diffusion kernels designed for categorical state spaces.

The last option is conceptually cleaner. Multinomial diffusion and argmax flows explicitly model categorical variables (Hoogeboom et al. 2021). Structured denoising diffusion models in discrete state spaces extend the corruption process beyond uniform replacement kernels and provide a general D3PM framework (Austin et al. 2021). TabDDPM combines Gaussian diffusion for numerical features with multinomial diffusion for categorical features and shows that diffusion can be competitive for general tabular data (Kotelnikov et al. 2023).

For econometric simulation, the choice should be governed by the role of each variable. Rounding a binary demographic indicator may be acceptable in a first benchmark. Rounding the binary treatment \(W\) is usually not acceptable if the estimand conditions on the realized assignment mechanism. In that case, treatment should be fixed as context or generated by a separate assignment model.

Comparison with WGANs and Flows

The AIMM objective is to build realistic simulation designs from empirical datasets. Both WGANs and DDPMs are plausible tools. They differ mainly in what is hard.

Feature	WGAN simulation	DDPM simulation
Training problem	Minimax game between generator and critic	Supervised denoising objective
Main tuning risk	Critic strength, Lipschitz penalty, optimizer dynamics, mode collapse	Noise schedule, denoiser capacity, sampling steps, preprocessing
Sampling speed	Fast single generator pass	Slower iterative denoising
Diagnostics	Critic loss plus downstream distribution checks	Denoising loss plus downstream distribution checks
Mode coverage	Can be fragile when generator collapses	Usually strong because training sees noisy neighborhoods at all scales
Conditional designs	Feed context to generator and critic	Feed context to denoiser at every reverse step
Ablations	Possible but can destabilize the adversarial game	Natural: intervene on context, corruption path, or reverse constraints

This comparison explains why DDPMs are attractive for econometric practice. The fitting problem is closer to nonparametric regression than to adversarial equilibrium computation. Every training example supplies a known target \(\epsilon\). The model is trained across many local perturbations of each row, which tends to preserve support and neighborhood structure. For a simulation study where the cost of a missed subpopulation can be high, that stability is valuable.

The tradeoff is sampling cost. A WGAN can generate \(100{,}000\) rows with one forward pass per row. A DDPM may need tens or hundreds of denoising steps. Improved DDPMs learn reverse variances and improve the likelihood/sample-quality tradeoff (Nichol and Dhariwal 2021). DDIMs reduce the number of sampling steps by using a deterministic or partially deterministic non-Markovian sampler with the same training objective (J. Song, Meng, and Ermon 2021). In practice, a tabular simulation design can usually tolerate slower sampling because generating the synthetic benchmark is not the bottleneck relative to fitting many downstream estimators.

Normalizing Flows

Normalizing flows are another natural density-modeling candidate. A flow constructs an invertible map \(f_\theta\) such that

\[ z = f_\theta(x), \qquad z \sim p_Z, \]

and evaluates the likelihood by the change-of-variables formula:

\[ \log p_\theta(x) = \log p_Z(f_\theta(x)) + \log \left| \det \frac{\partial f_\theta(x)}{\partial x} \right|. \]

This gives exact likelihoods and direct sampling, which is attractive for statistical work (Papamakarios et al. 2021). The cost is architectural: \(f_\theta\) must be invertible and have a tractable Jacobian determinant. Those constraints can make flows awkward for mixed discrete-continuous tabular rows, although many useful constructions exist.

Diffusion models relax invertibility. They do not give exact likelihood as directly as a flow, although likelihood bounds and probability-flow ODEs are available in the broader score-based framework. For simulation design, exact likelihood is less important than faithful reproduction of the empirical row distribution and stable conditional sampling. That tilts the practical comparison toward DDPMs unless likelihood evaluation itself is part of the research design.

Trex API Surface

Trex exposes WGAN and DDPM simulators through the same minimal interface: fit on a matrix of transformed tabular rows, optionally condition on a context matrix, and sample synthetic transformed rows. The fitted rows are transformed back to data scale with TabularTransformer.

For the WGAN benchmark, the class is trex.TabularWGAN:

TabularWGAN(
    hidden_dims=(128, 128, 128),
    critic_hidden_dims=None,
    noise_dim=None,
    batch_size=128,
    max_steps=1000,
    critic_steps=5,
    lr=1e-4,
    optimizer="adam",        # or "optimistic_adam"
    gp_weight=5.0,
    generator_dropout=0.1,
    critic_dropout=0.0,
    binary_dims=(),
    lower_bounds=None,
    upper_bounds=None,
    seed=None,
    device=None,
)

wgan.fit(X, context=None)
wgan.partial_fit(X, context=None, steps=100)
fake_z = wgan.sample(n, context=None)

For the DDPM benchmark, the class is trex.TabularDiffusion:

TabularDiffusion(
    hidden_dims=(128, 128, 128),
    time_dim=32,
    n_timesteps=100,
    beta_start=1e-4,
    beta_end=0.02,
    batch_size=128,
    max_steps=1000,
    lr=1e-3,
    dropout=0.0,
    seed=None,
    device=None,
)

ddpm.fit(X, context=None)
ddpm.partial_fit(X, context=None, steps=100)
fake_z = ddpm.sample(n, context=None)

A typical tabular call is:

transformer = TabularTransformer(
    column_names=columns,
    binary_columns=binary_columns,
    nonnegative_columns=nonnegative_columns,
).fit(df[columns])

X = transformer.transform(df[columns])
context = df[["t"]].to_numpy()

ddpm = TabularDiffusion(max_steps=1000, seed=123, device="cuda")
ddpm.fit(X, context=context)

fake = transformer.inverse_transform(
    ddpm.sample(len(df), context=context).numpy()
)

This branch does not currently expose a TabularFlow class. A normalizing-flow simulator would fit the same interface naturally:

flow.fit(X, context=None)
fake_z = flow.sample(n, context=None)
logp = flow.log_prob(X, context=None)

The extra method is log_prob, since exact or tractable likelihood evaluation is the main reason to add a flow baseline. The implementation choice would then be a conditional spline flow, masked autoregressive flow, or coupling flow depending on whether tabular likelihood or fast sampling is the priority.

Distance Metrics

Let \(R_1,\ldots,R_n \in \mathbb{R}^d\) denote observed rows and \(\tilde R_1,\ldots,\tilde R_m \in \mathbb{R}^d\) denote synthetic rows. All diagnostics below are computed after the same preprocessing convention used for the simulation exercise; in the Lalonde application, earnings are measured in thousands of dollars. Lower values indicate closer agreement.

For coordinate \(j\), the one-dimensional Wasserstein distance is

\[ W_{1j} = \int_0^1 \left| \hat F_{j}^{-1}(u) - \hat G_{j}^{-1}(u) \right|du, \]

where \(\hat F_j\) and \(\hat G_j\) are the empirical CDFs of observed and synthetic column \(j\). The reported marginal_w1_mean and marginal_w1_max are

\[ \frac{1}{d}\sum_{j=1}^d W_{1j}, \qquad \max_{1 \leq j \leq d} W_{1j}. \]

The Kolmogorov-Smirnov distance for coordinate \(j\) is

\[ K_j = \sup_x \left| \hat F_j(x) - \hat G_j(x) \right|. \]

The reported marginal_ks_mean and marginal_ks_max average and maximize \(K_j\) across columns. The KS distance is scale-free, while Wasserstein distance preserves the units of each variable.

The mean-distance diagnostic is

\[ \text{mean\_l2} = \left\| \bar R - \bar{\tilde R} \right\|_2. \]

It is a location check. It can be small even when tails or dependence are wrong.

Let \(\hat\Sigma_R\) and \(\hat\Sigma_{\tilde R}\) be empirical covariance matrices. The covariance diagnostic is

\[ \text{cov\_frobenius} = \left\| \hat\Sigma_R - \hat\Sigma_{\tilde R} \right\|_F. \]

Because covariance is measured in original units, high-variance earnings columns can dominate it. The correlation diagnostic repeats the same calculation after converting covariance matrices into correlation matrices:

\[ \text{corr\_frobenius} = \left\| \hat C_R - \hat C_{\tilde R} \right\|_F. \]

This is the cleaner summary of linear dependence across mixed-scale columns.

Finally, sliced Wasserstein averages one-dimensional Wasserstein distances over random projections. For a unit vector \(u \in \mathbb{S}^{d-1}\), project rows to \(u'R_i\) and \(u'\tilde R_i\). The population target is

\[ \text{SW}_1(R,\tilde R) = \mathbb{E}_{u \sim \operatorname{Unif}(\mathbb{S}^{d-1})} \left[ W_1(u'R, u'\tilde R) \right]. \]

The reported value approximates this expectation with random directions. This is useful in tabular simulation because it can detect joint-distribution failures that are invisible in marginal diagnostics, while remaining much cheaper than exact high-dimensional optimal transport.

Metrics for Estimator Geometry

These distances are generic generator diagnostics. They are not, by themselves, econometric loss functions. For estimator benchmarking, the question is which parts of the empirical distribution make estimators fail.

For treatment-effect estimators, several features matter:

• overlap and common support in \(X\) across treatment states;
• leverage points that dominate regression or balancing weights;
• tail behavior in outcomes and pre-treatment earnings;
• treatment imbalance and small treated or control cells;
• dependence between treatment, covariates, and outcomes;
• nonlinear nuisance-function difficulty in \(m_w(x)=\mathbb{E}[Y\mid W=w,X=x]\) and \(e(x)=\Pr(W=1\mid X=x)\).

The reported metrics proxy these features imperfectly. Marginal Wasserstein and KS distances detect column-level failures such as missed tails, mass at zero, or wrong binary shares. Mean and covariance distances catch first- and second-moment errors that often affect regression adjustment. Correlation distance summarizes linear dependence across columns. Sliced Wasserstein is the broadest row-geometry diagnostic because it sees random linear combinations of variables and can detect joint-distribution errors missed by marginal checks.

The metrics do not directly measure overlap, propensity-score calibration, leverage, or nuisance-function difficulty. A complete simulation-design benchmark should add estimator-specific diagnostics: propensity-score distributions, nearest-neighbor distances across treatment arms, maximum leverage, effective sample size for weighting estimators, and tail-sensitive outcome summaries. The present memo uses generic distances as a first screen and downstream ATT replication as one applied estimand check.

Simulation Protocol

A useful DDPM simulator for econometrics should make the target distribution explicit. The following workflow is close to the AIMM proposal but swaps the fitting engine:

1. Choose the empirical target: unconditional rows, rows conditional on treatment, rows within market, or a panel transition distribution.
2. Define a preprocessing map \(S\) and record all inverse transformations.
3. Decide which variables are generated and which are fixed as context.
4. Fit a denoiser with a transparent noise schedule and held-out diagnostics.
5. Generate synthetic samples with the same sample size, treatment shares, panel length, or market composition as the empirical design.
6. Report distances between real and generated data: marginal Wasserstein, KS distances, mean distance, covariance and correlation Frobenius norms, sliced Wasserstein, and downstream estimand replication.
7. Run estimator comparisons on many generated datasets.
8. Add ablations by changing one empirical feature at a time: overlap, tails, treatment imbalance, covariance, policy rules, or outcome noise.

The ablation step is where diffusion can be especially useful. Because generation is a gradual reverse process, one can condition, guide, or project intermediate draws. Examples include:

• holding treatment fixed and generating only \((Y,X)\) conditional on \(W\);
• matching marginal treatment shares exactly while allowing covariates and outcomes to vary;
• projecting generated rows back into known support restrictions;
• increasing or decreasing tail noise to stress-test robust estimators;
• conditioning on policy environments to compare counterfactual regimes.

The relevant standard is practical: does the DDPM improve the credibility of finite-sample estimator comparisons?

Treatment as Target, Context, or Assignment

Treatment status plays different roles in different simulation designs. The distinction matters because the estimand changes when the simulated object changes.

In a joint row simulation, the target is

\[ P(Y,W,X). \]

The generator draws complete rows, including treatment. This is useful when the assignment mechanism is part of the empirical environment one wants to reproduce. The risk is that rare treatments may be poorly learned, and the realized treatment share will vary across synthetic samples.

In a conditional row simulation, the target is

\[ P(Y,X \mid W), \]

and the analyst supplies the treatment vector \(W_1,\ldots,W_n\). This is the design used in the Lalonde exercises below. It keeps the sample size and treatment count fixed, so differences across generators come from the learned conditional outcome/covariate distribution rather than from random variation in treatment assignment.

In an assignment-mechanism simulation, one separately models

\[ P(X), \qquad P(W\mid X), \qquad P(Y\mid W,X). \]

This is closer to a causal DGP. It is also more demanding because misspecification in the assignment model changes overlap, treatment shares, and the interpretation of ATT or ATE exercises. A learned simulator can be used for any of these objects, but the paper must state which one is being used. In the empirical sections below, Iris fixes species labels as context, the held-out scikit benchmarks are unconditional row simulations, and the Lalonde designs fix treatment as context.

Empirical Illustration

Iris Sanity Check

The Iris dataset is too small and too clean to be an econometric benchmark. It is useful here because it lets us see the mechanics without hiding behind a large application. The row is

\[ x_i = (\text{sepal length}, \text{sepal width}, \text{petal length}, \text{petal width}, \text{species}). \]

The exercise fixes the species counts and learns the conditional distribution of the four continuous measurements given species. This is the same design pattern one would use when a binary treatment indicator or site membership is part of the simulation design rather than a variable the generator is allowed to change.

from __future__ import annotations

import sys
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from scipy.stats import wasserstein_distance
from sklearn.datasets import load_breast_cancer, load_diabetes, load_iris, load_wine
from sklearn.model_selection import train_test_split

TREX_ROOT = Path("/home/alal/Desktop/code/econometrics/trex")
if str(TREX_ROOT) not in sys.path:
    sys.path.insert(0, str(TREX_ROOT))

from trex import TabularDiffusion, TabularTransformer, TabularWGAN, distribution_metrics

SEED = 20260526
rng = np.random.default_rng(SEED)
torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def one_hot(values: pd.Series | np.ndarray) -> np.ndarray:
    values = np.asarray(values, dtype=int)
    levels = np.sort(np.unique(values))
    return np.column_stack([(values == level).astype(float) for level in levels])

def tidy_metrics(real: pd.DataFrame, fake: pd.DataFrame, columns: list[str], label: str) -> pd.DataFrame:
    out = distribution_metrics(real[columns].to_numpy(), fake[columns].to_numpy(), seed=SEED)
    return pd.DataFrame([{"sample": label, **out}])

iris = load_iris(as_frame=True)
iris_df = iris.frame.copy()
iris_df.columns = [
    "sepal_length",
    "sepal_width",
    "petal_length",
    "petal_width",
    "species",
]
iris_df["species"] = iris_df["species"].astype(int)

iris_features = ["sepal_length", "sepal_width", "petal_length", "petal_width"]
iris_context = one_hot(iris_df["species"])

iris_transformer = TabularTransformer(
    column_names=iris_features,
    nonnegative_columns=iris_features,
).fit(iris_df[iris_features])
iris_z = iris_transformer.transform(iris_df[iris_features])
iris_lower, iris_upper = iris_transformer.transformed_bounds()

iris_wgan = TabularWGAN(
    hidden_dims=(64, 64),
    critic_hidden_dims=(64, 64),
    batch_size=64,
    max_steps=500,
    critic_steps=2,
    lr=1e-4,
    optimizer="optimistic_adam",
    lower_bounds=iris_lower,
    upper_bounds=iris_upper,
    seed=SEED,
    device=DEVICE,
)
start = time.perf_counter()
iris_wgan.fit(iris_z, context=iris_context)
iris_wgan_seconds = time.perf_counter() - start

iris_wgan_z = iris_wgan.sample(len(iris_df), context=iris_context).numpy()
iris_wgan_fake = pd.DataFrame(
    iris_transformer.inverse_transform(iris_wgan_z),
    columns=iris_features,
)
iris_wgan_fake["species"] = iris_df["species"].to_numpy()

iris_ddpm = TabularDiffusion(
    hidden_dims=(64, 64),
    n_timesteps=50,
    max_steps=700,
    batch_size=64,
    lr=1e-3,
    seed=SEED,
    device=DEVICE,
)
start = time.perf_counter()
iris_ddpm.fit(iris_z, context=iris_context)
iris_seconds = time.perf_counter() - start

iris_fake_z = iris_ddpm.sample(len(iris_df), context=iris_context).numpy()
iris_fake = pd.DataFrame(
    iris_transformer.inverse_transform(iris_fake_z),
    columns=iris_features,
)
iris_fake["species"] = iris_df["species"].to_numpy()

iris_metrics = pd.concat(
    [
        tidy_metrics(
            iris_df,
            iris_wgan_fake,
            [*iris_features, "species"],
            f"WGAN synthetic, {iris_wgan_seconds:.1f}s",
        ),
        tidy_metrics(
            iris_df,
            iris_fake,
            [*iris_features, "species"],
            f"DDPM synthetic, {iris_seconds:.1f}s",
        ),
    ],
    ignore_index=True,
)
iris_metrics[[
    "sample",
    "marginal_w1_mean",
    "marginal_ks_mean",
    "mean_l2",
    "cov_frobenius",
    "corr_frobenius",
    "sliced_wasserstein",
]].round(4).style.hide(axis="index")

sample	marginal_w1_mean	marginal_ks_mean	mean_l2	cov_frobenius	corr_frobenius	sliced_wasserstein
WGAN synthetic, 5.1s	0.531700	0.356000	0.918200	4.560700	2.512400	0.543400
DDPM synthetic, 1.2s	0.125200	0.112000	0.100800	1.448600	0.529700	0.141000

species_names = dict(enumerate(iris.target_names))
fig, axes = plt.subplots(1, 3, figsize=(9, 3.4), sharex=True, sharey=True, constrained_layout=True)

for ax, data, title in [
    (axes[0], iris_df, "Observed Iris"),
    (axes[1], iris_wgan_fake, "WGAN synthetic"),
    (axes[2], iris_fake, "DDPM synthetic"),
]:
    for species, group in data.groupby("species"):
        ax.scatter(
            group["sepal_length"],
            group["petal_length"],
            s=28,
            alpha=0.75,
            label=species_names[int(species)],
        )
    ax.set_title(title)
    ax.set_xlabel("sepal length")
axes[0].set_ylabel("petal length")
axes[2].legend(frameon=False, loc="lower right")
plt.show()

The Iris example is deliberately modest. Its diagnostic role is to check whether small conditional generators preserve species-conditional clusters and produce plausible continuous measurements. The WGAN panel shows the adversarial baseline; the DDPM panel shows the same task fit through supervised denoising.

Held-Out Distance Benchmarks

For reporting generator quality, a held-out comparison is usually preferable to comparing synthetic rows against the same empirical rows used for fitting. The in-sample distance answers whether the generator can reproduce the training distribution. The held-out distance asks a more relevant question: after fitting on one sample, how close are generated rows to new rows from the same empirical source?

The following benchmark uses small numeric datasets bundled with scikit-learn. For each dataset, the models fit only the training rows. They then generate the same number of rows as the test split, and distances are computed against the held-out test rows. A train-resample baseline draws rows with replacement from the training set; it is not a learned generator, but it gives a useful finite-sample reference.

def sklearn_frame(loader, name: str) -> pd.DataFrame:
    data = loader(as_frame=True)
    frame = data.frame.copy()
    if "target" in frame.columns:
        frame = frame.drop(columns=["target"])
    frame.columns = [f"x{j}" for j in range(frame.shape[1])]
    frame = frame.astype(float)
    frame.attrs["dataset_name"] = name
    return frame

sklearn_datasets = [
    sklearn_frame(load_iris, "iris features"),
    sklearn_frame(load_wine, "wine features"),
    sklearn_frame(load_breast_cancer, "breast cancer features"),
    sklearn_frame(load_diabetes, "diabetes features"),
]

def fit_small_generators(train: pd.DataFrame, test: pd.DataFrame, seed: int) -> pd.DataFrame:
    columns = list(train.columns)
    transformer = TabularTransformer(
        column_names=columns,
        nonnegative_columns=columns,
    ).fit(train)
    z_train = transformer.transform(train)
    lower, upper = transformer.transformed_bounds()
    n_test = len(test)

    wgan_model = TabularWGAN(
        hidden_dims=(48, 48),
        critic_hidden_dims=(48, 48),
        batch_size=min(128, len(train)),
        max_steps=250,
        critic_steps=2,
        lr=1e-4,
        optimizer="optimistic_adam",
        lower_bounds=lower,
        upper_bounds=upper,
        seed=seed,
        device=DEVICE,
    )
    ddpm_model = TabularDiffusion(
        hidden_dims=(48, 48),
        n_timesteps=50,
        batch_size=min(128, len(train)),
        max_steps=300,
        lr=1e-3,
        seed=seed,
        device=DEVICE,
    )

    wgan_model.fit(z_train)
    ddpm_model.fit(z_train)

    train_resample = train.sample(n_test, replace=True, random_state=seed).reset_index(drop=True)
    wgan_fake = pd.DataFrame(
        transformer.inverse_transform(wgan_model.sample(n_test).numpy()),
        columns=columns,
    )
    ddpm_fake = pd.DataFrame(
        transformer.inverse_transform(ddpm_model.sample(n_test).numpy()),
        columns=columns,
    )

    rows = []
    for label, fake in [
        ("train resample", train_resample),
        ("WGAN", wgan_fake),
        ("DDPM", ddpm_fake),
    ]:
        metrics = distribution_metrics(
            test[columns].to_numpy(),
            fake[columns].to_numpy(),
            n_projections=64,
            seed=seed,
        )
        rows.append({"model": label, **metrics})
    return pd.DataFrame(rows)

heldout_rows = []
for k, frame in enumerate(sklearn_datasets):
    train, test = train_test_split(frame, test_size=0.30, random_state=SEED + k)
    result = fit_small_generators(
        train.reset_index(drop=True),
        test.reset_index(drop=True),
        SEED + 100 * k,
    )
    result.insert(0, "dataset", frame.attrs["dataset_name"])
    result.insert(1, "n_train", len(train))
    result.insert(2, "n_test", len(test))
    heldout_rows.append(result)

heldout_table = pd.concat(heldout_rows, ignore_index=True)
heldout_table[[
    "dataset",
    "n_train",
    "n_test",
    "model",
    "marginal_w1_mean",
    "marginal_ks_mean",
    "mean_l2",
    "corr_frobenius",
    "sliced_wasserstein",
]].round(4).style.hide(axis="index")

dataset	n_train	n_test	model	marginal_w1_mean	marginal_ks_mean	mean_l2	corr_frobenius	sliced_wasserstein
iris features	105	45	train resample	0.185600	0.144400	0.092700	0.526400	0.180800
iris features	105	45	WGAN	0.767000	0.600000	1.352900	1.405300	0.712900
iris features	105	45	DDPM	0.236100	0.211100	0.117100	0.721300	0.227800
wine features	124	54	train resample	7.238800	0.170900	85.463200	2.266000	19.293800
wine features	124	54	WGAN	18.541000	0.535600	103.444700	6.568500	49.041100
wine features	124	54	DDPM	5.508700	0.208000	22.973800	3.244900	14.316400
breast cancer features	398	171	train resample	6.070300	0.092400	105.952600	3.840500	16.305200
breast cancer features	398	171	WGAN	20.242500	0.481300	164.367700	13.082300	54.865400
breast cancer features	398	171	DDPM	10.397400	0.165900	72.737000	12.204700	33.915000
diabetes features	309	133	train resample	0.006400	0.096200	0.013100	0.971300	0.006200
diabetes features	309	133	WGAN	0.038700	0.567700	0.040500	3.947800	0.037500
diabetes features	309	133	DDPM	0.026900	0.557100	0.063100	2.293700	0.025500

fig, ax = plt.subplots(figsize=(8, 3.6), constrained_layout=True)
plot_data = heldout_table.pivot(index="dataset", columns="model", values="sliced_wasserstein")
plot_data = plot_data[["train resample", "WGAN", "DDPM"]]
x = np.arange(len(plot_data))
width = 0.24
for j, model in enumerate(plot_data.columns):
    ax.bar(x + (j - 1) * width, plot_data[model], width=width, label=model)
ax.set_xticks(x)
ax.set_xticklabels(plot_data.index, rotation=20, ha="right")
ax.set_ylabel("sliced Wasserstein to held-out test rows")
ax.set_title("Held-out row-distribution distance")
ax.legend(frameon=False, ncols=3)
plt.show()

Held-out distances are noisy on these small datasets, but they are conceptually cleaner than in-sample distances. A generator should be judged against rows it did not see during fitting, and the train-resample baseline helps distinguish model failure from the ordinary sampling variability of a small empirical distribution.

Lalonde Designs

The AIMM paper uses WGANs to build simulation designs for treatment-effect estimators. The Lalonde setting is useful because it has two distinct empirical targets. The experimental target is the NSW randomized sample: treated and control units both come from the experiment. The observational target combines NSW treated units with a large CPS control pool. The second target is closer to the classic nonexperimental evaluation problem: controls come from a different empirical environment, treatment is rare, earnings are heavy-tailed and truncated, and covariate overlap is strained.

The exercise below evaluates both targets. In each case, earnings are scaled to thousands of dollars, the binary treatment indicator \(W\) is held fixed as context, and each generator learns the conditional distribution of \((Y,X)\) given \(W\). Thus the experimental and observational exercises differ in the empirical row distribution being learned, not in the simulation protocol.

LDW_COLUMNS = [
    "t", "age", "education", "black", "hispanic",
    "married", "nodegree", "re74", "re75", "re78",
]
EARNINGS_COLUMNS = ["re74", "re75", "re78"]
BINARY_COLUMNS = ["black", "hispanic", "married", "nodegree"]
NONNEGATIVE_COLUMNS = ["age", "education", "re74", "re75", "re78"]
COVARIATES = ["age", "education", "black", "hispanic", "married", "nodegree", "re74", "re75"]
OUTCOME = "re78"
TREATMENT = "t"

LOCAL_DATA = Path("data/original_data")
LALONDE_SOURCES = {
    "NSW experimental": LOCAL_DATA / "nsw_experimental.feather",
    "NSW/CPS observational": LOCAL_DATA / "cps_merged.feather",
}

missing_sources = [str(path) for path in LALONDE_SOURCES.values() if not path.exists()]
if missing_sources:
    raise FileNotFoundError(f"Missing local Lalonde data files: {missing_sources}")

def load_lalonde_sample(path: Path) -> pd.DataFrame:
    out = pd.read_feather(path)[LDW_COLUMNS].astype(float)
    out[EARNINGS_COLUMNS] = out[EARNINGS_COLUMNS] / 1000.0
    return out

lalonde_samples = {
    name: load_lalonde_sample(path)
    for name, path in LALONDE_SOURCES.items()
}

lalonde_summary = pd.DataFrame([
    {
        "design": name,
        "n": len(df),
        "n_treated": int(df[TREATMENT].sum()),
        "n_control": int(len(df) - df[TREATMENT].sum()),
        "treated_share": float(df[TREATMENT].mean()),
    }
    for name, df in lalonde_samples.items()
])

The experimental sample has 445 rows, with 185 treated observations and 260 experimental controls. The observational sample has 16,177 rows, with the same 185 NSW treated observations and 15,992 CPS controls. The treated share is therefore about 41.6 percent in the experimental sample and 1.1 percent in the observational sample.

lalonde_summary.round(4).style.hide(axis="index")

design	n	n_treated	n_control	treated_share
NSW experimental	445	185	260	0.415700
NSW/CPS observational	16177	185	15992	0.011400

def design_matrix(df: pd.DataFrame) -> np.ndarray:
    return np.column_stack([np.ones(len(df)), df[COVARIATES].to_numpy(dtype=float)])

def regression_adjustment_att(df: pd.DataFrame) -> dict[str, float | int | str]:
    data = df[LDW_COLUMNS].astype(float).copy()
    t = data[TREATMENT].to_numpy(dtype=int)
    y = data[OUTCOME].to_numpy(dtype=float)
    X = design_matrix(data)
    treated = t == 1
    control = t == 0
    if treated.sum() < 5 or control.sum() <= X.shape[1] + 1:
        return {
            "att_thousand": np.nan,
            "att_dollars": np.nan,
            "se_thousand": np.nan,
            "se_dollars": np.nan,
            "n": len(data),
            "n_treated": int(treated.sum()),
            "n_control": int(control.sum()),
            "status": "insufficient support",
        }
    beta0 = np.linalg.lstsq(X[control], y[control], rcond=None)[0]
    treated_effects = y[treated] - X[treated] @ beta0
    att = float(treated_effects.mean())
    se = float(treated_effects.std(ddof=1) / np.sqrt(treated.sum()))
    return {
        "att_thousand": att,
        "att_dollars": 1000.0 * att,
        "se_thousand": se,
        "se_dollars": 1000.0 * se,
        "n": len(data),
        "n_treated": int(treated.sum()),
        "n_control": int(control.sum()),
        "status": "ok",
    }

def result_row(design: str, name: str, df: pd.DataFrame) -> dict[str, float | int | str]:
    row = {"design": design, "sample": name, **regression_adjustment_att(df)}
    row["treated_share"] = row["n_treated"] / row["n"]
    return row

feature_columns = [column for column in LDW_COLUMNS if column != TREATMENT]
feature_binary = [column for column in BINARY_COLUMNS if column in feature_columns]
feature_nonnegative = [column for column in NONNEGATIVE_COLUMNS if column in feature_columns]

def rebuild_lalonde(generated_features: np.ndarray, generated_t: np.ndarray) -> pd.DataFrame:
    out = pd.DataFrame(generated_features, columns=feature_columns)
    out.insert(0, TREATMENT, generated_t.reshape(-1).astype(float))
    out = out[LDW_COLUMNS].astype(float)
    for col in [TREATMENT, *BINARY_COLUMNS]:
        out[col] = (np.clip(out[col], 0.0, 1.0) >= 0.5).astype(float)
    for col in NONNEGATIVE_COLUMNS:
        out[col] = np.maximum(out[col], 0.0)
    return out

def fit_lalonde_design(design: str, df: pd.DataFrame, seed_offset: int) -> dict[str, object]:
    features = df[feature_columns].copy()
    context = df[[TREATMENT]].to_numpy(dtype=np.float32)

    transformer = TabularTransformer(
        column_names=feature_columns,
        binary_columns=feature_binary,
        nonnegative_columns=feature_nonnegative,
    ).fit(features)
    train_z = transformer.transform(features)
    lower, upper = transformer.transformed_bounds()
    batch_size = min(512, len(df))

    wgan_model = TabularWGAN(
        hidden_dims=(64, 64),
        critic_hidden_dims=(64, 64),
        batch_size=batch_size,
        max_steps=450,
        critic_steps=2,
        lr=1e-4,
        optimizer="optimistic_adam",
        gp_weight=5.0,
        binary_dims=[feature_columns.index(col) for col in feature_binary],
        lower_bounds=lower,
        upper_bounds=upper,
        seed=SEED + seed_offset,
        device=DEVICE,
    )
    ddpm_model = TabularDiffusion(
        hidden_dims=(64, 64),
        n_timesteps=80,
        batch_size=batch_size,
        max_steps=650,
        lr=1e-3,
        seed=SEED + 1000 + seed_offset,
        device=DEVICE,
    )

    start = time.perf_counter()
    wgan_model.fit(train_z, context=context)
    wgan_seconds = time.perf_counter() - start

    start = time.perf_counter()
    ddpm_model.fit(train_z, context=context)
    ddpm_seconds = time.perf_counter() - start

    def synthesize(model: object, context_array: np.ndarray) -> pd.DataFrame:
        generated = model.sample(len(context_array), context=context_array).numpy()
        return rebuild_lalonde(
            transformer.inverse_transform(generated),
            context_array,
        )

    return {
        "design": design,
        "observed": df,
        "context": context,
        "transformer": transformer,
        "wgan": wgan_model,
        "ddpm": ddpm_model,
        "wgan_fake": synthesize(wgan_model, context),
        "ddpm_fake": synthesize(ddpm_model, context),
        "wgan_seconds": wgan_seconds,
        "ddpm_seconds": ddpm_seconds,
        "synthesize": synthesize,
    }

lalonde_fits = {
    design: fit_lalonde_design(design, df, 100 * k)
    for k, (design, df) in enumerate(lalonde_samples.items())
}

fit_times = pd.DataFrame([
    {"design": design, "model": "WGAN-GP", "seconds": fit["wgan_seconds"]}
    for design, fit in lalonde_fits.items()
] + [
    {"design": design, "model": "DDPM", "seconds": fit["ddpm_seconds"]}
    for design, fit in lalonde_fits.items()
])

The timing comparison is included only to show the practical cost of the two fitting routines in this small benchmark. It is not the main estimand of the exercise.

fig, ax = plt.subplots(figsize=(7.5, 2.8), constrained_layout=True)
plot_times = fit_times.pivot(index="design", columns="model", values="seconds")
x = np.arange(len(plot_times))
width = 0.34
for j, model in enumerate(plot_times.columns):
    ax.bar(x + (j - 0.5) * width, plot_times[model], width=width, label=model)
ax.set_xticks(x)
ax.set_xticklabels(plot_times.index, rotation=10, ha="right")
ax.set_ylabel("seconds")
ax.set_title("Generator fit time")
ax.legend(frameon=False, ncols=2)
plt.show()

att_rows = []
metric_rows = []
for design, fit in lalonde_fits.items():
    observed = fit["observed"]
    wgan_fake = fit["wgan_fake"]
    ddpm_fake = fit["ddpm_fake"]
    att_rows.extend([
        result_row(design, "Observed", observed),
        result_row(design, "WGAN-GP synthetic", wgan_fake),
        result_row(design, "DDPM synthetic", ddpm_fake),
    ])
    metric_rows.extend([
        tidy_metrics(observed, wgan_fake, LDW_COLUMNS, "WGAN-GP synthetic").assign(design=design),
        tidy_metrics(observed, ddpm_fake, LDW_COLUMNS, "DDPM synthetic").assign(design=design),
    ])

att_table = pd.DataFrame(att_rows)
metric_table = pd.concat(metric_rows, ignore_index=True)

att_table.round(4).style.hide(axis="index")

design	sample	att_thousand	att_dollars	se_thousand	se_dollars	n	n_treated	n_control	status	treated_share
NSW experimental	Observed	1.787800	1787.760600	0.573900	573.901900	445	185	260	ok	0.415700
NSW experimental	WGAN-GP synthetic	-0.131000	-131.021800	0.053400	53.417400	445	185	260	ok	0.415700
NSW experimental	DDPM synthetic	-0.295500	-295.484600	0.381200	381.243900	445	185	260	ok	0.415700
NSW/CPS observational	Observed	0.689900	689.858100	0.598500	598.457200	16177	185	15992	ok	0.011400
NSW/CPS observational	WGAN-GP synthetic	-0.323300	-323.318600	0.073300	73.313700	16177	185	15992	ok	0.011400
NSW/CPS observational	DDPM synthetic	2.411500	2411.506100	0.481300	481.329100	16177	185	15992	ok	0.011400

metric_table[[
    "design",
    "sample",
    "marginal_w1_mean",
    "marginal_ks_mean",
    "mean_l2",
    "cov_frobenius",
    "corr_frobenius",
    "sliced_wasserstein",
]].round(4).style.hide(axis="index")

design	sample	marginal_w1_mean	marginal_ks_mean	mean_l2	cov_frobenius	corr_frobenius	sliced_wasserstein
NSW experimental	WGAN-GP synthetic	1.073400	0.227200	2.477500	64.624500	3.284600	1.612000
NSW experimental	DDPM synthetic	0.326700	0.137300	0.467300	22.441700	0.596200	0.379600
NSW/CPS observational	WGAN-GP synthetic	2.980800	0.298400	7.960600	250.852000	2.983500	4.003800
NSW/CPS observational	DDPM synthetic	0.898400	0.087000	0.992500	77.744000	0.602400	0.884100

The ATT and distance tables separate two questions. The ATT table asks whether the downstream regression-adjustment estimand is similar. The distance table asks whether the generated rows resemble the empirical rows before any estimator is applied. The experimental design is a comparatively balanced benchmark. The observational design is a stress test with the same treated sample but a much larger and different control pool. A generator that looks acceptable in the experimental sample can still fail in the observational sample because the latter puts much more weight on overlap, earnings tails, and the dependence between treatment status and pre-treatment earnings.

The sliced Wasserstein distance is especially useful here because it is multivariate. It projects the full row vector onto many random directions and averages the one-dimensional Wasserstein distances. A generator can match each column separately while getting the joint geometry wrong; sliced Wasserstein is meant to catch that failure without solving a full high-dimensional optimal-transport problem. The Frobenius distances for covariance and correlation matrices provide a more classical second-moment version of the same check.

fig, axes = plt.subplots(
    len(lalonde_fits),
    3,
    figsize=(10.5, 5.8),
    sharex=True,
    sharey=True,
    constrained_layout=True,
)
for row, (design, fit) in enumerate(lalonde_fits.items()):
    plot_items = [
        ("Observed", fit["observed"]),
        ("WGAN-GP", fit["wgan_fake"]),
        ("DDPM", fit["ddpm_fake"]),
    ]
    for col, (title, df) in enumerate(plot_items):
        ax = axes[row, col]
        sample = df.sample(min(4000, len(df)), random_state=SEED + row)
        ax.scatter(sample["re75"], sample["re78"], s=5, alpha=0.16)
        ax.set_title(f"{design}: {title}")
        ax.set_xlabel("1975 earnings, thousands")
    axes[row, 0].set_ylabel("1978 earnings, thousands")
plt.show()

One synthetic dataset is only a first pass. A generator can match the ATT and still miss the row distribution. The distance table and scatter plot are therefore part of the simulation design, not decoration. Comparing the experimental and observational designs also clarifies what the synthetic-data task is supposed to preserve: randomized-sample balance in the first case, and nonexperimental control-pool geometry in the second.

Repeated Simulation Results

The single synthetic draw above answers whether one fitted generator can produce a plausible table for each Lalonde target. A simulation study also needs repeated draws from the fitted generator. The estimand and sample geometry are held fixed within each target: every synthetic dataset has the observed sample size and the observed treatment vector for that target. This isolates the fitted row distribution from changes in treatment prevalence.

The repeated design is:

\[ \tilde R_{ib}^{(m)} \sim \hat P_m(\cdot \mid W = w_i), \qquad i=1,\ldots,n,\quad b=1,\ldots,B, \]

where \(m \in \{\text{WGAN}, \text{DDPM}\}\) indexes the generator and \(w_i\) is the observed treatment assignment for row \(i\). Each replicate is evaluated by the same regression-adjustment ATT and by the same distribution distances used above.

def synthetic_from_fit(fit: dict[str, object], model_key: str, context_array: np.ndarray) -> pd.DataFrame:
    return fit["synthesize"](fit[model_key], context_array)

def simulation_row(
    design: str,
    model_name: str,
    replicate: int,
    observed: pd.DataFrame,
    df: pd.DataFrame,
) -> dict[str, float | int | str]:
    att = result_row(design, model_name, df)
    metrics = distribution_metrics(
        observed[LDW_COLUMNS].to_numpy(),
        df[LDW_COLUMNS].to_numpy(),
        n_projections=64,
        seed=SEED + replicate,
    )
    observed_att = regression_adjustment_att(observed)["att_dollars"]
    return {
        "design": design,
        "model": model_name,
        "replicate": replicate,
        "att_dollars": att["att_dollars"],
        "se_dollars": att["se_dollars"],
        "att_error_dollars": att["att_dollars"] - observed_att,
        "sliced_wasserstein": metrics["sliced_wasserstein"],
        "marginal_ks_mean": metrics["marginal_ks_mean"],
        "corr_frobenius": metrics["corr_frobenius"],
        "cov_frobenius": metrics["cov_frobenius"],
    }

N_REPLICATES = 12
simulation_rows = []
for design, fit in lalonde_fits.items():
    observed = fit["observed"]
    context = fit["context"]
    for b in range(N_REPLICATES):
        simulation_rows.append(
            simulation_row(design, "WGAN-GP", b, observed, synthetic_from_fit(fit, "wgan", context))
        )
        simulation_rows.append(
            simulation_row(design, "DDPM", b, observed, synthetic_from_fit(fit, "ddpm", context))
        )

simulation_table = pd.DataFrame(simulation_rows)
simulation_summary = (
    simulation_table
    .groupby(["design", "model"])
    .agg(
        att_mean_dollars=("att_dollars", "mean"),
        att_sd_dollars=("att_dollars", "std"),
        abs_att_error_mean=("att_error_dollars", lambda x: np.mean(np.abs(x))),
        sliced_wasserstein_mean=("sliced_wasserstein", "mean"),
        marginal_ks_mean=("marginal_ks_mean", "mean"),
        corr_frobenius_mean=("corr_frobenius", "mean"),
        cov_frobenius_mean=("cov_frobenius", "mean"),
    )
    .reset_index()
)
simulation_summary.round(4).style.hide(axis="index")

design	model	att_mean_dollars	att_sd_dollars	abs_att_error_mean	sliced_wasserstein_mean	marginal_ks_mean	corr_frobenius_mean	cov_frobenius_mean
NSW experimental	DDPM	476.925700	695.365300	1310.834900	0.448300	0.134800	0.633900	25.167400
NSW experimental	WGAN-GP	-156.049600	48.931700	1943.810300	1.641100	0.229300	3.241800	65.175700
NSW/CPS observational	DDPM	1717.733800	506.847000	1027.875600	0.842100	0.088400	0.620400	77.999600
NSW/CPS observational	WGAN-GP	-314.292800	55.946300	1004.150900	3.908400	0.298400	2.987500	250.750300

The repeated draws turn the generator comparison into a Monte Carlo object. The ATT columns ask whether the downstream estimator sees a similar treatment-effect problem. The distance columns ask whether the synthetic tables reproduce the empirical distribution that makes the treatment-effect problem hard. The comparison should be read within each design. The experimental target tests whether the models can reproduce a randomized small-sample environment. The observational target tests whether they can reproduce the severe imbalance and covariate geometry of the CPS control design. Estimand replication alone can reward a generator for getting one target right for the wrong distributional reasons. Conversely, a generator with better row geometry can still produce a noisier downstream estimand in a small or rare-treatment design.

fig, axes = plt.subplots(len(lalonde_fits), 2, figsize=(9.5, 5.8), constrained_layout=True)
for row, (design, group_design) in enumerate(simulation_table.groupby("design", sort=False)):
    observed_att = regression_adjustment_att(lalonde_samples[design])["att_dollars"]
    for model, group in group_design.groupby("model"):
        xpos = 0 if model == "WGAN-GP" else 1
        axes[row, 0].scatter(
            np.full(len(group), xpos) + rng.normal(0, 0.035, len(group)),
            group["att_dollars"],
            alpha=0.75,
            s=26,
            label=model,
        )
        axes[row, 1].scatter(
            np.full(len(group), xpos) + rng.normal(0, 0.035, len(group)),
            group["sliced_wasserstein"],
            alpha=0.75,
            s=26,
            label=model,
        )
    axes[row, 0].axhline(observed_att, color="black", lw=1, ls="--")
    axes[row, 0].set_title(f"{design}: ATT")
    axes[row, 0].set_ylabel("ATT, dollars")
    axes[row, 1].set_title(f"{design}: sliced Wasserstein")
    axes[row, 1].set_ylabel("distance")
    for ax in axes[row, :]:
        ax.set_xticks([0, 1])
        ax.set_xticklabels(["WGAN", "DDPM"])
plt.show()

Treatment-Share Ablation

The observational CPS target has only 185 treated observations out of 16,177 rows. That imbalance is part of the empirical problem. A learned conditional generator lets us change this feature while keeping the fitted conditional row distributions fixed. The experimental target is already much more balanced, so the ablation is mainly informative for the observational design.

The ablation below draws a balanced sample with 500 treated and 500 controls from each fitted generator. It does not estimate a new population parameter. It asks how the regression-adjustment estimand behaves when the rare-treatment geometry is relaxed while the conditional outcome and covariate distributions are generated by the same fitted model.

balanced_context = np.r_[np.ones((500, 1)), np.zeros((500, 1))].astype(np.float32)
ablation_rows = []
for design, fit in lalonde_fits.items():
    ablation_rows.extend([
        result_row(design, "WGAN-GP observed share", fit["wgan_fake"]),
        result_row(design, "DDPM observed share", fit["ddpm_fake"]),
        result_row(
            design,
            "WGAN-GP balanced 500/500",
            synthetic_from_fit(fit, "wgan", balanced_context),
        ),
        result_row(
            design,
            "DDPM balanced 500/500",
            synthetic_from_fit(fit, "ddpm", balanced_context),
        ),
    ])
ablation_table = pd.DataFrame(ablation_rows)
ablation_table.round(4).style.hide(axis="index")

design	sample	att_thousand	att_dollars	se_thousand	se_dollars	n	n_treated	n_control	status	treated_share
NSW experimental	WGAN-GP observed share	-0.131000	-131.021800	0.053400	53.417400	445	185	260	ok	0.415700
NSW experimental	DDPM observed share	-0.295500	-295.484600	0.381200	381.243900	445	185	260	ok	0.415700
NSW experimental	WGAN-GP balanced 500/500	-0.198500	-198.488900	0.032900	32.891300	1000	500	500	ok	0.500000
NSW experimental	DDPM balanced 500/500	-0.299800	-299.773300	0.244900	244.908500	1000	500	500	ok	0.500000
NSW/CPS observational	WGAN-GP observed share	-0.323300	-323.318600	0.073300	73.313700	16177	185	15992	ok	0.011400
NSW/CPS observational	DDPM observed share	2.411500	2411.506100	0.481300	481.329100	16177	185	15992	ok	0.011400
NSW/CPS observational	WGAN-GP balanced 500/500	-0.199300	-199.317500	0.045700	45.684500	1000	500	500	ok	0.500000
NSW/CPS observational	DDPM balanced 500/500	1.453200	1453.221500	0.281300	281.281500	1000	500	500	ok	0.500000

ablation_plot = ablation_table.copy()
fig, axes = plt.subplots(
    len(lalonde_fits),
    2,
    figsize=(10, 5.4),
    sharex="col",
    constrained_layout=True,
)
if len(lalonde_fits) == 1:
    axes = np.array([axes])

for row, (design, group) in enumerate(ablation_plot.groupby("design", sort=False)):
    group = group.sort_values("sample", ascending=True)
    y = np.arange(len(group))

    axes[row, 0].barh(y, group["att_dollars"])
    axes[row, 0].axvline(
        regression_adjustment_att(lalonde_samples[design])["att_dollars"],
        color="black",
        lw=1,
        ls="--",
    )
    axes[row, 0].set_title(f"{design}: ATT")
    axes[row, 0].set_xlabel("dollars")

    axes[row, 1].barh(y, group["se_dollars"])
    axes[row, 1].axvline(
        regression_adjustment_att(lalonde_samples[design])["se_dollars"],
        color="black",
        lw=1,
        ls="--",
    )
    axes[row, 1].set_title(f"{design}: SE")
    axes[row, 1].set_xlabel("dollars")

    for ax in axes[row, :]:
        ax.set_yticks(y)
        ax.set_yticklabels(group["sample"])
plt.show()

This is the kind of ablation that learned simulation designs should make routine. The empirical target is used to fit the conditional row distribution; the analyst then changes one design feature and watches the downstream estimator respond. More substantive ablations would alter overlap, tail thickness, panel persistence, market concentration, or the outcome equation. The same implementation pattern applies: hold the fitted conditional generator fixed, change the design object deliberately, and report both distributional diagnostics and estimator consequences.

The balanced-sample exercise is a controlled perturbation of the simulation environment rather than an estimate of either original Lalonde estimand. The observed-share synthetic rows keep the empirical treatment share of their source design, while the balanced ablation has a treated share of 50 percent. The SE panel is useful because treatment-share changes mechanically alter precision even when the fitted conditional row distributions are held fixed. Large changes in the ablation table should be interpreted as sensitivity to the simulated design, not as evidence that the target estimand itself changed in the original data.

Limitations

Diffusion models do not remove judgment from simulation design. They move the judgment to preprocessing, conditioning, diagnostics, and ablation design. That is still progress, because those choices are visible and auditable.

There are also limits:

• a DDPM trained on one empirical dataset cannot invent support that is absent from that dataset;
• synthetic draws inherit sampling noise and measurement artifacts from the source data;
• postprocessing can hide model failures if it is too aggressive;
• small treatment cells remain hard because the conditional distribution is weakly learned;
• privacy is not automatic and should not be claimed without a separate privacy analysis.

Finally, DDPMs are not always better than WGANs. If the table is small, smooth, and continuous, a WGAN or normalizing flow may be enough. If fast generation is essential, WGANs have a real advantage. If exact likelihood is required, flows are more natural. The evidence here is also limited: experimental and observational Lalonde variants, a few small scikit-learn datasets, a small number of repeated synthetic draws, and no systematic hyperparameter sensitivity. The case for DDPMs is strongest as a practical hypothesis for settings with complicated support, mixed variable types, and estimator-relevant dependence patterns that adversarial training may struggle to preserve.

Conclusion

The AIMM paper made an important methodological point: Monte Carlo studies can be made more credible by learning simulation designs from real applications. WGANs were a reasonable first tool for that job. DDPMs offer a cleaner fitting problem for the same econometric goal. They learn a denoising map across a controlled corruption path, connect directly to score-based density modeling, and appear well suited to conditional row generation and targeted ablation.

For econometric practice, the useful payoff is sharper finite-sample evidence: simulation designs that preserve the empirical geometry of the problem being studied, while allowing controlled interventions on the features that make estimators succeed or fail. The right next step is not to declare a universal generator winner, but to evaluate learned simulation designs by both distributional fidelity and estimator-specific stress tests across several empirical environments.

Appendix: Training-Step Sweep

The main examples fit each generator once. A useful diagnostic is to repeat the fit for progressively larger training budgets and compare generated samples to both the training rows and held-out rows. If train distances keep falling while test distances flatten or rise, the model is learning the empirical sample more than the underlying row distribution. If both distances improve, additional training is buying distributional fidelity rather than only memorization.

The exercise below uses the scikit-learn wine features. This is not an econometric stress test, but it is small enough to make the calculation transparent. Each row is standardized by a TabularTransformer fit on the training split. For each method, one model is initialized and then trained incrementally. At 50, 100, 200, 400, 800, 1600, and 3200 steps, the current checkpoint generates one synthetic sample with the same size as the training split and one with the same size as the held-out split. Distances are computed on the original feature scale. The horizontal axis should not be read as exactly equal compute across models: one WGAN step includes critic updates, while one DDPM step is one denoising-loss update. The seconds column gives cumulative fitting cost up to the checkpoint.

step_frame = sklearn_frame(load_wine, "wine features")
step_train, step_test = train_test_split(
    step_frame,
    test_size=0.30,
    random_state=SEED + 901,
)
step_train = step_train.reset_index(drop=True)
step_test = step_test.reset_index(drop=True)
step_columns = list(step_train.columns)

step_transformer = TabularTransformer(
    column_names=step_columns,
    nonnegative_columns=step_columns,
).fit(step_train)
step_train_z = step_transformer.transform(step_train)
step_lower, step_upper = step_transformer.transformed_bounds()
step_grid = [50, 100, 200, 400, 800, 1600, 3200]

def sample_with_seed(seed: int, draw_fn):
    cpu_state = torch.random.get_rng_state()
    cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
    torch.manual_seed(seed)
    try:
        return draw_fn()
    finally:
        torch.random.set_rng_state(cpu_state)
        if cuda_state is not None:
            torch.cuda.set_rng_state_all(cuda_state)

def checkpoint_metrics(
    model_name: str,
    max_steps: int,
    seconds: float,
    fake_train_z: np.ndarray,
    fake_test_z: np.ndarray,
    seed_offset: int,
) -> pd.DataFrame:
    fake_train = pd.DataFrame(
        step_transformer.inverse_transform(fake_train_z),
        columns=step_columns,
    )
    fake_test = pd.DataFrame(
        step_transformer.inverse_transform(fake_test_z),
        columns=step_columns,
    )
    rows = []
    for split, real, fake in [
        ("train", step_train, fake_train),
        ("test", step_test, fake_test),
    ]:
        metrics = distribution_metrics(
            real[step_columns].to_numpy(),
            fake[step_columns].to_numpy(),
            n_projections=96,
            seed=SEED + seed_offset,
        )
        rows.append({
            "model": model_name,
            "max_steps": max_steps,
            "split": split,
            "seconds": seconds,
            **metrics,
        })
    return pd.DataFrame(rows)

def incremental_generator_trace(model_name: str, model, seed_base: int) -> pd.DataFrame:
    rows = []
    current_steps = 0
    start = time.perf_counter()
    for checkpoint in step_grid:
        model.partial_fit(step_train_z, steps=checkpoint - current_steps)
        current_steps = checkpoint
        seconds = time.perf_counter() - start
        rows.append(
            checkpoint_metrics(
                model_name,
                checkpoint,
                seconds,
                sample_with_seed(
                    SEED + seed_base + checkpoint,
                    lambda: model.sample(len(step_train)).numpy(),
                ),
                sample_with_seed(
                    SEED + seed_base + 1000 + checkpoint,
                    lambda: model.sample(len(step_test)).numpy(),
                ),
                seed_base + checkpoint,
            )
        )
    return pd.concat(rows, ignore_index=True)

step_ddpm = TabularDiffusion(
    hidden_dims=(64, 64),
    n_timesteps=50,
    batch_size=min(128, len(step_train)),
    max_steps=50,
    lr=1e-3,
    seed=SEED + 2000,
    device=DEVICE,
)
step_wgan = TabularWGAN(
    hidden_dims=(64, 64),
    critic_hidden_dims=(64, 64),
    batch_size=min(128, len(step_train)),
    max_steps=50,
    critic_steps=2,
    lr=1e-4,
    optimizer="optimistic_adam",
    gp_weight=5.0,
    lower_bounds=step_lower,
    upper_bounds=step_upper,
    seed=SEED + 3000,
    device=DEVICE,
)

step_sweep_table = pd.concat(
    [
        incremental_generator_trace("DDPM", step_ddpm, 2000),
        incremental_generator_trace("WGAN-GP", step_wgan, 3000),
    ],
    ignore_index=True,
)

step_sweep_table[[
    "model",
    "max_steps",
    "split",
    "seconds",
    "marginal_w1_mean",
    "marginal_ks_mean",
    "mean_l2",
    "corr_frobenius",
    "sliced_wasserstein",
]].round(4).style.hide(axis="index")

model	max_steps	split	seconds	marginal_w1_mean	marginal_ks_mean	mean_l2	corr_frobenius	sliced_wasserstein
DDPM	50	train	0.101500	6.870000	0.143300	6.968500	4.609500	19.992400
DDPM	50	test	0.101500	12.935900	0.220800	55.241300	4.379100	39.203100
DDPM	100	train	0.238500	4.846500	0.142100	23.422800	3.443500	12.627600
DDPM	100	test	0.238500	8.443800	0.180900	36.026900	3.898800	22.585000
DDPM	200	train	0.466900	6.340500	0.148900	31.040800	2.667900	17.620100
DDPM	200	test	0.466900	8.062200	0.217900	8.402700	3.420200	22.194600
DDPM	400	train	0.904100	9.041500	0.158800	44.461600	2.675800	23.442100
DDPM	400	test	0.904100	9.356000	0.206600	71.153500	3.213800	23.921900
DDPM	800	train	1.784100	7.208500	0.152000	23.380900	2.025200	20.343400
DDPM	800	test	1.784100	8.088400	0.219400	49.916700	2.640600	22.349700
DDPM	1600	train	3.233500	6.368000	0.155100	60.426100	1.544100	18.845900
DDPM	1600	test	3.233500	6.997400	0.192300	35.611000	2.075400	20.654100
DDPM	3200	train	6.031900	6.487100	0.134600	50.122500	1.272400	18.576500
DDPM	3200	test	6.031900	6.392400	0.202300	16.876500	2.427300	17.805800
WGAN-GP	50	train	0.509400	20.709600	0.493200	8.314600	6.596200	62.256700
WGAN-GP	50	test	0.509400	20.589500	0.522800	64.721400	6.313200	61.738900
WGAN-GP	100	train	0.915200	20.776600	0.511800	9.436700	6.528200	50.832400
WGAN-GP	100	test	0.915200	20.106000	0.521400	55.797100	6.642400	49.064200
WGAN-GP	200	train	1.835500	20.286400	0.535400	15.802300	6.772300	60.035200
WGAN-GP	200	test	1.835500	19.480600	0.544200	26.826400	6.853200	57.330800
WGAN-GP	400	train	3.673400	18.199000	0.539700	172.203600	4.510800	48.870800
WGAN-GP	400	test	3.673400	16.034300	0.510000	120.118600	4.921500	42.715900
WGAN-GP	800	train	7.366900	17.497000	0.475200	65.533400	6.888500	48.569800
WGAN-GP	800	test	7.366900	18.218800	0.515700	84.487900	7.423700	50.548700
WGAN-GP	1600	train	14.750200	16.735900	0.255000	127.092400	5.203500	48.730000
WGAN-GP	1600	test	14.750200	22.570000	0.265000	217.314300	5.841200	67.337500
WGAN-GP	3200	train	29.515900	5.071900	0.133400	22.586900	2.886700	13.770300
WGAN-GP	3200	test	29.515900	5.376500	0.173800	17.955800	3.641900	13.885200

fig, axes = plt.subplots(2, 2, figsize=(10, 5.4), sharex=True, constrained_layout=True)
plot_metrics = [
    ("sliced_wasserstein", "sliced Wasserstein"),
    ("marginal_w1_mean", "mean marginal Wasserstein"),
    ("marginal_ks_mean", "mean marginal KS"),
    ("corr_frobenius", "correlation Frobenius"),
]
styles = {
    ("DDPM", "train"): {"color": "tab:blue", "linestyle": "-", "marker": "o"},
    ("DDPM", "test"): {"color": "tab:blue", "linestyle": "--", "marker": "o"},
    ("WGAN-GP", "train"): {"color": "tab:orange", "linestyle": "-", "marker": "s"},
    ("WGAN-GP", "test"): {"color": "tab:orange", "linestyle": "--", "marker": "s"},
}
for ax, (metric, title) in zip(axes.ravel(), plot_metrics):
    for (model, split), group in step_sweep_table.groupby(["model", "split"], sort=False):
        ax.plot(
            group["max_steps"],
            group[metric],
            label=f"{model} {split}",
            **styles[(model, split)],
        )
    ax.set_xscale("log", base=2)
    ax.set_title(title)
    ax.set_xlabel("training steps")
    ax.set_ylabel("distance")
axes[0, 0].legend(frameon=False)
plt.show()

The train/test split in the plot is the important part. A single in-sample distance can make extra training look mechanically attractive. The held-out curve asks whether those extra optimization steps produce synthetic rows that resemble data not used in fitting. The WGAN and DDPM curves also separate two issues: distributional fidelity as optimization proceeds, and stability of the optimization problem itself. In this run, the useful diagnostic is not monotonicity at every point. The useful diagnostic is whether the test curve tracks the training curve closely or whether it begins to separate as the number of gradient steps increases.

References

Arjovsky, Martin, Soumith Chintala, and Leon Bottou. 2017. “Wasserstein Generative Adversarial Networks.” In Proceedings of the 34th International Conference on Machine Learning, 214–23.

Athey, Susan, Guido W. Imbens, Jonas Metzger, and Evan Munro. 2024. “Using Wasserstein Generative Adversarial Networks for the Design of Monte Carlo Simulations.” Journal of Econometrics.

Austin, Jacob, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. 2021. “Structured Denoising Diffusion Models in Discrete State-Spaces.” In Advances in Neural Information Processing Systems.

Gulrajani, Ishaan, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. 2017. “Improved Training of Wasserstein GANs.” In Advances in Neural Information Processing Systems.

Ho, Jonathan, Ajay Jain, and Pieter Abbeel. 2020. “Denoising Diffusion Probabilistic Models.” In Advances in Neural Information Processing Systems.

Hoogeboom, Emiel, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. 2021. “Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions.” arXiv Preprint arXiv:2102.05379.

Kotelnikov, Akim, Dmitry Baranchuk, Ivan Rubachev, and Artem Babenko. 2023. “TabDDPM: Modelling Tabular Data with Diffusion Models.” In Proceedings of the 40th International Conference on Machine Learning, 17564–79.

Nichol, Alexander Quinn, and Prafulla Dhariwal. 2021. “Improved Denoising Diffusion Probabilistic Models.” In Proceedings of the 38th International Conference on Machine Learning, 8162–71.

Papamakarios, George, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. 2021. “Normalizing Flows for Probabilistic Modeling and Inference.” Journal of Machine Learning Research 22 (57): 1–64.

Song, Jiaming, Chenlin Meng, and Stefano Ermon. 2021. “Denoising Diffusion Implicit Models.” In International Conference on Learning Representations.

Song, Yang, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2021. “Score-Based Generative Modeling Through Stochastic Differential Equations.” In International Conference on Learning Representations.