Deep neural networks to register and annotate cells in moving and deforming nervous systems

When analyzing neuronal calcium imaging data, it is essential to accurately link neurons’ identities over time to construct reliable calcium traces. This task is challenging in freely moving animals where animal movement warps the nervous system on sub-second timescales. Therefore, we sought to develop a fast and accurate method to perform non-rigid image registration that can deal with these warping images. Previous studies have described such methods for non-rigid registration of point clouds (e.g. neuron centroid positions; Nguyen et al., 2017; Yu et al., 2021; Wu et al., 2022; Ma et al., 2016), but, as we describe below, we found that performing full image alignment allows for higher accuracy neuron alignment.

To solve this task, we used a previously described network architecture (Fu et al., 2020; Hu et al., 2018) that takes as input a pair of 3-D images (i.e. volumes of fluorescent imaging data of the head of the worm) from different timepoints of the same neural recording (Figure 1A). The network is tasked with determining how to warp one 3-D image (termed the ‘moving image’) so that it resembles the other 3-D image (termed the ‘fixed image’). Specifically, the network outputs a dense displacement field (DDF), a pixel-wise coordinate transformation function designed to indicate which points in the moving and fixed images are the same (see Materials and methods). The moving image is then transformed through this DDF to create a warped moving image, which should look like the fixed image. This network was selected because its LocalNet architecture (a modified 3-D U-Net) allows it to do the feature extraction and image reconstruction necessary to solve the task. To train and evaluate the network, we used data from freely moving animals expressing both pan-neuronal NLS-GCaMP and NLS-tagRFP, but only provided the tagRFP images to the network, since this fluorophore’s brightness should remain static over time. Since Euler registration of images (rotation and translation) is simple, we performed Euler registration on the images using a GPU-accelerated grid search prior to inputting them into the network. During training, we also provided the network with the locations of the centroids of matched neurons found in both images, which were available for these training and validation data since we had previously used gradient descent to solve those registration problems (‘registration problem’ here is defined as a single image pair that needs to be aligned) and link neurons’ identities (Atanas et al., 2023). The centroid locations are only used for network training and are not required for the network to solve registration problems after training. The loss function that the network was tasked with minimizing had three components: (1) image loss: negative of the Local squared zero-Normalized Cross-Correlation (LNCC) of the fixed and warped moving RFP images, which takes on a higher value when the images are more similar (hence making the image loss more negative); (2) centroid alignment loss: the average of the Euclidean distances between the matched centroid pairs, where lower values indicate better alignment; and (3) regularization loss: a term that increases the overall loss the more that the images are deformed in a non-rigid manner (in particular, penalizing image scaling and scrambling of adjacent pixels; see Materials and methods).

Figure 1 with 1 supplement see all

Download asset Open asset

We trained and validated the network on 5176 and 1466 image pairs (from 57 and 22 animals), respectively, over 300 epochs (Figure 1B). We then evaluated network performance on a separate set of 447 image pairs reserved for testing that were recorded from different animals. On average, the network improved the Normalized Cross-Correlation (NCC; an overall metric of image similarity, where 1.0 means images are identical; see Materials and methods for equation) from 0.577 in the input image pairs to 0.947 in the registered image pairs (Figure 1C shows example of centroid positions; Figure 1D shows image example; Figure 1E shows both; Figure 1F shows quantification of pre-aligned and registered). After alignment, the average distance between aligned centroids was 1.45 pixels (Figure 1F). These results were only modestly different depending on the animal or the exact registration problem being solved (Figure 1—figure supplement 1A–C).

To determine which features of the network were critical for its performance, we trained three additional networks, using loss functions where we omitted either the centroid alignment loss, the regularization loss, or the image loss. In the first case, the network would not be able to learn based on whether the neuron centroids were well-aligned; in the second case, there would be no constraints on the network performing any type of deformation to solve the task; in the third case, the deformations that the network learned to apply could only be learned from the alignment of the centroids, not the raw tagRFP images. Registration performance of each network was evaluated using the NCC and centroid distance, which quantify the quality of tagRFP image alignment and centroid alignment, respectively (Figure 1F). While the NCC scores were similar for the full network and the no-regularization and no-centroid alignment networks, other performance metrics like centroid distance were significantly impaired by the absence of centroid alignment loss or regularization loss (Figure 1E–F). This suggests that in the absence of centroid alignment loss or regularization loss, the network learns how to align the tagRFP images, but does so using unnatural deformations that do not reflect how the worm bends. In the case of the no-image loss network, all performance metrics, including both image and centroid alignment, were impaired compared to the full network (Figure 1F). This suggests that requiring the network to learn how to warp the RFP images enhances the network’s ability to learn how to align the neuron positions (i.e. centroids). Thus, the loss on RFP image alignment can serve to regularize point cloud alignment.

The finding that the centroid positions were precisely aligned by the full network indicates that the centers of the neurons were correctly registered by the network. However, it does not ensure that all of the pixels that comprise a neuron are correctly registered, which could be important for subsequent feature extraction from the aligned images. For example, it is formally possible to have perfect RFP image alignment in a context where the pixels from one neuron in the moving RFP image are scrambled to multiple neuron locations in the warped moving RFP image. In fact, we observed this in our first efforts to build such a network, where the loss function was only composed of the image loss. As an additional control to test for this possibility in our trained networks, we examined the network’s performance on data from a different strain that expresses pan-neuronal NLS-mNeptune (analogous to the pan-neuronal NLS-tagRFP) and eat-4::NLS-GFP, which is expressed in ~40% of the neurons in the C. elegans head (Figure 1G shows example image). If the pixels within the neurons are being correctly registered, then applying image registration to the GFP channel for these image pairs should result in highly correlated images (i.e. a high NCC value close to 1). If the pixels within neurons are being scrambled, then these images should not be well-aligned. This test is somewhat analogous to building a ground-truth labeled dataset, but does not require manual labeling and is less susceptible to human labeling errors. We used the DDF that the network output based on the pan-neuronal mNeptune data to register the corresponding eat-4::NLS-GFP images from the same timepoints and found that this resulted in high-quality GFP image alignment (Figure 1H). In contrast, while the no-centroid alignment and no-regularization networks output a DDF that successfully aligned the RFP images, applying this DDF to corresponding GFP images resulted in poor GFP image registration (Figure 1H shows that the no-centroid alignment network aligns the RFP channel, but not the GFP channel, in the eat-4::NLS-GFP strain). This further suggests that these reduced networks lacking centroid alignment or regularization loss are aligning the RFP images through unnatural image deformations. Together, these results suggest that the full Brain Alignment Neural Network (BrainAlignNet) can perform non-rigid registration on pairs of images from freely moving animals.

The registration problems included in the training, validation, and test data above were pulled from a set of registration problems that we had been able to solve with gradient descent (example images in Figure 1—figure supplement 1D). These problems did not include the most challenging cases, for example, when the two images to be registered had the worm’s head bent in opposite directions (although we note that they did include substantial non-rigid deformations). We next asked whether a network trained on arbitrary registration problems, including those that were not solvable with gradient descent (example images in Figure 1—figure supplement 1E), could obtain high performance. For this test, we also omitted the Euler registration step that we performed in advance of network training, since the goal was to test whether this network architecture could solve any arbitrary C. elegans head alignment problem. For this analysis, we used the same loss function as the successful network described above. We also increased the amount of training data from 5176 to 335,588 registration problems. The network was trained for 300 epochs. However, the test performance of the network was not high in terms of image alignment or centroid alignment (Figure 1—figure supplement 1F). This suggests that additional approaches may be necessary to solve these more challenging registration problems. Overall, our results suggest that, provided there is an appropriate loss function, a deep neural network can perform non-rigid registration problems to align neurons across the C. elegans head with high speed and single-pixel-level accuracy.

The above results suggest that BrainAlignNet can perform high-quality image alignments. The main value of performing these alignments is to enable accurate linking of neurons over time. To test whether performance was sufficient for this, we incorporated BrainAlignNet into our existing image analysis pipeline for brain-wide calcium imaging data and compared the results to our previously described pipeline, which used gradient descent to solve image registration (Atanas et al., 2023). This image analysis pipeline, the Automated Neuron Tracking System for Unconstrained Nematodes (ANTSUN), includes steps for neuron segmentation (via a 3D U-Net), image registration, and linking of neurons’ identities (Figure 2A). Several steps are required to link neurons’ identities based on registration of these high-complexity 3D images. First, image registration defines a coordinate transformation between the two images, which is then applied to the segmented neuron ROIs, warping them into a common coordinate frame. To link neuron identities over time, we then build an N-by-N matrix (where N is the number of all segmented neuron ROIs at all timepoints in a given recording, i.e. approximately # ROIs * # timepoints) with the following structure: (1) Enter zero if the ROIs were in an image pair that was not registered (we do not attempt to solve all registration problems, as this is unnecessary); (2) Enter zero if the ROIs were from a registered image pair, but the registration-warped ROI did not overlap with the fixed ROI; and (3) Otherwise, enter a heuristic value indicating the confidence that the ROIs are the same neurons based on several ROI features. These features include similarity of ROI positions and sizes, similarity of red channel brightness, registration quality, a penalty for overly nonlinear registration transformations, and a penalty if ROIs were displaced over large distances during alignment. Finally, custom hierarchical clustering is applied to the matrix to generate clusters consisting of the ROIs that reflect the same neuron recorded at different time points. Calcium traces are then constructed from all of these timepoints, normalizing the GCaMP signal to the tagRFP signal. We term the ANTSUN pipeline with gradient descent registration ANTSUN 1.4 (Atanas et al., 2023; Pradhan et al., 2025) and the version with BrainAlignNet registration ANTSUN 2.0 (Figure 2A).

Figure 2 with 1 supplement see all

Download asset Open asset

We ran a series of control datasets through both versions of ANTSUN to benchmark their results. The first was from animals with pan-neuronal NLS-mNeptune and eat-4::NLS-GFP, described above. The resulting GFP traces from these recordings allow us to quantify the number of timepoints where the neuron identities are not accurately linked together into a single trace (Figure 2B shows example dataset). Specifically, in this strain, this type of error can be easily detected since it can result in a low-intensity GFP neuron (eat-4-) suddenly having a high-intensity value when the trace mistakenly incorporates data from a high-intensity neuron (eat-4+), or vice versa. We computed this error rate, taking into account the overall similarity of GFP intensities (i.e. since we can only observe errors when GFP- and GFP+ neurons are combined into the same trace). For both versions of ANTSUN, the error rates were <0.5%, suggesting that >99.5% of timepoints reflect correctly linked neurons (Figure 2C).

We next estimated motion artifacts in the data collected from ANTSUN 2.0, as compared to ANTSUN 1.4. Here, we processed data from three pan-neuronal GFP datasets. GFP traces should ideally be constant, so the fluctuations in these traces provide an indication of any possible artifacts in data collection or processing (Figure 2—figure supplement 1A shows example dataset). Quantification of signal variation in GFP traces showed similar results for ANTSUN 1.4 and 2.0, suggesting that incorporating BrainAlignNet did not introduce artifacts that impair signal quality of the traces (Figure 2—figure supplement 1B).

We also quantified other aspects of performance on GCaMP datasets (example GCaMP datasets in Figure 2D). ANTSUN 2.0 successfully extracted traces from a similar number of neurons as ANTSUN 1.4 (Figure 2E). However, while ANTSUN 1.4 requires 250 CPU days per dataset for registration, ANTSUN 2.0 only requires 9 GPU hours, reflecting a >600 fold increase in computation speed (Figure 2F). These results suggest that ANTSUN 2.0, which uses BrainAlignNet, provides a massive speed improvement in extracting neural data from GCaMP recordings without compromising the accuracy of the data.

We next sought to examine whether BrainAlignNet could be purposed to align neurons from other animals in distinct imaging paradigms. The hydrozoan jellyfish Clytia hemisphaerica has recently been developed as a genetically and optically tractable model for systems and evolutionary neuroscience (Cunningham et al., 2024). It presents a valuable test of the generalizability of our methods as the neuronal morphology, radially symmetric body plan, and centripetal motor output are all distinct from our C. elegans preparations above.

To test the generalizability of BrainAlignNet, we collected a set of videos in which Clytia were gently restrained under cover glass: neurons were restricted to a single z-plane in an epifluorescence configuration. In this setting, Clytia warps significantly as it moves and behaves, necessitating non-rigid motion correction. This paradigm provides an opportunity to relate neural activity to behavior, but even in this context, there are non-rigid deformations that make neuron tracking over time very challenging. Further, in this preparation, neurons have the potential to overlap in the single z-plane. This makes it important to perform full image alignment on all frames of a video, that is aligning all data to a single frame (note that this is different from the approach in C. elegans, see Figure 2A). We tested whether BrainAlignNet could facilitate this demanding case of neuron alignment.

To assess BrainAlignNet’s performance, we utilized recordings of a sparsely labeled transgenic line where neuron tracking is solvable using classic point-tracking across adjacent frames. This allowed us to determine ground-truth neuron positions (i.e. centroids) that could be used to rigorously quantify alignment accuracy. Specifically, we collected videos of transgenic Clytia medusae expressing mCherry under control of the RFamide neuropeptide promoter, which expresses in roughly 10% of neurons distributed throughout the animal’s body (Weissbourd et al., 2021). We formulated the alignment problem as follows. All images were Euler-aligned to a single reference frame. While Euler registration removes much of the macro displacement of the animal, it is entirely insufficient for trace extraction as it was unable to account for the myriad non-rigid deformations occurring across the jellyfish nerve ring (Figure 3A shows example). The network was as described above, with only slight modifications to mask out DDF deformations in the z axis and ignore regularization across z slices (since the image was 2-D). We then trained on 300 total image pairs from three animals using only image loss and regularization loss (we found that centroid alignment loss was unnecessary in this more constrained 2D search space). We then evaluated image alignment (Figure 3B) and centroid alignment (Figure 3C) on 25,997 frames from three separate animals withheld from the training data. Compared to the Euler-aligned images, BrainAlignNet led to a dramatic increase in image alignment and centroid alignment, ultimately aligning matched centroids within 1.51 pixels of each other on average (Figure 3B–C). Additionally, we found comparable performance across the withheld frames of the training animals (Figure 3—figure supplement 1). These results obtained in jellyfish show that, with only light modification, a version of BrainAlignNet can produce neuron alignment in a radically different organism at a quality level matching that observed in our C. elegans data – aligning matched neurons within ~1.5 pixels of each other.

Figure 3 with 1 supplement see all

Download asset Open asset

We next turned our attention to annotating the identities of the recorded neurons in brain-wide calcium imaging data. C. elegans neurons have fairly stereotyped positions in the heads of adult animals, though fully accurate inference of neural identity from position alone has not been shown to be possible. Fluorescent reporter gene expression using well-defined genetic drivers can provide additional information. The NeuroPAL strain is especially useful in this regard. It expresses pan-neuronal NLS-tagRFP, but also has expression of NLS-mTagBFP2, NLS-CyOFP1, and NLS-mNeptune2.5 under a set of well-chosen genetic drivers (example image in Figure 4A Yemini et al., 2021). With proper training, humans can manually label the identities of most neurons in this strain using neuron position and multi-spectral fluorescence. For most of the brain-wide recordings collected using our calcium imaging platform, we used a previously characterized strain with a pan-neuronal NLS-GCaMP7F transgene crossed into NeuroPAL (Atanas et al., 2023). While freely moving recordings were conducted with only NLS-GCaMP and NLS-tagRFP data acquisition, animals were immobilized at the end of each recording to capture multi-spectral fluorescence. Humans could manually label many neurons’ identities in these multi-spectral images, and the image registration approaches described above could map the ROIs in the immobilized data to ROIs in the freely moving recordings in order to match neuron identity to GCaMP traces.

Figure 4 with 1 supplement see all

Download asset Open asset

Manual annotation of NeuroPAL images is time-consuming. First, to perform accurate labeling, the individual needs substantial training. Even after being trained, labeling all ROIs in one NeuroPAL animal can take 3–5 hr. In addition, different individuals have different degrees of knowledge or confidence in labeling certain cell classes. For these reasons, it was desirable to automate NeuroPAL labeling using datasets that had previously been labeled by a panel of human labelers. In particular, the labels that they provided with a high degree of confidence would be most useful for training an automated labeling network. Previous studies have developed statistical approaches for semi-automated labeling to label neural identity from NeuroPAL images, but the maximum precision that we are aware of is 90% without manual correction (Yemini et al., 2021).

We trained a 3-D U-Net (Wolny et al., 2020) to label the C. elegans neuron classes in a given NeuroPAL 3-D image. As input, the network received four fluorescent 3-D images from the head of each worm: pan-neuronal NLS-tagRFP, plus the NLS-mTagBFP2, NLS-CyOFP1, and NLS-mNeptune2.5 images that label stereotyped subsets of neurons (Figure 4A). During training, the network also received the human-annotated labels of which pixels belong to which neurons. Humans provided ROI-level labels and the boundaries of each ROI were determined using a previously described neuron segmentation network (Atanas et al., 2023) trained to label all neurons in a given image (agnostic to their identity). Finally, during training, the network also received an array indicating the relative weight to assign each pixel during training (Figure 4B). This was incorporated into a pixel-weighted cross-entropy loss function (lower values indicate more accurate labeling of each pixel), summing across the pixels in a weighted manner. Pixel weighting was adjusted as follows: (1) background was given extremely low weight; (2) ROIs that humans were not able to label were given extremely low weight; (3) all other ROIs received higher weight proportional to the subjective confidence that the human had in assigning the label to the ROI and the rarity of the label (see Materials and methods for exact details). Regarding this latter point, neurons that were less frequently labeled by human annotation received higher weight so that the network could potentially learn how to classify these neurons from fewer labeled examples.

We trained the network over 300 epochs using a training set of 81 annotated images and a validation set of 10 images (Figure 4C). Because the size of the training set was fairly small, we augmented the training data using both standard image augmentations (rotation, flipping, adding Gaussian noise, etc.) and a custom augmentation where the images were warped in a manner to approximate worm head bending (see Materials and methods). Overall, the goal was for this network to be able to annotate neural identities in worms in any posture provided that they were roughly oriented length-wise in a 284 x 120 x 64 (x, y, z) image. Because this Automatic Cell Labeling Network (AutoCellLabeler) labels individual pixels, it was necessary to convert these pixel-wise classifications into ROI-level classifications. AutoCellLabeler outputs its confidence in its label for each pixel, and we noted that the network’s confidence for a given ROI was highest near the center of the ROI (Figure 4D). Therefore, to determine ROI-level labels, we took a weighted average of the pixel-wise labels within an ROI, weighing the center pixels more strongly. The overall confidence of these pixel scores was also used to compute an ROI-level confidence score, reflecting the network’s confidence that it labeled the ROI correctly. Finally, after all ROIs were assigned a label, heuristics were applied to identify and delete problematic labels. Labels were deleted if (1) the network already labeled another ROI as that label with higher confidence; (2) the label was present too infrequently in the network’s training data; (3) the network labeled that ROI as something other than a neuron (e.g. a gut granule or glial cell, which we supplied as valid labels during training and had labeled in training data); or (4) the network confidently predicted different parts of the ROI as different labels (see Materials and methods for details).

We evaluated the performance of the network on 11 separate datasets that were reserved for testing. We assessed the accuracy of AutoCellLabeler on the subset of ROIs with high-confidence human labels (subjective confidence scores of 4 or 5, on a scale from 1 to 5). On these neurons, average network confidence was 96.8% and its accuracy was 97.1%. We furthermore observed that the network was more confident in its correct labels (average confidence 97.3%) than its incorrect labels (average confidence 80.7%; Figure 4E). More generally, AutoCellLabeler confidence was highly correlated with its accuracy (Figure 4F shows a breakdown of cell labeling at different confidences, with an inset indicating accuracy). Indeed, excluding the neurons where the network assigns low (<75%) confidence increased its accuracy to 98.1% (Figure 4—figure supplement 1A displays the full accuracy-recall tradeoff curve). Under this confidence threshold cutoff, AutoCellLabeler still assigned a label to 90.6% of all the ROIs that had high-confidence human labels, so we chose to delete the low-confidence (<75%) labels altogether (see Figure 4—figure supplement 1A for rationale for the 75% cutoff value).

We also examined model performance on data where humans had either low confidence or did not assign a neuron label. In these cases, it was harder to estimate the ground truth. Overall, model confidence was much lower for neurons that humans labeled with low confidence (87.3%) or did not assign a label (81.3%). The concurrence of AutoCellLabeler relative to low-confidence human labels was also lower (84.1%; we note that this is not really a measure of accuracy since these ‘ground-truth’ labels had low confidence). Indeed, overall the network’s concurrence versus human labels scaled with the confidence of the human label (Figure 4G).

We carefully examined the subset of ROIs where the network had high confidence (>75%), but humans had either low confidence or entered no label at all. This was quite a large set of ROIs: AutoCellLabeler identified significantly more high-confidence neurons (119/animal) than the original human labelers (83/animal), and this could conceivably reflect a highly accurate pool of network-generated labels exceeding human performance. To determine whether this was the case, we obtained new human labels by different human labelers for a random subset of these neurons. Whereas some human labels remained low-confidence, others were now labeled with high confidence (20.9% of this group of ROIs). The new human labelers also labeled neurons that were originally labeled with high confidence so that we could compare the network’s performance on relabeled data where the original data was unlabeled, low confidence, or high confidence. AutoCellLabeler’s performance on all three groups was similar (88%, 86.1%, and 92.1%, respectively), which was comparable to the accuracy of humans relabeling data relative to the original high-confidence labels (92.3%; Figure 4H). The slightly lower accuracy on these re-labeled data is likely due to the human labeling of the original training, validation, and testing data being highly vetted and thoroughly double-checked, whereas the re-labeling that we performed just for this analysis was done in a single pass. Overall, these analyses indicate that the high-confidence network labels (119/animal) have similar accuracy regardless of whether the original data had been labeled by humans as unlabelable, low confidence, or high confidence. We note that this explains a subset of the cases where human low-confidence labels were not in agreement with network labels (Figure 4G). Taken together, these observations indicate that AutoCellLabeler can confidently label more neurons per dataset than individual human labelers.

We also split out model performance by cell type. This largely revealed similar trends. Model labeling accuracy and confidence were variable among the neuron types, with highest accuracy and confidence for the cell types where there were higher confidence human labels and a higher frequency of human labels (Figure 4K). For the labels where there were high confidence network and human labels, we generated a confusion matrix to see if AutoCellLabeler’s mistakes had recurring trends (Figure 4—figure supplement 1B). While mistakes of this type were very rare, we observed that the ones that occurred could mostly be categorized as either mislabeling a gut granule as the neuron RMG, or mislabeling the dorsal/ventral categorization of the neurons IL1 and IL2 (e.g. mislabeling IL2D as IL2). Together, these categories accounted for 50% of all AutoCellLabeler’s mistakes. We also observed that across cell types, AutoCellLabeler’s confidence was highly correlated with human confidence (Figure 4—figure supplement 1C), suggesting that the main limitations of model accuracy are due to limited human labeling accuracy and confidence.

To provide better insights into which network features were critical for its performance, we trained additional networks lacking some of AutoCellLabeler’s key features. To evaluate these networks, we considered both the number of high-confidence labels assigned by AutoCellLabeler and the accuracy of those labels measured against high-confidence human labels. Surprisingly, a network that was trained with only standard image augmentations (i.e. lacking the custom augmentation to bend the images in a manner that approximates a worm head bend) had similar performance (Figure 4I). However, a network that was trained without a pixel-weighting scheme (i.e. where all pixels were weighted equally) provided far fewer high-confidence labels. This suggests that devising strategies for pixel weighting is critical for model performance, though our custom augmentation was not important. Interestingly, all trained networks had similar accuracy (Figure 4J) on their high-confidence labels, suggesting that the network architecture in all cases is able to accurately assess its confidence.

We examined whether the full group of fluorophores was critical for AutoCellLabeler performance. This is a relevant question because (i) it is laborious to make, inject, and annotate a large number of plasmids driving fluorophore expression and (ii) the large number of plasmids in the NeuroPAL strain has been noted to adversely impact the animals’ growth and behavior (Atanas et al., 2023; Yemini et al., 2021; Sharma et al., 2024). To test whether fewer fluorescent labels could still facilitate automatic labeling, we trained four additional networks: one that only received the pan-neuronal tagRFP image as input, and three that received pan-neuronal tagRFP plus a single other fluorescent channel (CyOFP, tag-mBFP2, or mNeptune). As we still had the ground-truth labels based on humans viewing the full set of fluorophores, the supervised labels were identical to those supplied to the full network.

We evaluated the performance of these models by quantifying the number of high-confidence labels that each network provided in each testing dataset (Figure 5A) and the accuracy of these labels measured against high-confidence human labels (Figure 5B). We found that all four networks had attenuated performance relative to the full AutoCellLabeler network, which was almost entirely explainable by these networks having lower confidence in their labels, since network accuracy was always consistent with its confidence (Figure 5—figure supplement 1A). Of the four attenuated networks, the tagRFP +CyOFP network performance (107 neurons per animal labeled at 97.4% accuracy) was closest to the full network in its performance. Given that there are >20 mTagBFP2 and mNeptune plasmids in the full NeuroPAL strain, these results raise the possibility that a smaller set of carefully chosen plasmids could permit training of a network with equal performance to the full network that we trained here.

Figure 5 with 1 supplement see all

Download asset Open asset

We did not expect the tagRFP-only network to perform well, since the task of labeling tagRFP-only images is nearly impossible for humans. Surprisingly, this network still exhibited relatively high performance, with an average of 94 high-confidence neurons per animal and 94.8% accuracy on those neurons. On most neuron classes, it behaved nearly as well as the full network, though there are 10–20 neuron classes that it is much worse at labeling, such as ASG, IL1, and RMG (Figure 5C). This suggests a high level of stereotypy in neuron position, shape, and/or RFP brightness for many neuron types. Since this network only requires the red channel fluorescence, it could be used directly on freely moving data, which has only GCaMP and tagRFP channel data. Since the tagRFP-only network was trained only on high-SNR images collected from immobilized animals, we first checked that the network was able to generalize outside its training distribution to single images with lower SNR (example images in Figure 5—figure supplement 1B). It was able to label 79 high-confidence neurons per animal at 95.2% accuracy on the lower SNR images (Figure 5A–B, right). We then investigated whether allowing the network to access different postures of the same animal improved its accuracy. Specifically, we evaluated the tagRFP-only network on 100 randomly selected timepoints in the freely moving data of each animal (example images in Figure 5—figure supplement 1B). We then related these 100 network labels to the human labels, which could be easily determined, since ANTSUN registers free-moving images back to the immobilized NeuroPAL images that had been labeled by humans. We averaged the 100 network labels to obtain the most likely network label for each neuron, as well as the average confidence for that label. To properly compare network versions, we determined how many neurons could be labeled at any given target labeling accuracy – for example, how many neurons the network can label and still achieve 95% accuracy (Figure 5D; changing the threshold network confidence value to include a given label allowed us to determine these full curves). This analysis revealed that averaging network labels across the 100 timepoints improved network performance, though only modestly. These results suggest that single-color labels can be used to train networks to a high level of performance, but additional fluorescence channels further improve performance.

The strong performance of the tagRFP-only network on out-of-domain lower SNR images suggests an impressive ability of the AutoCellLabeler network to generalize across different modalities of data. This raised the possibility that it may be possible to use this network architecture to build a foundation model of C. elegans neuron annotation that works across strains and imaging conditions. As a first step to explore this idea, we investigated to what extent the tagRFP-only network could generalize to other strains of C. elegans besides the NeuroPAL strain. We used our previously described SWF415 strain, which contains pan-neuronal NLS-GCaMP7F, pan-neuronal NLS-mNeptune2.5, and sparse tagRFP expression (Atanas et al., 2023). Notably, the pan-neuronal promoter used in this strain for NLS-mNeptune expression (Primb-1) is distinct from the pan-neuronal promoter that drives NLS-tagRFP expression in NeuroPAL (a synthetic promoter), so the levels of red fluorescence are likely to be different across these strains. Since humans do not know how to label neurons in SWF415 (it lacks the NeuroPAL transgene), we did a more limited analysis by analyzing network labels for a subset of neurons that have highly reliable activity dynamics with respect to behavior (AVA, AVE, RIM, and AIB encode reverse locomotion; RIB, AVB, RID, and RME encode forward locomotion; SMDD encodes dorsal head curvature; and SMDV and RIV encode ventral head curvature; Atanas et al., 2023; Kato et al., 2015; Li et al., 2014; Chalfie et al., 1985; Gordus et al., 2015; Luo et al., 2014; Wang et al., 2020; Ben Arous et al., 2010; Lim et al., 2016; Hallinen et al., 2021; Ray and Gordus, 2025). Specifically, we asked whether neurons labeled in SWF415 recordings with high confidence by the network had the behavior encoding properties typical of the neuron, assessed via analysis of the GCaMP traces from that neuron. Our previously described CePNEM model (Atanas et al., 2023) was used to determine whether each labeled neuron encoded forward/reverse locomotion or dorsal/ventral head curvature. The network provided high-confidence labels for an average of 7.4/21 of these neurons per animal, and the encoding properties of these neurons matched expectations 68% of the time (randomly labeled neurons had a match of 19%). However, it was possible for the network to (i) incorrectly label a neuron as another neuron that happened to have the same encoding; or (ii) correctly label a neuron that CePNEM lacked statistical power to declare an encoding for. We accounted for these effects via simulations (see Materials and methods), which estimated that the actual labeling accuracy of the network on SWF415 was 69% (Figure 5E). This is substantially lower than this network’s accuracy on similar images from the NeuroPAL strain (i.e. the strain used to train the network), where an average of 12.5 of these neurons were labeled per animal with 97.1% accuracy. Nevertheless, this analysis indicates that AutoCellLabeler has a reasonable ability to generalize to strains with different genetic drivers and fluorophores, suggesting that in the future it may be worthwhile to pursue building a foundation model that labels C. elegans neurons across many strains.

Annotation of cell types via supervised learning (using human-provided labels) is fundamentally limited by prior knowledge and human throughput in labeling multi-spectral imaging data. In principle, unsupervised approaches that can automatically identify stereotyped cell types would be preferable. Thus, we next sought to train a neural network to perform unsupervised discovery of the cell types of the C. elegans nervous system (Figure 6A). If successful, these approaches could be useful for labeling of mutant genotypes, new permutations of NeuroPAL, or even related species, where it might be laborious to painstakingly characterize new multispectral lines and generate >100 labeled datasets. In addition, such an approach would be useful in more complex animals that do not yet have complete catalogs of cell types.

Figure 6 with 1 supplement see all

Download asset Open asset

To facilitate unsupervised cell type discovery, we trained a network to register different animals’ multi-spectral NeuroPAL imaging data to one another. Successful alignment of cells across all recorded animals would amount to unsupervised cell type annotation, since the cells that align across animals would comprise a functional ‘cluster’ corresponding to the same cell type identified in different animals (we note though that each cell type, or cluster, here would still need to be assigned a name). The architecture of this network was similar to BrainAlignNet, but the training data here consisted of pairs of four-color NeuroPAL images from two different animals and the network was tasked with aligning all four fluorescent channels (Figure 6B). No cell type positions (i.e. centroids) or neuronal identities were provided to the network during training. Regularization and augmentation were similar to that of BrainAlignNet (see Materials and methods). Training and validation data were comprised of 91 animals’ datasets, which gave rise to 3285 unique pairs for alignment; 11 animals were withheld for testing (the same test set as for AutoCellLabeler). The network was trained over 600 epochs (Figure 6C). In the analyses below, we characterize performance on training data and withheld testing data, describing any differences. We note that, in contrast to the networks described above, high performance on training data is still useful in this case, since the only criterion for success in unsupervised learning is successful alignment (i.e. even if all data need to be used for training to do so). Strong performance on testing data is still more desirable though, since it is less efficient to train different networks over and over as new data are incorporated into the full dataset.

We first characterized the ability of this Unsupervised Cell Discovery Network (CellDiscoveryNet) to align images of different animals. Image alignment was reasonably high for all four fluorescent NeuroPAL channels with a median NCC of 0.80 overall (Figure 6D). Alignment accuracy was nearly equivalent in training and testing data (Figure 6D). We also examined how well the centroid positions of defined cell types were aligned, utilizing our prior knowledge of neurons’ locations from human labels (Figure 6E). We computed this metric only on cell types that were identified with high confidence in both of the images of a given registration problem. The median centroid distance was 7.2 pixels, with similar performance on training and testing data. This was initially rather disappointing, as it suggested that the majority of neurons were not being placed at their correct locations during registration. However, we observed two important properties of the centroid alignments. First, the distribution of centroid distances was bimodal (Figure 6E) – the 20th percentile centroid distance was only 1.4 pixels, which corresponds to a correct neuron alignment. Second, the median centroid distance decreased to 3.3 for registration problems with high (> 90th percentile = 0.85) NCC scores on the images. Together, these observations suggest that CellDiscoveryNet correctly aligns neurons some of the time.

We next sought to differentiate the neuron alignments where CellDiscoveryNet was correct from those where it was incorrect. To accomplish this, we adapted our ANTSUN pipeline (described in Figure 2) to use CellDiscoveryNet instead of BrainAlignNet. This modified ANTSUN 2 U (Unsupervised) takes as input multi-spectral data from many animals instead of monochrome images from different time points of the same animal. This approach then allows us to effectively cluster neurons that might be the same neuron found across different animals. We ran CellDiscoveryNet on pairs of images and used the resulting DDFs to align the corresponding segmented neuron ROIs. We then constructed an N-by-N matrix (N is the number of all segmented neuron ROIs in all animals, i.e. approximately # ROIs * # animals). Entries in the matrix are zero if the two neurons were in an image pair that was never registered or if the two neurons did not overlap at all in the registered image pair. Otherwise, a heuristic value indicating the likelihood that the neurons are the same was entered into the matrix. This heuristic included the same information as in ANTSUN 2.0 (described above), such as registration quality and ROI position similarity. The only difference was that the heuristic for tagRFP brightness similarity was replaced with a heuristic for four-channel color similarity (see Methods). Custom hierarchical clustering of the rows of this matrix then identified groups of ROIs hypothesized to be the same cell type identified in different animals.

To determine the performance of this unsupervised cell type discovery approach, we quantified both the number of cell types that were discovered (i.e. number of clusters) and the accuracy of cell type labeling within each cluster. Here, accuracy was computed by first determining the most frequent neuron label for each cell type, based on the human labels. We then determined the number of correct versus incorrect detections of this cell type for all cells that fell within the cluster. A correct detection was defined to be when the human label for that cell matched the most frequent label for that cell’s cluster. The number of cell types identified and the labeling accuracy are directly related: more permissive clustering identifies more cell types, but at the cost of lower accuracy. A full curve revealing this tradeoff is shown in Figure 6F (see Materials and methods). Based on this curve, we selected the clustering parameters that identified 125 cell types with 93% labeling accuracy. Not every cell type is detected in every animal. On the testing data, the CellDiscoveryNet-powered ANTSUN 2 U roughly matched human-level performance in terms of accuracy and number of neurons labeled per animal (Figure 6G–H). However, it fell slightly short of AutoCellLabeler (Figure 6G–H). Overall, this analysis reveals that CellDiscoveryNet facilitates unsupervised cell type discovery with a high level of performance, matching trained human labelers. Specifically, this network, combined with our clustering approach, can identify 125 canonical cell types across animals and assign cells to these cell type classes with 93% accuracy.

We examined whether the accuracy of cell identification was different across cell types or across animals. Figure 6I shows the accuracy of labeling for each cell type (see Figure 6—figure supplement 1A for per-animal accuracy). Indeed, results were mixed: some cell types had highly accurate detections across animals (eg: OLQD and RME), whereas a smaller subset of cell types was detected with lower accuracy (eg: AIZ and ASG), and yet other cell types were harder to assess accuracy due to a smaller number of human labels (e.g.: AIM and I4). In addition, there were five clusters that did not contain a sufficient number of human-labeled ROIs to be given a cell type label (<3 cells in these clusters had matching human labels; these are labeled ‘NEW 1’ through ‘NEW 5’). To examine which neurons these might correspond to, we examined the high-confidence AutoCellLabeler labels for ROIs in these clusters. This produced enough labels to categorize four of these five clusters as SAAD, SMBD, VB02, and VB02. The repeated VB02 label is likely an indication of under-clustering (i.e. the two VB02 clusters should have been merged into the same cluster). The identity of the fifth cluster was unclear, as the ROIs in that cluster were not well labeled by either humans or AutoCellLabeler.

Finally, we examined whether CellDiscoveryNet was able to label cells not detected via AutoCellLabeler. Specifically, we determined the fraction of the cells detected by CellDiscoveryNet that were labeled by AutoCellLabeler, which was 86%. The new unsupervised detections (the remaining 14%) included: new labels for cells that were otherwise well-labeled by AutoCellLabeler (e.g.: M3); the detection and labeling of several cell types that were uncommonly labeled by AutoCellLabeler (e.g.: RMEV); and the previously mentioned cell type that could not be identified based on human labels. This suggests that the unsupervised approach that we describe here is able to provide cell annotations that were not possible via human labeling or AutoCellLabeler. Overall, these results show how CellDiscoveryNet, together with our clustering approach, can identify >120 cell type classes that occur across animals and assign cells to these classes with 93% accuracy. This unsupervised discovery of cell types occurs in the absence of any human labels.

All data were collected from 1-day-old adult hermaphrodite C. elegans animals raised at 22 C on nematode growth medium (NGM) plates with E. coli strain OP50.

For the GCaMP-expressing animals without NeuroPAL, two transgenes were present: (1) flvIs17: tag-168::NLS-GCaMP7F+NLS-tagRFPt expressed under a small set of cell-specific promoters: gcy-28.d, ceh-36, inx-1, mod-1, tph-1(short), gcy-5, gcy-7; and (2) flvIs18: tag-168::NLS-mNeptune2.5. This resulting strain, SWF415, has been previously characterized (Atanas et al., 2023).

For the GCaMP-expressing animals with NeuroPAL, two transgenes were present in the strain: (1) flvIs17: described above; and (2) otIs670: low-brightness NeuroPAL. This resulting strain, named SWF702, has been previously characterized (Atanas et al., 2023).

The animals with eat-4::NLS-GFP and tag-168::NLS-GFP were also previously described (Atanas et al., 2023). As is described in the strain list, tag-168::NLS-mNeptune2.5 was also co-injected with each of these plasmids to generate the two strains: SWF360 (eat-4::NLS-GFP; tag-168::NLS-mNeptune2.5) and SWF467 (tag-168::NLS-GFP; tag-168::NLS-mNeptune2.5).

We provide here a list of these four strains:

• SWF415 flvIs17[tag-168::NLS-GCaMP7F, gcy-28.d::NLS-tag-RFPt, ceh-36:NLS-tag-RFPt, inx-1::tag-RFPt, mod-1::tag-RFPt, tph-1(short)::NLS-tag-RFPt, gcy-5::NLS-tag-RFPt, gcy-7::NLS-tag-RFPt]; flvIs18[tag-168::NLS-mNeptune2.5]; lite-1(ce314); gur-3(ok2245)

• SWF702 flvIs17; otIs670 [low-brightness NeuroPAL]; lite-1(ce314); gur-3(ok2245)

• SWF360 flvEx450[eat-4::NLS-GFP, tag-168::NLS-mNeptune2.5]; lite-1(ce314); gur-3(ok2245)

• SWF467flvEx451[tag-168::NLS-GFP, tag-168::NLS-mNeptune2.5]; lite-1(ce314); gur-3(ok2245)

For jellyfish recordings, we utilized the following strain, which has been previously described (Weissbourd et al., 2021):

Clytia hemisphaerica: RFamide::NTR-2A-mCherry.

Data used to train and evaluate the models include previously published datasets (Atanas et al., 2023; Pradhan et al., 2025; Dag et al., 2023) and newly collected data. These animals were recorded under similar recording conditions to those described in our previous study (Atanas et al., 2023). There were two types of datasets collected, relevant to this study: freely moving GCaMP/TagRFP data and immobilized NeuroPAL data.

Briefly, all neural data (free-moving and NeuroPAL) were acquired on a dual light-path microscope that was previously described (Atanas et al., 2023). The light path used to image GCaMP, mNeptune, and the fluorophores in NeuroPAL at single-cell resolution is an Andor spinning disk confocal system built on a Nikon ECLIPSE Ti microscope. Laser light (405 nm, 488 nm, 560 nm, or 637 nm) passes through a 5000 rpm Yokogawa CSU-X1 spinning disk unit (with Borealis upgrade and dual-camera configuration). For imaging, a 40 x water immersion objective (CFI APO LWD 40 X WI 1.15 NA LAMBDA S, Nikon) with an objective piezo (P-726 PIFOC, Physik Instrumente [PI]) was used to image the worm’s head (Newport NP0140SG objective piezo was used in some recordings). A dichroic mirror directed light from the specimen to two separate sCMOS cameras (Zyla 4.2 PLUS sCMOS, Andor) with in-line emission filters (525/50 for GCaMP/GFP, and 570 longpass for tagRFP/mNeptune in freely moving recordings; NeuroPAL filters described below). The volume rate of acquisition was 1.7 Hz (1.4 Hz for the datasets acquired with the Newport piezo).

For recordings, L4 worms were picked 18–22 hr before the imaging experiment to a new NGM agar plate seeded with OP50 to ensure that we recorded 1-day-old adults. Animals were recorded on a thin, flat NGM agar pad (2.5 cm x 1.8 cm x 0.8 mm). On the four corners of this agar pad, we pipetted a single layer of microbeads with diameters of 80 µm to alleviate the pressure of the coverslip (#1.5) on the worm. Animals were transferred to the agar pad in a drop of M9, after which the coverslip was added.

For NeuroPAL data collection, animals were immobilized via cooling to 1 °C, after which multi-spectral information was captured. We obtained a series of images from each recorded animal while the animal was immobilized (this has been previously described Atanas et al., 2023):

(1-3) Isolated images of mTagBFP2, CyOFP1, and mNeptune2.5: We excited CyOFP1 with a 488 nm laser at 32% intensity using a 585/40 bandpass filter. mNeptune2.5 was then recorded with a 637 nm laser at 48% intensity under a 655LP-TRF filter (to not contaminate this recording with TagRFP-T emissions). mTagBFP2 was then isolated with a 405 nm laser at 27% intensity and a 447/60 bandpass filter.

An image with TagRFP-T, CyOFP1, and mNeptune2.5 (all ‘red’ markers) in one channel, and GCaMP7f in another channel. As described in our previous study, this image was used for neuronal segmentation and registration to both the freely moving recording and individually isolated marker images. We excited TagRFP-T and mNeptune2.5 with a 561 nm laser at 15% intensity and CyOFP1 and GCaMP7f with a 488 nm laser at 17% intensity. TagRFP-T, mNeptune2.5, and CyOFP1 were imaged using a 570LP filter, and GCaMP7f was data collection used a 525/50 bandpass filter.

Each of these images was recorded for 60 time points. We functionally increased the SNR for each of the images by registering all 60 time points for a given image to one another and averaging the transformed images, creating an average image. To facilitate manual labeling of these datasets, we made a composite, three-dimensional RGB image where we set the mTagBFP2 image to blue, CyOFP1 image to green, and mNeptune2.5 image to red as done by Yemini et al., 2021. We manually adjusted channel intensities to optimally match their manual.

Clytia medusae were mounted on a glass slide in artificial seawater (ASW Weissbourd et al., 2021) and gently compressed under a glass coverslip, using a small amount of Vaseline as a spacer, so that swimming and margin folding motions were visible but limited. The imaging ROI included a segment of the nerve rings and subumbrella. Widefield epifluorescent images of transgenic mCherry fluorescence were acquired at 20 Hz, with 5 ms exposure times, using an Olympus BX51WI microscope, a photometrics Prime95B camera, and Olympus Cellsens software. Videos without visible motion, or with significant lateral drift, were discarded.

All data are freely and publicly available at the following link:

• https://doi.org/10.7910/DVN/8UE0H9

All code is freely and publicly available (use main/master branches):

• BrainAlignNet for worms: https://github.com/flavell-lab/BrainAlignNet, flavell-lab, 2024a and https://github.com/flavell-lab/DeepReg, flavell-lab, 2024b.

• BrainAlignNet for jellyfish: https://github.com/flavell-lab/BrainAlignNet_Jellyfish, flavell-lab, 2025a.

• GPU-accelerated Euler registration: https://github.com/flavell-lab/euler_gpu, flavell-lab, 2024c.

• ANTSUN 2.0: https://github.com/flavell-lab/ANTSUN, flavell-lab, 2024d branch v2.1.0; see also https://github.com/flavell-lab/flv-c-setup, and https://github.com/flavell-lab/FlavellPkg.jl/blob/master/src/ANTSUN.jl, for auxiliary package installation.

• AutoCellLabeler: https://github.com/flavell-lab/pytorch-3dunet, flavell-lab, 2020 and https://github.com/flavell-lab/AutoCellLabeler, flavell-lab, 2024g.

• CellDiscoveryNet: https://github.com/flavell-lab/DeepReg, flavell-lab, 2024b.

• ANTSUN 2 U: https://github.com/flavell-lab/ANTSUN-Unsupervised, flavell-lab, 2024h.

All materials (strains, etc) are freely available upon request.

Our main custom modifications to the DeepReg network focus on the design of the loss function. In particular, we implemented a new supervised centroid alignment loss component and new regularization loss sub-components. Overall, the loss function consists of three major components:

• Image loss $L_I$ captures the difference between the warped moving image and the ground-truth fixed image.

• Centroid alignment loss $L_C$ is a supervised portion of the loss function. Given pre-labeled centroids corresponding to ground-truth information about neuron positions in the fixed and moving images, this loss component captures the difference between the predicted moving centroids and the ground-truth moving centroids.

• Regularization loss $L_R$ captures the prior that the “simplest” DDF that achieves the desired transform outcome is the best. For example, it’s implausible that a pair of neurons that start close together end up on opposite sides of the worm, so a DDF that generates such a transformation would have a high value of regularization loss.

The total loss is then computed as $Loss=w_I{L}_i+w_C{L}_C+w_R{L}_R$. We set $w_I=1$, $w_C=0.1$, and $w_R=1$.

During training, input data was subject to augmentation. We used random affine transformations for augmentation. Each transformation was generated by perturbing the corner points of a cube by random amounts and computing the affine transformation resulting in that perturbation. The same transformation was then applied to the moving image, fixed image, moving centroids, and fixed centroids.

BrainAlignNet was trained using the Adam optimizer with a learning rate of ${10}^{-4}$.

The full configuration file we used during network training is available at https://github.com/flavell-lab/BrainAlignNet/tree/main/configs, flavell-lab, 2024i.

We integrated BrainAlignNet into our previously described ANTSUN pipeline (Atanas et al., 2023; Pradhan et al., 2025; also applied in Dag et al., 2023). Briefly, the pipeline: (i) performs some image pre-processing such as shear correction and cropping; (ii) segments the images into neuron ROIs via a 3D U-Net; (iii) finds time points where the worm postures are similar; (iv) performs image registration to define a coordinate mapping between these time points; (v) applies that coordinate mapping to the ROIs; (vi) constructs an ROI similarity matrix storing how likely different ROIs are to correspond to the same neuron; (vii) clusters that matrix to extract neuron identity; (viii) maps the linked ROIs onto the GCaMP data to extract neural traces; and (ix) performs some postprocessing such as background subtraction and bleach correction to extract neural traces.

The differences in ANTSUN 2.0 compared with our previously published version of this pipeline, ANTSUN 1.4, are that in ANTSUN 2.0 we use BrainAlignNet to perform image registration rather than the gradient descent-based elastix, and we modified the heuristic function used to construct the ROI similarity matrix. We only replaced the freely moving registration with BrainAlignNet; the immobilized registrations, channel alignment registration, and freely moving to immobilized registration are still performed with elastix. These remaining elastix-based registrations are much less computationally expensive, taking only about 2% of the total computation time of the original ANTSUN 1.4 pipeline. They will also likely be replaced with BrainAlignNet in a future release of ANTSUN, after further diagnostics and controls are run.

The heuristic function used to compute the ROI similarity matrix was updated to add additional terms specific to BrainAlignNet, including regularization and an additional ROI displacement term that serves to implement our prior that ROIs which moved less far in the registration are more likely to be correctly registered. Letting $i$ and $j$ be two different ROIs in our recording at time points $t_i$ (moving) and $t_j$ (fixed), the full expression for the ROI similarity matrix is:

where:

$R_{t_i{t}_j}$ is 1 if there exists a registration mapping $t_i$ to $t_j$, and 0 otherwise.

$d_i$ is the displacement of the centroid of ROI $i$ under the DDF registration between $t_i$ and $t_j$.

$q_{t_i{t}_j}$ is the registration quality, computed as the NCC of warped moving image $t_i$ vs fixed image $t_j$.

$r_{ij}$ is the fractional overlap of warped moving ROI $i$ and fixed ROI $j$ (intersection/max size).

$a_{ij}$ is the absolute difference in marker channel activity (i.e. tagRFP brightness) between ROIs $i$ and $j$, normalized to mean activity at the corresponding timepoints $t_i$ and $t_j$.

$c_{ij}$ is the distance between the centroids of warped moving ROI $i$ and fixed ROI $j$.

$n_{t_i{t}_j}$ is the (unweighted) nonrigid penalty loss of the DDF registration from $t_i$ to $t_j$.

$w_i$ are weights controlling how important each variable is.

Additionally, the matrix is forced to be symmetrical by setting $M_{ji}=M_{ij}$ whenever $M_{ji}=0$ and $M_{ij}\ne 0$. It is also sparse since $R_{t_i{t}_j}$ and $r_{ij}$ are usually 0. Finally, there are two additional hyperparameters in the clustering algorithm, $w_7$ and $w_8$. $w_7$ controls the minimum height the clustering algorithm will reach (effectively, $w_7$ is a cap on how low $M_{ij}$ values can get, or how low the heuristic value can fall before determining that the ROIs are not the same neuron) and $w_8$ controls the acceptable collision fraction (a collision is defined by a cluster containing multiple ROIs from the same timepoint, which should not happen since each neuron should correspond to only one ROI at each time point).

We determined the weights $w_i$ by performing a grid search through 2912 different combinations of weights on three eat-4::NLS-GFP datasets. To evaluate the outcome of each combination, we computed the error rate (rate of incorrect neuron linkages) and number of detected neurons. The error rate was computed as previously described by Atanas et al., 2023: since the strain eat-4::NLS-GFP expresses GFP in some but not all neurons, we can quantify registration errors as instances where a GFP-positive neuron lacked GFP in a time point and vice versa, as these correspond to neuron mismatches. We then selected the combination of parameters that maximizes the number of detected neurons while minimizing the error rate. One eat-4::NLS-GFP dataset (the one shown in Figure 2) was used as a withheld testing animal to determine this optimal set of parameters. The pan-neuronal GFP and pan-neuronal GCaMP animals were not included in this parameter search.

The values of the parameters we used were:

To assess the performance of our AutoCellLabeler network on the SWF415 strain, we could not compare its labels to human labels since humans do not know how to label neurons in this strain. Therefore, we used functional information about neuron activity patterns to assess accuracy of the network. We used our previously described CePNEM model to do this (Atanas et al., 2023). Briefly, CePNEM fits a single neural activity trace to the animal’s behavior to extract parameters about how that neuron represents information about the animal’s behavior. CePNEM fits a posterior distribution for each parameter, and statistical tests run on that posterior are used to determine encoding of behavior. For example, if nearly all parameter sets in the CePNEM posterior for a given neuron have the property that they predict the neuron’s activity is higher when the animal is reversing, then CePNEM would assign a reversal encoding to that neuron.

By doing this analysis in NeuroPAL animals where the identity of each neural trace is known, we have previously created an atlas of neural encoding of behavior (Atanas et al., 2023). This atlas revealed a set of neurons that have consistent encodings across animals: AVA, AVE, RIM, and AIB encode reverse locomotion; RIB, AVB, RID, and RME encode forward locomotion; SMDD encodes dorsal head curvature; and SMDV and RIV encode ventral head curvature. Based on this prior knowledge, we decided to quantify the fraction $f$ of labeled neurons with the expected activity-behavior coupling. For example, if AutoCellLabeler labeled 10 neurons as AVA and 7 of them encoded reverse locomotion when fit by CePNEM, this fraction would be 0.7.

However, this fraction is not necessarily an accurate estimate of AutoCellLabeler’s accuracy. For example, it might have been possible for AutoCellLabeler to mislabel a neuron as AVA that happened to encode reverse locomotion in that animal, thus making the incorrect label appear accurate. On the other hand, CePNEM is limited by statistical power and can sometimes fail to detect the appropriate encoding. This could make a correct label appear inaccurate.

To correct for these factors, we ran a simulation analysis to try to estimate the fraction $p$ of labels that were correct. To do this, we iterated through every one of AutoCellLabeler’s labels that was one of the consistent-encoding neuron classes (i.e. one of the neurons listed above). In each simulation, we assign labels to neurons in the following manner: with probability $p_{sim}$ (i.e. the fraction of labels estimated by our simulation to be correct), the label was reassigned to a random neuron that was given that label by a human in a NeuroPAL animal (at confidence 3 or greater); with probability $1-p_{sim}$, the label was reassigned to a random neuron in the same (SWF415) animal. In this way, the simulation controls for both of the possible inaccuracies outlined above. Then the fraction $f_{sim}$ of labeled neurons with the expected encoding was computed for each simulation. 1000 simulation trials were run for each value of $p_{sim}$, which ranged from 0 to 100 – the mean and standard deviation of these trials are shown in Figure 5E. We then computed the probability $p_{sim}$ for which $f_{sim}$ was in closest agreement to $f$, which was 69% (dashed vertical line). This is our estimate for the ground-truth correct label probability $p$.

To convert the CellDiscoveryNet registration outputs into neuron labels across animals, we created a modified version of our ANTSUN image processing pipeline. We skipped the pre-processing steps since the images were already pre-processed, used 102 four-channel images instead of 1600 one-channel images, set the registration graph to be the complete graph (except each pair of images is only registered once and not once in each direction), substituted BrainAlignNet with CellDiscoveryNet for the nonrigid registration step, and skipped the trace extraction steps of ANTSUN (stopping after it computed linked neuron IDs).

We also modified the heuristic function in the matrix that was subject to clustering to better account for the nature of this multi-spectral data. Specifically, we removed the marker channel brightness heuristic $a_{ij}$ since brightness of neurons relative to the mean ROI is not likely to be well conserved across different animals. We replaced it with a more problem-specific heuristic: color. Specifically, the color $C_i$ of an ROI $R_i$ was defined as the 4-vector of its brightness in each of the four channels, normalized to the average brightness of that ROI across the four channels. We then define

where $i,j$ each indicates an ROI label.

In this way, $a_{ij}$ will be small if the ROIs have similar colors and large if they have different colors. We use this new color-based $a_{ij}$ in the same way in the heuristic function that we used the original brightness-based $a_{ij}$, except that we set its weight $w_4=7$.

We did not run hyperparameter search on any of the other weight parameters $w_i$ for this dataset to avoid overfitting to the 102 animals included in it, instead leaving them all at their default values from the original ANTSUN 2.0 pipeline (with the one exception of $w_8$ which we set to 0 here in light of having much fewer animals than we did timepoints). We hypothesize that performance may increase even further upon hyperparameter search, although this would likely require considerably more data for testing. The only exception was that we varied parameter $w_7$, which controls the precision vs recall tradeoff. Larger values of $w_7$ result in more, but less accurate, detected clusters; each cluster corresponds to a single neuron class label. We elected to use a value of $w_7={10}^{-9}$ for all displayed results; the full tradeoff curve is available in Figure 6f.