# Tree stick-breaking priors for covariate-dependent mixture models

Akira Horiguchi  Department of Statistical Science, Duke University
214 Old Chemistry, Box 90251, Durham, NC 27708  Cliburn Chan  Department of
Biostatistics and Bioinformatics, Duke University
2424 Erwin Road, Suite 1102, Hock Plaza Box 2721, Durham, NC 27710  Department
of Statistical Science, Duke University
214 Old Chemistry, Box 90251, Durham, NC 27708  Li Ma Send correspondence to
<EMAIL_ADDRESS>Department of Statistical Science, Duke University
214 Old Chemistry, Box 90251, Durham, NC 27708  Department of Biostatistics
and Bioinformatics, Duke University
2424 Erwin Road, Suite 1102, Hock Plaza Box 2721, Durham, NC 27710

###### Abstract

Stick-breaking priors are often adopted in Bayesian nonparametric mixture
models for generating mixing weights. When covariates influence the sizes of
clusters, stick-breaking mixtures can leverage various computational
techniques for binary regression to ease posterior computation. Existing
stick-breaking priors are typically based on continually breaking a single
remaining piece of the unit stick. We demonstrate that this “single-piece”
scheme can induce three highly undesirable behaviors; these behaviors are
circumvented by our proposed model which continually breaks off all remaining
pieces of a unit stick while keeping posterior computation essentially
identical. Specifically, the new model provides more flexibility in setting
cross-covariate prior correlation among the generated random measures,
mitigates the impact of component label switching when posterior simulation is
performed using Markov chain Monte Carlo, and removes the imposed artificial
decay of posterior uncertainty on the mixing weights according to when the
weight is “broken off” the unit stick. Unlike previous works on covariate-
dependent mixtures, which focus on estimating covariate-dependent
distributions, we instead focus on inferring the effects of individual
covariates on the mixture weights in a fashion similar to classical regression
analysis, and propose a new class of posterior predictives for summarizing
covariate effects.

Keywords: Discrete random measure, Bayesian nonparametric, tail-free process,
clustering analysis, flow cytometry.

##  1 Introduction

Mixture models are a popular approach to clustering analysis. Stick-breaking
(SB) processes are widely used in the Bayesian literature as priors on mixture
models (Sethuraman, 1994; Ishwaran and James, 2001) largely due to their
tractable posterior computation. SB priors are particularly well-suited when
covariates influence cluster sizes, as these priors can achieve efficient
posterior computation by incorporating computational strategies involving
latent variable representation and data augmentation techniques originally
designed for binary outcome regression models. Notable examples include
Rodriguez and Dunson (2011)’s probit SB whose Markov chain Monte Carlo (MCMC)
algorithm relies on Albert and Chib (1993)’s truncated Gaussian data
augmentation for probit regression, and Rigon and Durante (2021)’s logit SB
whose various posterior sampling methods rely on Polson et al. (2013)’s Pólya-
gamma data augmentation for logistic regression. SB priors have also been
applied in generalizations of mixture models such as mixed-membership models
(Linderman et al., 2015).

Covariate-dependent clustering applications range widely; one motivating
example for this paper is flow cytometry (FCM), a technique whose uses include
identifying cells according to their type (e.g. B-cell, T-cell), in which the
lack of existing nonparametric mixture models based on SB priors become
apparent. In FCM, a question of interest is to see how cell populations change
across covariate values. For example, one might wish to detect cell-type
population differences between healthy and sick patients. Very often, the
covariate values influence relatively small cell clusters.

Existing SB priors are based on continually breaking a single remaining piece
of the unit stick. As we shall demonstrate through both theoretical
investigation and numerical experiments, this “single-piece” scheme can induce
three highly undesirable properties that severely impair posterior inference
in the presence of covariates.

The first is that under common default hyperprior values, “single-piece” stick
breaking induces a rather strong positive prior correlation in the random
measures over covariate values, which can cause excessive borrowing of
information across the covariates even when the corresponding cluster sizes
vary sharply across covariate values, thereby deprecating the clustering at
each covariate value. More concretely, Corollary 2.1 in Section 2.3 states
that if the stick breaks have prior mean $0.5$ with non-negative pairwise
correlations, then the prior cross-predictor correlation between random
measures is lower-bounded by $1/3$. Rodriguez and Dunson (2011) noticed this
lower bound and explained it as “a consequence of our use of a common set of
atoms at every [covariate]; even if the set of weights are independent from
each other, the fact that the atoms are shared means that the distributions
cannot be independent.” We show that the large size of this baseline
correlation is determined largely by the particular way the weights are
constructed from the stick breaks.

Another undesirable property involves the long-standing problem of component
label switching. In a Bayesian mixture model with $K$ components under a prior
with exchangeable mixture components, the prior and the resulting posterior
cannot distinguish the $K!$ permutations of the mixture component labels and
hence are symmetric between the corresponding $K!$ regions of the parameter
space $\Theta$. If a posterior simulation stays in a single region, the
simulation’s outputs can be interpreted, but the more a simulation travels
between these regions, the more difficult the interpretation becomes. Section
3 contains a simulation study in which the MCMC simulation travels between
eight of these regions even after a long burn-in. We find that the ease with
which this MCMC simulation travels between so many regions is not inherent to
SB priors; rather, it is specific to the “single-piece” approach to generating
weights. We also find that the simulation’s interpretation cannot in general
be ”fully restored” by Cron and West (2011)’s relabeling algorithm.

A third undesirable behavior is that a weight’s variance tends to decrease
according to the order in which the weight was created during stick breaking.
To intuit why this phenomenon occurs, note that any weight
$W_{k}=V_{k}\prod_{l<k}(1-V_{l})$ must follow

$\displaystyle\texttt{Var}W_{k}+[\mathbb{E}W_{k}]^{2}=\mathbb{E}W_{k}^{2}=\mathbb{E}V_{k}^{2}\mathbb{E}R_{k-1}^{2}$

where $R_{k-1}\coloneqq\prod_{l<k}(1-V_{l})$ is the length of the remaining
right piece after $k-1$ stick breaks. Each $\texttt{Var}W_{k}$ is bounded
above by $\mathbb{E}R_{k-1}^{2}=\prod_{l<k}\mathbb{E}(1-V_{l})^{2}$ which is a
product of $k-1$ quantities strictly between zero and one. So in the likely
scenario that $\mathbb{E}R_{k-1}^{2}$ shrinks to zero (or to a very small
positive number) as $k$ increases, then $\texttt{Var}W_{k}$ must also.
Although this artificial variance decay is easily unnoticed in our simulation
study in Section 3, where all of the weight variances are small (likely
because we are able to perfectly specify the weights model), it is strikingly
present in our analysis of real-world FCM data in Section 4 (see Figure 8(a)),
where a few of the weights have large variances (likely because it is often
impossible to impose a perfectly specified model on real-world data). Seeing
as this variance decay is a product of this “single-piece” stick-breaking
construction rather than of any particular qualities of the data, it becomes
difficult to trust the uncertainty quantification of these weights.

We aim to construct new SB priors for covariate-dependent mixture modeling
while maintaining its desirable computational properties. To this end, we
adopt a dyadic tree perspective on stick-breaking models. Specifically,
starting from the whole unit stick as the root of a tree, we view each piece
of the stick that arises during the stick-breaking as a node in a bifurcating
tree. The traditional “single-piece” stick breaking then corresponds to a
special tree topology—which is “lopsided” on one side. See Figure 1. The
issues of traditional stick-breaking processes become clear under the tree
perspective. We show that the lopsided structure lies behind the
aforementioned issues of SB models.

Lopsided tree: $K=16$ important/.style=fill=red,edge=red,ultra thick, delay=
aggregate postparse=int, for tree children-first= if n children=0
circle,draw=black,minimum size=16pt rectangle,draw=black,minimum size=16pt
[1,important,for tree=s sep=8mm,level distance=2mm
[][2,important[][3,important[][4,important[][5,important[][6,important[][7,important[][8,important[][9,important[9,important]
[¡-,red,ultra thick] (.south)–++(0em,-8ex)–++(-2em,0pt)
node[anchor=east,align=center]$(1-V_{1})(1-V_{2})\cdots(1-V_{8})\,V_{9}=W_{9}$;
[[][[][[][[][[][[][]]]]]]]]]]]]]]]]

Balanced tree: $K=16$ important/.style=fill=red,edge=red,ultra thick, delay=
aggregate postparse=int, for tree children-first= if n children=0
circle,draw=black,minimum size=16pt rectangle,draw=black,minimum size=16pt
[1,important,for tree=level distance=2mm
[[[[][]][[][]]][[[][]][[][]]]][3,important[6,important[12,important[9,important]
[¡-,red,ultra thick] (.south)–++(0em,-8ex)–++(-2em,0pt)
node[anchor=east,align=center]$(1-V_{1})\,V_{3}\,V_{6}\,V_{12}=W_{9}$;
[]][[][]]][[[][]][[][]]]] ]

Figure 1: Binary tree representation of stick-breaking schemes. Each scheme’s
initial stick break is represented as a root node; each subsequent stick break
is represented by an internal node. The mass of a stick-break’s “left (right)
piece” is sent to its left (right) child. Thus, the $K-1$ stick breaks
distribute the initial unit mass to the tree’s $K$ leaf nodes which represent
the scheme’s $K$ weights. As an example, the filled leaf in each tree
represents weight $W_{9}$ whose value is determined by the stick breaks whose
representative internal nodes are the filled leaf’s ancestors. This weight is
a product of nine stick breaks in the lopsided-tree scheme but only four stick
breaks in the balanced-tree scheme.

We instead propose a new class of SB priors based on an alternative tree
topology—namely a “balanced” one (shown in Figure 1), which corresponds to the
continual breaking of all remaining pieces of the unit stick at each stage of
the stick breaking. Under this scheme, each weight becomes a multiplicative
product of the same (small) number of stick breaks. No weight is a product of
just one stick break, which allows this alternative scheme to offer a much
wider range of prior cross-covariate correlation in the resulting random
measures. Furthermore, this tree topology mitigates the impact of label
switching as shown empirically in Section 3.

Posterior computation under our new covariate-dependent SB model remains
essentially identical to that of existing SB mixtures. In particular, this
scheme can easily adapt existing (and possibly future) algorithms for
tractable posterior inference, e.g., Rodriguez and Dunson (2011)’s Markov
chain Monte Carlo (MCMC) algorithm for probit SB and Rigon and Durante
(2021)’s MCMC, variational Bayes, and expectation-maximization algorithms for
logistic SB. In this work, we shall demonstrate using Pólya-gamma augmentation
for existing logistic SB priors under the new balanced-tree SB mixture models.
We note that our new SB prior for mixing weights along a balanced bifurcating
tree is a tail-free process, and in the presence of covariates, it becomes a
covariate-dependent tail-free model introduced by Jara and Hanson (2011).
According to our knowledge, such processes have not been applied as a discrete
prior for mixing distributions except in Cipolli and Hanson (2017) who propose
the finite Pólya tree as a prior for the mixing distributions and Stefanucci
and Canale (2021) whose multiscale SB method generalizes our approach to
generating weights. Both works consider only univariate response data with no
covariate dependence, and neither work provides theoretical or methodological
justification for using their respective method over an approach based on
traditional SB.

Previous works (e.g. Rigon and Durante, 2021) on covariate-dependent random
measures have mostly evaluated the “goodness” of the models based on posterior
predictive densities. While this might be appropriate when density estimation
or prediction is the main objective, they are not suitable for applications
such as flow cytometry in which the key scientific interest lies in
quantifying the impact of a covariate (e.g., treatments or patient
characteristics) on the clusters (e.g., immune cell types). For the latter
purpose, we instead utilize the posterior distributions of covariate effects,
in the spirit of classical regression analysis, to examine model performance.
One complication is that under stick-breaking models (both existing and our
new model), covariate effects are parameterized on the stick-breaking
proportions rather than on the cluster weights, which is typically of direct
scientific interest. As such, we propose a strategy to translate posterior
covariate effects on stick-breaking proportions onto that on the weights.
While not the major focus of this paper, our strategy presents a practical
solution to practitioners who wish to more intuitively interpret covariates.

We end the introduction by reviewing several works that also utilize tree
structures to construct SB models in ways different from ours. To model
hierarchical data, Ghahramani et al. (2010) interleave two “lopsided” SB
processes to allow a wide range of latent tree structures (e.g. nodes with any
number of children) to be inferred, and Dubey et al. (2014)’s extension
incorporates time information into weights. In a different context, Nassar et
al. (2018) introduce what is essentially a generalization of our “balanced” SB
approach to generating covariate-dependent weights and even utilize PG
augmentation in their Gibbs sampler, but do so to model nonlinear dynamical
systems and hence do not consider its use in mixture models.

The paper is structured as follows. Section 2 introduces our SB scheme without
covariate dependence. Section 2.3 extends the model to incorporate covariate
dependence. Section 2.6 considers posterior computation. Section 3 showcases a
numerical example where tree structure exacerbates label switching on
posterior inference for weights. Section 4 applies our method to the GRIFOLS
FCM data set. Section 5 provides concluding remarks.

##  2 Method

###  2.1 Basic setup and notation

We start by introducing our tree-based stick-breaking construction in the
absence of covariates. Throughout we assume independent and identically
distributed (i.i.d.) observations $y_{1},\ldots,y_{n}\in\mathbb{R}^{d}$ from a
sampling density of the following mixture form

$f(\cdot;G)=\int_{\Theta}h(\cdot;\theta)\mathop{}\\!\mathrm{d}G(\theta).$ (1)

where $h\colon\mathbb{R}^{d}\times\Theta\rightarrow\mathbb{R}^{+}$ is a
parametric mixing kernel density, e.g., Gaussian density
$h(y;\theta)=(2\pi\sigma^{2})^{-1/2}\exp\left\\{-(y-\mu)^{2}/(2\sigma^{2})\right\\}$
with $\theta=(\mu,\sigma^{2})\in\mathbb{R}\times\mathbb{R}^{+}$.

We consider a mixing measure $G$ on $\Theta$ in (1) with SB priors. Thus each
realization of $G$ is a discrete measure of the following form with
probability one:

$G=\sum_{k=1}^{\infty}W_{k}\delta_{\theta_{k}}\qquad\theta_{k}\stackrel{{\scriptstyle
ind}}{{\sim}}G_{0}$

where $G_{0}$ is a base measure on $\Theta$ and the nonnegative weights
$W_{k}$, which are generated through breaking a unit stick, are independent of
the atoms $\theta_{k}$. This turns the mixture density (1) into
$f(\cdot;G)=\sum_{k=1}^{\infty}W_{k}h(\cdot;\theta_{k})$.

The remainder of this section examines the traditional SB scheme from a tree
perspective, and then introduces a tree-based generalization to SB priors.

###  2.2 Stick-breaking priors from a tree perspective

A weight-generation process aims to produce nonnegative quantities that sum to
unity. A SB process achieves this by successively breaking off pieces of a
stick of initially unit length, where the lengths of the resulting pieces
become the desired weights. In the traditional SB scheme, after each stick
break produces a piece that is untouched afterwards while the other piece—the
“remaining stick”—breaks again and again. More formally, to begin, a stick of
unit-length breaks at location $V_{1}\in[0,1]$ so that the piece that breaks
off has length $V_{1}$. The remaining stick, of length $1-V_{1}$, then breaks
at (relative) location $V_{2}\in[0,1]$ so that the new piece that breaks off
has length $(1-V_{1})V_{2}$. The remaining stick, now of length
$(1-V_{1})(1-V_{2})$, then breaks at location $V_{3}\in[0,1]$ so that the new
piece that breaks off has length $(1-V_{1})(1-V_{2})V_{3}$, so on and so
forth. For general $k\in\mathbb{N}$, the $k^{\text{th}}$ stick break at
location $V_{k}\in[0,1]$ produces a piece that breaks off with length
$\left[\prod_{l=1}^{k-1}(1-V_{l})\right]V_{k}$, which sets the value of weight
$W_{k}$; the remaining stick then breaks at location $V_{k+1}\in[0,1]$. If the
process could either continue infinitely or terminate after a preset finite
number of breaks, in which case the final remaining stick sets the last
weight.

The SB scheme can be identified with a bifurcating tree whose nodes correspond
to the pieces that arise during the stick breaking procedure. We construct the
binary tree by first assigning the initial unit-length stick to the tree’s
root node, and iteratively apply the following steps. If a piece of the stick
does not break further, then this piece corresponds to a leaf node, i.e., a
node with no children. If a piece is further broken, then we use two children
nodes to represent the two resulting pieces. The traditional SB scheme’s
identified with a bifurcating tree (see the top right of Figure 2), which is
the most “lopsided” as only no nodes in the tree can have both children that
are further divided. It is the deepest tree possible for generating a given
number of pieces. We shall thus refer to the traditional SB strategy as the
lopsided-tree (LT) SB.

Given this tree perspective of stick-breaking, it is natural to consider
alternative SB schemes that correspond to different tree topologies than the
lopsided one. To this end, it will be convenient to index each binary tree
node $\bm{\varepsilon}$ by a finite string of 0s and 1s, where each digit
indicates whether a node along the path from root to $\bm{\varepsilon}$ is a
left (0) or right (1) child of its parent. For notational ease, the root node
corresponds to the empty string $\emptyset$. Examples are shown in Figure 2.
In general, if a node $\bm{\varepsilon}$ is at level $m$ of the tree (the root
node is at level 0, its two children are at level 1, and so on), then
$\bm{\varepsilon}$ is the $m$-length string
$\varepsilon_{1}\varepsilon_{2}\cdots\varepsilon_{m}$, where each
$\varepsilon_{i}\in\\{0,1\\}$ and the string
$\bm{\varepsilon}\bm{\varepsilon}^{\prime}$ denotes the concatenation of any
two finite strings $\bm{\varepsilon}$ and $\bm{\varepsilon}^{\prime}$. The set
of all finite strings (including $\emptyset$) of 0s and 1s is
$E^{*}\coloneqq\cup_{m=0}^{\infty}\\{0,1\\}^{m}$.

With this notation, we can now formally introduce a SB strategies based on
general tree topologies. In particular, we give particular attention to the SB
scheme corresponding to a “balanced” bifurcating tree, in which all nodes (up
to a maximum level) are split into two children as illustrated at the bottom
of Figure 2. We refer to this particular SB scheme as a balanced-tree (BT) SB.
Opposite to the lopsided-tree scheme, the balanced-tree scheme results in the
most shallow tree topology needed to generate any given number of weights.

Specifically, in this new scheme, after each stick break _both_ pieces break
again until the total number of breaks exceeds a preset threshold (we argue
against allowing infinitely many breaks in Section 2.2.1). The unit-length
stick $I_{\emptyset}$ breaks according to $V_{\emptyset}$ so that the left
piece $I_{0}$ and right piece $I_{1}$ have lengths

$|I_{0}|=|I_{\emptyset}|V_{\emptyset}=V_{\emptyset},\qquad|I_{1}|=|I_{\emptyset}|(1-V_{\emptyset})=1-V_{\emptyset}.$

These two pieces $I_{0}$ and $I_{1}$ then break similarly according to
respective splitting variables $V_{0}$ and $V_{1}$ so that the resulting four
pieces $I_{00}$, $I_{01}$, $I_{10}$, and $I_{11}$ have lengths

$|I_{00}|=V_{\emptyset}V_{0},\qquad|I_{01}|=V_{\emptyset}(1-V_{0}),\qquad|I_{10}|=(1-V_{\emptyset})V_{1},\qquad|I_{11}|=(1-V_{\emptyset})(1-V_{1}).$

These four pieces $I_{00}$, $I_{01}$, $I_{10}$, and $I_{11}$ then break
similarly according to respective splitting variables $V_{00}$, $V_{01}$,
$V_{10}$, and $V_{11}$ to produce the eight pieces $I_{000}$, $I_{001}$,
$I_{010}$, $I_{011}$, $I_{100}$, $I_{101}$, $I_{110}$, and $I_{111}$ with
lengths defined similarly. If we preset the number of such collections of
stick breaks to $m\in\mathbb{N}$, this recursive procedure produces $2^{m}$
leaf nodes each with a stick piece. The lengths of these pieces define the
desired weights: $W_{\bm{\varepsilon}}\coloneqq|I_{\bm{\varepsilon}}|$ for
each $\bm{\varepsilon}\in\\{0,1\\}^{m}$.

The value of any SB weight at a node
$\bm{\varepsilon}=\varepsilon_{1}\cdots\varepsilon_{m}$ can be expressed as

$W_{\varepsilon_{1}\cdots\varepsilon_{m}}=\prod_{l=1}^{m}V_{\varepsilon_{1}\cdots\varepsilon_{l-1}}^{1-\varepsilon_{l}}(1-V_{\varepsilon_{1}\cdots\varepsilon_{l-1}})^{\varepsilon_{l}}\quad\quad
V_{\bm{\varepsilon}}\sim F_{\bm{\varepsilon}}\text{ for all
}\bm{\varepsilon}\in E^{*}$ (2)

where by convention $\varepsilon_{1}\cdots\varepsilon_{l-1}=\emptyset$ if
$l=1$, and the splitting variables are distributed according to a countable
sequence $\\{F_{\bm{\varepsilon}}\colon\bm{\varepsilon}\in E^{*}\\}$ of
distributions each with full support on $[0,1]$ and the tail-free condition
(Freedman, 1963; Fabius, 1964), i.e.

$\displaystyle\text{supp}(V_{\bm{\varepsilon}})=[0,1]\text{ for all
}\bm{\varepsilon}\in E^{*}$ (3) $\displaystyle
V_{\emptyset}\perp\\{V_{0},V_{1}\\}\perp\\{V_{00},V_{01},V_{10},V_{11}\\}\perp\cdots.$
(4)

$\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}{0}$$\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>
$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{\emptyset}}$

$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{0}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{1}}$

$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{10}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{11}}$

$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{110}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{111}}$

$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{\emptyset}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{0}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{1}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{10}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{11}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{110}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{111}}$

$\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}{0}$$\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{.5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color<EMAIL_ADDRESS>
$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{\emptyset}}$

$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{0}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{1}}$

$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{0}$$\color[rgb]{1,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{1,1,1}\pgfsys@color@gray@stroke{1}\pgfsys@color@gray@fill{1}{1}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{00}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{01}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{10}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{11}}$

$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{\emptyset}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{0}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{00}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{01}}$$\color[rgb]{1,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{1,0,0}{V_{1}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{10}}$$\color[rgb]{0,1,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,1,1}\pgfsys@color@cmyk@stroke{1}{0}{0}{0}\pgfsys@color@cmyk@fill{1}{0}{0}{0}{W_{11}}$

Figure 2: Illustration of LT (top row) and BT (bottom row) that terminate with
$K=4$ leaves. Shown are the unit-length stick breaks (left column) and binary
trees (right column). In either SB scheme, the weights sum to unity.

With this tree view, we can now define the tree stick-breaking (treeSB) class
of priors, which includes traditional SB but also admits more general tree
structures. We say that a probability measure $G$ is endowed with a treeSB
prior with parameters $G_{0}$, $\\{F_{\bm{\varepsilon}}\\}$, and binary tree
structure $\tau$, denoted by $G\sim
treeSB(G_{0},\\{F_{\bm{\varepsilon}}\\},\tau)$, if it can be constructed as

$\displaystyle G=\sum_{\bm{\varepsilon}\in
B(\tau)}W_{\bm{\varepsilon}}\delta_{\theta_{\bm{\varepsilon}}},\qquad\theta_{\bm{\varepsilon}}\stackrel{{\scriptstyle
ind}}{{\sim}}G_{0}$ (5)

where the set $B(\tau)\subset E^{*}$ indexes $\tau$’s leaf nodes, and the
random weights $\\{W_{\bm{\varepsilon}}\colon\bm{\varepsilon}\in B(\tau)\\}$
are constructed according to (2) where the distributions
$F_{\bm{\varepsilon}}$ follow the full-support and tail-free conditions (3)
and (4). As with traditional SB, this turns the mixture density (1) into
$f(\cdot;G)=\sum_{\bm{\varepsilon}\in
B(\tau)}W_{\bm{\varepsilon}}h(\cdot;\theta_{\bm{\varepsilon}})$.

####  2.2.1 Choice of the number of leaves

In real-life applications, a practitioner will use only finitely many stick
breaks (i.e. use $\tau$ with finite depth), in which case the treeSB
construction (5) is entirely valid. Special consideration must be taken if
$\tau$ has infinite depth. In the following we avoid these technical
considerations by assuming $\tau$ has large finite depth, which suffices for
practical uses of the model. We note that a similar simplication is often
adopted through truncation of traditional SB priors (Ishwaran and James,
2001).

How should a practitioner set the number of leaves $K$? For LT, the finite
approximation theorems of Ishwaran and James (2002) imply that the value of
$K$ beyond say $64$ seems to affect inference by at most a trivial amount.
This statement is supported by the moment expressions (8) and (9) to be
introduced in Theorem 1 and by Theorem 2. For BT, however, Theorem 2 implies
the possibility of setting $K$ so large that the variance of the random
measure $G$ shrinks to almost zero. This danger is based on identically
distributed splitting variables and hence might possibly be avoided by
choosing splitting variables whose mean and/or variance increase with tree
level. In any case, a rule of thumb that still allows “room to breathe” is to
set $K$ to be the smallest power of two greater than or equal to twice the
prior expected number of clusters.

###  2.3 Tree stick-breaking with covariates

Following the well-known strategy of MacEachern (2000), we next extend the
covariate-independent treeSB prior (5) by replacing the splitting variables
$\\{V_{\bm{\varepsilon}}\\}_{\bm{\varepsilon}\in E^{*}}$ in (2) with a
sequence of stochastic processes
$\\{V_{\mathbf{x},\bm{\varepsilon}}\colon\mathbf{x}\in\mathcal{X}\\}_{\bm{\varepsilon}\in
E^{*}}$, where $\mathcal{X}$ is a set of covariates and each splitting
variable $V_{\mathbf{x},\bm{\varepsilon}}$ has distribution
$F_{\mathbf{x},\bm{\varepsilon}}$. The resulting random probability measure
$G_{\mathbf{x}}\sim treeSB(G_{0},\\{F_{\mathbf{x},\bm{\varepsilon}}\\},\tau)$
now depends on $\mathbf{x}$ through its weights:

$\displaystyle G_{x}=\sum_{\bm{\varepsilon}\in
B(\tau)}W_{\mathbf{x},\bm{\varepsilon}}\delta_{\theta_{\bm{\varepsilon}}},\quad\quad\theta_{\bm{\varepsilon}}\stackrel{{\scriptstyle
ind}}{{\sim}}G_{0}.$ (6)

There are a number of possible strategies to incorporate covariate dependence
on weights, including utilizing probit and logit transform on the weights as
is commonly done for traditional SB. In particular, let us consider the logit
approach due to the particular computational convenience for posterior
inference that follows from the Pólya-Gamma augmentation technique, to be
detailed in Section 2.6. (We note that the theoretical properties we establish
in Section 2.4 do not assume this model choice and apply more generally.)

Specifically, we adopt the following logit-normal model on each splitting
variable:

$V_{\mathbf{x},\bm{\varepsilon}}=\text{logistic}(\eta_{\mathbf{x},\bm{\varepsilon}}),\qquad\eta_{\mathbf{x},\bm{\varepsilon}}=\bm{\psi}(\mathbf{x})^{T}\bm{\gamma}_{\bm{\varepsilon}},\qquad\bm{\gamma}_{\bm{\varepsilon}}\sim
N_{R}(\bm{\mu}_{\bm{\gamma}},\bm{\Sigma}_{\bm{\gamma}})$ (7)

with hyperparameters $\bm{\mu}_{\bm{\gamma}}$ and $\bm{\Sigma}_{\bm{\gamma}}$,
where $\text{logistic}(z)=\exp[z/(1+z)]$ and
$\eta_{\mathbf{x},\bm{\varepsilon}}$ is a linear combination of selected
functions of the covariates
$\bm{\psi}(\mathbf{x})=[\psi_{1}(\mathbf{x}),\ldots,\psi_{R}(\mathbf{x})]^{T}$.
Thus, $\\{\eta_{\mathbf{x},\bm{\varepsilon}}\colon\mathbf{x}\in\mathcal{X}\\}$
is a Gaussian process (GP) with mean
$\mu_{\mathbf{x},\bm{\varepsilon}}=\bm{\psi}(\mathbf{x})^{T}\bm{\mu}_{\bm{\gamma}}$
and covariance
$\texttt{Cov}(\eta_{\mathbf{x},\bm{\varepsilon}},\eta_{\mathbf{x}^{\prime},\bm{\varepsilon}})=\bm{\psi}(\mathbf{x})^{T}\bm{\Sigma}_{\bm{\gamma}}\bm{\psi}(\mathbf{x}^{\prime})$.
The remainder of the paper (with the exception of Section 2.4) assumes the
logit-normal prior (7).

###  2.4 Moments of random measures

Here we provide the mean, variance, and correlation properties of the random
measure $G_{\mathbf{x}}$ and compare these properties between the LT and BT SB
schemes. Proofs can be found in the Supplementary Materials. We note that
logit-normal prior (7) satisfies the conditions that Theorems 1 and 2 place on
the splitting variable distributions.

######  Theorem 1.

Suppose for any (vector of) covariates $\mathbf{x}$ that the atomic random
measure

$G_{x}=\sum_{k=1}^{K}W_{\mathbf{x},k}\,\delta_{\theta_{k}},\qquad
K\in\\{1,2,\ldots,\infty\\}$

on $\Theta$ is constructed by drawing each $\theta_{k}$ independently (of each
other and of $\mathbf{x}$) from some base measure $G_{0}$ on $\Theta$, and
drawing a weight vector $(W_{\mathbf{x},1},\ldots,W_{\mathbf{x},K})$ according
to some distribution (that may or may not depend on $\mathbf{x}$) on the
probability simplex $\mathbb{S}_{K}$.

For any measurable sets $A,A^{\prime}\in\mathcal{B}$ and covariates
$\mathbf{x},\mathbf{x}^{\prime}$, we have

$\displaystyle\mathbb{E}G_{\mathbf{x}}(A)$ $\displaystyle=G_{0}(A)$ (8)
$\displaystyle\texttt{Var}G_{\mathbf{x}}(A)$
$\displaystyle=\left[G_{0}(A)-G_{0}^{2}(A)\right]a_{\mathbf{x},\mathbf{x}}$
(9) $\displaystyle\texttt{Cov}(G_{\mathbf{x}}(A),G_{\mathbf{x}}(A^{\prime}))$
$\displaystyle=\left[G_{0}(A\cap
A^{\prime})-G_{0}(A)G_{0}(A^{\prime})\right]a_{\mathbf{x},\mathbf{x}}$ (10)
$\displaystyle\texttt{Cov}(G_{\mathbf{x}}(A),G_{\mathbf{x}^{\prime}}(A))$
$\displaystyle=\left[G_{0}(A)-G_{0}^{2}(A)\right]a_{\mathbf{x},\mathbf{x}^{\prime}}$
(11)

where
$a_{\mathbf{x},\mathbf{x}^{\prime}}=\sum_{k=1}^{K}\mathbb{E}W_{\mathbf{x},k}W_{\mathbf{x}^{\prime},k}$.
If also the quantities $G_{0}(A)$, $\left[1-G_{0}(A)\right]$,
$G_{0}(A^{\prime})$, and $\left[1-G_{0}(A^{\prime})\right]$ are all nonzero,
then

$\displaystyle\texttt{corr}(G_{\mathbf{x}}(A),G_{\mathbf{x}}(A^{\prime}))$
$\displaystyle=\frac{G_{0}(A\cap
A^{\prime})-G_{0}(A)G_{0}(A^{\prime})}{\sqrt{G_{0}(A)[1-G_{0}(A)]G_{0}(A^{\prime})[1-G_{0}(A^{\prime})]}}$
(12)
$\displaystyle\texttt{corr}(G_{\mathbf{x}}(A),G_{\mathbf{x}^{\prime}}(A))$
$\displaystyle=a_{\mathbf{x},\mathbf{x}^{\prime}}/\sqrt{a_{\mathbf{x},\mathbf{x}}a_{\mathbf{x}^{\prime},\mathbf{x}^{\prime}}}.$
(13)

Theorem 1 states that each of these various moments can be factorized into a
function of the base measure $G_{0}$ and a function of the quantity
$a_{\mathbf{x},\mathbf{x}^{\prime}}$. As functions of only $G_{0}$, the mean
(8) implies $G_{0}$ can be viewed as the mean of the random measure
$G_{\mathbf{x}}$ while the correlation (12) does not depend on tree depth. On
the other hand (and of greater interest), the cross-covariate correlation (13)
depends not on the base measure $G_{0}$ but rather on
$a_{\mathbf{x},\mathbf{x}^{\prime}}$, which Theorem 2 expresses as a function
of $\tau$ and the splitting variables’ mean and cross-covariate covariance. In
particular, this theorem states that $a_{\mathbf{x},\mathbf{x}^{\prime}}$ for
BT approaches zero as $K\rightarrow\infty$ whereas
$a_{\mathbf{x},\mathbf{x}^{\prime}}$ for LT approaches a positive limit, which
means LT induces a baseline cross-covariate covariance value (11) between
random measures that does not vanish as the number of weights approaches
infinity. Corollary 2.1 posits that these statements also apply to their
cross-covariate correlation counterparts (13) if also assumed are splitting
variables with mean $1/2$ and nonnegative cross-covariate correlation.

######  Theorem 2.

Suppose weights $(W_{\mathbf{x},\bm{\varepsilon}}:\bm{\varepsilon}\in
B(\tau))$ are constructed by SB according to tree topology $\tau$. Also assume
that at any $\mathbf{x}$ the splitting variables
$\\{V_{\mathbf{x},\bm{\varepsilon}}:\bm{\varepsilon}\in E^{*}\\}$ are
identically distributed. Let
$a_{\mathbf{x},\mathbf{x}^{\prime}}=\sum_{\bm{\varepsilon}\in
B(\tau)}\mathbb{E}W_{\mathbf{x},\bm{\varepsilon}}W_{\mathbf{x}^{\prime},\bm{\varepsilon}}$
and $K=|B(\tau)|\in\mathbb{N}$. If $\tau$ is a LT, then

$a_{\mathbf{x},\mathbf{x}^{\prime}}=\frac{\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}\\{1-[1-\mathbb{E}V_{\mathbf{x}}-\mathbb{E}V_{\mathbf{x}^{\prime}}+\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}]^{K}\\}}{\mathbb{E}V_{\mathbf{x}}+\mathbb{E}V_{\mathbf{x}^{\prime}}-\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}}$

which equals $\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}$ if $K=1$ and
approaches
$\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}(\mathbb{E}V_{\mathbf{x}}+\mathbb{E}V_{\mathbf{x}^{\prime}}-\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}})^{-1}$
as $K\rightarrow\infty$. If instead $\tau$ is a BT and $m=\log_{2}K$, then

$a_{\mathbf{x},\mathbf{x}^{\prime}}=\left(1-\mathbb{E}V_{\mathbf{x}}-\mathbb{E}V_{\mathbf{x}^{\prime}}+2\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}\right)^{m}$

which approaches zero as $m\rightarrow\infty$.

######  Corollary 2.1 (Lower bounds).

If $\mathbb{E}V_{\mathbf{x}}=1/2$ for any $\mathbf{x}$ and
$\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})\geq 0$ for any
$\mathbf{x},\mathbf{x}^{\prime}\in\mathcal{X}$, lower bounds for (13) are
$1/3$ for LT and $2^{-m}$ for BT. These lower bounds are strict if
$\texttt{Var}_{V}(\mathbf{x})>0$ for all $\mathbf{x}\in\mathcal{X}$.

Finally, the following theorem establishes the important property that the
random measure $G_{x}$ changes smoothly with respect to $\mathbf{x}$. This
property, expressed as the cross-covariate correlation (13) between random
measures approaching unity as $\mathbf{x}^{\prime}\rightarrow\mathbf{x}$,
requires a moment condition which is satisfied if the stochastic process
$\\{V_{\mathbf{x}}:\mathbf{x}\in\mathcal{X}\\}$ is second-order stationary.

######  Theorem 3 (Smoothness).

Given the assumptions in Theorem 2, if weights are either finite LT or BT and
if
$\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}\rightarrow\mathbb{E}V_{\mathbf{x}}^{2}$
as $\mathbf{x}^{\prime}\rightarrow\mathbf{x}$, then (13) $\rightarrow 1$.

###  2.5 Numerical illustration on prior cross-covariate correlation

Here we wish to explore the impact of the cross-covariate correlation between
splitting variables on the cross-covariate correlation between random
measures: given covariates $\mathbf{x}$ and $\mathbf{x}^{\prime}$, how does
$\texttt{corr}(V_{\mathbf{x}},V_{\mathbf{x}^{\prime}})$ affect
$\texttt{corr}\left(G_{\mathbf{x}}(A),G_{\mathbf{x}^{\prime}}(A)\right)$ for
any measurable set $A\in\mathcal{B}$? We provide insight into this question
through the following example.

######  Example 1.

We reduce the number of levers that influence
$\texttt{corr}(V_{\mathbf{x}},V_{\mathbf{x}^{\prime}})$ by making the
following simplifying assumptions:

$\bm{\mu}_{\bm{\gamma}}=\bm{0},\quad\quad\bm{\Sigma}_{\bm{\gamma}}=\text{diag}(\sigma^{2}_{1},\sigma^{2}_{2}),\quad\quad\bm{\psi}(\mathbf{x})=(1,0)^{T}\text{
and }\bm{\psi}(\mathbf{x}^{\prime})=(1,1)^{T}$

where $\sigma^{2}_{1}>0$. Though seemingly strict, these assumptions encompass
a large class of scenarios. The mean-zero assumption is reasonable if no prior
information is given. These assumptions also imply
$\texttt{corr}(V_{\mathbf{x}},V_{\mathbf{x}^{\prime}})=\texttt{corr}(\eta_{\mathbf{x}},\eta_{\mathbf{x}^{\prime}})=(1+\tilde{\sigma}^{2}_{2})^{-1/2}$,
which is a strictly decreasing function of
$\tilde{\sigma}^{2}_{2}\coloneqq\sigma^{2}_{2}/\sigma^{2}_{1}\geq 0$ whose
image is $(0,1]$. Thus, any positive value of
$\texttt{corr}(V_{\mathbf{x}},V_{\mathbf{x}^{\prime}})$ can be achieved by
using the appropriate $\tilde{\sigma}^{2}_{2}$ value if $\mathbf{x}\neq
x^{\prime}$. Similarly, these assumptions reduce
$\texttt{corr}\left(G_{\mathbf{x}}(A),G_{\mathbf{x}^{\prime}}(A)\right)$
(whose expression is provided by (13) and Theorem 2) to a function of just
$\sigma_{1}^{2}$, $\sigma_{2}^{2}$, $K$, and $\tau$.

Figure 3: Simulated values of (13) as a function of
$\tilde{\sigma}^{2}_{2}\coloneqq\sigma^{2}_{2}/\sigma^{2}_{1}$ for LT and BT
each with $K=64$ and various values of $\sigma^{2}_{1}$. Lower bounds are
shown as solid horizontal lines.

Figure 3 explores the behavior of the correlation
$\texttt{corr}\left(G_{\mathbf{x}}(A),G_{\mathbf{x}^{\prime}}(A)\right)$,
which by Corollary 2.1 has lower bounds of $1/3$ and $1/K$ for, respectively,
LT and BT. These lower bounds hold regardless of the degree of correlation
between splitting variables, which in this scenario is controlled by choice of
$\tilde{\sigma}^{2}_{2}$. Thus, the LT scheme always imposes a nontrivial
baseline correlation between random measures while the BT scheme can achieve
both large and small correlation values, which provides more flexibility in
setting prior correlation values.

###  2.6 Posterior sampling

This section details the posterior computation of the mixture model

$f(\cdot;G_{\mathbf{x}})=\int_{\Theta}h(\cdot;\theta)\mathop{}\\!\mathrm{d}G_{\mathbf{x}}(\theta).$
(14)

with a treeSB prior on the mixing measure $G_{\mathbf{x}}$. If $\tau$ is LT’s
tree structure with $K<\infty$ leaves, we call the resulting mixture a finite
LT mixture (LTM). We similarly define a finite BT mixture (BTM), where $K$
must be a power of $2$.

For posterior computation, we generalize Rigon and Durante (2021)’s Gibbs
sampler to admit any SB scheme. Their Gibbs step for $\bm{\gamma}$ relies on
Polson et al. (2013)’s Pólya–Gamma (PG) data augmentation technique, which
allows efficient posterior sampling of a Bayesian logistic regression. For
this technique, Polson et al. (2013) carefully construct the PG family of
distributions to allow conditionally conjugate updating for the coefficient
parameter and provide a fast, exact way to simulate PG random variables. Let
$Z_{i}$ be the number of successes, $n_{i}$ the number of trials, and
$\mathbf{x}_{i}$ be a $p$-tuple of regressors for observation $i\in[n]$. Also
suppose
$Z_{i}\sim\text{Binom}(n_{i},[1+\exp(-\mathbf{x}_{i}^{T}\bm{\gamma})]^{-1})$
and $\bm{\gamma}\sim N(\mathbf{b},\mathbf{B})$. They show that a posterior
sampler for $\bm{\gamma}$ is obtained by iterating the two steps

$\displaystyle\left(\omega_{i}\mid\bm{\gamma}\right)$
$\displaystyle\sim\text{PG}(n_{i},\mathbf{x}_{i}^{T}\bm{\gamma})$
$\displaystyle\left(\bm{\gamma}\mid\omega,Z_{1},\ldots,Z_{n}\right)$
$\displaystyle\sim\text{N}(\mathbf{m}_{\omega},\mathbf{V}_{\omega})$

where

$\displaystyle\mathbf{m}_{\omega}$
$\displaystyle=\mathbf{V}_{\omega}\left[\mathbf{X}^{T}\bm{\kappa}+\mathbf{B}^{-1}\mathbf{b}\right]\quad\text{with
}\bm{\kappa}^{T}=(Z_{1}-n_{1}/2,\ldots,Z_{n}-n_{n}/2),$
$\displaystyle\mathbf{V}_{\omega}$
$\displaystyle=\left[\mathbf{X}^{T}\bm{\Omega}\mathbf{X}+\mathbf{B}^{-1}\right]^{-1}\quad\text{with
}\bm{\Omega}=\text{diag}(\omega_{1},\ldots,\omega_{n}),$

and $\mathbf{X}$ is the matrix with row entries $\mathbf{x}_{i}^{T}$.

Result: Update $\bm{\gamma}_{\bm{\varepsilon}}$ for each internal node
$\bm{\varepsilon}$ in binary tree $\tau$.

for  _each internal node $\bm{\varepsilon}$ in $\tau$_ do

Let $S_{\bm{\varepsilon}}$ be the set of indices $i\in[n]$ where $C_{\tau}(i)$
is a descendant of node $\bm{\varepsilon}$ ;

for  _every $i\in S_{\bm{\varepsilon}}$_ do

Sample Pólya-gamma data $[\omega_{i\bm{\varepsilon}}\mid\cdots]\sim
PG(1,\,\mathbf{x}_{i}^{T}\bm{\gamma}_{\bm{\varepsilon}})$ ;

end for

Let $L_{\bm{\varepsilon}}\subset S_{\bm{\varepsilon}}$ be the set of indices
corresponding to “left descendants” of node $\bm{\varepsilon}$ ;

Update $\bm{\gamma}_{\bm{\varepsilon}}$ by drawing from full conditional
$[\bm{\gamma}_{\bm{\varepsilon}}\mid\cdots]\sim
N_{R}(\bm{\mu}_{\bm{\gamma}_{\bm{\varepsilon}}},\Sigma_{\bm{\gamma}_{\bm{\varepsilon}}})$
where  $\displaystyle\bm{\mu}_{\bm{\gamma}_{\bm{\varepsilon}}}$
$\displaystyle=\bm{\Sigma}_{\bm{\gamma}_{\bm{\varepsilon}}}\left[\mathbf{X}_{\bm{\varepsilon}}^{T}\bm{\kappa}_{\bm{\varepsilon}}+\bm{\Sigma}_{\bm{\gamma}}^{-1}\bm{\mu}_{\bm{\gamma}}\right]\quad\text{with
}\bm{\kappa}_{\bm{\varepsilon}}=(1_{i\in L_{\bm{\varepsilon}}}-0.5)_{i\in
S_{\bm{\varepsilon}}},$
$\displaystyle\bm{\Sigma}_{\bm{\gamma}_{\bm{\varepsilon}}}$
$\displaystyle=\left[\mathbf{X}_{\bm{\varepsilon}}^{T}\bm{\Omega}_{\bm{\varepsilon}}\mathbf{X}_{\bm{\varepsilon}}+\bm{\Sigma}_{\bm{\gamma}}^{-1}\right]^{-1}\quad\text{with
}\bm{\Omega}_{\bm{\varepsilon}}=\text{diag}(\omega_{1{\bm{\varepsilon}}},\ldots,\omega_{|S_{\bm{\varepsilon}}|\bm{\varepsilon}}),$
and $\mathbf{X}_{\bm{\varepsilon}}$ is the $|S_{\bm{\varepsilon}}|\times R$
matrix with row entries $\mathbf{x}_{i}^{T}$ for only those $i\in
S_{\bm{\varepsilon}}$.

end for

Algorithm 1 Gibbs step to update logit-normal coefficient under any binary
tree.

Our generalization of this Gibbs step is stated explicitly in Algorithm 1.
Given a posterior draw, for all $i\in[n]\coloneqq\\{1,\ldots,n\\}$, let
$C_{\tau}(i)$ be $\tau$’s leaf node assigned to observation $i$. For each
internal node $\bm{\varepsilon}$, the update for the coefficient
$\bm{\gamma}_{\bm{\varepsilon}}$ relies on a Bayesian logistic regression with
responses

$Z_{i\bm{\varepsilon}}=\begin{cases}1\quad&\text{if leaf node
}C_{\tau}(i)\text{ is a ``left descendant'' of node $\bm{\varepsilon}$}\\\
0\quad&\text{if leaf node }C_{\tau}(i)\text{ is a ``right descendant'' of node
$\bm{\varepsilon}$}\\\ \end{cases}$

for all $i$ corresponding to a descendant node of $\bm{\varepsilon}$. Here, a
“left descendant” of $\bm{\varepsilon}$ is either the left child or a
descendant of the left child of $\bm{\varepsilon}$, and “right descendant” is
defined similarly. For example, in LT the left child of any internal node is a
leaf (see Figure 2) whereas for any internal node in BT the number of left
descendant nodes equals the number of right descendant nodes. This formulation
easily applies to any finite treeSB scheme.

Next we compare between LT and BT the computational cost of this Gibbs step
for some arbitrarily fixed MCMC iteration. In this case, each tree has the
same number of internal nodes and hence their corresponding mixture models
perform the same number of regressions. From this viewpoint, the computational
cost should be the same between the two SB schemes. However, we find from
experience that for the data sets in Sections 3 and 4, training a BTM model
takes less time than training its LTM counterpart, which is supported by the
theoretical discussion in Section 6.5 in the Supplementary Materials.

##  3 Simulation study

Figure 4: Data generated by 20 skew-normal distributions. The boxed numbers
are the number of observations generated from each distribution. The triangles
indicate the observations whose covariate assignment differs between the two
data sets.

This example explores how tree structure affects MCMC fitting. We create two
data sets $\mathscr{D}^{indep}$ and $\mathscr{D}^{dep}$, whose cluster
proportions respectively do not and do depend on the observations’ covariate
values. The two data sets share the same $n=8000$ observations shown in Figure
4. Each of the 20 skew-normal (SN) distributions generates some number of
observations; these cluster counts are listed as the $20$-tuple

$\mathbf{n}=8\cdot(200,170,130,100,80,70,50,30,28,22,20,18,16,14,12,10,9,8,7,6).$

For each $i\in\mathbb{N}_{20}$, each data set assigns a covariate
$\mathbf{x}_{i}$ in the space $\mathcal{X}\coloneqq\\{1\\}\times\\{0,1\\}^{3}$
to observation $i$. Let $\mathbf{n}_{\mathbf{x}}^{indep}$ and
$\mathbf{n}_{\mathbf{x}}^{dep}$ be the cluster counts for all observations
assigned the covariates $\mathbf{x}$ by, respectively, $\mathscr{D}^{indep}$
and $\mathscr{D}^{dep}$. The covariate-specific cluster counts
$\mathbf{n}_{\mathbf{x}}^{indep}$ are the same for all covariates, i.e.
$\mathbf{n}_{\mathbf{x}}^{indep}=\mathbf{n}/8$ for all
$\mathbf{x}\in\mathcal{X}$. Thus, we dub the data set $\mathscr{D}^{indep}$
“covariate-independent.” In contrast, the “covariate-dependent” data set
$\mathscr{D}^{dep}$ has six clusters whose covariate-specific counts do depend
on $\mathbf{x}$:

$\displaystyle\mathbf{n}_{\mathbf{x}}^{dep}=(\mathbf{n}/8)+20(2x_{1}-1)(\mathbf{e}_{8}-\mathbf{e}_{5})+20(2x_{2}-1)(\mathbf{e}_{10}-\mathbf{e}_{6})+20(2x_{3}-1)(\mathbf{e}_{9}-\mathbf{e}_{7}),$

where $\mathbf{x}=(1,x_{1},x_{2},x_{3})$ and each $\mathbf{e}_{j}$ is the
$j$th standard basis vector.

We fit a dependent LTM model and a dependent BTM model to each data set using
the Gibbs sampler in Section 2.6. Each model uses $K=32$, prior mean
$\bm{\mu}_{\gamma}=\bm{0}_{p}$, and prior covariance
$\bm{\Sigma}_{\gamma}=10\bm{I}_{p}$. Thus, four MCMC chains are run. Each
chain burns in $100,000$ steps before sampling every $100$ steps to ultimately
keep $10,000$ posterior draws.

Figure 5 shows for each chain the Jaccard distance
$[1-J(\mathbf{c}^{(i)},\mathbf{c}_{0})]\in[0,1]$ between the chain’s
clustering $\mathbf{c}^{(i)}$ at the $i$th posterior draw (for all
$i\in\mathbb{N}_{10000}$) and the true clustering $\mathbf{c}_{0}$ (Jaccard,
1912). These distances provide a sense of a chain’s clustering performance and
its mixing behavior. For each data set, LT and BT have roughly the same median
Jaccard distance, but the mean Jaccard distance is slightly larger for LT. The
LT chains also seem to get stuck in local modes for longer than the BT chains
do.

Figure 5: Jaccard distances at each posterior draw.

What accounts for this difference in clustering behavior? Figure 6 shows
selected mixture-weight inference for $\mathbf{x}=1010$. The results are
similar for the other covariate values; “full” versions of these plots are in
the Supplementary Materials. In Figure 6(a)’s $(LT,\mathscr{D}^{indep})$
panel, none of the credible intervals (CIs) are centered around the second
largest weight value of $0.17$ (this is more easily seen in Figure 6(b)),
which might foster the belief that this LTM model fails to capture this
weight, but upon closer inspection, the heights of the eight curiously wide
$95\%$ CIs are roughly that same value. This observation together with Figure
6(c) imply that the eight wide LT CIs are a result of label switching in the
MCMC chain. Similarly, Figure 6(d) implies that the wide CIs in Figure 6(a)’s
$(BT,\mathscr{D}^{dep})$ panel also happen because of label switching. Indeed,
Figure 7 shows that these wide CIs disappear and the data’s weights are
captured if we “resolve” the label switching by performing the following post-
hoc sorting of weights: for each posterior draw $j$ we rearrange (in
decreasing order) the elements of the chain’s corresponding weight vector
$\mathbf{w}_{1010}^{(j)}$, then for each newly rearranged weight index
$\tilde{k}$ we compute $95\%$ CIs for the values in the set
$\\{w_{1010,\tilde{k}}^{(j)}:j\in\mathbb{N}_{10000}\\}$.

(a) Pointwise $95\%$ CIs of weights inferred by models trained on the two data
sets.

(b) Weight indices are sorted by median weight. Points show the data’s true
weights.

(c) Weights inferred by the LTM trained on $\mathscr{D}^{indep}$ for indices
corresponding to wide CIs.

(d) Weights inferred by the BTM trained on $\mathscr{D}^{dep}$ for indices
corresponding to wide CIs.

Figure 6: Selected mixture-weight inference for $\mathbf{x}=1010$. Figure 7:
Inferred weights that are sorted before taking pointwise $95\%$ CIs. Points
indicate the data’s true weight values.

Although this post-hoc sorting seems to resolve the wide CIs for this example,
the procedure requires several assumptions to work and hence cannot be applied
generally. Until a general solution is created, it seems prudent to “nip the
problem in the bud” by mitigating the impact of label switching on inference.
But why does label switching seem to affect LT’s inference more than BT’s? We
conjecture that the behavior difference stems from two tree-related
mechanisms.

First, consider how much friction there is to label switch between two leaf
nodes (for either tree structure). The amount of friction seems to increase
with the number of splitting variables affected by the proposed switch. If the
two leaves are graphically far apart, the proposed switch would affect many
splitting variables and hence be unlikely to occur. Under this reasoning,
switching is most likely to occur between two leaves that are (graphically)
close to each other. In a BT, this means two leaves that share the same
parent. In a LT, this means two leaves separated by at most one tree level.
For conceptual simplicity, we consider only such pairs and call these leaves
adjacent to each other. Note that LT has $K$ adjacent pairs whereas BT has
$K/2$. Thus, LT has more adjacent pairs from which switching can occur.

Second, a label switch between two adjacent LT leaves can create a “chain
reaction” of label switches either up or down the entire tree. If a switch
occurs between adjacent LT leaf-pair $(\varepsilon 0,\varepsilon 10)$ for some
internal node $\varepsilon$, the adjacent pair $(\varepsilon 10,\varepsilon
110)$ becomes “newly eligible” for switching, and if that switch eventually
occurs, the adjacent pair $(\varepsilon 110,\varepsilon 1110)$ becomes “newly
eligible” for switching, and so on, which creates the possibility of a chain
reaction in the direction of increasing tree level (though a chain reaction
can just as easily occur in the other tree direction). This might explain the
behavior in Figure 6(c), where the $0.17$ mass moves almost sequentially
through eight nodes. In contrast, because any BT leaf is adjacent to exactly
one other leaf and because adjacency is a symmetric binary relation, any
switch between two adjacent BT leaves will likely stay localized i.e. such a
switch will likely not create any “newly eligible” pairs. This BT conjecture
is supported by the chain behavior in Figure 6(d).

These two proposed mechanisms seem to explain the observed MCMC chain behavior
difference between LT and BT. Overall, our aforementioned inference
observations indicate that if we consider the parameter space $\Theta$ and
posterior local modes corresponding to label permutations, the regions between
these local modes appear to be deeper valleys for the BT posterior than they
are for the LT posterior, making it more difficult to jump between different
labeling for BT than for LT.

##  4 Case study: flow cytometry data

The goal of this analysis is to quantify the impact of a patient’s age, race,
and sex on their proportions of T-cell types. The FCM data comes from 127
healthy plasma donors used to establish normative ranges for immune cell
subsets enumerated by flow cytometry, as described in Yi et al. (2019). Donors
were selected to be equally represented across the following age cohorts:
18–29 years, 30–39, 40–49, and 50–66. Of the 127 donors, 51 are Caucasian, 69
are African American, and 7 are classified as Other. Regarding sex, 65 are
male and 62 are female. We focus our analysis on the T cell “panel,” a
collection of custom antidotes designed to identify maturational and
functional subsets of T cells developed by the Human ImmunoPhenotyping
Consortium (HIPC) (Finak et al., 2016). Although the panel is designed to make
it possible to identify multiple predefined T cell subsets, the sample is of
all peripheral blood mononuclear cells, and clusters may include non-T cell
subsets.

We subset the available data in the following way. Of the 127 donors, we use
the 120 whose race is listed as either Caucasian or African-American and omit
the remaining 7 subjects to avoid conflating effects attributed to the Other
race label with subject-specific effects. Table 1 contains contingency tables
for the subjects used in our analysis. Because these 120 subjects produce too
many viable cells for the model to train on in a reasonable amount of time, we
subset the data in a way that uses all 120 subjects while representing each
covariate combination $\mathbf{x}$ by the same number $n_{\mathbf{x}}$ of
viable cells. First, we partition the subjects according to their age group,
sex, and race. For each of the $4\times 2\times 2=16$ covariate combinations
$\mathbf{x}$, we then subset $n_{\mathbf{x}}/|\mathscr{D}_{\mathbf{x}}|$
viable cells from each subject with covariates $\mathbf{x}$, where
$|\mathscr{D}_{\mathbf{x}}|$ is the number of subjects that share covariates
$\mathbf{x}$ – see Table 1 for all values of $|\mathscr{D}_{\mathbf{x}}|$ –
and $n_{\mathbf{x}}\coloneqq 12600=2^{3}\times 3^{2}\times 5^{2}\times 7$ is
divisible by all possible values of $|\mathscr{D}_{\mathbf{x}}|$. Thus, the
training data consists of
$n=\sum_{\mathbf{x}}n_{\mathbf{x}}=16n_{\mathbf{x}}=201600$ viable cells.

(Caucasian) | Male | Female | Total
---|---|---|---
$18-29$ | 4 | 8 | 12
$30-39$ | 7 | 8 | 15
$40-49$ | 5 | 6 | 11
$50-65$ | 8 | 5 | 13
Total | 24 | 27 | 51

(African-American) | Male | Female | Total
---|---|---|---
$18-29$ | 10 | 6 | 16
$30-39$ | 9 | 9 | 18
$40-49$ | 10 | 7 | 17
$50-65$ | 9 | 9 | 18
Total | 38 | 31 | 69

Table 1: Contingency tables for subjects used to analyze data from the GRIFOLS
study.

To the data we fit a LTM model and a BTM model each with hyperparameter values
of $K=64$, prior mean $\bm{\mu}_{\gamma}=\bm{0}_{p}$, and prior covariance
$\bm{\Sigma}_{\gamma}=100\bm{I}_{p}$ where $p=7$. Each chain burns in $10,000$
steps before sampling every $50$ steps to ultimately keep $1,000$ posterior
draws. We also incorporate two statistical methodologies to address features
common to FCM data. First, for both models we use cross-sample calibration to
account for subject-data being collected in different batches (Gorsky et al.,
2020). Second, it is well known that FCM clusters do not match the shape of
skew-normal distributions. To address this kernel misspecification, we
incorporate coarsening (Miller and Dunson, 2018) into our mixture models. In
particular, we use the coarsening-parameter value of $\zeta=0.1$, which is
strong enough to combine similar-looking clusters but weak enough to avoid
creating many multimodal clusters.

We use the following strategy to translate posterior covariate effects of
stick-breaking proportions onto those of the weights. For any given MCMC chain
and each of its posterior draws, each internal node $\bm{\varepsilon}$ of the
tree corresponds to a vector of regression coefficients
$\bm{\gamma}_{\bm{\varepsilon}}=(\gamma_{0},\gamma_{age1},\gamma_{age2},\gamma_{age3},\gamma_{race},\gamma_{sex})^{T}$
and each leaf $\bm{\delta}$ of the tree corresponds to a weight
$W_{\mathbf{x},\bm{\delta}}$ created according to (7) and (2) for each vector
of binary covariates $\mathbf{x}\in\mathcal{X}_{GRIFOLS}$ where
$\mathcal{X}_{GRIFOLS}=\\{(1,x_{age1},x_{age2},x_{age3},x_{race},x_{sex})\in\\{0,1\\}^{6}:x_{age1}\geq
x_{age2}\geq x_{age3}\\}$. As an example, we examine the posterior effect of
the race covariate on the weights by computing differences
$W_{(1,x_{age1},x_{age2},x_{age3},1,x_{sex}),\bm{\delta}}-W_{(1,x_{age1},x_{age2},x_{age3},0,x_{sex}),\bm{\delta}}$
for various shared combinations of covariate values
$x_{age1},x_{age2},x_{age3},x_{sex}$ with $x_{age1}\geq x_{age2}\geq x_{age3}$
at each leaf $\bm{\delta}$. The effect of the other covariates on the weights
are computed in the same manner.

###  4.1 Trustworthiness of LT’s credible intervals

As mentioned in Section 1, Figure 8(a) showcases a disturbing dependence of a
LTM model weight’s CI length on the order in which the weight was created
during stick breaking. Indeed, cluster 01’s CI lengths are at least $5\times$
many other CI lengths with larger weight indices. It now becomes difficult to
trust how a LTM model quantifies a weight’s uncertainty in light of its strong
artificial dependence on whichever leaf node the weight happens to belong to.
This is particularly true if the model on the weights is misspecified to some
degree (which is likely the case in any analysis of real data), in which case
a trustworthy inference should retain a nontrivial amount of variance in the
weights regardless of how large $n$ is. In contrast, each weight constructed
by BT SB is a product of the same number ($\log_{2}K=6$) of stick breaks and
thus does not share this artificial dependence. This assertion is supported by
Figure 8(b) whose CIs show no obvious relationship to weight index and are
roughly as wide as the CIs of cluster 05 and 07 from the LTM model.

(a) LTM model.

(b) BTM model.

Figure 8: Pointwise $95\%$ CIs of various weight differences inferred by a
mixture model.

###  4.2 BTM model inference

In light of Section 4.1, we provide the results from only the BTM model.
Figure 12 shows the salient clusters inferred by the BTM model. Figures 13,
14, 15, and 16 show weight differences for various pairs of covariate-
combinations that differ by exactly one covariate value. For space
considerations, we relegate them to the Supplementary Materials, but here we
present how a practitioner might interpret these results.

##### Changes with age

Two cell subsets, corresponding to clusters 20 and 25, are over-represented in
older age groups (30-39, 40-49, 50-65) compared with the youngest cohort
(18-29) across both sexes and races. Cluster 20 has the phenotype of a central
memory CD4+ T cell (CCR7+/CD45RA-), which arises after antigen encounter and
hence is expected to increase over time. Cluster 25 has the phenotype of a
pro-inflammatory monocyte/macrophage and suggests aging is associated with
increased inflammation. This is not unexpected, as many diseases of aging are
associated with chronic and excessive inflammation.

The cell subset corresponding to cluster 6 is over-represented in the second-
youngest age group (30-39) compared with the youngest cohort (18-29) across
both sexes and races. Cluster 6 has the phenotype of a naïve CD4+ T cell, and
as naïve T cells are traditionally considered antigen-inexperienced T cells,
this is counter-intuitive as thymic generation of new T cells drops
precipitously after childhood. However, there is evidence that in adult
humans, naïve T cells are maintained by cell division (den Braber et al.,
2012), and hence while the diversity of the naïve T cell repertoire might be
decreasing, the numbers may not.

Finally, the cell subset corresponding to cluster 24 is over-represented in
older age groups (30-39, 40-49, 50-65) compared with the youngest cohort
(18-29) across Caucasian males and females. Cluster 24 has the phenotype of a
transitional cell moving from naïve to memory.

##### Race-related differences

All three mentioned CD4+ T cell types (naïve, central memory, and
transitional) show similar race-related differences in the 30-39 age cohort.
Namely, they are predicted to be less abundant for African Americans relative
to Caucasians. The same can also be said for both naïve and transitional CD4+
T cells in the oldest (50-65) age cohort, but not for the central memory CD4+
T cells.

##### Sex-related differences

Only subtle sex differences in immune cell cluster are observed. In the oldest
(50-65) age cohort, transitional cells are predicted to be moderately less
abundant in males relative to females, whereas in the 30-39 age cohort,
central memory cells are predicted to be less abundant in males relative to
females.

##  5 Conclusion

The focus of this paper has been to propose a new stick-breaking procedure
that circumvents several undesirable behaviors of LT stick-breaking models.
First, Section 2.5 shows that the LT induces a strong prior cross-covariate
correlation between random measures whereas the BT allows more flexibility.
Second, Section 3 details how the impact of label switching on posterior
inference is stronger for LTM models than it is for BTM models, as well as
conjectures as to the mechanisms that cause this behavior. Third, Section 4
illustrates the presence of artificial variance decay in weights constructed
from LT SB as well as its absence in weights constructed from BT SB. We also
provide a Gibbs sampler for the BTM model that either matches or beats the
efficiency of that for the LTM model.

The model specified on the splitting variables can be tailored according to
the nature of the given data. The splitting variables in (7) depend on
covariates via a fixed-effects model, but a practitioner may wish to also
allow random effects (e.g. corresponding to subject-specific idiosyncrasies in
a FCM data set) to reduce posterior uncertainty and hence improve covariate-
effect inference on the mixture weights. To achieve posterior computation with
a mixed-effects model on the splitting variables, Wang and Roy (2018)’s two-
block Gibbs sampler can be adapted to work for a BTM model.

## Acknowledgements

We gratefully acknowledge Scott White for assistance in processing flow
cytometry data.

This research was supported by the Translating Duke Health Initiative
(Immunology) [AH, CC, LM], the Duke University Center for AIDS Research
(CFAR), an NIH funded program (5P30 AI064518) [SW, CC], the National Institute
of General Medical Sciences grant R01-GM135440 [AH, LM], and the National
Science Foundation grant DMS-2013930 [LM].

## References

  * Albert and Chib (1993) Albert, J. H. and S. Chib (1993).  Bayesian analysis of binary and polychotomous response data.  Journal of the American statistical Association 88(422), 669–679.
  * Cipolli and Hanson (2017) Cipolli, W. and T. Hanson (2017).  Computationally tractable approximate and smoothed polya trees.  Statistics and Computing 27(1), 39–51.
  * Cron and West (2011) Cron, A. J. and M. West (2011).  Efficient classification-based relabeling in mixture models.  The American Statistician 65(1), 16–20.
  * den Braber et al. (2012) den Braber, I., T. Mugwagwa, N. Vrisekoop, L. Westera, R. Mögling, A. B. de Boer, N. Willems, E. H. Schrijver, G. Spierenburg, K. Gaiser, et al. (2012).  Maintenance of peripheral naive t cells is sustained by thymus output in mice but not humans.  Immunity 36(2), 288–297.
  * Dubey et al. (2014) Dubey, K. A., Q. Ho, S. A. Williamson, and E. P. Xing (2014).  Dependent nonparametric trees for dynamic hierarchical clustering.  Advances in Neural Information Processing Systems 27.
  * Fabius (1964) Fabius, J. (1964).  Asymptotic behavior of bayes’ estimates.  The Annals of Mathematical Statistics 34(4), 846–856.
  * Finak et al. (2016) Finak, G., M. Langweiler, M. Jaimes, M. Malek, J. Taghiyar, Y. Korin, K. Raddassi, L. Devine, G. Obermoser, M. L. Pekalski, et al. (2016).  Standardizing flow cytometry immunophenotyping analysis from the human immunophenotyping consortium.  Scientific reports 6(1), 1–11.
  * Freedman (1963) Freedman, D. A. (1963).  On the asymptotic behavior of bayes’ estimates in the discrete case.  The Annals of Mathematical Statistics 34(4), 1386–1403.
  * Ghahramani et al. (2010) Ghahramani, Z., M. Jordan, and R. P. Adams (2010).  Tree-structured stick breaking for hierarchical data.  Advances in neural information processing systems 23.
  * Gorsky et al. (2020) Gorsky, S., C. Chan, and L. Ma (2020).  Coarsened mixtures of hierarchical skew normal kernels for flow cytometry analyses.  arXiv preprint arXiv:2001.06451.
  * Ishwaran and James (2001) Ishwaran, H. and L. F. James (2001).  Gibbs sampling methods for stick-breaking priors.  Journal of the American Statistical Association 96(453), 161–173.
  * Ishwaran and James (2002) Ishwaran, H. and L. F. James (2002).  Approximate dirichlet process computing in finite normal mixtures: smoothing and prior information.  Journal of Computational and Graphical statistics 11(3), 508–532.
  * Jaccard (1912) Jaccard, P. (1912).  The distribution of the flora in the alpine zone. 1.  New phytologist 11(2), 37–50.
  * Jara and Hanson (2011) Jara, A. and T. E. Hanson (2011).  A class of mixtures of dependent tail-free processes.  Biometrika 98(3), 553–566.
  * Linderman et al. (2015) Linderman, S., M. J. Johnson, and R. P. Adams (2015).  Dependent multinomial models made easy: Stick-breaking with the pólya-gamma augmentation.  Advances in Neural Information Processing Systems 28.
  * MacEachern (2000) MacEachern, S. N. (2000).  Dependent dirichlet processes.  Unpublished manuscript, Department of Statistics, The Ohio State University, 1–40.
  * Miller and Dunson (2018) Miller, J. W. and D. B. Dunson (2018).  Robust bayesian inference via coarsening.  Journal of the American Statistical Association 114(527), 1113–1125.
  * Nassar et al. (2018) Nassar, J., S. W. Linderman, M. Bugallo, and I. M. Park (2018).  Tree-structured recurrent switching linear dynamical systems for multi-scale modeling.  arXiv preprint arXiv:1811.12386.
  * Polson et al. (2013) Polson, N. G., J. G. Scott, and J. Windle (2013).  Bayesian inference for logistic models using pólya–gamma latent variables.  Journal of the American statistical Association 108(504), 1339–1349.
  * Rigon and Durante (2021) Rigon, T. and D. Durante (2021).  Tractable bayesian density regression via logit stick-breaking priors.  Journal of Statistical Planning and Inference 211, 131–142.
  * Rodriguez and Dunson (2011) Rodriguez, A. and D. B. Dunson (2011).  Nonparametric bayesian models through probit stick-breaking processes.  Bayesian analysis (Online) 6(1).
  * Sethuraman (1994) Sethuraman, J. (1994).  A constructive definition of dirichlet priors.  Statistica sinica 4(2), 639–650.
  * Stefanucci and Canale (2021) Stefanucci, M. and A. Canale (2021).  Multiscale stick-breaking mixture models.  Statistics and Computing 31(2), 1–13.
  * Wang and Roy (2018) Wang, X. and V. Roy (2018).  Analysis of the pólya-gamma block gibbs sampler for bayesian logistic linear mixed models.  Statistics & Probability Letters 137, 251–256.
  * Yi et al. (2019) Yi, J. S., M. Rosa-Bray, J. Staats, P. Zakroysky, C. Chan, M. A. Russo, C. Dumbauld, S. White, T. Gierman, K. J. Weinhold, et al. (2019).  Establishment of normative ranges of the healthy human immune system with comprehensive polychromatic flow cytometry profiling.  PLoS One 14(12), e0225512.

SUPPLEMENTARY MATERIAL

Proofs:

Proofs of Theorems 1, 2, 3, and Corollary 2.1.

Discussion:

Discussion on computational cost of Gibbs step from Algorithm 1.

Plots:

Plots of results from Sections 3 and 4.

##  6 Proofs

###  6.1 Proof of Theorem 1

###### Proof of Theorem 1.

This proof follows closely from Appendix 2 of Rodriguez and Dunson (2011). To
prove the mean result (8), we have

$\displaystyle\mathbb{E}G_{\mathbf{x}}(A)$
$\displaystyle=\mathbb{E}\left\\{\sum_{k=1}^{K}W_{\mathbf{x},k}\delta_{\theta_{k}}(A)\right\\}=\sum_{k=1}^{K}\left[\mathbb{E}W_{\mathbf{x},k}\right]Pr(\theta_{1}\in
A)=G_{0}(A).$

To prove the remaining results, let $X=G_{\mathbf{x}}(A)$ and
$Y=G_{\mathbf{x}^{\prime}}(A^{\prime})$ and note that

$\displaystyle\mathbb{E}XY$
$\displaystyle=\mathbb{E}\left[\left\\{\sum_{k=1}^{K}W_{\mathbf{x},k}\delta_{\theta_{k}}(A)\right\\}\left\\{\sum_{k=1}^{K}W_{\mathbf{x}^{\prime},k}\delta_{\theta_{k}}(A^{\prime})\right\\}\right]$
$\displaystyle=\sum_{k=1}^{K}\mathbb{E}\left[W_{\mathbf{x},k}W_{\mathbf{x}^{\prime},k}\delta_{\theta_{k}}(A)\delta_{\theta_{k}}(A^{\prime})\right]+\sum_{k=1}^{K}\sum_{k^{\prime}\neq
k}\mathbb{E}\left[W_{\mathbf{x},k}W_{\mathbf{x}^{\prime},k^{\prime}}\delta_{\theta_{k}}(A)\delta_{\theta_{k^{\prime}}}(A^{\prime})\right]$
$\displaystyle=G_{0}(A\cap
A^{\prime})a_{\mathbf{x},\mathbf{x}^{\prime}}+G_{0}(A)G_{0}(A^{\prime})b_{\mathbf{x},\mathbf{x}^{\prime}}$

where
$a_{\mathbf{x},\mathbf{x}^{\prime}}=\sum_{k=1}^{K}\mathbb{E}W_{\mathbf{x},k}W_{\mathbf{x}^{\prime},k}$
and $b_{\mathbf{x},\mathbf{x}^{\prime}}=\sum_{k=1}^{K}\sum_{k^{\prime}\neq
k}\mathbb{E}W_{\mathbf{x},k}W_{\mathbf{x}^{\prime},k^{\prime}}$. Because
$a_{\mathbf{x},\mathbf{x}^{\prime}}+b_{\mathbf{x},\mathbf{x}^{\prime}}=1$, we
have
$\mathbb{E}X\mathbb{E}Y=G_{0}(A)G_{0}(A^{\prime})(a_{\mathbf{x},\mathbf{x}^{\prime}}+b_{\mathbf{x},\mathbf{x}^{\prime}})$,
and thus

$\displaystyle\texttt{Cov}(X,Y)=\mathbb{E}XY-\mathbb{E}X\mathbb{E}Y$
$\displaystyle=\left[G_{0}(A\cap
A^{\prime})-G_{0}(A)G_{0}(A^{\prime})\right]a_{\mathbf{x},\mathbf{x}^{\prime}}.$

The variance and covariances (9), (10), and (11) result from appropriately
setting $x=x^{\prime}$ or $A=A^{\prime}$ (or both). The correlations (12) and
(13) result from the definition
$\texttt{corr}(X,Y)=(\mathbb{E}XY-\mathbb{E}X\mathbb{E}Y)/\sqrt{\texttt{Var}X\texttt{Var}Y}$,
which requires the denominator to be nonzero. ∎

###  6.2 Proof of Theorem 2

###### Proof of Theorem 2.

The LT result comes from Appendix 2 of Rodriguez and Dunson (2011). For the BT
result, let $m=|B(\tau)|$. For any $\bm{\varepsilon}\in B(\tau)$,

$\displaystyle\mathbb{E}W_{\mathbf{x},\bm{\varepsilon}}W_{\mathbf{x}^{\prime},\bm{\varepsilon}}$
$\displaystyle=\prod_{l=1}^{m}\mathbb{E}\left[V_{\mathbf{x},\varepsilon_{1}\cdots\varepsilon_{l-1}}V_{\mathbf{x}^{\prime},\varepsilon_{1}\cdots\varepsilon_{l-1}}\right]^{1-\varepsilon_{l}}\mathbb{E}\left[\left(1-V_{\mathbf{x},\varepsilon_{1}\cdots\varepsilon_{l-1}}\right)\left(1-V_{\mathbf{x}^{\prime},\varepsilon_{1}\cdots\varepsilon_{l-1}}\right)\right]^{\varepsilon_{l}}.$

Because the splitting variables are identically distributed, the quantity
$\mathbb{E}W_{\mathbf{x},\bm{\varepsilon}}W_{\mathbf{x}^{\prime},\bm{\varepsilon}}$
becomes the product
$\prod_{l=1}^{m}c_{\mathbf{x},\mathbf{x}^{\prime}}^{1-\varepsilon_{l}}d_{\mathbf{x},\mathbf{x}^{\prime}}^{\varepsilon_{l}}=c_{\mathbf{x},\mathbf{x}^{\prime}}^{m-\sum_{l=1}^{m}\varepsilon_{l}}d_{\mathbf{x},\mathbf{x}^{\prime}}^{\sum_{l=1}^{m}\varepsilon_{l}}$
where
$c_{\mathbf{x},\mathbf{x}^{\prime}}=\mathbb{E}\left[V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}\right]$
and
$d_{\mathbf{x},\mathbf{x}^{\prime}}=\mathbb{E}\left[\left(1-V_{\mathbf{x}}\right)\left(1-V_{\mathbf{x}^{\prime}}\right)\right]$.
For any $k\in\\{0,1,\ldots,m\\}$, the set $B(\tau)$ has exactly
$\binom{m}{k}$-many leaves $\varepsilon_{1}\cdots\varepsilon_{m}$ with the
property $\sum_{l=1}^{m}\varepsilon_{l}=k$. Thus, the desired sum
$\sum_{\bm{\varepsilon}\in
B(\tau)}\mathbb{E}W_{\mathbf{x},\bm{\varepsilon}}W_{\mathbf{x}^{\prime},\bm{\varepsilon}}$
is then
$\sum_{k=0}^{m}\binom{m}{k}c_{\mathbf{x},\mathbf{x}^{\prime}}^{m-k}d_{\mathbf{x},\mathbf{x}^{\prime}}^{k}=\left(c_{\mathbf{x},\mathbf{x}^{\prime}}+d_{\mathbf{x},\mathbf{x}^{\prime}}\right)^{m}$.
∎

###  6.3 Proof of Theorem 3

###### Proof of Theorem 3.

In general, as $\mathbf{x}^{\prime}\rightarrow\mathbf{x}$, we have
$\mathbb{E}V_{\mathbf{x}^{\prime}}\rightarrow\mathbb{E}V_{\mathbf{x}}$. But if
also
$\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}\rightarrow\mathbb{E}V_{\mathbf{x}}^{2}$,
then $a_{\mathbf{x},\mathbf{x}^{\prime}}\rightarrow a_{\mathbf{x},\mathbf{x}}$
by Theorem 2. Thus, (13) $\rightarrow 1$. ∎

###  6.4 Proof of Corollary 2.1

###### Proof of Corollary 2.1.

We begin by proving the result for BT. If $\mathbb{E}V_{\mathbf{x}}=1/2$ for
any $x$, then for BT we have

$\displaystyle\eqref{eq:corrxx}$
$\displaystyle=\left[\frac{\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}}{\sqrt{\mathbb{E}V_{\mathbf{x}}^{2}\mathbb{E}V_{\mathbf{x}^{\prime}}^{2}}}\right]^{m}=\left[\frac{\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})+1/4}{\sqrt{(\texttt{Var}_{V}(\mathbf{x})+1/4)(\texttt{Var}_{V}(\mathbf{x}^{\prime})+1/4)}}\right]^{m}.$

Under the condition $\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})\geq 0$,
(13) is minimized when $\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})=0$
whenever $x\neq x^{\prime}$. Setting this condition results in

$\displaystyle\eqref{eq:corrxx}=\left[(4\texttt{Var}_{V}(\mathbf{x})+1)(4\texttt{Var}_{V}(\mathbf{x}^{\prime})+1)\right]^{-m/2}\geq
2^{-m}.$

This lower bound is strict if $\texttt{Var}_{V}(\mathbf{x})>0$ for all
$x\in\mathcal{X}$.

Now we prove the result for LT. If $\mathbb{E}V_{\mathbf{x}}=1/2$ for any $x$,
then for LT we have

$\displaystyle\eqref{eq:corrxx}$
$\displaystyle=\frac{\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}}{\sqrt{\mathbb{E}V_{\mathbf{x}}^{2}\mathbb{E}V_{\mathbf{x}^{\prime}}^{2}}}\frac{\sqrt{(1-\mathbb{E}V_{\mathbf{x}}^{2})(1-\mathbb{E}V_{\mathbf{x}^{\prime}}^{2})}}{(1-\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}})}\frac{1-[\mathbb{E}V_{\mathbf{x}}V_{\mathbf{x}^{\prime}}]^{K}}{\sqrt{(1-[\mathbb{E}V_{\mathbf{x}}^{2}]^{K})(1-[\mathbb{E}V_{\mathbf{x}^{\prime}}^{2}]^{K})}}$
$\displaystyle=\left(\frac{1}{3/4-\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})}-1\right)\sqrt{\frac{1}{\texttt{Var}_{V}(\mathbf{x})+1/4}-1}\sqrt{\frac{1}{\texttt{Var}_{V}(\mathbf{x}^{\prime})+1/4}-1}$
$\displaystyle\times\frac{1-[\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})+1/4]^{K}}{\sqrt{(1-[\texttt{Var}_{V}(\mathbf{x})+1/4]^{K})(1-[\texttt{Var}_{V}(\mathbf{x}^{\prime})+1/4]^{K})}}.$

If we fix $K$, $\texttt{Var}_{V}(\mathbf{x})$ and
$\texttt{Var}_{V}(\mathbf{x}^{\prime})$, then (13) strictly increases in
$\texttt{corr}_{V}(\mathbf{x},\mathbf{x}^{\prime})$. So under the condition
$\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})\geq 0$, (13) is minimized
when $\texttt{Cov}_{V}(\mathbf{x},\mathbf{x}^{\prime})=0$ whenever $x\neq
x^{\prime}$. Setting this condition results in

$\displaystyle\eqref{eq:corrxx}$
$\displaystyle=\frac{1}{3}\sqrt{\left(\frac{1}{\texttt{Var}_{V}(\mathbf{x})+1/4}-1\right)\left(\frac{1}{\texttt{Var}_{V}(\mathbf{x}^{\prime})+1/4}-1\right)}$
$\displaystyle\times\frac{1-[1/4]^{K}}{\sqrt{(1-[\texttt{Var}_{V}(\mathbf{x})+1/4]^{K})(1-[\texttt{Var}_{V}(\mathbf{x}^{\prime})+1/4]^{K})}}.$

Because $\texttt{Var}_{V}(\mathbf{x})\leq 1/4$ for all $x\in\mathcal{X}$, we
have $\eqref{eq:corrxx}\geq\frac{1}{3}\times\frac{1-4^{-K}}{1-2^{-K}}$. This
lower bound is a decreasing function of $K$. At $K=\infty$, it is $1/3$. This
lower bound is also strict if $\texttt{Var}_{V}(\mathbf{x})>0$ for all
$x\in\mathcal{X}$. ∎

###  6.5 Computational cost of Gibbs step

Next we compare between LT and BT the computational cost of this Gibbs step
for some arbitrarily fixed MCMC iteration. In this case, each tree has the
same number of internal nodes and hence their corresponding mixture models
perform the same number of regressions. From this viewpoint, the computational
cost should be the same between the two SB schemes.

However, we find from experience that when LTM and BTM models are trained on
the same data set, the BTM model usually finishes first. To see why, we note
that for a tree $\tau$ and leaf-assignment $C_{\tau}$, the computational cost
of the Gibbs step can be roughly identified with the sum

$\sum_{i=1}^{n}\big{(}\\#\text{ of $C_{\tau}$'s regressions involving obs
}i\big{)}=\sum_{i=1}^{n}\big{(}\\#\text{ of ancestors of leaf node
}C_{\tau}(i)\big{)}.$

For BT, all observations have $\log_{2}K$ ancestors, which means the sum
equals $n\log_{2}K$. The sum for LT, however, decreases as observations are
allocated to leaf nodes closer to root (meaning the observations are involved
in fewer regressions) and increases as observations are allocated further from
root. We can push this reasoning to the extreme to get the sum’s minimum and
maximum values for LT. If $K^{+}$ clusters are inferred, then the minimum is
achieved when the root’s left-child leaf node contains all but $(K^{+}-1)$
observations and $(K^{+}-1)$ other leaf nodes each contain just one
observation, which results in the sum equaling

$(n-K^{+}+1)+(2+3+4+\cdots+K^{+})=n+K^{+}(K^{+}-1)/2.$

The maximum value is similarly achieved if it is the nonempty leaf node
furthest from root that contains the $(n-K^{+}+1)$ observations, which results
in the sum equaling

$K^{+}(n-K^{+}+1)+(1+2+\cdots+(K^{+}-1))=K^{+}(n-K^{+}+1)+K^{+}(K^{+}-1)/2.$

If $n\gg K^{+}$ and $n$ is fixed, then the former quantity is roughly constant
in $K^{+}$ whereas the latter quantity is roughly linear in $K^{+}$, as shown
in Figure 9.

Though these extreme scenarios provide useful computational cost bounds for
LT, seldom will they be even remotely close to how a leaf assignment
$C_{\tau}$ will allocate observations. On the other hand, what exactly is a
more “typical” allocation, and can we find an analytical representation for it
or for a reasonable approximation of it? The answer to the former question
(and hence also the latter question) depends on, among other things, the size
of the data’s clusters and how well the mixture model infers these sizes. For
simplicity, in the following exposition we will assume the mixture model
infers $K^{+}$ leaf nodes each containing $n/K^{+}$ observations, in which
case the sum for LT becomes $n(K^{+}+1)/2$, which again is linear in $K^{+}$.

The final hurdle for this computational cost comparison between LT and BT is
that our three sum quantities for LT — $n+K^{+}(K^{+}-1)/2$,
$K^{+}(n-K^{+}+1)+K^{+}(K^{+}-1)/2$, and $n(K^{+}+1)/2$ — are functions of $n$
and $K^{+}$ while the sum quantity for BT — $n\log_{2}K$ — is a function of
$n$ and $K$. How can we relate $K$ to $K^{+}$? We might optimistically assume
that the practitioner has chosen $K$ to be the smallest power of two greater
than or equal to $K^{+}$, but we will “leave room for error” by instead
assuming $K$ is the smallest power of two greater than or equal to $2K^{+}$,
which implies $\log_{2}K=(1+\lceil\log_{2}K^{+}\rceil)$. Now all four sum
quantities can be expressed as functions of $n$ and $K^{+}$ as shown in Figure
9. The (scaled) sum for BT grows logarithmically in $K^{+}$ whereas the
(scaled) sum for the “typical” LT scenario grows linearly, which explains why
we find that training a BTM model is usually faster than training its LTM
counterpart.

Figure 9: Scaled sum values of $(1+n^{-1}K^{+}(K^{+}-1)/2)$,
$n^{-1}(K^{+}(n-K^{+}+1)+K^{+}(K^{+}-1)/2)$, $(K^{+}+1)/2$, and
$(1+\lceil\log_{2}K^{+}\rceil)$ for various combinations of $n$ and $K^{+}$.
The first three expressions are displayed as ‘L’ while the fourth expression
is displayed as ‘B’.

##  7 Simulation study plots

Figure 10: Weights for various covariate combinations. Figure 11: Weights for
various covariate combinations.

##  8 GRIFOLS plots

Figure 12: Salient clusters inferred by BTM model with $\zeta=0.1$. Figure 13:
Weight differences between adjacent age groups for various baseline race/sex.
Figure 14: Weight differences between age groups and the youngest age group
for various baseline race/sex. Figure 15: Weight differences between race for
various baseline sex and age group. Figure 16: Weight differences between sex
for various baseline race and age group.