Fast deep learning correspondence for neuron tracking and identification in C. elegans using synthetic training
Xinwei Yu1, Matthew S. Creamer2, Francesco Randi1, Anuj K. Sharma1, Scott W. Linderman3,4, Andrew M. Leifer1,2,*
1 Department of Physics, Princeton University, Princeton, NJ, United States of America
2 Princeton Neuroscience Institute, Princeton University, Princeton, NJ, United States of America
3 Department of Statistics, Stanford University, Stanford, CA, United States of America
4 Wu Tsai Neurosciences Institute, Stanford University, Stanford, CA, United States of America
We present an automated method to track and identify neurons in C. elegans, called “fast Deep Learning Correspondence” or fDLC, based on the transformer network architecture. The model is trained once on empirically derived synthetic data and then predicts neural correspondence across held-out real animals via transfer learning. The same pre-trained model both tracks neurons across time and identifies corresponding neurons across individuals. Performance is evaluated against hand-annotated datasets, including NeuroPAL . Using only position information, the method achieves 80.0% accuracy at tracking neurons within an individual and 65.8% accuracy at identifying neurons across individuals. Accuracy is even higher on a published dataset . Accuracy reaches 76.5% when using color information from NeuroPAL. Unlike previous methods, fDLC does not require straightening or transforming the animal into a canonical coordinate system. The method is fast and predicts correspondence in 10 ms making it suitable for future real-time applications.
The nervous system of the nematode C.elegans is well characterized, such that each of the 302 neurons are named and have stereotyped locations across animals [28, 27, 34]. The capability to find corresponding neurons across animals is essential to investigate neural coding and neural dynamics across animals. Despite the worm’s overall stereotypy, the variability in neurons’ spatial arrangement is sufficient to make predicting neural correspondence a challenge. For whole-brain calcium imaging [32, 22], identifying neurons across animals is additionally challenging because the nuclear localized markers that are used tend to obscure morphological features that would otherwise assist in neural identification.
An ideal method for finding neural correspondence in C. elegans should accommodate two major use cases. The first is tracking neurons within an individual across time as the animal’s head moves and deforms. Here the goal is to be able to say with confidence that a neuron imaged in a volume taken at time is the same as another neuron taken from a volume imaged at time . Tracking across time is needed to extract calcium dynamics from neurons during freely moving population calcium imaging [32, 22, 11]. Additionally, very fast real-time tracking will be needed to guide closed-loop optogenetic stimulation of neurons in moving animals as those methods move to larger neural populations [14, 26, 9, 25].
The second and more general use case is finding neural correspondence across individuals. Often this is to identify the name of a neuron with respect to the connectome  or a gene expression atlas . Even when a neuron’s name cannot be ascertained, being able to identify which neurons are the same across recordings allows researchers to study neural population codes common across individuals.
For both use cases, a method to find neural correspondence is desired that is accurate, fast, requires minimal experimental training data and that generalizes across animal pose, orientation, imaging hardware, and conditions. Furthermore, an ideal method should not only perform well when restricted to neural positioning information but, should also be flexible enough to leverage genetically encoded color labeling information or other features for improved accuracy when available. Multicolor strains are powerful new tools that use multiple genetically encoded fluorescent labels to aid neural identification [35, 29] (we use one of those strains, NeuroPAL , for validating our model). However, some applications, like whole-brain imaging in moving worms, are not yet easily compatible with the multicolor imaging required by these new strains, so there remains a need for improved methods that use position information alone.
A variety of automated methods for C. elegans have been developed that address some, but not all of these needs. Most methods developed so far focus on finding the extrinsic similarity  between one neuron configuration, called a test, and another neuron configuration called a template. Methods like these deform space to minimize distances between neurons in the template and neurons in the test and then attempt to solve an assignment problem . For example, a simple implementation would be to use a non-rigid registration model, like Coherent Point Drift (CPD)  to optimize a warping function between neuron positions in the test and template. More recent non-rigid registration algorithms like PR-GLS  also incorporate relative spatial arrangement of the neurons .
Models can also do better by incorporating the statistics of neural variability. NeRVE registration and clustering , for example, also uses a non-rigid point set registration algorithm  to find a warping function that minimizes the difference between a configuration of neurons at one time point and another. But NeRVE further registers the test neurons onto multiple templates to define a feature vector and then finds neural correspondence by clustering those feature vectors. By using multiple templates, the method implicitly incorporates more information about the range and statistics of that individual animal’s poses to improve accuracy.
A related line of work uses generative models to capture the statistics of variability across many individual worms. These generative models specify a joint probability distribution over neural labels and the locations, shapes, sizes, or appearance of neurons identified in the imaging data of multiple individuals [4, 30, 19, 20]. These approaches are based on assumptions about the likelihood of observing a test neural configuration, given an underlying configuration of labeled neurons. For example, these generative models often begin with a Gaussian distribution over neuron positions in a canonical coordinate system and then assume a distribution over potentially non-rigid transformations of the worm’s pose for each test configuration. Then, under these assumptions, the most likely neural correspondence is estimated via approximate Bayesian inference.
The success of generative modeling hinges upon the accuracy of its underlying assumptions, and these are challenging to make for high-dimensional data. An alternative is to take a discriminative modeling approach . For example, recent work  has used conditional random fields (CRF) to directly parameterize a conditional distribution over neuron labels, rather than assuming a model for the high-dimensional and complex image data. CRF allows for a wide range of informative features to be incorporated in the model, such as the angles between neurons, or their relative anterior-posterior positions, which are known to be useful for identifying neurons . Ultimately, however, it is up to the modeler to select and hand curate a set of features to input into the CRF.
The next logical step is to allow for much richer features to be learned from the data. Artificial neural networks are ideal for tackling this problem, but they require immensely large training sets. Until now, their use for neuron identification has been limited. For example, in one tracking algorithm, artificial neural networks provide only the initialization, or first guess, for non-rigid registration .
Our approach is based on a simple insight: it is straightforward to generate very large synthetic datasets of test and template worms that nonetheless are derived from measurements. We use neural positions extracted from existing imaging datasets, and then apply known, nonlinear transformations to warp those positions into new shapes for other body postures. Furthermore, we simulate the types of noise that appear in real datasets, such as missing or spurious neurons. Using these large-scale synthetic datasets, we can train an artificial neural network to map the simulated neural positions back to the ground truth. Given sufficient training data (which we can generate at will), the network learns the most informative features of the neural configurations, rather than requiring the user to specify them by hand. Realistic synthetic or augmented datasets like these have been key to cracking other challenging problems in neural and behavioral data analysis [13, 17], and have already shown promising potential for tracking neurons .
In this work, we use synthetic data to train a Transformer network, an artificial neural network architecture that has shown great success in natural language processing tasks . Transformers incorporate an attention mechanism that can leverage similarities between pairs of inputs to build a rich representation of the input sequence for downstream tasks like machine translation and sentiment prediction. We reasoned this same architecture would be well-suited to extract spatial relationships between neurons in order to build a representation that facilitates finding correspondence to neurons in a template worm.
Not only is the Transformer well-suited to learning features for the neural correspondence problem, it also obviates the need to straighten  the worm in advance. Until now, existing methods have either required the worm to be straightened in preprocessing [4, 5] or explicitly transformed them during inference [30, 19]. Straightening the worm is a non-trivial task, and it is especially error-prone for complicated poses such as when the worm rolls along its centerline.
Finally, one of the main advantages of the Transformer architecture is that it permits parallel processing of the neural positions using modern GPU hardware. In contrast to existing methods, which have not been optimized for speed, the Transformer can make real-time predictions once it has been trained. This speed is a necessary step toward future experiments with real-time, closed-loop, targeted delivery of optogenetic stimulation of large populations of neurons in freely moving animals.
Materials and methods
Our model was trained on a synthetic dataset derived from recordings of 12 freely moving animals. The model’s performance was evaluated on a different set of 12 held-out recordings: one moving recording and 11 immobile recordings. The recording of moving animals used strain AML32 wtfIs5[Prab-3::NLS::GCaMP6s; Prab-3::NLS::tagRFP] and was first reported in . Immobile recordings to evaluate performance used strain AML320 (otIs669[NeuroPAL] V 14x; wtfIs145 [pBX + rab-3::his-24::GCaMP6::unc-54] ) derived from NeuroPAL strain OH15262 . All strains and datasets in this study used nuclear localized fluorescent reporters.
Neural configurations acquired as part of this study have been posted in an Open Science Foundation repository with DOI:10.17605/OSF.IO/T7DZU available at https://dx.doi.org/10.17605/OSF.IO/T7DZU.
To image neurons in the head of freely moving worms, we used a dual-objective spinning-disk based tracking system  (Yokogawa CSU-X1 mounted on a Nikon Eclipse TE2000-S). Fluorescent images of the head of a worm were recorded through a 40x objective with both 488- and 561-nm excitation laser light as the animal crawled. The 40x objective translated up and down along the imaging axis to acquire 3D image stacks at a rate of 6 head volumes/s.
To image neurons in the immobile multi-color NeuroPAL worms  we modified our setup by adding emission filters in a motorized filter wheel (Prior ProScan-II), and adding a Stanford Research Systems SR474 shutter controller (with SR475 shutters) to programmatically illuminate the worm with different wavelength laser light. We use three lasers with light of different wavelengths: 405 nm (Coherent OBIS-LX 405 nm 100 mW), 488 nm (Coherent SAPPHIRE 488 nm 200 mW), and 561 nm (Coherent SAPPHIRE 561 nm 200 mW). Only one laser at a time reached the sample, through a 40x oil-immersion objective (1.3 NA, Nikon S Fluor). The powers measured at the sample, after spinning disk and objective, were 0.14 mW (405 nm), 0.35 mW (488 nm), and 0.36 mW (561 nm). In the spinning disk unit, a dichroic mirror (Chroma ZT405/488/561tpc) separated the excitation from the emission light. The latter was relayed to a cooled sCMOS camera (Hamamatsu ORCA-Flash 4.0 C11440-22CU), passing through the filters mounted on the filter wheel (Table 1). Fluorescent images were acquired in different “channels”, i.e. different combinations of excitation wavelength, emission filter, and camera exposure time (Table 2). The acquisition was performed using a custom software written in LabVIEW that specifies the sequence of channels to be imaged, and controls shutters, filter wheel, piezo translator, and camera. After setting the z position, the software acquires a sequence of images in the specified channels.
|Filter label||Filters (Semrock part n.)|
|F3||FF02-675/67 + FF01-692/LP|
|Channel||Excitation (nm)||Emission window (nm) [filter]||Primary Fluorophore|
Preprocessing and Segmentation
We extracted the position of individual neurons from 3D fluorescent images to generate a 3D point cloud, (Fig. 1A). This process is called segmentation and the fDLC model is agnostic to the specific choice of the segmentation algorithm.
For recordings of strains AML32, we used a segmentation algorithm adopted from . We first applied a threshold to find pixels where the intensities are significantly larger than the background. Then, we computed the 3D Hessian matrix and its eigenvalues of the intensity image. Candidate neurons were regions where the maximal eigenvalue was negative. Next, we searched for the local intensity peaks in the region and spatially disambiguated peaks in the same region with a watershed separation based on pixel intensity.
Generating synthetic point clouds for training
We developed a simulator to generate a large training set of synthetic animals with known neural correspondence. The simulator takes as its input the point clouds collected from approximately 4,000 volumes spread across recordings of 12 freely moving animals. For each volume, the simulator performs a series of stochastic deformations and transformations to generate 64 new synthetic individuals where the ground truth correspondence between neurons in the individuals and the original point cloud is known. A total of synthetic point clouds were used for training.
The simulator introduces a variety of different sources of variability and real-world deformations to create each synthetic point cloud (Fig. 1B,E). The simulator starts by straightening the worm in the XY plane using its centerline so that it now lies in a canonical worm coordinate system. Before straightening, Z is along the optical axis and XY are defined to be perpendicular to the optical axis and are arbitrarily set by the orientation of the camera. After straightening, the animal’s posterior-anterior axis lies along the X axis. To introduce animal-to-animal variability in relative neural position, a non-rigid transformation is applied to the neuron point cloud against a template randomly selected from recordings of the real observed worms using coherent point drift (CPD) . To add variability associated with rotation and distortion of the worm’s head in the transverse plane, we apply a random affine transformation to the transverse plane. To simulate missing neurons and segmentation errors, spurious neurons are randomly added, and some true neurons are randomly removed, for up to 20% of the observed neurons. To introduce variability associated with animal pose, we randomly deform the centerline of the head. Lastly, to account for variability in animals’ size and orientation, a random affine transformation in XY plane is applied that rescaled the animal’s size by up to 5%. With those steps, the simulator deforms a sampled worm and generates a new synthetic worm with different orientation and posture while maintaining known correspondence.
Whenever practical, the magnitude of the simulator’s perturbations was informed by observation of real worms. For example the centerlines generated by the simulator were directly sampled from recordings of real individuals. The magnitude of added Gaussian noise was set to have a standard deviation of , as this roughly matched our estimate of variability observed by eye.
Deep Learning Correspondence Model
Overview and input
The deep learning correspondence model (fDLC) is an artificial neural network based on the Transformer  architecture (Fig. 1C) and is implemented in the automatic differentiation framework PyTorch . The fDLC model takes as input the positional coordinates of a pair of worms, a template worm , and test worm, (Fig. 1F). For each worm, approximately 120 neurons are segmented and passed to the fDLC model.
The model works as an encoder, which maps the input neuron coordinates to continuous embeddings . The model is composed of a stack of identical layers. Each layer consists of two sub-layers: a multi-head self-attention mechanism , and a fully connected feed-forward network. The multi-head attention mechanism is the defining feature of the transformer architecture and makes the architecture well-suited for finding relations in sequences of data, such as words in a sentence or, in our case, spatial locations of neurons in a worm. Each head contains a one-to-one mapping between the nodes in the artificial network and the C. elegans neurons. In the transformer architecture, features of a previous layer are mapped via a linear layer into three attributes of each node, called the query, the key and the value pairs. These attributes of each node contain high dimensional feature vectors which, in our context, represent information about the neuron’s relative position. The multi-head attention mechanism computes a weight for each pair of nodes (corresponding to each pair of C. elegans neurons). The weights are calculated by performing a set computation on the query and key. The output is calculated by multiplying this resultant weight by the value. In our implementation, we set the number of heads in the multi-head attention module to be 8 and we set the dimension of our feature vectors to be 128. We choose the best set of the hyperparameters by evaluating on a validation set, which is distinct from the training set and also from any data used for evaluation. A residual connection  and layer normalization  are employed for each sub-layer, as is widely used in artificial neural networks.
Calculating probabilities for potential matches
The fDLC model generates a high dimensional () embedding for neuron from the template worm and for the neuron from the test worm. The similarity of a pair of embeddings, as measured by the inner product , determines the probability that the pair is a match. Specifically, we define the probability that neuron in the template worm matches neuron in the test worm as , where
Equivalently, the vector is modeled as the “softmax” function of the inner products between the embedding of neuron and the embeddings of all candidate neurons . The softmax output is non-negative and sums to one so that can be interpreted as a discrete probability distribution over assignments of neuron .
We also find the most probable correspondence between the two sets of neurons by solving a maximum weight bipartite matching problem where the weights are given by the inner products between test and template worm embeddings. This is a classic combinatorial optimization problem, and it can be solved in polynomial time using the Hungarian algorithm .
The fDLC model returns two sets of outputs to the end user. One is the algorithm’s estimate of the most probable labels for each neuron in the test worm; i.e. the solution to the maximum weight bipartite matching problem described above. The other is an ordered list of alternative candidate labels for each individual neuron in the test worm and their probabilities ranked from most to least probable.
The model was trained on synthetic animals derived from recordings of 12 individuals. The model was trained only once and the same trained model was used throughout this work.
Training is as follows. We performed supervised learning with ground truth labels of neuron names provided by the synthetically generated data. A cross-entropy loss function was used. If neuron and neuron has the same name, the cross-entropy loss function favors the model to output . If neuron and neuron have different names, the loss function favors the model to output . The model was trained for 12 hours on a 2.40 GHz Intel machine with NVIDIA Tesla P100 GPU.
Evaluating model performance and comparing against other models
Minimum confidence threshold
Accuracy of the fDLC model is calculated by including all neural matches between the test and template regardless of confidence, except for in the case of Fig. 3. In that analysis only, we sought to reduce the number of matches to better compare to the NeRVE method which only matches 80% of the neurons. We therefore excluded matches below a minimum confidence threshold of 0.05 from our calculation of accuracy. The tradeoff between accuracy and coverage is explored further in Fig. 4G.
Coherent Point Drift
We use Coherent Point Drift (CPD) Registration  as a baseline with which to compare our model’s performance. In our implementation, CPD is used to find the optimal non-rigid transformation to align the test worm with respect to the template worm. We then calculated the distance for each pair of the neurons from the transformed test worm and the template worm. We used the Hungarian algorithm  to find the optimal correspondence that minimizes the total squared distance for all matches.
The recently developed NeuroPAL strain  expresses four different genetically encoded fluorescent proteins in specific expression patterns so that a human can better identify neurons across animals. Manual human annotation based on these expression patterns serves as ground truth in evaluating our model’s performance throughout this work. In Fig. 5B we also explored combining color information with our fDLC model. To do so we developed a simple color matching model that operated in parallel to our position based fDLC model. Outputs of both models were then combined to predict the final correspondence between neurons.
Our color matching model consists of two steps: First, the intensity of each of the color channels is normalized by the total intensity. Then the similarity of color for each pair of neurons is measured as the inverse of the Kullback–Leibler divergence between their normalized color features.
To calculate the final combined matching matrix, we add the color similarity matrix to the position matching log probability matrix from our fDLC model. The similarity matrix of color is multiplied by a factor . We chose so that the amplitude of values in the similarity matrix of color is comparable to our fDLC output. We note the matching results are not particularly sensitive to the choice of . The most probable matches are obtained by applying Hungarian algorithm on the combined matching matrix.
Source code in Python is provided for the model, for the simulator, and for training and evaluation. A jupyter notebook with a simple example is also provided. Code is available at https://github.com/XinweiYu/fDLC_Neuron_ID
Fast deep learning correspondence accurately labels neurons across synthetic individuals
We developed a fast deep learning correspondence (fDLC) model that seeks to find the correspondence between configurations of C. elegans neurons in different individuals or in the same individual across time (Fig. 1). We used a deep learning artificial neural network architecture, called the transformer architecture , that specializes at finding pairs of relations in datasets. The transformer architecture identified similarities across spatial relations of neurons in a test and a template to identify correspondences between the neurons.
Within a single individual, neural positions vary as the worm moves, deforms, and changes its orientation and pose. Across isogenic individuals, there is an additional source of variability that arises from the animal’s development. In practice, further variability also arises from experimental measurements: neuron positions must first be extracted from fluorescent images, and slight differences in label expression, imaging artifacts, and optical scattering all contribute to errors in segmenting individual neurons.
We created a simulator to model these different sources of variability and used it to generate realistic pairs of empirically derived synthetic animals with known correspondence between their neurons for training our model (Fig. 1B, E). The simulator took configurations of neuron positions from real worms as inputs and then scaled and deformed them, forced them to adopt different poses sampled from real worms, and then introduced additional sources of noise to generate many new synthetic individuals. We then trained our fDLC model on these experimentally derived synthetic individuals of different sizes and poses.
Training our model on the empirically derived synthetic data offered advantages compared to experimentally acquired data. First, it allowed us to train on larger datasets than would otherwise be practical. We trained on synthetic individuals, but only seeded our simulator with unlabeled neural configurations from experimentally acquired recordings of 12 individuals. Second, we did not need to provide ground truth correspondence because the simulator instead generates its own ground truth correspondence between synthetic individuals, thereby avoiding a tedious and error prone manual step. We note we still use human annotation for validation, but this requires much smaller data sets. Third, by using large and varied synthetic data, we force the model to generalize its learning to a wide range of variabilities in neural positions and we avoid the risks of overtraining on idiosyncrasies specific to our imaging conditions or segmentation. Overall, we reasoned that training with synthetic data should make the model more robust and more accurate across a wider range of conditions, orientations and animal poses than would be practical with experimentally acquired datasets.
We trained our fDLC model on synthetic individuals (Fig. 1C and methods) and then evaluated its performance on 200 additional held-out synthetic pairs of individuals which had not been accessible to the model during training (Fig. 2). Model performance was evaluated by calculating the accuracy of the models’ predicted correspondence with respect to the ground truth in pairs of synthetic individuals. One individual is called the “test” and the other is the “template”. Accuracy is reported as the fraction of neurons, present in both the test and the template, that were correctly matched. Our fDLC model achieved 96.6% average accuracy on the 200 pairs of held-out synthetic individuals. We compared this performance to that of Coherent Point Drift (CPD) , a classic registration method used for automatic cell annotation. CPD achieved 30.8% mean accuracy on the same held-out synthetic individuals. Our measurements show that the fDLC model significantly outperforms CPD at finding correspondence in synthetic data. For the rest of the work, we use experimentally acquired human annotated data to evaluate performance.
fDLC accurately tracks neurons within an individual across time
We next evaluated the fDLC model’s performance at tracking neurons within an individual over time, as is needed, for example, to measure calcium activity in moving animals [32, 22]. We evaluated model performance on an experimentally acquired calcium imaging recording of a freely moving C. elegans (strain AML32) from  in which a team of human experts had manually tracked and annotated neuron positions over time. This recording had been excluded from the set of recordings used by the simulator. We sampled neuron configurations from 201 different time points during this recording to form 200 pairs of configurations upon which to evaluate the fDLC model. Each pair consisted of a test and template. The template was always from the same time point , while the test was taken to be any of the other 200 time points. We applied the pre-trained fDLC model to the 200 pairs of neuron configurations and compared the model’s predicted correspondence to the ground truth from manual human tracking (Fig. 3). Across the 200 pairs, the fDLC model showed an average accuracy of 80.0%. We emphasize that the fDLC model achieved this high accuracy on tracking a real worm using only neuron position information even though it is trained exclusively on synthetic data.
We compared the performance of our fDLC model to that of CPD Registration, and to Neuron Registration Vector Encoding and clustering (NeRVE), a classical computer vision model that we had previously developed specifically for tracking neurons within a moving animal over time  (Fig. 3B). NeRVE only matches approximately 80% of neurons, so for the sake of comparison, in this analysis (but not others), we only consider comparable number of matches from the fDLC, see “minimum confidence threshold” in “Materials and Methods.”) fDLC clearly outperformed CPD achieving 80.0% accuracy compared to CPD’s 58.5%.
Both CPD and fDLC predict neural correspondence of a test configuration by comparing only to a single template. In contrast, the NeRVE method takes 100 templates, where each one is a different neuron configuration from the same individual, and uses them all to inform its prediction. The additional templates give the NeRVE method extra information about the range of possible neural configurations made by the specific individual whose neurons are being tracked. We therefore compared the fDLC model both to the full NeRVE method and also to a restricted version of the NeRVE method in which NeRVE had access only to the same single template as the fDLC model. (Under this restriction the NeRVE method no longer clusters and the method collapses to a series of gaussian mixture model registrations .) In this way, we could compare the two methods when given the same information. fDLC’s mean performance of 80.0% was very similar but statistically significantly more accurate than the restricted NeRVE model (mean 79.1%, , Wilcoxon signed rank test). The full NeRVE model that had access to additional templates outperformed the fDLC model slightly (82.0% , Wilcoxon signed rank test). We conclude that the fDLC model is suitable for tracking individual neurons across time and performs similarly to the NeRVE method.
In the following sections, we further show that the fDLC method is orders of magnitude faster than NeRVE. Moreover, unlike NeRVE which can only be used within an individual, fDLC is also able to predict the much more challenging neural correspondence across individuals.
fDLC is fast enough for future real-time tracking
Because it relies on an artificial neural network, the fDLC model finds correspondence for a set of neurons faster than traditional methods (Table 3). From the time that a configuration of segmented neurons is loaded onto a GPU, it takes only an average of 10 ms for the fDLC model to label all neurons on a 2.4 GhZ Intel machine with an NVIDIA Tesla P100 GPU. If not using a GPU, the model labels neurons in 50 ms. In contrast, on the same hardware it takes CPD 930 ms and it takes NeRVE on average over 10 seconds. The fDLC model may be a good candidate for potential closed-loop tracking applications because its speed of 100 volumes per second is an order of magnitude faster than the 6 to 10 volumes per second recording rate typically used in whole-brain imaging of freely moving C. elegans [22, 32]. We note that for a complete closed-loop tracking system, fast segmentation algorithms will also be needed in addition to the fast registration and labeling algorithms presented here. The fDLC model is agnostic to the details of the segmentation algorithm so it is well suited to take advantage of fast segmentation algorithms when they are developed.
|fDLC [this work]||0.01|
Table shows the measured time per volume required for different models to predict neural correspondence of a single volume. Time required is measured after neuron segmentation is complete and a configuration of neural positions has been loaded into memory. The same hardware is used for all models.
The fDLC model uses built-in libraries to parallelize the computations for labeling a single volume, and this contributes to its speed. In particular, each layer of the neural network contains thousands of artificial neurons performing the same computation. Computations for each neuron in a layer can all be performed in parallel and modern GPUs have as many as 3,500 CUDA cores.
In practice, the method is even faster for post-processing applications (not-realtime) because it is also parallelizable at the level of each volume. Labeling one volume has no dependencies on any previous volumes and therefore each volume can be processed simultaneously. The number of volumes to be processed in parallel is limited only by number the of volumes that can be loaded onto the memory of a GPU. When tracking during post-processing in this work, we used 32 volumes simultaneously.
fDLC accurately finds neural correspondence across individuals
Having shown that fDLC performs well at identifying neurons within the same individual, we wanted to address its capability to identify neurons across different animals. Identifying corresponding neurons across individuals is crucial for studying the nervous system. However, finding neural correspondence across individuals is more challenging than within an individual because there is variability in neuronal position from both the animal’s movement as well as from development. To evaluate the fDLC model’s performance at finding neural correspondence across individuals, we applied the same synthetically-trained fDLC model to a set of 11 manually annotated NeuroPAL worms.
NeuroPAL worms contain genetically encoded multicolor fluorescently labeled neural landmarks to aid neural identification . For each of the 11 Neuropal recordings, neurons were automatically segmented and manually annotated based on the neuron’s position and color features as described in  (see Fig. 4A,B). Across the 11 animals, a human assigned a ground-truth label to a mean of 43% of segmented head neurons, providing approximately 58 labeled neurons per animal (Fig. 4C). The remaining neurons were not confidently labeled and thus not included in this analysis. We selected as template the recording that contained the largest of confidently labeled human annotated neurons. We evaluated our model by comparing its predicted correspondence between neurons in the other 10 test datasets and this template, using only position information. All 11 ground-truth recordings were held-out in that they were not involved in the generation of the synthetic data that had been used to train the model.
We applied the synthetically-trained fDLC model to each pair of held-out NeuroPAL test and template recordings and calculated the accuracy as the fraction of labeled neurons present in both the test and the template that was correctly matched. Across the 10 pairs of NeuroPAL recordings using position information alone, the fDLC model had an accuracy of 65.8%, significantly higher than the CPD method’s accuracy of 54.7% (, Wilcoxon signed-rank test).
For each neuron in the test recording, the fDLC model returns the probability with which that neuron corresponds to each possible neuron in the template, . We wondered whether we could better use the likelihood information about potential matches generated by the algorithm. A Hungarian algorithm finds the most probable match by considering all s for all neurons in the test. By default we use this best match in evaluating performance. But the s also provide the user with a list of alternative matches ranked by their likelihood. We therefore also assessed the accuracy for the top 3 most likely candidate neurons. In this context, we define accuracy as the fraction of neurons for which the correct match is included in the list of the top 3 candidates. As before, the denominator is taken to be the number of ground-truth labeled neurons that appear in both test and template. When considering the top 3 neurons, the fDLC model achieves an accuracy of 84.4% using only position information.
Validating on an alternative dataset
Data quality, selection criteria, human annotation, hardware and preprocessing can all vary from lab-to-lab making it challenging to directly compare methods. To validate our model against different measurement conditions and to allow for a direct comparison with another recent method, we applied our fDLC model to a previously published dataset of 9 NeuroPAL individuals  (Fig. 4F). This public dataset used different imaging hardware and conditions and was annotated by human experts from a different group. On this public dataset, our method achieved 78.9% accuracy while CPD achieved 59.4%. When assessing the top 3 candidate accuracy, the fDLC model performance was 91.3%. The fDLC model performance was overall higher on the published dataset compared to our newly collected dataset presented here. This suggests that our method performs well when applied to real-world datasets in the literature.
|CRF (open atlas)||%||9|||
|CRF (data driven atlas)||74%||9|||
Table lists reported mean accuracy of different models evaluated on the same publicly accessible dataset from . indicates the number of template-test pairs used to calculate accuracy. (CRF method uses an atlas as the template, whereas we randomly take one of the 9 individuals and designate that as the template). CPD and fDLC performance on this dataset are also shown in Fig. 4F.
We further compared the fDLC model to a recent model called Conditional Random Fields (CRF) from . We compared the reported performance of the CRF model on the published dataset to the performance of the fDLC model evaluated on the same dataset (Table 4). The CRF model has two variants an “open atlas” and a “data-driven atlas.” fDLC accuracy is notably higher than the reported CRF performance for either variant, although we are unable to test for statistical significance. We note the fDLC method also offers other advantages compared to the CRF approach in that the fDLC method is optimized for speed and avoids the need to transform the worm into a canonical coordinate system. Taken together, we conclude that the fDLC model’s accuracy compares favorably to that of the CRF model while also providing other advantages.
Tradeoff between performance and coverage
Our model predicts correspondence for all pairs of neurons, but the confidence of the predicted label varies by pair. Therefore, there is an inherent tradeoff between the accuracy of matched neurons (performance) and the fraction of neurons matched (coverage).The probability in Equation 1 (in Materials and methods) serves as an estimate of the confidence level. By default the fDLC model matches all neurons regardless of confidence and, with the exception of Fig. 3, that is how performance is evaluated in this work. (In Fig. 3 a minimum confidence was imposed to better compare to the NeRVE approach, which also only matches a subset of neurons. Details are described in methods section “Minimum confidence threshold.”)
In practice, it is often desirable to focus on only those neurons that are matched with very high confidence. We therefore tried restricting the model to only include matches above certain confidence thresholds. As the confidence threshold increased, model accuracy increased at the cost of the number of labeled neurons (Fig. 4G). For a threshold of we observed an increase in the accuracy on our Neuropal dataset from 65.8% to 73.0% while only reducing coverage from 100% to 78.0% of the total neuron matches (intersection of test and template). By increasing the threshold to 0.99, the accuracy increases to 79.2% but now only 38.1% of the neurons are matched. This feature of the model allows the experimenter to transparently tune the trade-off between prediction accuracy and coverage.
Incorporating color information
Our method only takes positional information as input to predict neural correspondence. However, when additional features are available, the position-based predictions from the fDLC model can be combined with predictions based on other features to improve overall performance. As demonstrated in , adding color features from a NeuroPAL strain can reduce the ambiguity of predicting neural correspondence. We applied a very simple color model to calculate the similarity of color features between neuron in the test recording to every possible neuron in the template. The color model returns matching probabilities, based on the Kullback-Liebler divergence of the normalized color spectra in a pair of candidate neurons (details described in Materials and methods). The color model is run in parallel to the fDLC model (Fig. 5A). Overall matching probabilities that incorporate both color and position information are calculated by combining the color matching probabilities with the position probabilities . The Hungarian algorithm is run on the combined matching algorithm to predict the best matches.
Adding color information increased the fDLC model’s accuracy from 65.8% to 76.5% (Fig. 5B) when evaluated on our dataset, and improved the accuracy in every recording pair. The top 3 candidate labels attained 92.4% accuracy. Accuracy was calculated from a comparison to human ground truth labeling, as before.
We chose a trivially simple color model in part to demonstrate the flexibility with which the fDLC model framework can integrate information about other features. Since our simple color model utilized no prior knowledge about the distributions of colors in the worm, we would expect a more sophisticated color model, for example, the statistical model used in , to do better. And indeed that model evaluated on a different dataset is reported to have a higher performance with color than our model on our dataset (86% reported accuracy in  compared to 77% for the fDLC evaluated here). But that model also performs much worse than fDLC when both are restricted to use only neural position information (50% reported accuracy for  compared to 66% for the fDLC). Together, this suggests the fDLC model framework can take advantage of additional feature information like color and still perform well when such information is missing.
The fDLC model finds neural correspondence within and across individuals with an accuracy that compares favorably to other methods. The model focuses primarily on identifying neural correspondence using position information alone. The fDLC model framework also makes it easy to integrate other features. We demonstrated that color information could be added by integrating the fDLC model with a simple color model to increase overall accuracy. We expect that performance would improve further with a more sophisticated color model that takes into account the statistics of the colors in a NeuroPAL worm .
The fDLC model framework offers a number of additional advantages beyond accuracy. First, it is versatile and general. The same pre-trained model performed well at both tracking neurons within a freely moving individual across time and at finding neural correspondence across different individuals. It achieved even higher accuracy on a publicly accessible dataset acquired on different hardware with different imaging conditions from a different group. This suggests that the framework should be applicable to many real-world datasets.
Second, the model reports the relative confidence in its estimate of each neuron’s correspondence. An experimenter can therefore tune the overall accuracy of the model’s predictions, at the cost of leaving some neurons unmatched, by simply setting a desired minimum confidence. Similarly, the model also provides probability estimates of all possible matches for each neuron, not just the most likely. This allows an experimenter to consider a collection of possible matches such as the top 3.
In contrast to previous methods, an advantage of the fDLC method is that it does not require the worm to be straightened, axis aligned, or otherwise transformed into a canonical coordinate system. This eliminates an error-prone and often manual step. Instead, the fDLC model finds neural correspondence directly from neural position information even in worms that are in different poses or orientations.
Importantly, the model is trained entirely on synthetic data, which avoids the need for large experimentally acquired ground truth datasets to train the artificial neural network. Acquiring ground truth neural correspondence in C. elegans is time consuming, error prone, and often requires manual hand annotation. The ability to train the fDLC model with synthetic data derived from measurements alleviates this bottleneck and makes the model attractive for use with other organisms with stereotyped nervous systems where ground truth datasets are similarly challenging to acquire.
The model is also fast and finds neural correspondence of a new neural configuration in 10 ms. This speed is sufficient to keep up with real-time calcium imaging applications and will be valuable for future efforts combining large scale population calcium imaging with closed loop targeted optogenetic stimulation in freely moving animals [14, 26, 9, 25]. We note that to be used in a real-time closed loop application, our fDLC model would need to be combined with faster segmentation algorithms because current segmentation algorithms are too slow for real-time use. Because segmentation can be easily paralellized, we expect that faster segmentation algorithms will be developed soon.
Many of the advantages listed here stem from the fDLC model’s use of the transformer architecture  in combination with supervised learning. The transformer architecture, with its origins in natural language processing, is well suited to find spatial relationships within a configuration of neurons. By using supervised learning on empirically-derived synthetic training data of animals in a variety of different poses and orientations, the model is forced to learn relative spatial features within the neurons that are informative for finding neural correspondence across many postures and conditions. Finally, the transformer architecture leverages recent advances in GPU parallel processing for speed and efficiency, which will help pave the way for real-time, closed-loop optogenetic experiments in freely moving worms.
We thank Ev Yemini and Oliver Hobert of Columbia University for strain OH15262. We acknowledge productive discussions with John Murray of University of Pennsylvania. This work used computing resources from the Princeton Institute for Computational Science and Engineering. Research reported in this work was supported by the Simons Foundation under awards SCGB #543003 to AML and SCGB #697092 to SWL; by the National Science Foundation, through an NSF CAREER Award to AML (IOS-1845137) and through the Center for the Physics of Biological Function (PHY-1734030); and by the National Institute of Neurological Disorders and Stroke of the National Institutes of Health under award numbers R21NS101629 to AML and 1R01NS113119 to SWL.
Xinwei Yu, Conceptualization, Software, Formal analysis, Investigation, Methodology, Visualization, Validation, Writing-original draft, Writing—review and editing; Matthew S. Creamer, Investigation, Collected Data, Writing—review and editing; Francesco Randi, Resources, Designed optics and related software libraries, Writing—review and editing; Anuj Sharma, Resources, Writing—review and editing, Performed all transgenics; Scott W. Linderman, Conceptualization, Funding acquisition, Writing—review and editing; Andrew M Leifer, Conceptualization, Supervision, Funding acquisition, Writing—original draft, Project administration, Writing—review and editing
- (2016) Layer normalization. External Links: Cited by: Architecture.
- (2006) Pattern recognition and machine learning. . Cited by: Introduction.
- (2007) Rock, Paper, and Scissors: extrinsic vs. intrinsic similarity of non-rigid shapes. Cited by: Introduction.
- (2019-03) A probabilistic atlas for cell identification. arXiv:1903.09227 [q-bio]. External Links: Cited by: Introduction, Introduction.
- (2020-03) Automated annotation of cell identities in dense cellular images. bioRxiv, pp. 2020.03.10.986356 (en). External Links: Cited by: Abstract, Introduction, Introduction, Datasets, Preprocessing and Segmentation, Fig 4, Validating on an alternative dataset, Validating on an alternative dataset, Table 4.
- (2018-08) The CeNGEN Project: The Complete Gene Expression Map of an Entire Nervous System. Neuron 99 (3), pp. 430–433 (English). External Links: Cited by: Introduction.
- (2016-06) Deep Residual Learning for Image Recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Las Vegas, NV, USA, pp. 770–778. External Links: Cited by: Architecture.
- (2011-08) Robust Point Set Registration Using Gaussian Mixture Models. IEEE Transactions on Pattern Analysis and Machine Intelligence 33 (8), pp. 1633–1645. External Links: Cited by: Introduction, fDLC accurately tracks neurons within an individual across time.
- (2012-10) Controlling interneuron activity in Caenorhabditis elegans to evoke chemotactic behaviour. Nature 490 (7419), pp. 273–277 (eng). External Links: Cited by: Introduction, Discussion.
- (1955-03) The Hungarian method for the assignment problem. Naval Research Logistics Quarterly 2 (1-2), pp. 83–97 (en). External Links: Cited by: Calculating probabilities for potential matches, Coherent Point Drift.
- (2020-07) Robust single neuron tracking of calcium imaging in behaving Hydra. bioRxiv, pp. 2020.06.22.165696 (en). External Links: Cited by: Introduction.
- (2018-07) Tracking Activity In a Deformable Nervous System With Motion Correction and Point-Set Registration. bioRxiv, pp. 373035 (en). External Links: Cited by: Introduction.
- (2020) YASS: Yet Another Spike Sorter applied to large-scale multi-electrode array recordings in primate retina. bioRxiv. Cited by: Introduction.
- (2011-02) Optogenetic manipulation of neural activity in freely moving Caenorhabditis elegans. Nature Methods 8 (2), pp. 147–152. External Links: Cited by: Introduction, Discussion.
- (2009-09) A 3D digital atlas of C. elegans and its application to single-cell analyses. Nature methods 6 (9), pp. 667–672. External Links: Cited by: Introduction.
- (2016) Non-rigid point set registration by preserving global and local structures. IEEE Transactions on Image Processing 25 (1), pp. 53–64. External Links: Cited by: Introduction.
- (2020) Deep learning tools for the measurement of animal behavior in neuroscience. Current opinion in neurobiology 60, pp. 1–11. Cited by: Introduction.
- (2010-12) Point Set Registration: Coherent Point Drift. IEEE Transactions on Pattern Analysis and Machine Intelligence 32 (12), pp. 2262–2275. External Links: Cited by: Introduction, Generating synthetic point clouds for training, Coherent Point Drift, Fast deep learning correspondence accurately labels neurons across synthetic individuals, Table 3.
- (2020-07) Extracting neural signals from semi-immobilized animals with deformable non-negative matrix factorization. (en). Cited by: Introduction, Introduction.
- (2021) Neuron Matching in C. elegans With Robust Approximate Linear Regression Without Correspondence. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 2837–2846 (en). Cited by: Introduction.
- (2017-05) Automatically tracking neurons in a moving and deforming brain. PLOS Computational Biology 13 (5), pp. e1005517. External Links: Cited by: Introduction, Datasets, Preprocessing and Segmentation, Fig 3, fDLC accurately tracks neurons within an individual across time, fDLC accurately tracks neurons within an individual across time, Table 3.
- (2016-02) Whole-brain calcium imaging with cellular resolution in freely behaving Caenorhabditis elegans. Proceedings of the National Academy of Sciences 113 (8), pp. E1074–E1081 (en). External Links: Cited by: Introduction, Introduction, Imaging, fDLC is fast enough for future real-time tracking, fDLC accurately tracks neurons within an individual across time.
- (2017) Automatic differentiation in pytorch. Cited by: Overview and input.
- (2008) Straightening Caenorhabditis elegans images. Bioinformatics 24 (2), pp. 234. Cited by: Introduction.
- (2014) Simultaneous optogenetic manipulation and calcium imaging in freely moving C. elegans. Frontiers in Neural Circuits 8, pp. 28. External Links: Cited by: Introduction, Discussion.
- (2011-02) Real-time multimodal optical control of neurons and muscles in freely behaving Caenorhabditis elegans. Nature methods 8 (2), pp. 153–158. External Links: Cited by: Introduction, Discussion.
- (1976-08) Post-embryonic development in the ventral cord of Caenorhabditis elegans. Philosophical Transactions of the Royal Society of London. Series B, Biological Sciences 275 (938), pp. 287–297 (eng). External Links: Cited by: Introduction.
- (1986-11) The structure of the nervous system of the nematode Caenorhabditis elegans. Philosophical Transactions of the Royal Society of London. B, Biological Sciences 314 (1165), pp. 1–340 (en). External Links: Cited by: Introduction, Introduction.
- (2019-07) An annotation dataset facilitates automatic annotation of whole-brain activity imaging of C. elegans. bioRxiv, pp. 698241 (en). External Links: Cited by: Introduction.
- (2020) Statistical atlas of c. elegans neurons. In Medical Image Computing and Computer Assisted Intervention – MICCAI 2020, pp. 119–129. Cited by: Introduction, Introduction.
- (2017) Attention is all you need. External Links: Cited by: Introduction, Overview and input, Architecture, Fast deep learning correspondence accurately labels neurons across synthetic individuals, Discussion.
- (2016-02) Pan-neuronal imaging in roaming Caenorhabditis elegans. Proceedings of the National Academy of Sciences 113 (8), pp. E1082–E1088 (en). External Links: Cited by: Introduction, Introduction, fDLC is fast enough for future real-time tracking, fDLC accurately tracks neurons within an individual across time.
- (2018) Deep-learning-based flexible pipeline for segmenting and tracking cells in 3d image time series for whole brain imaging. bioRxiv. External Links: Cited by: Introduction, Introduction, Introduction.
- (2020) Connectomes across development reveal principles of brain maturation in c. elegans. bioRxiv. External Links: Cited by: Introduction.
- (2020-12) NeuroPAL: a multicolor atlas for whole-brain neuronal identification in C. elegans. Cell (en). External Links: Cited by: Abstract, Introduction, Datasets, Imaging, Preprocessing and Segmentation, Color Model, Fig 4, fDLC accurately finds neural correspondence across individuals, Incorporating color information, Incorporating color information, Discussion.