In-context learning emerges in chemical reaction networks without attention

TL;DR

This work shows theoretically and numerically that chemical processes can achieve in-context learning through a mechanism the authors call subspace projection, in which the entire input vector is mapped onto comparison subspaces, with the dominant projection determining the computational output.

Abstract

We investigate whether chemical processes can perform in-context learning (ICL), a mode of computation typically associated with transformer architectures. ICL allows a system to infer task-specific rules from a sequence of examples without relying solely on fixed parameters. Traditional ICL relies on a pairwise attention mechanism which is not obviously implementable in chemical systems. However, we show theoretically and numerically that chemical processes can achieve ICL through a mechanism we call subspace projection, in which the entire input vector is mapped onto comparison subspaces, with the dominant projection determining the computational output. We illustrate this mechanism analytically in small chemical systems and show numerically that performance is robust to input encoding and dynamical choices, with the number of tunable degrees of freedom in the input encoding as a key limitation. Our results provide a blueprint for realizing ICL in chemical or other physical media and suggest new directions for designing adaptive synthetic chemical systems and understanding possible biological computation in cells.

Figures (9)

• Figure 1: Schematic illustration of ICL by a chemical reaction network. A) Definition of the ICL task as matching a query element to an example within the context vector. These context elements $\mathbf{z}_i$ are mapped into label predictions $q_i$ by the chemical reaction network. B) The standard mechanism for ICL uses a pairwise dot-product attention (involving learnable key $\mathbf{W}_K$ and query $\mathbf{W}_Q$ matrices) between the query element and the members of the context. C) Schematic illustration of the subspace projection mechanism, which involves projections of the context vector $\mathbf{z}$ onto different vectors which determine the steady-state output of the chemical reaction network (such as $\mathbf{A}$). Training the network involves placing these vectors so that steady-state outputs can be decoded to solve the ICL task. This is accomplished by placing these learnable vectors near certain subspaces which correspond to pairwise attention comparisons with the query component.
• Figure 2: ICL performance across training conditions. A) Plots of the training and validation set accuracy as well as accuracy on novel classes (termed ICL accuracy) as training progress. Five runs and their average are shown for an example with context element dimension $D = 8$; the default parameters (see Supplementary Material Section \ref{['SIsec:params']}) are used throughout the paper unless otherwise specified. B) Plots of the different accuracy measures after training networks across different conditions, with the mean and standard deviation over five runs shown for each condition. In the top plot, the number of training samples $N_\text{samp}$ and $D$ are varied together, while in the bottom the number of GMM classes $N_\text{class}$ and within-class variation $\epsilon$ are varied. In the bottom plot we set $N_\text{samp} = 250$ to show the behavior with low training data. C) Heat map of the average ICL accuracy over five runs as the context length $N_\text{c}$ and number of species $N_\text{n}$ are varied, using $N_\text{samp} = 2.5 \times 10^5$. The dashed gray line is used to visualize the values $N_\text{n} = N_\text{c}$. The bottom panel shows along this diagonal the mean and standard deviation as well as the maximum of the ICL accuracy over five runs. Note that classification with $N_\text{c} =1$ is trivial.
• Figure 3: Demonstration of ICL mechanism in small networks. A) For a two-species network, the ICL accuracy as training progresses ($N_{\mathrm{c}} = 2$, $D = 1$). Five trials are shown for each condition. In the "one branch" condition, networks are trained only on context vectors $\mathbf{z} = (z_1, z_2, z_q)$ satisfying $z_q = \max(z_1, z_2)$, corresponding to the subspaces $\mathcal{M}_{1>}$ and $\mathcal{M}_{2>}$. In the "both branches" condition, no restriction is applied and contexts are drawn from all subspaces. B) Depiction of the context space of $\mathbf{z}$. Scatter points show context vectors sampled during training, colored according to the subspace branch they lie on. The linear subspaces $\mathcal{M}_{1}$ and $\mathcal{M}_{2}$ are shown as planes, with their intersection indicated by a solid black line. The learned vectors $\mathbf{A}$ and $\mathbf{B}$ that determine the steady state are shown as red and blue arrows (see main text for details); their lengths are not drawn to scale. C) Scatter plot of the steady-state concentration of species $B$ (here $C_{\text{tot}} = 1$) for context vectors sampled during training with one branch and with both branches. Small random vertical offsets are applied for visualization. D) Same as panel A, but for a three-species network. E) Same as panel B, but showing the additional learned vectors for the three-species network. See main text and Supplementary Material Section \ref{['SIsec:threespecies']} for details. F) Ternary plot showing the steady-state distribution over $\bar{\mathbf{C}}$ with $C_{\text{tot}} = 1$ for context vectors sampled during training with both branches.
• Figure 4: Sparsifying the input encoding across chemical models. A) The encoding vectors $\mathbf{K}_r$, shown horizontally for the 20 reactions in a trained fully connected linear reaction network with $N_\text{n} = 5$ species and context vectors with $N_\text{c} = 4$ and $D = 4$. White vertical lines separate the elements of the context vector. The corresponding encoding vectors trained after applying the sparsity masks $\rho_\text{edge}$ and $\rho_\text{all}$ are shown below. B) ICL accuracy for linear reaction networks of different sizes as $\rho_\text{all}$ is varied. For each value of $\rho_\text{all}$ and $N_\text{n}$, we run 30 trials with different random seeds controlling both sparsity-mask sampling and training stochasticity. We compute the number of degrees of freedom in the encoding vectors $\mathbf{K}_r$ and subtract $N_\text{req} = 2N_\text{c}(N_\text{c}+1)D$ to obtain the predicted difference $\Delta_\text{dof}$, which, when positive, indicates sufficient degrees of freedom to cover all ICL subspace branches. Running averages (solid lines), standard deviations (shaded regions), and maxima (dash–dot lines) are computed using windows of size 50. The inset shows data obtained using both $\rho_\text{all}$ and $\rho_\text{edge}$, plotting ICL accuracy directly against the sparsity parameter rather than against the sampled number of degrees of freedom. Here the mean and standard deviation over the 30 runs for each parameter value are shown. C,D) Same as the main plot of panel B, but for the autocatalytic and winner-take-all network models, respectively. Results using the $\rho_\text{edge}$ sparsity mask are shown in Supplementary Material Section \ref{['SIsec:sparsity']}.
• Figure 5: Heat map of the average ICL accuracy over five runs as the context length $N_\text{c}$ and number of species $N_\text{n}$ are varied, using different rate-encoding functions $\sigma$. Here Softplus refers to $\log\left(1 + \exp(x) \right)$ and Sigmoid refers to $10 / \left(1 + \exp(-x/10)\right)$.