Multi-layered maps of neuropil with segmentation-guided contrastive learning

Abstract

Maps of the nervous system that identify individual cells along with their type, subcellular components and connectivity have the potential to elucidate fundamental organizational principles of neural circuits. Nanometer-resolution imaging of brain tissue provides the necessary raw data, but inferring cellular and subcellular annotation layers is challenging. We present segmentation-guided contrastive learning of representations (SegCLR), a self-supervised machine learning technique that produces representations of cells directly from 3D imagery and segmentations. When applied to volumes of human and mouse cortex, SegCLR enables accurate classification of cellular subcompartments and achieves performance equivalent to a supervised approach while requiring 400-fold fewer labeled examples. SegCLR also enables inference of cell types from fragments as small as 10 μm, which enhances the utility of volumes in which many neurites are truncated at boundaries. Finally, SegCLR enables exploration of layer 5 pyramidal cell subtypes and automated large-scale analysis of synaptic partners in mouse visual cortex.

Fig. 1. SegCLR.

a, In SegCLR, positive pairs (blue double-headed arrows) are chosen from proximal but not necessarily overlapping 3D views (small blue boxes) of the same segmented cell, while negative pairs (red double-headed arrows) are chosen from different cells. The SegCLR network is trained to produce an embedding vector for each local 3D view such that embeddings are more similar for positive pairs than negative pairs (cartoon of clustered points). b, The input to the embedding network is a local 3D view (4.1 × 4.1 × 4.3 μm at 32 × 32 × 33 nm resolution for human data; 4.1 × 4.1 × 5.2 μm at 32 × 32 × 40 nm resolution for mouse) from the electron microscopy volume, masked by the segmentation for the object at the center of the field of view. An encoder network based on a ResNet-18 is trained to produce embeddings, via projection heads and a contrastive loss that are used only during training. c,d, Visualization via UMAP projection of the SegCLR embedding space for the human temporal cortex (c) and mouse visual cortex (d) datasets. Points for a representative sample of embeddings are shown, colored via 3D UMAP RGB, with the corresponding 3D morphology illustrated for six locations (network segmentation mask input in black, surrounded by 10 × 10 × 10 μm context in gray; masked electron microscopy input data not shown). e,f, Embeddings visualized along the extent of representative human (e) and mouse (f) cells. Each mesh rendering is colored according to the 3D UMAP RGB of the nearest embedding for the surrounding local 3D view. Some axons are cut off to fit. Scale bars: c,d, 5 μm; e,f, 100 μm.

Fig. 2. Subcompartment classification of SegCLR embeddings.

a, Embedding vectors computed across the extent of the electron microscopy datasets can be used as compact inputs to downstream tasks, such as subcompartment classification. Each embedding represents a single local view (~4–5 μm on a side). b, Ground truth examples of axon and dendrite subcompartment classes from the human temporal cortex dataset. The local 3D views for single embeddings are indicated by the wireframe cubes. The embedding network also receives the electron microscopy image data from within the segment mask. c, Embedding clusters from the human cortical dataset visualized via 2D UMAP. Each point is an embedding, colored by its ground truth subcompartment class as judged without reference to the embeddings. d, Evaluation of linear classifiers trained for the subcompartment task on the human dataset. The mean F1 score across classes was computed for networks trained using varying sized subsets of the full available training data. For each training set sample size, mean and standard deviation of multiple subset resamplings are shown (error bars are obscured by the points for larger sample sizes). Light gray points show the best class-wise mean F1 score obtained for any training subset sampled at a given size. The light blue line indicates the performance of a fully supervised ResNet-18 classifier (a convolutional neural network, CNN) trained on the full and subsets of the available training data. Error bars are s.d. (n = 20 subsamples). See Supplementary Table 1 for the number of training samples per class. e, As in c, for the mouse visual cortex dataset and three ground truth classes (axon, dendrite, soma). f, As in d, for the mouse dataset. Error bars are s.d. (n = 20 subsamples). See Supplementary Table 1 for the number of training samples per class.

Fig. 3. Cell type classification of large and small cell fragments via aggregated embeddings.

a, 3D renderings of representative proofread neuron and glia cells, for a selected subset of the types used in the mouse and human datasets. The pyramidal cell axon is cut off to fit. Cells are oriented from white matter (WM) to pia. b, Rendering of representative cutouts from a pyramidal dendrite (top inset) and axon (bottom inset). Different size cutouts are defined by the skeleton node aggregation radius R. c, Cell type classifiers are trained on top of SegCLR embeddings after aggregation into a mean embedding over the cutout. d, Cell typing performance of shallow ResNet classifiers over different aggregation radii for the six labeled cell types in the human dataset. Zero radius corresponds to a single unaggregated embedding node. Error bars are s.d. (n = 20 subsamples). See Supplementary Table 2 for the number of training samples per class. e, Confusion matrix for the 6-class human cell type task at a 10 μm aggregation radius. GT, ground truth. f, Illustration of SegCLR cell type predictions over the extent of a single basket cell from the mouse test set. The orange areas are predicted correctly, while the sparse black areas show mispredictions. g, Cell typing performance for the mouse dataset. The 13-class task (black) uses all of the ground truth-labeled classes, while the 10-class task (green) combines all pyramidal cell labels into a single class. The 6-class task (blue) further reduces the neuronal labels into excitatory and inhibitory groups, comparable to the labels available on the human dataset (d). Error bars are s.d. (n = 20 subsamples). See Supplementary Table 2 for the number of training samples per class. h, Confusion matrix for the mouse 13-class cell type task at a 25 μm aggregation radius. Colored boxes indicate the group of four pyramidal cell types that were collapsed into the 10-class task, and the five excitatory and four inhibitory types collapsed into the 6-class task in g. Abbreviations: AC, astrocyte; BC, basket cell; BPC, bipolar cell; E, excitatory neuron; I, inhibitory interneuron; MC, Martinotti cell; MGC, microglia cell; NGC, neurogliaform cell; OGC, oligodendrocyte cell; OPC, oligodendrocyte precursor cell; P2–6, cortical layer 2–6 pyramidal cell; THLC, thalamocortical axon. Scale bars: a, neuronal 100 μm, glia 25 μm; b,f, 100 μm;.

Fig. 4. Unsupervised exploration of mouse layer 5 pyramidal dendrite embeddings.

a, SegCLR embeddings projected to 3D UMAP space, with two selected axes displayed. Each point represents an embedding (aggregation distance 50 μm) sampled from only the dendrites of mouse layer 5 pyramidal cells. The UMAP data separate into three main clusters. b, Renderings of selected cells, colored to match a for locations for which the nearest embedding fell within cluster 1 (blue) or cluster 2 (red). Projections falling within cluster 1 are strongly associated with the apical dendrite subcompartment, while cluster 2 is strongly associated with a subset of basal dendrites corresponding to cells with a distinct ‘near-projecting’ (NP) morphology (inset). c, Of the cells for which the projections fall within cluster 3, a subset have distinctive axon trajectories consistent with their axons joining a major output tract (left). These ‘tract’ cells also occupy a distinct subregion within cluster 3 (middle, green). Cells occupying the remainder of cluster 3 (middle, purple) consistently lack the axon tract morphology (right). The tract and no-tract groups are also able to be separated in both primary visual area V1 (left group of cells for both tract and no-tract) and higher visual areas (right group of cells for both tract and no-tract). Scale bars, 100 μm.

Fig. 5. OOD input detection via Gaussian processes.

a, We handled OOD inputs by computing prediction uncertainties alongside class labels, and calibrated the uncertainties to reflect the distance between each test example and the training distribution. b, To evaluate OOD detection, we trained classifiers on glial cell type labels, and then evaluated the classifiers on a 50–50 split between glial and OOD neuronal cell types. c, UMAP of locally aggregated embeddings (radius 10 μm) from the human cortical dataset, colored by ground truth-labeled cell type. d, Confusion matrix for a ResNet-2 classifier trained on only the four glia types, with OOD neuronal examples mixed in at test time. e, As in c with the UMAP embeddings now colored by their SNGP uncertainty. The colormap transitions from blue to red at the threshold level used to reject OOD examples in our experiment. f, Confusion matrix for the SNGP-ResNet-2, assembled from 20-fold cross-validations. Examples that exceed the uncertainty threshold are now treated as their own OOD predicted class. g, Spatial distribution of local uncertainty over an unproofread segment that suffers from reconstruction merge errors between the central OPC glia and several neuronal fragments. The uncertainty signal distinguishes the merged neurites (red, high uncertainty or OOD) from the glia cell (blue, low uncertainty) with a spatial resolution of approximately the embedding aggregation distance. Scale bar, 25 μm. See Fig. 3 for cell type abbreviations.

Fig. 6. Quantitative analysis of pre- and post-synaptic partner cell type frequencies.

a, The distribution of synapses by predicted inhibitory versus excitatory presynaptic axonal cell type, over the depth of the mouse visual cortex. Lines L1–6 mark the boundaries of the six cortical layers. b, Representative layer 2–3 pyramidal cell. c, Pyramidal cell input synapse locations, annotated (dots) and colored by predicted inhibitory versus excitatory presynaptic partner cell type. d, Distribution of upstream presynaptic partners for layer 2–3 pyramidal cells in V1 with proofread dendrites (n = 69). Each cell is represented by a set of points describing its ratios of input cell types. The mean and standard deviation of the ratios are also shown (darker point and line). UNC refers to synapses for which classification was uncertain. I-UNC refers to synapses for which only a coarse classification (excitatory versus inhibitory) could be made. Error bars are s.d. e, The ratio of thalamocortical to pyramidal innervation onto pyramidal targets in different cortical layers shows major thalamic input to cortical layer 4 (P23, n = 69; P4, n = 23; P5, n = 136; P6, n = 127). Error bars are s.d. f. The thalamocortical synapse counts over the cortical depth. The thalamocortical synapse counts over the cortical depth. Lines L1–6 as in a. g, The thalamocortical synapse count along the lateral axis through the dataset drops at the boundary between V1 and the HVA. The shaded area shows the approximate projection of the V1–HVA border onto the lateral axis. h, As in c, for downstream postsynaptic partners. i, As in d, distribution of downstream postsynaptic partners for layer 2–3 pyramidal cells in V1 with proofread axons (n = 12). P-UNC refers to synapses where SegCLR was not able to classify a pyramidal subtype with sufficient certainty. Error bars are s.d. j,k, Inhibitory versus excitatory balance of downstream postsynaptic partners with increasing distance along P23 axons (n = 12) (j) and P4 axons (n = 19) (k). Mean and s.e.m. are shown for each distance bucket. l, As in k, divided into individual inhibitory subtypes (n = 19). I, inhibitory; P, pyramidal; UNC, uncertain classification. Scale bars: b, 100 μm; c,h, 10 μm. Source data

Extended Data Fig. 1. Subcompartment classification of SegCLR embeddings trained on a different dataset.

a. The SegCLR model in Fig. 2f was trained on the MICrONS dataset using large training batches. b. The model from (a) was further trained (refined) on the h01 dataset with a small batch size on a single machine. c. The linear classifier was trained with subcompartment labels for the h01 dataset. d. Comparison of linear classifiers trained for the subcompartment task on the h01 dataset using embeddings from the MICrONS model (gray), the MICrONS model with refinement on the h01 dataset (red), and the h01 model (black). For each training set sample size, mean and standard deviation of multiple subset resamplings is shown (error bars are obscured by the points for larger sample sizes). The light blue line indicates the performance of a fully supervised ResNet-18 classifier trained on the full and subsets of the available training data.

Extended Data Fig. 2. Cell type classification of large and small cell fragments via aggregated embeddings.

a. Confusion matrix for the mouse 13-class cell type task at 25 μm aggregation radius. Colored boxes indicate the group of four pyramidal cell types that were collapsed into the 10-class task, and the five excitatory and four inhibitory types collapsed into the 6-class task (similar to Fig. 3h). We restricted this evaluation to dendrites based on the automated subcompartment classification (Fig. 2) using the same classifier as in Fig. 3 that was trained on the entire ground truth (13-class). Consequently, the test set did not contain any examples of glia subtypes and thalamocortical cells. b. As in (a), but restricted to axons. There were no examples of glia subtypes in this test set.

Extended Data Fig. 3. Quantitative analysis of pre- and post-synaptic partner cell type frequencies.

a. Distribution of the distance between the farthest embedding node and the selected embedding node at the synapse (R_agg=25 μm) for presynaptic segments. In most cases, the distance to the farthest embedding node is close to the maximal permitted distance. In a few cases the presynaptic segment was small, with only a few embedding nodes. b. As in (a) for postsynaptic segments. c.d. Inhibitory versus excitatory balance of downstream postsynaptic partners with increasing distance along P5 axons (N = 34) (c) and P6 axons (N = 6) (d). Source data