# Negative binomial count splitting
for single-cell RNA sequencing data

Anna Neufeld1,2∗, Joshua Popp3, Lucy L. Gao4, Alexis Battle3,5, and Daniela
Witten1,6
1 Department of Statistics, University of Washington, Seattle, WA, USA
2 Public Health Sciences Division, Fred Hutchinson Cancer Center, Seattle, WA,
USA
3 Department of Biomedical Engineering, Johns Hopkins University, Baltimore,
Maryland, USA
4 Department of Statistics, University of British Columbia, Vancouver, BC,
Canada
5 Department of Computer Science, Department of Genetic Medicine, and Malone
Center
for Engineering in Healthcare, Johns Hopkins University, Baltimore, Maryland,
USA
6 Department of Biostatistics, University of Washington, Seattle, WA, USA
*email<EMAIL_ADDRESS>
###### Abstract

The analysis of single-cell RNA sequencing (scRNA-seq) data often involves
fitting a latent variable model to learn a low-dimensional representation for
the cells. Validating such a model poses a major challenge. If we could
sequence the same set of cells twice, we could use one dataset to fit a latent
variable model and the other to validate it. In reality, we cannot sequence
the same set of cells twice. _Poisson count splitting_ was recently proposed
as a way to work backwards from a single observed Poisson data matrix to
obtain independent Poisson training and test matrices that could have arisen
from two independent sequencing experiments conducted on the same set of
cells. However, the Poisson count splitting approach requires that the
original data are exactly Poisson distributed: in the presence of any
overdispersion, the resulting training and test datasets are not independent.
In this paper, we introduce _negative binomial count splitting_ , which
extends Poisson count splitting to the more flexible negative binomial
setting. Given an $n\times p$ dataset from a negative binomial distribution,
we use Dirichlet-multinomial sampling to create two or more independent
$n\times p$ negative binomial datasets. We show that this procedure
outperforms Poisson count splitting in simulation, and apply it to validate
clusters of kidney cells from a human fetal cell atlas.

##  1 Introduction

A single-cell RNA sequencing (scRNA-seq) dataset involving $n$ cells and $p$
genes can be written as a matrix $X\in\mathbb{Z}_{\geq 0}^{n\times p}$, where
entry $X_{ij}$ is the number of unique molecular identifiers from the $i$th
cell that map to the $j$th gene. It is common to model $X$ as a realization
from a random variable $\mathbb{X}$, and assume that

$\operatorname{\text{E}}\left[\mathbf{X}_{ij}\right]=\gamma_{i}\Lambda_{ij},\text{
with }g(\Lambda)=L\beta^{\top}\text{ for }L\in\mathbb{R}^{n\times
K^{*}},\beta\in\mathbb{R}^{p\times K^{*}},$ (1)

for some link function $g(\cdot)$ [Sarkar and Stephens, 2021]. In (1),
$\gamma=(\gamma_{1},\ldots,\gamma_{n})^{T}$ stores cell-specific _size
factors_ , which reflect technical variation in sequencing depth between
cells. The matrix $\Lambda$ represents the biological variation of interest,
which, after applying a link function $g(\cdot)$, has rank $K^{*}$ for some
$K^{*}\leq\min(n,p)$.

Fitting the model (1)–that is, obtaining estimates $\hat{L}(X)$ and
$\hat{\beta}(X)$ of $L$ and $\beta$–may be of interest for a number of
reasons. For example, we may wish to denoise the data by replacing the size-
factor normalized dataset $\mathrm{diag}(\gamma)^{-1}X$ with its low-rank
estimate $g^{-1}\left(\hat{L}(X)\hat{\beta}(X)^{T}\right)$ [Eraslan et al.,
2019, Townes et al., 2019], or we may wish to interpret $\hat{L}(X)$ as a
measure of an unobserved aspect of cell state, e.g. cell type or position
along a developmental trajectory [Aizarani et al., 2019, Van den Berge et al.,
2020]. After fitting the model (1), we typically want to perform some type of
model validation or inference.

######  Example 1.

We want to assess the quality of our low-rank approximation
$\mathrm{diag}(\gamma)^{-1}X\approx
g^{-1}\left(\hat{L}(X)\hat{\beta}(X)^{T}\right)$ [Sarkar and Stephens, 2021,
Batson et al., 2019].

######  Example 2.

We want to identify genes that are associated with $\hat{L}(X)$ [Aizarani et
al., 2019, Van den Berge et al., 2020, Zhang et al., 2019].

######  Example 3.

We want to know whether we would obtain a similar estimate $\hat{L}(X)$ on an
independent realization of $\mathbb{X}$ drawn from the same distribution [Cao
et al., 2020, Lange et al., 2004, Ullmann et al., 2022]. We refer to this as
“reproducibility.”

In Example 1, _because we estimated $L$ and $\beta$ on the data $X$, we cannot
re-use $X$ to assess model fit_ [Hastie et al., 2009]. In Example 2, _because
we estimated $L$ on the data $X$, we cannot re-use $X$ to test for
association_ [Gao et al., 2022]. In Example 3, _we only have access to one
dataset_ , so it is unclear how to proceed. While very specialized approaches
are available to overcome these challenges in specific instantiations of
Example 1 [Fu and Perry, 2020, Grabski et al., 2023], Example 2 [Gao et al.,
2022, Chen and Witten, 2023, Zhang et al., 2019, Chung and Storey, 2015], and
Example 3 [Tibshirani and Walther, 2005, Lange et al., 2004], in this paper we
will provide a much more flexible framework for model validation or inference
after fitting (1), which will be applicable to all three examples.

To illustrate the problem, we generate a toy data matrix $X\in\mathbb{Z}_{\geq
0}^{100\times 2}$ with elements drawn independently from a negative binomial
distribution with mean $5$ and variance $10$. This is a special case of (1)
with $\gamma_{1}=\ldots=\gamma_{n}=1$, $K^{*}=1$, $L=\mathbb{1}_{100}$, and
$\beta=[5,5]^{T}$. The data are shown in Figure 1(a). To illustrate Example 1,
we apply $k$-means clustering to the data to estimate $K$ clusters for a range
of values of $K$. Though there is one true cluster in this example (all 100
cells are homogenous), the mean squared error (MSE; defined in (2) in Section
4) computed on the same data used for clustering is monotone decreasing in $K$
(Figure 1(c)), incorrectly suggesting that a larger value of $K$ always leads
to a better fit. To illustrate Example 2, we fit a negative binomial
generalized linear model (GLM) to test whether the expected expression of the
first gene is associated with the cluster labels when we estimate $K=2$
clusters. When we perform this test on the same data used for clustering, we
obtain p-values that are much smaller than the $\mathrm{Unif}(0,1)$
distribution (Figure 1(d)), and thus do not control the Type 1 error rate
(recall that no true clusters are present, and thus there is no true
association between the first gene and the estimated clusters).

The solution here might seem obvious: to split our 100 cells into a training
set, used to fit (1), and a test set, used for model validation.
Unfortunately, this _sample splitting_ approach does not work. The issue is
that fitting (1) using the cells in the training set yields latent variable
coordinates for cells in the training set only. To use the test set for
validation or inference, we must obtain latent variable coordinates of the
cells in the test set. This step involves using the test set data itself,
which invalidates downstream evaluation or inference. In our toy example in
Figure 1, we apply $k$-means clustering to the cells in the training set, and
then assign cluster labels to the cells in the test set using 3-nearest
neighbor classification. We see in Figure 1(c) that, over 1000 simulated
datasets, the within-cluster MSE computed on the test set (see Appendix A for
details) decreases monotonically with the number of clusters, because we used
the test set both to compute latent variable coordinates for the cells in the
test set and to compute the within-cluster MSE. Similarly, Figure 1(d) shows
that, over 1,000 simulated datasets, the p-values from a negative binomial GLM
that regresses the first gene from the test set onto the test set cluster
assignments do not control the Type 1 error rate. We refer the reader to Owen
and Perry [2009] for more discussion of the inadequacy of sample splitting in
the setting of Example 1, and Gao et al. [2022], Chen and Witten [2023], and
Neufeld et al. [2022] for a related discussion in the setting of Example 2.

In the setting of Example 3, Cao et al. [2020] implement a procedure that they
call “intradataset cross-validation” (see Algorithm 5). Inspired by the
general framework of Abdelaal et al. [2019], their procedure involves
estimating clusters using all of the data and then performing 5-fold cross-
validation to assess the accuracy of a classifier fit to predict these
clusters. Low cross-validation error is treated as evidence of cluster
reproducibility, because it means that a given cell’s cluster assignment can
be reproduced by a classifier, even when that cell itself was not used to
train the classifier. Figure 1(e) shows a confusion matrix comparing the cell
types estimated via clustering (with $K=5$) to those predicted using cross-
validation (with five folds and a support vector machine classifier) for the
cells in the toy dataset from Figure 1(a). Despite the fact that all cells are
homogenous in this dataset (and thus the estimated clusters are driven by
random noise), 95% of the cells fall on the diagonal of the confusion matrix,
falsely suggesting reproducibility of the clusters. The issue is that, since
all of the data from all of the cells was used for clustering, any downstream
model evaluation is compromised, even if the downstream task makes use of
cross-validation.

Figure 1:  _(a)_ Data $X\in\mathbb{Z}_{\geq 0}^{100\times 2}$, where each
entry $X_{ij}$ is drawn independently from a negative binomial distribution
with mean $5$ and variance $10$. _(b)_ The same data $X$, colored by the
clusters estimated when $k$-means with $K=2$ is applied to $\log(X+1)$. _(c)_
The within-cluster MSE computed after fitting k-means with $K=1,2,\ldots,10$,
averaged over 1000 datasets. We consider three approaches: we cluster using
all of the data and then compute the MSE using all of the data (purple), we
cluster using 50 observations and compute the MSE on the other 50 (gray), and
we apply our proposed negative binomial count splitting method (blue). Details
are given in Section 4.2. _(d)_ Uniform QQ-plot of negative binomial GLM
p-values for testing for differential expression of gene 1 across the
estimated clusters for $1,000$ realizations of $X$. We consider three
approaches: we cluster and fit the GLM using all of the data (purple), we
cluster using 50 observations and fit the GLM on the other 50 (gray), and we
apply our proposed negative binomial count splitting method (blue). Details
are given in Section 4.3. _(e)_ Normalized confusion matrix resulting from the
intradataset cross-validation procedure of Cao et al. [2020] (see Algorithm 5
in Section 5) that uses the same data for both clustering and cross-
validation. _(f)_ Normalized confusion matrix resulting from our modified
version of intradataset cross-validation (see Algorithm 6 in Section 5).

Now, suppose that we were able to sequence the same set of cells twice to
obtain two independent datasets $X^{\mathrm{train}}\in\mathbb{Z}_{\geq
0}^{n\times p}$ and $X^{\mathrm{test}}\in\mathbb{Z}_{\geq 0}^{n\times p}$
generated from $\mathbb{X}$ in (1) (with the same true underlying $L$ and
$\beta$ matrices). We could estimate $L$ and/or $\beta$ using only
$X^{\mathrm{train}}$, and could then validate the results or conduct inference
using $X^{\mathrm{test}}$. Thus, the challenges associated with Examples 1 and
2 displayed in Figure 1(c) and 1(d) would be entirely avoided. Similarly, we
could estimate one set of clusters on $X^{\mathrm{train}}$ and another set of
clusters on $X^{\mathrm{test}}$ and compare the two clusterings using a metric
such as the adjusted Rand Index [Hubert and Arabie, 1985], entirely avoiding
the challenge of Example 3.

In practice, we cannot sequence the same set of cells twice. Instead, we
propose using our single dataset $X$ to reverse engineer two datasets
$X^{\mathrm{train}}$ and $X^{\mathrm{test}}$ that function like independent
sequencing experiments performed on the same sets of cells. Our proposal is an
extension of the ideas of Batson et al. [2019], Sarkar and Stephens [2021],
and Neufeld et al. [2022], who perform this reverse engineering under the
assumption that
$X_{ij}\overset{\mathrm{ind.}}{\sim}\mathrm{Poisson}(\Lambda_{ij})$. However,
scRNA-seq data are typically overdispersed relative to the Poisson or the
binomial distribution, and so the _Poisson count splitting_ procedure
developed in these earlier papers will fail to produce independent training
and test sets.

In recent work, Neufeld et al. [2023] and Dharamshi et al. [2023] developed
_data thinning_ , a vast generalization of Poisson count splitting that
enables us to split a random variable drawn from a number of well-known
distributional families into two or more independent components. In this
paper, we focus on the special case of data thinning for negative binomial
random variables, which allows us to obtain independent matrices
$X^{\mathrm{train}}$ and $X^{\mathrm{test}}$ under the assumption that the
elements of $X$ are independent draws from negative binomial distributions. We
refer to this procedure as _negative binomial count splitting_. Critically,
$X^{\mathrm{train}}$ and $X^{\mathrm{test}}$ are drawn from the same
distribution as $X$, up to a parameter scaling.

Figure 1(c) and Figure 1(d) show that, in our toy example, negative binomial
count splitting correctly determines that $K^{*}=1$ and controls the Type 1
error rate. Figure 1(f) shows that a modified version of “intradataset cross-
validation” that makes use of negative binomial count splitting yields a
confusion matrix that accurately reflects the absence of signal.

Negative binomial count splitting requires a negative binomial assumption to
ensure independence between the training and test sets. However, once the data
has been split, we are free to use any latent variable estimation method or
inferential technique.

In Section 2, we review the Poisson count splitting procedure of Neufeld et
al. [2022] and its generalization to data thinning [Neufeld et al., 2023,
Dharamshi et al., 2023]. In Section 3, we introduce negative binomial count
splitting and provide some theoretical results. In Section 4, we apply
negative binomial count splitting to Example 1 and Example 2 on simulated
data. In Section 5, we revisit the intradataset cross-validation procedure of
Cao et al. [2020], and assess the reproducibility (Example 3) of cell types
and subtypes from their human cell atlas using negative binomial count
splitting. We close with a brief discussion in Section 6.

##  2 Background

###  2.1 A review of data thinning

Our goal is to decompose an scRNA-seq dataset $X\in\mathbb{Z}_{\geq
0}^{n\times p}$ into independent datasets
$X^{\mathrm{train}}\in\mathbb{Z}_{\geq 0}^{n\times p}$ and
$X^{\mathrm{test}}\in\mathbb{Z}_{\geq 0}^{n\times p}$ drawn from the same
model as $X$, up to a parameter scaling. We first review _Poisson count
splitting_ , which accomplishes this goal if
$\mathbb{X}_{ij}\overset{\mathrm{ind.}}{\sim}\mathrm{Poisson}(\Lambda_{ij})$.

######  Algorithm 1 (Poisson count splitting).

Let $X\in\mathbb{Z}_{\geq 0}^{n\times p}$. For a chosen $M\in\mathbb{Z}^{+}$
and $\epsilon_{1},\ldots,\epsilon_{M}\in(0,1)$ such that
$\sum_{m=1}^{M}\epsilon_{m}=1$, draw
$\left(\mathbb{X}_{ij}^{(1)},\mathbb{X}_{ij}^{(2)},\ldots,\mathbb{X}_{ij}^{(M)}\right)\mid\mathbb{X}_{ij}=X_{ij}\sim\mathrm{Multinomial}\left(X_{ij},\epsilon_{1},\ldots,\epsilon_{M}\right)$.

######  Theorem 1.

Let $X\in\mathbb{Z}_{\geq 0}^{n\times p}$ be a dataset with entries $X_{ij}$
drawn from
$\mathbb{X}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{Poisson}}(\mu_{ij})$.
If we apply Algorithm 1 to each element $X_{ij}$ of this data, then (1)
$\mathbf{X}^{(m)}_{ij}\overset{\mathrm{ind.}}{\sim}\mathrm{Poisson}(\epsilon_{m}\mu_{ij})$
for $i=1,\ldots,n$, $j=1,\ldots,p$, and (2) the folds
$\mathbf{X}^{(1)},\ldots,\mathbf{X}^{(M)}$ are mutually independent.

Theorem 1 follows from the well-known binomial thinning property of the
Poisson distribution (see Durrett 2019, Section 3.7.2). For any fold
$m\in\\{1,\ldots,M\\}$, we define
$\mathbf{X}^{(-m)}=\mathbf{X}-\mathbf{X}^{(m)}$. The mutual independence
between folds given in Theorem 1 ensures that, for all $m\in\\{1,\ldots,M\\}$,
$\mathbf{X}^{(m)}$ is independent of $\mathbf{X}^{(-m)}$. Thus, by treating
$\mathbf{X}^{(-m)}$ as a training set and $\mathbf{X}^{(m)}$ as a test set, we
arrive at an alternative to cross-validation that uses Poisson count
splitting, rather than sample splitting, to create training and test sets.

Sarkar and Stephens [2021] and Neufeld et al. [2022] apply Poisson count
splitting to overcome the challenges arising in Examples 1 and 2 of Section 1,
under the assumption that the scRNA-seq data follows a Poisson distribution.
Unfortunately, the Poisson assumption in Theorem 1 is necessary to achieve
independence between the folds $\mathbf{X}^{(1)},\ldots,\mathbf{X}^{(M)}$. If
the data instead follow a negative binomial distribution, then
$\mathbf{X}^{(1)},\ldots,\mathbf{X}^{(M)}$ are correlated, and Poisson count
splitting will fail to provide a valid approach for model evaluation or
inference.

Neufeld et al. [2023] and Dharamshi et al. [2023] expand upon Poisson count
splitting to describe a general recipe for decomposing a single random
variable $\mathbb{Y}\sim F_{\theta}$, for some distribution $F_{\theta}$
indexed by an unknown parameter $\theta$, into independent pieces
$\mathbb{Y}^{(m)}\overset{\mathrm{ind.}}{\sim}Q^{(m)}_{\theta}$, where
$Q^{(m)}_{\theta}$ is a (possibly different) distribution indexed by the same
parameter $\theta$. They refer to this general framework as _data thinning_.
In particular, the framework developed by Neufeld et al. [2023] enables us to
decompose a negative binomial random variable into $M\geq 2$ independent
negative binomial random variables, without knowledge of the mean parameter.
While the property of the Poisson distribution that allows for Poisson count
splitting is well-known, its negative binomial counterpart is less well-known.
In fact, prior to Neufeld et al. [2023], thinning the negative binomial
distribution appears to be unexplored outside of the time series literature of
the 1990s [Joe, 1996].

###  2.2 Negative binomial models for scRNA-seq data

Throughout this paper, we let $\operatorname{\mathrm{NB}}(\mu,b)$ denote the
negative binomial distribution with mean $\mu$ and variance
$\mu+\frac{\mu^{2}}{b}$, for $\mu>0$ and $b>0$. This parameterization is
commonly used when modeling scRNA-seq data. Note that if
$\mathbf{y}\mid\tau\sim\mathrm{Poisson}(\mu\tau)$ and
$\tau\sim\mathrm{Gamma}(b,b)$, then
$\mathbf{y}\sim\operatorname{\mathrm{NB}}(\mu,b)$. Because the variance
$\mu+\frac{\mu^{2}}{b}$ is always strictly larger than the mean $\mu$, the
negative binomial model is overdispersed relative to the Poisson model. This
motivates its use in the analysis of RNA sequencing data, where the data are
non-negative integers with excess variance relative to the Poisson
distribution [Choudhary and Satija, 2022, Hafemeister and Satija, 2019, Sarkar
and Stephens, 2021, Lopez et al., 2018]. We refer to the parameter $b$ as the
overdispersion parameter. As $b\rightarrow\infty$, the negative binomial
distribution approaches the Poisson distribution. In the scRNA-seq literature,
it is common to assume that each gene, but not each cell, has its own
overdispersion parameter [Hafemeister and Satija, 2019, Love et al., 2014,
Lopez et al., 2018]. Thus, in what follows, we will assume that
$\mathbb{X}_{ij}\overset{\mathrm{ind.}}{\sim}\mathrm{NB}\left(\mu_{ij},b_{j}\right)$.

##  3 Negative binomial count splitting

In this section, we consolidate some theoretical results from Neufeld et al.
[2022] and Neufeld et al. [2023] to facilitate their immediate application to
scRNA-seq data. Results in this section are proven in Appendix B.

###  3.1 Algorithm and main result

We now introduce negative binomial count splitting and state the key result.

######  Algorithm 2 (Negative binomial count splitting).

Let $X\in\mathbb{Z}_{\geq 0}^{n\times p}$. For a chosen $M\in\mathbb{Z}^{+}$,
$b_{j}^{\prime}\geq 0$ for $j=1,\ldots,p$, and
$\epsilon_{1},\ldots,\epsilon_{M}\in(0,1)$ such that
$\sum_{m=1}^{M}\epsilon_{m}=1$, draw
$\left(\mathbf{X}_{ij}^{(1)},\ldots,\mathbf{X}_{ij}^{(M)}\right)\mid\mathbf{X}_{ij}=X_{ij}\sim\mathrm{DirichletMultinomial}\left(X_{ij},\epsilon_{1}b_{j}^{\prime},\ldots,\epsilon_{M}b_{j}^{\prime}\right)$.

The marginals of a Dirichlet-multinomial distribution are beta-binomial, i.e.
$\mathbf{X}^{(m)}_{ij}\mid\mathbf{X}_{ij}=X_{ij}\sim\text{BetaBinomial}(X_{ij},\epsilon_{m}b_{j}^{\prime},(1-\epsilon_{m})b_{j}^{\prime})$.
While binomial thinning has appeared in numerous papers as a way to construct
training and test sets from count-valued data [Neufeld et al., 2022, Sarkar
and Stephens, 2021, Leiner et al., 2022], to our knowledge beta-binomial
thinning has only been used for this purpose by Neufeld et al. [2023]. The
beta-binomial thinning operator has appeared in the time series literature as
far back as McKenzie [1986] for constructing autoregressive processes with
negative binomial marginals. The following result, which appeared for $M=2$ in
the context of autoregressive processes in Joe [1996], tells us what happens
when $\mathbf{X}_{ij}\sim\mathrm{NB}\left(\mu_{ij},b_{j}\right)$ and we apply
Algorithm 2 with $b_{j}^{\prime}=b_{j}$.

######  Theorem 2.

Let $X\in\mathbb{Z}_{\geq 0}^{n\times p}$ be a dataset such that the entries
$X_{ij}$ are realizations of
$\mathbf{X}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{NB}}\left(\mu_{ij},b_{j}\right)$.
If we apply Algorithm 2 to each element $X_{ij}$ of this data using
$b_{j}^{\prime}=b_{j}$, then (1)
$\mathbf{X}^{(m)}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{NB}}(\epsilon_{m}\mu_{ij},\epsilon_{m}b_{j})$
for $i=1,\ldots,n$, $j=1,\ldots,p$, and (2) the folds $X^{(1)},\ldots,X^{(M)}$
are mutually independent.

Drawing
$(\mathbb{X}_{ij}^{(1)},\ldots,\mathbb{X}_{ij}^{(M)})\mid\mathbb{X}_{ij}=X_{ij}$
from a Dirichlet-multinomial distribution, as in Algorithm 2, is the same as
first drawing
$\left(\boldsymbol{\epsilon}_{1},\ldots,\boldsymbol{\epsilon}_{M}\right)$ from
a
$\mathrm{Dirichlet}(\epsilon_{1}b_{j}^{\prime},\ldots,\epsilon_{m}b_{j}^{\prime})$
distribution and then letting
$(\mathbb{X}_{ij}^{(1)},\ldots,\mathbb{X}_{ij}^{(M)})\mid\mathbb{X}_{ij}=X_{ij}\sim\mathrm{Multinomial}(X_{ij},\boldsymbol{\epsilon}_{1},\ldots,\boldsymbol{\epsilon}_{M})$.
Thus, to accommodate the overdispersion in $\mathbb{X}_{ij}$ relative to the
Poisson distribution, we add additional randomness to the sampling process by
making the parameters $(\epsilon_{1},\ldots,\epsilon_{M})$ from Algorithm 1
random variables.

Theorem 2 implies that $\mathbf{X}^{(m)}$ is independent of
$\mathbf{X}^{(-m)}:=\mathbf{X}-\mathbf{X}^{(m)}$ for $m=1,\ldots,M$. To
address Examples 1, 2, and 3 from Section 1, we will use $\mathbf{X}^{(-m)}$
as a training set and $\mathbf{X}^{(m)}$ as a test set. The bottom line is
that if we believe that
$\mathbb{X}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{NB}}\left(\mu_{ij},b_{j}\right)$
and we know the true values $b_{j}$, then a direct extension of Poisson count
splitting is available.

###  3.2 The role of the parameter $b_{j}^{\prime}$

Theorem 2 requires that we apply Algorithm 2 with the correct value of the
overdispersion parameter; i.e. that
$\mathbf{X}_{ij}\sim\mathrm{NB}(\mu_{ij},b_{j})$ and we choose
$b_{j}^{\prime}=b_{j}$. In this section, we consider what happens when
$b_{j}^{\prime}\neq b_{j}$.

When $b^{\prime}_{j}=\infty$, drawing
$\left(\mathbf{X}_{ij}^{(1)},\ldots,\mathbf{X}_{ij}^{(M)}\right)\mid\mathbf{X}_{ij}=X_{ij}\sim\mathrm{DirichletMultinomial}\left(X_{ij},\epsilon_{1}b_{j}^{\prime},\ldots,\epsilon_{M}b_{j}^{\prime}\right)$
is equivalent to drawing
$\left(\mathbf{X}_{ij}^{(1)},\ldots,\mathbf{X}_{ij}^{(M)}\right)\mid\mathbf{X}_{ij}=X_{ij}\sim\mathrm{Multinomial}\left(X_{ij},{\epsilon}_{1},\ldots,{\epsilon}_{M}\right)$,
and so Algorithm 2 reduces to Algorithm 1.

######  Theorem 3 (Neufeld et al. [2022]).

If
$\mathbf{X}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{NB}}\left(\mu_{ij},b_{j}\right)$
and we apply Algorithm 2 with $b_{j}^{\prime}=\infty$ to $\mathbf{X}_{ij}$,
then (1)
$\mathbf{X}^{(m)}_{ij}\overset{\mathrm{ind.}}{\sim}\operatorname{\mathrm{NB}}(\epsilon_{m}\mu_{ij},b_{j})$
and (2)
$\mathrm{Cor}(\mathbf{X}^{(m)}_{ij},\mathbf{X}^{(-m)}_{ij})=\frac{\sqrt{\epsilon_{m}(1-\epsilon_{m})}}{\sqrt{\frac{b_{j}^{2}}{\mu_{ij}^{2}}+\frac{b_{j}}{\mu_{ij}}+\epsilon_{m}(1-\epsilon_{m})}}$,
where $\mathbf{X}^{(-m)}:=\mathbf{X}-\mathbf{X}^{(m)}$.

Theorem 3 says that while applying Poisson count splitting (or negative
binomial count splitting with $b_{j}^{\prime}=\infty$) on data from a negative
binomial distribution yields training and test sets that follow the same model
as the full data up to a parameter scaling, these datasets are positively
correlated. The positive correlation increases as the true value of $b_{j}$
decreases, and decreases to $0$ as $b_{j}\rightarrow\infty$. Moreover, we see
from Theorem 3 that the overdispersion parameters (and thus the variances) of
$\mathbf{X}^{(m)}_{ij}$ are too small relative to Theorem 2. Thus, by failing
to put enough noise into our sampling process, applying Poisson count
splitting to negative binomial data results in training and test sets that are
not as noisy as they should be, leading to positive correlation between them.

We now consider the more general case of finite $b_{j}^{\prime}$. The
following result is included (under a different parameterization) in Neufeld
et al. [2023].

######  Theorem 4.

If $\mathbf{X}_{ij}\sim\operatorname{\mathrm{NB}}(\mu_{ij},b_{j})$ and we
apply Algorithm 2 with parameter $b_{j}^{\prime}$:

  1. 1.

$\operatorname{\text{E}}[\mathbf{X}^{(m)}_{ij}]=\epsilon_{m}\mu_{ij}$,

  2. 2.

$\operatorname{\mathrm{Var}}(\mathbf{X}^{(m)}_{ij})=\epsilon_{m}\operatorname{\mathrm{Var}}(\mathbf{X}_{ij})+\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(\frac{b_{j}+1}{b_{j}^{\prime}+1}-1\right)$,

  3. 3.

$\mathrm{Cov}(\mathbf{X}^{(m)}_{ij},\mathbf{X}^{(-m)}_{ij})=\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(1-\frac{b_{j}+1}{b_{j}^{\prime}+1}\right)$,
where $\mathbf{X}^{(-m)}:=\mathbf{X}-\mathbf{X}^{(m)}$.

Unlike in Theorem 3, the folds of data that result from applying Algorithm 2
with arbitrary values for $b_{j}^{\prime}$ do not necessarily follow negative
binomial distributions. The first statement of Theorem 4 says that, regardless
of the value of $b^{\prime}_{j}$ used, the expected value of
$\mathbf{X}^{(m)}$ and $\mathbf{X}^{(-m)}$ are scaled by
$\operatorname{\epsilon}_{m}$ and $(1-\operatorname{\epsilon}_{m})$ compared
to $\operatorname{\text{E}}[\mathbb{X}]$ in (1). Thus, we can estimate the
latent space using the training set $\mathbf{X}^{(-m)}$, and validate these
estimates using the held out set $X^{(m)}$. The second statement of Theorem 4
tells us that using the wrong value for $b_{j}^{\prime}$ affects the variance
of $\mathbf{X}^{(m)}$. The third statement says that if $b^{\prime}_{j}\neq
b_{j}$, then the training and test sets are correlated, and the magnitude of
the correlation grows with the magnitude of the discrepancy between $b_{j}$
and $b_{j}^{\prime}$. The result is displayed and empirically confirmed in
Figure 2. In order to use $\mathbf{X}^{(m)}$ to validate a model fit to
$\mathbf{X}^{(-m)}$ or to do valid inference on latent variables fit to
$\mathbf{X}^{(-m)}$, we need independence between $\mathbf{X}^{(m)}$ and
$\mathbf{X}^{(-m)}$. Thus, the major takeaway from Theorem 4 is that, when the
true $b_{j}$ are unknown, it is important to estimate them well. In Section 4,
we use the well-known R package sctransform [Hafemeister and Satija, 2019] to
estimate each $b_{j}$.

Figure 2: We generate $100,000$ independent realizations of
$\mathbf{X}_{11}\sim\mathrm{NB}(25,8)$. Then, for 50 values of
$b_{1}^{\prime}$ ranging from $10^{-6}$ to $10^{6}$, we split each of these
realizations into $\mathbf{X}^{(1)}_{11}$ and $\mathbf{X}^{(2)}_{11}$ by
applying Algorithm 2 with
$b_{1}^{\prime},M=2,\epsilon_{1}=0.3,\epsilon_{2}=0.7$. We display the sample
correlation between the $100,000$ realizations of $\mathbf{X}_{11}^{(1)}$ and
$\mathbf{X}_{11}^{(2)}$, as a function of $b^{\prime}_{1}$. The pink line
shows the theoretical values computed using Theorem 4. The horizontal green
line shows the asymptote at $b^{\prime}_{1}=\infty$ given by Theorem 2. The
vertical red line displays where $b^{\prime}_{1}=b_{1}$. As expected, both the
empirical and theoretical correlations are $0$ at this point.

###  3.3 The role of the parameters $\epsilon_{1},\ldots,\epsilon_{M}$

In this section, we consider the case where we used the “correct” value of
$b_{j}^{\prime}$ and thus the results of Theorem 2 hold. In this setting, it
is simple to show that, for a given fold $m$, the parameter $\epsilon_{m}$
governs a tradeoff between the amount of information in the training set
$\mathbf{X}^{(-m)}$ and in the test set $\mathbf{X}^{(m)}$. This result is
summarized in the following theorem.

######  Theorem 5 (Information tradeoff as we vary $\operatorname{\epsilon}$).

If $\mathbb{X}_{ij}\sim\operatorname{\mathrm{NB}}(\mu_{ij},b_{j})$, then the
Fisher information contained in a single datapoint $X_{ij}$ for the parameter
$\mu_{ij}$ is
$I_{\mu_{ij}}(\mathbf{X}_{ij})=\frac{b_{j}}{(b_{j}+\mu_{ij})\mu_{ij}}$. If we
apply Algorithm 2 with $b^{\prime}_{j}=b_{j}$, then for $m=1,\ldots,M$, the
Fisher information contained in $\mathbb{X}^{(m)}_{ij}$ for the parameter
$\mu_{ij}$ is $\epsilon_{m}I_{\mu_{ij}}(\mathbf{X}_{ij})$, and the
corresponding Fisher information contained in $\mathbb{X}^{(-m)}_{ij}$ is
$(1-\epsilon_{m})I_{\mu_{ij}}(\mathbf{X}_{ij})$.

We will see in Section 4 that the ideal choice of
$\epsilon_{1},\ldots,\epsilon_{M}$ depends on the application.

##  4 Simulation Study

In this section, we apply negative binomial count splitting to Example 1 and
Example 2 from Section 1, instantiated to the setting where there are $K^{*}$
true discrete latent variables which represent cell types that we wish to
estimate using $k$-means clustering. More specifically, we consider the
following settings:

  * Example 1, instantiated to clustering: We fit models with $K$ clusters for $K=1,\ldots,10$, evaluate each model, then let the best value of $K$ be our estimate of $K^{*}$.

  * Example 2, instantiated to clustering: We cluster the data into $K=2$ estimated cell types and test each gene for differential expression across the two estimated clusters.

After introducing our simulation setup in Section 4.1, we show that negative
binomial count splitting can be easily applied to Example 1 (Section 4.2) and
Example 2 (Section 4.3).

###  4.1 Data generating mechanism

We generate datasets with $n$ cells, $p$ genes, and $K^{*}$ true cell types.
Each gene has a baseline expression level $\exp(\beta_{j0})$ where
$\beta_{j0}\overset{\mathrm{i.i.d.}}{\sim}N(0,1)$ for $j=1,\ldots,p$.

For each dataset, we assign each cell to one of the $K^{*}$ clusters with
equal probability. The first column of the latent variable matrix
$L\in\mathbb{R}^{n\times K^{*}}$ contains ones, and the remaining columns are
indicators for clusters $k=2,\ldots,K^{*}$. The matrix
$\beta\in\mathbb{R}^{p\times K^{*}}$ stores $\beta_{10},\ldots,\beta_{p0}$ in
the first column. When $K^{*}>1$, $\beta_{j2}=\beta^{*}$ for the first $5\%$
of the genes, and the rest of the entries in the second column of $\beta$ are
$0$. If $K^{*}>2$, then $\beta_{j3}=\beta^{*}$ for the next $5\%$ of the
genes, and all other entries are $0$. We continue filling in the $\beta$
matrix in this manner until all $K^{*}$ columns have been filled. We consider
different values of $\beta^{*}$ for different datasets, but within a dataset
we always use the same value of $\beta^{*}$ such that all $K^{*}$ clusters are
equally easy to detect. Finally, we let
$\log\left(\Lambda\right)=L\beta^{\top}$.

We let the overdispersion parameter $b_{j}$ for the $j$th gene be a function
of the average expression
$\bar{\Lambda}_{j}=\frac{1}{n}\sum_{i=1}^{n}\Lambda_{ij}$ for that gene
[Choudhary and Satija, 2022, Hafemeister and Satija, 2019, Love et al., 2014].
More specifically, we set $b_{j}=\frac{\bar{\Lambda}_{j}}{\tau}$ for either
$\tau=1$ (“mild overdispersion”) or $\tau=5$ (“severe overdispersion”). We
then let $\mathbf{X}_{ij}\sim\mathrm{NB}(\Lambda_{ij},b_{j})$, such that
$\operatorname{\mathrm{Var}}(\mathbb{X}_{ij})=\Lambda_{ij}\left(1+\frac{\Lambda_{ij}}{b_{j}}\right)\approx\Lambda_{ij}\left(1+\tau\right)$.
We omit size factors from this simulation study, as they are not the focus of
this paper.

###  4.2 Selecting the number of clusters

####  4.2.1 Methods

We now introduce the general algorithm used in this section.

######  Algorithm 3 (Estimating the number of clusters).

Start with datasets $X^{\mathrm{train}}\in\mathbb{Z}_{\geq 0}^{n\times p}$ and
$X^{\mathrm{test}}\in\mathbb{Z}_{\geq 0}^{n\times p}$ and parameter
$\operatorname{\epsilon}\in(0,1)$, where
$\operatorname{\text{E}}\left[X^{\mathrm{test}}\right]=\frac{1-\operatorname{\epsilon}}{\operatorname{\epsilon}}\operatorname{\text{E}}[X^{\mathrm{train}}]$.
For $K=1,\ldots,10$:

  1. 1.

Run k-means clustering to estimate $K$ clusters on
$\log(X^{\mathrm{train}}+1)$. This yields a cluster assignment
$\hat{c}_{i}\in\\{1,\ldots,K\\}$ for $i=1,\ldots,n$.

  2. 2.

For $i=1,\ldots,n$ and $j=1,\ldots,p$, estimate
$\operatorname{\text{E}}[\mathbb{X}_{ij}^{\mathrm{train}}]$ with the sample
mean of the training set points assigned to the same cluster:
$\hat{\mu}^{\mathrm{train}}_{ij}=\frac{1}{\sum_{i^{\prime}=1}^{n}\mathbb{1}\\{\hat{c}_{i^{\prime}}=\hat{c}_{i}\\}}\sum_{i^{\prime}=1}^{n}X^{\mathrm{train}}_{i^{\prime}j}\mathbb{1}\\{\hat{c}_{i^{\prime}}=\hat{c}_{i}\\}$.

  3. 3.

Estimate $\operatorname{\text{E}}[\mathbb{X}_{ij}^{\mathrm{test}}]$ as
follows:
$\hat{\mu}^{\mathrm{test}}_{ij}=\frac{1-\operatorname{\epsilon}}{\operatorname{\epsilon}}\hat{\mu}^{\mathrm{train}}_{ij}$.

  4. 4.

Compute the within-cluster mean squared error on the test set as follows:

$MSE(K)=\frac{1}{n\times
p}\sum_{i=1}^{n}\sum_{j=1}^{p}\left(\log\left(X^{\mathrm{test}}_{ij}+1\right)-\log\left(\hat{\mu}^{\mathrm{test}}_{ij}+1\right)\right)^{2}.$
(2)

We apply Algorithm 3 as follows.

  * Naive method: Run Algorithm 3 with $X^{\mathrm{train}}=X^{\mathrm{test}}=X$ and $\operatorname{\epsilon}=0.5$.

  * Poisson count splitting (PCS): Apply Algorithm 1 to $X$ with $M=2$ and some $(\epsilon_{1},\epsilon_{2})$ to obtain $X^{(1)}$ and $X^{(2)}$. Run Algorithm 3 with $X^{\mathrm{train}}=X^{(1)},X^{\mathrm{test}}=X^{(2)}$, and $\epsilon=\epsilon_{1}$.

  * Negative binomial count splitting, known $b$ (NBCS-known): Apply Algorithm 2 to $X$ with $M=2$, some choice of $(\epsilon_{1},\epsilon_{2})$, and $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(b_{1},\ldots,b_{p})$ to obtain $X^{(1)}$ and $X^{(2)}$. Run Algorithm 3 using $X^{\mathrm{train}}=X^{(1)},X^{\mathrm{test}}=X^{(2)}$, and $\epsilon=\epsilon_{1}$.

  * Negative binomial count splitting, estimated $b$ (NBCS-estimated): First, use the R package sctransform [Hafemeister and Satija, 2019] to obtain estimates $\hat{b}_{1},\ldots,\hat{b}_{p}$ of $b_{1},\ldots,b_{p}$. Details are given in Appendix C. Then run Algorithm 2 on matrix $X$ with $M=2$, some choice of $(\epsilon_{1},\epsilon_{2})$, and $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(\hat{b}_{1},\ldots,\hat{b}_{p})$ to obtain $X^{(1)}$ and $X^{(2)}$. Then run Algorithm 3 using $X^{\mathrm{train}}=X^{(1)},X^{\mathrm{test}}=X^{(2)}$, and $\epsilon=\epsilon_{1}$.

We now extend both versions of NBCS to perform cross-validation.

  * Negative binomial cross-validation, known $b$ (NBCV-known): Obtain $(X^{(1)},\ldots,X^{(M)})$ by running Algorithm 2 on $X$ with $M=10$, $\epsilon_{m}=\frac{1}{M}$ for $m=1,\ldots,M$, and $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(b_{1},\ldots,b_{p})$. For $m=1,\ldots,M$, apply Algorithm 3 with $X^{\mathrm{train}}=X^{(-m)}$, $X^{\mathrm{test}}=X^{(m)}$, and $\operatorname{\epsilon}=\frac{M-1}{M}$. For each value of $K$, record the total MSE summed across the $M$ folds.

  * Negative binomial cross-validation, estimated $b$ (NBCV-estimated):  First obtain estimates $\hat{b}_{1},\ldots,\hat{b}_{p}$ of $b_{1},\ldots,b_{p}$ using sctransform on the full data $X$. Then proceed as in NBCV-known, but apply Algorithm 2 with $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(\hat{b}_{1},\ldots,\hat{b}_{p})$.

We already showed in Section 1 that the naive method fails to provide a viable
solution to Examples 1 and 2; nonetheless we include it here for the sake of
comparison. We do not consider sample splitting in this section, as we already
saw in Section 1 that it fails to provide a viable solution to our problems of
interest.

####  4.2.2 Results

We first generate 1,000 datasets with $n=1000$ and $p=1000$ and
$\beta^{*}=1.5$ for $K^{*}\in\\{1,3,5\\}$ and the two overdispersion settings
described in Section 4.1. For each dataset, we consider the naive method,
Poisson count splitting, NBCS-known, and NBCS-estimated (all with
$\operatorname{\epsilon}=0.5$). We plot the average MSE over the 1000
datasets, defined in (2), as a function of $K$. To facilitate comparisons
between methods, the y-axis of Figure 3 has been scaled such that the MSE for
each method ranges from $0$ to $1$.

Figure 3 shows that, regardless of the true value of $K^{*}$ or the amount of
overdispersion, the MSE for the naive method decreases monotonically with $K$.
Thus, we cannot simply select the value of $K$ that minimizes the loss
function. We must instead search for a bend or an “elbow” in the MSE plot,
which fails to provide a clear answer for the number of clusters we should
select. We see that PCS performs well (the loss function is minimized when
$K^{*}=K$) under mild overdispersion, but under severe overdispersion its
performance approaches that of the naive method. This is as expected from
Theorem 3, since the correlation between $X^{\mathrm{train}}=X^{(1)}$ and
$X^{\mathrm{test}}=X^{(-1)}$ under PCS increases as the amount of
overdispersion increases. Both versions of NBCS have loss functions that are
minimized when $K=K^{*}$, regardless of the true value of $K^{*}$ or the
amount of overdispersion in the data. Thus, we can select the number of
clusters by selecting the value of $K$ that minimizes the loss function. As
the naive method and PCS do not lead to loss functions that are minimized at
the true number of clusters, we do not consider them for the remainder of this
section.

Figure 3: The average within-cluster MSE (see (2) in Algorithm 3) over 1000
datasets for four of the methods described in Section 4.2.1, scaled to have
range between 0 and 1.

We next explore the role of the parameter $\epsilon$ in NBCS. We generate
2,000 datasets with mild overdispersion with $n=500$ and $p=40$ for $K=5$ and
for $\beta^{*}$ values ranging from $2$ to $6$. For values of $\epsilon$
ranging from $0$ to $1$, we perform NBCS-known. For each dataset and each
value of $\epsilon$, we consider three metrics. The left panel of Figure 4
displays the adjusted Rand index (ARI) between the true clusters and those
estimated using $X^{\mathrm{train}}$ when $K=K^{*}$, as a function of
$\epsilon$. For a given signal strength, the average ARI increases with
$\epsilon$ because a large value of $\operatorname{\epsilon}$ means that we
use more of the information in our data in the cluster estimation phase of
Algorithm 3 (see Theorem 5). The center panel of Figure 4 shows the proportion
of times that the MSE is minimized at $K=K^{*}$ (i.e. that we select the
correct value of $K$), given that the ARI between the true clusters and the
estimated clusters when $K=K^{*}$ exceeds $0.8$. This metric decreases with
$\epsilon$, as large values of $\operatorname{\epsilon}$ leave less
information in the test set for us to validate the estimated clusters. The
right panel of Figure 4 shows the overall proportion of datasets for which the
loss function is minimized at $K=K^{*}$. The optimal value of
$\operatorname{\epsilon}$ depends on the true signal strength, but it always
involves a tradeoff between choosing $\operatorname{\epsilon}$ large enough to
estimate good clusters on $X^{\mathrm{train}}$, but not so large that we
cannot accurately validate the clusters.

Figure 4:  We generate 2,000 datasets with mild overdispersion, $n=500$,
$p=40$, $K^{*}=5$ for each of five values of $\beta^{*}$ ranging from $2$ to
$6$. For values of $\epsilon$ ranging from $0$ to $1$, we perform NBCS-known.
_Left:_ The average adjusted Rand index between the true clusters and those
estimated with on the training set when $K=K^{*}$, as a function of
$\epsilon$. _Center:_ The proportion of datasets for which the test set MSE is
minimized at $K=K^{*}$, only considering datasets for which the adjusted Rand
index between the true clusters and the clusters estimated on the training set
when $K=K^{*}$ exceeds 0.8. _Right:_ The overall proportion of datasets for
which the loss function is minimized at $K=K^{*}$.

Finally, we compare NBCS-known with $\operatorname{\epsilon}=0.9$ to NBCV-
known with 10 folds. We generate 2,000 datasets where $n=500$ and $p=40$ for
values of $\beta^{*}$ ranging from $1$ to $6$ and $K^{*}=1,3,5$. As both
methods use 90% of the information in the data for training and 10% for
testing, these methods have the same average MSE curves over many datasets.
However, for a given dataset, NBCV-known selects the correct value of $K$ more
often than NBCS-known, because averaging the MSE over 10 folds reduces the
variance in our validation step (Figure 5).

Figure 5: We fix $n=500$ and $p=40$ and we generate 200 datasets for each
combination of $K\in\\{1,3,5\\}$, both overdispersion settings, and
$\beta^{*}$ values ranging from $0.1$ to $3$. For each value of $\beta^{*}$,
we show the proportion of datasets for which the loss function was minimized
at $K=K^{*}$.

###  4.3 Testing for differential expression

####  4.3.1 Methods

In this section we let $K^{*}=2$ and we always estimate two clusters on the
data. Our focus is no longer on estimating the number of clusters, but rather
on studying differential expression across a given set of estimated clusters.
We use the following algorithm.

######  Algorithm 4 (Testing for differential expression.).

Start with datasets $X^{\mathrm{train}}\in\mathbb{Z}_{\geq 0}^{n\times p}$ and
$X^{\mathrm{test}}\in\mathbb{Z}_{\geq 0}^{n\times p}$.

  1. 1.

Apply $k$-means clustering with $K=2$ to estimate clusters on
$\log(X^{\mathrm{train}}+1)$. This yields a cluster assignment
$\hat{c}_{i}\in\\{0,1\\}$ for $i=1,\ldots,n$.

  2. 2.

For $j=1,\ldots,p$, fit a negative binomial GLM of $X^{\mathrm{test}}_{j}$ on
$\hat{c}$. Report the Wald p-value for the slope coefficient.

We consider the following ways to obtain $X^{\mathrm{train}}$ and
$X^{\mathrm{test}}$ in Algorithm 4.

  * Naive method: Let $X^{\mathrm{train}}=X^{\mathrm{test}}=X$.

  * Poisson count splitting (PCS): Obtain $X^{\mathrm{train}}=X^{(1)}$ and $X^{\mathrm{test}}=X^{(2)}$ by running Algorithm 1 on the data $X$ with $M=2$ and $(\epsilon_{1},\epsilon_{2})=(\epsilon,1-\epsilon)$.

  * Negative binomial count splitting, known $b$ (NBCS-known): Obtain $X^{\mathrm{train}}=X^{(1)}$ and $X^{\mathrm{test}}=X^{(2)}$ by running Algorithm 2 on $X$ with $M=2$, $(\epsilon_{1},\epsilon_{2})=(\epsilon,1-\epsilon)$, and $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(b_{1},\ldots,b_{p})$.

  * Negative binomial count splitting, estimated $b$ (NBCS-estimated:) Use the R package sctransform [Hafemeister and Satija, 2019] to obtain estimates $\hat{b}_{1},\ldots,\hat{b}_{p}$ of $b_{1},\ldots,b_{p}$ using the full dataset $X$. Then apply Algorithm 2 on matrix $X$ with $M=2$, $(\epsilon_{1},\epsilon_{2})=(\epsilon,1-\epsilon)$, and $(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(\hat{b}_{1},\ldots,\hat{b}_{p})$.

Unlike in Section 4.2, we do not consider splitting with $M>2$ folds and we do
not aggregate results across folds by interchanging the roles of the train and
test sets. We leave the possibility of aggregating differential expression
test statistics across multiple folds to future work. Once again, we do not
consider sample splitting because we already saw in Figure 1 that it fails to
provide a viable solution to in Example 2.

####  4.3.2 Results

We generate datasets using the mechanism described in Section 4.1 with
$K^{*}=2$, $n=500$, and $p=40$. Under this mechanism, the first two genes are
differentially expressed across the two true clusters. We refer to the
remaining $38$ genes as null genes because they have the same expected
expression across all cells. Figure 6 shows uniform QQ plots of the p-values
obtained from the four variations of Algorithm 4 for the null genes,
aggregated across 1000 datasets for each of 16 $\beta^{*}$ values. We see from
Figure 6 that both the naive method and Poisson count splitting fail to
control the Type 1 error rate. Poisson count splitting performs worse when
overdispersion is severe. On the other hand, both versions of NBCS control the
Type 1 error rate.

Figure 6: We generate $1,000$ datasets with $K^{*}=2,n=500,$ and $p=40$ for
each overdispersion setting and for each of 16 values of $\beta^{*}$ ranging
from $0$ to $3$. For each dataset, we carry out the four variations of
Algorithm 4 given in Section 4.3. The QQ-plots indicate that only NBCS
controls the Type 1 error rate across both overdispersion settings.

Finally, we explore the role of $\operatorname{\epsilon}$ in this setting. We
generate $1000$ datasets where $n=1000$ and $p=1000$ for each of 16 values of
$\beta^{*}$ ranging from $0$ to $3$. For each dataset, we consider the
adjusted Rand index between the true clusters and those estimated on the
training set. As in Section 4.2, we expect that larger values of $\epsilon$
will lead to higher adjusted Rand indices, on average, for a given signal
strength $\beta^{*}$. This is confirmed in the left panel of Figure 7. On the
other hand, given the clusters that we estimated on the training set, smaller
values of $\epsilon$ leave us with more power to detect differential
expression on the test set. We define $\hat{\beta}_{j}^{*}$ to be the
estimated GLM coefficient for a gene $X_{j}$ if we regress its mean vector
$\Lambda_{j}$ onto the clusters estimated on the training set; note that
$\hat{\beta}_{j}^{*}=\beta^{*}$ only if the estimated clusters are exactly
equal to the true clusters. The right panel of Figure 7 plots the proportion
of times that the differential expression p-value for a non-null gene was less
than $0.05$, as a function of $\hat{\beta}_{j}^{*}$. We see that, for a given
value of $\hat{\beta}_{j}^{*}$, the proportion of null hypotheses rejected is
highest when $\epsilon$ is small, because smaller values of $\epsilon$ leave
more information in the test set.

Figure 7: _Left:_ The average adjusted Rand index between the true clusters
and those estimated on the training set, plotted as a function of $\beta^{*}$
for both overdispersion settings and for each value of $\epsilon$. _Right:_
The proportion of times that the differential expression p-value for a non-
null gene was less than $0.05$, as a function of the association between the
gene’s expected expression and the estimated clusters.

##  5 Application to fetal cell atlas data

In this section, we apply negative binomial count splitting to solve the
problems associated with Example 3 from Section 1. Cao et al. [2020] sequenced
more than 4 million cells from 121 human fetal samples to create a fetal cell
atlas, with a goal of organizing the cells from each of 15 different organs
into cell types and cell subtypes. In one analysis, after preprocessing, they
clustered the expression data from the kidney cells into nine main cell types.
They then subsetted the data to include only the cluster thought to correspond
to metanephric kidney cells, and clustered that subset into 10 cell subtypes.
As there is no ground truth available for the main cell types or the cell
subtypes, it is difficult to validate their results. Cao et al. [2020] use a
procedure that they call _intradataset cross-validation_ ; see their Figure 2.
We describe a simplified version of this procedure in Algorithm 5.

######  Algorithm 5 (Intradataset cross-validation, Cao et al. [2020]).

Input: a cell-by-gene expression matrix $X\in\mathbb{Z}_{\geq 0}^{n\times p}$.

  1. 1.

After preprocessing, cluster the dataset $X$ to obtain estimated cell types
$\hat{L}(X)^{\mathrm{cluster}}$.

  2. 2.

Divide the $n$ cells into $5$ folds and, for $m=1,\ldots,5$:

    1. (a)

Let all cells in the $m$th fold be the test set; the remaining cells are the
training set.

    2. (b)

Train a classifier to predict $\hat{L}(X)^{\mathrm{cluster}}_{i}$ using
$X_{i}$, where $i$ indexes the cells in the training set.

    3. (c)

Use this trained classifier to predict
$\hat{L}^{\mathrm{cluster}}_{i^{\prime}}$ using $X_{i^{\prime}}$, where
$i^{\prime}$ indexes the cells in the test set. Let
$\hat{L}_{i^{\prime}}^{\mathrm{classifier}}$ denote the prediction for the
$i^{\prime}$th cell.

  3. 3.

For $i=1,\ldots,n$, the $i$th cell now has a value
$\hat{L}(X)^{\mathrm{cluster}}_{i}$, obtained from clustering, and
$\hat{L}^{\mathrm{classifier}}_{i}$, obtained via prediction when it belonged
to the test set. Make a confusion matrix comparing
$\hat{L}(X)^{\mathrm{cluster}}$ and $\hat{L}^{\mathrm{classifier}}$. A
diagonal confusion matrix is treated as evidence that the clusters are
reproducible, since the cluster assignment for a given cell can be recovered
by a classifier that was trained without knowledge of that cell’s cluster
assignment. The adjusted Rand index (ARI) numerically summarizes the degree of
agreement between $\hat{L}(X)^{\mathrm{cluster}}$ and
$\hat{L}^{\mathrm{classifier}}$.

In Section 1 (Figure 1(e)), we showed using a toy dataset that this procedure
is problematic. Because the entire dataset $X$ is used in Step 1 to estimate
the clusters $\hat{L}(X)^{\mathrm{cluster}}$, the test set in Step 2(c) has
not truly been held out of the training process. Thus, despite the fact that
the clusters estimated on the toy dataset are driven by random noise, the
confusion matrix output by Algorithm 5 is close to diagonal and the ARI is
close to $1$ (Figure 1(e)). Negative binomial count splitting provides a
simple alternative (Figure 1(f)), which we outline below.

######  Algorithm 6 (Intradataset cross-validation via count splitting).

Input: a cell-by-gene expression matrix $X\in\mathbb{Z}_{\geq 0}^{n\times p}$.

  1. 1.

Obtain an estimate $\hat{b}_{j}$ of the gene-specific overdispersion parameter
$b_{j}$ for $j=1,\ldots,p$.

  2. 2.

Apply Algorithm 2 (negative binomial count splitting) with $M=2$,
$\epsilon_{1}=\epsilon_{2}=0.5$, and
$(b_{1}^{\prime},\ldots,b_{p}^{\prime})=(\hat{b}_{1},\ldots,\hat{b}_{p})$ to
create two folds of data, $X^{(1)}$ and $X^{(2)}$.

  3. 3.

After preprocessing, apply a clustering algorithm to $X^{(1)}$ to obtain an
estimated cluster $\hat{L}(X^{(1)})_{i}$ for the $i$th cell for
$i=1,\ldots,n$.

  4. 4.

After preprocessing, apply the same clustering algorithm to $X^{(2)}$ to
obtain an estimated cluster $\hat{L}(X^{(2)})_{i}$ for the $i$th cell for
$i=1,\ldots,n$.

  5. 5.

Make a confusion matrix comparing $\hat{L}(X^{(1)})$ and $\hat{L}(X^{(2)})$. A
diagonal confusion matrix (up to a permutation of the columns) is treated as
evidence of cluster reproducibility, since the cells are reliably assigned to
the same cluster on independent realizations of the data. The ARI numerically
summarizes the degree of agreement between $\hat{L}(X^{(1)})$ and
$\hat{L}(X^{(2)})$.

We now compare Algorithm 5 and Algorithm 6 on data from the human fetal cell
atlas. We consider two versions of Algorithm 6. We apply a version where, in
Step 1, we let each $\hat{b}_{j}=\infty$, which corresponds to assuming that
the data are Poisson. We also apply a version where, in Step 1, we estimate
each $\hat{b}_{j}$ using the full dataset using sctransform; details are given
in Appendix C. In both cases, we use $\epsilon_{1}=\epsilon_{2}=0.5$ such that
$X^{(1)}$ and $X^{(2)}$ are identically distributed. Our goal is not to
discover the optimal clusters in this dataset, but rather to compare Algorithm
5 and Algorithm 6 as strategies for validating clusters. As such, we do not
attempt to reproduce the exact analysis from Cao et al. [2020]; we use a
simplified implementation, which is described in Appendix D.

To start, we apply Algorithm 5 and both versions of Algorithm 6 to all
$178,603$ kidney cells from the human fetal cell atlas. The results are shown
in panels (a), (b), and (c) of Figure 8. Regardless of the algorithm used, we
see diagonal confusion matrices and high ARIs, suggesting reproducibility of
the clusters. For panels (b) and (c), we permute the columns of the matrices
to make them appear as diagonal as possible, but we note that the ARI is
invariant to these permutations.

We next apply Algorithm 5 and both versions of Algorithm 6 to the $90,876$
kidney cells that Cao et al. [2020] annotated as metanephric cells. The
results are shown in panels (d), (e), and (f) of Figure 8. Overall, panels (e)
and (f) show lower ARIs than panels (b) and (c), suggesting that these
supposed cell subtypes are somewhat less reproducible than the main cell
types. We note that the difference between the main cell type analysis and the
cell subtype analysis is least stark for Algorithm 5, where the double use of
data causes the cell subtypes to appear more reproducible than they are in
both analyses. Similarly, the difference between the main cell type analysis
and the cell subtype analysis is less stark for the Poisson version of
Algorithm 6 than the sctransform version, suggesting that the dependence
between the training set and the test set induced by setting
$\hat{b}_{1}=\ldots=\hat{b}_{p}=\infty$ is also enough to make the cell
subtypes appear slightly more reproducible than they are.

We repeat each version of Algorithm 6 ten times, using ten different random
splits of the data. The resulting ARIs are shown in Figures 8(d) and 8(h). The
ARIs from the cell subtype analysis are consistently lower than those from the
main cell type analysis.

Figure 8: The results of applying (a) Algorithm 5, (b) Algorithm 6 assuming
the data is Poisson, and (c) Algorithm 6 with overdispersion parameters
estimated via sctransform, to the full kidney dataset. Panels (d), (e), and
(f) show the results of applying those algorithms to only the metanephric
cells. Panels (d) and (h) show the adjusted Rand indices (ARIs) that result
from 10 replications of Algorithm 6.

In summary, due to the double use of data, Algorithm 5 overestimates the
reproducibility of the metanephric cell subtypes. On the other hand, Algorithm
6 allows us to see that the metanephric cell subtypes are less reproducible
than the main kidney cell types.

##  6 Discussion

In this paper, we introduced an algorithm to split a negative binomial random
variable into two or more independent negative binomial random variables, and
we applied it to the analysis of scRNA-seq data. This algorithm is also
applicable in other settings; data are modeled as negative binomial in a wide
variety of fields. Furthermore, while we were particularly motivated by
unsupervised settings in which sample splitting is not an option, negative
binomial count splitting can also be used in supervised settings such as
inference after variable selection in negative binomial regression. For
further discussion, see Neufeld et al. [2023] and Dharamshi et al. [2023]. An
implementation of the techniques in this paper is available in the R package
countsplit, and tutorials showing how this can be integrated with existing
scRNA-seq software packages, are available at anna-
neufeld.github.io/countsplit.tutorials.

## Acknowledgements

Anna Neufeld and Daniela Witten were supported by the Simons Foundation, the
National Institutes of Health (NIH), and the Keck Foundation. Lucy Gao was
supported by the National Sciences and Engineering Council of Canada. Joshua
Popp was supported by the NIH. Alexis Battle was supported by the NIH and the
Chan Zuckerberg Initiative.

## References

  * Abdelaal et al. [2019] Tamim Abdelaal, Lieke Michielsen, Davy Cats, Dylan Hoogduin, Hailiang Mei, Marcel JT Reinders, and Ahmed Mahfouz.  A comparison of automatic cell identification methods for single-cell RNA sequencing data.  _Genome Biology_ , 20:1–19, 2019.
  * Aizarani et al. [2019] Nadim Aizarani, Antonio Saviano, Laurent Mailly, Sarah Durand, Josip S Herman, Patrick Pessaux, et al.  A human liver cell atlas reveals heterogeneity and epithelial progenitors.  _Nature_ , 572(7768):199–204, 2019.
  * Batson et al. [2019] Joshua Batson, Loïc Royer, and James Webber.  Molecular cross-validation for single-cell RNA-seq.  _BioRxiv_ , page 786269, 2019.
  * Cao et al. [2020] Junyue Cao, Diana R O’Day, Hannah A Pliner, Paul D Kingsley, Mei Deng, Riza M Daza, et al.  A human cell atlas of fetal gene expression.  _Science_ , 370(6518):eaba7721, 2020.
  * Chen and Witten [2023] Yiqun T Chen and Daniela M Witten.  Selective inference for k-means clustering.  _Journal of Machine Learning Research_ , 2023.
  * Choudhary and Satija [2022] Saket Choudhary and Rahul Satija.  Comparison and evaluation of statistical error models for scRNA-seq.  _Genome Biology_ , 23(1):1–20, 2022.
  * Chung and Storey [2015] Neo Christopher Chung and John D Storey.  Statistical significance of variables driving systematic variation in high-dimensional data.  _Bioinformatics_ , 31(4):545–554, 2015.
  * Dharamshi et al. [2023] Ameer Dharamshi, Anna Neufeld, Keshav Motwani, Lucy L Gao, Daniela Witten, and Jacob Bien.  Generalized data thinning using sufficient statistics.  _arXiv preprint arXiv:2303.12931_ , 2023.
  * Durrett [2019] Rick Durrett.  _Probability: Theory and Examples_ , volume 49.  Cambridge University Press, 2019.
  * Eraslan et al. [2019] Gökcen Eraslan, Lukas M Simon, Maria Mircea, Nikola S Mueller, and Fabian J Theis.  Single-cell RNA-seq denoising using a deep count autoencoder.  _Nature Communications_ , 10(1):1–14, 2019.
  * Fu and Perry [2020] Wei Fu and Patrick O Perry.  Estimating the number of clusters using cross-validation.  _Journal of Computational and Graphical Statistics_ , 29(1):162–173, 2020.
  * Gao et al. [2022] Lucy L Gao, Jacob Bien, and Daniela Witten.  Selective inference for hierarchical clustering.  _Journal of the American Statistical Association_ , pages 1–11, 2022\.
  * Grabski et al. [2023] Isabella N Grabski, Kelly Street, and Rafael A Irizarry.  Significance analysis for clustering with single-Cell RNA-sequencing data.  _Nature Methods_ , pages 1–7, 2023.
  * Hafemeister and Satija [2019] Christoph Hafemeister and Rahul Satija.  Normalization and variance stabilization of single-cell RNA-seq data using regularized negative binomial regression.  _Genome Biology_ , 20(1):1–15, 2019.
  * Harremoës et al. [2010] Peter Harremoës, Oliver Johnson, and Ioannis Kontoyiannis.  Thinning, entropy, and the law of thin numbers.  _IEEE Transactions on Information Theory_ , 56(9):4228–4244, 2010.
  * Hastie et al. [2009] Trevor Hastie, Robert Tibshirani, Jerome H Friedman, and Jerome H Friedman.  _The Elements of Statistical Learning: Data Mining, Inference, and Prediction_.  Springer, 2009.
  * Hubert and Arabie [1985] Lawrence Hubert and Phipps Arabie.  Comparing partitions.  _Journal of Classification_ , 2(1):193–218, 1985\.
  * Joe [1996] Harry Joe.  Time series models with univariate margins in the convolution-closed infinitely divisible class.  _Journal of Applied Probability_ , 33(3):664–677, 1996.
  * Lange et al. [2004] Tilman Lange, Volker Roth, Mikio L Braun, and Joachim M Buhmann.  Stability-based validation of clustering solutions.  _Neural Computation_ , 16(6):1299–1323, 2004\.
  * Leiner et al. [2022] James Leiner, Boyan Duan, Larry Wasserman, and Aaditya Ramdas.  Data fission: splitting a single data point.  _arXiv preprint arXiv:2112.11079_ , 2022.
  * Lopez et al. [2018] Romain Lopez, Jeffrey Regier, Michael B Cole, Michael I Jordan, and Nir Yosef.  Deep generative modeling for single-cell transcriptomics.  _Nature Methods_ , 15(12):1053–1058, 2018.
  * Love et al. [2014] Michael I Love, Wolfgang Huber, and Simon Anders.  Moderated estimation of fold change and dispersion for RNA-seq data with DESeq2.  _Genome Biology_ , 15(12):1–21, 2014.
  * McKenzie [1986] Ed McKenzie.  Autoregressive moving-average processes with negative-binomial and geometric marginal distributions.  _Advances in Applied Probability_ , 18(3):679–705, 1986.
  * Neufeld et al. [2022] Anna Neufeld, Lucy L Gao, Joshua Popp, Alexis Battle, and Daniela Witten.  Inference after latent variable estimation for single-cell RNA sequencing data.  _Biostatistics_ , 2022.
  * Neufeld et al. [2023] Anna Neufeld, Ameer Dharamshi, Lucy L Gao, and Daniela Witten.  Data thinning for convolution-closed distributions.  _arXiv preprint arXiv:2301.07276_ , 2023.
  * Owen and Perry [2009] Art B Owen and Patrick O Perry.  Bi-cross-validation of the svd and the nonnegative matrix factorization.  _The Annals of Applied Statistics_ , 3(2):564–594, 2009.
  * Sarkar and Stephens [2021] Abhishek Sarkar and Matthew Stephens.  Separating measurement and expression models clarifies confusion in single-cell RNA sequencing analysis.  _Nature Genetics_ , 53(6):770–777, 2021.
  * Tibshirani and Walther [2005] Robert Tibshirani and Guenther Walther.  Cluster validation by prediction strength.  _Journal of Computational and Graphical Statistics_ , 14(3):511–528, 2005.
  * Townes et al. [2019] F. W. Townes, S. C. Hicks, M. J. Aryee, and R. A. Irizarry.  Feature selection and dimension reduction for single-cell RNA-Seq based on a multinomial model.  _Genome Biology_ , 20(1):1–16, 2019.
  * Ullmann et al. [2022] Theresa Ullmann, Christian Hennig, and Anne-Laure Boulesteix.  Validation of cluster analysis results on validation data: A systematic framework.  _Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery_ , 12(3):e1444, 2022.
  * Van den Berge et al. [2020] Koen Van den Berge, Hector Roux De Bezieux, Kelly Street, Wouter Saelens, Robrecht Cannoodt, Yvan Saeys, Sandrine Dudoit, and Lieven Clement.  Trajectory-based differential expression analysis for single-cell sequencing data.  _Nature Communications_ , 11(1):1–13, 2020.
  * Zhang et al. [2019] Jesse M Zhang, Govinda M Kamath, and N Tse David.  Valid post-clustering differential analysis for single-cell RNA-Seq.  _Cell Systems_ , 9(4):383–392, 2019.

##  Appendix A Implementation details for Figure 1

To generate Figure 1, we first generate a toy dataset $X\in\mathbb{Z}_{\geq
0}^{100\times 2}$, where each element $X_{ij}$ is drawn independently from the
$\mathrm{NB}(5,5)$ distribution, which has mean $5$ and variance $10$ (see
Section 2.2 for parameterization details). In panels (c) and (d), the naive
method that uses the data twice and our proposed method are implemented as
described in Section 4. Here, we provide details about the implementation of
sample splitting used in Figure 1.

We begin by splitting the cells such that the first $n/2$ cells belong to the
training set, and the remaining cells belong to the test set, where $n=100$.

We first describe Figure 1(c). For values of $k$ ranging from $1$ to $10$, we
cluster the first $n/2$ rows of the matrix $\log(X+1)$ using k-means
clustering with $K$ clusters. This yields estimated cluster assignments
$\hat{c}_{i}$ only for cells $i=1,\ldots,n/2$. We then apply $3$-nearest
neighbors classification to label each cell in the test set with the majority-
label from its three nearest neighbors in the training set (using Euclidean
distance on the log-transformed data). We then compute an estimated mean
$\hat{\mu}_{ij}$ for each datapoint $X_{ij}$ as

$\hat{\mu}_{ij}=\frac{1}{\sum_{i^{\prime}=1}^{n/2}\mathbb{1}\\{\hat{c}_{i}^{\prime}=\hat{c}_{i}\\}}\sum_{i^{\prime}=1}^{n/2}X_{i^{\prime},j}\mathbb{1}\\{\hat{c}_{i}^{\prime}=\hat{c}_{i}\\},$
(3)

which is the sample mean of the training set data points belonging to this
cluster. Finally, we compute the within-cluster mean-squared error as

$\frac{1}{n/2\times
2}\sum_{i=n/2+1}^{n}\sum_{j=1}^{2}\left(\log(X_{ij}+1)-\log(\hat{\mu}_{ij}+1)\right)^{2}.$
(4)

Sample splitting is implemented in a similar way in Figure 1(d). We run
$k$-means clustering with $k=2$ on the logged training data. This yields
cluster assignments $\hat{c}_{i}$ only for $i=1,\ldots,n/2$. We once again
obtain cluster assignments $\hat{c}_{i}$ for $i=n/2+1,\ldots,n$ by applying
$3$-nearest neighbors. Finally, for $j=1$ and for $j=2$, we fit a negative
binomial generalized linear model where the response is $X_{ij}$ and the
covariate is $\hat{c}_{i}$, for $i=n/2+1,\ldots,n$.

In both cases, sample splitting fails because $\hat{c}_{i}$ for
$i=n/2+1,\ldots,n$ is obtained using the data $X_{ij}$ for $i=n/2+1,\ldots,n$,
via the $3$-nearest neighbors classification step. Thus, the test set is not
truly “held out” in computing $\hat{c}_{i}$ for $i=n/2+1,\ldots,n$.

To create panels (e) and (f) of Figure 1, we apply Algorithms 5 and 6 to the
single realization of toy data shown in Figure 1. In both algorithms, we
estimate $k=5$ clusters by running $k$-means on the log-transformed data. For
Algorithm 5, we use a support vector machine (SVM) with a linear kernel as the
classifier.

##  Appendix B Proofs for Section 3

###  B.1 Proof of Theorem 2

This result follows from applying Theorem 3 from [25] to each individual
element $X_{ij}$ in the matrix $X$, in the specific case where the $X_{ij}$
are independent negative binomial random variables. We note that [25] use a
different parameterization of the negative binomial: for
$\mathbf{y}\sim\operatorname{\mathrm{NB}}(r,p)$, they have
$\operatorname{\text{E}}[\mathbf{y}]=r\frac{1-p}{p}$ and
$\operatorname{\mathrm{Var}}(\mathbf{y})=r\frac{1-p}{p^{2}}$. This
parameterization of the negative binomial is convolution-closed in the
parameter $r$ if $p$ is held fixed, and corresponds to our parameterization if
$\mu=r\frac{1-p}{p}$ and $b=r$.

###  B.2 Proof of Theorem 3

We note that, when $b_{j}=\infty$, for $m=1,\ldots,M$,
$\mathbf{X}_{ij}^{(m)}\mid\mathbf{X}_{ij}=X_{ij}\sim\mathrm{Binomial}(X_{ij},\epsilon_{m})$
and
$X^{(-m)}\mid\mathbf{X}_{ij}=X_{ij}\sim\mathrm{Binomial}(X_{ij},1-\epsilon_{m})$.
Armed with these two facts, the first statement of Theorem 3 is proved in [15]
and [20], and the second statement of Theorem 3 is proved in [24].

###  B.3 Proof of Theorem 4

While parts of this theorem were proved in [24], we prove this theorem in full
directly here so that the notation matches that of this paper.

The first statement of Theorem 4 follows directly from the law of total
expectation. Regardless of the value of $b^{\prime}_{j}$,

$\displaystyle\operatorname{\text{E}}\left[\mathbf{X}^{(m)}_{ij}\right]$
$\displaystyle=\operatorname{\text{E}}\left[\operatorname{\text{E}}\left[\mathbf{X}^{(m)}_{ij}\mid\mathbf{X}=X_{ij}\right]\right]=\operatorname{\text{E}}\left[\epsilon_{m}\mathbf{X}_{ij}\right]=\epsilon_{m}\mu_{ij},$
$\displaystyle\operatorname{\text{E}}\left[\mathbf{X}^{(-m)}_{ij}\right]$
$\displaystyle=\operatorname{\text{E}}\left[\operatorname{\text{E}}\left[\mathbf{X}^{(-m)}_{ij}\mid\mathbf{X}=X_{ij}\right]\right]=\operatorname{\text{E}}\left[(1-\epsilon_{m})\mathbf{X}_{ij}\right]=(1-\epsilon_{m})\mu_{ij}.$

This is because the parameters $b_{j}^{\prime}$ do not affect the expected
values of the
$\mathrm{DirichletMultinomial}\left(X_{ij},\epsilon_{1}b_{j}^{\prime},\ldots,\epsilon_{M}b_{j}^{\prime}\right)$
distribution, only the variance.

The second statement of Theorem 4 uses the law of total variance. We start by
deriving the marginal variance of
$\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right)$, as follows:

$\displaystyle\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right)$
$\displaystyle=\operatorname{\text{E}}\left[\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\mid\mathbf{X}_{ij}=X_{ij}\right)\right]+\operatorname{\mathrm{Var}}\left(\operatorname{\text{E}}\left[\mathbf{X}^{(m)}_{ij}\mid\mathbf{X}_{ij}=X_{ij}\right]\right).$

As
$\mathbf{X}^{(m)}_{ij}\mid\mathbf{X}_{ij}=X_{ij}\sim\operatorname{\mathrm{BetaBinomial}}\left(X_{ij},\epsilon_{m}b_{j}^{\prime},(1-\epsilon_{m})b_{j}^{\prime}\right)$,
we plug in the (known) mean and variance of the beta-binomial distribution.

$\displaystyle\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right)$
$\displaystyle=\operatorname{\text{E}}\left[\frac{\mathbf{X}_{ij}\epsilon_{m}(1-\epsilon_{m})(b_{j}^{\prime}+\mathbf{X})}{(b_{j}^{\prime}+1)}\right]+\operatorname{\mathrm{Var}}\left(\epsilon_{m}\mathbf{X}_{ij}\right)$
$\displaystyle=\frac{\epsilon_{m}(1-\epsilon_{m})b_{j}^{\prime}}{(b_{j}^{\prime}+1)}\operatorname{\text{E}}\left[\mathbf{X}\right]+\frac{\epsilon_{m}(1-\epsilon_{m})}{(b_{j}^{\prime}+1)}\operatorname{\text{E}}\left[\mathbf{X}^{2}\right]+\epsilon_{m}^{2}\operatorname{\mathrm{Var}}\left(\mathbf{X}_{ij}\right).$

We next plug in the known mean and variance of
$\mathbf{X}_{ij}\sim\operatorname{\mathrm{NB}}(\mu_{ij},b_{j})$.

$\displaystyle\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right)$
$\displaystyle=\frac{\epsilon_{m}(1-\epsilon_{m})b_{j}^{\prime}\mu_{ij}}{(b^{\prime}+1)}+\frac{\epsilon_{m}(1-\epsilon_{m})}{(b_{j}^{\prime}+1)}\left(\mu_{ij}+\frac{\mu_{ij}^{2}}{b_{j}}+\mu_{ij}^{2}\right)+\epsilon_{m}^{2}\left(\mu_{ij}+\frac{\mu_{ij}^{2}}{b_{j}}\right)$
$\displaystyle=\epsilon_{m}\mu_{ij}+\frac{\epsilon_{m}^{2}\mu_{ij}^{2}}{b_{j}}+\frac{\epsilon_{m}(1-\epsilon_{m})\mu^{2}}{(b_{j}^{\prime}+1)}\left(\frac{1}{b_{j}}+1\right).$

To compare the magnitude of this variance to
$\epsilon_{m}\operatorname{\mathrm{Var}}(\mathbf{X})$, we add and subtract
$\epsilon_{m}\frac{\mu_{ij}^{2}}{b_{j}}$.

$\displaystyle\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right)$
$\displaystyle=\left[\epsilon_{m}\mu_{ij}+\epsilon_{m}\frac{\mu_{ij}^{2}}{b_{j}}\right]-\left[\epsilon_{m}\frac{\mu_{ij}^{2}}{b_{j}}-\epsilon_{m}^{2}\frac{\mu_{ij}^{2}}{b_{j}}\right]+\frac{\epsilon_{m}(1-\epsilon_{m})\mu_{ij}^{2}}{(b_{j}^{\prime}+1)}\left(\frac{1}{b_{j}}+1\right)$
$\displaystyle=\epsilon_{m}\operatorname{\mathrm{Var}}(\mathbf{X})-\left[\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\right]+\frac{\epsilon_{m}(1-\epsilon_{m})\mu_{ij}^{2}}{(b_{j}^{\prime}+1)}\left(\frac{1}{b_{j}}+1\right)$
$\displaystyle=\epsilon_{m}\operatorname{\mathrm{Var}}(\mathbf{X})+\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(\frac{b_{j}+1}{b_{j}^{\prime}+1}-1\right),$

as claimed in Theorem 4. We omit the derivation of
$\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(-m)}_{ij}\right)$, as it is
identical to the derivation above after noting that
$\operatorname{\mathrm{Var}}(\mathbf{X}^{(-m)}_{ij}\mid\mathbb{X}_{ij}=X_{ij})=\operatorname{\mathrm{Var}}(\mathbf{X}_{ij}-\mathbf{X}^{(m)}_{ij}\mid\mathbb{X}_{ij}=X_{ij})=\operatorname{\mathrm{Var}}(\mathbf{X}^{(m)}_{ij}\mid\mathbb{X}_{ij}=X_{ij})$.

Finally, for the third statement of Theorem 4, we use the fact that

$2\times\mathrm{Cov}(\mathbf{X}^{(-m)}_{ij},\mathbf{X}^{(m)}_{ij})=\operatorname{\mathrm{Var}}(\mathbb{X}_{ij})-\operatorname{\mathrm{Var}}(\mathbf{X}^{(-m)}_{ij})-\operatorname{\mathrm{Var}}\left(\mathbf{X}^{(m)}_{ij}\right).$

We then plug in the known values of these variances, and simplify, as follows:

$\displaystyle
2\times\mathrm{Cov}(\mathbf{X}^{(-m)}_{ij},\mathbf{X}^{(m)}_{ij})$
$\displaystyle=\operatorname{\mathrm{Var}}(\mathbb{X}_{ij})-\epsilon_{m}\operatorname{\mathrm{Var}}(\mathbb{X}_{ij})-\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(\frac{b_{j}+1}{b_{j}^{\prime}+1}-1\right)$
$\displaystyle\ \ \ \ \
-(1-\epsilon_{m})\operatorname{\mathrm{Var}}(\mathbb{X}_{ij})-\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(\frac{b_{j}+1}{b_{j}^{\prime}+1}-1\right)$
$\displaystyle=-2\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(\frac{b_{j}+1}{b_{j}^{\prime}+1}-1\right).$

Thus,
$\mathrm{Cov}(\mathbf{X}^{(-m)}_{ij},\mathbf{X}^{(m)}_{ij})=\epsilon_{m}(1-\epsilon_{m})\frac{\mu_{ij}^{2}}{b_{j}}\left(1-\frac{b_{j}+1}{b_{j}^{\prime}+1}\right)$.

###  B.4 Proof of Theorem 5

We first derive the Fisher information contained in $X_{ij}$ for the parameter
$\mu_{ij}$:

$\displaystyle I_{\mu_{ij}}\left(\mathbb{X}_{ij}\right)$
$\displaystyle=-\operatorname{\text{E}}\left[\frac{d^{2}}{d\mu_{ij}^{2}}\log\left(f\left(\mathbb{X}_{ij}\mid\mu_{ij},b_{j}\right)\right)\right]$
$\displaystyle=-\operatorname{\text{E}}\left[\frac{d^{2}}{d\mu_{ij}^{2}}\left(\log\left(\frac{\Gamma(\mathbb{X}_{ij}+b_{j})}{\Gamma(b_{j})\mathbb{X}_{ij}!}\right)+\mathbb{X}_{ij}\log\left(\frac{\mu_{ij}}{\mu_{ij}+b_{j}}\right)+b_{j}\log\left(\frac{b_{j}}{\mu_{ij}+b_{j}}\right)\right)\right]$
$\displaystyle=-\operatorname{\text{E}}\left[\mathbb{X}_{ij}\left(\frac{-1}{\mu_{ij}^{2}}+\frac{1}{(\mu_{ij}+b_{j})^{2}}\right)+b_{j}\left(\frac{1}{(\mu_{ij}+b_{j})^{2}}\right)\right]$
$\displaystyle=\frac{1}{\mu_{ij}}-\frac{\mu_{ij}+b_{j}}{(\mu_{ij}+b_{j})^{2}}=\frac{b_{j}}{\mu_{ij}(\mu_{ij}+b_{j})}.$

Now we derive the Fisher information contained in $\mathbf{X}^{(m)}_{ij}$ for
the parameter $\mu_{ij}$:

$\displaystyle I_{\mu_{ij}}\left(\mathbf{X}^{(m)}_{ij}\right)$
$\displaystyle=-\operatorname{\text{E}}\left[\frac{d^{2}}{d\mu_{ij}^{2}}\log\left(f\left(\mathbf{X}^{(m)}_{ij}\mid\operatorname{\epsilon}\mu_{ij},\operatorname{\epsilon}b_{j}\right)\right)\right]$
$\displaystyle=-\operatorname{\text{E}}\left[\frac{d^{2}}{d\mu_{ij}^{2}}\left(\log\left(\frac{\Gamma(\mathbb{X}_{ij}^{(m)}+\epsilon_{m}b_{j})}{\Gamma(\epsilon_{m}b_{j})(\mathbb{X}_{ij}^{(m)})!}\right)+\mathbb{X}^{(m)}_{ij}\log\left(\frac{\epsilon_{m}\mu_{ij}}{\epsilon_{m}\mu_{ij}+\epsilon_{m}b_{j}}\right)+\epsilon_{m}b_{j}\log\left(\frac{\epsilon_{m}b_{j}}{\epsilon_{m}\mu_{ij}+\epsilon_{m}b_{j}}\right)\right)\right]$
$\displaystyle=-\operatorname{\text{E}}\left[\mathbb{X}_{ij}^{(m)}\left(\frac{-1}{\mu_{ij}^{2}}+\frac{1}{(\mu_{ij}+b_{j})^{2}}\right)+\epsilon_{m}b_{j}\left(\frac{1}{(\mu_{ij}+b_{j})^{2}}\right)\right]$
$\displaystyle=\frac{\epsilon_{m}}{\mu_{ij}}-\frac{\epsilon_{m}\mu_{ij}+\epsilon_{m}b_{j}}{(\mu_{ij}+b_{j})^{2}}=\epsilon_{m}I_{\mu_{ij}}\left(\mathbb{X}_{ij}\right).$

Theorem 2 guarantees independence between $\mathbf{X}^{(m)}_{ij}$ and
$\mathbf{X}^{(-m)}_{ij}$, and therefore
$I_{\mu_{ij}}\left(\mathbf{X}^{(-m)}_{ij}\right)+I_{\mu_{ij}}\left(\mathbf{X}^{(m)}_{ij}\right)=I_{\mu_{ij}}\left(\mathbb{X}_{ij}\right)$.
It follows that
$I_{\mu_{ij}}\left(\mathbf{X}^{(-m)}_{ij}\right)=(1-\epsilon_{m})I_{\mu_{ij}}\left(\mathbb{X}_{ij}\right)$.

##  Appendix C Implementation details for Section 4

We use the R package sctransform in the following manner to estimate gene-
specific overdispersion parameters in the simulations for Section 4.

Briefly, for each gene $j=1,\ldots,p$, sctransform begins by fitting a
negative binomial GLM with $\mathbb{X}_{j}$ as the response, and the logged
total number of unique molecular identifiers (UMIs) for the $n$ cells as the
covariate. (In our simulations, the total number of UMIs for a cell is the row
sum for that cell.) This yields a maximum likelihood estimate
$\hat{b}_{j}^{\mathrm{MLE}}$ for each gene $j=1,\ldots,p$. These maximum
likelihood estimates are known to be quite noisy for sparse negative binomial
data. Furthermore, the “null model” that only includes the total number of
UMIs as a covariate may be the correct model for the majority of the genes,
but will be incorrect for any genes that exhibit true differential expression
across unknown latent variables. Thus, as a second step, sctransform fits a
smooth kernel regression to estimate a relationship between the average
expression of each gene and the gene-specific overdispersion. These smoothed
estimates are used as the gene-specific overdispersion parameters.

We run the vst() function from the sctransform package in R with its default
settings in Section 4. We note that we simulated data in which the two main
assumptions of sctransform are met: most genes are not differentially
expressed, and there is a smooth relationship between the average expression
of a gene and its parameter $b_{j}$. We use the same strategy in Section 5,
although in the real data setting we do not know for sure that the modeling
assumptions of sctransform are met.

##  Appendix D Implementation details for Section 5

In Section 5, we analyze a publicly-available dataset that is associated with
[4], which can be downloaded from https://descartes.brotmanbaty.org/. For the
main cell type analysis, we used all cells from this dataset that were
collected from the kidney. For the cell subtype analysis, we used all of the
kidney cells that were labeled as metanephric cells in the original analysis
by [4]. For both analyses, we filtered to genes with non-zero counts in at
least 10 cells. After this subsetting, the kidney dataset has dimension
$178,603\times 34,714$ and the metanephric dataset has dimension $90,876\times
31,385$.

To carry out Step 1 of Algorithm 5 on our two datasets, we used the Monocle3
package. To preprocess each dataset, we generated a 50-dimensional principal
components embedding of each dataset and subsequently a 2-dimensional UMAP
embedding using the default settings of the preprocessing functions in the
Monocle3 package. Next, we performed Leiden clustering using the the Monocle3
clustering function. We chose a resolution parameter that gave a similar
number of clusters to those obtained in the original paper. More specifically,
we set the resolution parameter to $1\times 10^{-6}$ for the full kidney
dataset and $1\times 10^{-5}$ for the metanephric cell subset. To carry out
Step 2 of Algorithm 5, we split the 50-dimensional principal components
embedding of the full count data from kidney and metanephric cells into five
folds containing equal numbers of cells. We used the cluster labels, inferred
as described above (on the full counts), as the “true” cluster labels. For
each of the five folds, we then trained a linear SVM model to predict the
cluster assignment from 80% of the embedded expression data. We generated a
confusion matrix by comparing the “true” labels to this trained model’s
predictions on the held-out subset. We note that this SVM is slightly
different from that of [4], who trained their SVM using the whole
transcriptome rather than the reduced-dimension embedding.

To carry out the “assume Poisson” version of Algorithm 6 on our two datasets,
we performed Poisson count splitting (Algorithm 1, or, equivalently, Algorithm
2 with $b^{\prime}_{j}=\infty$) with $M=2$ folds and
$\epsilon_{1}=\epsilon_{2}=0.5$ on each dataset to obtain an $X^{(1)}$ and an
$X^{(2)}$ for each of the datasets. We then followed the preprocessing and
clustering procedures outlined above on each fold for each dataset, to obtain
two clusterings for each dataset. To produce Figures 8(b) and 8(f), we re-
ordered the test set labels on the $y$-axis to make the confusion matrix as
diagonal as possible (since the clusters are invariant to re-labeling).

Finally, to carry out the negative binomial version of Algorithm 6 on our two
datasets, we applied the sctranstorm package in R with its default parameters
to estimate overdispersion values for each gene (see Appendix C). We then
performed negative binomial count splitting (Algorithm 2) with $M=2$ folds and
$\epsilon_{1}=\epsilon_{2}=0.5$ on each dataset to obtain an $X^{(1)}$ and an
$X^{(2)}$ for each of the datasets. We then proceeded as in the Poisson case.