Differentiable Programming of Indexed Chemical Reaction Networks and Reaction-Diffusion Systems

The winner-take-all DNA neural network from Cherry and Qian [11] is a prime example of a CRN with neural-network-like structure, and thus will be used both for illustrating the ICRN formalism here and for the first simulation and training example in Section 5. It can be described by the following indexed reactions.

In an indexed reaction, base species, such as $X,W,XF,\mathrm{\dots }$, and base rate constants, such as ${\alpha }^1,{\alpha }^2,\mathrm{\dots }$ are associated with index symbols, $i,j,k$. Associating specific values to all the index symbols in an indexed reaction describes a concrete reaction. For example, assigning values $i=1$ and $j=2$ in the first indexed reaction gives the concrete reaction $X_1+W_{12}+X{F}_1\stackrel{{\alpha }_{12}^1}{\to }{X}_1+P_{12}$. The enumeration of the first indexed reaction is the collection of concrete reactions formed by substituting every combination of $i,j\in \{1,\mathrm{\dots }n\}\times \{1,\mathrm{\dots },m\}$. The enumeration of the ICRN is the collection of enumerations for each of its indexed reactions. This enumeration is a multiset: the same reaction may appear multiple times, possibly with different rate constants, and in such cases the effective rates constants are summed.

By discussing ICRNs as representations of chemical systems, we have considered meaning, or semantics. Turning our attention to syntax, what are the rules for constructing ICRNs? The reader should pause and consider whether the examples below are valid ICRNs (${\gamma }_i,{\gamma }_k\in {ℝ}_{\ge 0}$).

We will define the ICRN, covering its syntax and its relationship to its semantics. Just as the various CRN conventions can be understood in terms of Feinberg’s formalism, our formalism serves as a common way to precisely describe ICRNs.

An ICRN is specified by five collections: $(𝒰,𝒮,𝒦,ℐ,ℛ)$. The base species $𝒮$ and base rate constants $𝒦$ are multidimensional analogs of species and rate constants from concrete CRNs. Indexed reactions $ℛ$ use index symbols $𝒰$ to describe interactions between the multidimensional objects. The index sets $ℐ$ describe the index values that are valid to associate with base species, base rate constants, or indexed reactions.

Index symbols are essential for the abstraction provided by ICRNs. In the winner-take-all system, the set of index symbols $𝒰$ contains index symbols $i,j,k$. Each index symbol $u\in 𝒰$ has an associated index set, denoted $ℐ(u)$, which is the set of values the index symbol can take on. For example, $ℐ(i)=\{1,\mathrm{\dots },n\}$. A tuple of index symbols has an index set that is the product of the index sets of its constituent index symbols $ℐ(𝐮)=ℐ(u_1)\times \mathrm{\cdots }\times ℐ(u_w)$.

Defining index symbols as objects that exist globally throughout the system, as opposed to each indexed reaction having its own notion of index symbols, has two advantages. Firstly, global index symbols avoid the confusion that arises when the same index symbol appears in multiple reactions but takes on different values across the reactions. Secondly, global index symbols directly represent entities of the system. In the winner-take-all system, there are $n$ inputs and $m$ outputs to a winner-take-all computation. Thus, the index symbol $i$ represents a single input while $j$ and $k$ are used to represent a single output.

Base species and base rate constants are multidimensional objects representing groups of concrete species and concrete rate constants, respectively. Both base species and base rate constants have index sets, denoted $ℐ(s)$ for base species $s\in 𝒮$ and $ℐ(k)$ for base rate constants $k\in 𝒦$. A base species maps from its index set to a concrete species, while a base rate constant maps from its index set to a concrete rate constant. For example, the base species $W$ has an index set $\{1,\mathrm{\dots },n\}\times \{1,\mathrm{\dots },m\}$ and represents a group of concrete species, $W_{11},\mathrm{\dots },W_{mn}$. Base species can be thought of as dictating the overall structure of a group of biomolecules that vary at designated regions, or domains. Each concrete species is a specific variation of the overall structure and specifies the chemical identity of the domains.

Fixing $𝒰$, $𝒮$, and $𝒦$, and their index sets, there is a notion of a set of valid indexed species $V_{𝒮}$ and a set of valid indexed rate constants $V_{𝒦}$. An indexed object is valid if the index set of the base object and index set of the index symbols are the same.

An indexed reaction $r\in ℛ$ is a triple in ${\mathbb{N}}^{V_{𝒮}}\times {\mathbb{N}}^{V_{𝒮}}\times V_{𝒦}$. Reactants and products are multisets of $V_{𝒮}$ and are accompanied by an indexed rate constant from $V_{𝒦}$. The set of index symbols present in the indexed reaction implies the index set of the indexed reaction, denoted $ℐ(r)$. Specifying values for each index symbol specifies a concrete reaction, so indexed reactions are functions from the index set to concrete reactions. The enumeration of an indexed reaction is the image of the index set under $r$, $\text{Im}(r)$. The enumeration of an ICRN is the multiset union of the enumeration of each indexed reaction: ${\cup }_{r\in ℛ}\text{Im}(r)$.

At this point, the reader should be able to identify $\{A\stackrel{𝐵}{\to }C,B\stackrel{}{\to }\mathrm{\varnothing }\}$ as the ill-defined ICRN. The other three ICRNs are valid given that the index sets of their objects are in agreement. Additionally, the reader may have noticed indexed species written as $s_{𝐮}$ rather than $(s,𝐮)$. For instance, we wrote $W_{ij}$ instead of $(W,(i,j))$. The former notation emphasizes the concrete object that results from enumeration. However, indexed bases and indexed rate constants are objects in their own right and are, strictly speaking, the building blocks of indexed reactions. When there is room for confusion, the latter notation will be used to refer to indexed species and indexed rate constants.

ICRNs can be enumerated to capture a wide variety of concrete CRNs. There are systems of interest that are well suited to being described by indexed reactions but with restricted index sets. In many of these cases, the system of interest can be captured by creative use of the base species and base rate constants. For example, in the winner-take-all system, the implementation of the annihilation reaction in the DNA substrate produces a non-functioning annihilator molecule for the case where $j=k$. This restriction in indexing can be addressed by setting ${\alpha }_{jk}^3=0$ when $j=k$, so the concrete reaction does not occur for $j=k$. Symmetric systems in particular are discussed further in appendix A.

The dynamics of an ICRN follows the dynamics of the enumerated CRN. Concentrations of base species and base rate constants are naturally represented as tensors. Accordingly, the dynamics of an ICRN can be expressed with tensor operations, which can be derived through symbolic manipulation of the ICRN.

The tensor representing the concentrations for base species $s$ is written as $[s]\in {ℝ}^s$, where ${ℝ}^s={ℝ}_{\ge 0}^{|J_1|\times \mathrm{\cdots }\times |J_{N_s}|}$ with $ℐ(s)=J_1\times \mathrm{\cdots }\times J_{N_s}$. To distinguish between the concentrations associated with the ICRN and the concrete CRN that it abstracts, we will use ${[s]}_{𝐮}$ to refer to concentrations in the ICRN representation and $[s_{𝐮}]$ to refer to the concentration of the concrete species, $s_{𝐮}$, where $ℐ(𝐮)=ℐ(s)$. The time-dependent state of an ICRN is $x=([s_1],\mathrm{\dots },[s_{|𝒮|}])\in {ℝ}^{𝒮}$ where ${ℝ}^{𝒮}={ℝ}^{s_1}\times \mathrm{\cdots }\times {ℝ}^{s_{|𝒮|}}$. Thus, $x$ specifies $[s]$ for all $s\in 𝒮$. The dynamics of an ICRN describes the time derivative of $x$, which is the tuple of time derivatives for each $[s]$: $\frac{dx}{dt}=(\frac{d[s_1]}{dt},\mathrm{\dots },\frac{d[s_{|𝒮|}]}{dt})$

The dynamics for the ICRN follows the dynamics of the enumerated concrete CRN, so the time derivative for $[s]$ is the collection of time derivatives for each concrete species represented by $s$, where the time derivative of ${[s]}_{𝐮}$ is defined by the time derivative of $[s_{𝐮}]$ from the enumerated concrete CRN: $\frac{d{[s]}_{𝐮}}{dt}:=\frac{d[s_{𝐮}]}{dt}$

As seen in Fig 2, each indexed reaction $r$ has an associated indexed reaction flux $\mathrm{\Phi }(r)$, a real-valued tensor. For each appearance of an indexed species $(s,𝐮)$, sums are taken along the dimensions of $\mathrm{\Phi }(r)$ that correspond to index symbols that are present in the indexed reaction but not present in the indexed species. The result is the flux contribution to $s$, denoted ${\delta }_r(s,𝐮)$, which is a real-valued tensor with the same shape as $[s]$. A base species $s$ receives contributions from every appearance of any indexed species that derives from $s$ among all the indexed reactions.

The indexed reaction flux, $\mathrm{\Phi }(r)$ is a real-valued tensor with the same shape as the index set of $r$. Let $𝐩$ be index symbols such that $ℐ(𝐩)=ℐ(r)$ and $a_r,b_r\in {\mathbb{N}}^{V_{𝒮}}$ be the multisets of reactants and products, respectively. Then $a_r(y,𝐮),b_r(y,𝐮)\in \mathbb{N}$, where $(y,𝐮)$ is an indexed species. Additionally, let $(k_r,𝐪)\in V_{𝒦}$ be the indexed rate constant used in reaction $r$. Here, $k_{r,𝐪}\in {ℝ}_{\ge 0}$ is the real-valued rate constant. The following is an element-wise specification of $\mathrm{\Phi }(r)$, indexing with $𝐩$, since all index symbols in $𝐩$ are present in the right hand expression.

The flux contribution to $s$ from $r$ depends on which indexed species $(s,𝐮)$ is participating in the reaction. The following equation is an element-wise specification of the tensor ${\delta }_r(s,𝐮)$ using the index symbols $𝐮$.

On the right hand side, $𝐩$ is the tuple of index symbols in $r$ and $𝐩-𝐮$ refers to the index symbols that are in $𝐩$ but not $𝐮$. Thus, the right hand side is an expression that is parameterized by $𝐮$. Moreover, the summation is an Einstein summation, a kind of indexed summation which can be efficiently computed [58].

The totality of the contributions to $s$ from each appearance of a derived indexed species can be formulated as the following element-wise specification.

The left hand side describes the rate of change of an element of $[s]$ corresponding to index symbols $𝐯$. On the right, ${\delta }_r{(s,𝐮)}_{𝐯}$ refers to an element of the tensor ${\delta }_r(s,𝐮)$ indexed by $𝐯$. Since $\left(b_r(y,𝐮)-a_r(y,𝐮)\right)\in \mathbb{Z}$, the right hand side is an expression parameterized by $𝐯$. An example of the dynamics derivation appears in appendix B.

ICRNs can be extended naturally from the well-mixed to a spatial setting, becoming reaction-diffusion systems. For classical CRNs, reaction-diffusion dynamics are described mathematically by a PDE: ${\partial }_t𝑿=𝑫{\nabla }^2𝑿+R(𝑿)$. For reaction-diffusion with species $𝒮$ in $d$ spatial dimensions, the concentration field $𝑿$ is a function $𝑿:\mathrm{\Omega }\subset {ℝ}^d\to {ℝ}^{𝒮}$, with $\mathrm{\Omega }$ a bounded subset of the spatial dimension. $𝑫$ is a diagonal matrix of diffusion constants, ${\nabla }^2$ is the Laplacian operator encoding diffusion, and $R$ represents the derivatives due to the CRN’s reactions. ICRNs require only a tensor generalization of these semantics.

The abstract winner-take-all computation can classify hand written digits. Following Cherry and Qian [11], our dataset is a reduced-resolution subset of the MNIST [18] dataset and consists of input-output pairs $({𝐱}_d,{𝐲}_d)$ where $d$ indicates a specific pair, ${𝐱}_d\in {\{0,1\}}^{100}$ specifies a $10$ by $10$ image of a handwritten digit, and ${𝐲}_d\in {\{0,1\}}^3$ specifies the digit that the image represents (seven, eight, or nine). A digit classifier is a function from images to classes. The abstract winner-take-all computation acts as a classifier by computing weighted sums for each class and taking the class with the highest sum to be the output class. The weights are specified by a matrix $\widehat{𝐖}\in {ℝ}_{\ge 0}^{100\times 3}$, where a column contains the weights used in the weighted sum per class.

Cherry and Qian set the weights of the classifier by pruning an average-digit matrix. Average-digit weights are the averages of examples in the training data. A column of the weight matrix is populated by the average of $𝐱$ from the digit class corresponding to the column. Due to experimental limitations from having many nonzero elements, the pruned matrix takes the $20$ highest elements of each class.

The winner-take-all ICRN can be trained to optimize different qualities. One view is that the fluorescence signal for the correct class should completely saturate (at 100 nM) while for other classes it is $0$. A large difference in the output fluorescence is important experimentally for being able to resolve the output signal from the classifier. The superiority score $\text{SUP}=[F]\cdot (2y-\mathrm{𝟏})$ captures this idea, for a given target input-output pair $(x,y)$ that results in endpoint fluorescence vector $[F]$. Ideal behavior gives a superiority metric of $100$, and lower values reflect deviation from the ideal (Fig 4B). The loss function can be set to produce a weight matrix with a high superiority metric. Another possibility is to optimize the ICRN’s ability to produce the correct output without a concern for how low the fluorescence is for the other output classes. The endpoint fluorescence can be interpreted as a probability by normalizing with the softmax function, $p_j=\exp ({[F]}_j)/{\sum }_k\exp ({[F]}_k)$. The loss function computes a cross entropy loss $-{\sum }_j{y}_j\log (p_j)$ with an $L_{1/2}$ regularization term, ${\sum }_{i,j}{\widehat{𝐖}}_{ij}^{1/2}$, that penalizes high weights, which encourages the training process to produce fewer nonzero weights, increasing experimental feasibility. The results of applying these training methods are shown in Fig 4D.

We will begin with the Gray-Scott model [25] to illustrate the simulation and training of reaction-diffusion systems within our software framework. The Gray-Scott model is a reaction-diffusion system with just two species, $U$ and $V$, and four reactions: $U+2V\stackrel{1}{⟶}3V$, $V\stackrel{F+k}{⟶}0$, $U\stackrel{𝐹}{⟶}0$, and $0\stackrel{𝐹}{⟶}U$. As such, it does not stand to gain simplicity from indexing, but nonetheless it is a (trivial) instance of an ICRN. The PDE can be written as:

This reaction-diffusion system has just four parameters: rate constants $F$ and $k$, and diffusion constants $D_U$ and $D_V$. However, tuning these parameters, the system generates a wide variety of patterns, both stable and time-varying, explored and categorized by Pearson [49].

Our task is to use differentiable programming [16] to find Gray-Scott parameters that yield a particular type of pattern – e.g. mazes, spots, waves. As in Pearson’s paper, we will fix the diffusion constants (with $U$ diffusing twice as fast as $V$), and vary only $F$ and $k$. We set up the task as follows: first, we choose $F$ and $k$ to yield a certain type of pattern when simulated forwards from initial conditions randomly sampled from a distribution (see Fig 5). The Gray-Scott PDE is fairly stiff, necessitating many steps with small $dt$ to let its characteristic patterns develop from unstructured initial conditions. We thus take a single long simulation output using these parameters, and call this our target pattern. During training, we will run shorter simulations that are more amenable to gradient descent, but which will introduce other challenges.

Next, we need to define a loss function measuring how closely simulations during training match the target. We would like the loss function to be relatively stable to different initial conditions, so we cannot use a simple pixel-by-pixel comparison of simulation output and target: rather, we want something that captures and compares the qualitative features of the patterns. Taking inspiration from work on neural transfer of artistic style [23] and matching textures with self-organizing neural cellular automata [43], we evaluate pattern matches using the low-level (early layer) features from VGG-16, a convolutional neural network trained for image classification [57]. We feed our target pattern to VGG-16 as an input, and extract the activations from the first through third max-pooling layers. We then compute Gram matrices from these features, capturing some sense of how different elements of the pattern vary with each other over space. The Gram matrix $G^l$ element $G_{i,j}^l={\sum }_{x,y}{f}_{i,x,y}^l{f}_{j,x,y}^l$, where $x$ and $y$ are the spatial dimensions of the layer and $f_{i,x,y}^l$ is feature $i$ at position $(x,y)$ in layer $l$.

Finally, we can compute a stylistic loss function, defined as the mean squared distance between the target Gram matrices and those extracted from the simulated output by the same process: ${ℒ}_{\text{texture}}=\frac{1}{\mid l\times i\times j\mid }{\sum }_{l,i,j}{(G_{i,j}^{l,\text{ target}}-G_{i,j}^{l,\text{ output}})}^2$. For stability, we augment the loss by adding a “moments-matching” term, comparing the means and log-variances of the simulation outputs with reference values, drawn by simulating the Gray-Scott model in different parts of the pattern-forming parameter space:

Here, $U$ is the simulated output, and $\widehat{U}$ is an average over ten reference simulation outputs, and ${\sigma }^2(X)$ is the variance of $X$. Finally, we add a “potential well” regularization term, penalizing parameter values outside of the pattern-forming region of $F\in [0,0.08],k\in [0.03,0.08].$ We then have loss $ℒ={ℒ}_{\text{texture}}+{\alpha }_1{ℒ}_{\text{moments}}+{\alpha }_2{ℒ}_{\text{regularization}}$.

We train the parameters using gradient descent on this loss function applied to the output of short simulations that start with the target pattern rather than a random noise state. Thus we are accommodating the stiffness of the Gray-Scott PDE by training for stability rather than for de novo generation of the texture. This has an additional benefit in that for some parameters, pattern formation in the Gray-Scott system can fail to “kick start” properly for random initial conditions. As in the winner-take-all example, both the simulation (using icrn) and loss function computation are written using differentiable JAX primitives, meaning that the derivatives may be calculated automatically using backpropagation-through-time (BPTT) [2, 53]. We integrate the gradients using mini-batch gradient descent, adding a small amount of randomness to the parameters when the learning converges to a local minimum.

The results shown in Fig 5, including a final long simulation from the lowest loss endpoint among the training runs, are perhaps deceptively well-behaved: training of this system was finicky and challenging. The vast bulk of parameter combinations give rise to spatially homogeneous outputs, and only in a sliver of parameter space are interesting patterns to be found – and there, the sensitivity to initial state was often confounding. Making this problem even harder, learning can turn just two knobs, $F$ and $k$. The success of modern machine learning in many contexts is thanks to, among other things, the high dimensionality of deep neural networks. With so many parameters, there is almost always some direction in parameter space we can move along to reduce the loss, meaning that even simple, first-order optimization algorithms like gradient descent perform well. In the next section we will consider an ICRN with tens of thousands of parameters in the reaction-diffusion setting. Though small by modern standards, it is big enough to enjoy the blessings of dimensionality.

Reaction-diffusion systems can be trained to grow into significantly more complex patterns. Here, we take the Hopfield network-inspired CRN and training approach used in [8], and train an ICRN to grow from simple initial conditions into a rendering of a photograph.

The Hopfield CRN consists of a set of neurons $𝒙$, with variables $x_i\in ℝ$, and $i,j\in \{1,\mathrm{\dots },N\}$ for an $N$-neuron system. In the well-mixed setting, the Hopfield CRN can be trained to store a set of “memories”, similar to the classical Hopfield network [29]. The neurons are connected all-to-all. Catalytic reactions between pairs of neurons correspond to weighted sums, and a trimolecular degradation reaction for each neuron corresponds to an activation function. To allow negative values for the variables, and negative weights between neurons, we use a dual-rail representation: the variable $x_i$ is the difference in the concentrations of two species, $x_i=x_i^{+}-x_i^{-}$, with $x_i^{+},x_i^{-}\in {ℝ}_{\ge 0}$. A fast annihilation reaction is added between the positive and negative species for each variable. The ICRN is defined as

where $\text{ReLU}(x)=\text{max}(0,x)$, ensuring that all rate constants are non-negative. The parameters for the Hopfield ICRN consist of two $N\times N$ weight matrices $W^{+}$ and $W^{-}$ (with $W_{ij}^{\pm }\in ℝ$), two $N$-vectors of degradation rate constants $k_{\text{deg}}^{+}$ and $k_{\text{deg}}^{-}$ (with $k_{\text{deg},i}^{\pm }\in {ℝ}_{\ge 0}$), and two $N$-vectors of diffusion constants $D^{+}$ and $D^{-}$ (with $D_i^{\pm }\in {ℝ}_{\ge 0}$). When the weight matrices are symmetric and degradation rate constants are equal, the Hopfield ICRN has equivalent attractor-basin structure as the classical Hopfield associative memory, but we find that allowing asymmetric and unequal parameters facilitates improved pattern formation [8]. The ICRN software admits an efficient representation of dual-rail ICRNs.

The training approach and loss function for the pattern forming task was inspired by Mordvintsev’s work on neural cellular automata [45]. We define the first three $x$ variables ($i\in \{0,1,2\}$) to be the visible units, corresponding to red, green, and blue, and the remaining $N-3$ are hidden. We pick an RGB target image $𝑻$, and train the parameters of the ICRN such that the visible state of the system matches the target as closely as possible, when it is grown from specified initial conditions for a specific duration of time, $t_{\text{grow}}$. The target in Fig 6 is grayscale, but we treat it in RGB for generality. We can write the pixel-wise loss as

where $h\in \{1,\mathrm{\dots },H\}$ and $w\in \{1,\mathrm{\dots },W\}$ are the discretized points in space. The loss function considers only the visible species. This loss is simpler than that used in the Gray-Scott example, but is also less flexible – a pattern which to the eye appears similar in style may have a high loss, if the pixel values are different. As in the Gray-Scott training task, we train the rate constants, but for this task we also train the diffusion constants. Unlike in the winner-take-all example, the initial conditions are fixed, consisting of a simple asymmetric pattern in the visible units. We use the Optax implementation of the Adam optimizer [37, 17].

Fig 6 shows the training of a Hopfield ICRN with $N=256$ neurons. The corresponding CRN has 512 species (two per dual-rail variable), 131072 ($=2N^2$) catalytic rate constants, 512 ($=2N$) degradation rate constants, and 512 ($=2N$) diffusion constants, all of which may be varied during training. This amounts to considerably more real values than the ${10}^4$ grayscale pixels of the target image, helping explain the accuracy of pattern formation. In general, we observe that increasing $N$ allows the training procedure to find a set of parameters with a lower loss.

The Hopfield CRN has several other notable differences from the Gray-Scott model. First, it is much less sensitive to how the parameters are initialized, and the choice of target pattern. Most training runs are at least somewhat successful, matching the target to some level of fidelity. Second, the Hopfield reaction-diffusion system is less stiff, and so can grow into the pattern in fewer steps without encountering numerical errors. With so many species, the state of the system occupies significant memory, and so we use gradient checkpointing to reduce the memory overheads from backpropagation-through-time [26]. Finally, note that the target image is a correct snapshot at the specified time, but is not stable – more sophisticated training procedures are discussed in [8] that improve the stability and robustness of pattern formation in Hopfield CRNs.