The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.
51,052 characters
A sliced Wasserstein and diffusion approach to random coefficient models
\setlength{\abovedisplayskip}{5pt}
\setlength{\belowdisplayskip}{5pt}
\setlength{\abovedisplayshortskip}{5pt}
\setlength{\belowdisplayshortskip}{5pt}
\hypersetup{colorlinks,breaklinks,urlcolor=blue,linkcolor=blue}
\title{\LARGE A sliced Wasserstein and diffusion approach to random coefficient models}
\author{
Keunwoo Lim\thanks{Department of Applied Mathematics, University of Washington, Seattle, WA 98195, USA. E-mail: \tt{[email removed]}}, ~~Ting Ye\thanks{Department of Biostatistics, University of Washington, Seattle, WA 98195, USA. E-mail: \tt{[email removed]}}, ~ and~ Fang Han\thanks{Department of Statistics, University of Washington, Seattle, WA 98195, USA; e-mail: {\tt [email removed]}}}
\date{\today}
\maketitle
\begin{abstract}
We propose a new minimum-distance estimator for linear random coefficient models. This estimator integrates the recently advanced sliced Wasserstein distance with the nearest neighbor methods, both of which enhance computational efficiency. We demonstrate that the proposed method is consistent in approximating the true distribution. Moreover, our formulation naturally leads to a diffusion process-based algorithm and is closely connected to treatment effect distribution estimation—both of which are of independent interest and hold promise for broader applications.
\end{abstract}
{\bf Keywords:} sliced Wasserstein distance, minimum-distance estimation, nearest neighbor methods, diffusion process.
\section{Introduction} \label{section: introduction}
Consider a linear random coefficient model (RCM) involving a set of random variables $(X, Y, \beta)$ with the underlying structure
\begin{align}\label{eq:linear-RCM}
Y = \langle \beta, X \rangle := \text{the inner product of }X \text{ and }\beta.
\end{align}
Here the covariates $X$ and the random coefficients $\beta$ are both supported on $\mathbb{R}^d$. Throughout this paper, we assume that $X$ and $\beta$ are independent. This paper focuses on estimating the distribution of $\beta$ based solely on the observations $\{(X_i, Y_i); i \in [n]\}$, where $[n] := \{1, 2, \ldots, n\}$ and $\{(X_i, Y_i, \beta_i); i \in [n]\}$ are $n$ independent realizations of $(X, Y, \beta)$.
Linear RCMs of the form \eqref{eq:linear-RCM} play a fundamental role in many quantitative studies; see, for instance, \cite{lewbel2017unobserved} and \cite{Bonhomme2024} for some recent reviews on the applications of RCMs in economics. Traditionally, research on RCMs has primarily focused on estimating the mean and variance of $\beta$ \citep{rubin1950note, hildreth1968some, swamy1970efficient}. While these two moments summarize average marginal effects and variances, they often fail to capture all information of interest to researchers.
To address this limitation, Beran and Hall \citep{beran92} pioneered the study of nonparametrically identifying and estimating the distribution of $\beta$. Since then, various approaches have been proposed, including minimum-distance estimation \citep{Beran76,beran1993semiparametric}, series expansion methods for (conditional) density functions \citep{beran1996nonparametric, Hohmann2016, Dunker19, Dunker25, gaillac2021nonparametric, Gaillac22}, kernel methods \citep{Hoderlein2010, Holzmann20}, nearest neighbor (NN) methods \citep{holzmann2024multivariate}, and discrete-grid-based constrained least squares methods \citep{fox2011, fox2016simple, heiss2022nonparametric}.
In this paper, we revisit the minimum-distance estimators proposed by Beran and Millar \citep{Beran76} and enhance their framework by incorporating nearest neighbor (NN) methods \citep{lin2023estimation,holzmann2024multivariate} and the sliced Wasserstein (SW) metric, the latter of which has recently gained prominence in generative artificial intelligence. Compared to generative modeling approaches based on the original Wasserstein metric \citep{arjovsky2017wasserstein, tolstikhin2017wasserstein, bousquet2017optimal}, the SW metric has been observed to provide similar empirical performance while being computationally more attractive and less sensitive to the dimension of the data \citep{kolouri2018sliced, liutkus2019sliced,fan2024minimum}. Computationally, optimization algorithms towards minimizing the SW metric have also been introduced and validated in \cite{bonneel2015sliced} and \cite{tanguy2024properties}, among many others.
Building on these insights, this paper proposes a novel SW-based minimum-distance estimator for the distribution of $\beta$, where weights are adaptively chosen over an NN graph to reduce bias. To the best of our knowledge, this formulation is new. Computationally, following the framework of \cite{tanguy2024properties}, we develop block gradient descent algorithms that achieve efficient optimization with a time complexity that scales polynomially in $d$, whereas existing methods typically require an exponentially large number of grid points in $d$. We further establish theoretical guarantees for the proposed estimator and derive its rates of convergence. Lastly, drawing inspiration from parallel research on SW flows \citep{bonnotte2013unidimensional, liutkus2019sliced} and causal random coefficient models \citep{heckman1997making}, we bridge our formulation with diffusion processes and causal inference. For the former, we introduce diffusion process-based methods to approximate the distribution of $\beta$. For the latter, we incorporate entropic regularization into the model of \cite{heckman1997making}, enhancing the identification and estimation of the treatment effect distribution.
\vspace{0.2cm}
\noindent {\bf Paper organization.} Section \ref{section: minimum-distance estimator} presents the proposed minimum-distance estimator and Section \ref{section: algorithm} illustrates the (approximate) block coordinate decent algorithms for implementation. Section \ref{section: theory} presents the main theory concerning estimation and computation. Sections \ref{section: diffusion} and \ref{sec: causal} connect the studied problem to diffusion processes and casual inference, respectively. The empirical performance of the algorithms is reported in Section \ref{section: simulations}. Proofs of all theoretical results are relegated to the Appendix.
\vspace{0.2cm}
\noindent{\bf Notation.} Let $(\mathbb{R}^{d}, \|\cdot\|_2)$ represent the $d$-dimensional real space coupled with the Euclidean metric. Write $\mathbb{S}^{d-1} = \{V \in \mathbb{R}^{d}\,\vert\, \Vert V \Vert_{2}=1\}$, $B_{r}(a) = \{ U\in \mathbb{R}^{d} \,\vert \,\Vert U - a\Vert_{2}< r \}$, and $\widebar{B}_{r}(a)$ to represent the sphere, the radius-$r$ open ball, and the corresponding closed ball in $(\mathbb{R},\|\cdot\|_2)$, respectively. Let $\{e_{1},\dots, e_{d}\}$ be the standard bases in $\mathbb{R}^{d}$ with regard to the Euclidean inner product $\langle\cdot,\cdot\rangle$, and let $\sigma$ be the Haar measure in $\mathbb{S}^{d-1}$. For any vectors $v_1,\ldots,v_m\in\mathbb{R}^d$, write $(v_1,\ldots,v_m)=(v_1^\top,\ldots,v_m^\top)^\top$. For vectors $U, \tilde{U}\in \mathbb{R}^{d}$, we also use $P^{U}(\tilde{U}) = \langle U, \tilde{U}\rangle$ to represent their inner product. Let $\mathcal{P}(K)$ denote the set of all probability measures on $K$ and $\delta_{x}\in \mathcal{P}(K)$ be the Dirac measure on $x \in K$. Lastly, for arbitrary measure $\mu \in \mathcal{P}(\mathbb{R}^{d})$, we introduce its projected measure with respect to the direction $V\in\mathbb{S}^{d-1}$ as $ \mu^{V} = \mu \circ (P^{V})^{-1}$.
\section{Method}\label{section: minimum-distance estimator}
Consider $\{(X_i, Y_i, \beta_i); i \in [n]\}$ to be $n$ realizations of \eqref{eq:linear-RCM}. The objective of interest is to estimate the distribution of $\beta$, denoted as $\mu_\beta$, based solely on the observations $(X_i, Y_i)$'s.
This paper is concerned with minimum-distance estimation \citep{Wolfowitz53, Wolfowitz57}, which involves an objective functional and a domain of minimization. In our formulation, the objective functional is based on the Wasserstein $W_2$ distance, defined as
\begin{align*}
W_2(\mu, \tilde{\mu}) = \bigg( \inf_{\pi \in \Pi(\mu, \tilde{\mu})} \mathbb{E}_{(x, y) \sim \pi} \Vert x - y \Vert_2^2 \bigg)^{\frac{1}{2}},
\end{align*}
where $\mu$ and $\tilde{\mu}$ are two probability measures on $\mathbb{R}^d$, and $\Pi(\mu, \tilde{\mu})$ denotes the set of all couplings of $\mu$ and $\tilde{\mu}$, i.e., probability measures on $\mathbb{R}^d \times \mathbb{R}^d$ with marginal distributions $\mu$ and $\tilde{\mu}$, respectively.
For each direction $V \in \mathbb{S}^{d-1}$, we introduce a $k$-NN empirical measure $\widebar{\mu}^V_k$ for the objective functional. Specifically, for each $i \in [n]$, first let
\[
\tilde{X}_{i} = \frac{X_{i}}{\Vert X_{i} \Vert_{2}}, \quad
\tilde{Y}_{i} = \frac{Y_{i}}{\Vert X_{i} \Vert_{2}}
\]
be the {\it normalized} counterparts of $X_i$ and $Y_i$, respectively. Let ${\mathds 1}(\cdot)$ represent the indicator function and
\[
\overline{S}_X(V, k) := \Big\{ j \in [n]: \sum_{i=1}^n {\mathds 1}\big(\| \tilde{X}_i - V \|_2 < \| \tilde{X}_j - V \|_2 \big) < k \Big\}
\]
denote the set of indices corresponding to the $k$-NNs of $V$ in $\{\tilde{X}_i; i \in [n]\}$. Since $X_i$ may take discrete values, ties can occur, causing the cardinality of $\overline{S}_X(V, k)$ to exceed $k$. In such cases, consistent with \cite{lin2024failure, lin2024consistency}, we select an {\it arbitrary} subset
\[
S_X(V, k) \subset \overline{S}_X(V, k) \quad \text{such that its cardinality}~ \big| S_X(V, k) \big| = k.
\]
When $\tilde{X}_i$ is continuous, as discussed in Section~\ref{subsection: estimation}, we have $S_X(V,k) = \overline{S}_X(V,k)$.
Using this subset and noting that, for any direction $V \in \mathbb{S}^{d-1}$ and $i \in [n]$,
\begin{align}\label{eq:insight}
\beta_i^\top V \approx \tilde{Y}_i \tilde{X}_i^\top V \quad \text{as } \tilde{X}_i \approx V,
\end{align}
we define the NN-induced random measure as
\[
\widebar{\mu}^V_k = \frac{1}{k} \sum_{i \in S_X(V, k)} \delta_{P^V(\tilde{Y}_i \tilde{X}_i)}
= \frac{1}{k} \sum_{i \in S_X(V, k)} \delta_{\tilde{Y}_i \tilde{X}_i^\top V}.
\]
In light of \eqref{eq:insight}, for arbitrary $V$ and all sufficiently small $k$, the laws of $\widebar{\mu}^V_k$ and $\beta^\top V$ should be close to each other.
Encouraged by the above insight, we introduce the following objective functional
\begin{align}\label{eq:fk}
\mathcal{F}_k(\mu) = \int_{\mathbb{S}^{d-1}} W_2^2(\widebar{\mu}^V_k, \mu^V) \, \mathrm{d} \sigma(V),
\end{align}
which integrates the squared $W_2$ distance between the NN random measure $\widebar{\mu}^V_k$ and the projected measure $\mu^V$ over all directions $V \in \mathbb{S}^{d-1}$.
Lastly, we propose the following {\it constrained} $k$-NN-based minimum-distance estimator
\begin{align}\label{eq:mde}
\hat{\mu}_\beta \in \argmin_{\mu \in \mathcal{P}_N(\widebar{B}_R(0))} \mathcal{F}_k(\mu).
\end{align}
Here $R$ is a pre-specified radius to enforce a compact space, and the domain $\mathcal{P}_N(\widebar{B}_R(0))$ consists of all size-$N$ discrete measures in $\mathcal{P}(\widebar{B}_R(0))$, defined as
\begin{align*}
\mathcal{P}_N(\widebar{B}_R(0)) = \Big\{ \eta(w) \in \mathcal{P}(\widebar{B}_R(0)) : \eta(w) = \frac{1}{N} \sum_{i=1}^N \delta_{w_i} \text{ with } w_1, \ldots, w_N \in \widebar{B}_R(0)\\
\text{and } w= (w_1, \ldots, w_N)\in \widebar{B}_R(0)^N \Big\}.
\end{align*}
The output is then a discrete approximation to any $\mu_\beta\in \mathcal{P}(\widebar{B}_R(0))$.
\begin{remark}
There are connections between \eqref{eq:fk} and the sliced Wasserstein $W_2$ distance. Notice that the sliced $W_2$ distance between any two probability measures, $\mu, \tilde{\mu} \in \mathcal{P}(\mathbb{R}^d)$, is defined as
\begin{align*}
SW_2(\mu, \tilde{\mu}) = \bigg( \int_{\mathbb{S}^{d-1}} W_2^2(\tilde\mu^V, \mu^V) \, \mathrm{d} \sigma(V) \bigg)^{\frac{1}{2}}.
\end{align*}
Thus, $\mathcal{F}_k(\cdot)$ can be interpreted as a $k$-NN-based revision of the original $SW_2$ distance, encouraged by the approximate equality in \eqref{eq:insight}. In particular, when $k = n$, $\mathcal{F}_k(\cdot)$ reduces to the $SW_2$ distance between the empirical measure and the target.
\end{remark}
\section{Algorithm}\label{section: algorithm}
The computation of $\hat{\mu}_{\beta}$ in \eqref{eq:mde} reduces to finding a size-$N$ set of elements in $\widebar{B}_R(0)$:
\begin{align}\label{eq:points}
\hat{w} \in \eta^{-1}(\hat{\mu}_{\beta}).
\end{align}
Optimizing \eqref{eq:mde} can thus also be interpreted as minimizing $\mathcal{F}_{k} \circ \eta$ over $\widebar{B}_{R}(0)^{N}$, which is identifiable up to a permutation. However, solving this minimization problem requires continuous integration over the Haar measure, which is computationally infeasible.
To address this issue, we adopt a common trick and propose a discrete Monte Carlo approximation to the original $\mathcal{F}_k$:
\begin{align}\label{eq:approximate-mde}
\widebar{\mathcal{F}}_{k}(\mu) = \frac{1}{m} \sum_{i = 1}^{m} W_{2}^{2}(\widebar{\mu}_{k}^{V_{i}}, \mu^{V_{i}}),
\end{align}
where $V_{1}, \dots, V_{m}$ are sampled independently from the Haar measure $\sigma$, independent of the system.
We accordingly introduce two algorithms for optimizing \eqref{eq:mde} using the Monte Carlo approximation \eqref{eq:approximate-mde}, namely, Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd}. Without loss of generality, in these two algorithms we set $N = k$ since otherwise, we could choose the larger of the two. In addition, in these algorithms, we use functions for {\it sorting, argsorting}, and {\it projecting} and define them as follows. For a real sequence $a = (a_{q})_{q\in[k]} = (a_{1}, \dots, a_{k})$, we define
\[
\text{sort}(a) = (a_{\nu(q)})_{q\in[k]} = (a_{\nu(1)}, \dots, a_{\nu(k)})
\]
and
\[
\text{argsort}(a) = (\nu(q))_{q\in[k]} = (\nu(1), \dots, \nu(k)),
\]
where $(\nu(1), \dots, \nu(k))$ is an arbitrary permutation of $[k]$ such that $a_{\nu(1)} \leq \cdots \leq a_{\nu(k)}$. Additionally, the projection of a vector $\psi \in \mathbb{R}^{kd}$ onto the set $\widebar{B}_{R}(0)^{k}$ in Euclidean space is defined as
\[
\text{Proj}_{\widebar{B}_{R}(0)^{k}}(\psi) = \argmin_{w \in \widebar{B}_{R}(0)^{k}} \Vert \psi - w \Vert_{2}.
\]
Algorithm \ref{algorithm: bcd} is a block coordinate descent method for computing $\hat{w}$. This algorithm is a revised version of \citet[Algorithm 1]{tanguy2024properties}, incorporating the $k$-NN graph into the computation. In each iteration, the local optimal transport law is updated using Lemma \ref{lemma: Bobkov} in the Appendix, leveraging the vector $\psi_{\ell}$ and the $k$-NN matrix $D$. The vector $\psi_{\ell+1}$ is obtained as the minimizer of the local optimal transport law within the domain $\widebar{B}_{R}(0)^{k}$.
The computation of $\psi_{\ell+1}$, however, involves convex optimization over a constrained domain, which can be computationally expensive. Algorithm \ref{algorithm: abcd} addresses this computational challenge by introducing a projected gradient descent variant. In this algorithm, the optimal transport law from Algorithm \ref{algorithm: bcd} is approximated using an $L^{2}$-distance formulation. This enables the minimization problem to be solved using linear algebra, significantly reducing computational complexity. In addition, it is possible to check that the computational complexity is $O(mdn\log n+tmk \log k)$; see Appendix Section \ref{sec:cc-calculation} ahead for details. This is substantially faster than the existing grid-based algorithms that has computational complexity exponentially growing with the dimension.
\begin{algorithm}[!ht]
\caption{Block Coordinate Descent}\label{algorithm: bcd}
\LinesNotNumbered
\KwInput{$(Y_{i}, \beta_{i}, X_{i})$, $V_{j}$, $R$, $i\in[n]$, $q\in[m]$}
\KwInitialization{$\psi_{1} \sim \text{Uniform}\,(\widebar{B}_{R}(0)^{k})$, $\psi_{0}$, $\ell \leftarrow 1$, $D\in \mathbb{R}^{k \times m}$}
\For{$j = 1$ to $m$}
{Choose the $k$-NNs of $V_{j}$ as $\{\tilde{X}_{j(q)}; q\in[k]\}$ with ties broken arbitrarily\\
$D_{\cdot,\, j} \leftarrow \text{sort}\,(\, \tilde{Y}_{j(q)}\,\langle\, V_{j},\,\tilde{X}_{j(q)}\,\rangle\,)$ for each $q\in[k]$
}
\While{$\Vert \psi_{\ell} - \psi_{\ell-1}\Vert_{2} > 0$}
{
\For{$j = 1$ to $m$}
{$\nu_{j}\leftarrow \text{argsort}\,(\langle V_{j}, \psi_{\ell, 1}\rangle,\ldots,\langle V_{j}, \psi_{\ell, k}\rangle)$
}
$\psi_{\ell+1}\leftarrow\argmin_{w\in \widebar{B}_{R}(0)^{k}}\sum_{j = 1}^{m}\sum_{q = 1}^{k}(\langle V_{j}, w_{q}\rangle - D_{\nu_{j}^{-1}(q),\, j})^{2}$\\
$\ell \leftarrow \ell+1$}
\KWReturn{$\psi_{\ell}$}
\end{algorithm}
\begin{algorithm}[!ht]
\caption{Approximate Block Coordinate Descent}
\label{algorithm: abcd}
\LinesNotNumbered
\KwInput{$(Y_{i}, \beta_{i}, X_{i})$, $V_{j}$, $R$, $t$, $i\in[n]$, $q\in[m]$}
\KwInitialization{$\psi_{1} \sim \text{Uniform}\,(\widebar{B}_{R}(0)^{k})$, $D\in \mathbb{R}^{k \times m}$}
\For{$j = 1$ to $m$}
{Choose the $k$-NNs of $V_{j}$ as $\{\tilde{X}_{j(q)}; q\in[k]\}$ with ties broken arbitrarily\\
$D_{\cdot,\, j} \leftarrow \text{sort}\,(\, \tilde{Y}_{j(q)}\,\langle\, V_{j},\,\tilde{X}_{j(q)}\,\rangle\,)$ for each $q\in[k]$
}
\For{$\ell = 1$ to $t$}
{\For{$j = 1$ to $m$}
{$\nu_{j}\leftarrow \text{argsort}\,(\langle V_{j}, \psi_{\ell, 1}\rangle,\ldots,\langle V_{j}, \psi_{\ell, k}\rangle)$
}
\For{$q = 1$ to $k$}
{$\psi_{\ell+1, q}\leftarrow m^{-1}d\sum_{j=1}^{m}D_{\nu_{j}^{-1}(q), \,j}V_{j}$
}
$\psi_{\ell+1}\leftarrow \text{Proj}_{\widebar{B}_{R}(0)^{k}}(\psi_{\ell+1})$\\
}
\KWReturn{$\psi_{t+1}$}
\end{algorithm}
\section{Theory}\label{section: theory}
This section is divided into three parts. The first part discusses the assumptions, while the other two address the rates of convergence of the estimator and the convergence of the proposed algorithms.
\subsection{Assumptions}
We focus on the case $d \geq 2$, as the case $d=1$ is trivial. Define the normalized versions of $X$ and $Y$ as $\tilde{X} = X / \Vert X \Vert_{2}$ and $\tilde{Y} = Y / \Vert X \Vert_{2}$, respectively, such that $\tilde{X} \in \mathbb{S}^{d-1}$. This normalization preserves the inner product structure, as $\tilde{Y} = \langle \beta, \tilde{X} \rangle$.
The first assumption specifies the data-generating structure, which implies that $\{(\tilde{X}_i, \tilde{Y}_i, \beta_i); i \in [n]\}$ are independent and identically distributed as $(\tilde{X}, \tilde{Y}, \beta)$.
\begin{assumption}\label{assumption:1}
Assume $\{(X_i, Y_i, \beta_i); i \in [n]\}$ are $n$ independent realizations of $(X, Y, \beta)$, which satisfies the linear structure \eqref{eq:linear-RCM} and that $X$ is independent of $\beta$.
\end{assumption}
The next assumption imposes constraints on the distribution of $\beta$, ensuring that it has a compact support. This setting is similar to those made in \cite{Beran76}, \cite{Hoderlein2010}, and \cite{Holzmann20}, and also common in other nonparametric mixture model analyses \citep{Miao24,Han23,lim2024smoothed}.
\begin{assumption}\label{assumption: coefficient}
There exists a finite constant $R < \infty$ such that $\mu_{\beta}$ is supported on $\widebar{B}_{R}(0)$.
\end{assumption}
The following two assumptions provide restrictions on the distribution of $\tilde{X}$.
\begin{assumption}\label{assumption: bounded covariates}
The law of $\tilde{X}$ is absolutely continuous with respect to the Haar measure $\sigma$ on $\mathbb{S}^{d-1}$. Furthermore, the corresponding Haar density, denoted by $f_{\tilde{X}}$, is lower bounded by some positive constant $\tau_{0} > 0$.
\end{assumption}
Assumption \ref{assumption: bounded covariates} is equivalent to Assumption 4 in \cite{Hoderlein2010}, which naturally holds when $X$ has a full-dimensional support. In particular, this assumption is valid when \eqref{eq:linear-RCM} does not include an intercept term, i.e., $X_{i,1} = 1$ for all $i \in [n]$. Such scenarios are plausible in many cases, including those discussed in \citet[Section 4.3]{Hoderlein2010} and \cite{lewbel2017unobserved}. However, it no longer holds if the model includes an intercept term. To address such cases, we introduce the following general assumption.
\begin{assumption}\label{assumption: unbounded covariates}
The law of $\tilde{X}$ is absolutely continuous with respect to the Haar measure $\sigma$ on $\mathbb{S}^{d-1}$. Additionally, there exist constants $\tau_{0}$, $\rho_{0}$, $C_{\tilde{X}}$, and $\alpha > 0$ such that for all $0 < \tau \leq \tau_{0}$ and $0 < \rho \leq \rho_{0}$, there exist closed sets $D(\tau)$, $L(\rho, \tau) \subset \mathbb{S}^{d-1}$ satisfying:
\begin{align*}
\sigma\Big(\Big\{\, V \in \mathbb{S}^{d-1} \,\Big|\, f_{\tilde{X}}(V) \leq \tau \,\Big\} \setminus D(\tau) \Big) = 0, \quad
L(\rho, \tau) = \Big\{\, V \in \mathbb{S}^{d-1} \,\Big|\, B_{\rho}(V) \cap D(\tau) = \emptyset \,\Big\},
\end{align*}
and
\[
\sup_{\tilde{V} \in \mathbb{S}^{d-1} \setminus L(\rho, \tau)} \min_{V \in L(\rho, \tau)} \Vert V - \tilde{V} \Vert_{2} \leq C_{\tilde{X}} \tau^{\alpha} + \rho.
\]
\end{assumption}
Assumption~\ref{assumption: unbounded covariates} is designed to be weak, although at the expense of some technical complexity. To aid understanding, we present the following result, which simplifies Assumption~\ref{assumption: unbounded covariates} to a tail condition on $X$.
\begin{proposition}\label{prop:heavy-tail}
Suppose there exist constants $C_{f} > 0$ and $\kappa > 1$ such that $X = (1, X_{2:d})$ and the density $f_{X_{2:d}}$ of $X_{2:d}$ satisfies:
\[
f_{X_{2:d}}(T) \geq C_{f}(1 + \Vert T \Vert_{2})^{-\kappa},
\]
for all $T \in \mathbb{R}^{d-1}$. Assumption \ref{assumption: unbounded covariates} then holds with $\alpha = 1/(\kappa + 1)$, $\tau_{0} = \mathcal{S}(\mathbb{S}^{d-1})C_{f}/2^{3\kappa + 3}$, and $\rho_{0} = 1/4$, where $\mathcal{S}(\mathbb{S}^{d-1})$ is a surface area of $\mathbb{S}^{d-1}$.
\end{proposition}
\begin{remark}
Proposition \ref{prop:heavy-tail} refines Assumption 1 in \cite{Holzmann20} and similar assumptions in \cite{Dunker19} and \cite{Dunker25} (see, e.g., the discussions in \citet[Section 4.3]{Dunker25}) by accommodating a broader range of tail conditions. In the literature, such tail conditions are often introduced for technical convenience, as we also do here. Notably, in an interesting recent paper, \cite{Gaillac22} examined some cases where $X$ is even allowed to have a compact support, albeit at the cost of imposing certain supersmoothness conditions on the density of $\beta$.
\end{remark}
\subsection{Estimation}\label{subsection: estimation}
Our first main theoretical result addresses the landscape of the objective function and the existence of minimizers to \eqref{eq:mde}.
\begin{theorem}\label{proposition: convex}
Under Assumption \ref{assumption: coefficient}, the random function $\mathcal{F}_{k} \circ \eta$ over the domain $\widebar{B}_{R}(0)^{N}$ is $4RN^{-1/2}$-Lipschitz continuous. Additionally, with probability one, there exist global minimizer(s) for \eqref{eq:mde}, which is unique (up to a permutation) if $N = k = n$.
\end{theorem}
The next theorem establishes the rates of convergence of $\hat\mu_\beta$ to $\mu_\beta$ under the $SW_2$ distance. The first part pertains to the case under Assumption \ref{assumption: bounded covariates}, and the second part addresses the case under Assumption \ref{assumption: unbounded covariates}.
\begin{theorem}\label{theorem: bounded estimator}
Assume Assumptions \ref{assumption:1} and \ref{assumption: coefficient}. Then, \textit{any} minimum-distance estimator $\hat{\mu}_{\beta}$ of \eqref{eq:mde} satisfies the following.
\begin{itemize}
\item[(i)] Under Assumption \ref{assumption: bounded covariates}, there exists a positive constant $C = C(R, d, \tau_{0})$ such that
\begin{align*}
\begin{cases}
\mathbb{E}[SW_{2}(\mu_{\beta}, \hat{\mu}_{\beta})] \leq C n^{-1/(d+5)}, &\text{for } 2 \leq d \leq 5\text{ with } N = k = n^{\frac{6}{d+5}}, \\
\mathbb{E}[SW_{2}(\mu_{\beta}, \hat{\mu}_{\beta})] \leq C n^{-1/(2d-1)}, &\text{for } d \geq 6 \text{ with } N = k = n^{\frac{d}{2d-1}}.
\end{cases}
\end{align*}
\item[(ii)] Under Assumption \ref{assumption: unbounded covariates}, there exists a positive constant $C = C(R, d, \tau_{0}, \rho_{0}, C_{\tilde{X}})$ such that
\begin{align*}
\begin{cases}
\mathbb{E}[SW_{2}(\mu_{\beta}, \hat{\mu}_{\beta})] \leq C n^{-\frac{\alpha}{2d\alpha+10\alpha+4}}\log n, &\text{for } 2 \leq d \leq 5 \text{ with } N = k = n^{\frac{6\alpha}{d\alpha+5\alpha+2}}, \\
\mathbb{E}[SW_{2}(\mu_{\beta}, \hat{\mu}_{\beta})] \leq C n^{-\frac{\alpha}{4d\alpha-2\alpha+4}}\log n, &\text{for } d \geq 6 \text{ with } N = k = n^{\frac{d\alpha}{2d\alpha-\alpha+2}}.
\end{cases}
\end{align*}
\end{itemize}
\end{theorem}
It is worth noting that Theorem \ref{proposition: convex} guarantees the existence of global minimizer(s) for \eqref{eq:mde}, but it does not provide a general guarantee of uniqueness unless $N = k = n$.
Theorem \ref{theorem: bounded estimator} further demonstrates that \textit{any} global minimizer of \eqref{eq:mde} achieves a polynomial rate of convergence to the true parameter $\mu_\beta$. For related discussions, see \citet[Remark 3.1]{Miao24}, which observe similar phenomena in nonparametric Poisson mixtures where multiple global minimizers (in that case, nonparametric maximum likelihood estimators) may exist but all converge to the estimand.
At this stage, it remains unclear whether the derived rate is minimax optimal or not. Moreover, as different metrics are employed, the obtained results are not directly comparable to those in the literature. Nevertheless, it is worth highlighting that, unlike prior works including \cite{Hoderlein2010}, \cite{Holzmann20}, and \cite{Gaillac22}, our analysis avoids imposing {\it any} smoothness conditions on the densities of $X$ or $\tilde{X}$, which underscores an advantage of the proposed distance-based approach.
\subsection{Computation}\label{subsection: computation}
Section \ref{subsection: estimation} discusses the statistical properties of the estimator $\hat\mu_\beta$ in \eqref{eq:mde}. However, from a practical standpoint, it is also crucial to understand the procedure from a computational perspective, which is the focus of this section. In this section, the algorithms of interest for approximating $\hat\mu_{\beta}$ are Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd} presented in Section \ref{section: algorithm}.
Our first theoretical result in this section examines the landscape of the approximate objective function \eqref{eq:approximate-mde}, paralleling Theorem \ref{proposition: convex}.
\begin{proposition}\label{proposition: discrete continuity}
Under Assumption \ref{assumption: coefficient}, the random function $\widebar{\mathcal{F}}_{k}\circ \eta$ over the domain $\widebar{B}_{R}(0)^{N}$ is $4RN^{-1/2}$-Lipschitz continuous. Additionally, with probability one, there exist global minimizer(s) for \eqref{eq:approximate-mde}, which is unique (up to a permutation) if $N = k = n$.
\end{proposition}
Next, we aim to quantify the algorithmic convergence. To this end, we introduce the following proposition, which analyzes the approximation of $\widebar{\mathcal{F}}_{k}\circ\eta$ to $\mathcal{F}_{k}\circ\eta$.
\begin{proposition}\label{proposition: uniform convergence}
Under Assumption \ref{assumption: coefficient} and conditional on the data:
\begin{itemize}
\item[(i)] $\widebar{\mathcal{F}}_{k}\circ \eta$ converges to $\mathcal{F}_{k}\circ \eta$ uniformly over $\widebar{B}_{R}(0)^{N}$ in the sense that
\begin{align*}
\mathbb{P}_{\sigma}\Big(\lim_{m\to\infty} \Big\Vert \widebar{\mathcal{F}}_{k}\circ\eta - \mathcal{F}_{k}\circ \eta \Big\Vert_{\infty} = 0\Big) = 1,
\end{align*}
where
\[
\Big\Vert \widebar{\mathcal{F}}_{k}\circ\eta - \mathcal{F}_{k}\circ \eta \Big\Vert_{\infty} := \sup_{w\in \widebar{B}_{R}(0)^{N}}\Big\Vert \widebar{\mathcal{F}}_{k}\circ\eta(w) - \mathcal{F}_{k}\circ \eta(w) \Big\Vert;
\]
\item[(ii)] furthermore,
\begin{align*}
\sqrt{m} \Big\Vert \widebar{\mathcal{F}}_{k}\circ\eta - \mathcal{F}_{k}\circ\eta \Big\Vert_{\infty} ~~\text{converges in distribution to}~~ \Big\Vert \mathbb{G} \Big\Vert_{\infty},
\end{align*}
where $\mathbb{G}$ is a centered Gaussian process on $\widebar{B}_{R}(0)^{N}$ with covariance structure
\begin{align*}
\textnormal{cov}\,\mathbb{G}(w, \tilde{w}) = \int_{V \in \mathbb{S}^{d-1}} W_{2}^{2}(\widebar{\mu}_{k}^{V}, \eta(w)^{V})W_{2}^{2}(\widebar{\mu}_{k}^{V}, \eta(\tilde{w})^{V}) \dd \sigma(V) - \mathcal{F}_{k}(\eta(w))\mathcal{F}_{k}(\eta(\tilde{w})).
\end{align*}
\end{itemize}
\end{proposition}
The main theorem of this section establishes the validity of Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd}.
\begin{theorem}\label{theorem: computation}
Assuming $N = k$ and under Assumption \ref{assumption: coefficient}, Algorithm \ref{algorithm: bcd} outputs a local minimizer of $\widebar{\mathcal{F}}_{k}(\eta(w))$. Furthermore, in the case of $\psi_{\ell} = \tilde{\psi}_{\ell}$, where $\psi_{\ell}$ and $\tilde{\psi}_{\ell}$ denote the $\ell$-th iteration of Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd}, respectively, we have
\begin{align*}
\lim_{m\to\infty} \mathbb{P}_{\sigma}\Big( W_{2}(\eta(\psi_{\ell+1}), \eta( \tilde{\psi}_{\ell+1})) > \epsilon \Big) = 0
\end{align*}
for all $\epsilon > 0$.
\end{theorem}
\section{A diffusion process approach}\label{section: diffusion}
In this section, building on a recent breakthrough connecting optimal transport-based generative modeling to diffusion processes \citep{liutkus2019sliced}, we propose a diffusion process-based approximation to the optimization problem \eqref{eq:mde}. While rigorous theoretical support for this approximation is currently unavailable, its empirical effectiveness has been studied and will be demonstrated in Section \ref{section: simulations}.
Specifically, consider the following minimization problem involving a regularized functional $\tilde{\mathcal{F}}_{k}^{\lambda}(\mu)$ over the space of probability measures $(\mathcal{P}(\widebar{B}_{R}(0)), W_{2})$:
\begin{align*}
\tilde{\mathcal{F}}_{k}^{\lambda}(\mu) := \frac{1}{2} \mathcal{F}_{k}(\mu) + \lambda \mathcal{H}(\mu),
\end{align*}
where $\mathcal{F}_{k}(\cdot)$ is introduced in \eqref{eq:fk}, $\lambda > 0$ is a tuning parameter, and $\mathcal{H}(\mu)$ is the regularization functional defined as:
\begin{align*}
\mathcal{H}(\mu) =
\begin{cases}
\int \varrho \log \varrho\,\dd x, & \text{if the density $\varrho$ of $\mu$ exists}, \\
+\infty, & \text{otherwise}.
\end{cases}
\end{align*}
The following proposition provides the theoretical background for formulating the associated continuity equation in this setting. For functions \(f\colon \mathbb{R}^{d} \to \mathbb{R}^{d}\) and \(\tilde{f}\colon \mathbb{R}^{d} \to \mathbb{R}\), define:
\[
\textnormal{div} f \coloneqq \sum_{i=1}^{d} \frac{\partial f_{i}}{\partial x_{i}}, \quad \Delta \tilde{f} \coloneqq \sum_{i=1}^{d} \frac{\partial^{2} \tilde{f}}{\partial x_{i}^{2}}.
\]
\begin{proposition}[Theorem 2, \cite{liutkus2019sliced}]
\label{proposition: flow}
Assume that \(\lambda > 0\), \(R > \sqrt{d}\), and let \(\mu_{0}\) be a probability measure in \(\mathcal{P}(\widebar{B}_{R}(0))\) with density \(\varrho_{0} \in L^{\infty}(\widebar{B}_{R}(0))\). Then, for the minimization problem:
\begin{align*}
\tilde{\mathcal{F}}_{\zeta}^{\lambda}(\mu) = \frac{1}{2}SW_{2}^{2}(\mu, \zeta) + \lambda \mathcal{H}(\mu),
\end{align*}
where \(\zeta\) is a probability measure with positive smooth density, there exists a generalized minimizing movement scheme such that the density \((\varrho_{t})_{t}\) of \((\mu_{t})_{t}\) satisfies the continuity equation:
\begin{align}\label{eq: continuity}
\frac{\partial \varrho_{t}}{\partial t} = -\textnormal{div}(v_{t} \varrho_{t}) + \lambda \Delta \varrho_{t},
\end{align}
in the weak sense \citep[Page 123]{santambrogio2015optimal}, where \(v_{t}\) is associated with the Kantorovich potential \(\Upsilon_{t}^{V}\) \citep[Page 13]{santambrogio2015optimal} between \(\mu_{t}^{V}\) and \(\zeta^{V}\), given by:
\begin{align*}
v_{t}(x) = v(x, \mu_{t}) = -\int_{\mathbb{S}^{d-1}} (\Upsilon_{t}^{V})'\,(\langle x, V \rangle)\,V \dd \sigma(V).
\end{align*}
\end{proposition}
\begin{remark}
The regularization functional \(\mathcal{H}(\mu)\) plays a crucial role in generative modeling, as it is believed to mitigate overfitting. In contrast, when \(\lambda = 0\), the problem reduces to the original estimation task, with the corresponding continuity equation provided in, e.g., \citet[Theorem 5.6.1]{bonnotte2013unidimensional}. Furthermore, by defining \(\mathcal{H}(\mu)\), we restrict our focus to those probability measures \(\mu\) that are absolutely continuous with respect to the Lebesgue measure. While this excludes the discrete cases considered in Section~\ref{section: theory}, it aligns with the smoothness requirements of optimal transport maps \citep[Chapter 5.4.2]{bonnotte2013unidimensional}.
\end{remark}
Equation \eqref{eq: continuity} resembles the Fokker-Planck equation \citep[Theorem 2.2]{pavliotis2014stochastic} for the stochastic differential equation
\begin{align}\label{eq: sde}
\dd Q_{t} = v(Q_{t}, \mu_{t})\, \dd t + \sqrt{2\lambda}\,\dd W_{t},
\end{align}
where \((W_{t})_{t}\) is a standard Brownian motion. We then propose Algorithm \ref{algorithm: flow} to approximate \eqref{eq: sde} using particle systems \citep[Algorithm 1]{liutkus2019sliced} and Euler-Maruyama discretization \citep[Section 5.2]{pavliotis2014stochastic}. For \(L\) particles at step \(\ell\), the iteration is given by:
\begin{align}\label{eq: particle}
Q_{(\ell+1)h}^{i} = Q_{\ell h}^{i} + v(Q_{\ell h}^{i}, \tilde{\mu}_{\ell h}^{L})\,h + \sqrt{2\lambda h}\,Z_{\ell h}^{i}, \quad 1 \leq i \leq L,
\end{align}
where \(Z_{\ell h}^{i}\) is standard Gaussian noise. The drift \(v(Q_{\ell h}^{i}, \tilde{\mu}_{\ell h}^{L})\) is computed as:
\begin{align}\label{eq: drift}
v(Q_{\ell h}^{i}, \tilde{\mu}_{\ell h}^{L}) = -\frac{1}{m}\sum_{j=1}^{m} \left(\langle Q_{\ell h}^{i}, V_{j} \rangle - F_{\widebar{\mu}^{V}_{k}}^{-1} \circ F_{\tilde{\varrho}_{\ell h}^{V}}(\langle Q_{\ell h}^{i}, V_{j} \rangle)\right)V_{j},
\end{align}
where \(F_{\tilde{\varrho}_{\ell h}^{V}}\) is the empirical CDF of \((\tilde{\mu}_{\ell h}^{L})^{V}\), and \(F_{\widebar{\mu}^{V}_{k}}^{-1}\) is the empirical quantile function for \(\widebar{\mu}^{V}_{k}\). The algorithm outputs \(Q_{(t+1)h}^{i}\) as samples of the unobserved coefficient. If \(m\) scales polynomially with \(d\), the algorithm has polynomial complexity in \(d\).
\begin{algorithm}[!ht]
\caption{Diffusion-based Generative Modeling}
\label{algorithm: flow}
\LinesNotNumbered
\KwInput{$(Y_{i}, \beta_{i}, X_{i})$, $V_{j}$, $L$, $R$, $\lambda$, $h$, $t$, $i \in [n]$, $ j \in [m]$}
\KwInitialization{$\psi_{1} \sim \text{Uniform}\,(\widebar{B}_{R}(0)^{L})$, $D\in \mathbb{R}^{k \times m}$}
\For{$j = 1$ to $m$}
{Choose the $k$-NNs of $V_{j}$ as $\{\tilde{X}_{j(q)}; q\in[k]\}$ with ties broken arbitrarily\\
$D_{\cdot,\, j} \leftarrow \text{sort}\,(\, \tilde{Y}_{j(q)}\,\langle\, V_{j},\,\tilde{X}_{j(q)}\,\rangle\,)$ for each $q\in[k]$
}
Construct $F_{\widebar{\mu}^{V}_{k}}^{-1}$ with $D$\\
\For{$\ell = 1$ to $t$}
{Construct $F_{\tilde{\varrho}_{\ell h}^{V}}$ with $\psi_{\ell}$\\
\For{$i = 1$ to $L$}
{$v(\psi_{\ell, i}, \tilde{\mu}_{\ell h}^{L}) \leftarrow -\frac{1}{m}\sum_{j = 1}^{m} (\, \langle \psi_{\ell, i}, V_{j}\rangle - F_{\widebar{\mu}^{V}_{k}}^{-1} \circ F_{\tilde{\varrho}_{\ell h}^{V}}(\langle \psi_{\ell, i}, V_{j} \rangle)\,)\,V_{j}$\\
$\psi_{\ell+1, i} \leftarrow \psi_{\ell, i}+ v(\psi_{\ell, i}, \tilde{\mu}_{\ell h}^{L})\,h+ \sqrt{2\lambda h}\,Z_{\ell h}^{i}$
}
}
\KWReturn{$\psi_{t+1}$}
\end{algorithm}
\section{Application to causal inference}\label{sec: causal}
In this section, we connect the model \eqref{eq:linear-RCM} to the causal random coefficient model introduced in \citet[Section 7]{heckman1997making}:
\begin{align}\label{eq:causal-RCM}
Y(w) = \langle Z, \beta_Z \rangle + w \cdot R + U, ~~~\text{where } (\beta_Z,R,U) \text{ is independent of } Z.
\end{align}
Here $Z \in \mathbb{R}^d$ represents pretreatment covariates, $w \in \mathcal{W} \subset \mathbb{R}$ is the (continuous or discrete) treatment, $U$ is an unspecified noise term, and $\{Y(w); w \in \mathcal{W} \}$ denotes the potential outcomes corresponding to different treatment levels. It follows that
\[
R = Y(w+1) - Y(w), \quad \text{for any } w \in \mathcal{W}.
\]
Thus, $R$ denotes the unit-level treatment effect. In the special case where $\mathcal{W} = \{0,1\}$, we obtain
\[
R = Y(1) - Y(0),
\]
which represents the unit-level treatment effect of receiving the treatment compared to control.
Notably speaking, we have made two key assumptions in Model \eqref{eq:causal-RCM}. First, we have assumed that the untreated potential outcome satisfies the random coefficient structure: $ Y(0) = \langle Z, \beta_Z \rangle + U$.
Second, we have assumed that the unit-level treatment effect $R$ is independent of the observed covariates $Z$, though, it may be correlated with the unobserved noise $U$.
In practice, however, not all potential outcomes $\{Y(w); w \in \mathcal{W} \}$ are observable. In the context of a {\it randomized controlled trial}, we introduce a random treatment variable $W \in \mathcal{W}$ and only observe the outcome $Y : = Y(W)$, which corresponds to the potential outcome under the assigned treatment. The observed outcome follows the model:
\begin{align}\label{eq:causal-RCM-1}
Y(W) = \langle Z, \beta_Z \rangle + W \cdot R + U, \quad \text{where } W \text{ is independent of } (Z,U,\beta_Z,R).
\end{align}
Clearly, Model \eqref{eq:causal-RCM-1} is a special case of \eqref{eq:linear-RCM}, with
\begin{align}\label{eq:causal-notation}
\beta = (\beta_Z, R, U), \quad X = (Z, W, 1).
\end{align}
This facilitates the estimation of $R$'s distribution in view of the results established in previous sections.
However, to make the arguments formal, we still need to verify Assumption \ref{assumption: unbounded covariates}, which is the focus of the following results.
\begin{assumption}\label{assumption: causal tail}
There exist constants $\kappa > 1$ and $C_{Z} > 0$ such that the Lebesgue density of $Z$, denoted by $f_Z$, satisfies
\begin{align*}
f_Z(T) \geq C_{Z}(1 + \Vert T \Vert_{2})^{-\kappa}, \quad \text{for all } T \in \mathbb{R}^d.
\end{align*}
\end{assumption}
The following proposition addresses the case of a continuous treatment $W$.
\begin{proposition}\label{prop: causal_heavy_tail}
Assume Assumption \ref{assumption: causal tail}, and further assume that the Lebesgue density of $W$, denoted by $f_W$, satisfies
\[
f_W(t) \geq C_{W}(1 + |t|)^{-\kappa}, \quad \text{for any } t \in \mathbb{R}.
\]
Then, the random vector $X = (Z, W, 1)$ defined in \eqref{eq:causal-notation} satisfies
\[
f_{(Z, W)}(T, t)\geq C_{Z}C_{W} (1+\Vert(T, t)\Vert_{2})^{-2\kappa}.
\]
Thus, Assumption \ref{assumption: unbounded covariates} holds with $\alpha = 1/(2\kappa + 1)$, $\tau_{0} = \mathcal{S}(\mathbb{S}^{d-1})C_{Z}C_{W}/2^{6\kappa + 3}$, and $\rho_{0} = 1/4$, as established in Proposition \ref{prop:heavy-tail}.
\end{proposition}
This indicates that the Assumption \ref{assumption: unbounded covariates} holds for the case of continuous and unbounded treatment, which aligns with the support assumption regarding the identification in \citet[Theorem 1]{heckman1997making}.
However, assuming a continuous and unbounded treatment may be restrictive in practice. To address this, we introduce a regularized \textit{working model}, analogous to ridge regression in linear models and entropic regularization in optimal transport:
\begin{align}\label{eq:working-model}
\text{(Regularized Working Model)} \quad Y_{\epsilon}(W) = \langle Z, \beta_Z \rangle + R \cdot W_{\epsilon} + U,
\end{align}
where
\[
W_{\epsilon} = W + \tilde{W}_{\epsilon}, \quad \tilde{W}_{\epsilon} \sim \text{Cauchy}(0, \epsilon) \text{ and independent of the system}.
\]
As $\epsilon \to 0$, the working model \eqref{eq:working-model} converges to the original causal random coefficient model \eqref{eq:causal-RCM-1}, ensuring that the estimated distribution of $R$ remains approximately unchanged. Meanwhile, the additional noise term $\tilde{W}_{\epsilon}$ enhances model flexibility.
The following proposition formalizes this idea.
\begin{proposition}\label{prop: causal_heavy_tail_2}
Assume Assumption \ref{assumption: causal tail}, $\vert W \vert \leq M_{W}$ for constant $M_{W}$, and density $f_{W_{\epsilon}}$ exists. Then, for any given $\epsilon > 0$, the random vector $X_{\epsilon} = (Z, W_{\epsilon}, 1)$ satisfies
\[
f_{(Z, W_{\epsilon})}(T, t) \geq C_{Z} \pi^{-1}\epsilon\, \max\{1, \epsilon + M_{W}\}^{-2}(1+\Vert (T, t)\Vert_{2})^{-2\max\{2, \kappa\}}.
\]
Therefore, with $\alpha = 1/(2\max\{2, \kappa\} + 1)$, $\tau_{0} = \mathcal{S}(\mathbb{S}^{d-1})C_{Z}\epsilon / \pi \max\{1, \epsilon + M_{W}\}^{2}2^{6\max\{2, \kappa\} + 3}$, and $\rho_{0} = 1/4$, Assumption \ref{assumption: unbounded covariates} holds as established in Proposition \ref{prop:heavy-tail}.
\end{proposition}
Proofs are provided in Section \ref{subsection: proofs of causal inference} and empirical results for this causal approach are put in Section \ref{subsection: causal rcm}.
\section{Numerical experiments}\label{section: simulations}
\subsection{Simulation studies}
This section evaluates the performance of the algorithms using the folllowing numerical simulations. In the following, covariates are drawn from the von Mises-Fisher distribution with a concentration parameter $0.1$ and a mean direction $(d^{-1/2},\dots,d^{-1/2})$, satisfying Assumption \ref{assumption: bounded covariates}. We consider three types of underlying uniform distributions for the coefficient samples:
\[
\mu_{\text{sph}},~~ \mu_{\text{deg}},~~ {\rm and}~~ \mu_{\text{dis}},
\]
with supports of different dimensions and $R = 10$ as
\begin{gather*}
\text{supp} \,\mu_{\text{sph}} = \Big\{\, U\in\mathbb{R}^{d}~\,\vert\, ~\Vert U - \delta R/2\, e_{2}\Vert_{2}\leq R/4,\, \delta \in \{1, -1\}\,\Big\},\\
\text{supp}\, \mu_{\text{deg}} = \Big\{\, U\in\mathbb{R}^{d}~\,\vert\, ~\Vert U - \delta R/2\, e_{2}\Vert_{2}= R/4, \, \delta \in \{1, -1\}\,\Big\},\\
\text{supp}\, \mu_{\text{dis}} = \Big\{\,U \in \mathbb{R}^{d}~\,\vert\, ~U = \delta R/2 \, e_{i},\, 1\leq i \leq d,\, \delta \in \{1, -1\}\,\Big\}.
\end{gather*}
The simulations are performed on an 11th Gen Intel Core i7-11700 @ 2.50GHz CPU, and all reproducible codes are available at \url{https://github.com/keunwoolim/rcm_sw}.
The performance of the algorithms is measured by computation time (in seconds) and the empirical distance (in $SW_{2}$ distance) between the original samples and the output, averaged over $100$ experiments. The empirical $SW_{2}$ distance using the Monte Carlo method is computed with the POT package \citep{flamary2021pot} with $100$ unit vector samples.
For implementing Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd}, we perform $20$ iterations, setting $m = 50$ and $m = 1000$ respectively. We set $k = \lceil{n^{d/(2d-1)}} \rceil$, which corresponds to the general $d\geq 6$ case in Theorem \ref{theorem: bounded estimator} and is empirically observed to be more robust than the suggested diverging rates in Theorem \ref{theorem: bounded estimator}. For Algorithm \ref{algorithm: flow}, we set $L = k$, $m = 50$, $t = 20$, $h = 1$, and $\lambda = 0.01$. The convex optimization problem for Algorithm \ref{algorithm: bcd} is solved using the CVXPY package \citep{diamond2016cvxpy}.
Table \ref{table: simulation_dimension} reports the results of Algorithms \ref{algorithm: bcd}-\ref{algorithm: flow} for dimensions $2 \leq d \leq 5$ with $n = 500$ samples from $\mu_{\text{sph}}$. This demonstrates the efficiency of Algorithm \ref{algorithm: abcd}, as it is averagely $2.1$ times faster than Algorithm \ref{algorithm: bcd} while maintaining similar empirical accuracy. The approximation accuracy of Algorithm \ref{algorithm: flow} is lower but not substantially worse than the first two algorithms. This justifies its validity and also aligns with the purpose of generative modeling.
\begin{table}[!ht] \caption{Average performance of Algorithms \ref{algorithm: bcd} - \ref{algorithm: flow} with regard to the dimension}
\centering
\begin{tabular}{ccccccc}
\toprule
\multirow{2}{*}{Dimension} &
\multicolumn{2}{c}{Algorithm \ref{algorithm: bcd}} &
\multicolumn{2}{c}{Algorithm \ref{algorithm: abcd}} & \multicolumn{2}{c}{Algorithm \ref{algorithm: flow}} \\
& {Distance} & {Time} & {Distance} & {Time} & {Distance} & {Time} \\
\midrule
2 & 0.697 & 4.344 & 0.690 & 2.008 & 0.717 & 5.574\\
3 & 0.762 & 2.935 & 0.752 & 1.366 & 0.808 & 2.580\\
4 & 0.832 & 2.512 & 0.787 & 1.156 & 0.937 & 1.808\\
5 & 0.943 & 2.271 & 0.910 & 1.061 & 1.057 & 1.528\\
\bottomrule
\end{tabular}
\label{table: simulation_dimension}
\end{table}
Table \ref{table: simulation_sample size} reports the performance of Algorithm \ref{algorithm: abcd} for $n = 500, 1000, 1500, 2000$ samples from $\mu_{\text{sph}}$, $\mu_{\text{deg}}$, and $\mu_{\text{dis}}$, with $d = 2$. For all types of probability measures, the average empirical distance decreases as the sample size increases, while maintaining empirically robust computation time.
\begin{table}[!ht] \caption{Average performance of Algorithm \ref{algorithm: abcd} with regard to the sample size}
\centering
\begin{tabular}{ccccccc}
\toprule
\multirow{2}{*}{Sample Size} &
\multicolumn{2}{c}{$\mu_{\text{sph}}$} &
\multicolumn{2}{c}{$\mu_{\text{deg}}$} &
\multicolumn{2}{c}{$\mu_{\text{dis}}$} \\
& {Distance} & {Time} & {Distance} & {Time} & {Distance} & {Time} \\
\midrule
500 & 0.690 & 2.008 & 0.639 & 1.998 & 0.877 & 2.006 \\
1000 & 0.543 & 3.162 & 0.462 & 3.218 & 0.805 & 3.170 \\
1500 & 0.440 & 4.250 & 0.424 & 4.256 & 0.777 & 4.262 \\
2000 & 0.375 & 4.982 & 0.373 & 5.035 & 0.725 & 5.060 \\
\bottomrule
\end{tabular}
\label{table: simulation_sample size}
\end{table}
Lastly, Figure \ref{figure: introduction} plots the coefficient samples and outputs (i.e., $\hat w$'s introduced in \eqref{eq:points}) of Algorithms \ref{algorithm: bcd}-\ref{algorithm: flow} for $100$ experiments. For Algorithms \ref{algorithm: bcd} and \ref{algorithm: abcd}, we performed $20$ iterations with $n = 2000$ and $k = 159$, setting $m = 50$ and $m = 1000$ respectively. For Algorithm \ref{algorithm: flow}, the parameters are $L = 159$, $m = 10$, $t = 20$, $h = 1$, and $\lambda = 0.01$.
\begin{figure}[h]
\centering\hspace{-0cm}
\includegraphics[scale=0.6]{images/introduction.png}
\caption{Coefficient samples and outputs of Algorithms \ref{algorithm: bcd}-\ref{algorithm: flow}.}
\label{figure: introduction}
\end{figure}
Figure \ref{figure: generation} plots the outputs of $100$ experiments of Algorithm \ref{algorithm: flow} depending on $L$ and $\lambda$, with fixed parameters $n = 2000$, $m = 10$, $t = 20$, and $h = 1$. In Figure \ref{figure: generation_particles}, $\lambda$ is fixed as $0.01$, and $L$ varies as $10$, $20$, $40$, and $80$. Meanwhile, in Figure \ref{figure: generation_regularization}, $L$ is fixed as $20$, and $\lambda$ varies as $0.01$, $0.02$, $0.04$, and $0.08$. Intuitively, the particle systems approximation becomes more precise as $L$ increases, and the generated samples become more diverse as $\lambda$ increases.
\begin{figure}\hspace{-0.5cm}
\begin{subfigure}[h]{0.5\textwidth}
\centering
\includegraphics[width=\linewidth]{images/particles.png}
\caption{Varying number $L$ of particles.}
\label{figure: generation_particles}
\end{subfigure}
\hfill
\begin{subfigure}[h]{0.5\textwidth}
\centering
\includegraphics[width=\linewidth]{images/regularization.png}
\caption{Varying regularization parameter $\lambda$.}
\label{figure: generation_regularization}
\end{subfigure}
\caption{Outputs of Algorithm \ref{algorithm: flow} with varying numbers of particles and regularization parameter.}
\label{figure: generation}
\end{figure}
\subsection{Causal RCM}\label{subsection: causal rcm}
Based on the regularized working model \eqref{eq:working-model} for causal inference, we analyze the AIDS Clinical Trials Group Study 175 (ACTG 175) data provided in the R \texttt{speff2trial} package \citep{speff2trial} to evaluate the treatment effects of therapy for human immunodeficiency virus type I. Specifically, using the monotherapy of zidovudine (Therapy 1) as the control group, we examine the effects of combination therapies—zidovudine and didanosine (Therapy 2) and zidovudine and zalcitabine (Therapy 3)—as well as monotherapy with didanosine (Therapy 4), as considered in \citet[Table 2]{hammer1996trial}.
We analyze each investigational treatment (Therapies 2-4) separately against the control. In each analysis, the random treatment variable $W$ is assigned as $W = 0$ for the control group and $W = 1$ for the treatment group. The pretreatment covariates $Z$ include the baseline CD4 T cell count and the number of days of prior antiretroviral therapy, analogous to the setting in \citet[Section 6]{10.1093/biomet/asad045}. After standard normalization, this design produces three sets of covariate samples, each with $d = 4$ and respective sample sizes of $n = 1054, 1056, 1093$ for Therapies 2–4. The outcome variable $Y$ is the CD4 T cell count measured at $20 \pm 5$ weeks, which is also normalized.
We implemented a modified version of Algorithm \ref{algorithm: abcd} and conducted 100 experiments for each of the three cases. The parameter $k = \lceil{n^{d/(2d-1)}} \rceil$ was computed as $k = 54, 54, 55$, respectively, and additional parameters were set as $\epsilon = 0.005$, $R = 10$, and $m = 1000$, with 20 iterations.
\begin{figure}[h]
\centering
\includegraphics[scale=0.72]{images/Rplot02.png}
\caption{Treatment coefficient samples for Therapies 2–4.}
\label{figure: causal}
\end{figure}
Figure \ref{figure: causal} presents the box plots for the estimated treatment coefficients corresponding to Therapies 2–4. The estimation results also show the estimated percentiles of the unit treatment effect. For example, the estimated median treatment effects for Therapies 2-4 are $0.113$, $0.047$, and $0.045$, respectively. These results suggest that all three therapies are effective compared to zidovudine monotherapy, with the combination of zidovudine and didanosine showing particularly notable effectiveness. These findings are generally consistent with the results of \cite{hammer1996trial}.