C ONNECTOME - CONSTRAINED LATENT VARIABLE MODELS OF WHOLE - BRAIN NEURAL ACTIVITY
A BSTRACT
The availability of both anatomical connectivity and brain-wide neural activity measurements in C. elegans make the worm a promising system for learning detailed, mechanistic models of an entire nervous system in a data-driven way.
However, one faces several challenges when constructing such a model.
We often do not have direct experimental access to important modeling details such as single-neuron dynamics and the signs and strengths of the synaptic connectivity.
Further, neural activity can only be measured in a subset of neurons, often indirectly via calcium imaging, and signiﬁcant trial-to-trial variability has been observed.
To address these challenges, we introduce a connectome-constrained latent variable model (CC-LVM) of the unobserved voltage dynamics of the entire C. elegans nervous system and the observed calcium signals.
We used the framework of variational autoencoders to ﬁt parameters of the mechanistic simulation constituting the generative model of the LVM to calcium imaging observations.
A variational approximate posterior distribution over latent voltage traces for all neurons is efﬁciently inferred using an inference network, and constrained by a prior distribution given by the biophysical simulation of neural dynamics.
We applied this model to an experimental whole-brain dataset, and found that connectomic constraints enable our LVM to predict the activity of neurons whose activity were withheld signiﬁcantly better than models unconstrained by a connectome.
We explored models with different degrees of biophysical detail, and found that models with realistic conductancebased synapses provide markedly better predictions than current-based synapses for this system.
I NTRODUCTION
The anatomical connectivity of the entire C. elegans nervous system, including both chemical and electrical synapses, has been known for several decades [26; 24; 27].
However, well-calibrated and predictive connectome-constrained mechanistic models of this nervous system have yet to be demonstrated [12; 25; 9; 6].
This is partially because currently available experimental whole-brain data are insufﬁcient to completely constrain computational models.
First, the single-neuron and synapse dynamics are generally unknown.
Second, the connectome does not directly inform the signs and strengths of individual synapses.
Third, the response properties of sensory neurons are incompletely known.
Further, it is unclear what level of biophysical detail is necessary to reproduce the essential computational function of the C. elegans nervous system.
Here, we used whole-brain calcium imaging data to constrain the missing parameters in a connectomeconstrained, biophysically detailed model of the C. elegans nervous system.
We started with a simpliﬁed non-spiking passive point-neuron model of the voltage dynamics of individual neurons in the circuit.
We modeled inputs to the neurons from both electrical and chemical synapses. The chemical synapses are modeled nonlinearly, with either current-based or conductance-based biophysics.
The challenge of ﬁtting such a model to data is two-fold.
First, the voltage dynamics of the neurons are not directly observed, but rather indirectly measured via their slow calcium dynamics.
Second, there is signiﬁcant trial-to-trial variability in neural activity, suggesting strong initial state dependence in the neural responses to the same sensory stimulus [8].
These issues can both be addressed by treating the collective voltage signals of all neurons in the nervous system as an unobserved latent variable whose dynamics are determined by the simpliﬁed connectome-constrained biophysical model with unknown neuronal and synaptic parameters.
Our connectome-constrained latent variable model (CC-LVM) of voltage dynamics of the C. elegans nervous system is a large-scale latent variable model with a very high-dimensional latent space consisting of voltage dynamics of 300 neurons over 5 minutes of time.
The generative model for these latent variables is described by stochastic differential equations modeling the nonlinear dynamics of the network activity, and a novel differential equation describing how the calcium signals signals are generated from the voltage dynamics of individual non-spiking neurons.
We developed a variational autoencoder-based framework for inferring the unobserved voltage dynamics from the observed calcium dynamics of only a subset of the neurons in the nervous system.
The C. elegans nervous system consists of 300 neurons, divided into 118 distinct classes (often bilaterally symmetric pairs) [23; 13; 22].
These neurons can be categorized into sensory neurons, interneurons, and motor neurons.
The majority of these neurons, about 200, are concentrated in the head of the animal.
The synaptic connectivity of C. elegans has been mapped with electron microscopy, providing researchers with a complete connectome [26; 27].
We applied the CC-LVM to whole-brain calcium imaging data which captured 170 of the 300 neurons in multiple worms as they responded to chemosensory stimuli [28].
In principle, an accurate model of the nervous system constrained by incomplete activity measurements but complete connectivity measurements can enable accurate predictions of activity in neurons which were not recorded.
We tested this hypothesis by using the CC-LVM to predict the activity of neurons which were measured but whose activity was withheld during model training.
We also used the model to predict the activity of entire worms, by holding out single trials during training.
The CC-LVM predicted activity of withheld neurons and withheld worms signiﬁcantly better than models unconstrained by the connectome, demonstrating the utility of the connectome even when little is known about the signs and strengths of individual connections.
Further, we found that models with conductance-based synapses provide superior predictions to models with current-based synapses.
This is surprising because while conductance-based synapses are more mechanistically accurate, they do not necessarily support efﬁcient inference [7].
CC-LVMs thus provide a new tool for connectome and activity constrained modeling of neural circuits and for discovering the appropriate level of detail of biophysical model for a given system.
This work provides a framework not only for modeling the C. elegans nervous system in particular, but also the neuronal networks of other biological systems.
1. P RIOR W ORK
Previous work has focused on creating network simulations based on anatomical connectome data with unknown single neuron synaptic biophysics [16; 6; 25; 17; 12].
These network models provide a holistic view of the entire nervous system and were validated by comparing simulated locomotion with movements of live animals [10].
However, they were not able ﬁt against prerecorded calcium ﬂuorescence data, so their simulated neuronal activities were not conﬁrmed.
Work has also been done on modeling unlabeled neuron populations but these models cannot predict activity at a single neuron resolution because the voltage data is not mapped to speciﬁc cells [4].Recent advances in machine learning have enabled increasingly sophisticated latent variable models that uncover structure from data generated by advanced neural interfacing technologies [19].
This framework has been used in other contexts to infer neuronal voltage dynamics from high-dimensional calcium ﬂuorescence recordings [18; 1; 21].
Despite their utility, these models do not account for connectomic data.
Our work employs both LVMs and connectomic constraints to create a model that is informed both by neural population dynamics and anatomical network structure.
Recently, Bayesian models have been widely applied in neuronal datasets to infer properties such as spiking activity, neural dynamics, and connectivity [25; 1; 21].
In particular, Linderman et al.
[15] created a hierarchical state space model of C. elegans neural activity. However, because such a model [Figure] [Figure]
[Figure] is not mechanistic, it cannot predict the activity of neurons that were never measured. Warrington et al. [25] used sequential Monte Carlo (SMC) to impute the intracellular voltage potentials of C. elegans neurons from 49 recorded calcium traces [11].
Their work combined neuronal, body and calcium observation simulators in order to model locomotion [3].
Their simulator produced a series of exemplar neuronal voltage traces but they did not test their model’s ability to predict the activity of unrecorded neurons, so the accuracy of their inferences remain unknown.
Our CC-LVM model uses a variational auto-encoder (VAE) to perform inference instead of SMC which is more computationally efﬁcient because it allows direct sampling from the approximate posterior.
The CC-LVM was trained on a large number of neurons and we validated our model by evaluating its ability to predict the activity of neurons held out from the training data.
Finally, we searched a space of generative models and optimized voltage predictions by comparing them to real calcium activity data.
C ONNECTOME - CONSTRAINED LATENT VARIABLE MODEL
We constructed a connectome-constrained latent variable model (CC-LVM) of the C. elegans nervous system, where each node in our network represents a speciﬁc neuron in the animal.
The activity of each neuron was modeled with a latent variable analogous to voltage.
Since the neurons in C. elegans do not spike [2], the dynamics of each neuron in the network is represented with a stochastic non-spiking passive leaky integrator equation with learned time constants and resting membrane potentials [5].
The neurons were coupled by both chemical and electrical synapses with learned weights [5].
We allowed these weights to be non-zero only where the connectome indicates the existence of synapses.
Given a set of learned parameters (weights, time constants, and resting membrane potentials), the dynamics of the network deﬁnes a prior distribution over neural activity trajectories.
Importantly, the stochastic nature of the neuron voltages causes deviations from perfectly deterministic dynamics, which allows us to model the observed variability in single-neuron dynamics.
This variability has several potential sources, including the unmeasured initial states of the neurons and our incomplete knowledge of the sensory inputs driving the nervous system.
A latent variable model of this scale with a nonlinear generative model deﬁned by the stochastic dynamics is difﬁcult to ﬁt because of the high number of parameters.
To address this challenge, we used the probabilistic inference framework of variational autoencoders (VAE)
[20] to train a black-box voltage inference network to predict a posterior distribution over neural activity trajectories.
We then used this inference network to train the parameters of the CC-LVM.
The resulting LVM has a biologically realistic generative model of the nonlinear neural dynamics of the C. elegans nervous system, and a black-box temporal convolutional inference network which, given sensory stimulus and calcium imaging measurements, predicts a factorized Gaussian distribution over the voltages of all the neurons in the network.
Several variants of the LVM were developed: we tested models in which the synaptic connections were modeled as either current-based or conductance-based [5], and evaluated different types of connectome constraint.
The LVM was optimized with the ELBO (evidence lower bound) objective.
2. N ETWORK MODEL WITH PASSIVE POINT - NEURON VOLTAGE DYNAMICS
Neurons in the C. elegans nervous system are largely non-spiking [2], so we model the voltage dynamics for these neurons as a passive point neurons with a single electrical compartment.
Let v ∈ R N denote the voltages of the N neurons.
The voltage v i ( t ) for each post-synaptic neuron i at the time t was calculated using a ﬁrst-order leaky integrator equation given by
[Equation]
where τ i is the voltage time constant,
o i is the chemosensory input provided to only the sensory neurons, s ci is the chemical synaptic input, s ei is the electrical synaptic input, v resti is the resting neuron voltage.
We studied two variations of the model, the current-based model and the conductance-based model , which differ in their formulations of chemical synaptic input s ci .
Since neurons in the C. elegans nervous system are largely non-spiking, we model the chemical synapses as having graded release of neurotransmitter, rather than the all-or-none quantal release seen in spiking neurons.
In both models, we model the amount of neurotransmitter released, W cji g ( v j ( t )) , in proportion to the pre-synaptic voltage v j followed by a softplus activation g ( · ) which sets a minimum voltage below which there is no synaptic release.
We use g ( · ) , to denote a softplus function for the rest of the paper.
In our current-based model, synaptic input s ci to a post-synaptic neuron i is directly proportional pre-synaptic neurotransmitter concentration:
[Equation]
where W cji represents the chemical synaptic weight between pre-synaptic neuron j and post-synaptic neuron i .
W cji can be positive or negative depending on if the synaptic connection is excitatory or inhibitory.
If neurons j and i are not connected, W cji is set to zero.
In the conductance-based model, we model the synaptic current entering the post-synaptic neuron with more biophysical detail as
[Equation]
Here, the pre-synaptic neurotransmitter concentration W cji g ( v j ( t )) is more accurately modeled as proportional to the conductance at the post-synaptic terminal.
The post-synaptic current is then given by the product of the synaptic conductance and the difference between the post-synaptic voltage v i ( t ) and the synaptic reversal potential E ji .
In contrast to the current-based synapse whose input is independent of the post-synaptic voltage v i ( t ) , the more accurate conductance-based synapse model also has a dependence on the post-synaptic voltage.
Additionally, this model decouples the sign of the synapse from the strength.
The reversal potential of a synapses E ji dictates whether a synapse is excitatory or inhibitory.
A large and positive E ji corresponds to an excitatory synapse causing depolarization of the post-synaptic neuron.
Conversely, an inhibitory synapse will have a negative E ji causing hyperpolarization.
In this model, we can now independently train the sign of a synapse and its non-negative strength W c ji , which is not easily possible with current-based synapses.
In both the current-based and conductance-based models, the following equation was used to represent electrical synaptic inputs:
[Equation]
where W eji is restricted to be non-negative and v j − v i is the potential difference between presynaptic and postsynaptic neurons.
We also restrict W eji
= W eij because the potential differences between electrical synapses are symmetric.
To directly compare the outputs of the LVM to neural activity measurements, our model must generate calcium signals from the voltage traces.
We model the calcium concentration [Ca] i of each neuron i as a ﬁrst-order leaky integrator, driven by voltage-gated calcium channels with the same nonlinear current-voltage (I-V) function g ( v i ) :
[Equation]
where g ( · ) represents softplus activation, τ [Ca] is a time constant shared across all neurons.
We then map calcium concentration
[Ca] into the measured calcium ﬂuorescence signals f via an afﬁne transform with scalar α f and bias β f ,
[Equation]
with measurement noise represented by a noise amplitude σ f and a noise term (cid:15) fi ( t ) ∼ N (0 , 1) .These
equations are simulated in discrete time using Euler integration.
2. V ARIATIONAL INFERENCE
Note that voltage dynamics in equation 1 is a deterministic equation.
We provide further ﬂexibility to the model in order to model the signiﬁcant trial-to-trial variability we observed in Figure 1c.
Such variability could be caused by different initial state of neurons, or by unknown, unmeasured inputs from environment.
In order to address this, we allow for the voltage dynamics to deviate from the deterministic dynamics, and introduce a Gaussian prior over these deviations.
This results in the following stochastic version of the voltage dynamics,
[Equation]
Where σ vi is the noise amplitude, and the standard normal noise term is (cid:15)
vi ∼ N (0 , 1) .
With this approach, we re-formulated the neuronal voltage as a latent variable v with the generative model given by 7 and must perform Bayesian inference to estimate their value.
Exact Bayesian inference is intractable for this high-dimensional nonlinear stochastic dynamical system.
For this reason, we developed a computationally efﬁcient strategy for inferring an approximate posterior distribution
P ( v | f ) over the latent variable v , given the measured ﬂuorescence data f using variational Bayes in the framework of the variational autoencoder [14], see derivation in Appendix A.1.
2. I NFERENCE NETWORK
The inference network Q φ ( v | f , o ) generates the approximate posterior distribution of the latent variable v (voltage for all neurons at all time points) given the measured ﬂuorescence data f and the sensory inputs o .
It consists of multiple 1D convolutional layers that transform the input data into a Gaussian distributions v i ( t ) ∼ N ( µ i ( t ) , σ i ( t )) for each neuron i = 1 , ..., N at every time point t = 0 , ..., D .
The approximate posterior distribution Q φ ( v | f , o ) is factorized as
[Equation]
D is the length of the time window of data fed into the inference network, and φ is its parameters.
2. G ENERATIVE MODEL
The generative model outputs the calcium ﬂuorescence trace f and neuronal voltage v given sensory inputs o .
It can be formulated as
[Equation]
where θ contains parameters of the generative model.
P θ ( v | o ) is a biophysically realistic connectomeconstrained network with passive point-neuron dynamics.
It outputs a prior distribution over voltage v given sensory input o .
We use a nonlinear sensory mapping H to transform our chemosensory input to neuronal stimulus to the sensory neurons.
The stimuli is represented as a set of binary vectors h k ( t ) , which are 1 given the presence of stimulus k at time t and 0 otherwise.
P θ ( f | v , o ) is another realistic model that maps the voltage v into reconstructed ﬂuorescence f , given sensory inputs o .
The voltages v are sampled from the approximate posterior distributions generated by the inference network.
Details of both distributions are described in Appendix A.2.
[Table]
2. O BJECTIVE FUNCTION FOR GENERATIVE MODEL AND INFERENCE NETWORK
The variational objective of our latent variable model uses the ELBO.The reconstructed calcium ﬂuorescence f is then ﬁt to the measured ﬂuorescence data by maximizing the ELBO for each measured time point.
The objective function is a combination of reconstruction loss between the reconstructed and measured ﬂuorescence traces and the KL divergence between voltage prior and posterior.
(details in Appendix A.4, A.5):
[Equation]
Optimization of the KL divergence with the reconstruction loss allows our inference network to learn the time dynamics of our neuronal voltages because the difference between the posterior and prior distributions is made small
2. C ONNECTOME CONSTRAINT FOR MODEL WEIGHTS
To build the model using chemical and electrical synapse count, C c and C e , as a connectome constraint we factorized the chemical synaptic weight W cji as W cji =
α c T cji M cji , where α c represents a global scaling factor which scales synapse counts to per-synapse currents or conductances depending on the model class.
T cji is a binary value which indicates the connectivity, and M cji reﬂects the magnitude of connection strength.
The electrical synaptic weights matrix W eji is similarly deﬁned as W eji = α e T eji M eji .
We compared the performance of the connectome constrained models to unconstrained ones (Table 1, Appendix A.6).
The anatomical connectome can be represented mathematically as two N × N matrices: one containing the chemical synapse counts ( C c ), and one containing the sizes of electrical synapses ( C e ).
We designed the connectome sparsity constraint by ﬁxing the sparsity of our model to that of the connectome.
For chemical connections, the model sparsity T cji = 1 [ C cji > 0] , and likewise for electrical connections, T eji =
[ C eji > 0] .
For both chemical and electrical connections, the magnitudes M cji and M eji are trainable.
We applied the connectome sparsity constraint to both current and conductance-based versions of the
LVM.
In the conductance-based model, we were able 1 to apply the even stronger connectome count constraint: directly assigning the synapse counts to the chemical connection magnitudes ( M cji =
C cji ) and the synapse sizes to the electrical connection magnitudes ( M eji = C eji ).
Under this constraint, both the sparsity and magnitude are ﬁxed, and only the global scaling factors α c and α e are trainable.
Additionally, we evaluated another variant of this constraint, connectome count 2 .
This constraint uses parameters of connectome count 1 to initialize M cji .
M cji is then allowed to deviate from C cji during re-training.
We compared our connectome-constrained models to two unconstrained models.
The learned dense model used fully connected sparsity matrices T c and T e , with the magnitudes M c and
M e as trainable parameters.
The learned sparse model uses L 1 regularization on both magnitude ) , and because the signs of the synaptic connections are unknown, we could not directly apply the connectome count constraint to it. matrices: (cid:107) M e (cid:107) 1 and (cid:107) M c (cid:107) 1 .
The learned total count model preserves the total synapse count for both magnitude matrices: (cid:107)(cid:107) M c (cid:107) 1 − (cid:107) C c (cid:107) 1 (cid:107)
and (cid:107)(cid:107) M e (cid:107) 1 − (cid:107) C e
(cid:107) 1 (cid:107) 22 .
α c and α e are 0.002 and 0.065 respectively, which is the same in the conne.count 1 model.
Thus, it regularizes the sum of elements in M c and M e equal to that of C c and C e .
E XPERIMENTS 3. D ATASET
We applied the CC-LVM to a calcium imaging dataset in which immobilized, pan-neuronally labeled
C.elegans were presented with a panel of chemosensory stimuli (2-butanone, 2,3-pentanedione, and
NaCl)
The activity of 170 neurons in the head of the animal were measured across 21 individuals.
Each recording lasts for 5 min, with an acquisition rate of 4 Hz.
In this dataset, every neuron in the worm brain could be identiﬁed in vivo using a deterministic multicolor landmark strain.
Thus, it was the ﬁrst experiment in which whole brain dynamics could be mapped onto the anatomical connectome at the resolution of single neurons.
The connectome constraints we applied utilized the anatomical connectivity data from [26].
The complete hermaphrodite C.elegans was reconstructed using electron microscopy, and the connections between its 300 neurons were mapped.
The connectome contains 3464 chemical synapses and 1031 electrical synapses.
More details can be found in appendix A.6.
3. N EURON HOLDOUT EVALUATION
We hypothesized that the CC-LVM which incorporates the complete connectome but has access only to the measured activity of a subset of the neurons may nevertheless be able to accurately predict the activity of unmeasured neurons.
This would constitute experimentally testable predictions from our mechanistic model.
We tested this hypothesis by performing neuron holdout evaluations, withholding a single bilateral pair of measured neurons from the model during both training and testing.
Note that our latent variable model always infers the voltages of all neurons regardless of how many neurons are observed in any experiment, since the latent states are the voltages of the entire nervous system.
The goal of this evaluation is to test the model’s ability to predict neural activity of neurons which were never observed at any time.
We ﬁt a version of the model with a pair of neurons withheld, and then compare the model predictions for the same pair of neurons with the ground truth calcium ﬂuorescence measurements.
We repeated this evaluation for all 107 measured neuron pairs.
Because many of the neurons in C. elegans are bilaterally symmetric, the symmetric pairs tend to have high correlations with each other.
In order to prevent the model from predicting one neuron solely based on its sibling, we removed the symmetric neurons in groups of two.
For every neuron pair, we performed the neuron holdout evaluation for both current and conductance models with different levels of connectome constraints.
For these experiments, we calculated the correlation coefﬁcient (details in Appendix A.4) between the predicted and measured calcium ﬂuorescence magnitudes to quantify the predictive performance (Figure 2a and Table 2, more predicted traces in Appendix A.9).
For each evaluation, we trained 4 models with random initialization, and reported both the mean value and standard error (SE) for each model.
We found connectome sparsity and connectome count constrained models were better at predicting traces of held out neurons than unconstrained models.
3. W ORM HOLDOUT EVALUATION
Another evaluation method we performed was to hold out the data from a handful of individual worms, train the model on the remaining worms, and predict the activity of the neurons of the held out individuals.
For each of the 9 model variants, we trained on 15 worms, and tested the model onwithheld worms.
In each case, we trained 4 models with different random initialization.
For the worm holdout traces, we calculated the correlation coefﬁcient between the predicted and measured calcium ﬂuorescence.
We also reported ELBO , reconstruction loss and KL divergence (Appendix A.4) on tested worms to represent the ﬁtting performance of each model.
We reported the mean and standard error (SE) from models trained with 4 different random initializations.
We found that models using the connectome count constraints performed better at predicting calcium traces from held out worms in Table 2.
[Table]
3. C OMPARING THE DIFFERENT MODEL VARIANTS
The gold standard for comparing generative models in machine learning is to compare the model evidence P ( f ) on a test dataset.
However, this is intractable for many generative models, and so the evidence lower bound ( ELBO ) is used for model ﬁtting and sometimes also for model comparison.
However, the ELBO is not a tight lower bound, and so is not guaranteed to provide the same ordering over models as the model evidence P ( f ) .
In this paper, exploiting the interpretable nature of our mechanistic generative model, we developed a stricter criterion for comparing models based on the neuron holdout evaluation.
Neuron holdout predictive performance is biologically meaningful, as it evaluates the ability of a model to make testable predictions about neurons which were not directly measured.
We found that overall, the conductance-based CC-LVMs under connectome count constraint make the best neuron holdout predictions as quantiﬁed by the correlation coefﬁcient .
We found that model accuracies at neuron holdout prediction are not very tightly correlated with their test ELBO .
While connectome count 2 still achieves the best performance, the learned dense model achieves the second best performance, despite performing worst in the neuron holdout evaluation.
This suggests that the learned dense model has potentially discovered an equally accurate but mechanistically incorrect way to model the calcium imaging data.
Stronger connectome constraint improves the neuron holdout predictions .
We found that our
CC-LVM with the connectome count constraint produced signiﬁcantly better predictions than CCLVM with the connectome sparsity constraint, and both connectome constrained models outperformed unconstrained LVMs with the same dynamics (Table 2).
This suggests that the neuron activity predictions are improved by stronger connectome constraint.
Note that connectome count 2 achieves slightly better performance than connectome count 1 , as its parameters get initialized from connectome count 1 , then are allowed to deviate during re-training.
We also explored whether CC-LVMs achieve better performance due solely to the smaller number of trainable parameters.
Comparing the connectomeconstrained models to the learned sparse LVM with sparsity regularization, we found that the learned sparse model performed worse.
This indicates that the connectome constraint improves predictability due to its topology, not due to the number of trainable parameters.
Perhaps relatedly, we found superior predictions for interneurons and motor neurons, compared to sensory neurons.
While the inputs to motor neurons and especially interneurons are largely contained in the connectome, the inputs to sensory neurons are not as well constrained by the connectome since they also respond to external stimuli.
Instead, we trained a blackbox model to predict the tunings of sensory neurons to experimentally presented odors, but this likely does not fully characterize all the sensory inputs to the worm.
This result additionally suggests the utility of the connectome, for neurons where the connectome provides information about the dominant sources of input.
Conductance-based models outperform current-based models.
We found that conductance-based models achieve better performance in neuron holdout evaluations, despite needing to estimate the reversal potential E in addition to the non-negative strength of a chemical synapse W c .
This could be for several reasons.
The biophysical model of a conductance based synapse might be a more accurate description of reality.
There are also two technical advantages to this model.
First, conductance based synapses allow us to make the fullest use of the connectomic data by using the non-negative synapse counts.
Second, separating the optimization of the sign of a synapse from its magnitude might allow for an easier optimization of both.
In contrast, the sign and magnitude of a current-based synapse are coupled in the same parameter, and we have observed strong local minima preventing the efﬁcient optimization of the sign and the magnitude.
Accordingly, we found that conductance-based models converge to a connectome closer to the counts given by the original data with less variance due to the removal of the neuron sign parameter (Appendix A.7).
We found that a synapse initialized as excitatory is unlikely to change its sign after optimization.
For the same reasons, we are unable to effectively use the non-negative synapse counts contained in the connectomic data to inform the magnitude of the current-based synapse while leaving the sign unconstrained.
D ISCUSSION
In this work, we proposed a connectome-constrained latent variable model (CC-LVM) of the C.elegans nervous system.
Our latent variable framework enables the estimation of the initial state of the network dynamics and the capture of modeling imperfections, which together enable us to ﬁt models to calcium imaging data with signiﬁcant trial-to-trial variability.
In contrast to popular black box generative models used in LVMs in neural data modeling [19], we developed a biophysics based, connectome constrained, mechanistically detailed generative model for the neural dynamics.
We evaluated two current-based and conductance-based models for synaptic connection.
We also explored using connectomic data to constrain the model to different degrees of detail.
Since CC-LVMs enable the prediction of neurons which were not used to ﬁt the model, we were able to simulate and validate the process of making biologically relevant and experimentally testable predictions about neurons which have not yet been experimentally characterized.
We found the conductance-based model produces much better-correlated predictions in neuron holdout evaluation and worm holdout evaluation.
And using the knowledge of the connectome can aid in silico predictions of individual neurons whose activity were not measured.
We hope our mechanistic model can be used to predict the effects of perturbations or ablations, leading to a causal understanding of neural computation in neural circuits.
The research in this paper was enabled by recent advances in whole brain calcium imaging in identiﬁed neurons [28].
While we have only analyzed the ﬁrst datasets to be collected in this manner, we hope that a rich database of such recordings across a diversity of conditions will lead to achieving the dream of standard models of the C. elegans nervous system [10] which generalize across experiments.
We trained each of our LVMs on 1 Quadro RTX 8000.
Training procedures are described in Appendix A.3, A.4, A.5.
We released our software and datasets ( https://github.com/
TuragaLab/wormvae ) for reproducibility.
A CKNOWLEDGEMENT
This work was funded by the Howard Hughes Medical Institute.
Lu Mi is supported by Intel and Mathworks. We thank William Bishop, Roman Vaxenburg, Core Francisco Park, Helena Casademunt for the constructive comments to improve the paper A A PPENDIX A. V ARIATIONAL INFERENCE
We used bayesian inference (VAE) to compute the posterior probability over the latent variables v (neuron voltages) given the data f (measured calcium ﬂuorescence),
[Equation]
To solve this problem, we needed to compute the model evidence,
[Equation]
However, computing P ( f ) is a mathematically intractable problem, so we instead used variational inference to approximate the value of P ( v | f )
Denoting the approximate posterior Q ( v | f ) , we can formulate evidence lower bound (ELBO) as
[Equation]
which tightens the lower bound on log P ( f ) as L is maximized
This allowed us to improve the inference network and ﬁnd the posterior such that Q ( v | f ) ≈ P ( v | f ) , since the Kullback-Leibler (KL) divergence, which measures the distance between the distributions Q and P , is minimized by maximizing ELBO.
A. P RIOR AND POSTERIOR DISTRIBUTION AT A DISCRETE TIME STEP
Using Euler discretization with a time step of ∆ t , the ﬁrst-order leaky integrator equation in Equationin the network model is
[Equation]
Then the voltage v i ( t + ∆ t ) of neuron i at time step t + ∆ t is formulated as
[Equation]
[Equation]
This equation indicates that the voltage state v i ( t ) at the current time step t is determined by a function over multiple components from the last time step t − ∆ t , including the past voltage state, the chemical and electric synaptic inputs, and the sensory input:
[Equation]
For the posterior distribution Q φ ( v ( t ) | f
( t ) , o ( t )) , the inference network outputs µ i ( t ) and σ i ( t ) for every neuron i at each time step t .
The posterior distribution for the voltage generated by the inference network at each time step is v i ( t ) ∼ N ( µ i ( t ) , σ i ( t )) , given a sequence of measured ﬂuorescence data f (0) , ..., f ( D ∆ t ) and sensory input o (0) , ..., o ( D ∆ t ) with total number of steps D + 1 .
We then used the reparameterization trick to sample the voltage v i ( t ) from this posterior distribution to generate a sequence of voltage sample from v i (0) , ..., v i ( D ∆ t ) .
For the prior distribution, we used the voltage samples v i ( t ) from the posterior distribution at the last step as the past voltage state in F to calculate the prior distribution P θ ( v i ( t + ∆ t )
| o
i ( t + ∆ t )) at time step t + ∆ t .
The prior distribution is deﬁned as N ( F ( v i ( t ) , s c i ( t ) , s e i ( t ) , o i ( t + ∆ t )) , σ v i ) .
σ v i is the constant standard deviation of the prior distribution that is trainable during optimization.
At the initial time step 0 , we deﬁned another trainable parameter µ i (0) to represent the mean of the voltage prior distribution.
Similarly, we used Euler integration to discretize the ﬁrst-order leaky integrator equation (Equation 5).
The calcium concentration [ Ca ] i ( t + ∆ t ) of neuron i at time step t + ∆ t is given by
[Equation]
We used a sequence of voltage samples v i ( t ) , t = 0 , ..., D ∆ t as the inputs to this formula to get a sequence of calcium concentration [ Ca ] i (1) , ..., [ Ca ] i ( D ∆ t ) .
The initial concentration is estimated from the measured ﬂuorescence trace at time step 0 using [ Ca ] i (0) =
( F i (0) − β f ) /α f .
Equationin the generative model P θ generates the output distributions P θ ( f ( t ) | v ( t ) , o ) of ﬂuorescence trace at time step t as N ( α f [ Ca ] i ( t ) + β f , σ f ) .
The values α f , β f , σ f are trainable ﬂuorescence parameters.
A. M ODEL ARCHITECTURE
The calcium ﬂuorescence data is given in a matrix of dimensions ( N, T ) where N is the total number of neurons in the connectome. and a matrix of dimensions (3 , T ) which consists of 3 one-hot encoding vectors that are 1 when the chemosensory stimulus is present and 0 otherwise.
Each of the 3 vectors represents one of the three chemical stimuli presented: NaCl, 2-butanone and 2,3-pentanedione.
The calcium ﬂuorescence matrix is passed to the input layer of our inference network Q ( v | f , o ) which is a 1 D convolutional layer.
The output of the input layer then gets passed through an upsampling layer, a ReLU nonlinearity, another 1 D convolutional layer, and then another upsampling layer.
The ﬁrst convolutional layer has 2 N input channels, 2 N output channels and a kernel size of 11; the second has 2 N input channels, 2 N output channels and a kernel size of 21.
The ﬁrst upsampling layer upsamples by a factor of 4 and the second upsamples by a factor of 10. These factors account for the fact that we simulate our the network r times faster with upsampling factor r = 40 .
The sensory input is passed through a 1 D convolutional layer with N input channels, 2 N output channels and a kernel size of 21.
This is then passed through a ReLU and then concatenated to the output of the 2nd convolution on our calcium data.
The merged matrix is then 1 D convolved again with 2 N inputs, 2 N outputs and a kernel size of 41.
This is one of the ﬁnal outputs of our inference network and gives a ( N, T ) matrix that represents the means of our posterior distribution.
To generate the standard deviation of our factorized posterior distribution, we pass our merged data through another convolutional layer with the same dimensions but different parameters.
We pass that output through a nonlinear afﬁne transform (using a softplus nonlinearity) to rescale the output to an appropriate level of noise.
The parameters of this inference network are all the weights and biases of the convolutions and upsampling layers described above, which we denote by φ in the paper.
The input to the generative model is the same sensory input that we pass to our inference network.
The ﬁrst layer of our generative model is a single layer nonlinear transform that converts ourdimensional input to an N dimensional input that is 0 to any non-sensory neurons.
The next layer is our recurrent neural dynamics, given by Equation 1, which passes information to itself in the next time step based on the connectivity given by the connectome.
Although some sources say that C. elegans has 302 neurons, we only simulate 300 because two cells originally classiﬁed as neurons are not connected to the rest of the connectome Witvliet et al.
The simulated voltage for each neuron in the connectome is passed to our calcium simulation layer (Equation 5).
The ﬁnal layer of our model is a calcium to ﬂuorescence afﬁne transformation (Equation 6).
The output of this ﬁnal layer is then compared to the ground truth ﬂuorescence data to calculate the reconstruction loss.
Below, we have included some psuedocode describing the architecture of our network.
Inference_Network ( calcium_fluor , missing_data_mask , sensory_input ) : calcium_input = concatenate ( calcium_fluor , missing_data_mask ) conv
[Equation]
conv
[Equation]
calcium
[Equation]
input
[Equation]
up
[Equation]
upsample
[Equation]
conv
[Equation]
out ) conv2_
[Equation]
conv
[Equation]
) up
[Table]
[Equation] upsample ( conv [Equation] out )
[Table] [Table] [Table] ) mean_latent_neuron_voltage = conv4 ( merged_calc_sensory ) s
t
d _ l a t e n t _ n e u r o n _ v o l t a g e = s o f t p l u s ( conv5 ( merged_calcim_sensory ) ) sample_latent_neuron_voltage = mean_latent_neuron_voltage + rand_norm * s t
d
_ l a t e n t _ n e u r o n _ v o l t a g e r e t u r n sample_latent_neuron_voltage Generative_Model ( sample_latent_neuron_voltage , sensory_input ) : # Equation 1 neuron_voltage_dynamics = leaky_integrator_connectome_dynamics ( sample_latent_neuron_voltage , sensory_input ) # Equation 5 calcium_concentration_dynamics = leaky_integrator_calcium_model ( neuron_voltage_dynamics )
# Equation 6 f l u o r e s c e n c e _ t r a c e = n o n l i n e a r _ a f f i n e _ t r a n s f o r m ( calcium_concentration_dynamcis ) r e t u r n f l u o r e
s c e n c e _ t r a c e
A. F ORMULATIONS FOR OBJECTIVE FUNCTIONS AND METRICS
Here we provide the mathematical formulations for the objective functions and metrics we used in this work.
The evidence lower bound (ELBO) is a combination of reconstruction log likelihood and the negative
KL divergence.
The reconstruction log likelihood, given the predicted ﬂuorescence trace ˆ f i ( t ) of neuron i at time step t as output distribution N ( ˆ f i ( t ) , σ f ) .
And ˆ f i is formulated as
[Equation]
Given the measured ﬂuorescence trace f i ( t ) , number of neurons N , and total number of time steps
D + 1 , we calculated the reconstruction negative log likelihood as the reconstruction loss
[Equation]
For KL divergence between voltage prior and posterior distribution.
We represent the posterior distribution generated from the inference network as N ( µ i ( t ) , σ i ( t )) for neuron i at each time step t .
We then sample v i ( t ) from the posterior distribution.
According to Equation 16, we represent the mean ˆ v i ( t ) of prior distribution N ( ˆ v i ( t ) , σ v i ) as
[Equation]
where σ vi is a trainable parameter.
We then use a closed form of KL divergence deﬁned between two univariate Gaussian distributions
[Equation]
We minimize the objective function L = − ELBO during optimization given by
[Equation]
In the worm holdout evaluation, for each neuron pair, we calculated the ELBO , reconstruction loss and KL divergence metrics on each holdout worm, for a total 6 worms.
We also trained each model setting 4 times under the same constraints and different initializations.
We calculated the mean and standard error of ELBO , reconstruction loss and KL divergence based on 4 × 6 = 24 samples for each model.
For both neuron and worm holdout evaluations, we calculated the correlation coefﬁcients between the measured f i ( t ) and reconstructed ﬂuorescence traces ˆ f i ( t ) . The correlation coefﬁcient for each neuron is simply
[Equation]
In the neuron holdout evaluation, for each neuron pair, we trained a model in which the ground truth of said neuron pair was removed from the training data.
There were 107 pairs and 170 measured neurons in total.
We trained each model 4 times with different initializations under each of 8 connectome constraints, for a total of 8 × 107 × 4 = 3424 different models.
We calculated the mean and standard error of r i based on 4 × 170 = 680 samples for each model.
In the worm holdout evaluation, we calculated these metrics on each holdout worm, for a total of 6 worms.
Each worm has 170 measured neurons.
We trained each model 4 times under the same connectomic constraints and with different initializations.
We calculated the mean and standard error of r i based on 6 × 170 × 4 = 4080 samples for each model as our metrics.
A. G RADIENT - BASED OPTIMIZATION
We used PyTorch to minimize the ELBO using gradient descent with respect to the parameters of both the generative and inference models,
For LVMs which used connectome count constraint with α c and α e trainable.
The objective function is
[Equation]
For LVMs which used other constraints including connectome count and learned dense , learned sparse , where M c and M e trainable.
The objective function is
[Equation]
(25) where φ is all the parameters of the inference network.
We used a clamp in PyTorch to constrain the weight magnitude M c and M e to be non-negative in several models.
We restrict the M e to be symmetric by deﬁning it as the sum of an upper triangular matrix and its transpose.
Our simulation time step is 6 .
25 ms.
Each window size D lasts for 7 . 5 s with 30 imaging steps.
We used weights offor both the reconstruction loss and KL divergence.
We used Adam with a learning rate scheduler to perform the optimization.
We used a initial learning rate of 3 e − 4 , with a learning rate schedular with step size of 50, and gamma of 0.5.
We also set a gradient clip value of 1.
Each model was trained for 300 epochs in which each epoch is one full pass through all the training data.
In the paper, we use the variable θ to denote the trainable parameters of our generative model.
The parameters in θ are given by
[Equation]
where τ and τ [ Ca ] are the time scales of the neuron and calcium dynamics.
M c and M e are the magnitudes of neuronal connections for chemical and electrical synapses respectively.
α f and β f are the scaling and bias parameters of the ﬂuorescence afﬁne transform.
σ v and σ f are the noise magnitudes in the neural dynamics equation (Equation 1) and the ﬂuorescence equation (Equation 6).
The variable H represents the parameters of the sensory mapping that transforms chemosensory data into neural stimulus.
The variable φ denotes all the trainable parameters of the inference network.
These are just all the weights and biases of the 5 different convolutional layers described in section A.3.
A. C ONNECTOME CONSTRAINTS FOR MODEL WEIGHTS
We constructed CC-LVMs on both current-based and conductance-based synapses with different levels of connectomic constraints.
Given anatomical connectome data with two N × N matrices: chemical synapse counts ( C c ), and sizes of electrical synapses ( C e ).
The chemical synaptic weight W cji as
W cji =
α c T cji M cji , where α c represents a global scaling factor.
T cji is a binary value which indicates the connectivity, and M cji reﬂects the magnitude of connection strength.
The electrical synaptic weights matrix W eji is similarly deﬁned as W eji = α e T eji M eji .
We provide more implementation details of each constraint below.
conductance-based + connectome count 1 We directly assign the synapse counts to the chemical connection magnitudes ( M cji =
C cji ) and the synapse sizes to the eletrical connection magnitudes ( M eji = C eji ).
Under this constraint, both the sparsity T cji and T eji are ﬁxed to the connectome sparsity, magnitude M cji and M eji are ﬁxed to the connectome synapse count, and only the global scaling factors α c and α e are trainable.
M cji and M eji are restricted to be non-negative.
conductance-based + connectome count 2 We initialize all model parameters with the parameters from the trained conductance-based + connectome count 1 model.
And then the sparsity T cji and
T eji are ﬁxed to the connectome sparsity, global scaling factors α c and α e are ﬁxed to 0 .
And now magnitude M cji and M eji are trainable during the re-training.
M cji and M eji are restricted to be non-negative.
conductance-based + connectome sparsity For chemical and electrical connections, the model sparsities are T cji =
[ C cji > 0] and T eji
= 1 [ C eji > 0] respectively.
For both chemical and electrical connections, the magnitudes M cji and M eji are trainable.
The sparsity T cji and T eji are ﬁxed to the connectome sparsity and global scaling factors, α c and α e , are ﬁxed to 0 .
M cji and M eji are restricted to be non-negative.
conductance-based + learned total count For chemical and electrical connections, the model is fully connected with the sparsity T cji = 1 and T eji = 1 respectively.
For both chemical and electrical connections, the magnitudes M cji and M eji are trainable.
The model is constrained to preserve the same total synapse count as connectome count (cid:107) C c (cid:107) for both magnitude matrices adding (cid:107)(cid:107) M c (cid:107) 1 − (cid:107) C c (cid:107) 1 (cid:107)
and (cid:107)(cid:107) M e (cid:107) 1 − (cid:107) C e
(cid:107) 1 (cid:107) 22 in the objective function.
α c and α e are deﬁned as 0.002 and 0.065, which are the same as the trained conne.count 1 model, so that regularizes the sum of elements in M c equal to that of C c , and likewise for M e and C e .
The sparsity T cji and T eji are ﬁxed to 1 and
M cji and M eji are restricted to be non-negative.
conductance-based + learned sparse For chemical and electrical connections, the model is fully connected with the sparsity T cji = 1 and T eji = 1 respectively.
For both chemical and electrical connections, the magnitudes M c ji and M e ji are trainable.
Meanwhile, it uses an additional L [Table] regularization on both magnitude matrices: (cid:107) M c (cid:107) 1 and (cid:107) M e (cid:107) 1 , so that the fraction of nonzero elements in M c is small, and likewise for M e .
The sparsity T cji and T eji are ﬁxed to 1 and global scaling factors α c and α e are ﬁxed to 0 .
M cji and M eji are restricted to be non-negative.
conductance-based + learned dense For chemical connections and electrical connections, the model is a fully connected with sparsity T cji = 1 and T eji = 1 respectively.
For both chemical and electrical connections, the magnitudes M cji and M eji are trainable.
The sparsity matrices, T cji and T eji are ﬁxed.
Global scaling factors α c and α e are also ﬁxed to 0 .
M cji and M eji are restricted to be non-negative.
current-based + connectome sparsity For chemical and electrical connections, the model sparsity is
T cji =
[ C cji > 0] and T eji
= 1 [ C eji > 0] respectively.
For both chemical and electrical connections, the magnitudes M cji and M eji are trainable.
The sparsity T cji and T eji are ﬁxed to the connectome sparsity and global scaling factors, α c and α e , are ﬁxed to 0 .
M eji are restricted to be non-negative, while the sign of M cji is not restricted.
current-based + learned sparse For chemical and electrical connections, the model is fully connected with sparsity T cji = 1 and T eji = 1 .
For both chemical and electrical connections, the magnitudes
M cji and M eji are trainable.
For both chemical and electrical connections, the magnitudes M cji and
M eji are trainable.
Meanwhile, it uses an additional L 1 regularization on both magnitude matrices: || M c
||[-
-]
and || M e || 1 , so that the fraction of nonzero elements in M c is small, and likewise for M e .
The sparsity T cji and T eji are ﬁxed to 1 , global scaling factors, α c and α e , are ﬁxed to 0 .
M eji are restricted to be non-negative, while the sign of M cji is not restricted.
current-based + learned dense For chemical and electrical connections, the model is fully connected with T c ji = 1 , and T e ji = 1 .
For both chemical and electrical connections, the magnitudes M c ji and
M e ji are trainable.
The sparsity T c ji and T e ji are ﬁxed to 1 and global scaling factors α c and α e are ﬁxed to 0 .
M eji are restricted to be non-negative, while the sign of M cji is not restricted.
A. M ODEL WEIGHT CONSISTENCY ANALYSIS
We used the correlation coefﬁcient to quantify the inter-consistency of models trained with different random initializations, as well as the connectome-consistency, which is the correlation coefﬁcient between model weights and connectome counts.
In particular, for the connectome consistency, we compare the consistency of our model weights and connectome synapse count.
We calculate the correlation coefﬁcient r
i between the absolute value of model weights | W ij | and connectome synapse count C ij along each row i .
Then it is averaged across
N rows and 4 models with identical constraints but different random initializations to get ¯ r i as our metric for connectome-consistency,
[Equation]
For inter-consistency, we evaluate how model weights are affected by random initialization after training.
We calculate the correlation coefﬁcient r
i between the weights W 1 ij of model 1 and weights
W 2 ij of model 2 along each row i .
Model 1 and model 2 have the same constraints but different random initializations.
The correlation coefﬁcient between them is
[Equation]
This coefﬁcient, r i , is then averaged across N rows and each combination of 2 models within themodels with random initialization to get ¯ r i as our metric for inter-consistency.
As shown in Table 3, the conductance-based model parameters are more consistent under different random initializations.
Optimization of the conductance-based model is easier because the chemical weights W c of the conductance-based synapses are always positive and we do not need a trainable sign variable.
Meanwhile, the model parameters are better constrained with the connectome sparsity or counts, which enables it to be consistent across different initialization.
A. P REDICTIVE PERFORMANCE FOR WORM HOLDOUT V . S . N EURON HOLDOUT
We ﬁrstly demonstrate the violin plot for worm holdout evaluation in Figure 4.
We found that the results from worm holdout are different from those in the neuron holdout experiment in Figure 2.
Here, all 9 LVM variants perform similarly across different categories.
Our connectome-constrained models do not outperform other unconstrained baselines.
We also studied whether the predictive performance of worm holdouts and neuron holdouts were correlated.
In Figure 5, we show the scatter plot of predictive performance for worm holdouts v.s. neuron holdouts.
Each point represents one neuron from the same worm.
We also report the correlation coefﬁcient between the predictive performance for worm holdout and neuron holdout for each LVM variant.
We found that our connectome constrained model achieves the highest correlation, indicating highest consistency for neuron sets which gets better performance in both evaluations.
A. P REDICTED HOLDOUT TRACES WITH VARIOUS CONSTRAINTS
For each neuron, we performed both neuron holdout evaluation and worm holdout evaluation for current-based and conductance-based models with different levels of connectomic constraint.
For neuron holdout evaluation, we plot additional predicted traces from 15 representative neurons, including 12 with good performance (top), and 3 with worse performance (bottom) in Figure 6, model with connectome count constraint outperforms other baselines.
For worm holdout evaluation, we show the results on the same 15 neurons from 9 LVMs in Figure 7, all models achieves comparable results.
A. P REDICTED TRACES AND VOLTAGE LATENT SPACE FOR NEURON HOLDOUT EVALUATION
For neuron holdout evaluation, we also show the predicted traces and voltage latent space for all measured neurons.
We show 189 neurons in Figure 8, Figure 9 and Figure 10.
In each subplot, we reported the measured ﬂuorescence trace ground truth, predicted ﬂuorescence trace ground truth, and the mean and standard deviation (std) of voltage latent space (posterior distribution) from the inference network.
A. P REDICTED TRACES AND VOLTAGE LATENT SPACE FOR WORM HOLDOUT EVALUATION
For worm holdout evaluation, similarly, we also show the predicted traces and voltage latent space for all neurons.
We show 300 neurons from a train worm in Figure 11–Figure 15.
Meanwhile, we also show 300 neurons from a test worm in Figure 16–Figure 20.
In each subplot, we reported the measured ﬂuorescence trace ground truth (for recorded neurons), predicted ﬂuorescence trace ground truth, and the mean and standard deviation (std) of voltage latent space (posterior distribution) from the inference network.
For the train worm, we ﬁnd most of the neurons have good predictive performance compared to the measured ground truth.
That indicates optimization is saturated on train set.
However, test worm has a relatively worse predictive performance, that indicates the generalization gaps between each individual trial (animal).
While some neurons still achieve good results, that indicates our model is still generalized to unseen individuals to some degree.
Interestingly, for both train and test worms in the worm holdout evaluation, we ﬁnd those unmeasured neurons still have reasonable patterns.
That indicates our model successfully avoids the degeneracy of solutions for those unmeasured neurons.