# Partially Specified Causal Simulations

Alireza Zamanian
Technical University of Munich,
School of Computer, Information and Technology;
Fraunhofer IKS, Germany
<EMAIL_ADDRESS>
&Leopold Mareis
Technical University of Munich,
School of Computer, Information and Technology;
Fraunhofer IKS, Germany
<EMAIL_ADDRESS>
&Narges Ahmidi
Fraunhofer IKS, Germany
<EMAIL_ADDRESS>

(12 September 2023)

###### Abstract

Simulation studies play a key role in the validation of causal inference
methods. The simulation results are reliable only if the study is designed
according to the promised operational conditions of the method-in-test. Still,
many causal inference literature tend to design over-restricted or
misspecified studies. In this paper, we elaborate on the problem of improper
simulation design for causal methods and compile a list of desiderata for an
effective simulation framework. We then introduce partially randomized causal
simulation (PARCS), a simulation framework that meets those desiderata. PARCS
synthesizes data based on graphical causal models and a wide range of
adjustable parameters. There is a legible mapping from usual causal
assumptions to the parameters, thus, users can identify and specify the subset
of related parameters and randomize the remaining ones to generate a range of
complying data-generating processes for their causal method. The result is a
more comprehensive and inclusive empirical investigation for causal claims.
Using PARCS, we reproduce and extend the simulation studies of two well-known
causal discovery and missing data analysis papers to emphasize the necessity
of a proper simulation design. Our results show that those papers would have
improved and extended the findings, had they used PARCS for simulation. The
framework is implemented as a Python package, too. By discussing the
comprehensiveness and transparency of PARCS, we encourage causal inference
researchers to utilize it as a standard tool for future works.

_Keywords_ Simulation $\cdot$ Structural Causal Models $\cdot$ Causal DAGs
$\cdot$ Causal Inference $\cdot$ AI Validation

##  1 Introduction

Causal inference and discovery (CI), a crucial sub-field of statistics and
artificial intelligence, aims to unravel cause-and-effect relationships from
observational data. CI methods tackle a wide array of research questions of
not associational, but causal nature, such as estimating the causal effect of
a variable on another in a complex system, or discovering the underlying
causal connections among study variables (See Pearl (2009)). These methods
cannot be validated using observational data sets due to the unavailability of
the ground truth. To mitigate this problem, simulation studies are widely used
in the literature (Sofrygin et al., 2017; Reisach et al., 2021), where the
ground truth is manually generated by the researchers and, therefore,
accessible. At the core of simulation lies a data-generating process (DGP)
comprising a set of equations and sampling procedures for variables, which
determines the exact distribution of the synthesized data. Clearly, to
generate reliable validation results, a DGP must be designed according to the
promised operational conditions of the CI method, namely the mathematical
assumptions. See Figure 1 for examples of method assumptions and DGPs.

Method 1 assumptions • $Z_{1}$ follows the standard normal distribution. •
$Z_{2}$ is a linear function of $Z_{1}$ and Gaussian additive noise. • Define
$Z_{3}=Z_{1}^{2}$ deterministically. Method 2 assumptions • $(Z_{1},Z_{2})$
follows a zero-centered bivariate normal distribution. • $Z_{3}$ is a non-
linear function of $Z_{1},Z_{2}$. Method 3 assumptions • $Z_{3}$ is a function
of $Z_{2}$. • $Z_{2}$ and $Z_{3}$ have at least two common causes. • The
sparsity of relation among the nodes is $0.5$. restricted parameters •
$\mu_{Z_{1}},\sigma_{Z_{1}}=0,1$ • $\beta=(1,0,0,0)$ free parameters •
$(\alpha_{0},\alpha_{1})\stackrel{{\scriptstyle\text{i.i.d}}}{{\sim}}U(-5,5)$
• $\sigma_{\varepsilon}\sim U(1,3)$ restricted parameters •
$\mu_{Z_{1}}=\alpha_{0}=0$ free parameters • $\alpha_{1}\sim U(-5,5)$ •
$\sigma_{\varepsilon}\sim U(1,3)$ •
$\beta\stackrel{{\scriptstyle\text{i.i.d.}}}{{\sim}}U(-10,10)$ restricted
parameters ?  free parameters ? Meta-algorithm 1. sample the free parameters N
times; 2. per sample, parameterize the following DGP: •
$Z_{1}\sim\mathcal{N}(\mu_{Z_{1}},\sigma_{Z_{1}})$ •
$\varepsilon\sim\mathcal{N}(0,\sigma_{\varepsilon})$ •
$Z_{2}:=\alpha_{1}Z_{1}+\alpha_{0}+\varepsilon$ •
$Z_{3}:=\beta_{1}Z_{1}^{2}+\beta_{2}Z_{1}Z_{2}+\beta_{3}Z_{2}^{2}+\beta_{0}$
Method 1 DGPs DGP 1 • $Z_{1}\sim\mathcal{N}(0,1)$ •
$\varepsilon\sim\mathcal{N}(0,\mathbf{1.3})$ •
$Z_{2}:=\mathbf{-3}Z_{1}+\mathbf{2.1}+\varepsilon$ • $Z_{3}:=Z_{1}^{2}$ DGP 2
• $Z_{1}\sim\mathcal{N}(0,1)$ • $\varepsilon\sim\mathcal{N}(0,\mathbf{2.9})$ •
$Z_{2}:=\mathbf{0.8}Z_{1}\mathbf{-2}+\varepsilon$ • $Z_{3}:=Z_{1}^{2}$ ⋮ DGP N
Method 2 DGPs DGP 1 • $Z_{1}\sim\mathcal{N}(0,\mathbf{1.7})$ •
$\sigma_{\varepsilon}\sim(0,\mathbf{2.5})$ •
$Z_{2}:=\mathbf{0.2}Z_{1}+\sigma_{\varepsilon}$ •
$Z_{3}:=\mathbf{1.7}Z_{1}^{2}\mathbf{-4}Z_{1}Z_{2}+\mathbf{0.6}Z_{2}^{2}+\mathbf{3.1}$
DGP 1 • $Z_{1}\sim\mathcal{N}(0,\mathbf{3})$ •
$\sigma_{\varepsilon}\sim(0,\mathbf{1.9})$ •
$Z_{2}:=\mathbf{-4.7}Z_{1}+\sigma_{\varepsilon}$ •
$Z_{3}:=\mathbf{-3}Z_{1}^{2}+\mathbf{2.4}Z_{1}Z_{2}+\mathbf{0}Z_{2}^{2}+\mathbf{1}$
⋮ DGP N Figure 1: Schematic of 3 methods for validation (top row), a meta-
algorithm (middle row) and generated DGPs for methods 1 and 2 (bottom row);
The $Z$ variables are simulated study variables. Meta-parameters are:
$\mu_{Z_{1}},\sigma_{Z_{1}},\sigma_{\varepsilon},\alpha,\beta$. Restricted and
free parameters are specified below each method. The bold face values in DGPs
refer to the free parameters which are sampled. The meta-algorithm design does
not suit the assumptions of method 3.

Generally, an assumption may lead to many DGPs that induce different data
distributions. For instance, the “non-linearity” assumption for the variable
relations (e.g., assumption over the variable $Z_{3}$ in Figure 1-Method 2)
can take various forms, such as in second-order polynomial or trigonometric
functions. For a complete validation, the method must be tested with all
complying DGPs. The question is, how do we identify all complying DGPs? One
solution is to devise a parametric algorithm (called “meta-algorithm” for
brevity), with its internal parameters (“meta-parameters”) that outputs one
DGP per each meta-parameter combination. If there is a clear mapping from the
pool of assumptions to the meta-parameters (as is for Figure 1-Method 1), we
identify and fix the subset of meta-parameters that relate to the assumptions
(restricted parameters). Then we sample the remaining ones (free parameters) N
times to generate N complying DGPs. This approach is adopted in the majority
of CI literature for method validation; as notable attempts to design
comprehensive meta-algorithms, we name the Simcausal R package (Sofrygin et
al., 2017), the simulation design for the 2017 Atlantic causal inference
conference (ACIC) data analysis callenge (Hahn et al., 2019) and time series
causal simulation framework by Lawrence et al. (2021).

Designing a meta-algorithm for CI methods is a nontrivial and demanding task:
firstly, meta-parameters must be relevant to a wide range of common
assumptions in the field of application; otherwise, they cannot be adjusted
accordingly. For example, in Figure 1, it is unclear how the meta-parameters
can be adjusted to satisfy the assumptions of Method 3 about common causes.
Secondly, the meta-parameters must be defined such that they can be fixed or
varied independently; otherwise, the distinction between restricted and free
parameters is difficult, if not impossible. For instance, the sparsity
assumption in Figure 1-Method 3 requires the function coefficients such as
$\alpha$ and $\beta$ to be adjusted only jointly. Lastly, the design must be
comprehensive and inclusive so that the meta-algorithm itself does not cause a
harmful restriction on the validation results. For example, the meta-algorithm
in Figure 1 is improperly restricted to only three variables, and the non-
linearity of the relations is only defined for one variable ($Z_{3}$) and only
via the second-order polynomial form.

Scanning the literature, we noticed that simulation studies often suffer from
improperly designed meta-algorithms. This causes the generated DGPs to be
considerably limited to few specific scenarios and do not cover the entire
operational space of the method. Consequently, the validation results become
incomplete, unreliable, and, in extreme cases, misleading.

In this paper, we gather a list of common causal assumptions found in
literature and compile a set of desiderata for a potentially effective
simulation framework. We then introduce PARCS (partially randomized causal
simulation), a simulation meta-algorithm aimed to satisfy those desiderata and
address the aforementioned limitations. We explain how PARCS relates to the
common causal assumptions, controls disentangled and independent simulation
parameters, and provides comprehensive complying DGPs for various collections
of assumptions. We believe PARCS allows researchers to overcome the issues of
simulation design, run more in-depth analyses of their methods, and report
their simulation setup more transparently and reproducibly. PARCS is
implemented and available as a Python package111The Python package is
available at https://pypi.org/project/pyparcs/. A detailed documentation is
available at https://fraunhoferiks.github.io/parcs/ for researchers.

The main contributions of this paper are listed below.

  1. 1.

We present the causal methods simulation desiderata based on the common
simulation studies found in the CI literature.

  2. 2.

We introduce PARCS simulation meta-algorithm. We show how common causal
assumptions map to PARCS parameters. We present the governing equations,
design choices, and theoretical details for the simulation process.

  3. 3.

We demonstrate the utility of PARCS by reproducing and extending the
simulation studies done by two SoA CI papers by Jarrett et al. (2022) and
Shimizu et al. (2006). Throughout experiments, we show how PARCS generates
additional CI validation scenarios which were untested and ignored by the
papers, and thus allows for more extensive and precise validation.

The rest of the paper is organized as follows: A review of common causal
assumptions in literature is presented in Section 2. Section 3 provides the
theoretical background and explains the design of PARCS. Simulation case
studies can be found in Section 4. We revisit the planned simulation
desiderata and present the conclusion in Section 5.

##  2 Simulation desiderata

Let the random vector $Z\in\mathbb{R}^{d}$ represent the study variables in a
data set. The joint density of $(Z_{1},\dots,Z_{d})$ can be factorized as a
product of one-dimensional conditional densities, as

$\displaystyle P(z)=\prod\nolimits_{i}P(z_{\gamma(i)}\;|\;s^{\gamma}(i)),\quad
s^{\gamma}(i)=\\{z_{\gamma(1)},\dots,z_{\gamma(i-1)}\\}$ (1)

for any permutation $\gamma$ over the components $\\{1,\dots,d\\}$ where
$S^{\gamma}(i)$ is the set of all variables preceding $i$ in the $\gamma$
ordering. Furthermore, we define intervention on a component $Z_{j}\in Z$ as
altering the conditional density $P(Z_{j}|.)$ in Equation 1 to $Q(Z_{j}|.)$,
where $Q$ represents the probability distributions after intervention. In a
basic case, $Q(Z_{j}|.)$ represents a degenerate distribution; this is equal
to setting $Z_{j}$ to a fixed value $Z_{j}=z^{*}_{j}$. It is clear that after
intervention, the joint density over $Z$ (and perhaps some of the one-
dimensional conditional densities) changes, i.e., $P(z)\neq Q(z)$.

Equation 1 offers a simulation procedure, by which we specify and sequentially
sample from the conditional densities along the order of $\gamma$. Although
this procedure is applicable to any permutation, we are interested in
permutations with a specific causal property; in particular, a permutation
$\gamma^{c}$, under which interventions on $Z_{\gamma^{c}(j)}$ does not change
the distributions of the upstream nodes $S^{\gamma^{c}}(j)$. As a result, the
corresponding distribution for an intervention defined by $Q$ is given as

$\displaystyle
Q(z)=\prod_{i=1}^{j-1}P(z_{\gamma^{c}(i)}|s^{\gamma^{c}}(i))\cdot\prod_{i=j}^{d}Q(z_{\gamma^{c}(i)}|s^{\gamma^{c}}(i))\cdot$
(2)

In terms of simulation procedure, Equation 2 reads as “sample for the upstream
variables from the pre-intervention, and for the intervened and the downstream
variables from the post-intervention densities”.

This leads us to the definition of a causal directed acyclic graph (DAG): a
DAG $\mathcal{G}$ is defined over a node set ${N=\\{Z_{1},\dots,Z_{d}\\}}$
corresponding to the study variables, and a directed edge set
${E\subseteq\\{(i_{1},i_{2})\in N^{2}:i_{1}\in S^{\gamma^{c}}(i_{2})\\}}$.
Each edge represents a causal effect from one node (parent) to another
(child). By definition, only the parent nodes affect the child node. Denoting
the set of all parent nodes for $Z_{i}$ in a DAG by
$\mathrm{pa}_{\mathcal{G}}(i)$, we rewrite Equation 1 as
$P(z)=\prod_{i=1}^{d}P(z_{i}\;|\;\mathrm{pa}_{\mathcal{G}}(i))$. This means
that given the parent set, a node is independent of the rest of the
graph222known as the Markov property of DAGs.

By the definition of causal DAG, the causal ordering $\gamma^{c}$ is the
topological sorting in DAG. This implies
${\mathrm{pa}_{\mathcal{G}}(i)\subseteq S^{\gamma^{c}}(i)}$, namely, the
parents of a node precede it in the sorting. To sample from a DAG, based on
the general sampling procedure for Equation 1, we sample the nodes along its
topological sorting, and use the realizations at each step as the inputs of
the next step. An illustrative example of a simple DAG is depicted in 2(a). In
this DAG, $Z_{1}$ is a parent of $Z_{2}$, and $Z_{1},Z_{2}$ are parents of
$Z_{3}$. An arbitrary permutation can be $\gamma=\\{1,3,2\\}$ which leads to
the factorization $P(Z_{1})P(Z_{3}|Z_{1})P(Z_{2}|Z_{1},Z_{3})$. On the other
hand, given the causal permutation $\gamma^{c}=\\{1,2,3\\}$, the distribution
factorizes as $P(Z_{1})P(Z_{2}|Z_{1})P(Z_{3}|Z_{1},Z_{2})$. Assuming an
intervention on $Z_{3}$, the last two terms under $\gamma$, and only the last
term under $\gamma^{c}$ changes.

In order to fully specify the statistical model for a DAG, it is sufficient to
assume a working model for each conditional density. These models are defined
in form of a structural causal model (SCM) as

$\displaystyle
Z_{i}=h_{i}(Z_{\mathrm{pa}_{\mathcal{G}}(i)},\varepsilon_{i}),\quad
i\in\\{1,\dots,d\\},$ (3)

via a deterministic functions
$h_{i}(.):\mathbb{R}^{|\mathrm{pa}_{\mathcal{G}}(i)|+1}\to\mathbb{R}$ and
independent random errors $\varepsilon_{i}$. To sample from an observational
distribution as in Equation 3, we sample error terms and calculate the
variables’ realization along the topological order $\gamma$ of the DAG.

Equations 1 and 3 introduce the necessary and sufficient set of parameters
that define a causal model of DAG and SCM: the graph’s nodes $Z$ and edges
$E$, deterministic functions $h(.)$ and the distributions of the noise terms
$\varepsilon$. Causal methods make their assumptions on one or more of these
parameters, partially or fully, directly or indirectly. Table 1 outlines a
(non-exhaustive) list of assumptions which are central to the CI literature,
and their relationship to the four introduced parameters. Based on Table 1, we
present the following list of desiderata which must be statisfied by a
potential CI simulation framework in order to be effective and useful for the
CI community.

Parameter |  Assumptions |  Example literature
---|---|---
$Z$ |  number of nodes∗ |
|  expected value, hidden/observed |  Shpitser (2013); Mohan et al. (2013); Mohan and Pearl (2022)
$E$ |  DAG structure |  Spirtes and Glymour (1991); Uhler et al. (2013); Shpitser (2013); Pfister and Peters (2022); Xu et al. (2021); Mohan et al. (2013); Mohan and Pearl (2022)
|  sparsity |  Spirtes and Glymour (1991); Uhler et al. (2013); Peters and Bühlmann (2015); Pfister and Peters (2022)
|  linearity of relations∗∗ |  Spirtes and Glymour (1991); Shimizu et al. (2006); Hoyer et al. (2008),
$h$ |  governing distribution |  Hoyer et al. (2008)
|  model coefficients |  Pfister and Peters (2022)
$\varepsilon$ |  noise distribution and model |  Hoyer et al. (2008); Peters and Bühlmann (2015)
|  noise independence∗∗∗ |
Table 1: Examples of DAG-related assumptions in CI literature. The lists of
sub-categories and papers are non-exhaustive.
* The number of nodes is often not a direct assumption of the methods, but almost all the simulation studies wrongly fix the number of nodes in DGP.
** The mentioned papers explicitly put an assumption on linearity/non-
linearity of the relations.
*** Almost all CI methods make an implicit noise-independence assumption.

##### A. Underlying causality

As the first and most fundamental requirement, the framework must be able to
embody the notion of causality and govern the simulation based on causal
relations. To satisfy this requirement, it is sufficient to take the causal
DAG and SCM models (Equation 1 and 3) as the DGP.

##### B. Data dimensions and modality

The framework must provide means to control the statistical characteristics of
data such as sample size, number of variables, data types and conditional
densities. Furthermore, selecting the distribution of any variable must not be
limited by other simulation parameters.

##### C. Controlled complexity

Data complexity must be controlled in terms of number of variables, causal
relations, non-linearity of relations, higher-order polynomial terms, and
causal parents interactions of the causes of a variable. Different sources of
complexity should be independently controllable, such that the framework
allows us to induce complexity by adjusting each source as desired.

##### D. Noise models

The framework must provide means to induce various types of noise in the data.
The noise profile must be free to be affected and controlled by graph
variables.

##  3 PARCS simulation meta-algorithm

To design a meta-algorithm that satisfies Desideratum A, we assume the
underlying DGP to be a causal DAG and an accompanying SCM, and adopt the
sampling procedures defined by Equations 1-3. In addition, we employ a set of
design choices for $h$ and $\varepsilon$ to satisfy Desiderata B, C, and D.
These choices lead to disentanglement of the parameters and thus allow for
restricted and free parameters distinction. We introduce this concept in the
current section denoted by partial randomization.

###  3.1 DAG sampling procedure

Take an arbitrary node $Z_{i}$ and let the $j$ index selects the node indices
in $\mathrm{pa}(i)$. Given the samples of $Z_{\mathrm{pa}(i)}$, the sampling
procedure for $Z_{i}$ is as follows.

##### Edge functions

Samples of each $Z_{j}$ are transformed by a so-called edge function
$e_{ji}:\mathbb{R}\to\mathbb{R}$ as

$\displaystyle Z_{j\rightarrow i}=e_{ji}(Z_{j};\phi_{ji}),\quad\text{ for
}j\in\mathrm{pa}(i),$ (4)

where $\phi_{ji}$ parameterizes $e_{ji}$. Edge functions may introduce the
first source of non-linearity in the simulation process. Being reminiscent of
the generalized linear model’s link function concept, each type of functions
(such as log, logistic or Gaussian radial basis functions) induces a
particular pattern in the parent-child relation. Following the notation in
Equation 4, we denote the vector of the transformed samples of all parents (by
their corresponding edge functions) for $Z_{i}$ as
$Z_{\mathrm{pa}(i)\rightarrow i}$.

##### Input library

We assume a library $\zeta(x)$ consisting of any desired form of linear and
non-linear combinations of elements of a vector $x$. For instance,
$\zeta:x\mapsto(x_{1}^{2},x_{2}^{2},\dots,x_{d}^{2}),\;x\in\mathbb{R}^{d}$ is
a library that returns the 2nd-order polynomials. We apply this library to the
edge-transformed samples of the parents, i.e. $Z_{\mathrm{pa}(i)\rightarrow
i}$, to obtain an input vector for $Z_{i}$ (hence the name, input library). As
we will show in the next step, this vector is used to specify the distribution
of $Z_{i}$.

The complexity of the input library $\zeta$ is predicated on the variety of
the functions in $\zeta$, which may induce trigonometric or nth-order
polynomial terms. These transformations represent the additional source of
non-linearity in the simulation process. Including each form must be justified
by the designer of the simulation study. As a basic choice, we consider the
following input library for the rest of the paper333Equation 5a is the choice
also in the PARCS Python library, at least up to version 1.0.0. It may change
in future versions.:

$\displaystyle\zeta\;:$
$\displaystyle\;x\mapsto\left(1,\zeta_{\text{linear}}\left(x\right),\zeta_{\text{quad.}}\left(x\right)\right),\quad
x\in\mathbb{R}^{d},$ (5a) where $\displaystyle\zeta_{\text{linear}}\;:$
$\displaystyle\;x\mapsto x,$ (5b) $\displaystyle\zeta_{\text{quad.}}\;:$
$\displaystyle\;x\mapsto(x_{i}x_{j}:x_{i},x_{j}\in x).$

##### Parameterized distributions

By the probability integral transform theorem, a random variable $X$ follows a
distribution $P$ with the inverse cumulative distribution function (iCDF)
$F_{P}^{-1}$, if $X=F_{P}^{-1}(\varepsilon)$ and an error $\varepsilon$
follows a continuous uniform distribution $\text{Unif}(0,1)$. We employ this
theorem for choosing the structural equation function $h$ for the nodes. More
specifically, we choose $h_{i}(.)$ to be the iCDF of a probability
distribution $P^{i}_{\theta}$, parameterized by $\theta$. The parent nodes
affect $Z_{i}$ via controlling the distribution parameters: for the parameter
space $\Theta^{i}$, a coefficient matrix
$W^{i}\in\mathbb{R}^{|\Theta^{i}|\times|\zeta^{i}|}$ is specified and the
distribution parameters are given by

$\displaystyle\theta^{i}=W^{i}\zeta(Z_{\mathrm{pa}(i)\to i})^{\top}.$ (6)

By sampling the error term as $\varepsilon_{i}\sim\text{Unif}(0,1)$ and
calculating the $Z_{i}$ realizations as
$Z_{i}=h_{i}(\varepsilon_{i};\theta^{i})$, we have

$\displaystyle Z_{i}\;|\;Z_{\mathrm{pa}(i)}\sim P^{i}_{\theta},$ (7)

namely, conditioned on its parents, a node follows the distribution whose iCDF
is selected for $h_{i}$. The choice of iCDF results in different statistical
characteristics, data types and noise models in the synthesized data. The
values of $W^{i}$ also determine which of the distribution parameters (and
how) the parents affect in the child node. For instance, all-zero coefficients
for one of the distribution parameters means that the parents do not affect
the distribution via that parameter.

In summary, a DGP in PARCS is specified by sets of output distributions and
parameter coefficient vectors of its nodes, and transforming functions of its
edges. The sampling procedure repeats for all nodes in the topological
(causal) order of the graph. For source nodes (nodes with no parents), the
corresponding coefficient vector $W$ is a vector of length $|\theta^{i}|\times
1$, including the bias terms (considering $\zeta$ in Equation 5a). Table 2
summarizes the parameters of the PARCS simulation meta-algorithm.

Parameter | Example value | Description | Controls
---|---|---|---
$|Z|$ | 5 | number of nodes | DAG structure
$E$ | $\\{(Z_{1},Z_{2}),\dots\\}$ | edges | DAG structure, sparsity
$P^{i}_{\theta}$ | Bern, Unif, $\mathcal{N}$ , … | node distributions | data type, noise model, expectations
$\zeta$ | $(1,Z^{2}_{j\rightarrow i},\dots)$ | input library | non-linearity
$W^{i}$ | $(1,0.2,\dots)$ | coefficient vectors for $\theta$ | causal relations
$e_{ji}$ | $\arctan(Z;\alpha,\beta,\gamma)$ | edge functions for edges | non-linearity
$\phi_{ji}$ | $\alpha=1.5,\dots$ | edge function parameters | non-linearity
Table 2: PARCS meta-parameters. The last 5 parameters must be specified for
all nodes/edges.

######  Example.

2(a) illustrates a simple DAG with three nodes ($Z_{1},Z_{2},Z_{3}$): a
normally distributed node $Z_{3}$ and its two parents, $Z_{1},Z_{2}$
(following arbitrary distributions). Assume all edge functions are identity
functions. For $Z_{3}$, we have $Z_{\mathrm{pa}(3)}=(Z_{1},Z_{2})$. We
transform the parents’ samples as $Z_{1\rightarrow
3}=I(Z_{1})=Z_{1},\;Z_{2\rightarrow 3}=I(Z_{2})=Z_{2}$. The input vector
(following Equation 5a) is $\zeta(Z_{\\{1,2\\}\rightarrow
3})=\left(1,Z_{1},Z_{2},Z_{1}^{2},Z_{1}Z_{2},Z_{2}^{2}\right)^{\top}$. We
choose $P^{3}$ to be the Gaussian normal $\mathcal{N}(\mu,\sigma)$. We have
$\Theta^{3}=\mathbb{R}\times\mathbb{R}_{>0},\theta^{3}=(\mu,\sigma)\in\Theta^{3}$.
Let the coefficient matrix be $W^{3}=\left(\begin{smallmatrix}W^{3}_{\mu}\\\
W^{3}_{\sigma}\end{smallmatrix}\right)=\left(\begin{smallmatrix}1&1&1&0&0&0\\\
2&0&0&0&1&0\end{smallmatrix}\right),$ which gives the following for $Z_{3}$ in
the DGP:

$Z_{3}|z_{1},z_{2}\sim\mathcal{N}(\mu=z_{1}+z_{2}+1,\;\;\sigma=z_{1}z_{2}+2).$

For example $Z_{3}|Z_{1}=0.2,Z_{2}=0.3\sim\mathcal{N}(\mu=1.5,\sigma=2.06)$.
Note that sampling from $\mathcal{N}$ is done by sampling the error terms
uniformly and passing it to $F^{-1}_{\mathcal{N}}$, parameterised by
$\theta^{3}$. The described process is depicted in 2(b).

$Z_{1}$$Z_{2}$$Z_{3}$ (a)

$Z_{1}$$Z_{2}$$e_{13}$$e_{23}$$\zeta$$W^{3}$.$F^{-1}_{P^{3}}(\varepsilon;\theta^{3})$Unif(0,1)$Z_{3}$Sample$\theta^{3}$$\varepsilon$
(b)

Figure 2: Example of sampling procedure for a simple DAG with three nodes
$Z_{1},Z_{2},Z_{3}$: (a) the corresponding causal DAG, and (b) PARCS mechanics
to sample $Z_{3}$ values.

The sampling procedure extends to interventional distributions by a simple
adjustment according to Equation 2. Intervention on a causal model is defined
changing the conditional distribution of a node away from its natural
(observational) conditional distribution. Intervention can be done by setting
the realizations of a node to a constant value (degenerate distribution), or
changing the SCM parameters, output distribution or the parent set of the
node. Regardless of the type of intervention, we repeat the same sampling
procedure as above, along the causal ordering of the DAG, until we reach a
node (or nodes) $Z_{i}$ on which we intervene. Then, the normal sampling is
interrupted, and the realizations are set according to the intervention.
Afterwards, the normal procedure again takes over; however, note that the
intervened values now are passed to the descendants444PARCS Python package
allows for in-place intervention actions for the DGP. With the help of
reusable error terms, counterfactual analysis is possible..

###  3.2 Partial randomization

As described in Section 1 and depicted by Figure 1, generating DGPs based on
method assumptions is done by identifying the restricted meta-parameters, and
varying the free meta-parameters. We argued why independence among meta-
parameters is crucial in this process. In this section, we explain how this
process, called partial randomization for brevity, is implemented in the PARCS
meta-algorithm.

The meta-parameters in PARCS follow a natural order, from the broadest to the
narrowest scope, which starts by the overall DAG structure and ends with all
the edge functions and node distributions along with their parameters.
However, users do not need to follow this order when restricting the meta-
parameters. they can leave any of the meta-parameters up to randomization, and
restrict the rest, regardless of the rank of the parameter in the order.
Figure 3 depicts a schematic of the partial randomization process via for a
hypothetical scenario. The partially specified DGP (top right) represents the
model assumptions. In this DAG, the existence of three nodes (1, 2, 5) is
assumed. While node 1 and 2 have a fixed distribution, node 5 awaits the
randomization process for its distribution (hence depicted by dashed line).
Furthermore, assumptions imply that there may exist a node 3 (hence depicted
by dotted line) which might receive an edge from 1 (dotted $1\rightarrow 3$
edge), but if it exists, then it definitely have outgoing edges to 2 and 5
(dashed $3\rightarrow 2$ and $3\rightarrow 5$ edges), namely it can be a
confounder for 2 and 5. Finally, a node 4 may exist, and it may be the parent
of node 1, node 5, both or none (dotted $4\rightarrow 1$ and $4\rightarrow 5$
edges).

According to the meta-parameter order, we first randomly decide the existence
of nodes 3 and 4. If the assumptions concern the chance of having these nodes,
they must be reflected in the randomization guidelines. For instance, the
guidelines might impose a success probability of $0.5$ for node 3, and $0.8$
for node 4. After each randomization, we have a new DGP where the existence of
nodes 3 and 4 has been decided. The three sampled DAGs in the Figure 3 (bottom
row) present different results for these nodes after randomization.

Next, we place the new edges based on the status of the newly added nodes and
the randomization guidelines. Edges $3\rightarrow 5$ and $3\rightarrow 2$ must
be added if node $3$ exists (as mandated by the assumptions), while we must
randomize the existence of $1\rightarrow 3$, $4\rightarrow 1$ and
$4\rightarrow 5$. This can be done, for instance, by randomly deciding on the
existence of edges according to a sparsity parameter. Sparsity can be either
determined by the model assumptions, or vary freely in a given range.
Subsequently, the edge functions and their parameters are randomized (or fixed
by the user).

Finally, we randomize the distribution for node $5$ and the newly-added nodes.
As before, we choose the distributions ($P$) from a given list of options in
the randomization guideline, and then sample the coefficients ($W$) based on
the number of parents and the input library ($\zeta$) that is specified by the
guideline. Note that this step of randomization is not limited by the meta-
parameters of the previous steps by any means. Even though the number of
sampled coefficients (trivially) depend on the number of parent nodes, type of
the distribution, complexity of the input library, and ranges of $W$ can be
freely randomized. Furthermore, every coefficient can be specified directly by
the user, not to be left for randomization. For example, the user can fix the
causal coefficient of node $1$ for the mean parameter of $2$ beforehand, and
randomize the rest of the coefficients.

Model assumptions partially specified DAG and SCM 12345 Randomization
guidelines PARCS randomization Figure 3: A schematic of graph randomization in
PARCS. thick nodes and edges are fully specified by the user. Dashed and
dotted nodes and edges are partially specified and randomized. In the
resulting three graphs, distributions and edge functions may differ
(represented by different colors). All graphs comply with the specified sub-
graph, and vary in the undefined sub-parts.

While the example in Figure 3 might seem disconnected to the actual examples
in the CI literature, it showcases several scenarios that may arise in real
use cases. As hinted in Table 1, the actual CI assumptions translate to the
four DAG/SCM parameters whatsoever, and those are being handled by PARCS meta-
parameters. To clarify the partial randomization process further, we revisit
the method assumptions in Figure 1 and explain briefly, how PARCS generates
DGPs for each:

######  Example.

Simulation studies for the three method assumptions in Figure 1 can be
implemented by PARCS as follows ($\zeta$ is selected based on Equation 5a):

  * M1

    1. 1.

Set $Z=\\{Z_{1},Z_{2},Z_{3}\\},E=\\{(Z_{1},Z_{2}),(Z_{1},Z_{3})\\}$;

    2. 2.

set $e_{12},e_{13}$ to identity function;

    3. 3.

set $P^{1}$ to normal; $W^{1}_{\mu}=0$, $W^{1}_{\sigma}=1$;

    4. 4.

set $P^{2}$ to normal; randomize question marks in $W^{2}_{\mu}=(?,?,0)$
(suppressing the $Z_{1}^{2}$ term); randomize $W^{2}_{\sigma}=(?,0,0)$ (std of
the error independent of $Z_{1}$);

    5. 5.

set $P^{3}$ to normal; set $W^{3}_{\mu}=(0,0,1)$ and $W^{3}_{\sigma}=(0,0,0)$.

  * M2

    1. 1.

Set
$Z=\\{Z_{1},Z_{2},Z_{3}\\},E=\\{(Z_{1},Z_{2}),(Z_{1},Z_{3}),(Z_{2},Z_{3})\\}$;

    2. 2.

set $Z_{1}$, $Z_{2}$ nodes and $(Z_{1},Z_{2})$ edge parameters as in M1;

    3. 3.

randomize $e_{13}$ and $e_{23}$ by selecting from available edge functions;
randomize $\phi_{13}$ and $\phi_{23}$;

    4. 4.

set $P^{3}$ to normal with zero std; randomize $W^{3}_{\mu}$ fully.

  * M3

    1. 1.

Randomize the number of common causes for $Z_{2},Z_{3}$ and create the new
nodes.

    2. 2.

Randomize an adjacency matrix for the new nodes according to the desired
sparsity level.

    3. 3.

Set and randomize $Z_{2}$, $Z_{3}$, newly sampled nodes and all edges, similar
to the process explained for M1 and M2.

###  3.3 Refined simulations

A challenge in the partial randomization process is to produce valid and
meaningful DAG configurations. In particular, the randomized DGPs shall not
violate the valid range of the distribution parameters, must lead to the
desired variability in the samples, and shall not produce simulation artifacts
which open a door to spurious inference on the synthesized data. Below, we
explain two possible options in PARCS design that users can choose to rely on,
in order to address these challenges.

####  3.3.1 Parameter restrictions and controlled variability

Many distribution parameters, such as Bernoulli’s success probability, are
defined in bounded ranges. Since, the partial randomization process samples
the free parameters independently and ignores the statistical characteristics
of the resulting structural equations, it is conceivable how a combination of
parameters might lead to invalid ranges for the distribution parameters. For
example, suppose ${Z_{2}\sim\text{Bern}(p=2Z_{1})}$. Any DGP that does not
impose ${Z_{1}\in[0,0.5]}$, will result in invalid success probability values
for $Z_{2}$. The issue can be resolved by imposing extra restrictions on the
SCM, e.g. forcing $Z_{1}$ to follow a uniform distribution with ranges smaller
than $[0,0.5]$; however, this defies the very purpose of PARCS, which is
limiting the DGP only according to the model assumptions.

As a solution, we propose the use of post-processing correction functions.
Applying to nodes (distributions), a correction operator transforms the values
of the distribution parameters to a valid range. This is done via a Sigmoid
function, applied to the raw values of the distribution parameters as

$\displaystyle f_{\text{corr.}}(\theta;L,U)=\frac{U-L}{1+e^{-\theta}}+L,$ (8a)

where $L$ and $U$ are the desired lower and upper bounds of the transformed
values. Equation 8a guarantees that for all the sampled meta-parameters, the
distribution parameters receive the inputs in their respective valid ranges.
Note that by applying correction, the parent-child nodes relations are
inevitably non-linear. Nevertheless, such transformations are common in
statistical modeling, e.g. in the logistic models, where the binary variable
is defined as $y=\sigma(w^{\top}x)$.

While Equation 8a guarantees the validity of the sampled parameters, the
quality of resulting data is not controlled. Following the argument of
correction function, consider the child node with Bernoulli distribution
${Z_{2}\sim\text{Bern}(p=2Z_{1})}$ and the parent node defined as
$Z_{1}\sim\mathcal{N}(\mu=10,\sigma=1)$. After correction, the success
probability is spread around $0.9995$ with the standard deviation of
$1.54\times 10^{-5}$ (empirical estimation over $1000$ samples). This means
that $Z_{2}$ is heavily imbalanced toward $1$. While this scenario could be
desired and intended in some designs, in general, we want to simulate data
with reasonable (controlled) variability. This can be achieved by introducing
an offset parameter in the correction function, and revising Equation 8a as

$\displaystyle f_{\text{corr.}}(\theta;L,U,O)=\frac{U-L}{1+e^{-\theta+O}}+L.$
(8b)

By controlling the offset, we shift the inputs to the correction function, and
hence change the mean of the output. In PARCS, rather than choosing $O$
explicitly, we set a target mean parameter. This way, the corresponding offset
is estimated empirically over a selected number of burn-in samples upon
instantiating the graph; it will remain fixed for the next sampling rounds.

####  3.3.2 Mitigating the simulation artifacts

Distribution parameters are directly affected by the statistics of the parent
nodes. In particular, the measurement unit and variance of the parent
determines the characteristics of the child. This is an undesired simulation
effect which leads to invalidity of the simulation process. Reisach et al.
(2021) shows how data scale and marginal variance might undesirably carry
information about the simulation DAG structure. Moreover, in an extreme case
where the support of a parent is considerably larger than the others, it will
dominate the matrix product in Equation 6. As a result, smaller magnitude
parents lose their effect solely due to unit of measurement. To address the
issue, we introduce a post-processing edge correction function which is a
simple standardization of the edge inputs by applying

$\displaystyle f_{\text{edge
corr.}}(Z_{j};\mu,\sigma)=\frac{Z_{j}-\mu}{\sigma},$ (9)

where $\mu$ and $\sigma$ are empirical mean and standard deviation of $Z_{j}$.
Similar to the offset argument in node correction function, the moments are
calculated empirically over the burn-in samples upon instantiating the graph
and will remain fixed for the next sampling rounds. By enabling correction,
sampled causal distributions are prevented from establishing invalid
structural equations and produce more meaningful simulations.

####  3.3.3 More complex simulations

No meta-algorithm design covers all possible scenarios; the goal must be to
design the capabilities and limitations according to the field of research.
PARCS is no exception to this rule. While it is designed to be effective for
the assumptions in Table 1, there are certain limitations in the design.
Particularly, the polynomial form of the distribution parameters, i.e.
Equations 5a and 6, cannot contain functions of multiple parents, e.g.,
${\theta_{k}=\sin{(Z_{1}+Z_{2})}+Z_{3}}$ (the form could have been implemented
via the choice of the edge function if the arbitrary function is applied to
one parent). This form, however, is still possible to induce for PARCS, using
an intermediate ‘dummy’ variable. In this case, we can imagine a new
(deterministic) variable is defined as $X=Z_{1}+Z_{2}$, while the $e_{XZ}$
edge function is the sine edge function. The message is, despite the apparent
incapability to simulate complex DGPs, PARCS provides indirect methods which
we can leverage in certain situations, while retaining the same randomization
options.

###  3.4 Comparison with other frameworks

The majority of CI research papers which run simulation studies for empirical
validations, employ a hard-coded and limited DGP, similar to the example in
Figure 1. In addition, benchmarking data sets have been introduced for causal
methods, such as the pairwise causal discovery data sets (Guyon et al., 2019).
Regardless of the complexity and diversity of scenario, both options lack the
flexibility for extension beyond a pre-defined DGP. Thus, for a meaningful
comparison of our method with existing solutions, we draw redears’ attentions
to two simulation frameworks which, similar to PARCS, are designed to be
applicable in various CI scenarios, but have limited funcitonalities.

Simcausal (Sofrygin et al., 2017) is a programming package, designed to
facilitate defining and sampling procedures of SCMs. Simcausal can prove
effective in designing causal simulations, not only in static, but also in
temporal settings. The drawback of the simcausal is that it requires full
specification of the structural equations for all variables. This means that
researchers must take extra steps and implement a parameterization and
randomization procedure in order to simulate partially specified DAGs and
SCMs. PARCS overcomes this limitation via the partial randomization
functionality, which is possible because of its meta-parameter design. In
summary, simcausal’s meta-algorithm does not go beyond the basic SCM
simulation and sampling procedures.

The next framework, and the closest to the PARCS design, is the causal
simulation framework for time series, proposed by Lawrence et al. (2021)
(which we call ‘temporal framework’ for ease of reference). In this framework,
time series simulation is based on sequential DAGs, in which the state of each
temporal variable at each time point is represented by a distinct node. The
framework’s meta-parameters are mainly tailored for a unique temporal
structure with ‘feature’, ‘target’, and ‘latent’ variables. For all nodes,
users can define (or randomize) the number of parents or children, complexity
of the distributions, and noise variances. Furthermore, the temporal aspect of
the DAG is controllable via parameters such as the ‘maximum temporal lag’ in
the parent-child pairs. Temporal framework also introduces refinement
procedures to prevent exploding of variables after recursive matrix
multiplications.

Similar to PARCS, users of the temporal framework can partially specify a
sequential DAG, and randomize the undefined meta-parameters. We, however,
argue that PARCS’ meta-parameters and randomization procedure proves more
flexible and adjustable to various method assumptions. Firstly, the temporal
framework cannot be adjusted to DAG structures beyond the feature-target-
latent structure, e.g. the instrumental variable setting or missing graphs
(see Section 4.1). For many structures, the framework does not provide
relevant meta-parameters, while users are forced to specify or randomize a
number of existing parameters which are irrelevant, e.g., the lag parameter
for static DAGs).

Secondly, the simulation complexity is adjusted via only one meta-parameter,
while, as argued in the desideratum C, there are many sources of complexity
such as non-linear equations, non-linear edge functions or non-additive non-
Gaussian dependent noise terms. In summary, the temporal framework can prove
effective for a range of simulation studies, especially for time series
methods, but will fall short in a wide range of studies which can be handled
by PARCS555Beyond the meta-algorithm design, the programming package for PARCS
and the temporal framework can be compared in terms of utility, user-
friendliness, transparency, and documentation. We leave this comparison to
future works, as it is out of the scope of this paper. It is worthy of note,
that the PARCS package provides temporal causal simulation based on the same
‘sequential DAG’ idea..

##  4 Case studies

In this section, we reproduce and extend the simulation studies of two SoA
papers to showcase the advantages of PARCS and its power to enhance the CI
studies. First paper by Jarrett et al. (2022) introduces Hyperimpute, an
missing data imputation method for ignorable missingness mechanisms in complex
datasets. Next, we look in the field of causal inference at the effective and
widely-used statistical LiNGAM method (Shimizu et al., 2006). This method
recovers the causal order and the causal linear coefficients. Beside the
popularity and wide usage, we selected these two works because the validity of
their methods depends significantly on their assumptions, i.e. DAG structure
and non-linearity assumptions for Hyperimpute and linearity assumption for
LiNGAM.

###  4.1 Hyperimpute: a missing data analysis case study

Missing data can be posed as a causal inference problem via m-graph models
(Mohan et al., 2013; Shpitser et al., 2015; Mohan and Pearl, 2022). An m-graph
is a causal DAG with $Z$ and $R$ nodes, representing the main variables and
missingness indicators respectively. The $Z\rightarrow R$ relations determine
the missingness mechanism. Figure 4 presents five examples of m-graphs,
inducing various missing data mechanisms. Being a causal DAG, m-graphs can be
simulated using PARCS. To showcase this ability, we analyze the simulation
study in Jarrett et al. (2022) for the Hyperimpute imputation algorithm. We
briefly describe the paper’s simulation setting, then continue with
reproducing and extending some of the presented results using PARCS. We
demonstrate that authors have correctly reported the results of a subset of
scenarios which complies with their assumptions, but ignore a wide range of
complying scenarios still. We show that the proposed method achieves
considerably different performances in some of the untested scenarios, which
means the analysis in the paper is incomplete due to simulation limitations.
For ease of reference, we refer to this paper as ‘Hyperimpute paper’.

$Z_{1}$$Z_{2}$$Z_{3}$$R_{1}$$R_{2}$$R_{3}$ (a) MCAR mechanism

$Z_{1}$$Z_{2}$$Z_{3}$$R_{2}$$R_{3}$ (b) MAR mechanism

$Z_{1}$$Z_{2}$$Z_{3}$$R_{1}$$R_{2}$$R_{3}$ (c) MNAR mechanism

$Z_{1}$$Z_{2}$$Z_{3}$$R_{1}$$R_{2}$$R_{3}$ (d) self-censoring (MNAR) mechanism

$Z_{1}$$Z_{2}$$Z_{3}$$R_{1}$$R_{2}$$R_{3}$ (e) no-self-censoring (MNAR)
mechanism*

Figure 4: Five examples of m-graphs, modeling different missingness
mechanisms: MCAR (a), MAR (b), MNAR (c, d, e). *The term no-self-censoring is
coined by Malinsky et al. (2022) for a similar graph as in (e), but with
undirected edges among $R$ nodes, hence modeling a chain graph. In theory, the
graph in (e) is considered a submodel.

####  4.1.1 Original simulation setting

In the Hyperimpute paper simulation study, missingness is induced artificially
in real-world datasets, according to a mechanism. Experiments are conducted
using 12 datasets from the UCI machine learning repository (Dua and Graff,
2017) and their proposed method is compared with 8 imputation methods. The
performance is measured using the RMSE calculated over imputed entries. Four
missingness mechanisms are simulated, which are described below:

  1. 1.

MCAR: By randomly masking data entries (4(a)).

  2. 2.

MAR: A special version of Missing-At-Random mechanism where a number of
variables are selected to be fully observed, and missingness in the remaining
ones is induced as a function of them (based on a logistic model) (4(b))

  3. 3.

MNAR: Taking the MAR setting and randomly inducing missingness in fully-
observed variables, now the missingness in remaining ones is dependant to
partially observed data, hence the mechanism is Missing-Not-At-Random (4(c))

  4. 4.

self-censoring MNAR: a logistic function of the each variable is used to
induce missingness in itself (4(d))

For the logistic functions, the selected coefficients satisfy a ‘unity vector’
restriction to avoid meaningless data synthesis. Analysis is done for
different missingness ratios (10%, 30%, 50%), different observed nodes for MAR
and MNAR scenario, and different sample sizes. For all experiments, the number
of iterations is $N_{\text{it}}=10$.

####  4.1.2 PARCS simulation setting

Table 3 describes the extension scenarios for the Hyperimpute simulation
study. To focus more on showcasing PARCS utility, we run the experiments for
only the compression dataset (Yeh, 1998), covering MAR and MNAR settings with
50% missingness ratio, and for two imputation methods, Hyperimpute and
Missforest (Stekhoven and Bühlmann, 2012). We follow the paper’s choice of
iteration number, unless mentioned otherwise.

Table 3: Hyperimpute extension scenarios Mechanism | Experiment no. | Extension | Description
---|---|---|---
MAR | 1 | coefficient vector | Same mechanism, Wider range of coefficient vector
2 | edge density | Same mechanism, fewer MAR edges
3 |  $R\rightarrow R$ edges | Different mechanism by adding new MAR edges
4 | nonlinearity | Same mechanism, nonlinear structural equations
MNAR | 5 | mechanism | Inducing other MNAR mechanisms

Implementation is done in three steps:

  1. 1.

$\mathcal{G}_{Z}$ subgraph: We represent the dataset variables as graph nodes,
which may only have out-going edges..

  2. 2.

$\mathcal{G}_{R}$ subgraph: We specify the DAG and SCM models for the
missingness indicator subgraph, partially and based on the intended
missingness scenario.

  3. 3.

$\mathcal{G}_{Z}$-to-$\mathcal{G}_{R}$ edges: We create missingness mechanism
edges by randomizing the $Z\rightarrow R$ edges. We control the pattern of the
connection by a particular ‘mask’ per each mechanism.

(a) Reproducing the original MAR simulation with PARCS

(b) MCAR-to-MAR transition performance analysis

(c) Inducing $R\rightarrow R$ edges in MAR

(d) linear and nonlinear MAR scenarios

(e) Inducing various MNAR scenarios

Figure 5: Simulation results for imputation methods, as done in the
Hyperimpute paper (Jarrett et al., 2022). Boxplots are obtained via 10
experiment iterations, except for the ‘*PARCS: general’ column in Figure (e)
with 30 iterations.

##### Experiment 1: Reproducing MAR results

Firstly, we produce the same MAR scenario in the paper (4(b)). The only
difference between the experiments is the selected ranges for coefficient
vectors $W_{p}$ (weights of the success probability of the Bernoulli
distribution for $R$ nodes): In the Hyperimpute paper, the coefficients are
selected such that the $W$ becomes a unit vector. We randomize the
coefficients in $[-5,-1]\cup[1,5]$ (correction applied). 5(a) presents the
RMSE boxplots for two simulation methods, where a considerable difference in
both mean and standard deviation of the obtained RMSE is visible. This
difference shows that the simulation in paper was incomplete due to free
parameter limitations.

##### Experiment 2: MCAR-to-MAR transition

In the MAR m-graph in 4(b), removing $Z\rightarrow R$ edges changes the
mechanism to MCAR. Both MCAR and MAR scenarios are studied in the Hyperimpute
paper. However, we can extend the analyses by visualizing the gradual increase
in RMSE when the mechanism alters to MAR. This is achieved by varying the
sparsity parameter of the randomization guideline from $0.0$ (MCAR) to $1.0$
(MAR). 5(b) presents the resulting boxplots for two imputation methods. Only
first and last columns are reported by the paper (reproduced here).

##### Experiment 3: Extending to different MAR mechanism

5(c) presents the boxplot result for another MAR scenario which is not tested
in the Hyperimpute paper. Here, the MAR mechanism is induced when $R$ nodes
have edges among them too. This scenario is achieved by updating the
$\mathcal{G}_{R}$ description with new edges. The results are generated for
varying graph sparsity values. For this experiment, the $Z\rightarrow R$
density is set to $0.0$. Except the first column (corresponding to the paper’s
MCAR), the plot provides new performance measures in scenarios. Imputation
methods show to be relatively robust to this extension of scenario.

##### Experiment 4: Non-linearity extension for MAR

Artificial missingness in the Hyperimpute paper is induced using a simple
logistic function. from PARCS perspective, this means the coefficients for
non-linear components in $\zeta$ are zero. In addition, only the identity edge
function is used. We test the imputation methods in the presence of non-linear
relations and interactions. 5(d) presents the boxplot result of two imputation
methods for linear and nonlinear relations. For this experiment, $R\rightarrow
R$ and $Z\rightarrow R$ edge densities are selected $0.0$ and $0.6$
respectively. The first column (linear) presents the reproduced results, while
last two columns demonstrate the methods’ performance in untested scenarios.
Similar to experiment 1, extension of simulation has lead to different RMSE
scores, meaning that the reported performance in paper was limited by the
simulation studies.

##### Experiment 5: Extending MNAR simulation

In the Hyperimpute paper, MNAR mechanisms are induced by either self-censoring
or applying MCAR missingness to fully observed nodes of a MAR mechanism. Using
PARCS, we simulate many known MNAR mechanisms by their corresponding m-graphs.
These mechanisms are specified by the $Z\rightarrow R$ mask parameter. In this
experiment we tested three scenarios: a general MNAR scenario (all-one mask)
and two specific self-censoring and no-self-censoring scenarios. Note that SC
and NSC scenarios are especial cases of the general scenario. 5(e) presents
the boxplot results of 3 scenarios alongside the result of the simulation done
in the Hyperimpute paper. For general MNAR scenario, we ran the simulations
for 30 iterations to achieve more reliable results, as the graph search space
in this scenario is broader compared to other scenarios. The significant
difference between the first boxplot (reproduced) and the other three
(extended) points to the fact that the reported performance in the paper for
MNAR missingness is limited by the incomplete simulation study.

###  4.2 LiNGAM: a causal inference showcase

LiNGAM is a well studied statistical causal model introduced by Shimizu et al.
(2006) alongside a causal inference and discovery method DirectLiNGAM. The
model can shortly be written as $Z=BZ+\varepsilon$ with a lower diagonal
matrix $B$ and independent, non-Gaussian error terms $\varepsilon$. Due to the
linearity of the statistical model, DirectLiNGAM determines the causal order
of the variables by recursively running linear regressions and removing the
source nodes. The parameters of the estimated adjacency matrix $\hat{B}$ are
then straightforward set as the estimated parameters of the corresponding
regressions.

####  4.2.1 Original Simulation Setting

Using PARCS, we first implemented a class of distributions following the
original paper the statistical model. In particular, we set the dimension of
$X$ to $p=5$, selected the adjacency matrix weights uniformly from the
intervals $[-2,-0.5]$ and $[0.5,2]$, forced the independent error terms to be
log-exponentially distributed and permuted the entries of $X$ uniformly over
$S_{5}$ to achieve a random causal ordering. As for the following two
experiments, we then generated $500$ datasets with a dataset size of $1000$
samples each. As an objective to quantify the performance of DirectLiNGAM, we
chose the Frobenius norm between the estimated and true adjacency matrix. The
results can be seen in Figure 6 and there, smaller values on the y-axis
indicate a better performance. As expected and shown in the green boxplot,
DirectLiNGAM recovered the $500$ scenarios generated under the setting of the
original paper very well.

####  4.2.2 PARCS simulation setting

DirectLiNGAM works good on data following the LiNGAM statistical model, which
is however hardly ever satisfied for real data. Still, practitioners use and
thus it is important to study its performance under model misspecification.
With PARCS, we conduct two experiments altering the original setup resulting
in datasets which are violating the linearity assumption.

##### Experiment 1:

For the first alteration, we uniformly sample a constant $\phi\in[0.75,1.25]$
and alter the structural equations from
$Z_{j}=B_{j,:}(Z_{1},\dots,Z_{p})^{\top}+\varepsilon_{j}$ to

$Z_{j}=B_{j,:}(g_{\phi}(Z_{1}),\dots,g_{\phi}(Z_{p}))^{\top}+\varepsilon_{j}$

for $j\in[p]$. The function $g_{\varphi}(z)=\text{sgn}(z)\cdot|z|^{\phi}$
introduces a nonlinear behaviour. In this setup, the value $|\phi-1|$
quantifies the distance to the LiNGAM assumption and only datasets with
$\phi=1$ do follow the original statistical model. One implication by this
data generating process is, that in all cases with $\phi<1$, the variance of
the variables will shrink with the causal order. The performance results of
DirectLiNGAM is presented by the blue visualizations in Figure 6. There we
find a clear increase in recovery performance for $\phi$ near $1$. The further
away $\phi$ is from $1$, the higher the model violation is, the worse is the
recovery performance.

##### Experiment 2:

To counteract the unstable variance behaviour from the proceeding experiment,
we add here a PARCS’ edge correction as in Equation 9 to all edges of the
nonlinear data generating process.We find in formulae for all $j\in[p]$ the
structural equations

$Z_{j}=B_{j,:}(\frac{g_{\phi}(Z_{1})-\mathrm{E}[g_{\phi}(Z_{1})]}{\sqrt{\mathrm{Var}(g_{\phi}(Z_{1}))}},\dots,\frac{g_{\phi}(Z_{p})-\mathrm{E}[g_{\phi}(Z_{p})]}{\sqrt{\mathrm{Var}(g_{\phi}(Z_{p}))}})^{\top}+\varepsilon_{j}.$

Although this setup stabilizes the variance, the relation between the
variables is still nonlinear and thus, the data does not follow LiNGAM. This
can in particular be also seen in the orange part of Figure 6 as we miss the
drop for $\phi$ values near $1$. Interestingly however, the general
performance of DirectLiNGAM is better on the corrected datasets for $\phi>1.1$
as the marginal distributions of e.g. $Z_{5}$ do not have a large variances
and are instead compact.

Figure 6: Frobenius norm of the difference between the estimated and true
adjacency matrix $\hat{B}-B$ against the nonlinearity parameter $\phi$. Each
boxplot in the right panel summarizes the performance over 500 datasets.

Regarding the role of PARCS, it provided us a algorithmic-inexpensive option
to test a causal discovery and inference method on a considerably extended
distribution space. Especially for practitioners, it is important to also be
able to trust a method when the underlying model is not correctly specified,
and knowing how to analyse the observed distribution to assess the potential
model performances.

##  5 Conclusion

For causal inference and discovery methods, simulation studies provide
reliable and comprehensive results only if they are designed according to the
promised operational conditions of the methods. Many causal inference
literature tend to over-restrict the design of their simulation study; this
leads to unreliable and misleading simulation results.

PARCS simulation framework is designed to be effortlessly adjusted by method
assumptions. It adopts the causal DAG and SCM models and provides a
parameterization scheme to control the statistical and causal characteristics,
and complexity of the data, from various aspects transparently and
independently. As a result of this design, PARCS allows the users to partially
fix a subset of parameters according to their assumptions, and sample the free
parameters in iteration, to generate series of complying data generating
processes.

In this paper, we introduced the simulation meta-algorithm of PARCS, and
exemplified its effectiveness by reproducing and extending the simulation
studies in two missing data analysis and causal discovery research papers. We
showed that using PARCS, their empirical validation results would be more
comprehensive.

We also compiled a list of desiderata for an effective causal simulation
framework. We revisit the list and discuss how PARCS satisfies all the items:

  * •

Underlying causality: PARCS models the underlying DGP with the widely-used
concept of causal DAG and SCM models. The simulation studies are defined by
specifying the DAG structure and SCM parameters.

  * •

Data dimension and modality: Arbitrary distributions can be assigned in order
to simulate different data modalities. In addition, more complex SCMs can be
simulated using intermediate nodes. On the data set level, sample size and
number of features are adjustable, and directly controlled by the user.

  * •

Controlled complexity: PARCS controls the data complexity using the number of
graph edges (sparsity), diverse output distributions, non-linear edge
functions, and extended input vector dictionaries. Each of these sources can
be adjusted separately. In all cases, partial randomization can be applied.

  * •

Noise models: The noise models are easily and independently adjusted by
adopting different node distributions. Any sort of distribution can be
simulated as long as the (parametric or non-parametric) iCDF of the
distribution is defined.

By introducing the theoretical background of PARCS and publishing the Python
package for public use, we propound the framework to the CI community for the
future research in this field. PARCS adheres to the widely-used causal
inference models, and facilitates transparent, reproducible, systematic, and
comprehensive empirical investigations; thus we believe it is capable of
playing an effective role as a benchmarking tool for CI methodologies.

##  6 Acknowledgements

This work was funded by the Bavarian Ministry for Economic Affairs, Regional
Development and Energy as part of a project to support the thematic
development of the Institute for Cognitive Systems.

We thank Octavia Ciora and Elisabeth Pachl for proofreading and providing
instructive feedback.

## References

  * (1)
  * Dua and Graff (2017) Dheeru Dua and Casey Graff. 2017.  UCI Machine Learning Repository.  http://archive.ics.uci.edu/ml
  * Guyon et al. (2019) Isabelle Guyon, Alexander R. Statnikov, and Berna Bakir Batu (Eds.). 2019.  _Cause Effect Pairs in Machine Learning_.  Springer, Cham Switzerland.  https://doi.org/10.1007/978-3-030-21810-2
  * Hahn et al. (2019) P. Richard Hahn, Vincent Dorie, and Jared S. Murray. 2019\.  Atlantic Causal Inference Conference (ACIC) Data Analysis Challenge 2017.  arXiv:1905.09515 [stat.ME]
  * Hoyer et al. (2008) Patrik O. Hoyer, Dominik Janzing, Joris M. Mooij, Jonas Peters, and Bernhard Schölkopf. 2008\.  Nonlinear causal discovery with additive noise models. In _Advances in Neural Information Processing Systems 21, Proceedings of the Twenty-Second Annual Conference on Neural Information Processing Systems, December 8-11, 2008_ , Daphne Koller, Dale Schuurmans, Yoshua Bengio, and Léon Bottou (Eds.). Curran Associates, Inc., Vancouver, British Columbia, Canada,, 689–696.
  * Jarrett et al. (2022) Daniel Jarrett, Bogdan Cebere, Tennison Liu, Alicia Curth, and Mihaela van der Schaar. 2022\.  HyperImpute: Generalized Iterative Imputation with Automatic Model Selection. In _International Conference on Machine Learning, ICML 2022, 17-23 July 2022_ _(Proceedings of Machine Learning Research, Vol. 162)_ , Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (Eds.). PMLR, Baltimore, Maryland, USA, 9916–9937.
  * Lawrence et al. (2021) Andrew R. Lawrence, Marcus Kaiser, Rui Sampaio, and Maksim Sipos. 2021.  Data Generating Process to Evaluate Causal Discovery Techniques for Time Series Data.  arXiv:2104.08043 [stat.ML]
  * Malinsky et al. (2022) Daniel Malinsky, Ilya Shpitser, and Eric J Tchetgen Tchetgen. 2022\.  Semiparametric inference for nonmonotone missing-not-at-random data: the no self-censoring model.  _J. Amer. Statist. Assoc._ 117, 539 (2022), 1415–1423.
  * Mohan and Pearl (2022) Karthika Mohan and Judea Pearl. 2022.  Graphical models for recovering probabilistic and causal queries from missing data.  In _Probabilistic and Causal Inference: The Works of Judea Pearl_. Association for Computing Machinery, New York, NY, United States, 413–432.
  * Mohan et al. (2013) Karthika Mohan, Judea Pearl, and Jin Tian. 2013.  Graphical Models for Inference with Missing Data. In _Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013_ , Christopher J. C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (Eds.). Neural Information Processing Systems Foundation, Inc. (NeurIPS), Lake Tahoe, Nevada, United States, 1277–1285.
  * Pearl (2009) Judea Pearl. 2009\.  Causal inference in statistics: An overview.  _Statistics surveys_ 3 (2009), 96–146.
  * Peters and Bühlmann (2015) Jonas Peters and Peter Bühlmann. 2015.  Structural intervention distance for evaluating causal graphs.  _Neural computation_ 27, 3 (2015), 771–799.
  * Pfister and Peters (2022) Niklas Pfister and Jonas Peters. 2022.  Identifiability of sparse causal effects using instrumental variables. In _Uncertainty in Artificial Intelligence, Proceedings of the Thirty-Eighth Conference on Uncertainty in Artificial Intelligence, UAI 2022, 1-5 August 2022_ _(Proceedings of Machine Learning Research, Vol. 180)_ , James Cussens and Kun Zhang (Eds.). PMLR, Eindhoven, The Netherlands, 1613–1622.  https://proceedings.mlr.press/v180/pfister22a.html
  * Reisach et al. (2021) Alexander G. Reisach, Christof Seiler, and Sebastian Weichwald. 2021.  Beware of the Simulated DAG! Causal Discovery Benchmarks May Be Easy to Game. In _Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021_ , Marc’Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (Eds.). Neural Information Processing Systems Foundation, Inc. (NeurIPS), virtual, 27772–27784.
  * Shimizu et al. (2006) Shohei Shimizu, Patrik O. Hoyer, Aapo Hyvärinen, and Antti J. Kerminen. 2006.  A Linear Non-Gaussian Acyclic Model for Causal Discovery.  _J. Mach. Learn. Res._ 7 (2006), 2003–2030.
  * Shpitser (2013) Ilya Shpitser. 2013\.  Counterfactual graphical models for longitudinal mediation analysis with unobserved confounding.  _Cognitive science_ 37, 6 (2013), 1011–1035.
  * Shpitser et al. (2015) Ilya Shpitser, Karthika Mohan, and Judea Pearl. 2015\.  Missing Data as a Causal and Probabilistic Problem. In _Proceedings of the Thirty-First Conference on Uncertainty in Artificial Intelligence, UAI 2015_ , Marina Meila and Tom Heskes (Eds.). AUAI Press, Amsterdam, The Netherlands, 802–811.  http://auai.org/uai2015/proceedings/papers/204.pdf
  * Sofrygin et al. (2017) Oleg Sofrygin, Mark J van der Laan, and Romain Neugebauer. 2017\.  Simcausal R package: conducting transparent and reproducible simulation studies of causal effect estimation with complex longitudinal data.  _Journal of statistical software_ 81, 2 (2017), 1–47.
  * Spirtes and Glymour (1991) Peter Spirtes and Clark Glymour. 1991.  An algorithm for fast recovery of sparse causal graphs.  _Social science computer review_ 9, 1 (1991), 62–72.
  * Stekhoven and Bühlmann (2012) Daniel J Stekhoven and Peter Bühlmann. 2012.  MissForest—non-parametric missing value imputation for mixed-type data.  _Bioinformatics_ 28, 1 (2012), 112–118.
  * Uhler et al. (2013) Caroline Uhler, Garvesh Raskutti, Peter Bühlmann, and Bin Yu. 2013.  Geometry of the faithfulness assumption in causal inference.  _The Annals of Statistics_ 41, 2 (2013), 436–463.
  * Xu et al. (2021) Liyuan Xu, Yutian Chen, Siddarth Srinivasan, Nando de Freitas, Arnaud Doucet, and Arthur Gretton. 2021\.  Learning Deep Features in Instrumental Variable Regression. In _9th International Conference on Learning Representations, ICLR 2021_. OpenReview.net, Virtual Event, Austria, 1.
  * Yeh (1998) I-C Yeh. 1998.  Modeling of strength of high-performance concrete using artificial neural networks.  _Cement and Concrete research_ 28, 12 (1998), 1797–1808.


11institutetext: University of Alabama at Birmingham 22institutetext:
University of Oxford33institutetext: National University of
Singapore44institutetext: Carnegie Mellon University
44email<EMAIL_ADDRESS>
#  Uncertainty-Aware Adapter: Adapting Segment Anything Model (SAM) for
Ambiguous Medical Image Segmentation

Anonymous  Mingzhou Jiang1,* Jiaying Zhou1,* Junde Wu2,*† Tianyang Wang 11
Yueming Jin 33 Min Xu4,🖂
$*$ equal technical contribution 1 $\dagger$ project lead

###### Abstract

The Segment Anything Model (SAM) gained significant success in natural image
segmentation, and many methods have tried to fine-tune it to medical image
segmentation. An efficient way to do so is by using Adapters, specialized
modules that learn just a few parameters to tailor SAM specifically for
medical images. However, unlike natural images, many tissues and lesions in
medical images have blurry boundaries and may be ambiguous. Previous efforts
to adapt SAM ignore this challenge and can only predict distinct segmentation.
It may mislead clinicians or cause misdiagnosis, especially when encountering
rare variants or situations with low model confidence. In this work, we
propose a novel module called the Uncertainty-aware Adapter, which efficiently
fine-tuning SAM for uncertainty-aware medical image segmentation. Utilizing a
conditional variational autoencoder, we encoded stochastic samples to
effectively represent the inherent uncertainty in medical imaging. We designed
a new module on a standard adapter that utilizes a condition-based strategy to
interact with samples to help SAM integrate uncertainty. We evaluated our
method on two multi-annotated datasets with different modalities: LIDC-IDRI
(lung abnormalities segmentation) and REFUGE2 (optic-cup segmentation). The
experimental results show that the proposed model outperforms all the previous
methods and achieves the new state-of-the-art (SOTA) on both benchmarks. We
also demonstrated that our method can generate diverse segmentation hypotheses
that are more realistic as well as heterogeneous.

###### Keywords:

SAM Adapter uncertainty samples.

##  1 Introduction

Medical image segmentation plays a vital role in healthcare, offering crucial
insights for various downstream clinician applications, like disease diagnosis
and evaluation. Recently, many proposed to adapt the pre-trained nature image
segmentation model, like Segment Anything Model (SAM) [1] to medical image
segmentation [2, 3, 4]. A popular cost-efficient approach is the Adapter
technique, which involves inserting a bottleneck module with only a few
parameters into the model. Through fine-tuning these small-size adapters, SAM
can bridge the domain gap between medical and natural images while retaining
superior performance. For instance, MSA [5], SAM-Med2D [6], SAM-adapter [7],
and others employ an Adapter strategy to transfer SAM to medical imaging,
achieving superior segmentation results.

However, unlike natural images, medical images have the unique features that
many organs and tissues in medical images are ambiguous. For example, when
segmenting lesions from lung abnormality images, different clinicians are
likely to provide varied annotations, and it is common to combine these
different annotations to represent the final ground truth with inherent
uncertainty. Conventional computer vision models, like SAM, are challenging to
apply directly in such cases. They tend to output the one-to-one mapping from
image to ground truth, which may lead to mispredictions and potentially
mislead clinicians’ diagnoses. Previous efforts to adapt SAM to the medical
field have also overlooked this critical challenge, rendering them
inapplicable for many real-world clinical scenarios. Therefore, for fine-
tuning SAM to medical images, it is essential to present a new fine-tuning
method that helps the model understand and calibrate uncertainty-aware
segmentation.

In this paper, we propose a novel module called the Uncertainty-aware Adapter
for fine-tuning SAM to ambiguous medical image segmentation. The basic idea is
to construct a latent space for sampling the possible segmentation variants
following the previous uncertainty works like Probabilistic U-Net (Prob U-Net)
[8, 9, 10] so that SAM can interact with the stochastic samples. Unlike most
previous works in the stochastic sample, which is merely applied to the output
layer of the model by concatenation, we design a new interaction method
between the model and samples to overcome the issue where the model might
ignore samples. In order to adapt this idea of sampling to our specific case
that the Adapters only contain a few parameters, we propose the Uncertainty-
aware Adapter, which utilizes a condition-based strategy to interact with
uncertainty samples from latent space. We designed a novel module called
Condition Modifies Sample Module (CMSM) in the Uncertainty-aware Adapter to
learn the interaction between the uncertainty sample and the Adapter. The
learnable position variant of the Uncertainty-aware Adapter matches the
feature-extracting process and is treated as the condition to calculate with
the uncertainty sample. Our method enhances interaction with uncertainty
samples, yielding more diverse and accurate segmentations.

Our contributions can be summarized as follows:

  * $\bullet$

We present the Uncertainty-aware Adapter SAM(UA-SAM), which combines a
probabilistic model to produce diverse likely segmentation hypotheses. It is
crucial to provide clinicians with reliable diagnostic assistance and reduce
the risk of misdiagnosis.

  * $\bullet$

We proposed the Condition Modifies Sample Module (CMSM) in the Uncertainty-
aware Adapter, which can help SAM to capture the uncertainty in medical images
while segmenting ambiguous medical images.

  * $\bullet$

We have evaluated our proposed UA-SAM model on the LIDC-IDRI dataset and the
REFUGE dataset. Our method demonstrates superior segmentation performance, a
significant step in ambiguous medical image segmentation.

##  2 Method

Figure 1: Overview of UA-SAM. We froze the parameters of SAM and only updated
the Adapter parameters. Note that we did not show the prompt encoder.

Before introducing UA-SAM, we provide an overview of the SAM architecture. SAM
consists of three main components: a large-scale image encoder, a prompt
encoder, and a lightweight mask decoder. The image encoder utilizes a Vision
Transformer (ViT) pre-trained by MAE [11] to process high-resolution images
(default $1024\times 1024$). The output embeddings of the image encoder are at
1/16 scale of the input image. The prompt encoder has two prompt styles:
sparse (points, boxes, texts) and dense (masks). Convolutional down-sampling
and GELU activation functions are applied to dense prompts. The mask decoder
updates image embeddings and prompt embeddings via two-way cross attention,
from prompt to image and image to prompt. This SAM framework allows it to
segment different targets based on different prompts.

In our work, we propose the Uncertainty-aware Adapter, learning a few
parameters while integrating uncertainty, which helps fine-tune SAM to the
uncertainty-aware medical image segmentation effectively. We use the
probabilistic model to model latent space. The Uncertainty-aware Adapter
utilizes a condition-based interaction method between the latent space
uncertainty sample and SAM, ensuring the model can output accurate and diverse
likely segmentation hypotheses. Our proposed UA-SAM architecture is shown in
Fig. 1.

Overview of UA-SAM. To output diverse plausible segmentation masks, UA-SAM
utilizes a probabilistic model to produce unlimited uncertainty samples, and a
segmentation model interacts with samples to output segmentation masks. The
probabilistic model inputs the images and models the distribution of images,
then outputs uncertainty samples. During the training and inference processes,
the probabilistic model provides an uncertainty sample z for the segmentation
model. The segmentation model is the SAM model that is fine-tuned by inserting
Adapters into the image encoder. The i-th ( i = 0,1,2…) Uncertainty-aware
Adapter is inserted behind the i-th Transformer block. All Adapter inputs
position variant ${p}_{i}$ and uncertainty sample z, then outputs ${p}_{i+1}$
for the next Adapter (Except for the last adapter, its output is directly
added to the embeddings reconstructed from the sample z to the same
dimensionality.).

Probabilistic model. Following the previously most commonly used paradigm [8,
9, 10], we utilize a Prior Net and a Posterior Net to separately construct
low-dimensional latent spaces P and Q, which obey the Gaussian distribution
$\mathcal{N}\left(\mu,diag\left(\sigma\right)\right)$. For the training
process, the uncertainty sample z comes from the latent space Q, and the
inference process from the latent space P, which is a random sample. At the
same time, there is a Kullback-Leibler divergence, which is a part of the loss
function (in Eq. 6) to penalize differences between the posterior distribution
Q and the prior distribution P.

Architecture of Uncertainty-aware Adapter. In order to reduce the number of
fine-tuning parameters as much as possible, the Uncertainty-aware Adapter (as
shown in Fig. 1) consists of two small-size components: one is a bottleneck
model that sequentially uses a down-projection, ReLU activation, and up-
projection, and another is the Condition Modifies Sample Module (CMSM), which
designed for interaction between uncertainty sample z and Adapter. Fig. 1
illustrates CMSM schematically. Given the learnable position variant
${p}_{i}\in{\mathbb{R}}^{1\times L}$ (L is the number of Adapter ) of current
Adapter and uncertainty sample $z\in{\mathbb{R}}^{B\times C}$ (B is the batch
size C is the latent space dimension), CMSM utilizes ${p}_{i}$ as the
condition to modify the state of z. It is similar to the dot-product attention
[12] using ${p}_{i}$ as the query and z as the key. This ensures the
uncertainty sample matches the feature extraction process so that the model
effectively responds to incorporating the uncertainty sample. Concretely, the
computational steps of the CMSM mechanism are demonstrated as follows:

${{p}_{i}}^{\prime}=Reshape\left(MLP\left({{p}_{i}}\right)\right)$ (1)
${z}^{\prime}={w}_{22}\left(z\right)\otimes
ReLU\left(Trans({w}_{11}\left({{p}_{i}}^{\prime}\right))\otimes{w}_{21}\left(z\right)\right)$
(2)
$f=ReLU\left({z}^{\prime}\right)\odot{w}_{12}\left({{p}_{i}}^{\prime}\right)$
(3) ${f}_{i}=MLP\left(Reconstruct\left(f\right)\right)$ (4)
${p}_{i+1}=flatten(MLP\left({w}_{12}\left({{p}_{i}}^{\prime}\right)\right))$
(5)

where ${w}_{11}$, ${w}_{12}$, ${w}_{21}$ and ${w}_{22}$ are linear projection
function. ‘Trans’ refers to the transposing to transform feature tensors into
proper shapes for calculation. Here, ‘$\otimes$’ and ‘$\odot$’ represent
matrix multiplication and element-wise multiplication at corresponding
positions. And the ‘Reconstruct’ denotes transforming feature tensors from
$f\in{\mathbb{R}}^{B\times C}$ to $f\in{\mathbb{R}}^{B\times H\times W\times
dim}$ (H and W are the height and width of patch embeddings processed by image
encoder. dim is a value of multiplying a ratio and embeddings dimension). The
uncertainty sample z is ultimately modified to the uncertainty feature
${f}_{i}$ that can be directly concatenated with down-projection embeddings.
And CMSM also outputs the next position variant ${p}_{i+1}$ for the next
Adapter.

$\mathcal{L}={E}_{z\sim Q\left(\cdotp\mid
Y,X\right)}\left[-\log{{P}_{\theta}}\left(Y\mid
S\left(X,z\right)\right)\right]+\beta\times{D}_{KL}\left(Q\left(z\mid
Y,X\right)\parallel P\left(z\mid X\right)\right)$ (6)

In the training process, we compute the loss, as shown in Eq. 6. Given the raw
image X and the ground truth segmentation Y, S is the predicted segmentation.
A DiceCEloss from the Monai library penalizes differences between S and Y (the
DiceCEloss arises from treating the output S as the parameterization of a
pixel-wise categorical distribution ${P}_{\theta}$)

##  3 Experiment

###  3.1 Dataset

We trained and evaluated our model on two multi-annotated datasets with
different modalities: the LIDC-IDRI dataset [13] and the REFUGE2 dataset [14].
For the LIDC-IDRI dataset, we use a pre-processed version of the dataset,
cropped 128*128 patches around lesions, which comprises 15096 thoracic CT
images with lesions annotated by four radiologists. Moreover, we conducted
experiments on the REFUGE2 optic-cup dataset, consisting of 1200 fundus
images, each annotated by seven radiologists. Both datasets are publicly
available.

For the LIDC-IDRI dataset, we divided the dataset into a training set and a
testing set at a ratio of 80:20. The training set comprises 12076 images,
while the testing set shall consist of 3018 (drop last), each measuring
128*128 pixels. For the REFUGE2 optic-cup dataset, we cropped the fundus
images around the center of the optic disc and then resized them to the size
of 512*512. Subsequently, we combined the training and validation set into a
new one, resulting in 800 images in the training set and 400 images in the
testing set.

###  3.2 Implementation details

In this study, we implemented our model on the LIDC-IDRI and REFUGE2 datasets.
We utilized the PyTorch and MONAI libraries for our project and employed the
Adam optimizer with an initial learning rate set to 1e-4. Meanwhile, we
incorporated the StepLR strategy for learning rate decay. During training, the
labels were randomly sampled from the multiple annotated labels corresponding
to the images. To ensure the highest model efficacy, we implemented a simple,
early-stopping mechanism. We use the "vit/b" version of SAM. Additionally, in
our training process, SAM was employed with single-point prompts.

Table 1: The Comparison of UA-SAM with SOTA methods evaluated by Dice Score. The best results are denoted in bold.  Method |  LIDC-IDRI |  REFUGE2
---|---|---
Deterministic Method | Unet [15] | 0.516 | 0.726
nnU-Net [16] | 0.846 | 0.829
Attention U-Net [17] | 0.866 | 0.846
TransUNet [18] | 0.879 | 0.835
R2U-Net [19] | 0.851 | 0.778
| Deeplabv3+ [20] | - | 0.846
| SAM | 0.157 | 0.367
| Adapter-SAM | 0.861 | 0.823
Uncertainty Method | Ensemble U-Net | 0.557 | 0.744
Pro U-Net [8] | 0.602 | 0.682
UGMCS-Net [21] | 0.876 | -
MRNet [22] | 0.878 | 0.849
| UA-SAM(ours) | 0.887 | 0.856

###  3.3 Main results

We compared our UA-SAM model with SOTA segmentation methods, classified into
deterministic methods (U-Net, nnU-Net, Attention U-Net, TransUNet, R2U-Net,
Deeplabv3+, SAM, Adapter-SAM) and uncertainty methods (Ensemble U-Net, Pro
U-Net, UGMCS-Net, MRNet). The segmentation performance was evaluated using the
Dice score, with quantitative results in Table. 1. We got our results by using
the majority vote strategy to deal with samples and multiple-annotated masks
(our model performs the best when it samples 4 times in the LIDC-IDRI dataset
and 3 times in the REFUGE2 dataset). Our model outperformed others with the
highest Dice score of 88.7% on the LIDC-IDRI dataset and 85.6% on the REFUGE2
dataset.

Additionally, in Table. 1, we compared our method with three related methods,
including Pro U-Net, SAM, and Adapter-SAM. It is quite evident that our method
has a significant advantage, even though UA-SAM and probabilistic U-Net both
utilize the probabilistic model to learn uncertainty. Our method still shows
an improvement of 28.5% and 17.4% on two datasets compared with Pro U-Net.
It’s worth mentioning that SAM struggles to cross the domain of these two
datasets, even though it has already demonstrated excellent zero-shot
capabilities in natural image segmentation. Adapter-SAM(similar to MSA[5], but
the position of the Adapter for MSA is inside the Transformer block.) learns
the medical domain knowledge via the Adapter and makes a giant improvement of
70.4% and 45.6% on two datasets compared with SAM, but it is 2.6% and 3.3%
lower than ours. It demonstrates the Uncertainty-aware Adapter we proposed is
more efficient than the normal Adapter.

Table 2: The results of Ablation Study. Note the ’WMS’ represents sample z without being modified by p and means z directly concatenate with p.  | Uncertainty
---
sample z
| Position
---
variant p
WMS | CMSM | LIDC-IDRI | REFUGE2
✓ |  |  |  | 0.862 | 0.824
| ✓ |  |  | 0.866 | 0.827
✓ | ✓ | ✓ |  | 0.868 | 0.828
✓ | ✓ |  | ✓ | 0.887 | 0.856
Figure 2: Visualization results. On the left side is LIDC-IDRI, and on the
right is REFUGE2 (raw and cropped images). The red text indicates the combined
result with the uncertainty region for the best.

Furthermore, we compared several methods based on SAM from the Table. 2, we
can see that there is no significant change in the model’s performance by
utilizing a method similar to Prob U-Net, which directly concatenates the
uncertainty sample z into the Adapter. As we mentioned before, the model might
ignore the uncertain sample. After incorporating the learnable position
variant, the model only shows a slight performance improvement, with the
greatest enhancement occurring when the uncertainty sample and position
variant are added simultaneously. Our method utilizes position variant p as a
condition to modify uncertainty sample z, which achieves improvements of 1.9%
and 2.8%, respectively, on two datasets compared to the method of directly
concatenating uncertainty sample and position variant. This confirms the
effectiveness of the interaction method we proposed in UA-SAM. Fig. 2 shows
the visualization results of UA-SAM and other methods. Our method can output
more diverse segmentation masks closer to the label distribution than others,
which is why UA-SAM performs better than when we conduct a majority vote
strategy for both labels and prediction samples. It also demonstrates that the
Uncertainty-aware Adapter we proposed can help SAM understand the uncertainty
of medical images.

In addition, we retrained multiple UA-SAM models on the REFUGE2 dataset, only
changing the dimensions of the latent space (the previous optimal model on
REFUGE2 was the result of continuous training using the two datasets
mentioned). From Fig. 3(a), We can observe that the optimal performance occurs
when the latent space dimension is 6, which we chose before. Then, we compared
the parameter of the Uncertainty-aware Adapter with the full fine-tune
methods. The parameter count of the Uncertainty-aware Adapter is only 8.08M,
which is much smaller than the others while outperforming others. It shows
that our method is cost-efficient.

Figure 3: Analytical experiments. (a) is the result of changing the latent
space dimension. (b) is the comparison of models parameter.

##  4 Conclusion

In this paper, we propose a cost-efficient fine-tuning method called
Uncertainty-aware Adapter, which can integrate medical domain knowledge and
uncertainty inherent in medical images. By employing the condition-based
interaction method between the Adapter and the sample from probabilistic
latent space, our method significantly improves over the original SAM,
outperforming the SOTA on the LIDC-IDRI dataset and the REFUGE2 dataset.
Furthermore, we clarified our method can output multiple likely segmentation
hypotheses. We believe the uncertainty-aware method for fine-tuning SAM can
make Segment ’Anything’ more reliable, and it is indispensable for downstream
medical tasks such as clinical diagnosis.

## References

  * [1] Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A.C., Lo, W.Y., et al.: Segment anything. arXiv preprint arXiv:2304.02643 (2023)
  * [2] Ma, J., He, Y., Li, F., Han, L., You, C., Wang, B.: Segment anything in medical images. Nature Communications 15(1), 654 (2024)
  * [3] Deng, R., Cui, C., Liu, Q., Yao, T., Remedios, L.W., Bao, S., Landman, B.A., Wheless, L.E., Coburn, L.A., Wilson, K.T., et al.: Segment anything model (sam) for digital pathology: Assess zero-shot segmentation on whole slide imaging. arXiv preprint arXiv:2304.04155 (2023)
  * [4] Gao, Y., Xia, W., Hu, D., Gao, X.: Desam: Decoupling segment anything model for generalizable medical image segmentation. arXiv preprint arXiv:2306.00499 (2023)
  * [5] Wu, J., Fu, R., Fang, H., Liu, Y., Wang, Z., Xu, Y., Jin, Y., Arbel, T.: Medical sam adapter: Adapting segment anything model for medical image segmentation. arXiv preprint arXiv:2304.12620 (2023)
  * [6] Cheng, J., Ye, J., Deng, Z., Chen, J., Li, T., Wang, H., Su, Y., Huang, Z., Chen, J., Jiang, L., et al.: Sam-med2d. arXiv preprint arXiv:2308.16184 (2023)
  * [7] Chen, T., Zhu, L., Deng, C., Cao, R., Wang, Y., Zhang, S., Li, Z., Sun, L., Zang, Y., Mao, P.: Sam-adapter: Adapting segment anything in underperformed scenes. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. pp. 3367–3375 (2023)
  * [8] Kohl, S., Romera-Paredes, B., Meyer, C., De Fauw, J., Ledsam, J.R., Maier-Hein, K., Eslami, S., Jimenez Rezende, D., Ronneberger, O.: A probabilistic u-net for segmentation of ambiguous images. Advances in neural information processing systems 31 (2018)
  * [9] Bhat, I., Pluim, J.P., Viergever, M.A., Kuijf, H.J.: Effect of latent space distribution on the segmentation of images with multiple annotations. Machine Learning for Biomedical Imaging 2(UNSURE2022), 151–171 (Apr 2023)
  * [10] Viviers, C.G., Valiuddin, A.M., de With, P.H., van der Sommen, F.: Probabilistic 3d segmentation for aleatoric uncertainty quantification in full 3d medical data. arXiv preprint arXiv:2305.00950 (2023)
  * [11] He, K., Chen, X., Xie, S., Li, Y., Dollár, P., Girshick, R.: Masked autoencoders are scalable vision learners. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 16000–16009 (2022)
  * [12] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., Polosukhin, I.: Attention is all you need. Advances in neural information processing systems 30 (2017)
  * [13] Knegt, S.: Probabilistic-unet-pytorch (nov 2019), https://github.com/stefanknegt/Probabilistic-Unet-Pytorch
  * [14] Fang, H., Li, F., Wu, J., Fu, H., Sun, X., Son, J., Yu, S., Zhang, M., Yuan, C., Bian, C., et al.: Refuge2 challenge: A treasure trove for multi-dimension analysis and evaluation in glaucoma screening. arXiv preprint arXiv:2202.08994 (2022)
  * [15] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. pp. 234–241. Springer (2015)
  * [16] Isensee, F., Jaeger, P.F., Kohl, S.A., Petersen, J., Maier-Hein, K.H.: nnu-net: a self-configuring method for deep learning-based biomedical image segmentation. Nature methods 18(2), 203–211 (2021)
  * [17] Oktay, O., Schlemper, J., Folgoc, L.L., Lee, M., Heinrich, M., Misawa, K., Mori, K., McDonagh, S., Hammerla, N.Y., Kainz, B., et al.: Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv:1804.03999 (2018)
  * [18] Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., Zhou, Y.: Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04306 (2021)
  * [19] Alom, M.Z., Hasan, M., Yakopcic, C., Taha, T.M., Asari, V.K.: Recurrent residual convolutional neural network based on u-net (r2u-net) for medical image segmentation. arXiv preprint arXiv:1802.06955 (2018)
  * [20] Chen, L.C., Zhu, Y., Papandreou, G., Schroff, F., Adam, H.: Encoder-decoder with atrous separable convolution for semantic image segmentation. In: Proceedings of the European conference on computer vision (ECCV). pp. 801–818 (2018)
  * [21] Yang, H., Wang, Q., Zhang, Y., An, Z., Chen, L., Zhang, X., Zhou, S.K.: Lung nodule segmentation and uncertain region prediction with an uncertainty-aware attention mechanism. IEEE Transactions on Medical Imaging (2023)
  * [22] Ji, W., Yu, S., Wu, J., Ma, K., Bian, C., Bi, Q., Li, J., Liu, H., Cheng, L., Zheng, Y.: Learning calibrated medical image segmentation via multi-rater agreement modeling. In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. pp. 12341–12351 (2021)