Neural and Mixture CTMCs

• Neural and Mixture CTMCs are nonparametric extensions of classical CTMCs that use neural networks to learn complex, nonlinear transition rate functions.
• They replace fixed parametric models with softplus-activated neural mappings, ensuring nonnegativity while capturing arbitrary state and covariate dependencies.
• Empirical benchmarks show that Neural CTMCs significantly reduce error rates compared to GLM baselines in applications like chemical kinetics, population dynamics, and gene regulation.

Neural and Mixture Continuous-Time Markov Chains (CTMCs) refer to a generalization of classical CTMC frameworks in which the transition rate functions are nonparametrically parameterized, allowing nonlinear dependencies on the current discrete state and covariates. Neural CTMCs employ neural networks to learn the propensity or transition-rate functions directly from fully observed sample paths, with the objective of inferring or accurately reconstructing complex stochastic dynamics underlying systems such as chemical reaction networks, population dynamics, and gene regulatory networks (Reeves et al., 2022).

1. Definition and Mathematical Setup

A CTMC is defined over a (potentially infinite) discrete state space $\Omega$ with piecewise constant sample paths $X(t) \in \Omega$ and transitions occurring at random times. The primary object is the transition rate $\lambda_{ij}(x,c)$, representing the instantaneous rate at which the process jumps from state $i$ to $j$ given covariates $c$, where $x$ collects state and covariate information. All possible transitions are assumed to correspond to a finite set of known reactions $ρ=1,\ldots,r$ with prescribed stoichiometric changes $\Delta S_{ρ}$:

$$\lambda_{i,i+\Delta S_ρ}(x, c) = α_ρ(x, c), \quad \lambda_{ij} = 0 \text{ otherwise}.$$

The total exit rate from state $X(t) \in \Omega$0 is $X(t) \in \Omega$1, and the sojourn time in each state is exponentially distributed with this rate. The process is constructed such that for infinitesimal $X(t) \in \Omega$2,

$X(t) \in \Omega$3

The state vector may comprise multiple species counts $X(t) \in \Omega$4 and continuous covariates $X(t) \in \Omega$5, concatenated as $X(t) \in \Omega$6.

2. Neural Parameterization of Rate Functions

Unlike mass-action or linear log-linear forms, Neural CTMCs do not impose any fixed parametric dependency. Each $X(t) \in \Omega$7 is parameterized as an output of a neural network, enabling arbitrary nonlinear dependence on both the discrete state and covariates. Collectively,

$X(t) \in \Omega$8

and each component is modeled as a neural function with nonlinearity and output constraining ($X(t) \in \Omega$9) to guarantee nonnegativity:

• In mass-action kinetics: $\lambda_{ij}(x,c)$0
• In Neural CTMCs: $\lambda_{ij}(x,c)$1, with $\lambda_{ij}(x,c)$2 a neural network, $\lambda_{ij}(x,c)$3 softplus.

This methodological extension enables the model to represent non-mass-action kinetics, encompassing a broader class of stochastic processes observed in empirical domains (Reeves et al., 2022).

3. Neural Network Architectures and Implementation

Two principal network configurations are described for $\lambda_{ij}(x,c)$4:

• Population Dynamics Example (Lotka–Volterra, $\lambda_{ij}(x,c)$5, no covariates):
  • Input: $\lambda_{ij}(x,c)$6
  • Fully connected neural network, 5 layers of 128 neurons with SELU activation, final output layer of $\lambda_{ij}(x,c)$7 neurons (one per reaction) with softplus.
• Chemical Reaction Network Example (2 species + covariate $\lambda_{ij}(x,c)$8, $\lambda_{ij}(x,c)$9):
  • Input passed to a dense layer (3$i$096), reshaped to (3$i$132) array
  • Batch-wise 1D convolution across rows (10 filters of size 4), SELU activation
  • Flattened and passed through dense layers (290$i$232$i$332), then output layer (32$i$4), softplus.

No additional regularization (e.g., dropout or weight decay) was employed beyond SELU/nonlinearity. Neural weights were set using framework default initialization. These choices ensure that the model has sufficient representational capacity to model complex propensities from observed trajectories.

4. Likelihood-Based Training and Loss Function

Given fully observed trajectories, the negative log-likelihood (NLL) is derived exactly from the CTMC process:

• Let $i$5 denote the dwell time in state $i$6 before the $i$7th reaction $i$8.
• The single-step likelihood is

$i$9

• The trajectory-wise NLL becomes

$j$0

or, regrouped by reaction $j$1,

$j$2

where $j$3 is the number of occurrences of reaction $j$4, $j$5 the pre-reaction states.

Gradient-based optimization (using, e.g., Adam) is then performed on this NLL, requiring only $j$6 cost per iteration, a critical advantage for systems with large or infinite state spaces as no explicit $j$7 matrix is constructed (Reeves et al., 2022).

5. Training Workflow and Computational Aspects

The supervised estimation procedure for N-CTMC can be summarized as:

1. Extract all tuples $j$8 from observed trajectories.
2. Group by reaction $j$9 to build state matrices $c$0 and dwell-time vectors $c$1.
3. Initialize neural parameters $c$2 with default framework initialization.
4. Iterate until convergence:
   • Evaluate the networks $c$3 for all states and reactions.
   • Compute the loss per reaction as $c$4.
   • Aggregate and compute total loss across reactions.
   • Use automatic differentiation to obtain gradients, and update parameters.
   • Monitor loss change or gradient norm for stopping.

Gradients propagate efficiently through all network branches and sum operations. Efficient vectorization by reaction type enables scalability for complex reaction networks.

6. Empirical Benchmarks and Comparative Results

N-CTMCs were empirically evaluated against log-linear generalized linear model (GLM) baselines of the form $c$5 on multiple synthetic systems:

• Birth–Death Process with Covariate: Across $c$6–$c$7 transitions, N-CTMC mean absolute error (N-MAE) matches that of a counting-based MLE baseline, with both outperforming the GLM approach for sufficiently large datasets.
• Population Dynamics (Lotka–Volterra Type, $c$8 reactions): With increasing numbers of trajectories ($c$9, $x$0, $x$1), N-CTMC error (MAE, MSE) decreases with data volume and surpasses log-linear GLM baselines by over an order of magnitude at $x$2 trajectories.
• Chemical Reaction Network with Temperature Covariate: For $x$3–$x$4 trajectories, N-CTMC achieves low MAE/W-MAE and MSE/W-MSE, while GLM baselines exhibit large bias, especially on states rarely visited in training.

Setting	N-CTMC MAE	GLM MAE	N-CTMC MSE	GLM MSE
Birth–Death (500k transitions)	$x$5	—	$x$6	—
Population (10,000 trj., 9 rxn)	$x$7	$x$8	$x$9	$ρ=1,\ldots,r$0
ChemNet (500 trj., 4 rxn)	$ρ=1,\ldots,r$1	$ρ=1,\ldots,r$2	$ρ=1,\ldots,r$3	$ρ=1,\ldots,r$4

In all cases, N-CTMC converges to ground truth at a rate comparable to counting-based MLE (which is only feasible when the state space is not prohibitively large), and demonstrates substantial improvement over log-linear models, particularly in denser state-regions and in settings with nonlinear effects or covariates (Reeves et al., 2022).

7. Significance and Generalization

By parameterizing CTMC transition rates with neural functions, N-CTMCs replace the explicit $ρ=1,\ldots,r$5 rate matrix with a shared, state-dependent mapping $ρ=1,\ldots,r$6. This parameterization enables the learning of arbitrarily complex, nonlinear, and covariate-dependent transition structures. The framework generalizes both mass-action and generalized linear models and applies naturally in birth–death processes, open-loop CTMC control, chemical kinetics with non-mass-action propensities, and temperature- or context-dependent population models. The ability to fit such models directly from fully observed trajectories, with rigorous likelihood-based gradient learning and empirical performance shown across diverse regimes, positions N-CTMCs as a scalable and expressive alternative to previous parametric approaches (Reeves et al., 2022).