A hierarchical dynamical low-rank algorithm for the stochastic description of large reaction networks

TL;DR

The paper tackles the curse of dimensionality in stochastic chemical kinetics by introducing a hierarchical dynamical low-rank method based on binary tree tensor networks (TTNs) to solve the chemical master equation (CME). By recursively partitioning the reaction network and representing the probability distribution with leaves (low-rank factors) and internal connection tensors, the authors develop a projector-splitting TTN (PS-TTN) integrator that evolves the low-rank structure in time, avoiding full state-space storage. Numerical experiments on a 5-species lambda phage model and a 20-species reaction cascade demonstrate dramatic memory reductions and substantial speedups while maintaining high accuracy relative to SSA or full CME references. The approach is noise-free and scalable, with clear avenues for enhancement through rank adaptation, sliding-window truncation, and automated partitioning for large biochemical networks.

Abstract

The stochastic description of chemical reaction networks with the kinetic chemical master equation (CME) is important for studying biological cells, but it suffers from the curse of dimensionality: The amount of data to be stored grows exponentially with the number of chemical species and thus exceeds the capacity of common computational devices for realistic problems. Therefore, time-dependent model order reduction techniques such as the dynamical low-rank approximation are desirable. In this paper we propose a dynamical low-rank algorithm for the kinetic CME using binary tree tensor networks. The dimensionality of the problem is reduced in this approach by hierarchically dividing the reaction network into partitions. Only reactions that cross partitions are subject to an approximation error. We demonstrate by two numerical examples (a 5-dimensional lambda phage model and a 20-dimensional reaction cascade) that the proposed method drastically reduces memory consumption and shows improved computational performance and better accuracy compared to a Monte Carlo method.

Figures (10)

• Figure 1: Population number depending on time for the Schlögl model with $k_0 = 2.5 \cdot 10^{-4}$, $k_1 = 0.18$, $k_2 = 37.5$ and $k_3 = 2200$ in the deterministic (top) and stochastic setting (bottom, left), with the steady state distribution $\Phi(x) = \lim_{t\to \infty}P(t, x)$ (bottom, right). The deterministic solution was calculated from Equation (\ref{['eq:schloegl_deterministic']}) and the single sample $x(t) \sim P(t, x)$ with SSA. Parameters were taken from Erban_2020. Note that the steady ensemble average is $\lim_{t \to \infty} \langle x(t) \rangle \approx 169.46$ for the stochastic approach, whereas the rate equation would predict, depending on the initial condition, either $\lim_{t \to \infty} x(t) = 100$ or $400$.
• Figure 2: (a) Graphical representation of the right-hand side of Equation (\ref{['eq:example-tree']}) as a binary tree. (b) Graphical representation of an internal node $\tau$ (gray) with child nodes $\tau_0$ and $\tau_1$ (blue). The child nodes can either be leaves, where the actual low-rank factors $X^{\tau_0}_{i_{\tau_0}}$ and $X^{\tau_1}_{i_{\tau_1}}$ are stored, or they are also internal nodes, and the low-rank factors are recursively defined by Equation (\ref{['eq:DLR-X']}).
• Figure 3: Dependency graphs (top) and tree structures (bottom) for the three partitions of the lambda phage example. The notation $A\longrightarrow B$ in the dependency graphs indicates that species $A$ depends on species $B$ due to a specific reaction propensity. Small numbers at the leaves of the tree structure denote on which species the corresponding low-rank factors depend.
• Figure 4: Top: Time-dependent 2-norm error for the PS-TTN integrator (left panel) and SSA (right panel) for the lambda phage model. The reference solution was obtained by solving the full CME on the truncated state space with scipy.solve_ivp. For the PS-TTN we used explicit Euler with time step size $\Delta t=10^{-3}$. For convenience we show the best-approximation in the matrix case for a truncation at $r=5$, which was obtained by a truncated SVD of the reference solution (dotted line). Bottom: Wall time comparison (in seconds) for the PS-TTN integrator, SSA and the reference solution ("exact").
• Figure 5: Left: 2-norm error for the PS-TTN solution at time $t=10$ depending on the time step size $\Delta t$ for the lambda phage model. For the solid lines, an explicit Euler scheme was used for the time integration. Right: Maximum mass error over the time integration interval $[0, 10]$ depending on the time step size $\Delta t$. For sake of comparison we show also a solution for the PS-TTN integrator with an implicit Euler scheme and rank $r=(6,6)$ (green, dash-dotted line).