# AcceleratedLiNGAM:
Learning Causal DAGs at the speed of GPUs

Victor Akinwande  Carnegie Mellon University  J. Zico Kolter  Carnegie Mellon
University  Bosch Center for AI

###### Abstract

Existing causal discovery methods based on combinatorial optimization or
search are slow, prohibiting their application on large-scale datasets. In
response, more recent methods attempt to address this limitation by
formulating causal discovery as structure learning with continuous
optimization but such approaches thus far provide no statistical guarantees.
In this paper, we show that by efficiently parallelizing existing causal
discovery methods, we can in fact scale them to thousands of dimensions,
making them practical for substantially larger-scale problems. In particular,
we parallelize the LiNGAM method, which is quadratic in the number of
variables, obtaining up to a 32-fold speed-up on benchmark datasets when
compared with existing sequential implementations. Specifically, we focus on
the causal ordering sub-procedure in DirectLiNGAM and implement GPU kernels to
accelerate it. This allows us to apply DirectLiNGAM to causal inference on
large-scale gene expression data with genetic interventions yielding
competitive results compared with specialized continuous optimization methods,
and VarLiNGAM for causal discovery on U.S. stock data.

##  1 Introduction

Causal discovery or inference aims to learn causal interactions in a data-
driven way [Pearl et al., 2000, Imbens and Rubin, 2015]. However, most causal
discovery methods need to establish the causal relationships between every
pair of variables in the data, which results in a complexity that is at least
quadratic in the number of variables [Glymour et al., 2019]. For instance,
DirectLiNGAM Shimizu et al. [2011], Hyvärinen and Smith [2013] works by
recursively performing regression and conditional independence tests between
pairs of variables to establish the causal ordering and then iteratively
regressing each variable on other variables higher in the established causal
order. Although DirectLiNGAM allows for full determination of the causal
structure under minimal assumptions of linearity, non-Gaussian errors,
acyclicity, and the absence of hidden confounders, it has a complexity of
O($d^{3}$) where $d$ is the number of variables.

To address this scalability limitation, a recent trend in the literature has
been to formulate heuristics combined with continuous optimization or deep
learning methods to learn the underlying structure of the data [Lee et al.,
2019, Lopez et al., 2022, Montagna et al., 2023]. However, care must be taken
when ascribing any causal interpretation to such methods. While these methods
enhance computational efficiency, and allow us to apply causal discovery to
large datasets, they rely on restrictive assumptions, are sensitive to hyper-
parameters and more importantly, their identifiability guarantees are often
not established. In-fact, as we show in the sections that follow, NOTEARS
Zheng et al. [2018] \- a widely used method cannot recover the underlying
structure in a simple causal DAG. As such, there is a compelling need to
explore other avenues to achieve scalability of causal discovery methods.

In this paper, we describe an implementation of the LiNGAM analysis, which is
accelerated by parallelization on consumer GPUs. We focus on the DirectLiNGAM
and VarLiNGAM methods, but the ideas presented are easily applicable to other
LiNGAM variants. Parallelization leads to a 32-fold speed-up over the
sequential version, enabling us to apply these methods to large-scale
datasets. We apply DirectLiNGAM to gene expression data with genetic
interventions ($d\approx 1000$) and VarLiNGAM to U.S. stock data ($d=487$). We
observe that there is still a need to explore improved ways to speed up LiNGAM
on GPUs, such as better I/O awareness and increasing the proportion of the
algorithm that is parallelized. Our implementation is open-sourced on
GitHub111https://github.com/Viktour19/culingam.

##  2 Background

Causal Directed Acyclic Graphs (DAGs) are a fundamental structure in causal
inference and graphical models. A causal DAG consists of nodes, each
representing a random variable. These variables can be anything of interest,
such as different attributes, events, or states. The edges in a causal DAG are
acyclic, directed and represent causal relationships between the variables. An
edge from variable $X_{i}$ to $X_{j}$ implies $X_{i}$ has a direct causal
influence on $X_{j}$. Similarly, the absence of a direct edge from $X_{i}$ to
$X_{j}$ indicates that, according to the model, $X_{i}$ does not have a direct
causal effect on $X_{j}$ given the other variables and edges in the model.
Causal DAGs encode the assumptions about the conditional independence of the
variables meaning the joint probability distribution over all possible values
of the variables can be factorized into a product of conditional distributions
given by:

$P(X_{1},X_{2},...,X_{n})=\prod_{i=1}^{n}P(X_{i}|\text{Parents}(X_{i}))$

In this section, we provide a brief overview of Functional Causal Models based
methods for learning causal DAGs, we discuss the GPU execution model and how
it is amenable to such methods, we present the basic LiNGAM analysis and
discuss recent methods in the literature that seek to learn DAGs using
heuristics and continuous optimization.

###  2.1 Causal discovery based on Functional Causal Models (FCM)

FCMs represent a causal effect $Y$ as a function of direct causes $X$, noise
$\epsilon$, and parameter sets $\theta$, assuming that the transformation from
causes to effects is invertible e.g Eqn. 1. For instance, in the linear, non-
Gaussian acyclic model (LiNGAM) Shimizu et al. [2006], the causal direction
can be determined when at most one of the noise term or cause is Gaussian. The
post-nonlinear (PNL) causal model extends this by considering nonlinear
effects such as sensor or measurement errors, and is identifiable in most
cases Zhang and Hyvarinen [2012]. FCMs provide a flexible framework for causal
discovery by capturing the causal asymmetry in the data-generating process and
accommodates various data distributions. When the exact functional form of the
data-generating process is unknown or in cases with discrete variables and
small cardinality, properties of the causal process may be obscured and
precise identification of causal directions remains challenging.

$Y=f(X,\epsilon;\theta_{1})$ (1)

###  2.2 Standard LiNGAM implementation

LiNGAM Shimizu et al. [2006] is characterized by properties that enable
identification of causal relationships between variables. The observed
variables $x_{i}$, $i\in\\{i,\ldots m\\}$ in the dataset can be arranged in a
causal order denoted by $k(i)$. This order implies that no later variable in
the sequence causes any earlier variable, forming a recursive structure that
can be represented as a DAG. This property also ensures there are no cyclic
relationships between the variables. In addition, the value of $x_{i}$ is
determined linearly by the values of the variables before it in the causal
order i.e. $x_{i}=\sum_{k(j)<k(i)}\theta_{ij}x_{j}+\epsilon_{i}+c_{i}$ where
$\theta_{ij}$ represent the strength of the causal effect of $x_{j}$ on
$x_{i}$, $\epsilon_{i}$ is the noise term and $c_{i}$ is an optional constant
term. The final property is that the noise terms $\epsilon_{i}$ are mutually
independent, continuous random variables with non-Gaussian distributions.

Constrained by these assumptions on the data-generating process, the
independence and non-Gaussianity of $\epsilon$ provides valuable information
needed to establish the causal direction say given a pair of variables. This
is because we can easily identify independence between the cause and the error
terms. As such, any causal model that is inconsistent with this identification
is discarded. We illustrate this principle in Figure 1. Given data generated
according a LiNGAM functional causal model as in Eqn. 1, the regression
residual can only be independent of the independent variable in the correct
causal direction.

In practice, datasets consist of more than two variables, so we need a general
method for estimating the causal order for datasets of arbitrary dimensions.
In Shimizu et al. [2011], DirectLiNGAM is presented as a method to achieve
this. DirectLiNGAM begins by identifying a variable not caused by any other
variable (termed exogenous) in the data through its independence from the
residuals of multiple pairwise regressions. The effects of this variable are
then removed from the other variables using least squares regression. Shimizu
et al. [2011] showed that a LiNGAM structure also holds for the residuals of
the regression, meaning that the residuals themselves can be treated as
variables in a LiNGAM model. In addition, the causal ordering of the residuals
corresponds to the causal ordering of the original observed variables.
DirectLiNGAM therefore iteratively applies the same independence and removal
steps to the residuals, finding the “exogenous” residual at each iteration. By
repeatedly identifying and removing the effects of exogenous variables (or
residuals), the model estimates the causal order of the original variables. We
show the implementation to identify the variable at position $k(i)$ in
Algorithm 1.

Algorithm 1 Causal ordering pseudo-implementation.

⬇

1# X : dataset [m, dim]

2# U : indices of variables [dim]

3# _residual: computes the regression residual

4# _diff_mutual_info: computes the difference of the mutual information

5

6def search_causal_order(X, U):

7k_list = np.zeros(len(U))

8for i in U:

9 k = 0

10 for j in U:

11 if i != j:

12 xi = X[:, i]

13 xj = X[:, j]

14 xi_std = (xi - np.mean(xi)) / np.std(xi)

15 xj_std = (xj - np.mean(xj)) / np.std(xj)

16

17 ri_j = _residual(xi_std, xj_std)

18 rj_i = _residual(xj_std, xi_std)

19

20 mi_diff = _diff_mutual_info(xi_std, xj_std, ri_j, rj_i)

21 k += np.minimum(0, mi_diff) ** 2

22

23 k_list[i] = -k

24return U[np.argmax(k_list)]

Figure 1: Illustration of the causal asymmetry principle underpinning LiNGAM.
Given data generated according a LiNGAM functional causal model as in Eqn. 1,
the regression residual can only be independent of the independent variable in
the correct causal direction (top figure). This holds for any distribution of
the noise except Gaussian. Independence is measured using the Mutual
Information (MI).

###  2.3 GPU execution model

The causal ordering procedure (Algorithm 1) is the main computational
bottleneck of the DirectLiNGAM algorithm, accounting for up to $96$% of the
overall wall-clock time (see Figure 2). The pseudo-implementation of the
procedure also makes the limitation of DirectLiNGAM, in practice, glaring. The
algorithm needs to compute statistical measures in an inner loop for each pair
of variables in the data. Such computational structure results in a complexity
quadratic in the number of variables making it difficult to apply on large-
scale datasets. DirectLiNGAM takes $7$ hours to process a dataset with $1$
million observations and $100$ variables on a high-performance AMD EPYC server
CPU (see Figure 2).

On a second look, we observe that each variables pair computation is
independent of the others. This independence is a key characteristic that
makes an algorithm suitable for parallel processing. In-fact, based on this,
successful attempts to parallelize the algorithm on a super-computer have been
made [Matsuda et al., 2022]. However, GPUs present a more accessible approach
and have not be explored. GPUs can perform the same operation on multiple
data-points simultaneously (vectorization), and also handle high arithmetic
intensity computations. Operations involving accumulations can also be
efficiently parallelized on GPUs by using techniques like parallel reduction,
where the work is divided among multiple threads and then combined to get the
final result. This implies vectorized computations like the mean and standard
deviation can be done much faster and computation of residuals and mutual
information differences for different pairs can be done in parallel batches.
Finally, recent work Shahbazinia et al. [2023] has demonstrated that LiNGAM is
amenable to GPU acceleration, although the authors propose using a heuristic
to prune the search procedure. This effectively modifies the algorithm, and
the implementation is not available for us to benchmark.

Figure 2: Benchmark of CPU (sequential) implementation of DirectLiNGAM. Given
data with specified number of samples and dimensions, the causal ordering sub-
procedure accounts for up to $96$% of overall runtime (top-left). It takes $7$
hours on a CPU to process a dataset of $1$ million samples with $100$
variables (top-right). Benchmark of GPU (parallel) implementation of
DirectLiNGAM (bottom-left) and VarLiNGAM (bottom-right). Given data with
specified number of samples and dimensions, the parallel implementation
achieves up to $32$ times speed-up when compared to the sequential
implementation. The benchmark is obtained using an NVIDIA RTX 6000 Ada with
$18\,176$ cores.

###  2.4 Continuous optimization based structure learning

Learning DAGs with GPUs is certainly not a new idea. Continuous optimization
methods for learning DAGs are amenable to acceleration on GPUs using packages
like PyTorch or Tensorflow. Algorithms such as NOTEARS Zheng et al. [2018] and
GOLEM Ng et al. [2020] simultaneously optimize over the DAGs structure and its
parameters by defining a differentiable acyclicity constraint and enabling
end-to-end optimization of a score function over graph adjacency matrices.
NOTEARS minimizes the MSE between the observations and the model predictions:

$\text{MSE}_{X}(\theta)=\frac{1}{m}\|X-X\theta\|_{F}^{2}$

where $\|\cdot\|_{F}$ denotes the Frobenius norm. NOTEARS includes a
constraint defined with a trace exponential function that equals zero if and
only if $\theta$ represents an acyclic graph, and a penalty term
$\lambda\|\theta\|_{1}$ where $\lambda\|.\|_{1}$ is defined element-wise and
$\lambda$ is a hyper-parameter.

GOLEM performs Maximum Likelihood Estimation (MLE) under the assumption of
Gaussian noise terms in the data, and includes both soft acyclicity and
sparsity constraints.

Unfortunately, both methods make restrictive assumptions about the data-
generating process such as equal noise variance across observations, or that
the marginal variance of each variable being strictly larger than its
ancestors in the DAG (termed varsortability) [Reisach et al., 2021, Ng et al.,
2023]. The non-convex nature of the optimization problem in GOLEM often
necessitates careful initialization and sophisticated optimization strategies
to ensure convergence to a meaningful solution. Moreover, neither NOTEARS nor
GOLEM provides identifiability guarantees; they may not perform reliably on
simulated datasets where the true underlying structure is known (see Section
3.1), and other such methods may not converge [Lopez et al., 2022].
Furthermore, the effectiveness of these algorithms is highly sensitive to the
choice of hyper-parameters, such as the sparsity threshold or the specific
loss function employed. Selecting these hyper-parameters is non-trivial and
may require extensive cross-validation or domain expertise, potentially
limiting the practicality of these methods.

##  3 AcceleratedLiNGAM: Analysis, and Extensions

In this section, we discuss the considerations that make acceleration of
Algorithm 1 efficient on GPUs. First, note that the outer loop over $i$ and
the inner loop over $j$ are independent and thus can be parallelized. The
dependency on the results of _residual and _diff_mutual_info for each pair
$(i,j)$ however, necessitates synchronization and memory management.
Therefore, we parallelize over $i$, each block (group of threads) handles a
different $i$ value, computing $\text{k\\_list}[i]$ and within each $i$, we
parallelize over $j$. We also ensure that within the inner loop, operations
are ordered using GPU abstractions.

This parallelization scheme requires up to $dim*(dim-1)$ cores. Profiling the
sequential implementation (Figure 2), Amdahl’s law suggests a theoretical
speed-up of $25$. In the limit of infinite processors, $speedup=1/(1-p)$,
where $p$ is the parallelizable portion of the algorithm ($0.96$ in our case)
but this does not account for the increase in workload size with the number of
processors [Gustafson, 1988]. Finally, we do not need synchronization for
updating $\text{k\\_list}[i]$ since the update order does not matter.

###  3.1 Efficient DirectLiNGAM implementation

We implemented AcceleratedLiNGAM and evaluate on an NVIDIA RTX 6000 Ada. The
implementation optimizes for memory use by performing parallel reductions in
shared memory. Shared memory refers to a fast on-chip memory space shared
among the threads of a block in the GPU. The results show a 32-fold speed-up
of DirectLiNGAM when compared with the sequential implementation (see Figure
2).

To validate that there are no logical errors in our parallel implementation,
we compare the results of applying the sequential implementation with those of
the parallel implementation on simulated data. We report the F1 score, recall,
and structural hamming distance (SHD) over $50$ simulations (different random
seed) in Figure 3.

Let $G=(V,E)$ be a DAG where $V$ is a set of vertices representing variables
and $E$ is a set of directed edges representing causal relationships between
variables. Vertices are connected such that each vertex $v_{i}$ at level $l$
may have parents from the set of nodes at level $l-1$: $\forall v_{i}\in
V,\exists l_{i}\in\\{0,1\\}\text{ such that }(v_{j},v_{i})\in E\Rightarrow
l_{j}=l_{i}-1$. Given $G$, we generate data such that the strength of the
causal effect $\theta\sim\mathcal{N}(0,1)$, and the noise terms
$\epsilon_{i}\sim\text{Uniform}(0,1),\quad\text{for }i=1,2,...,m$. Comparison
of the sequential and parallel implementation of DirectLiNGAM on this
simulated data show that they produce the exact same result, and recover the
true causal graph accurately (see Figure 3).

We evaluate NOTEARS on similarly simulated data selecting the best performance
across a grid: $\\{0.001,0.005,0.01,0.05,0.1\\}$ of $\lambda$ values. We
obtain an F1 score of $0.79$ $\pm$ $0.2$, Recall of $0.69$ $\pm$ $0.2$ and SHD
of $2.52$ $\pm$ $1.67$. This shows that even on data where the causal
influences are simple, NOTEARS does not perform well.

Figure 3: Comparison of parallel and sequential implementation of
DirectLiNGAM. We simulate data according to a linear FCM with $10\,000$
samples, and $10$ variables. Both implementations produce the exact same
result (top). Benchmark of CPU (sequential) implementation of VarLiNGAM. Given
data with specified number of samples and dimensions, the causal ordering sub-
procedure of DirectLiNGAM accounts for up to $96$% of overall runtime (bottom-
left). It takes $7$ hours on a CPU to process a dataset of $1$ million samples
with $100$ variables (bottom-right).

###  3.2 Extension: Efficient VarLiNGAM implementation

The basic LiNGAM analysis can be extended to auto-regressive modeling. In
Hyvärinen et al. [2010], VarLiNGAM is introduced. The key idea is to combine
FCMs and vector autoregressive (VAR) models to capture both instantaneous and
lagged influences among multiple time series. The model is expressed as:
$x(t)=\sum_{\tau=0}^{k}\theta_{\tau}x(t-\tau)+\epsilon(t)$ where $x(t)$ is the
observed time series at time $t$, $\theta_{\tau}$ is the matrix of causal
effects with time lag $\tau$ and $\epsilon(t)$ is the noise term (innovation).
Similar to all LiNGAM analyses, the model assumes the noise terms are mutually
independent, non-Gaussian, and the matrix $\theta_{\tau}$ corresponding to
instantaneous effects forms an acyclic graph.

On a high-level, VarLiNGAM works by first estimating coefficients of the VAR
model using standard auto-regressive modeling techniques, and then transforms
the estimated VAR coefficients based on the causal adjacency matrix estimated
using DirectLiNGAM thereby accounting for the direct causal relationships
while isolating the indirect effects captured by the VAR coefficients. We
refer the reader to Hyvärinen et al. [2010] for more details. The key thing to
note here is the same algorithm in Algorithm 1 dominates the runtime of
VarLiNGAM (See Figure 3) and so we obtain a similar speed-up of $30$ with the
GPU implementation.

###  3.3 Low level CUDA implementation details

All the benchmark results are run on an NVIDIA A6000 Ada GPU with $48$ GB of
memory, $18\,176$ cores and $91.1$ TFLOPS single-precision performance. The
register file size is 64K 32-bit registers per Streaming Multiprocessor (SM),
the maximum number of registers per thread is $255$, the maximum number of
thread blocks per SM is $16$, the shared memory capacity per SM is $100$ KB,
and the maximum shared memory per thread block is $99$ KB.

A key optimization is how the kernel uses shared memory for intermediate
reduction results within each block. This pattern is sensitive to the number
of threads because it assumes that the shared memory array is fully populated,
and after storing in shared memory, we perform a reduction within the block.
We set shared memory per thread block to $96$ KB. We implement warp tiling
(for optimal latency, and efficient synchronization) and notice a $20$% speed-
up but due to the non-associative nature of floating-point operations, naive
implementation of parallel reduction may lead to rounding errors which we
observe in the residual computation. As such, we leave this optimization for
future work.

Finally, optimizing GPU implementations of machine learning methods for I/O
awareness has achieved significant success recently, especially when combined
with the use of cores optimized for fast matrix multiplication (Tensor cores)
Dao et al. [2022], Fu et al. [2023]. A profile of our GPU kernels reveals that
the two most time-consuming operations are poll and pthread_cond_timedwait,
both accounting for about $50$% of the execution time. These operations are
associated with waiting for I/O operations or synchronization primitives,
indicating substantial room for optimization to enhance I/O awareness, and
efficient synchronization. The remainder of the DirectLiNGAM and VarLiNGAM
implementations, which are not parallelized, involve several regression
analysis. Although, we utilize heavily optimized libraries like numpy and
scikit-learn for these regressions, there remains potential for speed-up
through parallelism with Tensor cores.

##  4 Experiments

###  4.1 AcceleratedLiNGAM to gene expression data with genetic interventions

We experiment with the causal learning of gene regulatory networks from gene
expression data, with genetic interventions, as detailed in Friedman et al.
[2000], Pe er et al. [2001], Lopez et al. [2022]. This approach is enabled by
a single-cell RNA sequencing technique known as Perturb-Seq, as described by
Dixit et al. [2016] . Perturb-Seq allows for targeted genetic interventions
and the subsequent measurement of their effects on the complete gene
expression profiles in hundreds of thousands of individual cells using single-
cell RNA-seq.

Our experimental setup mirror that of Lopez et al. [2022], where the Perturb-
CITE-seq dataset Frangieh et al. [2021] contains expression profiles from
$218\,331$ melanoma (cancer) cells, after interventions targeting each of
$249$ genes. For each gene in the genome, measurement from a single-cell
combines the identity of the intervention (the target gene) along with a count
vector, that is the expression level of a particular gene. The dataset
includes patient-derived melanoma cells with the same genetic interventions
but exposed to three conditions: co-culture with T cells derived from the
patient’s tumor ($73\,114$ cells) (these can recognize and kill melanoma
cells), interferon (IFN)-$\gamma$ treatment ($87\,590$ cells) and control
($57\,627$ cells). We retain cells from $20$% of the interventions as a test
set. The dimensions (samples, dim) of the train set for co-culture, IFN, and
control datasets are ($65\,164$, $964$), ($75\,443$, $964$), and ($50\,539$,
$961$) respectively.

We applied AcceleratedLiNGAM to each of the three datasets. Since LiNGAM
analysis does not involve causal inference, after obtaining the weighted
adjacency matrix, we apply standard Variational Inference (VI) methods to
obtain both the interventional NLL (I-NLL) and the mean absolute error (I-MAE)
across held-out interventions. Specifically, we use Stein VI Liu and Wang
[2016] implemented in the Pyro package where we defined a model such that
variables without direct causal influence on another variable are leaf nodes
and otherwise are latent nodes with priors $\sim\mathcal{N}(0,1)$. We generate
$200$ posterior samples, and optimize for $5000$ iterations. DirectLiNGAM
itself is not sensitive to the random seed.

On a high-level, Stein VI involves approximating the target distribution
$p(x)$ by iteratively updating a set of particles. This process minimizes the
KL divergence between the approximating distribution $q(x)$ and the target
$p(x)$, which is achieved through the application of smooth transforms. The
key insight is to apply a perturbation to the identity map,
$T(x)=x+\epsilon\phi(x)$, where $\phi(x)$ is a smooth function that
characterizes the perturbation direction, and $\epsilon$ is a small scalar
representing the perturbation magnitude. This approach ensures $T$ is an
injective (one-to-one) map, maintaining the full rank of the Jacobian of $T$,
and thereby guaranteeing the invertibility required for the change of
variables formula to hold.

Table 1: Comparison of DirectLiNGAM with VI and DCD-FG method on the Perturb-CITE-seq datasets. We obtain the I-NLL (nll) and I-MAE (mae) on all three datasets. Lower values are better. | Co-culture | IFN | Control
---|---|---|---
| nll | mae | nll | mae | nll | mae
DirectLiNGAM | $1.5$ | $0.7$ | $1.5$ | $0.9$ | $3$ | $1.6$
DCD-FG ($\approx$) | $1.1$ | $0.7$ | $1.2$ | $0.7$ | $1.1$ | $0.7$

We compare the results from DirectLiNGAM, combined with VI, with those
obtained using DCD-FG introduced in Lopez et al. [2022]. DCD-FG works by
combining a parameterized distribution over factor directed graphs with a
hybrid likelihood model, optimized with an acyclicity constraint and was
applied to the Perturb-CITE-seq datasets. We observe the I-MAE to be lower or
about the same on the co-culture dataset (one leaf variable), and slightly
higher on the IFN and control datasets (one and two leaf variables
respectively) (see Table LABEL:tbl:infmetrics. We note DCD-FG is a continuous
optimization based structure learning method and prone to many of the issues
we have previously discussed but more pertinently, the results presented in
Lopez et al. [2022] do have quite a bit of variance. We also observe the I-NLL
of DirectLiNGAM to be slightly higher on all datasets. While this may seem
like worse performance, it is interesting to note how the control dataset
(i.e., no interventions) has the same I-NLL as the other two datasets with
DCD-FG. However, with DirectLiNGAM, the I-NLL of the co-culture and IFN
datasets are similar, but the control I-NLL is much higher. Since there is no
ground truth data, it is difficult to determine if DCD-FG is overfitting but
it seems likely.

###  4.2 AcceleratedLiNGAM to stock data with auto-regression modelling

We also experiment with learning the causal relationships among U.S. stock
indices. We obtain hourly adjusted closing price data for the stocks in the
S&P 500 from Yahoo Finance from January 2022 to December 2023. We preprocess
the data lightly, filling missing values using time-based linear
interpolation, removing indices with any remaining missing values, and
transforming the original series to a new series that is stationary with first
differencing before applying VarLiNGAM. This process leaves us with data from
$487$ indices. Finally, we consider a lag of $1$ in the VAR model and use the
estimator provided by the statsmodels package.

Figure 4: In and out degree distribution of the adjacency matrix obtained
using VarLiNGAM on S&P 500 hourly data. We observe some level of symmetry
between in-degree and out-degree, and the fairly uniform distribution across a
range of degrees indicates that there are no prominent hubs that significantly
stand out.

In Figure 4, we show the in and out degrees distribution of the adjacency
matrix we obtain ($\theta_{0}$). The distribution of in-degrees and out-
degrees is similar, which could imply that the instantaneous causal
relationships have some form of balance in terms of influences made and
received by the nodes. The graph also has two leaf nodes corresponding to USB:
U.S. Bancorp and FITB: Fifth Third Bancorp which is interesting because both
companies are holding companies i.e. they do not engage directly in business
operations themselves but rather hold the assets of and oversee the management
of their subsidiaries. Consequently, the discovery that they are influenced by
other indices without influencing any in return is particularly relevant.

We obtain the top 5 exerting casual influence nodes, and the top 5 receiving
influence nodes by measuring the total causal effects. We show the indices
that correspond to both categories in Table LABEL:tbl:infl. The top exerting
indices predominantly features companies involved in manufacturing, retail,
and services with a direct consumer focus, such as homebuilding (NVR), auto
parts retail (AZO), and food services (CMG). These companies are often more
directly impacted by consumer spending habits and economic cycles, given their
close ties to discretionary spending (e.g., travel, dining out, and home
purchases). Whereas, the top receiving indices operate in industries subject
to more stringent regulations or external factors, like CenterPoint Energy
(CNP) in the utilities sector, or that play a crucial role in the information
and packaging industries, such as News Corp (NWSA) and Amcor plc (AMCR), where
changes in consumer behavior or regulatory environments can significantly
impact their operations.

Table 2: We show the indices that correspond to both the top 5 exerting casual
influence nodes, and the top 5 receiving influence nodes from the discovered
causal graph. Top Exerting
---
NVRτ-1 | Homebuilding
AZOτ-1 | Auto Parts
CMGτ-1 | Restaurants
BKNGτ-1 | Online Travel
MTDτ-1 | Precision Instruments
Top Receiving
NWSAτ | Diversified Media
CNPτ | Utilities
FOXAτ | Broadcasting
AMCRτ-1 & AMCRτ | Packaging

##  5 Conclusion

Machine learning practitioners often encounter use-cases where predictive
performance alone does not suffice. Typical examples include in healthcare,
genomics, or predictive justice applications where the understanding of cause-
and-effect relationships is paramount. In healthcare, accurately predicting
patient outcomes is crucial, but understanding the causal factors behind
diseases can lead to more effective treatments and health policies [Räisänen
et al., 2006, Barros et al., 2022]. Similarly, in genomics, identifying the
causal relationships between genetic markers and diseases is vital for
developing targeted therapies and personalized medicine approaches [Burgess et
al., 2018, Lopez et al., 2022]. In predictive justice applications, it is not
enough to predict recidivism rates; understanding the causal factors can
inform more effective rehabilitation programs and fairer justice policies
[Khademi and Honavar, 2020, Loeffler and Nagin, 2022]. In all these domains,
the focus is not solely on prediction but also on uncovering the underlying
causal mechanisms. By addressing the scalability limitations of causal
discovery methods with statistical guarantees, we aim to enable the widespread
application of causal inference in large-scale data analysis.

In this paper, we accelerate LiNGAM methods, achieving significant speed-up
compared to sequential implementations. We do not change the algorithm; hence,
the identifiability guarantees of AcceleratedLiNGAM remain the same as those
of the original methods. We note that there is still room for further speed-up
through the use of specialized CUDA techniques and by leveraging the GPU
memory hierarchy. Also, the performance of GPUs continues to increase over
time. Thus, we anticipate that future iterations of AcceleratedLiNGAM could
see even greater improvements in computational efficiency. This work opens up
the application of LiNGAM to a wider range of applications, and by open-
sourcing our implementation, we aim to spur the development of these
improvements and applications by the broader community.

##  6 Acknowledgments

We thank Daiyaan Arfeen for valuable discussions. Victor Akinwande was
supported by funding from the Bosch Center for Artificial Intelligence.

## References

  * Barros et al. [2022] Vesna Barros, Itay Manes, Victor Akinwande, Celia Cintas, Osnat Bar-Shira, Michal Ozery-Flato, Yishai Shimoni, and Michal Rosen-Zvi.  A causal inference approach for estimating effects of non-pharmaceutical interventions during covid-19 pandemic.  _Plos one_ , 17(9):e0265289, 2022.
  * Burgess et al. [2018] Stephen Burgess, Christopher N Foley, and Verena Zuber.  Inferring causal relationships between risk factors and outcomes from genome-wide association study data.  _Annual review of genomics and human genetics_ , 19:303–327, 2018.
  * Dao et al. [2022] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  _Advances in Neural Information Processing Systems_ , 35:16344–16359, 2022.
  * Dixit et al. [2016] Atray Dixit, Oren Parnas, Biyu Li, Jenny Chen, Charles P Fulco, Livnat Jerby-Arnon, Nemanja D Marjanovic, Danielle Dionne, Tyler Burks, Raktima Raychowdhury, et al.  Perturb-seq: dissecting molecular circuits with scalable single-cell rna profiling of pooled genetic screens.  _cell_ , 167(7):1853–1866, 2016.
  * Frangieh et al. [2021] Chris J Frangieh, Johannes C Melms, Pratiksha I Thakore, Kathryn R Geiger-Schuller, Patricia Ho, Adrienne M Luoma, Brian Cleary, Livnat Jerby-Arnon, Shruti Malu, Michael S Cuoco, et al.  Multimodal pooled perturb-cite-seq screens in patient models define mechanisms of cancer immune evasion.  _Nature genetics_ , 53(3):332–341, 2021.
  * Friedman et al. [2000] Nir Friedman, Michal Linial, Iftach Nachman, and Dana Pe’er.  Using bayesian networks to analyze expression data.  In _Proceedings of the fourth annual international conference on Computational molecular biology_ , pages 127–135, 2000.
  * Fu et al. [2023] Daniel Y Fu, Hermann Kumbong, Eric Nguyen, and Christopher Ré.  Flashfftconv: Efficient convolutions for long sequences with tensor cores.  _arXiv preprint arXiv:2311.05908_ , 2023.
  * Glymour et al. [2019] Clark Glymour, Kun Zhang, and Peter Spirtes.  Review of causal discovery methods based on graphical models.  _Frontiers in genetics_ , 10:524, 2019.
  * Gustafson [1988] John L Gustafson.  Reevaluating amdahl’s law.  _Communications of the ACM_ , 31(5):532–533, 1988\.
  * Hyvärinen and Smith [2013] Aapo Hyvärinen and Stephen M Smith.  Pairwise likelihood ratios for estimation of non-gaussian structural equation models.  _The Journal of Machine Learning Research_ , 14(1):111–152, 2013.
  * Hyvärinen et al. [2010] Aapo Hyvärinen, Kun Zhang, Shohei Shimizu, and Patrik O Hoyer.  Estimation of a structural vector autoregression model using non-gaussianity.  _Journal of Machine Learning Research_ , 11(5), 2010.
  * Imbens and Rubin [2015] Guido W Imbens and Donald B Rubin.  _Causal inference in statistics, social, and biomedical sciences_.  Cambridge University Press, 2015.
  * Khademi and Honavar [2020] Aria Khademi and Vasant Honavar.  Algorithmic bias in recidivism prediction: A causal perspective.  In _34th AAAI Conference on Artificial Intelligence, AAAI 2020_ , pages 13839–13840. AAAI press, 2020.
  * Lee et al. [2019] Hao-Chih Lee, Matteo Danieletto, Riccardo Miotto, Sarah T Cherng, and Joel T Dudley.  Scaling structural learning with no-bears to infer causal transcriptome networks.  In _PACIFIC SYMPOSIUM ON BIOCOMPUTING 2020_ , pages 391–402. World Scientific, 2019.
  * Liu and Wang [2016] Qiang Liu and Dilin Wang.  Stein variational gradient descent: A general purpose bayesian inference algorithm.  _Advances in neural information processing systems_ , 29, 2016.
  * Loeffler and Nagin [2022] Charles E Loeffler and Daniel S Nagin.  The impact of incarceration on recidivism.  _Annual Review of Criminology_ , 5:133–152, 2022.
  * Lopez et al. [2022] Romain Lopez, Jan-Christian Hütter, Jonathan K. Pritchard, and Aviv Regev.  Large-scale differentiable causal discovery of factor graphs.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Matsuda et al. [2022] Kazuhito Matsuda, Kouji Kurihara, Kentaro Kawakami, Masafumi Yamazaki, Fuyuka Yamada, Tsuguchika Tabaru, and Ken Yokoyama.  Accelerating lingam causal discovery with massive parallel execution on supercomputer fugaku.  _IEICE TRANSACTIONS on Information and Systems_ , 105(12):2032–2039, 2022.
  * Montagna et al. [2023] Francesco Montagna, Nicoletta Noceti, Lorenzo Rosasco, Kun Zhang, and Francesco Locatello.  Scalable causal discovery with score matching.  _arXiv preprint arXiv:2304.03382_ , 2023.
  * Ng et al. [2020] Ignavier Ng, AmirEmad Ghassami, and Kun Zhang.  On the role of sparsity and dag constraints for learning linear dags.  _Advances in Neural Information Processing Systems_ , 33:17943–17954, 2020.
  * Ng et al. [2023] Ignavier Ng, Biwei Huang, and Kun Zhang.  Structure learning with continuous optimization: A sober look and beyond.  _arXiv preprint arXiv:2304.02146_ , 2023.
  * Pe er et al. [2001] Dana Pe er, Aviv Regev, Gal Elidan, and Nir Friedman.  Inferring subnetworks from perturbed expression profiles.  _BIOINFORMATICS-OXFORD-_ , 17:S215–S224, 2001.
  * Pearl et al. [2000] Judea Pearl et al.  Models, reasoning and inference.  _Cambridge, UK: CambridgeUniversityPress_ , 19(2):3, 2000.
  * Räisänen et al. [2006] Ulla Räisänen, Marie-Jet Bekkers, Paula Boddington, Srikant Sarangi, and Angus Clarke.  The causation of disease–the practical and ethical consequences of competing explanations.  _Medicine, Health Care and Philosophy_ , 9:293–306, 2006\.
  * Reisach et al. [2021] Alexander Reisach, Christof Seiler, and Sebastian Weichwald.  Beware of the simulated dag! causal discovery benchmarks may be easy to game.  _Advances in Neural Information Processing Systems_ , 34:27772–27784, 2021.
  * Shahbazinia et al. [2023] Amirhossein Shahbazinia, Saber Salehkaleybar, and Matin Hashemi.  Paralingam: Parallel causal structure learning for linear non-gaussian acyclic models.  _Journal of Parallel and Distributed Computing_ , 176:114–127, 2023.
  * Shimizu et al. [2006] Shohei Shimizu, Patrik O Hoyer, Aapo Hyvärinen, Antti Kerminen, and Michael Jordan.  A linear non-gaussian acyclic model for causal discovery.  _Journal of Machine Learning Research_ , 7(10), 2006.
  * Shimizu et al. [2011] Shohei Shimizu, Takanori Inazumi, Yasuhiro Sogawa, Aapo Hyvarinen, Yoshinobu Kawahara, Takashi Washio, Patrik O Hoyer, Kenneth Bollen, and Patrik Hoyer.  Directlingam: A direct method for learning a linear non-gaussian structural equation model.  _Journal of Machine Learning Research-JMLR_ , 12(Apr):1225–1248, 2011.
  * Zhang and Hyvarinen [2012] Kun Zhang and Aapo Hyvarinen.  On the identifiability of the post-nonlinear causal model.  _arXiv preprint arXiv:1205.2599_ , 2012.
  * Zheng et al. [2018] Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing.  Dags with no tears: Continuous optimization for structure learning.  _Advances in neural information processing systems_ , 31, 2018.


# AutoProSAM: Automated Prompting SAM for 3D Multi-Organ Segmentation

Chengyin Li 1, Rafi Ibn Sultan1, Prashant Khanduri1, Yao Qiang1, Indrin
Chetty2, Dongxiao Zhu1
1Department of Computer Science, Wayne State University<EMAIL_ADDRESS>
2Henry Ford Health System

###### Abstract

Segment Anything Model (SAM) is one of the pioneering prompt-based foundation
models for image segmentation and has been rapidly adopted for various medical
imaging applications. However, in clinical settings, creating effective
prompts is notably challenging and time-consuming, requiring the expertise of
domain specialists such as physicians. This requirement significantly
diminishes SAM’s primary advantage—its interactive capability with end
users—in medical applications. Moreover, recent studies have indicated that
SAM, originally designed for 2D natural images, performs suboptimally on 3D
medical image segmentation tasks. This subpar performance is attributed to the
domain gaps between natural and medical images and the disparities in spatial
arrangements between 2D and 3D images, particularly in multi-organ
segmentation applications. To overcome these challenges, we present a novel
technique termed AutoProSAM. This method automates 3D multi-organ CT-based
segmentation by leveraging SAM’s foundational model capabilities without
relying on domain experts for prompts. The approach utilizes parameter-
efficient adaptation techniques to adapt SAM for 3D medical imagery and
incorporates an effective automatic prompt learning paradigm specific to this
domain. By eliminating the need for manual prompts, it enhances SAM’s
capabilities for 3D medical image segmentation and achieves state-of-the-art
(SOTA) performance in CT-based multi-organ segmentation tasks.

##  1 Introduction

Recently, prompt-based foundation models in computer vision, such as the
Segment Anything Model (SAM) [15], has shown impressive performance and
versatility across various semantic segmentation tasks [41] based on user-
provided prompts. These foundation models offer new potential for multi-organ
segmentation in medical imagery, an area often limited by the availability and
quality of segmentation masks. Unlike custom-designed transformer models such
as UNETR [9], SwinUNETR [34], and FocalUNETR [18], which are typically trained
on limited patient samples and masks, foundation models like SAM benefit from
training on millions of 2D natural images and billions of masks. With
appropriate prompts, these foundation models can be adapted to a variety of
downstream tasks.

Figure 1: Challenges associated with using SAM for medical image segmentation
include (A) a T-SNE plot of embeddings encoded by SAM’s image encoder,
showcasing differences between medical image datasets such as AMOS [12] and
BTCV [16], and natural image datasets like ADE20K [45] and COCO [19]; (B) The
requirement of manually generated prompts from domain experts for SAM-based
medical image segmentation.

Despite their extensive training, large foundation models such as SAM struggle
to adapt to medical images. This is primarily due to the inherent differences
between medical and natural images and the challenges associated with
effective prompting (as shown in Fig. 1). Recent efforts have attempted to
extend the success of SAM to medical image segmentation tasks [43, 39, 21, 7,
33]. However, the demonstrated performance has exhibited reduced precision and
stability, particularly in more intricate segmentation tasks in 3D medical
images characterized by smaller sizes, irregular shapes, and lower contrast
properties compared to 2D natural images [7]. Fig. 1(A) illustrates this
oversight of substantial domain disparities between natural and medical
images. Furthermore, unlike natural images, medical images in various formats
are not easily interpretable by general users. The prompts, intended to guide
the foundation models to focus on a particular location, become a performance
bottleneck as users struggle to identify the region of interest. Therefore,
generating effective prompts for these models necessitates the involvement of
domain experts and physicians. This reliance on labor-intensive manually
generated prompts [5, 33] hampers its successful application, particularly in
multi-organ medical image segmentation tasks, as demonstrated in Fig. 1(B). As
a result, these foundation models have been shown to underperform in medical
image segmentation tasks compared to state-of-the-art (SOTA) models
specifically designed for medical imagery [7].

To address the issues mentioned above, we introduce a novel AutoProSAM method
for transitioning SAM from 2D natural image segmentation to 3D multi-organ
image segmentation without relying on human prompts. We start by designing
extensive modifications for the image encoder at the input level, allowing the
original 2D transformer to handle volumetric 3D inputs effectively. This
optimization enhances the reusability of pre-trained weights through a
parameter-efficient fine-tuning method, achieving a balance by leveraging
SAM’s vast knowledge from training on a large 2D natural image dataset and
adapting it with minimal tuning for 3D medical images. To solve the prompting
issue and adapt SAM for end users, we developed an Auto Prompt Generator (APG)
module that automatically learns the necessary prompts for the segmentation
process. This innovation eliminates the time-consuming and complex manual
prompt generation process, particularly for multi-organ medical image
segmentation tasks. Our extensive experimentation on CT-based multi-organ
segmentation datasets, which includes comprehensive comparisons with state-of-
the-art methods such as nnUNet [11] and recent adapters, shows a significant
performance improvement over existing techniques.

The main contributions are three-fold:

  * •

We introduce the Auto Prompt Generator to eliminate the complex and laborious
manual prompting process, simplifying multi-organ medical image segmentation
tasks for non-domain-expert users.

  * •

We adapt the 2D SAM model for 3D multi-organ segmentation with parameter-
efficient fine-tuning, bridging the gap between 2D natural images and 3D
medical images.

  * •

We conduct extensive experiments and analysis on three public and one private
multi-organ segmentation datasets, demonstrating that the proposed AutoProSAM
achieves superior performance in medical image segmentation tasks compared to
the SOTAs.

##  2 Related Work

###  2.1 Foundation Computer Vision Models

Recent advancements in computer vision have led to the development of numerous
pre-trained backbones using various algorithms and datasets [6]. As deep
learning models evolve, most modern vision frameworks now adhere to the pre-
training and fine-tuning paradigm [23]. Recently, there has been significant
interest among computer vision researchers in large, adaptable foundation
models leveraging pre-training techniques such as self-supervised learning
[14], contrastive learning [38], and language-vision pre-training [30].
Notably, the SAM model [15], pre-trained on a dataset of over 11 million
images, has emerged as a prompt-based foundation model for natural image
segmentation. SAM demonstrates impressive zero-shot capabilities, effectively
segmenting diverse subjects in real-world environments through an interactive
and prompt-driven approach. Additionally, SEEM [48], developed around the same
time as SAM, introduces a more comprehensive prompting scheme to facilitate
semantic-aware open-set segmentation. Furthermore, DINOv2 [26] focuses on
scaling up the pre-training of a Vision Transformer (ViT) model in terms of
data and model size, aiming to produce versatile visual features that simplify
the fine-tuning of downstream tasks.

###  2.2 Parameter-efficient Model Fine-Tuning

Given the extensive utilization of foundation models, the concept of
parameter-efficient fine-tuning has garnered significant attention. Existing
methods for efficient fine-tuning can be categorized into three groups [3].
First, addition-based methods incorporate lightweight adapters [27, 37] or
prompts [20, 13] into the original model, focusing solely on adjusting these
parameters. Second, specification-based methods [40, 8] concentrate on
selecting a small subset of the original parameters for tuning. Third,
reparameterization-based methods [10] leverage low-rank matrices to
approximate parameter updates. Recently, some researchers have extended pre-
trained image models to encompass video comprehension [27] or volumetric
segmentation [37]. However, these methods treat the additional dimension as a
”word group” and employ specialized modules to aggregate information along the
word dimension. In contrast, our work considers all three dimensions
simultaneously and directly adapts the trained transformer block to capture
patterns in 3D medical image inputs.

###  2.3 Adapting SAM to Medical Images

Among the efforts to adapt SAM for medical images, MedSAM [21] focused on
refining the SAM decoder using prompts generated from label masks across more
than 30 medical image datasets, resulting in improved performance over zero-
shot predictions. Zhang _et al_. [42] optimized a low-rank fine-tuning
approach for the SAM encoder and combined it with SAM decoder training for
abdominal segmentation tasks. Tal _et al_. [33] redesigned SAM’s prompt
encoder by using the original 2D medical image as input. SAM-Med2D [1] fine-
tuned SAM with a large collection of 2D medical image datasets. Conversely, Wu
_et al_. [39] pre-trained the entire encoder through self-supervised learning
on collected datasets for downstream medical image segmentation tasks. 3DSAM-
Adapter [7] extended SAM to 3D medical imagery, targeting tumor segmentation
tasks for each tumor type individually, a time- and labor-intensive process
due to the demanding prompting requirements. Similarly, Wang _et al_. [35]
addressed 3D medical imagery by retraining SAM from scratch, replacing the 2D
components to develop an in-house SAM for medical imagery. Despite some
progress, existing methods often overlook the 3D patterns of medical images or
rely heavily on extensive pre-training.

###  2.4 Automatic Prompts Generation

To automate the prompt generation process for SAM, research has primarily
taken two directions. One direction, exemplified by works such as [28, 17],
involves generating 2D bounding box prompts for either 2D image segmentation
or each 2D slice in 3D images. While these methods are a good starting point,
they do not address the 3D spatial relationships inherent in 3D images.
Another direction, seen in efforts such as [33, 24, 2], involves training an
auxiliary network to generate prompts for SAM from the input image. However,
these models fail to fully utilize SAM’s extensive pre-trained capabilities.
Our approach overcomes these limitations by leveraging features from SAM’s
encoder to generate automated prompts. Consequently, we address the
shortcomings of previous works and unlock SAM’s full potential in CT-based 3D
multi-organ segmentation tasks.

##  3 Method

Figure 2: (A) The overall architecture of the AutoProSAM, (B) the design of
the Depth Adapter module, which utilizes parameter-efficient model fine-
tuning, and (C) the architecture of the Auto Prompt Generator, featuring a
U-Net-like encoder-decoder design.

In this section, we explain how to modify the original SAM architecture,
initially developed for 2D natural images, to work with 3D volumetric medical
images for segmentation tasks. We begin by offering a brief overview of the
SAM framework (as depicted in Fig. 2), followed by a detailed explanation of
the adjustments made to the image encoder, prompt encoder, and mask decoder.

###  3.1 SAM Architecture

SAM [15] is a prompt-driven image segmentation framework, renowned for its
exceptional performance in segmenting natural images. The architecture of SAM
comprises three key components: an image encoder, a prompt encoder, and a mask
decoder. The image encoder utilizes the Vision Transformer (ViT) [4, 29] to
transform original images into one-time image embeddings. The prompt encoder
skillfully converts various types of prompts – including foreground/background
points, rough boxes or masks, clicks, text, or any information indicating the
target of segmentation – into compact embeddings. These embeddings from both
the image and prompt encoders are then seamlessly integrated by the mask
decoder to produce accurate segmentation masks.

Although SAM has been successful in 2D natural image segmentation, it faces
significant challenges when applied to 3D volumetric medical imagery. A key
issue is the model’s reliance on slice-wise predictions, which fail to
consider the inter-slice spatial context, thereby impacting its suitability
for complex medical tasks. Additionally, the inherent domain disparities
between medical and natural images contribute to their performance limitations
in medical applications. To effectively address these challenges and optimize
SAM for medical imaging tasks, tailored adaptation, and fine-tuning of the
model become essential.

###  3.2 Handling 3D Medical Inputs

To enhance SAM’s capability in processing 3D medical images, we propose an
adaptation strategy named AutoProSAM, as illustrated in Fig. 2A. This strategy
has two primary objectives: firstly, to enable the model to directly learn 3D
spatial patterns, and secondly, to ensure continuity by inheriting most
parameters from the pre-trained model while introducing easily adjustable
incremental parameters. The detailed design is elaborated as follows.

####  3.2.1 Positional Encoding Enhancement

In the pre-trained ViT model, there is a lookup table of size $C\times H\times
W$ for positional encoding, where $C$ represents the channel, $H$ the height,
and $W$ the width. Additionally, we initialize a tunable lookup table of size
$C\times D$ (with $D$ representing the depth of a volume patch) with zeros. To
obtain the positional encoding of a 3D point $(d,h,w)$, we add the embedding
from the frozen lookup table with $(h,w)$ to the embedding from the tunable
lookup table with $(d)$.

####  3.2.2 Patch Embedding Adjustments

We utilize a combination of $1\times k\times k$ and $k\times 1\times 1$ 3D
convolutions to approximate the effect of a $k\times k\times k$ convolution
(_e.g_., kernel size $k=14$). The $1\times k\times k$ convolution is
initialized with the weights from a pre-trained 2D convolution and remains
unchanged during the fine-tuning phase. As for the newly introduced $k\times
1\times 1$ 3D convolution, we apply depth-wise convolution to decrease the
number of parameters that need adjustment. This approach helps in managing the
complexity of the model.

####  3.2.3 Adapting Attention Block

The attention blocks can be directly adjusted to accommodate 3D features. In
the case of 2D inputs, the size of the queries is $[B,HW,C]$, which can be
effortlessly modified to $[B,DHW,C]$ for 3D inputs, while retaining all the
pre-trained weights. We implement sliding-window mechanisms akin to those in
the SwinUNETR [34] to mitigate the memory impact resulting from the increase
in dimensions. This approach aids in optimizing the model’s performance while
managing memory requirements.

####  3.2.4 Bottleneck Modifications

Given that convolution layers are generally easier to optimize than
transformers, we replace 2D convolutions in the bottleneck with 3D
counterparts and train them from scratch to improve performance.

By making the above adjustments, we can smoothly transition the 2D ViT into a
3D ViT, reusing most parameters. However, fully fine-tuning the 3D ViT can be
resource-intensive. To address this, we propose using a lightweight adapter
approach for efficient fine-tuning. The adapter comprises a down-projection
linear layer and an up-projection linear layer, represented as
$\operatorname{Adapter}(\mathbf{X})=\mathbf{X}+\operatorname{Act}(\mathbf{X}W_{\mathrm{Down}})W_{\mathrm{Up}}$.
Here, $\mathbf{X}\in\mathbb{R}^{N\times C}$ is the original feature
representation, $W_{\mathrm{Down}}\in\mathbb{R}^{C\times N^{\prime}}$ and
$W_{\mathrm{Up}}\in\mathbb{R}^{N^{\prime}\times C}$ are down-projection and
up-projection layers, and $\mathrm{Act}(\cdot)$ is the activation function. To
enhance 3D spatial awareness, we include a depth-wise 3D convolution after the
down-projection layer, as shown in Fig. 2B. This enhancement improves the
adapter’s utilization of 3D spatial cues.

Throughout the training phase, we exclusively adjust the parameters of
convolutions, depth adapters, and normalization layers, while maintaining all
other parameters in a frozen state. This frozen approach enhances memory
efficiency during training. Fine-tuning the adapter and normalization layers
aids in bridging the gap between natural images and medical images, enabling
the model to adapt more effectively to the medical image domain.

###  3.3 Auto Prompt Generator

SAM initially applies positional embedding to both prompt (_e.g_., points or
boxes) and image, ensuring that prompt and image embeddings at the same
position share identical positional encoding. Subsequently, the prompt
embedding engages in cross-attention with the image embedding, evolving from
positional to semantic attributes. However, this cross-attention mechanism,
while effective in 2D settings, can lead to over-smoothing issues when applied
to 3D feature maps [7]. Adapting to 3D settings can significantly increase the
number of tokens, potentially resulting in a uniform probability distribution.

Furthermore, prompt-based segmentation may not be well-suited for real-world
applications due to two primary reasons. Firstly, it is time-consuming for
multi-class scenarios, as seen in many medical image segmentation challenges
requiring the segmentation of multiple classes simultaneously. This becomes
particularly challenging for small or closely located organs. Secondly,
segmentation quality heavily relies on the precision of the prompts, yet
creating accurate prompts demands domain-specific expertise, often unavailable
to non-expert users. These limitations diminish the practicality of prompt-
based methods.

To overcome these limitations, we propose to use an Auto Prompt Generator
instead of positional encoding to represent the prompt. The whole process is
illustrated in Fig. 2C. Instead of using manually generated points or bounding
boxes, we directly take the output feature map after the last block of
attention and depth adapter operation. This APG follows a fully convolutional
neural (FCN) based encoder-decoder design that resembles 3D UNet [32]. This
generator boasts a lightweight structure, leveraging 3D-based convolution
operations, and can be effortlessly learned from scratch. This enables precise
prompt generation tailored to different medical segmentation tasks. Notably,
it eliminates the need for additional manually generated prompts, simplifying
and expediting the multi-class medical image segmentation tasks.

###  3.4 Mask Decoder

The mask decoder in SAM was originally designed to be lightweight, primarily
using stacks of 2D convolution layers. In our updated version, we have
replaced these 2D convolutions with 3D convolutions to enable direct 3D mask
generation. The initial design of the decoder, which lacks progressive
upsampling and skip connections, is effective for natural images that
typically feature large objects with distinct boundaries. However, in the
context of volumetric medical image segmentation, it is well-recognized that
U-shaped networks with skip connections at multiple levels are crucial for
success [11, 34]. This is due to the smaller size of objects in medical images
and their often indistinct boundaries, necessitating networks that can
preserve higher-resolution details for better discrimination. Therefore, the
use of U-shaped architectures with skip connections becomes imperative.

To address this need while keeping the design lightweight, we incorporate a
multi-layer aggregation mechanism (MLAM) [44, 7] in our decoder. We utilize
the intermediate feature maps from stages 1-4 (as depicted in Fig. 2A) of the
image encoder, along with the prompt embedding from the APG, to enrich the
mask feature map without compromising efficiency. For improved resolution
detail, we upsample the mask feature map to match the original resolution,
then concatenate it with the original image. This concatenated map is fused
using another 3D convolution to generate the final mask. This method
effectively combines high-resolution details with the original image data in
the mask generation process. We have streamlined the original SAM by focusing
solely on a targeted downstream task and omitting features like multi-task
generation and ambiguity awareness. The backbone of the mask decoder primarily
comprises lightweight 3D convolutional layers, which are known for their ease
of optimization. Consequently, we train all decoder parameters from scratch.

##  4 Experiments

###  4.1 Datasets

A total of 4 CT datasets, including 3 public and 1 institutional dataset, are
used to evaluate the performance of models for the 3D multi-organ segmentation
task.

BTCV: Beyond the Cranial Vault (BTCV) abdomen challenge dataset [16] includes
30 subjects with abdominal CT scans. In this dataset, 13 organs are annotated
by interpreters under the supervision of Vanderbilt University Medical Center
radiologists. The multi-organ segmentation task is framed as a 13-class
segmentation, which includes large organs such as liver, spleen, kidneys, and
stomach; vascular tissues of esophagus, aorta, inferior vena cava, splenic,
and portal veins; small anatomies of gallbladder, pancreas, and adrenal
glands. For the subjects, 24 scans are for training and 6 for testing.

AMOS: We actually use the CT parts in the original AMOS dataset [12], not
including their MRI parts. The publicly accessible AMOS-CT dataset [12]
consists of 200 multi-contrast abdominal CT scans for training and 100 for
testing. These scans are annotated for 15 anatomies, including spleen, right
kidney, left kidney, gallbladder, esophagus, liver, stomach, aorta, inferior
vena cava, pancreas, right adrenal gland, left adrenal gland, duodenum,
bladder, and prostate/uterus.

CT-ORG: CT Organ Segmentation Dataset (CT-ORG) [31] comprises 100 CT scans,
each of which includes ground-truth contours of 5 different organs (lung,
bones, liver, kidney and bladder). Of these, the first 19 CT scans are
exclusively reserved for testing purposes whereas the remaining 81 scans are
used for training.

Institutional Pelvic: In addition to public datasets, we also test the
performance of the segmentation models on a private dataset. An institutional
pelvic dataset of 300 prostate cancer patients is retrospectively selected,
and each contains 3 manually generated contours (prostate, bladder, and
rectum). The 300 cases are randomly split into a training set of 225 cases, a
validation set of 30, and a testing set of 45.

###  4.2 Implementation Details

The AutoProSAM model is trained using the AdamW optimizer with a warm-up
exponential decay scheduler of 200 epochs, in which we incorporate the first 5
epochs for warmup. The segmentation experiments use a batch size of 1 per GPU
with a patch size of $128\times 128\times 128$. Default initial learning rate
of $5e^{-4}$, momentum of 0.9 and decay of $1e^{-5}$. The framework is
implemented in MONAI 1.2 and PyTorch 2.0. Models are trained on a server with
eight NVIDIA A100 cards. All the baseline models are trained from scratch
following their default settings. For each dataset, we perform interpolation
for a fixed voxel spacing over all scans within this dataset, then intensity
normalization, sampling over positive and negative patches, and finally on-
the-fly data augmentation with flip, rotation, intensity scaling, and
intensity shifting. Foreground and background patches are randomly sampled at
a $1:1$ ratio. We select the best model with the highest Dice metric on the
validation dataset or the final epoch’s model for a dataset without a
predefined validation split, with empirical evidence of negligible gain via a
5-fold cross-validation strategy. Baseline models are trained from scratch
following their default settings. For inference, an overlapping area ratio of
0.75 is applied via the sliding window strategy.

We provide detailed pre-processing information about the interpolation to a
fixed voxel spacing, HU ranges, and intensity normalization in Tab. 1.

Dataset | Voxel Spacing (mm) | HU Range | Normalization
---|---|---|---
BTCV | $1.0\times 1.0\times 1.5$ | $[-125,275]$ | $[0,1]$
AMOS | $1.0\times 1.0\times 1.5$ | [-991, 362] | Subtract 50,
divide by 141
CT-ORG | Isotropic 2.0 | $[-1000,1000]$ | $[-1,1]$
Institutional | Isotropic 1.5 | $[-50,150]$ | $[0,1]$
Pelvic
Table 1: Pre-processing settings for each dataset.

A sampling of $128\times 128\times 128$ voxel patches is done for all four
datasets, with a positive-to-negative patches ratio of 1:1 sampling for AMOS,
CT-ORG, and BTCV datasets.

###  4.3 Loss Function

For training the AutoProSAM (as shown in Fig. 2A), a combination of Dice loss
and Cross-Entropy loss is used to assess the alignment between the predicted
mask and the ground truth on a pixel-wise basis. The objective function for
the segmentation head is defined as follows:

$\mathcal{L}_{\mathrm{Seg}}=\mathcal{L}_{\mathrm{Dice}}(\hat{p}_{i},g_{i})+\mathcal{L}_{\mathrm{CE}}(\hat{p}_{i},g_{i}),$
(1)

where $\hat{p}_{i}$ represents the predicted voxel probabilities from the main
task, and $g_{i}$ represents the ground truth mask for an input volume $i$.
The predicted probabilities, $\hat{p}_{i}$, result from applying the
AutoProSAM to the input 3D volume for the main task.

###  4.4 Evaluation Metrics

We utilize the Dice similarity coefficient (Dice) and Normalized Surface
Distance (NSD) [25] as metrics to evaluate the segmentation performance. Dice
metric measures the extent of agreement between two volumes with the following
formula:

$\mathrm{Dice}=\frac{2\sum_{i=1}^{I}g_{i}\hat{g}_{i}}{\sum_{i=1}^{I}g_{i}+\sum_{i=1}^{I}\hat{g}_{i}},$
(2)

where $g$ and $\hat{g}$ denote the ground truth and the predicted voxel
values. The NSD metric quantifies the overlap between ground truth and
predicted surfaces (with a fixed tolerance). For both metrics, 0 denotes the
worst match between two structures, and 100% means the perfect match.

Figure 3: Qualitative visualizations compare our AutoProSAM with baseline
methods using three subjects from public datasets (Rows 1-3) and two subjects
from the private Institutional Pelvic dataset (Rows 4-5). Enhanced areas in
these visualizations illustrate improvements in segmenting the left kidney
(light red) and pancreas (beige). Additionally, segmentation masks are shown
for the prostate (blue), bladder (green), and rectum (red).

##  5 Results & Discussion

Model | Tuned Params. | BTCV | AMOS | CT-ORG | Institutional Pelvic
---|---|---|---|---|---
(M) |  mDice $\uparrow$ |  mNSD $\uparrow$ |  mDice $\uparrow$ |  mNSD $\uparrow$ |  mDice $\uparrow$ |  mNSD $\uparrow$ |  mDice $\uparrow$ |  mNSD $\uparrow$
nnUNet [11] | 31.18 | 84.34 | 73.21 | 87.43 | 77.12 | 85.51 | 75.60 | 88.04 | 83.87
AttUNet [36] | 40.79 | 83.67 | 74.35 | 83.75 | 72.23 | 85.13 | 74.64 | 88.06 | 84.12
UNet++ [47] | 28.45 | 79.33 | 72.84 | 84.32 | 73.34 | 84.87 | 74.21 | 87.53 | 82.65
nnFormer [46] | 150.14 | 83.51 | 71.65 | 84.52 | 70.06 | 82.75 | 70.31 | 87.60 | 82.13
UNETR [9] | 93.02 | 85.47 | 74.35 | 77.24 | 60.58 | 83.13 | 71.56 | 86.68 | 81.54
SwinUNETR [34] | 62.83 | 86.58 | 75.26 | 86.19 | 74.83 | 84.34 | 72.32 | 89.12 | 83.54
Ours | 26.53 | 87.15 | 78.83* | 88.65* | 79.41* | 85.12 | 76.37* | 91.30* | 84.35*
Table 2: Comparison of the overall performance between four SOTAs and our
AutoProSAM on four datasets. The best results are highlighted in bold font.
(*: $p<0.01$, with Wilcoxon signed-rank test to all SOTAs)

In this section, we would like to evaluate the performance of our proposed
AutoProSAM concerning both the task-specific SOTAs and the SAM-adopting-based
methods. The effectiveness of different modules for our designs will be
evaluated as well.

###  5.1 Comparison with SOTAs

We extensively compare our model with the SOTA 3D medical image segmentation
approaches, including the most recent Transformer-based methods including
UNETR [9], SwinUNETR [34], and nnFormer [46], as well as CNN-based methods
such as nnUNet [11], AttUNet [36], and UNet++ [47]. As indicated in Tab. 2, we
observe that the proposed AutoProSAM generally outperforms other SOTA methods
in both Dice and NSD metrics across all four datasets.

Significant improvements are particularly noticeable in the BTCV dataset, with
up to a 3% increase in the average Dice score and a 3% to 7% improvement in
the average NSD metric. As the training sample size increases for the AMOS
dataset compared to BTCV, the AutoProSAM performs even better, achieving up to
a 14% improvement in Dice and a 2% to 19% improvement in NSD. For the CT-ORG
dataset, our model achieves the highest NSD and a competitive Dice score
compared to other baselines. Notably, the AutoProSAM consistently demonstrates
a robust performance in real-world scenarios, as seen in its consistently
superior performance on the private Institutional Pelvic dataset.

The superiority of our AutoProSAM model is further highlighted by qualitative
results. As depicted in Fig. 3, qualitative comparisons of the predicted masks
from various segmentation models reveal that our AutoProSAM yields visually
superior mask predictions, particularly in terms of more accurate boundary
delineation, when compared to its SOTA counterparts. In summary, while the
original SAM displayed relatively weaker performance in medical image
segmentation tasks as compared to SOTA methods [22], our adapted design
registers significant enhancements and generalization to 3D medical image
segmentation.

We also randomly selected samples from the CT-ORG dataset for more visual
comparisons of different methods. As depicted in Fig. 4, all methods generally
achieve precise prediction masks for organs that are relatively easy to
segment, such as lungs and bones. However, our AutoProSAM also provides more
consistently accurate predictions than other baselines for organs like the
liver, kidney, and bladder.

###  5.2 Comparison with SAM-based Methods

We further compare our method with existing SAM-based methods for multi-class
medical image segmentation, including the original SAM and MedSAM [21]. MedSAM
can be implemented with either full fine-tuning or partial fine-tuning, where
only the prompt encoder and mask decoder are updated using point prompts.
Other adaptation methods were not considered, as they either require
collecting a large number of medical images for fine-tuning (_e.g_., Medical-
SAM-Adapter [39], SAM-Med2D [1]), or they do not specifically target multi-
organ segmentation tasks (_e.g_., 3DSAM-Adapter [7]).

Model | Tuning | BTCV | AMOS
---|---|---|---
Option |  mDice $\uparrow$ |  mNSD $\uparrow$ |  mDice $\uparrow$ |  mNSD $\uparrow$
SAM [15] | None | 54.86 | - | 49.31 | -
MedSAM [21] | P&M | 80.65 | 66.82 | 70.28 | 59.61
MedSAM [21] | Full | 84.57 | 73.76 | 83.10 | 70.26
Ours | P&M | 87.15 | 78.83 | 88.65 | 79.41

Table 3: Comparison with SAM-based methods, where the best results are
highlighted in bold. ‘P&M’ indicates only fine-tuning the prompt encoder and
mask decoder, ‘Full’ denotes full fine-tuning, and ‘None’ represents no fine-
tuning.

The outcomes are meticulously outlined in Tab. 3, underscoring how our
adaptation strategy surpasses all existing methods. Notably, our approach
outshines the second-best technique by a margin of 3% in terms of BTCV
segmentation Dice, 5% for NSD, and 6% concerning AMOS segmentation Dice,
accompanied by a substantial 9% for NSD. Impressively, it even outperforms the
complete fine-tuning variant of MedSAM, even when considering parameter-
efficient fine-tuning. These results effectively validate our hypothesis that
parameters pre-trained on 2D images can be effectively harnessed to grasp 3D
spatial features with only minor adjustments. Moreover, our approach of
treating all dimensions equivalently emerges as a superior strategy compared
to interpreting the depth dimension as a distinct group in the context of
medical image segmentation.

Figure 4: Qualitative visualizations compare our AutoProSAM with baseline
methods over the CT-ORG dataset.

###  5.3 Ablation Study

####  5.3.1 Effects of APG and MLAM

The APG in our AutoProSAM refines the feature maps from the final stage of the
attention block in the image encoder, paralleling the way the original SAM’s
prompt encoder processes manual prompts. Without APG, these feature maps are
directly fed into the mask decoder. Beyond APG, a Multi-Layer Aggregation
Mechanism (MLAM) is also applied for the integration of information from
various stages of the encoder and APG to the final mask decoder. To evaluate
the effectiveness of both APG and MLAM, we conducted a comparative analysis
using the BTCV dataset by removing these components. The findings, presented
in Tab. 4, reveal significant improvements in both Dice and NSD metrics when
APG and MLAM are employed. These results highlight the substantial
contribution of APG and MLAM to enhancing the efficacy of medical image
segmentation.

APG | MLAM |  mDice $\uparrow$ |  mNSD $\uparrow$
---|---|---|---
✓ | ✓ | 87.15 | 78.83
✗ | ✓ | 85.23 | 74.12
✓ | ✗ | 84.37 | 73.54
✗ | ✗ | 81.14 | 70.33
Table 4: Ablation study evaluating the effectiveness of the Automated Prompt
Generator (APG) and Multi-Level Attention Mechanism (MLAM) designs in our
AutoProSAM.

####  5.3.2 Inference Time Analysis

As demonstrated in Table 5, our AutoProSAM, despite incorporating a heavy ViT
encoder, achieves an inference speed only marginally slower than that of
nnUNet [11] and SwinUNETR [34]. Notably, it significantly outperforms both
UNETR [9] and nnFormer [46] in terms of speed, highlighting its remarkable
efficiency among these SOTA methods.

Model | Inference Time (sec. per scan)
---|---
nnUNet [11] | 3.74
nnFormer [46] | 11.73
UNETR [9] | 58.93
SwinUNETR [34] | 4.24
AutoProSAM (ours) | 7.06
Table 5: Averaged inference time for different models in BTCV experiments with
a patch size of $128\times 128\times 128$.

##  6 Conclusion

In this paper, we introduced the AutoProSAM, a novel approach designed to
enhance SAM for 3D multi-organ medical image segmentation. By extending SAM
from 2D natural images to 3D medical images, our method addresses domain gaps
and spatial disparities through parameter-efficient fine-tuning and an Auto
Prompt Generator (APG) that automates prompt creation, eliminating the need
for manual input. Extensive experiments on public and private CT-based
datasets demonstrated the AutoProSAM’s superior performance over state-of-the-
art models, evidenced by higher Dice and NSD scores. Our ablation studies also
highlighted the importance of the APG and Multi-Layer Aggregation Mechanism
(MLAM) in improving segmentation accuracy.

Limitations. While our AutoProSAM achieved superior performance in multi-organ
segmentation tasks, certain limitations should also be noted. Our approach
does not consider the pre-training of the image encoder via collecting a huge
amount of medical images, nor does it focus on creating a universally
adaptable model for multiple different datasets. Further related research
would be worthwhile.

## References

  * [1] Junlong Cheng, Jin Ye, Zhongying Deng, Jianpin Chen, Tianbin Li, Haoyu Wang, Yanzhou Su, Ziyan Huang, Jilong Chen, Lei Jiang, et al.  Sam-med2d.  arXiv preprint arXiv:2308.16184, 2023.
  * [2] Zachery Morton Colbert, Daniel Arrington, Matthew Foote, Jonas Gårding, Dominik Fay, Michael Huo, Mark Pinkham, and Prabhakar Ramachandran.  Repurposing traditional u-net predictions for sparse sam prompting in medical image segmentation.  Biomedical Physics & Engineering Express, 10(2):025004, 2024.
  * [3] Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al.  Parameter-efficient fine-tuning of large-scale pre-trained language models.  Nature Machine Intelligence, 5(3):220–235, 2023.
  * [4] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [5] Yifan Gao, Wei Xia, Dingdu Hu, and Xin Gao.  Desam: Decoupling segment anything model for generalizable medical image segmentation.  arXiv preprint arXiv:2306.00499, 2023.
  * [6] Micah Goldblum, Hossein Souri, Renkun Ni, Manli Shu, Viraj Prabhu, Gowthami Somepalli, Prithvijit Chattopadhyay, Mark Ibrahim, Adrien Bardes, Judy Hoffman, et al.  Battle of the backbones: A large-scale comparison of pretrained models across computer vision tasks.  arXiv preprint arXiv:2310.19909, 2023.
  * [7] Shizhan Gong, Yuan Zhong, Wenao Ma, Jinpeng Li, Zhao Wang, Jingyang Zhang, Pheng-Ann Heng, and Qi Dou.  3dsam-adapter: Holistic adaptation of sam from 2d to 3d for promptable medical image segmentation.  arXiv preprint arXiv:2306.13465, 2023.
  * [8] Demi Guo, Alexander M Rush, and Yoon Kim.  Parameter-efficient transfer learning with diff pruning.  arXiv preprint arXiv:2012.07463, 2020.
  * [9] Ali Hatamizadeh, Yucheng Tang, Vishwesh Nath, Dong Yang, Andriy Myronenko, Bennett Landman, Holger R Roth, and Daguang Xu.  Unetr: Transformers for 3d medical image segmentation.  In Proceedings of the IEEE/CVF winter conference on applications of computer vision, pages 574–584, 2022.
  * [10] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.  Lora: Low-rank adaptation of large language models.  arXiv preprint arXiv:2106.09685, 2021.
  * [11] Fabian Isensee, Paul F Jaeger, Simon AA Kohl, Jens Petersen, and Klaus H Maier-Hein.  nnu-net: a self-configuring method for deep learning-based biomedical image segmentation.  Nature methods, 18(2):203–211, 2021.
  * [12] Yuanfeng Ji, Haotian Bai, Chongjian Ge, Jie Yang, Ye Zhu, Ruimao Zhang, Zhen Li, Lingyan Zhanng, Wanling Ma, Xiang Wan, et al.  Amos: A large-scale abdominal multi-organ benchmark for versatile medical image segmentation.  Advances in Neural Information Processing Systems, 35:36722–36732, 2022.
  * [13] Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim.  Visual prompt tuning.  In European Conference on Computer Vision, pages 709–727. Springer, 2022.
  * [14] Longlong Jing and Yingli Tian.  Self-supervised visual feature learning with deep neural networks: A survey.  IEEE transactions on pattern analysis and machine intelligence, 43(11):4037–4058, 2020.
  * [15] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al.  Segment anything.  arXiv preprint arXiv:2304.02643, 2023.
  * [16] Bennett Landman, Zhoubing Xu, J Igelsias, Martin Styner, T Langerak, and Arno Klein.  Miccai multi-atlas labeling beyond the cranial vault–workshop and challenge.  In Proc. MICCAI Multi-Atlas Labeling Beyond Cranial Vault—Workshop Challenge, volume 5, page 12, 2015.
  * [17] Wenhui Lei, Xu Wei, Xiaofan Zhang, Kang Li, and Shaoting Zhang.  Medlsam: Localize and segment anything model for 3d medical images.  arXiv preprint arXiv:2306.14752, 2023.
  * [18] Chengyin Li, Yao Qiang, Rafi Ibn Sultan, Hassan Bagher-Ebadian, Prashant Khanduri, Indrin J. Chetty, and Dongxiao Zhu.  Focalunetr: A focal transformer for boundary-aware segmentation of ct images, 2023.
  * [19] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick.  Microsoft coco: Common objects in context.  In Computer Vision–ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pages 740–755. Springer, 2014.
  * [20] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig.  Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing.  ACM Computing Surveys, 55(9):1–35, 2023.
  * [21] Jun Ma and Bo Wang.  Segment anything in medical images.  arXiv preprint arXiv:2304.12306, 2023.
  * [22] Maciej A Mazurowski, Haoyu Dong, Hanxue Gu, Jichen Yang, Nicholas Konz, and Yixin Zhang.  Segment anything model for medical image analysis: an experimental study.  Medical Image Analysis, page 102918, 2023.
  * [23] Bonan Min, Hayley Ross, Elior Sulem, Amir Pouran Ben Veyseh, Thien Huu Nguyen, Oscar Sainz, Eneko Agirre, Ilana Heintz, and Dan Roth.  Recent advances in natural language processing via large pre-trained language models: A survey.  ACM Computing Surveys, 2021.
  * [24] Saiyang Na, Yuzhi Guo, Feng Jiang, Hehuan Ma, and Junzhou Huang.  Segment any cell: A sam-based auto-prompting fine-tuning framework for nuclei segmentation.  arXiv preprint arXiv:2401.13220, 2024.
  * [25] Stanislav Nikolov, Sam Blackwell, Alexei Zverovitch, Ruheena Mendes, Michelle Livne, Jeffrey De Fauw, Yojan Patel, Clemens Meyer, Harry Askham, Bernardino Romera-Paredes, et al.  Deep learning to achieve clinically applicable segmentation of head and neck anatomy for radiotherapy.  arXiv preprint arXiv:1809.04430, 2018.
  * [26] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al.  Dinov2: Learning robust visual features without supervision.  arXiv preprint arXiv:2304.07193, 2023.
  * [27] Junting Pan, Ziyi Lin, Xiatian Zhu, Jing Shao, and Hongsheng Li.  St-adapter: Parameter-efficient image-to-video transfer learning.  Advances in Neural Information Processing Systems, 35:26462–26477, 2022.
  * [28] Sumit Pandey, Kuan-Fu Chen, and Erik B Dam.  Comprehensive multimodal segmentation in medical imaging: Combining yolov8 with sam and hq-sam models.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 2592–2598, 2023.
  * [29] Yao Qiang, Chengyin Li, Prashant Khanduri, and Dongxiao Zhu.  Interpretability-aware vision transformer.  arXiv preprint arXiv:2309.08035, 2023.
  * [30] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In International conference on machine learning, pages 8748–8763. PMLR, 2021.
  * [31] Blaine Rister, Darvin Yi, Kaushik Shivakumar, Tomomi Nobashi, and Daniel L Rubin.  Ct-org, a new dataset for multiple organ segmentation in computed tomography.  Scientific Data, 7(1):381, 2020.
  * [32] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  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, pages 234–241. Springer, 2015.
  * [33] Tal Shaharabany, Aviad Dahan, Raja Giryes, and Lior Wolf.  Autosam: Adapting sam to medical images by overloading the prompt encoder.  arXiv preprint arXiv:2306.06370, 2023.
  * [34] Yucheng Tang, Dong Yang, Wenqi Li, Holger R Roth, Bennett Landman, Daguang Xu, Vishwesh Nath, and Ali Hatamizadeh.  Self-supervised pre-training of swin transformers for 3d medical image analysis.  In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20730–20740, 2022.
  * [35] Haoyu Wang, Sizheng Guo, Jin Ye, Zhongying Deng, Junlong Cheng, Tianbin Li, Jianpin Chen, Yanzhou Su, Ziyan Huang, Yiqing Shen, et al.  Sam-med3d.  arXiv preprint arXiv:2310.15161, 2023.
  * [36] Sihan Wang, Lei Li, and Xiahai Zhuang.  Attu-net: attention u-net for brain tumor segmentation.  In International MICCAI brainlesion workshop, pages 302–311. Springer, 2021.
  * [37] Wenxuan Wang, Jiachen Shen, Chen Chen, Jianbo Jiao, Yan Zhang, Shanshan Song, and Jiangyun Li.  Med-tuning: Exploring parameter-efficient transfer learning for medical volumetric segmentation.  arXiv preprint arXiv:2304.10880, 2023.
  * [38] Xiao Wang and Guo-Jun Qi.  Contrastive learning with stronger augmentations.  IEEE transactions on pattern analysis and machine intelligence, 45(5):5549–5560, 2022.
  * [39] Junde Wu, Rao Fu, Huihui Fang, Yuanpei Liu, Zhaowei Wang, Yanwu Xu, Yueming Jin, and Tal Arbel.  Medical sam adapter: Adapting segment anything model for medical image segmentation.  arXiv preprint arXiv:2304.12620, 2023.
  * [40] Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  arXiv preprint arXiv:2106.10199, 2021.
  * [41] Chunhui Zhang, Li Liu, Yawen Cui, Guanjie Huang, Weilin Lin, Yiqian Yang, and Yuehong Hu.  A comprehensive survey on segment anything model for vision and beyond.  arXiv preprint arXiv:2305.08196, 2023.
  * [42] Kaidong Zhang and Dong Liu.  Customized segment anything model for medical image segmentation.  arXiv preprint arXiv:2304.13785, 2023.
  * [43] Yichi Zhang and Rushi Jiao.  How segment anything model (sam) boost medical image segmentation?  arXiv preprint arXiv:2305.03678, 2023.
  * [44] Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, et al.  Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6881–6890, 2021.
  * [45] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba.  Scene parsing through ade20k dataset.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 633–641, 2017.
  * [46] Hong-Yu Zhou, Jiansen Guo, Yinghao Zhang, Xiaoguang Han, Lequan Yu, Liansheng Wang, and Yizhou Yu.  nnformer: Volumetric medical image segmentation via a 3d transformer.  IEEE Transactions on Image Processing, 2023.
  * [47] Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang.  Unet++: Redesigning skip connections to exploit multiscale features in image segmentation.  IEEE transactions on medical imaging, 39(6):1856–1867, 2019.
  * [48] Xueyan Zou, Jianwei Yang, Hao Zhang, Feng Li, Linjie Li, Jianfeng Gao, and Yong Jae Lee.  Segment everything everywhere all at once.  arXiv preprint arXiv:2304.06718, 2023.