Fast Computation of Branching Process Transition Probabilities via ADMM

TL;DR

This paper introduces an efficient ADMM-based algorithm for computing transition probabilities in branching processes, significantly reducing computational costs and enabling scalable analysis in biological modeling.

Contribution

The paper presents a novel ADMM algorithm that leverages sparse optimization to compute branching process transition probabilities without matrix inversions, enhancing scalability.

Findings

Reduces computational cost by orders of magnitude.

Easily parallelizable and insensitive to tuning parameters.

Demonstrates scalability in biological models of blood cell production and transposon evolution.

Abstract

Branching processes are a class of continuous-time Markov chains (CTMCs) prevalent for modeling stochastic population dynamics in ecology, biology, epidemiology, and many other fields. The transient or finite-time behavior of these systems is fully characterized by their transition probabilities. However, computing them requires marginalizing over all paths between endpoint-conditioned values, which often poses a computational bottleneck. Leveraging recent results that connect generating function methods to a compressed sensing framework, we recast this task from the lens of sparse optimization. We propose a new solution method using variable splitting; in particular, we derive closed form updates in a highly efficient ADMM algorithm. Notably, no matrix products -- let alone inversions -- are required at any step. This reduces computational cost by orders of magnitude over existing methods, and the resulting algorithm is easily parallelizable and fairly insensitive to tuning parameters. A comparison to prior work is carried out in two applications to models of blood cell production and transposon evolution, showing that the proposed method is orders of magnitudes more scalable than existing work.

Tables

Table 1. Table S.1: This table provides complete details of the ADMM algorithm under our proposed framework, including the step size β 𝛽 \beta , the tolerances for the primal and dual feasibility conditions, ϵ p ​ r ​ i = D 1 ​ ϵ a ​ b ​ s + ϵ r ​ e ​ l ​ max ⁡ ( ‖ 𝒖 ( 𝒌 ) ‖ 2 , ‖ 𝒛 ( 𝒌 ) ‖ 2 ) subscript italic-ϵ 𝑝 𝑟 𝑖 subscript 𝐷 1 subscript italic-ϵ 𝑎 𝑏 𝑠 subscript italic-ϵ 𝑟 𝑒 𝑙 subscript norm superscript 𝒖 𝒌 2 subscript norm superscript 𝒛 𝒌 2 \epsilon_{pri}=D_{1}\epsilon_{abs}+\epsilon_{rel}\max(||\boldsymbol{u^{(k)}}||_{2},||\boldsymbol{z^{(k)}}||_{2}) and ϵ d ​ u ​ a ​ l = D 2 ​ ϵ a ​ b ​ s + ϵ r ​ e ​ l , ‖ 𝒚 ( 𝒌 ) ‖ 2 subscript italic-ϵ 𝑑 𝑢 𝑎 𝑙 subscript 𝐷 2 subscript italic-ϵ 𝑎 𝑏 𝑠 subscript italic-ϵ 𝑟 𝑒 𝑙 subscript norm superscript 𝒚 𝒌 2 \epsilon_{dual}=D_{2}\epsilon_{abs}+\epsilon_{rel},||\boldsymbol{y^{(k)}}||_{2} respectively, and ϵ a ​ b ​ s subscript italic-ϵ 𝑎 𝑏 𝑠 \epsilon_{abs} and ϵ r ​ e ​ l subscript italic-ϵ 𝑟 𝑒 𝑙 \epsilon_{rel} are the absolute and relative tolerances respectively. The regularization parameter λ 𝜆 \lambda was chosen to be 0.5 ​ ( log ⁡ M ) 0.5 𝑀 0.5(\log M) for both models. These were chosen in a simple manner across all trials but our framework still yielded promising performance results as compared to the PGD algorithm.

		HSC Model				BDS Model
	$N=64$	$N=128$	$N=256$	$N=512$	$N=1024$	$N=64$	$N=128$	$N=256$	$N=512$	$N=1024$
Stepsize $(\beta$)	$0.08$	$0.005$	$0.08$	$0.005$	$0.005$	$0.005$	$0.005$	$0.005$	$0.0005$	$0.0005$
$D_1$	$N^2$	$N^2$	$N^2$	$N^2$	$N^2$	$N^2$	$N^2$	$N^2$	$N^2$	$N$
$D_2$	$N^5$	$N^5$	$N^5$	$N^5$	$N^5$	$N^2$	$N^2$	$N^2$	$N^2$	$N$
${ϵ}_{abs}$	${10}^{-2}$	${10}^{-2}$	${10}^{-2}$	${10}^{-2}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$
${ϵ}_{abs}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$	${10}^{-3}$

Table 2. Table S.2: A comparison of median running times (percent change relative to PGD algorithm runtime) for different algorithms and scales for the HSC model. For all maximum population sizes, our proposed framework is faster in recovering the transition probabilities than the PGD algorithm for errors at least as good as PGD.

		Median runtime
Maximum population size	M	PGD	ADMM (% change)	ADMM (GPU) (% change)
$N=64$	$51$	$0.179$	$0.009$ $\left(-94.97\%\right)$	$0.016$ $\left(-91.06\%\right)$
$N=128$	$78$	$0.911$	$0.376$ $\left(-58.73\%\right)$	$0.141$ $\left(-84.52\%\right)$
$N=256$	$83$	$4.662$	$1.951$ $\left(-58.15\%\right)$	$0.176$ $\left(-96.22\%\right)$
$N=512$	$88$	$29.34$	$4.721$ $\left(-83.91\%\right)$	$0.227$ $\left(-99.23\%\right)$
$N=1024$	$93$	$869.49$	$20.44$ $\left(-97.65\%\right)$	$0.983$ $\left(-99.89\%\right)$

Table 3. Table S.3: A comparison of median running times (percent change relative to the runtime of the PGD algorithm) for different algorithms and scales for the BDS model. For all maximum population sizes, our proposed framework is faster in recovering the transition probabilities than the PGD algorithm for errors at least as good as PGD.

		Median runtime
Maximum population size	M	PGD	ADMM (% change)	ADMM (GPU) (% change)
$N=64$	$18$	$0.495$	$0.195$ $\left(-60.60\%\right)$	$0.104$ $\left(-78.99\%\right)$
$N=128$	$19$	$1.50$	$1.09$ $\left(-27.33\%\right)$	$0.196$ $\left(-86.93\%\right)$
$N=256$	$29$	$13.02$	$3.50$ $\left(-86.93\%\right)$	$0.216$ $\left(-98.34\%\right)$
$N=512$	$22$	$113.04$	$20.23$ $\left(-82.10\%\right)$	$0.524$ $\left(-99.54\%\right)$
$N=1024$	$28$	$459.93$	$132.85$ $\left(-71.12\%\right)$	$3.303$ $\left(-99.28\%\right)$

Equations

$$p_{\mathbf{x},\mathbf{y}}(t)=\mathbb{P}(\mathbf{X}(t+s)=\mathbf{y}|\mathbf{X}(s)=\mathbf{x}).$$

$${\varphi }_{jk}(t,s_1,s_2;\mathbf{\theta })$$

$$=\sum _{l=0}^{\infty }\sum _{m=0}^{\infty }{p}_{(jk),(lm)}(t;\mathbf{\theta })s_1^l{s}_2^m,$$

$$p_{(jk),(lm)}(t)=\frac{1}{l!m!}\frac{\partial^{l}}{\partial{s_{1}}}\frac{\partial^{m}}{\partial{s_{2}}}\phi_{jk}(t)\Bigr{|}_{\begin{subarray}{c}s_{1}=s_{2}=0\end{subarray}}.$$

$${\varphi }_{jk}(t,e^{2\pi i{\omega }_1},e^{2\pi i{\omega }_2})=\sum _{l,m=0}^{\infty }{p}_{(jk),(lm)}(t)e^{2\pi il{\omega }_1}{e}^{2\pi im{\omega }_2}.$$

$$p_{(jk),(lm)}(t)$$

$$\times e^{-2\pi il{\omega }_1}{e}^{-2\pi im{\omega }_2}d{\omega }_{1}d{\omega }_2$$

$$\approx \frac{1}{N^2}\sum _{u=0}^{N-1}\sum _{v=0}^{N-1}{\varphi }_{jk}(t,e^{2\pi iu/N},e^{2\pi iv/N})$$

$$\times e^{-2\pi ilu/N}{e}^{-2\pi imv/N}.$$

$$\mathbf{u}=\sum _{i=1}^N{\psi }_i{s}_i=\mathbf{\Psi s}.$$

$$\mathbf{b}=\mathbf{\Phi u}=\mathbf{As}\in {\mathbb{C}}^M,$$

$$\hat{\mathbf{s}}=\{\underset{\mathbf{s}}{\mathrm{argmin}}||s|{|}_0:\mathbf{As}=\mathbf{b}\}.$$

$$\hat{\mathbf{s}}=\underset{\mathbf{s}}{\mathrm{argmin}}\frac{1}{2}||\mathbf{As}-\mathbf{b}|{|}_2^2+\lambda ||\mathbf{s}|{|}_1,$$

$$\mathbf{B}={\mathbf{ASA}}^T\in {\mathbb{C}}^{M\times M}$$

$$\hat{\mathbf{S}}=\underset{\mathbf{S}}{\mathrm{argmin}}\frac{1}{2}||{\mathbf{ASA}}^T-\mathbf{B}|{|}_2^2+\lambda ||\mathbf{S}|{|}_1.$$

$$\text{vec}(\mathbf{S})=\tilde{\mathbf{s}}\in {\mathbb{C}}^{N^2},\text{vec}(\mathbf{B})=\tilde{\mathbf{b}}\in {\mathbb{C}}^{M^2},$$

$$\{\mathbf{S}{\}}_{l,m}=p_{jk,lm}(t).$$

$$\tilde{\mathbf{B}}=\{{\varphi }_{jk}(t,e^{2\pi iu/N},e^{2\pi iv/N}){\}}_{u,v=0}^{N-1}\in {\mathbb{C}}^{N\times N}.$$

$$\underset{\mathbf{U},\mathbf{Z}}{\min }$$

$$\mathbf{U}=\mathbf{Z}.$$

$$L_{\beta }(\mathbf{U},\mathbf{Z},\mathbf{Y})$$

$$+{\mathbf{Y}}^T(\mathbf{U}-\mathbf{Z})+\frac{\beta }{2}||\mathbf{U}-\mathbf{Z}|{|}_2^2,$$

$$\begin{array}{rl}{\mathbf{U}}^{(\mathbf{k}+1)}& :=\underset{\mathbf{U}}{\mathrm{argmin}}{L}_{\beta }(\mathbf{U},{\mathbf{Z}}^{(\mathbf{k})},{\mathbf{Y}}^{(\mathbf{k})})\\ {\mathbf{Z}}^{(\mathbf{k}+1)}& :=\underset{\mathbf{Z}}{\mathrm{argmin}}{L}_{\beta }({\mathbf{U}}^{(\mathbf{k}+1)},\mathbf{Z},{\mathbf{Y}}^{(\mathbf{k})})\\ {\mathbf{Y}}^{(\mathbf{k}+1)}& :={\mathbf{Y}}^{(\mathbf{k})}+\beta ({\mathbf{U}}^{(\mathbf{k}+1)}-{\mathbf{Z}}^{(\mathbf{k}+1)})\end{array}$$

$$\mathbf{Z^{(k+1)}}=\text{SoftThresh}\Big{(}\mathbf{U^{(k+1)}}+\frac{\mathbf{Y^{(k)}}}{\beta},\frac{\lambda}{\beta}\Big{)},$$

$$\displaystyle\implies\nabla_{U}\Big{[}\frac{1}{2}||\mathbf{F_{P}^{-1}U(F_{P}^{-1})}^{H}-\mathbf{B}||_{2}^{2}\Big{]}+\nabla_{U}(\lambda||\mathbf{Z}||_{1})$$

$$\displaystyle\qquad\qquad\qquad+\nabla_{U}(\mathbf{Y}^{T}(\mathbf{U}-\mathbf{Z}))+\nabla_{U}\Big{[}\frac{\beta}{2}||\mathbf{U}-\mathbf{Z}||_{2}^{2}\Big{]}=0,$$

$$⟹{({\mathbf{F}}_{\mathbf{P}}^{-1})}^H({\mathbf{F}}_{\mathbf{P}}^{-1})\mathbf{U}{({\mathbf{F}}_{\mathbf{P}}^{-1})}^H({\mathbf{F}}_{\mathbf{P}}^{-1})-{({\mathbf{F}}_{\mathbf{P}}^{-1})}^H\mathbf{B}({\mathbf{F}}_{\mathbf{P}}^{-1})$$

$$+\mathbf{Y}+\beta (\mathbf{U}-\mathbf{Z})=0,$$

$$⟹{({\mathbf{F}}_{\mathbf{P}}^{-1})}^H({\mathbf{F}}_{\mathbf{P}}^{-1})\mathbf{U}{({\mathbf{F}}_{\mathbf{P}}^{-1})}^H({\mathbf{F}}_{\mathbf{P}}^{-1})+\beta \mathbf{U}$$

$$={({\mathbf{F}}_{\mathbf{P}}^{-1})}^H\mathbf{B}({\mathbf{F}}_{\mathbf{P}}^{-1})+\beta \mathbf{Z}-\mathbf{Y}.$$

$$({\tilde{\mathbf{F}}}_{\mathbf{P}}^{-1}{)}^H{\tilde{\mathbf{F}}}_{\mathbf{P}}^{-1}{\tilde{\mathbf{u}}}^{(\mathbf{k}+1)}+\beta {\tilde{\mathbf{u}}}^{(\mathbf{k}+1)}$$

$$\mathbf{\tilde{u}^{(k+1)}}=\Big{[}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{F}_{P}^{-1}}+\beta\mathbf{\tilde{I}}\Big{]}^{-1}\Big{[}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{b}}+\beta\Big{(}\mathbf{\tilde{z}^{(k)}}-\frac{\mathbf{\tilde{y}^{(k)}}}{\beta}\Big{)}\Big{]}.$$

$$\mathbf{M}{\tilde{\mathbf{u}}}^{(\mathbf{k}+1)}=\mathbf{a},\text{where}$$

$$\mathbf{M}=\Big{(}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{F}_{P}^{-1}}+\beta\mathbf{\tilde{I}}\Big{)}\in\mathbb{C}^{N^{2}\times N^{2}}\quad\text{and}$$

$$\mathbf{a}=\Big{(}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{b}}+\beta\Big{(}\mathbf{\tilde{z}^{(k)}}-\frac{\mathbf{\tilde{y}^{(k)}}}{\beta}\Big{)}\Big{)}\in\mathbb{C}^{N^{2}}.$$

$$\hat{\mathbf{M}}{\tilde{\mathbf{F}}}_{\mathbf{P}}^{-1}{\tilde{\mathbf{u}}}^{(\mathbf{k}+1)}=\hat{\mathbf{a}},\text{where}$$

Full text

1 Introduction

Continuous time Markov Chains (CTMCs) have been widely used to stochastically model population dynamics with diverse applications in the fields of genetics, epidemiology, finance, cell biology, and nuclear fission [Renshaw, 2015; Allen, 2010; Paul and Baschnagel, 2013]. Their popularity is owed to their flexibility, interpretability, and desirable mathematical properties. From the perspective of statistical inference, one of the most fundamental quantities characterizing a CTMC is its transition probabilities, the conditional probabilities that a chain ends at a specific state given a starting state and finite time interval do not have closed form expressions. The set of all transition probabilities fully defines a process and forms the backbone of central quantities such as the likelihood function of discretely observed data from a CTMC [Guttorp, 2018]. Unfortunately, obtaining these transition probabilities is often computationally intensive, as it requires marginalization over an infinite set of endpoint-conditioned paths [Hobolth and Stone, 2009].

Classically, this is achieved by computing the matrix exponential of the infinitesimal generator of the CTMC at the expense of having $O(N^{3})$ runtime complexity, where $N$ is the size of the state space. Since this procedure is not suitable even for state spaces of moderate sizes, practitioners often rely on simplifying assumptions or sampling-based approaches via Markov Chain Monte Carlo (MCMC) algorithms [Rao and Teh, 2011; Ross, 1987; Grassmann, 1977], each with their own drawbacks. Many core frequentist and Bayesian inferential methods require an iterative evaluation of the observed data likelihood.

In certain structured classes within CTMCs, alternative approaches are possible. Transition probabilities can be calculated explicitly for simple examples such as the Poisson process, and the class of linear birth-death processes provides another special case that admits closed-form solutions [Crawford et al., 2014; Karlin and Taylor, 1975; Lange, 2010].

In this article, we focus on the case of multi-type branching processes, for which efficient numerical techniques have only recently been developed. Xu et al. [2015] make use of a method that expresses the probability generating function (PGF) as a Fourier series expansion. The generating functions are obtained as solutions to differential equation systems, and the transition probabilities can be recovered from fast numerical series inversion techniques, allowing maximum likelihood inference and Expectation-Maximization (EM) algorithms [Doss et al., 2013]). However, these methods can also become costly for large systems: for instance, approximating the Fourier inversion formula using a Riemann sum requires ${O}(N^{b})$ PGF evaluations, where $b$ is the number of types in the branching process and $N$ is the largest population size at the end of the desired transition probabilities, which entails millions of PGF evaluations even for a two-type process with populations in the thousands. As many likelihood-based procedures are iterative, scalable alternatives to this computation are critical to avoid a bottleneck that renders these methods out of reach. When sparsity is available in the probabilities, Xu and Minin [2015] suggest the use of compressed sensing to reduce the number of computations at the expense of solving a sparse optimization problem using proximal gradient descent (PGD).

To progress, we revisit the compressed sensing framework of Xu and Minin [2015] from the lens of variable splitting methods. We show how the sparse optimization problem can be solved by orders of magnitude more efficiently via a careful implementation of the alternating direction method of multipliers (ADMM) algorithm [Everett III, 1963; Eckstein and Fukushima, 1994; Boyd et al., 2011].

Our approach operates on a vectorized version of the problem, avoiding large matrix inversions that prevent the previous PGD approach from applying to large-scale settings. Surprisingly, not only does variable splitting avoid inversions, but can be accomplished without calls even to matrix multiplication by clever use of the Fast Fourier Transform. Our method inherits desirable convergence and recovery guarantees and admits easily parallelizable implementations. Moreover, it is surprisingly inelastic to tuning parameters — we find that successful performance is robust to a broad range of reasonable penalty parameters. We validate these merits in detailed empirical examples, showcasing significant advantages over prior algorithms applied to a two-type branching model of hematopoiesis, the process of blood cell formation, and a model of transposon evolution from genetic epidemiology.

2 Markov Branching Processes

A Markov branching process is a continuous-time Markov chain comprised of a collection of particles that proliferate independently, and whose reproduction or death follows a probability distribution. We consider continuous-time, multitype branching processes that take values over a discrete state space of non-negative integers where each particle type has its own mean lifetime and reproductive pattern [Lange, 2010; Karlin and Taylor, 1975]. For exposition, we focus on the two-type case, though the results generally apply.

Let $\boldsymbol{X(t)}$ denote a linear, two-type branching process that takes values in a discrete state space $\Omega$, with $X_{i}(t)$ representing the number of particles of type $i$ present at time $t\geq 0$, where each particle of type $i$ at the end of its lifespan can produce $k$ particles of type 1 and $l$ particles of type 2 with instantaneous rates $a_{i}(k,l)$. The overall rates are multiplicative in the number of particles as a result of the linearity of the rate, which follows from the independence assumption.

While they are defined by the instantaneous rates, the transition probabilities provide an alternate characterization of a branching process, also completely defining its dynamics:

[TABLE]

These transition probabilities play a central role in statistical inference, either in the form of maximum likelihood estimation or within Bayesian methods, where likelihood calculations often enter in schemes such as Metropolis-Hastings to sample from the posterior density. Motivated by this, we now focus our attention on computing them in the class of continuous-time branching process.

2.1 Generating function

The probability generating function (PGF) for a two-type process is

[TABLE]

with a natural extension existing for any $m$-type process. Given the instantaneous rates $a_{j}(k,l)$, the system of differential equations governing $\phi_{jk}$ can be derived using the Kolmogorov forward or backward equations as shown by Bailey [1991]. Once we have $\phi_{jk}$, the transition probabilities can be formally obtained by differentiating $\phi_{jk}$ in Equation 2 and normalizing by an appropriate constant,

[TABLE]

However, in practice repeated numerical differentiation is a computationally intensive procedure and often becomes numerically unstable, especially for large $l,m$. To avoid this issue, we use a technique by Lange [1982] that allows us to use the Fast Fourier transform (FFT) to calculate the transition probabilities $p_{(jk),(lm)}(t)$ that occur as coefficients of $s_{1}^{l}s_{2}^{m}$ in Equation 2. We can map the domain $s_{1},s_{2}\in[0,1]\times[0,1]$ to the boundary of the complex unit circle by setting $s_{1}=e^{2\pi i\omega_{1}},s_{2}=e^{2\pi i\omega_{2}}$ and view the PGF as a Fourier series

[TABLE]

We can now compute the transition probabilities via the Fourier inversion formula, approximating the integral by a Riemann sum:

[TABLE]

3 Compressed Sensing Framework

Compressed sensing (CS) is based on the principle that a sparse or compressible signal can be reconstructed, often perfectly, from a small number of its projections onto a certain subspace through solving a convex optimization problem. For our application, transition probabilities play the role of a target sparse signal of Fourier coefficients. CS enables us to restrict the necessary computations to a much smaller sample of PGF evaluations, which play the role of measurements used to recover the sparse signal.

In this setup, let $\mathbf{u}\in\mathbb{C}^{N}$ be an unknown sparse signal and let $\mathbf{\Psi}=[\psi_{1},\psi_{2},\ldots,\psi_{N}]\in\mathbb{C}^{N\times N}$ denote an orthonormal basis of $\mathbb{C}^{N}$. Then there exists a unique $\mathbf{s}\in\mathbb{C}^{N}$ such that

[TABLE]

If the number of non-zeros in $\mathbf{s}$, or $||s||_{0}$, is less than $K$, then $\mathbf{u}$ is said to be $K$-sparse under $\mathbf{\Psi}$. We are interested in cases where $K<<N$ or $\mathbf{u}$ is highly compressible.

Let $M\in\mathbb{Z}^{+}$ satisfy $K<M<<N$ and $\mathbf{u}$ be observed through a measurement

[TABLE]

where $\mathbf{\Phi}\in\mathbb{C}^{M\times M}$ denotes a nonadaptive sensing matrix and $\mathbf{A=\Phi\Psi}$. Here nonadaptiveness means that $\mathbf{\Phi}$ does not depend on $\mathbf{u}$. As $M<<N$, this system is underdetermined and the space of solutions is an infinite affine space. However, in certain sparse settings, Candès [2006] proves that reconstruction can be accurately achieved by finding the most sparse solution among all solutions of Equation 6, i.e.,

[TABLE]

Due to the combinatorial intractability of the non-convex objective function, it is impractical to solve this $\ell_{0}$ problem directly. Instead we use the $\ell_{1}$- relaxation as a proxy which results in a nice convex optimization problem, where we optimize the following unconstrained penalized objective

[TABLE]

with $\lambda$ serving as the regularization parameter to enforce the sparsity of $\mathbf{s}$.

Past works indicate that when $\mathbf{A}$ satisfies the Restricted Isometry Property (RIP) [Candes and Tao, 2005], Candès et al. [2006], the K-sparse signal $\mathbf{u}$ or equivalently $\mathbf{s}$ can be reconstructed using only $M=CK\log N$ measurements for some constant $C$. The result assumes that the columns of $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$ are not just incoherent pairwise, but K-wise incoherent. The coherence between $\boldsymbol{\phi}$ and $\boldsymbol{\psi}$ is defined as $\mu(\boldsymbol{\phi},\boldsymbol{\psi})=\sqrt{n}\mathop{\max}\limits_{1\leq i,j\leq n}|\left<\boldsymbol{\phi}_{i},\boldsymbol{\psi}_{j}\right>|$.

In practice, verifying RIP can be a demanding task; Candès et al. [2006] and Donoho [2006] show that RIP holds with high probability in some Gaussian random matrix settings. In the focus of this paper, $\mathbf{A}$ is comprised of a random “spike" basis playing the role of $\mathbf{\Phi}$ together with a Fourier domain representation $\mathbf{\Psi}$ [Rudelson and Vershynin, 2008; Candès et al., 2006], which form a maximally incoherent pair [Candes and Romberg, 2007]. We see how these components enter the derivation below, and later confirm the success this theory suggests via a thorough empirical study.

3.1 Higher dimensions

We can easily extend this exposition focused on the vector valued case to higher-dimensional signals [Candès, 2006]. To illustrate this, consider the $2D$ case where the sparse solution $\mathbf{S}\in\mathbb{C}^{N\times N}$ and the measurement

[TABLE]

are matrices instead of vectors. We may solve an equivalent problem

[TABLE]

Equivalently, this can also be represented in a vectorized framework with

[TABLE]

and we now pursue $\tilde{\mathbf{b}}=\tilde{\mathbf{A}}\tilde{\mathbf{s}}$, with $\tilde{\mathbf{A}}=\mathbf{A}\bigotimes\mathbf{A}$ being the Kronecker product of $\mathbf{A}$ with itself. It can be preferable to solve Equation 10 rather than the vectorized problem explicitly, as the number of entries in $\tilde{\mathbf{A}}$ grows rapidly. However, we will show how working with vectorized forms together with clever implementations of the Fast Fourier Transform will provide the best of both worlds, avoiding matrix multiplication and inversion entirely.

4 Transition Probabilities via ADMM

Recall that we wish to compute the transition probabilities $p_{jk,lm}(t)$ for any $t>0$ and initial $\mathbf{X}(0)=(j,k)$. Within the CS framework described in the previous section, $\mathbf{S}\in\mathbb{C}^{N\times N}$ is the matrix of transition probabilities with entries

[TABLE]

Without the CS framework, one can directly obtain the transition probabilities using Equation 4 by first computing an equally sized matrix of PGF solutions

[TABLE]

Obtaining $\mathbf{\tilde{B}}$ becomes computationally expensive for large $N$ values, especially when each PGF must be solved, for instance via numerically evaluating a differential equation. Given a way to compute $\mathbf{\tilde{B}}$, we recover the transition probabilities by taking the Fast Fourier Transform (FFT). We can better understand how this fits within the CS framework with the help of matrix operations. We have $\mathbf{S}=\mathbf{F}\mathbf{\tilde{B}}\mathbf{F^{T}}$, where $\mathbf{{F}}\in\mathbb{C}^{N\times N}$ denotes the Discrete Fourier Transform (DFT) matrix. Thus, the Inverse Discrete Fourier Transform (IDFT) matrix $\mathbf{{F}^{-1}}$ becomes the sparsifying basis $\mathbf{\Psi}$ mentioned in Equation 5, and we have $\mathbf{\tilde{B}}=\mathbf{\Psi S\Psi^{T}}$.

If we expect $\mathbf{S}$ to have a sparse representation, we can employ a method to reconstruct $\mathbf{S}$ using a much smaller set of PGF evaluations arranged in the matrix $\mathbf{B}\in\mathbb{C}^{M\times M}$, corresponding to a subset of entries from $\mathbf{\tilde{B}}$ selected uniformly at random. This much smaller matrix $\mathbf{\tilde{B}}$ is a projection $\mathbf{B}=\mathbf{ASA}^{T}\in\mathbb{C}^{M\times M}$ in accordance with Equation 9, where $\mathbf{A}\in\mathbb{C}^{M\times N}$ is obtained by selecting a subset of rows of $\mathbf{\Psi}$ that correspond to $\mathcal{J}$, the randomly sampled indices. This uniform sampling of rows is equivalent to multiplication by measurement matrix encoding the spike (or standard) basis. Mathematically, we have $\mathbf{A=\Phi\Psi}$, with the rows of the measurement matrix $\mathbf{\Phi}_{j}(l)=\delta(j-l)$. Thus, uniformly sampling the indices $\mathcal{J}$ is optimal in our setting because the spike and Fourier bases are maximally incoherent in any dimension [Candès and Wakin, 2008].

In the compressed sensing framework, we now require only computing this reduced matrix $\mathbf{B}$, which entails only a logarithmic proportion $|\mathbf{B}|\propto K\log|\mathbf{\tilde{B}}|$ of PGF evaluations compared to the original problem. Thus, the problem of computing the transition probabilities $\mathbf{S}$ has been reduced to a signal recovery problem that can be solved by solving the optimization problem in Equation 10.

4.1 Solving the convex problem using ADMM

We use the Alternating Direction Method of Multipliers (ADMM) algorithm [Everett III, 1963; Boyd et al., 2011] to solve this problem efficiently. ADMM is useful for optimization problems where the objective function is of the form $\{f(\mathbf{x})+g(\mathbf{y}):\mathbf{x}=\mathbf{y}\}$ with both $f(.)$ and $g(.)$ convex but not necessarily differentiable.

We first introduce some notations before detailing the optimization routine. Recall that for our application, $\mathbf{U}$ is a two-dimensional unknown sparse signal being observed through a measurement matrix $\mathbf{B}$, and $\mathbf{F_{P}}=\mathbf{PF}\in\mathbb{C}^{M\times N}$ denotes the sensing matrix, where $\mathbf{P}\in\mathbb{R}^{M\times N}$ is a selection matrix containing $M$ rows of the identity matrix of order $N$. Similarly, we will denote $\mathbf{F^{-1}_{P}}=\mathbf{PF^{-1}}\in\mathbb{C}^{M\times N}$ where $\mathbf{{F}^{-1}}$ is the IDFT matrix.

To solve this problem efficiently using ADMM, we first split the optimization variable by introducing a new variable $\mathbf{Z}$ under an equality constraint with the variable of interest:

[TABLE]

Next, we consider the Augmented Lagrangian

[TABLE]

where $\mathbf{Y}$ serves as the dual variable enforcing $\mathbf{U}=\mathbf{Z}$. Now, the ADMM algorithm entails the following iteration:

[TABLE]

Deriving the explicit update for the subproblem in $\mathbf{Z}$ here is straightforward, given by the soft-thresholding operator

[TABLE]

where SoftThresh$(\epsilon,\lambda/\beta)$ = $\max\{|\epsilon|-\lambda/\epsilon,0\}\cdot sign(\epsilon)$ for a constant $\epsilon\in\mathbb{R}^{+}$. To minimize $\mathbf{U}$, we set \nabla_{U}\Big{[}L_{\beta}(\mathbf{U},\mathbf{Z},\mathbf{Y})\Big{]} equal to [math],

[TABLE]

Notice that Equation 15 cannot be used to obtain a closed form update for $\mathbf{U}$. However, we can make progress by vectorizing, allowing us to rewrite Equation 15 as

[TABLE]

where $\mathbf{\tilde{F}_{P}^{-1}}=\mathbf{(F_{P}^{-1})}\bigotimes\mathbf{(F_{P}^{-1})}\in\mathbb{C}^{M^{2}\times N^{2}}$ and $\mathbf{\tilde{I}}=\mathbf{I}\bigotimes\mathbf{I}\in\mathbb{R}^{N^{2}\times N^{2}}$. This yields the update for $\mathbf{\tilde{u}}$

[TABLE]

Avoiding matrix operations: Updating $\mathbf{\tilde{u}}$ naïvely according to Equation 16 requires dense matrix multiplication and inversion, steps that would significantly add to the runtime of our algorithm in large-scale settings.

To circumvent this, we may instead think of Equation 16 as the solution to the linear system

[TABLE]

[TABLE]

[TABLE]

Multiplying $\mathbf{\tilde{F}^{-1}}=\mathbf{F^{-1}}\bigotimes\mathbf{F^{-1}}\in\mathbb{C}^{N^{2}\times N^{2}}$ to both sides of the above equation, we obtain

[TABLE]

[TABLE]

The complete details of these derivations appear step-by-step in the Supplement. Note that $\mathbf{\hat{M}}$ is a diagonal matrix of order $N^{2}$ because both $\beta\mathbf{\tilde{I}}$ and $\mathbf{\tilde{P}^{T}\tilde{P}}$ are diagonal matrices of order $N^{2}$. Of course in practice we only compute and store the vector $\mathbf{\hat{m}}$ containing the $N^{2}$ diagonal elements of $\mathbf{\hat{M}}$. For exposition, let $N=4,M=2,\beta=0.1$ and the uniformly sampled indices $\mathcal{J}=[1,2]$, then we can compute $\mathbf{\hat{m}}$ as,

[TABLE]

On the other hand, since $\mathbf{\hat{a}}=\text{vec}(\mathbf{\hat{A}})=\text{vec}(\mathbf{F^{-1}}\beta\mathbf{Z^{(k)}}+\mathbf{P^{T}BP})$, computing $\mathbf{\hat{a}}$ is equivalent to first computing $\mathbf{\hat{A}}$ and then vectorizing it. Thus, we can compute $\mathbf{\hat{m}}$, compute $\mathbf{\hat{a}}$ and then efficiently solve the linear system

[TABLE]

where $\oslash$ is the Hadamard (elementwise) division operator and $\text{vec}^{-1}$ is the inverse of the vectorization operator, i.e. reshaping $\mathbf{\hat{a}}\oslash\mathbf{\hat{m}}$ back into a $N\times N$ matrix. Thereafter, to isolate the optimization variable, we may recover $\mathbf{U^{(k+1)}}$ by “cancelling out” $\mathbf{F^{-1}}$ by applying the FFT on both sides. Doing so avoids explicitly multiplying matrices implied by 16; details are mentioned in line $8$ of Algorithm 1.

Computational complexity: We now we analyze the per-iteration computational complexity of the subproblems in minimizing $\mathbf{Z}$ and $\mathbf{U}$. First, note the minimization of $\mathbf{Z}$ requires evaluating the soft-thresholding operator elementwise, and thus has a complexity of $O(kN^{2})$, where k is the number of ADMM iterations and N is the maximum population size.

Next, we consider the complexity of naively minimizing $\mathbf{U}$ if one were to use Equation 16 directly. Doing so would involve dense matrix multiplications to calculate $\mathbf{M}$ and $\mathbf{a}$, as well as matrix inversion to obtain $\mathbf{U}$. Calculating $\mathbf{M}$ has a complexity of $O(N^{4}M^{2})$ as it involves matrix multiplication of $\mathbf{\tilde{F}_{P}^{-1}}\in\mathbb{C}^{M^{2}\times N^{2}}$ with its complex conjugate. Calculating $\mathbf{a}$ has a complexity of $O(\max\{(N^{2}M^{2},kN^{2}\})$ while inverting $\mathbf{M}$ has $O(N^{6})$ cost, thus incurring a total computational complexity of $O(\max\{N^{6},\max\{N^{2}M^{2},kN^{2}\}\})=O(N^{6})$ for the $\mathbf{U}$ update.

Our discussion on reducing this complexity by avoiding matrix operations reduces this significantly, making use of $\mathbf{\hat{m}}$ and $\mathbf{\hat{A}}$ instead of $\mathbf{M}$ and $\mathbf{a}$ as described above. Forming $\mathbf{\hat{m}}$ has $O(N^{2})$ cost as it involves the calculation of $N^{2}$ non-zero diagonal elements of $\mathbf{\hat{M}}$. Then, computing $\mathbf{\hat{A}}$ requires the calculation of $\mathbf{P^{T}BP}$, $\beta\mathbf{Z^{(k)}}$ and its IDFT, which have complexities $O(M^{2})$, $O(kM)$, and $O(kN^{2}\log_{2}N)$ respectively. Therefore, the calculation of $\mathbf{\hat{A}}$ has a complexity of $O(\max\{M^{2},kM,kN^{2}\log_{2}N\})=O(kN^{2}\log_{2}N)$ which matches the cost of computing two-dimensional FFT required to obtain $\mathbf{U^{(k+1)}}$. Since this is the dominant complexity in solving the linear system described by Equation 18, the overall complexity of the $\mathbf{U}$ update (and in turn the $\mathbf{Z}$ update) becomes $O(kN^{2}\log_{2}N)$, a dramatic reduction from $O(N^{6})$.

It is worth mentioning that the total complexity of the algorithm can be further improved by leveraging the sparsity of the transition probability matrix and using sparse FFT (sFFT) techniques [Hassanieh et al., 2012], which can reduce the complexity of the two-dimensional IDFT from $O(kN^{2}\log_{2}N)$ to $O(kM^{2}\log_{2}M)$. However, for our application, we use the Fastest Fourier Transform in the West (FFTW) [Frigo and Johnson, 1998] to implement FFT as the differences in execution times between implementing FFT using FFTW and sFFT are only appreciable for $N>2^{17}$ [Hassanieh et al., 2012], which is beyond the maximum population sizes considered in this paper.

Execution time of ADMM:

We further improve the wall-clock execution time of our algorithm by vectorizing intermediate steps (lines 2, 3 and 8 of Algorithm 1) whenever possible, and making use of optimized Numpy libraries Harris et al., 2020 and parallelized FFT in FFTW [Frigo and Johnson, 1998].

4.2 ADMM Convergence and Stopping Criteria

Having detailed the algorithm and having emphasized careful diagonalization tricks in the Fourier domain to significantly reduce runtime, we now discuss aspects related to convergence. We establish the convergence guarantees and stopping criteria of our ADMM algorithm. It is not difficult to show that our formulation inherits powerful standard convergence results. We begin by verifying assumptions behind one of the classical convergence theorems.

Proposition 4.1

The functions $f(\mathbf{U})$ and $g(\mathbf{Z})$ are closed, convex and proper.

The functions $f(\mathbf{U})=\frac{1}{2}||\mathbf{F_{P}^{-1}}\mathbf{U}\mathbf{(F_{P}^{-1}})^{H}-\mathbf{B}||^{2}_{2}$ and $g(\mathbf{Z})=\lambda||\mathbf{Z}||_{1}$ are closed, convex and proper by virtue of being the square of the $L_{2}$ norm and $L_{1}$ norm, respectively [Rudin et al., 1976; Folland, 1999].

Proposition 4.2

For $\beta=0$, $L_{0}=\lambda||\mathbf{Z}||_{1}+\frac{1}{2}||\mathbf{F_{P}^{-1}U\mathbf{(F_{P}^{-1}})^{H}}-\mathbf{B}||_{2}^{2}+\mathbf{Y}^{T}(\mathbf{U}-\mathbf{Z})$, the unaugmented Lagrangian has a saddle point for all $\mathbf{U},\mathbf{Z},\mathbf{Y}$.

Proposition 4.3

Under propositions 4.1 and 4.2, the ADMM algorithm achieves:

Primal residual convergence: $\mathbf{R^{(k+1)}}=\mathbf{U^{(k+1)}}-\mathbf{Z^{(k+1)}}\to\mathbf{0}$ as $k\to\infty$. 2. 2.

Dual residual convergence: $\mathbf{S^{(k+1)}}=\beta(\mathbf{Z^{(k+1)}}-\mathbf{Z^{(k)}})\to\mathbf{0}$ as $k\to\infty.$

Proofs of Propositions 4.2 and 4.3 are given in the Supplement.

4.2.1 Optimality Conditions and Stopping Criteria

The necessary and sufficient optimality conditions for the optimization problem described in Equation 12 are defined by primal feasibility,

[TABLE]

and dual feasibility,

[TABLE]

where $(\mathbf{U^{*},Z^{*},Y^{*}})$ is the saddle point of $L_{0}$, $\partial$ denotes the subdifferential operator [Rockafellar, 1970; Hiriart-Urruty and Lemaréchal, 2004; Borwein and Lewis, 2006] due to $\lambda||\mathbf{Z^{*}}||_{1}$ being non-differentiable. Thus, Equations 19-21 constitute the three optimality conditions for the optimization problem described in Equation 12. The last condition (21) always holds for $(\mathbf{U^{(k+1)},Z^{(k+1)},Y^{(k+1)}})$ whereas the other conditions give rise to the primal $\mathbf{R^{(k+1)}}$ and dual $\mathbf{S^{(k+1)}}$ residuals with proof in the Supplement; these residuals converge to zero as ADMM iterates according to Proposition 4.3.

We can thus use these residuals to put a bound on the objective suboptimality of the current point, $\frac{1}{2}||\mathbf{F_{P}^{-1}U^{(k)}}\mathbf{(F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}+\lambda||\mathbf{Z^{(k)}}||_{1}-q^{*}$, where $q^{*}$ is the optimal value that the objective function given in Equation 12 will converge to for $k\to\infty$. The following equation appearing in the proof of convergence in the Supplement,

[TABLE]

suggests that when the primal and dual residuals are small, the objective suboptimality must also be small, thereby motivating the use of the following stopping criterion:

[TABLE]

Here $\epsilon_{pri}>0$ and $\epsilon_{dual}>0$ are the tolerances for the primal (19) and dual feasibility (20) conditions respectively. We choose the following tolerances for robustness analysis,

[TABLE]

where N is the dimension of $\mathbf{U}$, $\epsilon_{abs}$ and $\epsilon_{rel}$ are the absolute and relative tolerances respectively, typically chosen between $10^{-2}$ to $10^{-4}$ depending on the application.

5 Application and Empirical Study

We propose to implement the framework described in the previous section to recover the transition probabilities for a two-type branching process; the idea and the framework can be easily extended to multi-type branching processes as well. As an outline, we calculate the full transition probability matrix $\mathbf{S}$ for a two-type branching process with known rates, randomly select a subset of the PGF evaluations $\mathbf{B}$, and use ADMM along with CS to recover the transition probabilities $\mathbf{\hat{S}}$ while also exploring the robustness of the ADMM algorithm with respect to the stepsize $\beta$, and the regularization parameter $\lambda$. Once satisfied with the robustness, we then proceed to compare the runtimes of recovering the transition probabilities using ADMM algorithm with the runtimes of PGD algorithm for similar errors.

We describe the two model applications used to generate data in our study.

5.1 Two-compartment hematopoiesis model

Hematopoiesis is a process of continuous formation and turnover of blood cells in order to fulfill the everyday demands of the body. At the origin of hematopoiesis is hematopoietic stem cell (HSC) which has two fundamental features: the first being the ability to self-renew, a divisional event which results in the formation of two HSCs from one HSC and the second being the ability for multipotent differentiation into all mature blood lineages.

We can stochastically model this biological phenomenon as a two-type branching process [Catlin et al., 2001]. Under such representation, $X_{1}$ and $X_{2}$, the type one and type two particle populations correspond to HSCs and progenitor cells, respectively. With the same parameters as denoted in Figure 2(a) of the Supplement, the non-zero instantaneous rates that define the process are,

[TABLE]

Now that we have the instantaneous rates, we can derive the solutions for its PGF, defined in Equation 2, and subsequently obtain the transition probabilities using Equation 4 with the details included in the Supplement. Note that the cell population in this application can reach the order of tens of thousands, thereby motivating the splitting methods approach to compute transition probabilities in an efficient manner.

5.2 Birth-death-shift model for transposons

For our second application, we look at the birth-death-shift (BDS) process proposed by Rosenberg et al. [2003] to model evolutionary dynamics of the genomic mobile sequence elements known as transposons. Each transposon can either produce a new copy that can move to a new genomic location, shift to a different genomic location, or be completely removed from the genome, independently of all other transposons with per-particle instantaneous rates $\gamma,\sigma,\delta$ and overall rates proportional to the total number of transposons. The $IS6110$ transposon in the Mycobacterium tuberculosis genome from a San Francisco community study dataset [Cattamanchi et al., 2006] was used to estimate these evolutionary rates by Rosenberg et al. [2003]. A two-type branching process can be used to model the BDS process over any finite observation interval [Xu et al., 2015], where $X_{1}$ and $X_{2}$ represent the numbers of initially occupied genomic locations and newly occupied genomic locations, respectively, thus capturing the full dynamics of the BDS process.

The two-type branching process representation of the BDS model has the following non-zero rates

[TABLE]

Proposition 5.1

The PGF of the BDS model described in Equation 25 is given by $\phi_{j,k}=\phi_{1,0}^{j}\phi_{0,1}^{k}$ by particle independence where,

[TABLE]

For our application, we calculate the transition probabilities $\mathbf{S}$ for maximum population sizes $N=2^{6},2^{7},\ldots,2^{10}$, given the time intervals $t$, branching process’ rate parameters $\boldsymbol{\theta}$ and initial population size $\mathbf{X(0)}$. Based on biologically sensible rates and observation times scales of data from previous hematopoiesis studies [Catlin et al., 2001; Golinelli et al., 2006; Fong et al., 2009], we set per-week branching rates and the observation time as $\boldsymbol{\theta}_{HSC}=[\rho,\nu,\mu]=[0.125,0.104,0.147]$ and $t=1$ week respectively. Similarly for the BDS application, based on previously estimated rates in Xu et al. [2015] and the average length between observations in the tuberculosis dataset from San Francisco [Cattamanchi et al., 2006], we set the per year event rates and the observation time as $\boldsymbol{\theta}_{BDS}=[\gamma,\sigma,\delta]=[0.016,0.004,0.019]$ and $t=0.35$ years respectively. In each case, we computed $M^{2}=(\min(\lfloor\sqrt{\log NK10}\rfloor,\lfloor\sqrt{N}-\sqrt{N}/5\rfloor)^{2}$ total random measurements to obtain $\mathbf{B}$, which will be used to recover the transition probabilities.

5.3 Robustness of ADMM algorithm

We begin by investigating the robustness of our ADMM algorithm under a broad range of parameter settings. We consider varying the stepsize $\beta$ and the regularization parameter $\lambda$ for each $N$, and repeat this procedure across all $N$ values for the HSC model while examining the effects of these variations on the relative $L_{2}$ errors $\epsilon^{L_{2}}_{rel}$ defined

[TABLE]

where $||.||$ indicates the $L_{2}$ norm, $\boldsymbol{S}$ and $\boldsymbol{\hat{S}}$ represent the “true” and recovered transition probabilities respectively.

In more detail, we start by fixing $\lambda=0.5(\log M)$ and vary $\beta$ in equal intervals from $10^{-4}$ to $9\times 10^{2}$ while recording and subsequently plotting $\epsilon^{L_{2}}_{rel}$ corresponding to each $\beta$. Thereafter, we fix $\beta=0.08$, vary $\lambda$ from $1$ to $10^{3}$ and plot $\epsilon^{L_{2}}_{rel}$ corresponding to each $\lambda$. The primal $\epsilon_{pri}$ and dual $\epsilon_{dual}$ feasibility tolerances are given by Equations 22 and 23 respectively with $\epsilon_{abs}=\epsilon_{rel}=10^{-3}$.

Figure 1 displays these results. The algorithm is fairly inelastic to both the tuning of the stepsize $\beta$ as well as the penalty parameter $\lambda$: it achieves low errors conveying successful signal recovery for a wide range of values, which is visually evident even on the log scale. This is promising as our method may enter as a subroutine in optimization frameworks when the scale of these parameters is unknown, and must be learned throughout an iterative scheme.

5.4 Comparing the performance of PGD and ADMM

Having confirmed that the method is not too sensitive to tuning, we compare its performance to the PGD algorithm proposed in Xu and Minin [2015]. In order to compare the performance we first compute sets of transition probabilities S of both HSC and BDS models using the full set of PGF solution measurements $\mathbf{\tilde{B}}$ as described in Equation 11.

Following Xu and Minin [2015] we set the regularization parameters as $\lambda^{PGD}_{HSC}=\sqrt{\log M}$ and $\lambda^{PGD}_{BDS}=\log M$ for the PGD algorithm while keeping $\lambda=0.5(\log M)$ for the ADMM algorithm for both the applications. It is promising that this simple heuristic for choosing the parameters leads to successful performance across all settings we consider, further validating our empirical validation of robustness to tuning. Figure 1(c) illustrates the sparse solution and provides a visual demonstration of the accuracy of reconstructed solutions using our ADMM approach, with the estimated matrix $\hat{\mathbf{S}}$ essentially identical to the ground truth.

For each value of N, we repeat this procedure five times under randomly sampled indices $\mathcal{J}$, and recover the transition probabilities $\mathbf{\hat{S}}$ using only a subset of PGF solution measurements $\mathbf{B}$ for both PGD and ADMM algorithms. We report median runtimes over the five trials in a fair conservative comparison, in that we ensure the measure of accuracy $\epsilon^{L_{2}}_{rel}$ under our proposed method is at least as good as PGD. To match errors across the algorithms, we choose

[TABLE]

[TABLE]

where $D_{1},D_{2}$ are $\{N^{2},N^{5}\}$ and $\{N^{2},N^{2}\}$ for the HSC and BDS models respectively, complete details summarized in Table S1 of the Supplement. Since our ADMM algorithm largely depends on the FFT, a GPU based algorithm was also implemented to show that the algorithm can be accelerated drastically through straightforward parallelization, thus making it highly scalable for large data; further details on the GPU version are included in the Supplement.

The median runtimes of PGD, ADMM (Vanilla), and ADMM (GPU) for both HSC and BDS models are reported in Figure 2 while the median relative $L_{2}$ norm errors, $\epsilon^{L_{2}}_{rel}$ are reported in the Supplement. As evident from Figure 2, the ADMM algorithm consistently outperforms the PGD algorithm across all $N$ values for both the HSC and BDS models. For instance, with $N=1024$, there is a $97.65\%$ and $71.12\%$ reduction in the runtimes of ADMM (Vanilla) algorithm for the HSC model and BDS models respectively while the GPU implementation of ADMM is about $21$ times faster than ADMM (vanilla) and about $885$ times faster than PGD algorithm for the HSC model. Similarly, for $N=1024$ the ADMM (GPU) algorithm is about $40$ times faster than ADMM (vanilla) and about $140$ times faster than PGD algorithm for the BDS model.

Computing Infrastructure

The experiment is conducted on Google Colaboratory111https://colab.research.google.com/ platform that claimed to have a RAM of 12.6 GB, a single-core Xeon CPU of 2.20 Ghz (No Turbo Boost) and a Tesla K80 with 2496 CUDA cores and 12GB GPU memory as well. The GPU version of FFT was implemented in CuPy [Okuta et al., 2017], a Numpy-like API library for CUDA with matched parameters in both the vanilla and GPU implementations of ADMM. The code and data used to implement the ADMM algorithm can be accessed from https://github.com/awasthi-12/Transition_Prob_via_ADMM/.

6 Discussion

We revisit the computational challenge of computing transition probabilities of large-scale multi-type branching processes. Focusing on the two-type setting, we derive a novel application of ADMM, showing how variable splitting can significantly accelerate methods to compute transition probabilities within a compressed sensing paradigm. Through a suite of experiments, we validate not only the robustness of ADMM with respect to the stepsize $\beta$ and the regularization parameter $\lambda$, but also its superior performance over previous attempts utilizing PGD. The advantages of the proposed method become especially pronounced as the population size $N$ increases. We also show through thoughtful algorithmic considerations that the primal variable “$\mathbf{U}$”-update in ADMM can be carried out without matrix inversions. In particular, mathematically manipulating the expressions to leverage the FFT, the dominant complexity requires $O(N^{2}\log_{2}N)$ flops, dramatically increasing the scale of problems the method can consider. Lastly, a GPU-based parallel implementation further reduces the runtimes over the standard version.

In many realistic data settings where these stochastic population models apply, it is natural to expect sparsity in the support of transition probabilities. We have been able to achieve precise results in two such relevant examples under realistic parameter settings from scientific literature. Importantly, not only can the relevant quantities be computed efficiently, but in a way that succeeds over a wide range of tuning parameters. This robustness is crucial when embedding these methods as subroutines within inferential schemes such as maximum likelihood estimation [Doss et al., 2013; Xu et al., 2015]. Recall that for a Markov process $\mathbf{X}(t)$ observed at times $\{t_{1},\ldots t_{T}\}$, the likelihood function of observed data is a product of its transitions

[TABLE]

Optimizing this likelihood numerically often entails iterative schemes, and similar bottlenecks arise under other estimators as well as Bayesian approaches where the likelihood appears in routines such as Metropolis-Hastings ratios [Guttorp, 2018; Stutz et al., 2022]. The robustness, efficiency, and accuracy of the proposed method make it well-suited to open the door to likelihood-based inference in previously intractable settings.

Acknowledgments

This work was partially supported by NSF grants DMS-2230074 and PIPP-2200047. We thank Galen Reeves for helpful discussions and Yiwen Wang for early contributions to the code.

Appendix A Derivations and Proofs

A.1 Discrete Fourier transform matrix

The $N\times N$ discrete Fourier transform (DFT) matrix $\mathbf{F}$ has entries $\{\mathbf{F}\}_{j,k}=\frac{1}{\sqrt{N}}(\omega)^{jk}$, where $\omega=e^{2\pi i/N}$ and $j,k=0,1,\dots N-1$. As mentioned in the main text, recall that the Hermitian is given by its conjugate transpose: \mathbf{F}^{H}=\big{(}\boldsymbol{\mathbf{\overline{F}}}\big{)}^{T}. Now let $\mathbf{x}$ be a vector with $\boldsymbol{\hat{x}}=\mathbf{F}\mathbf{x}$ and $\mathbf{x}=\mathbf{F}^{H}\boldsymbol{\hat{x}}$. Applying both $\mathbf{F}^{H}$ and $\mathbf{F}$ we get

[TABLE]

Thus, the DFT matrix $\mathbf{F}$ is unitary.

A.2 Extended Derivations

Here we provide details on obtaining $\mathbf{\hat{M}}$ and $\mathbf{\hat{a}}$ as described by Equation 17 of the main text. This enables us to update $\mathbf{u}$ efficiently, as mentioned in Section 4.2 of the main text.

We begin with the following equation

[TABLE]

where \mathbf{M}=\Big{(}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{F}_{P}^{-1}}+\beta\mathbf{\tilde{I}}\Big{)}\in\mathbb{C}^{N^{2}\times N^{2}} and \mathbf{a}=\Big{(}(\mathbf{\tilde{F}_{P}^{-1}})^{H}\mathbf{\tilde{b}}+\beta\Big{(}\mathbf{\tilde{z}^{(k)}}-\frac{\mathbf{\tilde{y}^{(k)}}}{\beta}\Big{)}\Big{)}\in\mathbb{C}^{N^{2}}. Multiplying both sides of Equation A.1 by $\mathbf{\tilde{F}^{-1}}=\mathbf{F^{-1}}\bigotimes\mathbf{F^{-1}}\in\mathbb{C}^{N^{2}\times N^{2}}$ and simplifying using the properties of the discrete Fourier transform matrix, the left-hand side of the equation becomes:

[TABLE]

Similarly, the right-hand side of Equation A.1 becomes:

[TABLE]

A.3 Proofs of propositions

We present the proofs of propositions 4.2 and 4.3 described in Section 4.2 of the main text, first recalling the statement of results:

Proposition A.1

$\!\mathbf{4.1}$* The functions $f(\mathbf{U})=\frac{1}{2}||\mathbf{F_{P}^{-1}}\mathbf{U}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||^{2}_{2}$ and $g(\mathbf{Z})=\lambda||\mathbf{Z}||_{1}$ are closed, convex, and proper.*

The proof of Proposition 4.1 has been provided in the main text. Proposition $\mathbf{4.1}$ implies that there exist $\mathbf{U}$ and $\mathbf{Z}$, not necessarily unique, that minimize the augmented Lagrangian, thereby making the $\mathbf{U}$ and $\mathbf{Z}$ updates

[TABLE]

solvable.

Proposition A.2

$\!\mathbf{4.2}$* For $\beta=0$, $L_{0}=\frac{1}{2}||\mathbf{F_{P}^{-1}U}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}+\lambda||\mathbf{Z}||_{1}+\mathbf{Y}^{T}(\mathbf{U}-\mathbf{Z})$, the unaugmented Lagrangian has a saddle point for all $\mathbf{U},\mathbf{Z},\mathbf{Y}$.*

Proof: A.1

A Lagrangian $L_{0}$ is said to have a saddle point when there exist $(\mathbf{U^{*}},\mathbf{Z^{*}},\mathbf{Y^{*}})$ not necessarily unique such that,

[TABLE]

From Proposition $\mathbf{4.1}$ it follows that $L_{0}(\mathbf{U^{*},Z^{*},Y^{*}})$ is finite for any saddle point $(\mathbf{U^{*},Z^{*},Y^{*}})$. This implies not only that $(\mathbf{U^{*},Z^{*}})$ is a solution to the primal objective function $\frac{1}{2}||\mathbf{F_{P}^{-1}U}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}+\lambda||\mathbf{Z}||_{1}$, thus $\mathbf{U^{*}}=\mathbf{Z^{*}}$ and $\frac{1}{2}||\mathbf{F_{P}^{-1}U^{*}}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}<\infty$ and $\lambda||\mathbf{Z^{*}}||_{1}<\infty$, but also that strong duality holds and $\mathbf{Y^{*}}$ is a dual optimal, i.e. it maximizes the dual function $\mathbf{Y}^{T}(\mathbf{U}-\mathbf{Z})$ [Boyd and Vandenberghe, 2004].

We can easily verify that

[TABLE]

where the last inequality follows from $(\mathbf{U^{*},Z^{*}})$ minimizing $\frac{1}{2}||\mathbf{F_{P}^{-1}U}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}+\lambda||\mathbf{Z}||_{1}$ among all $(\mathbf{U},\mathbf{Z})$ and $\mathbf{Y^{*}}$ maximizing the dual objective function among all $\mathbf{Y}$.

Proposition A.3

$\!\mathbf{4.3}$* Under propositions $\mathbf{4.1}$ and $\mathbf{4.2}$, the ADMM algorithm achieves the following:*

Primal residual convergence: $\mathbf{R^{(k+1)}}=\mathbf{U^{(k+1)}}-\mathbf{Z^{(k+1)}}\to 0$ as $k\to\infty$. 2. 2.

Dual residual convergence: $\mathbf{S^{(k+1)}}=\beta(\mathbf{Z^{(k+1)}}-\mathbf{Z^{(k)}})\to 0$ as $k\to\infty.$

Proof: A.2

The proof of proposition $\mathbf{4.3}$ [Gabay, 1983; Eckstein and Bertsekas, 1992; Boyd et al., 2011] is divided into three parts based on three central inequalities, all of which we prove as part of this proof.

Let $(\mathbf{U^{*}},\mathbf{Z^{*}},\mathbf{Y^{*}})$ be a saddle point for the unaugmented Lagrangian $L_{0}=\frac{1}{2}||\mathbf{F_{P}^{-1}U}(\mathbf{F_{P}^{-1}})^{H}-\mathbf{B}||_{2}^{2}+\lambda||\mathbf{Z}||_{1}+\mathbf{Y}^{T}(\mathbf{U}-\mathbf{Z})$. Equivalently let $(\mathbf{\tilde{u}^{*}},\mathbf{\tilde{z}^{*}},\mathbf{\tilde{y}^{*}})$ be the saddle point for the unaugmented Lagrangian $L_{0}=\frac{1}{2}||\mathbf{\tilde{F}_{P}^{-1}}\mathbf{\tilde{u}}-\mathbf{\tilde{b}}||_{2}^{2}+\lambda||\mathbf{\tilde{z}}||_{1}+\mathbf{\tilde{y}}^{T}(\mathbf{\tilde{u}}-\mathbf{\tilde{z}})$ and consider the first inequality

[TABLE]

where $q^{*}$ is the optimal value to which the objective function converges as $k\to\infty$.

A.3.1 Proof of Inequality A.3

The following inequality holds

[TABLE]

because $(\mathbf{\tilde{u}^{*}},\mathbf{\tilde{z}^{*}},\mathbf{\tilde{y}^{*}})$ is a saddle point for $L_{0}$. Using the constraint of the optimization problem $\mathbf{\tilde{u}^{*}}=\mathbf{\tilde{z}^{*}}$, the left-hand side of inequality A.4 reduces to $q^{*}$. The right-hand side of inequality A.4 can be simplified as

[TABLE]

Thus, we have

[TABLE]

Next, we consider the second inequality

[TABLE]

A.3.2 Proof of Inequality A.5

We know that, by definition, $\mathbf{\tilde{u}^{(k+1)}}$ minimizes $L_{\beta}(\mathbf{\tilde{u}},\mathbf{\tilde{z}^{(k)}},\mathbf{\tilde{y}^{(k)}})$. Since $f(\mathbf{\tilde{u}})=\frac{1}{2}||\mathbf{\tilde{F}_{P}^{-1}}\mathbf{\tilde{u}}-\mathbf{\tilde{b}}||^{2}_{2}$ is closed, proper, and convex by Proposition $\mathbf{4.1}$, it is differentiable, implying that $L_{\beta}$ is subdifferentiable. The necessary and sufficient optimality condition is

[TABLE]

where we use the fact that the subdifferential of the sum of a subdifferential function and a differentiable function with domain $\mathbb{R}^{N}$ is the sum of the subdifferential and the gradient [Rockafellar, 1970].

Since the update of $\mathbf{\tilde{y}}$ is

[TABLE]

we can substitute $\mathbf{\tilde{y}^{(k)}}=\mathbf{\tilde{y}^{(k+1)}}-\beta(\mathbf{\tilde{u}^{(k+1)}}-\mathbf{\tilde{z}^{(k+1)}})$ in A.6 and rearrange the terms to obtain

[TABLE]

This implies that $\mathbf{\tilde{u}^{(k+1)}}$ minimizes

[TABLE]

Analogously we can show that $\mathbf{\tilde{z}^{(k+1)}}$ minimizes $\lambda||\mathbf{\tilde{z}}||_{1}-(\mathbf{\tilde{y}^{(k+1)}})^{T}\mathbf{\tilde{z}}$ and it follows that,

[TABLE]

and

[TABLE]

Adding the inequalities A.7 and A.8, we obtain

[TABLE]

We can rearrange and simplify the terms on the left-hand side of inequality A.9,

[TABLE]

Similarly, we can also simplify the terms on the right-hand side of the inequality A.9,

[TABLE]

We can combine A.10 and A.11 and rearrange the terms to obtain inequality A.5

[TABLE]

Now that we have proved the inequalities A.4 and A.5), only the last inequality

[TABLE]

remains to be proven.

A.3.3 Proof of Inequality A.12

Recall that $(\mathbf{\tilde{u}^{*}},\mathbf{\tilde{z}^{*}},\mathbf{\tilde{y}^{*}})$ is a saddle point for the unaugmented Lagrangian $L_{0}$ and define

[TABLE]

Adding inequalities A.4 and A.5, rearranging the terms, and multiplying by 2 on both sides gives us

[TABLE]

The first term of inequality A.13 can be rewritten as

[TABLE]

and plugging $\mathbf{\tilde{r}^{(k+1)}}=\frac{1}{\beta}(\mathbf{\tilde{y}^{(k+1)}}-\mathbf{\tilde{y}^{(k)}})$ in the first two terms gives

[TABLE]

We can now rewrite the remaining terms of the inequality A.13

[TABLE]

where the term $\beta||\mathbf{\tilde{r}^{(k+1)}}||_{2}^{2}$ is extracted from Equation A.14. Adding and subtracting $\mathbf{\tilde{z}^{(k)}}$ from the last term of A.15 gives

[TABLE]

and by substituting $\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{(k)}}=(\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{*}})-(\mathbf{\tilde{z}^{(k)}}-\mathbf{\tilde{z}^{*}})$ in the last two terms, we get

[TABLE]

This implies that inequality A.13 can be written as

[TABLE]

We can expand $\beta||\mathbf{\tilde{r}^{(k+1)}}+(\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{(k)}})||_{2}^{2}$ as

[TABLE]

Thus, to prove the third inequality A.12, it suffices to show that $2\beta(\mathbf{\tilde{r}^{(k+1)}})^{T}(\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{(k)}})$ is positive. Recall that $\mathbf{\tilde{z}^{(k+1)}}$ minimizes $\lambda||\mathbf{\tilde{z}}||_{1}-(\mathbf{\tilde{y}^{(k+1)}})^{T}\mathbf{\tilde{z}}$ and similarly, $\mathbf{\tilde{z}^{(k)}}$ minimizes $\lambda||\mathbf{\tilde{z}}||_{1}-(\mathbf{\tilde{y}^{(k)}})^{T}\mathbf{\tilde{z}}$ which allows us to add

[TABLE]

and

[TABLE]

to get

[TABLE]

We substitute $\mathbf{\tilde{y}^{(k+1)}}-\mathbf{\tilde{y}^{(k)}}=\beta\mathbf{\tilde{r}^{(k+1)}}$ to get

[TABLE]

because $\beta>0$ and subsequently

[TABLE]

thus proving the third inequality. The convergence of ADMM is a consequence of the three inequalities A.3, A.7, and A.12. Finally, the inequality A.12 can be rewritten as

[TABLE]

which implies that $V^{(k)}$ decreases in each iteration, thus bounding $\mathbf{\tilde{y}^{(k)}}$ and $\mathbf{\tilde{z}^{(k)}}$ as $V^{(k)}<V^{0}$. In particular, note that inductively inequality A.12 leads to the following set of inequalities

[TABLE]

which, upon summation reveals

[TABLE]

Thus the sum of non-negative terms on the right-hand side \sum_{k=0}^{\infty}\Big{(}||\mathbf{\tilde{r}^{(k+1)}}||_{2}^{2}+||(\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{(k)}})||_{2}^{2}\Big{)} converges, as the $n^{th}$ partial sum is bounded above by

[TABLE]

*Thus, the primary residual, $\mathbf{\tilde{r}^{(k)}}\to 0$ and the secondary residual $\beta(\mathbf{\tilde{z}^{(k+1)}}-\mathbf{\tilde{z}^{(k)}})\to 0$ as $k\to\infty$ immediately from $\beta>0$, completing the proof of Proposition 3. *

Appendix B Additional figures and tables

Tables S.2 and S.3 show the median running times for different algorithms namely, PGD, ADMM (Vanilla), and ADMM (GPU) and the percentage change in runtimes relative to the runtime of the PGD algorithm,

[TABLE]

for the HSC and BDS models, respectively. The negative percent change indicates that the ADMM algorithm is faster than the PGD algorithm for the corresponding $N$ value.

Appendix C Details on Generating Functions and PGD Algorithm

C.1 Pseudocode for PGD

The following pseudocode summarizes the proximal gradient descent approach proposed in Xu and Minin [2015].

C.2 Derivation of the PGF for HSC model

The details for deriving the PGF are included here only for completeness, but are standard, following the “random variable technique" of [Bailey, 1991]. Given a two-type branching process with instantaneous rates $a_{i}(k,l)$, we can define the pseudo-generating function for $i=1,2$,

[TABLE]

The probability generating function can be expanded as

[TABLE]

Similarly, starting with one particle of type $2$ instead of type 1, we can obtain an analogous expression for $\phi_{01}(t,s_{1},s_{2})$. For brevity, we will write $\phi_{10}(t,s_{1},s_{2}):=\phi_{1}(t,s_{1},s_{2}),\phi_{01}(t,s_{1},s_{2}):=\phi_{2}(t,s_{1},s_{2})$.

Differentiating $\phi_{1}(t,s_{1},s_{2})$ and $\phi_{2}(t,s_{1},s_{2})$ with respect to time, we get

[TABLE]

We now derive the backward and forward equations with the Chapman-Kolmogorov equations, yielding the following symmetric relations:

[TABLE]

We expand around $t$ and apply Equation A.16 to derive the backward equations:

[TABLE]

We can apply an analogous argument for $\phi_{2}(t,s_{1},s_{2})$ to arrive at the following system of ODEs

[TABLE]

with initial conditions $\phi_{1}(0,s_{1},s_{2})=s_{1}$, $\phi_{2}(0,s_{1},s_{2})=s_{2}$.

Recall from the main text the rates defining the two-component HSC model are given by:

[TABLE]

Thus, the pseudo-generating functions become

[TABLE]

Plugging these into the backward equations yields,

[TABLE]

Notice that the differential equation for $\phi_{2}(t,s_{1},s_{2})$ corresponds to that of a pure death process and has a closed-form solution.

[TABLE]

Plugging in $\phi_{2}(0,s_{1},s_{2})=s_{2}$, we obtain $C=ln(1-s_{2})$, and we get

[TABLE]

On plugging Equation A.18 into the backward equation involving $\phi_{1}(t,s_{1},s_{2})$, we obtain

[TABLE]

Given the rates of the process and the values for the three arguments, Equation A.19 can be solved numerically, allowing the computation of $\phi_{i,j}(t,s_{1},s_{2})=\phi_{1}^{i}(t,s_{1},s_{2})\phi_{2}^{j}(t,s_{1},s_{2})$ which holds by particle independence.