Self-Supervised Graph Representation Learning for Neuronal Morphologies
Authors: Marissa A. Weis
Laura Pede
Timo Lüddecke
Alexander S Ecker
Venue: Accepted by TMLR ()

Abstract: Unsupervised graph representation learning has recently gained interest in several application domains such as neuroscience, where modeling the diverse morphology of cell types in the brain is one of the key challenges. It is currently unknown how many excitatory cortical cell types exist and what their defining morphological features are. Here we present GraphDINO, a purely data-driven approach to learn low-dimensional representations of 3D neuronal morphologies from unlabeled large-scale datasets. GraphDINO is a novel transformer-based representation learning method for spatially-embedded graphs. To enable self-supervised learning on transformers, we (1) developed data augmentation strategies for spatially-embedded graphs, (2) adapted the positional encoding and (3) introduced a novel attention mechanism, AC-Attention, which combines attention-based global interaction between nodes and classic graph convolutional processing. We show, in two different species and across multiple brain areas, that this method yields morphological cell type clusterings that are on par with manual feature-based classification by experts, but without using prior knowledge about the structural features of neurons. Moreover, it outperforms previous approaches on quantitative benchmarks predicting expert labels. Our method could potentially enable data-driven discovery of novel morphological features and cell types in large-scale datasets. It is applicable beyond neuroscience in settings where samples in a dataset are graphs and graph-level embeddings are desired.

Published in Transactions on Machine Learning Research (06/2023)
Self-Supervised Graph Representation Learning for Neuronal Morphologies
Marissa A. Weis Laura Hansel!, Timo Lüddecke', and Alexander S. Ecker
Institute of Computer Science and Campus Institute Data Science, University of Göttingen, Germany Institute for Theoretical Physics, University of Tübingen, Germany
3 Maa Planck Institute for Dynamics and Self- Organization, Göttingen, Germany
marissa.weiisunite Correspondeece:: marissa. wwwswwwwssnnee weeesss@nnnioocttneee..dd
Reviewed on OpenReview: http:://ppparrrreee...ee//////ad
Abstract
Unsupervised graph representation learning has recently gained interest in several applica- tion domains such neuroscience, where modeling the diverse morphology of cell types in the brain is one of the key challenges. It is currently unknown how many excitatory cortical cell types exist and what their defining morphological features are. Here we present GRAPUDINO, a purely data-driven approach to learn low-dimensional representations of 3D neuronal morphologies from unlabeled large-scale datasets. GRAPHDINO a novel transformer- based representation learning method for self-supervised learning on transformers, we (1) developed data augmentation strategics for spatially-embeedded graphs, (2) adapted the positional encoding and (3) introduced novel attention mechanism, AC-ATTENTION, which combines attention-basee between nodes and classic graph convolutional processing. We show, in two different species and across multiple brain arcas, that this method yiclds morphological cell type cluster- ings that are on par with manual feature-based classification by experts, but without using prior knowledge about the structural features of neurons. Moreover, it outperforms pre- vious approaches on quantitative benchmarks predicting expert labcls. Our method could polentially enable data-driven discovery of novel morphological features and cell types in large-scale datasets. It is applicable beyond neuroscience in settings where samples in dataset are graphs and graph-level embeddings are desired.
graphs. To enable
global interaction
1 Introduction
The brain is structured into different areas that contain diverse types of neurons (Ascoli et al. 2008). The morphology of cortical neurons is highly complex with widely varying shapes. Cell morphology has long been used to classify neurons into cell types (Ramón y Cajal [911], but characterizing neuronal morphologies is still a challenging open research question. Morphological analysis has traditionally been carried out by visual inspection (Ascoli et al. 2008: Defelipe et al.) 2013) or by computing a set of predefined, quantitatively
mcasurable fcatures such as number of branching points (Uylings & Van Polt 2002 Scorcioni al.] 2008
Oberlaender et al.) 2012 Polavaram
al.) 2014 Markram al.) 2015 Lu al. 2015 Gouwens el al. 2019). However, both approaches have deficits: expert assessments have high variance (Defelipe et al,
2013 and the manual definition of morphological features introduces biases (Wang 2018).
thus calling for
more unbiased, data-driven approaches to characterize the morphology of neurons.
Recent advances in recording technologies have greatly accelerated data collection and therefore the amount of data available (MICrONS Consortium et al.] 2023 Ramaswamy et al. 2015 Scala et al.] 2021 Allen Institute 2016 Peng et al. [2021] Winnubst et al.] 2019). These developments have opened the floor for data-driven approaches based unsupervised machine learning methods (Schubert et. al. 2019.

---

Published in Transactions on Machine Learning Research (06/2023)
Elabbady et al. 2022). One form of data representation that is particularly suitable for neurons is rep- resenting the skeleton of a neuron as tree. In such a tree, the root node represents the neuron's 's cell body and the node features are their 3D locations. The availability of a number of such skeleton datasets has recently sparked some work on graph-level representation learning of neuronal morphologies (Laturnus & Berens 2021: Zhao et al.) 2022 Chen et al. 2022a). Following this line of research and work from the graph learning community (Sun et al. 2020 You et al. 2020 we present an unsupervised graph-lovel representation
learning approach.
Our contributions in this paper are fourfold:
We propose a new self-supervised model learn graph-level embeddings for spatial graphs. Unlike previous methods, our approach does not require human annolation or manual feature definition.
We introduce a novel attention module that combines transformer-styl attention and message pass- ing between neighboring nodes as in graph neural networks.
3. We apply this approach to the classification of excitatory neuronal morphologies and show that it produces clusters that are comparable with known excitalory cell types obtained by manual urature- based classification and expert-labeling.
4. We outperform existing approaches based on manual feature engineering and auto-encoding in pre- dicting expert labels.
Our code is available al httpe://eekerrlaa.orr//ceeeeeeee
2 Related Work
2.1
Representation learning for neuronal morphologies
Morphology has been used for a long time to classify neurons by either letting experts visually inspect the cells (Ramón y Cajal 1911: Defelipe et al. 2013) or by specifying expert-defined features that can be extracted and used as input a classifier (Oberlaender
al., 2012 Markram et al., 2015; Kanari et all, Kanari et al. 2019; Gouwens et al.) 2019) (sce Armañanzas & Ascoli (0055) for review).
2017 Wang 2018
Ascoli et al. (2008) made an effort to unify the used expert-defined features.
With the advent of new technologies for microscopic imaging, electrical recording, and molecular analysis such as Patch-soq (Cadwell et al.) 2015) that allow the simultancous recording of transcriptomy, clectro- physiology and morphology of whole cells, several works have explored the prediction of cell types from
multiple modalities (Gala et al. 2021 or one modality from the other (Cadwell et al., 2015 Scala et al.,
2021 Gouwens et al. 2020).
Multiple previous works try to either hand-engineer or learn representation of neuronal morphologies. Laturnus & Berens (2021) propose a generative approach involving random walks in graphs to model neuronal morphologies. Schubert et al. (2019) process 2D projections of morphologies with a convolutional neural network (CNN) to learn low-dimensional representations. [Seshamani et al. (2020) extract local mesh features
around spines and combine them with traditional Sholl analysis (Sholl) 1953). Gouwens et al. (2019) define
a set of morphological fcatures based on graphs and perform hicrarchical clustering on them. We use the
latter as a baseline for classical approach with pre-defined features and Laturnus & Berens 20221 as a baseline of model with learned features.
Concurrent work (Zhao et al. 2022) proposes a contrastive graph neural network to learn neuronal embed- dings with a focus on retrieval efficiency from large-scale databases. Elabbady el. al. (2022) learn represen- tations of neurons based on subcellular features the somatic region of the neurons and show that those fcatures are sufficient for classifying cell types on largo-scale EM datasets. Chen et al. (022b) propose a combination of graph-based processing and manually-defin features to learn embeddings of neuronal mor- phologies using a L.STM-based network and contrastive learning. We compare to the latter in Section 5.8
2

---

Published in Transactions on Machine Learning Research (06/2023)
2.2 Graph Neural Networks (GNNs)
Graph neural networks learn node representations by recursively aggregating information from adjacent nodes as defined by the graph's structure. While early approaches date back over decade (Scarselli et al, 2009), recently numerous new variants were introduced for (semi-) supervised settings: relying on convolution over nodes (Duvenaud et al.) 2015; Hamilton et al.] 2017 Kipf & Welling 2017), using recurrence (Li et al. 2016), or making use of attention mechanisms (Veličković et al. 2018). A representation for the whole graph is often derived by a readout operation on the node representations, for instance averaging. See Dwivedi et al. (2020) for a recent. benchmark on graph neural network architectures.
Transformer-bssed GNNs. Similar to us, Zhang al. (2020) and Dwivedi & Bresson (2021) use trans- former attention to work with graphs. However, Zhang et al. [0020) compute transformer attention over the nodes of sampled subgraphs, while Dwivedi & Bresson (2021) compute the attention only over local neighbors of nodes, which boils down to a weighted message passing that is conditioned on node feature similarity, and trains with supervision. Unlike these previous approaches, we compute the attention between nodes of the global graph and adapt the transformer atiention to consider the adjacency matrix of the graph, which allows the model to take into account both the direct neighbors of a node as well as all other nodes in the graph. Mialon et al. (2021) consider encoding local sub-structures into their node features and leverage kernels on graphs in their attention as relative positional encodings. Their 1-step random walk (RW) kernel is similar to our AC-ATTENTION mechanism, except that the influence of the adjacency in their attention is not learnable. Ying et al. (2021) propose strategies to adapt positional encodings to graphs in order to leverage the structural information of the graphs with transformer attention. Specifically, they propose to use three different structural encodings: (1) a centrality encoding based on the node degree; (2) edge encod- ings based on the edge features and (3) a spatial encoding based on the shortest path between two nodes, For neural skeletons, the centrality encoding is not effective as all the nodes besides the soma have a node degree of two three. Furthermore, the edge encoding is not applicable since in the neuronal graphs do not have edge features. We use Taplacian positional encodings instead as it was shown that they are beneficial to capture structural and positional information (Dwivedi & Bresson 202I) and outperform previously pro- posed positional encodings (Zhang et al., 2020). We did not use any additional positional encodings such as shortest- path encodings (Ying et al. 2021), but they could be easily integrated into our model. Concurrent to our work, Rampášck et al. (2022) proposed a two-stream architccture, which transformer attention and message passing are computed in parallel and then combined after each block. In contrast, we propose one combined attention mechanism that subsumes transformer attention and message passing with a learned trade-off per node between the two settings. [Cen et al. 2022a) incorporate structural information into the transformer attention by extracting a subgraph representation around cach node before computing attention over nodes.
Self-supervised learning on graphs. Self-supervieed learning has proven to be a useful technique for
training image feature extractors (Oord al.) 2018 Chen et al. 2020 Chen & He 2021 Caron et al. 2021
and has been investigated for learning graph {Li et al. 2016 Hassani & Khasahmadi 2020 Qiu et al.
2020
You et al. 2020 Xu et. al. 2021b and node (Veličković et al. 2019) representations. Narayanan et al. 2017
learn graph representations through skip-gram with negative sampling by predicting present sub- graphs.
You et al. ) propose four data augmentations for contrastive learning of graph-level embeddings. [un
et al. (220)) learn graph-level representations in a contrastive way, by predicting a subgraph and a graph representation originate from the same graph. Similarly, Hassani Khasahmadi (2020) put node features of one view in contrast with the graph encoding of second view and vice versa. They build on graph diffusson networks (Klicpera et al. 2019) and only augment the structure of the graph but not the initial node features.
We use Sun et al. (2020) and You et al. (2020) as baseline for graph-level unsupervised representation lcarning. Quu ct al. (2020) propose a generic pre-training method which uses an InfoNCE objective (Oord et al.] 2018) to learn features by telling augmented versions of one subgraph from other subgraphs with random walks as augmentations. Xu et al. (2021) aim to capture local and global structures for whole- graph representation learning. They rely on an EM-like algorithm to jointly train the assignment of graphs to hicrarchical prototypes, the GNN paraneters and the prototypes. Zhu et al. (2021) propose adaptive augmentation, which considers node centrality and importance to generate graph views in contrastive
3

---

Published in Transactions on Machine Learning Research (06/2023)
framework. Similar to our approach, Thakoor et al. (2022) use two encoders of which only one is trained and the other is an exponential moving average of the first. In contrast to our approach, though, their training objective encourages the node embeddings of two augmented versions of the same graph to be similar not the graph-level embedding. Morcover, they use node feature and edge masking as graph augmentations.
Unlike most prior work, we contrast two global views of a graph order to learn a whole-graph representation. Our method operates on spatially embedded graphs, in which nodes correspond points in 3D space. We make use of this knowledge in the choice of augmentations.
3 GraphDINO
We propose GRAPHDINO, method for self- supervised representation learning of graphs. It is inspired by recent progress in sclf-superviced repre- sentation learning of images that has been shown to be competitive to supervised learning without re- lying on labels. The core idea is to enforce that the representations of two augmented versions the same image are close each other in latent space.
A. GraphDINO training
Objective
Au Augmentatins
Encoder
MIP
P1
attract ن
EMA weight update
similerity
Encader
MIP
Graph Input
P2
drop sub-branchess jitte mode
DINO (Caron al, 2021 an implementation of this self-supervised learning framework consist- ing of two encoders with transformer backbone, To avoid mode collapse, only one encoder is di- rectly trained through backpropagation (student) while the weights of the other encoder (teacher) are an exponential moving average (cma) of the stu- dent's weights. The latent representations E R given by the encoders distributions E Ds by multi-layer perceptron (MLP) and subsequent softmax operator over which the cross-entropy loss is computed (Fig. C.4). For further explanation of DINO see Appendix C.1
B. Embedding
Frrisedding
Encoder
Gropp Input
Figure 1: A. Sclf-supervieed learning of low di- mensional vector embeddings 21; #2 that capture the essence of the 3D morphology of individual neurons us- ing GraphDINO. Two augmented "views" of the neu- ron are input into the network, where the weights of one encoder (bottom) are exponential moving aver- age (EMA) of the other encoder (top). The resulting latent embeddings & are projected to probability distri- butions p by a MLP. The objective is to maximize the similarity between P1 and P2- B. An individual neu- ron is represented by its vector embedding as a point in the D-dimensional vector space.
mapped to probability
GRAPHDINO adapts this self-supervised framework to the data domain of graphs (Fig. D In order to use information given by the connectivity of the graph, modify the computation of the trans- former attention to take the graph adjacency matrix into account and use the graph Laplacian as posi- tional encoding.
More specifically, we introduce the following modifications: (1) we incorporate the graph's adjacency ma- trix into the attention computation; (2) we use the graph Laplacian as positional encoding; (3) we define augmentations suitable for spatial graphs.
Input. Input to the network is the 3D shape of a neuron which is represented as undirected graph G= EV. V is the set of nodes { and E the set of undirected edges E = {ee = (v,, ey)} that connect two nodes v; ny. The features of each node U; in the graph are encoded into a token using a linear transformation. These tokens are then used as input to the transformer model, which consists of l multi-head atlention modules with heads each.
Attention bias. Key-value query attention bocame popular in natural language modelling |Vaswani al. 2017) and is now used routinely also in image models (Dosovitskiy et al. 2020
4

---

Published in Transactions on Machine Learning Research (06/2023)
To make use of the information given by the adjacency matrix A RNN of the imput graph - L.. the neighborhood of nodes bias the attention towards A by adding a learned bias to the attention matrix that is conditioned on the input token values:
V, = V,A (o QK" Vik
(19K"
Allention
+ + V,
with Da pa - exp (Wx;),
(1)
where K, Q V are the keys, queries and values which are computed as a learned linear projection of the tokens. (.) denotes the softmax function. 2; E TRD is the token of node fi W R22D a learned weight matrix, X, E R are two factors per node that trade off how much weight is assigned to neighboring nodes versus all other nodes in the graph, and N is the number of nodes.
When = 0 and X = 1, the adjacincy-conddiione attention (AC-ATTENTION) reduces to regular trans- former attention. In the other extreme case (A =0), the attention matrix is dominated by A and the transformer attention computation akin to the message passing algorithm that is commonly used when working with graphs (Scarselli et al. 2009 Duvenaud al. 2015; Gilmer et al., 2017)- GRAPHDINO is more flexible than both regular message passing and point-cloud attention since can decide how much weight is given to the neighbors of node while maintaining the flexibility to allend to all other nodes in the graph as well.
Positional encoding. Following Dwivedi et al. (2020), we use the normalized graph Laplacian matrix L A
as positional encoding, which computed by L matrix, the N x N degree matrix, A the adjacency matrix, and U and A are the matrices of eigenvectors and cigenvalues, respectively. The positional encodings are the first 32 cigenvectors with largest cigenvalues. Positional encodings are added to the nodes features after tokenization.
U'AU, where I is the identity
Overview of data augmentations for spatially- embedded graphs such as neuronal skeletons.
Data augmentation. Data augmentation plays important role in self-supervised learning and needs to be adapted to the data, since it expresses which invariances should be imposed. Given the spatial neuronal data, wo apply the following augmentations: (1) Sub- sampling: We subsample the original graph a fixed number of nodes by randomly re- moving nodes that are not branching points (i.e. nodes connected to more than two other nodes), and connecting the two neighbors of the removed node. This facilitates batch pro- cessing. Furthermore, this augmentation re- tains the global structure of the neuron, while altering local structure in the two views. (2) Rotation: perform random 3D rotation around the y-axis, that is orthogonal to the pia. (3) Jittering:
Table 1:
Level
Augmentations
Graph
(1) Subsampling, (2) Rotation, (5) Translation
Subgraph
(3) Jittering, (4) Branch delction
randomly translate individual node positions by adding Gaussian noise with (0 Ø1). (4) Subgraph
delction: We identify branches that connect Icaf nodes to the last upstream parent node in the graph, i.e. terminal branches that do not split into further branches, and randomly delete n of them starting at
a random location along the branch, while maintaining the overall graph structure. (5) Graph position:
we randomly translate the graph as a whole by adding Gaussian noise with A(, 02) to all nodes. Unlike
Caron et al. (2021), we do not differentiate between the augmentations seen by the student and the teacher network.
5

---

Published in Transactions on Machine Learning Research (06/2023)
4 Data and Experiments
4.1 Synthetic graphs
To demonstrate that our novel attention mechanism is strictly more powerful than simple all-to-all attention on a graph, we generate a synthetic graph dataset. In this dataset, the five classes share similar node locations but differ how the nodes are connected. See Appendix A for the detailed generation process. We use this dataset to test the efficacy of our novel attention mechanism, AC-ATTENTION, and the positional encoding.
4.2 Neuronal and tree graphs
We apply GRAPHDINO to five publicly available neuronal datasets and one non-neuronal dataset.
Blue Brain Project (BBP): Rat somatosensory cortex. Available from the Neocortical Microcircuit Collaboration Portal of the Blue Brain Project (Ramaswamy et al. 2015), the datasel contains 1,389 neurons from juvenile rat somatosensory cortex. We train GRAPHDINO without supervision on the 3D dendritic morphologies of all neurons. For evaluation, we use the subset of 616 neurons which have been labeled by experts into cell types and cortical layer. Of these 616 neurons 286 are excitatory that have been assigned to 14 cell types (Markram et al., 2015). See Appendix [.5.1] for more details on the dataset. We use this dataset evaluate the capability of GRAPIDINO to learn useful representations of neuronal morphologies that align with known cell types, perform ablation experiments the novel graph augmentation strategics and compare to previous work using manually-defined features.
M1 PatchSeq: Mouse motor cortex. The MI PatchSeq dataset contains 275 excitatory and 371 inhibitory cells from MI in adult mouse primary motor cortex (Scala et al.) 2021 2 The excitatory cells (M1 EXC) have been classified into tufted, untufted and other neurons based on their morphology previous study (Laturnus & Berens 2021). We use this dataset compare to previous work that learns morphological embeddings in data-driven way. We train GRAPHDINO without supervision on the 3D dendritic morphologies of the 646 neurons. For evaluation,
a
follow the evaluation protocol and use the same datasct split as Laturnus & Berens (2021). We additionally report the 5-ncarest neighbour accuracy of three additional dataset splits to estimate the variance due to the chosen split, since the test set is very small (60 neurons) and the balanced accuracy is strongly influenced by the morphologically heterogeneous "other" class that is only represented by six samples in the test set (Laturnus Berens 2021
Allen Brain Atlas (ACT): Mouse visual cortex. As part of the Allen Cell Types Database, the dataset contains 510 neurons from the mouse visual cortex with broad coverage of types, layers and transgenic lines See Allen Institute (2016) for details on how the dataset was recorded. It comes with classification of each neuron into spiny, aspiny, or sparsely spiny, where spiny are assumed to be excitatory neurons and all else are inhibitory (Gouwens et al. 2019). Additionally, the cortical layer of each neuron is provided.
Brain Image Library (BIL): Whole mouse brain. The Brain Image Library contains 1,741 recon- structed neurons from cortex, claustrum, thalamus, striatum and other brain areas mice (Peng et al.
2021
Janelia MouseLight (JML): Whole mouse brain. The Janelia MouseLight platform contains 1,200 projection neurons from the motor cortex, thalamus, subiculum, and hypothalamus (Winnubst et al. 2019
Joint training on ACT, BIL and JML. Following Chen et al. 2022b), for joint training of the ACT, BIL and JML datasets, we rotate the neurons such that the first principal component aligned with the
http://miccrciireusss..pp ch/#/nain
http://oalliipess.rriinng
http://ousellgh..jamlla...
6

---

Published in Transactions on Machine Learning Research (06/2023)
y-axis. Chen et al. 2022b) group the neurons of the three datasets ACT, BIL and JML into eleven classes based on the cortical layer or brain region they originate from. They then evaluate their learned embeddings on a subset of six (for BIL) or four classes (for ACT and JML) that have a broad coverage across the datasets.
See Appendix C.5.4 4 for further details.
Botanical Trees. The Trees dataset (Scddl et al. 2021 is highly diverse dataset comprised of 391 skeletons of trees stemmiin from 39 different genuses and 152 species or breedings. The skeletons were extracted from LIDAR scans of the trees. Nodes of the skeletons have a 3D coordinate associated with them. We normalize the data such that the lowest point (start of the tree trunk) is normalized to (0, 0, 0).
4.2.1 Data Preprocessing
Since the objective of GRAPHDINO to learn purely from the 3D dendritic morphology of neurons, we normalize each graph such that the soma location is centered at (0, 0) (no cortical depth information is given to the model). Furthermore, axons are removed for all experiments in the paper, because the reconstruction of axonal arbors of excitatory neurons from light microscopy images is diffccult due to their small thickness and long ranges that they cover (Kanari et al.) 2019) and thus often unreliable. The imput nodes V have features v; = [r, y, z] where ot E R2 are the spatial xyz-coordinales in micrometers [um]]
4.2.2 Training details.
GRAPEDINO is implemented in PyTorch (Pasuke et al. 2019 and trained with the Adam optimizer (Kingma & Ba 2015). The latent dimensionality of 16 for the synthetic graphs and 32 for the neuronal and the botanical tree datasets. For M1 PatchSeq we the hyperparameters used for training on the different datasets. At inference time, the latent embeddings z are extracted from the student network for the unaugmented graphs. We use scipy for fitting Gaussian Mixture models (GMM) and k-nearest neighbor classifiers (kNN) (Pedregosa et al. 2011
latent dimensionality of 64. See C.2 for an overview of
5 Results
We first establish that GraphDINO works on the synthetic graph dataset and show that our novel AC-ATTENTION is necessary for exploiting information from graph connectiv- ity. Second, we show that our novel augmentation strate- gies are suitable for spatially-embeeddeed graphs that are tree- structured and that classical GNN message- sraeragg-passinggg is not suf- ficient when graphs have long-ranging branches. Then, we move to the gradually more complex, biological questions of aspnyy differentiation, cell type recovery and consistency with existing labels. To this end, we employ in total five neu- ronal datasets that encompass two species and range across multiple brain areas. Finally, we compare our model to sev- eral previous works based on manually-defined morphological features as well as approaches with learned features. See Ap- pendix B for the application of GRAPHDINO to non-neuronal dataset.
attention bias
attention bias positional enc.
95.0 3.6
29.3 6.0
20.0 1.0
Figure 2: t-SNE embedding of latent repre- sentations of 3D synthetic graphs shown for one example run per model. Accuracy av- eraged over five random seeds and given mean + standard deviation. " means re- moving one component from the full model.
5.1 AC-Attention recovers information encoded by graph connectivity
We start by demonstrating the efficacy of our novel AC-ATTENTION module. For this experiment, we use the synthetic graph dataset where classes differ in how nodes are connected whereas the distribution of node positions does not vary across classes, Therefore, considering the graph structure necessary to differentiate between the classes (more details Appendix A).. We train GRAPHDINO on the synthetic graph datasct without labels in three configurations: (1) with AC-ATTENTION, (2) with regular transformer attention and (3) with transformer attention and without positional encoding. We asses the quality of the learned
7

---

Published in Transactions on Machine Learning Research (06/2023)
embedding using the ground truth labels. A linear classifier on the learned embeddings achieves test set accuracy of (1) 95% + 4, (2) 29% + 6, and (3) 20% + 1, showing that AC-ATTENTION allows us to capture the structure of spatially embedded graphs when the location of the nodes alone does not provide sufficient information. Removing both AC-ATTENTION and the positional encoding results in the classifier performing al chance level. Using only the positional encoding performs slightly better than chance, because the positional encodings contain some information about node connections through the graph Laplacian. To make full use of the information given by the connectivity of the graphs, using AC-ATTENTION is essential (Fig. 2
5.2 Tailored graph augmentations are well-suited for spatially-embedddd graphs
In supervieed learning, data augmentation is used to obtain two views that define positive input pair. The augmentations here are chosen to encode invariances that should not change the under- lying sample idyntily. In previous contrastive learning for graphs, these augmentations were for example dropping random edges or masking node features (You et al. 2020 These augmentations are not appropriate for our spatially- embedded graphs that form a tree and whose only node features are their 3D location space. Thus we designed five novel augmentation techniques specifically for sp&- tially embedded graphs such neural morphologies or botantical trees: subsampling, rotation, node jittering, branch deletion and graph translation (sce Section 3).
Table Ablation results on the BBP datasct. Cell-type classification accu- racy [%] of our model and ablations averaged over three random seeds and given mean + standard deviation. a " indicates removal of an augmenta- tion or model component.
Model GraphDINO
Accuracy
65.8
3D rot.
55.4 i
node jitter
64.8
To Lest the importance of the individual graph augmentations we perform a set of ablation experiments using the BBP dataset. We remove one augmentation from our model at time and evaluate the leave-one-out 5-nearest neighbor accuracy when predicting the expert labels. For the subsampling augmentation vary the mum- ber of retained nodes. Our full model achieves average accuracy of 65.8% when classifying the excitatory cells into the 12 expert labels (Appendix C.5.1p. When removing individual data augmentations the accuracy decreases (Tab. 2
graph translation
55.6 =
drop branch
64.6
subsampling: 50 nodes
60.0
subsampling: 200 nodes
62.0
adjacency (7 = 0)
62.8 +
attention (A =
59.8 =
Especially 3D rotation and graph translation are important augmentation strategies whose removal lead to substantial performance deterioration.
5.3
Message-passing is not sufficient for long-range graphs
Next, we investigate whether classical message-passing is suf- ficient to process graphs with long-ranging branches such as neuronal morphologics. Therefore, we train GRAPHDINO once when only using message passing while removing the global attention (setting À O in Eq. D This decreases the performance to 59.8% (Tab. 2 Additionally, we train INFOGRAPH (Sun et al. 2020 as bascline for an unsu- pervised method that learns graph-level representat ions and INFOGRAPH achieves accuracy of 48.2% (Tab. 3). Thus, we conclude that using global at- tention is beneficial in situations where graphs contain long-range mangrange branches. Global attention enables infor- mation flow between distant (in terms of graph connectivity) nodes that might be close in space or function.
Table 3: Cell-type classification accuracy [%] on the BBP dataset. Performance of our model and INFOGRAPH (Sun et al. 2020 av- eraged over three random seeds and given mean + standard deviation.
Model
Accuracy
INFOGRAPH Sun et al. 2020 GraphDINO
48.2
uses GNN message-passingg
65.8
5.4 Morphological embeddings differentiate between spiny/aspiny cells and layers
To evaluate the capability of GRAPHDINO to capture essential features of 3D neuronal shapes purcly data- driven, we train GRAPHDINO on the BBP dataset and use t-distributed stochastic neighbor embedding (t- SNE) (van der Maaten & Hinton 2008 to map the learned embeddings of the BBP dataset into 2D (Fig. for
8

---

Published in Transactions on Machine Learning Research (06/2023)
A.
B.
Spiny
2
3
4
5
7
8
spiny
9
aspiny
10
11
C.
Cortical layer
D. Apical total length
14
15
Spiny
E.
Log likellhood
10
9.0
15
-10.0
L2/3
L
32
Number of clusters
F.
Spiny
2
3
5
Spiny
8
10
Spiny
14
Figure A. t-SNE embedding (perplexity=30) of latent representation of 3D neuronal morphologies of the BBP dataset showing a separation into spiny and aspiny neurons (n 616). B. t-SNE embedding (per of the latent representations of the morphologies of the spiny neurons colored by the cluster found by our model (n 286). C. Relative cortical layer distribution of neurons per cluster across L2/3 L6. Higher values are indicated by red. D. As B. but neurons colored by their total apical length revealing organization of the latent space in terms of morphological properties. E. Log- ligeliiel of Gaussian Mixture Model on held-out test set for spiny neurons used to select the optimal number of clusters. F. Example neurons for each cluster are shown with apical dendrites lighter color, while basal dendrites are colored darker. Soma is indicated by black circle.

---

Published in Transactions on Machine Learning Research (06/2023)
visualization. A clear separation between spiny and aspiny neurons can be observed (see Fig. 3 indicating that our learned representation captures meaningful biological differences of the neuronal morphologies.
Interestingly, some of the spiny neurons end up in the aspiny cluster (Fig. A bottom right). These are inverted L6 neurons (Fig. 3F Cluster 15), whose size and dendritic tree are morphologically similar to the surrounding aspiny neurons that also show a downwards bias the dendritic tree.
5.5 Morphological embeddings recover known excitatory cell types
To identify cell types, we fit a Gaussian mixture model (GMM) with a diagonal covariance matrix to our learned representation of the spiny neurons. To determine the number of clusters, we fit. 1,000 GMMs with different random seeds using five-fold vross-validation for 2 30 clusters. We average over the log-likelihood for each number of clusters over repetitions and folds. We find = 15 to be the optimal nunber of clusters (Fig. 3 E)
Having identified the optimal mumber of clusters, we ro-fit the GMM to the full dataset including all spiny neurons. To avoid picking a particularly good or bad random clustering, we 100 models and choose the one that has the highest average adjusted rand index (ARI) to all other clusterings.
The spiny neurons cluster nicely into different shapes and layers (Fig. 3F and Appendix Fig. D.6 retrieving known excitatory cell lypes. The first four spiny clusters contain mainly cells from layer 2/3
(L.2/3) (Fig 3 3C)
and group them by morphology: Cluster 1 contains wide and short neurons from layer 2/3, while L2/3 neurons in cluster are more elongated with a less pronounced apical tuft (Fig. 3F) Clusters 5-7 group cells from layer (L4) (Fig. 3C), differentiating between spiny stellate cells (cluster 5) and atufted L4 neurons (cluster 7) (Fig. F)) Within layer and 6, neurons are grouped by their size, amount of apical tuft and obliques, as well as the direction of the apical-like dendrites: For instance, cluster 10 groups thick-tufted pyramidal cells from layer 5 and cluster 15 contains inverted L6 neurons (Fig. 3F).
Most clusters show a strong preference for grouping cells whose soma position is in certain layer (Fig. even though the model in contrast to the experts who labeled the cells does not have access to anatomical knowledge such as cortical layer of origin. One exception are pyramidal L6 cells with upward- directed apicals that separate less well and get rather clustered with L4 and L5 neurons of the same size and similar morphological shape. This is to be expected, as the model only learns to differentiate between different morphologies but has no knowledge about anatomical features such as soma depth.
5.6 Data-driven clusters are consistent with expert labels
To compare our data-driven features to manually-designed features, we compute the adjusted rand index (ARI) between our clusters and the expert-identified cell types on the BBP dataset and compare the per- formance to the clusters based on morphometrics obtained by Gouwens et al. (2020). We achieve an ARI performance of 0.31 when clustering neurons across all cortical layers together while using significantly less prior information than Gouwens et al. 2019). In comparison, Gouwens et al. (2019 reached an ARI of 0.27 with a feature space specifically designed for spiny neurons and by splitting the neurons into their cortical layer of origin before performing the clustering. This approach reduces the complexity of the problem sig- nificantly, since misassignments across layers are excluded by construction. When performing the clustering like Gouwens et al. 2019) only within the layers, we achieve an ARI of 0.46 (Tab. 4
5.7 Morphological embeddings encode distinct morphological features
Laturnus & Berens 2021 classified the MI EXC dataset (Scala et al.) 2021 into three classes based presence of an apical tuft (tufted, untufted and others). Following their work, we train a 5-nearest-noighhor classifier on our learned embeddings and show that GRAPHDINO learns meaningful features to differentiate between the three classes (Tab. 5). Our method outperforms their MORPHVAE method as well as a baseline using density maps of the neurons (Laturnus Berens 2021). This datasct is rather small and Laturnus & Berens (2021) used only a single train/test split. To estimate how reliable the reported accuracy metrics are, we compute the cross- validated accuracy across multiple different. train /test splits, which show vari-
10

---

Published in Transactions on Machine Learning Research (06/2023)
Table Adjusted rand index (ARI) between identified clusters and expert labels for the learned embeddings from GRAPHDINO and manually-dcfined features by Gouwens et al. (2019) and expert labels, when performing the clustering within cortical layers and across the whole cortex.
Table 5: Balanced accuracy [%] on M1 EXC test set us- ing the learned embeddings from either GraphDINO or MorphVAR [Laturnus & Berens 2021 (mean + SEM across three runs and across three data splits, respec- tively). Percentages in brackets indicate the amount of labels used during training for MorphVAE. GraphDINO is trained without labels.
Accuracy over runs (mean SRM)
Accuracy over splits (mean sn)
Clustering
Features
ARI
Model
across layers
GRAPHDINO
0.31
MorphVAE (100 %)
70
within layers
Gouwens et (0019)
0.27
MorphVAE (0 %)
58+
GraphDINO
0.46
Density Map (0 %)
60
GraphDINO (0 %)
68 =
71 19
ability across splits of 1.9% (standard deviation; Tab. 5 We conclude that GraphDINO likely outperforms MorphVAE trained withoutsupervision and performs approximately on par with MorphVAE trained fully supervisod.
5.8 Morphological embeddings encode cortical regions
TREEMoCo (Chen et al. 2022b) is an LSTM-based model that concurrently proposed to perform un- supervised representation learning neuronal graphs. The model uses as input the simplified skeletons of neurons that only contain the branching points as nodes. They compute 26 manually-selected features in addition to the xyz-coordinates as node features to describe the morphology of the skeletons between branching points. TREEMoCo is trained on a combination the datasets BIL, JML and ACT and quan- titatively evaluated on the task of predicting the brain anatomical region or cortical layer of origin of the neurons subset of the neuronal classes. Chen et al. (202bb) remove 955 neurons from the dataset due to "reconstruction errors" and evaluate exact neurons used for training and evaluation both in terms of split and which neurons were removed, we trained unsupervised on the joint dataset and evaluated using 5-fold cross-validation, i.e. splitting the data into five folds and evaluating each fold, given the other four folds as training data and reporting the average performance across folds. For further details regarding the evaluation, see Appendix C.5.4
80-20% training-tes. split. Since we did not have access to the
GRAPHDINO performs on par or better than TREEMOCO and GRAPHCL when predict- ing the origin of neurons (Tab. 6). Note that GRAPHDINO is fully data-driven while TREEMoCo and GRAPUCL additionally em- ploy manually extracted node features.
Table 6: Cell-type classification on the TreeMoCo dataset. Performance of our model (GrApHDinO) averaged over three random seeds and given as mean = standard devi- ation. TREEMoCo and GRAPHCL performance given as the average accuracy over the last five epochs per dataset. *Results taken Fig. C1 of Chen et al. (2022b).
Note that the evaluation reported by Chen et al. (2022b) uses excitatory and inhibitory neurons at the same time. With this approach, morphologies of neurons of the "same" class can look very different (Fig. C.5). A better proxy lask to evaluate the encoding capabilities of the models would be to restrict the evaluation to only excitatory cells. For the ACT dataset this information available. We therefore repeated the evaluation only on this subset (Tab. 6 which should provide meaningful baseline for future studies.
Model
BIL-6
JML-4
ACT
ACT spiny
TrecMoCo" GraphCL
76.9
59.7
53.9
66.3
50.6
55.6
GRAPHDINO
79
63
544
73
11

---

Published in Transactions on Machine Learning Research (06/2023)
6 Limitations
GRAPHDINO is designed to learn graph-level representations of spatially-embeddeed tree-structured graphs using self-supervised learning. As we focus on graphs where each node has a location in 3D space and design the data augmentations accordingly, the approach is not expected to work out-of-the-bot box on graphs that have different node features. AC-ATTENTION is likely to be beneficial in many other sccnarios as well, since it can smoothly interpolate between message passing and global altention based on node similarity, but this hypothesis remains to be tested empirically. Data augmentations would need be adapted the respective data domain and the respective invariances that should be encoded or supervised learning to be used. The attention mechanism is not tied in any way to the self-supervised learning objective we use.
We encode the desired invariance for neuronal morphologies in GRAPHDINO via tailored data augmenta- tions. Rotation and translation equivariance could alternatively be built into the architecture of the encoders explicitly. Recent works have proposed such architectures for GNNs (Satorras et al. 2021), as well as for transformers (Fuchs et al. 2020 Adapting these for AC-ATTENTION would be an interesting future research direction.
Computing the full transformer attention matrix has quadratic complexity and might therefore be com- putationally infeasible for graphs with a large number of nodes. We solve this problem here by subsampling the neuronal skeletons to a smaller number of nodes, which has the added benefit of being strong data augmentation that keeps the global morphology of the neuron intact while altering the local structure be- tween the two views. However, this approach might not be suitable for all graph datasets. There has been some work in building attention mechanism that scale lincarly with the number of input tokens (Wang al., 2020 Kitaev et al. 2020; Choromanski al., 2021 but integrating them with the message passing might not be straightforward.
Self-supervieed learning has been shown to be most successful when training large datasets (Bao et al, 2022 Oquab et al., 2023). We equipped GRAPHDINO with appropriate inductive biases to make it possible to learn on the smaller publicly available neuronal datasets that have been used in previous studies. Nev- ertheless, applying GRAPHDINO to neuronal datasets with more samples will likely improve its learning capabilitics. With the continual development of better imaging techniques and initiatives like MICRONS (MICTONS Consortium et al. 2023) more large-scale datasets of neuronal morphologies will be available tost this hypothesis.
In terms of neuronal cell lype classification, we did not take some features into account that have been previously used to differentiate cell types, such as the shape of the soma (as formerly used for GABAergic interneurons) or spine donsitics (Ascoli et al. 2008). Future work could focus on incorporating them into our framework. Depending on the type of feature, they could be easily integrated by adding them as features of the graph or as additional node features.
7 Conclusion
Increasingly large and complex datasets of neurons have given rise to the need for unbiased and quantitative approaches to cell type classification. We have demonstrated one such approach that is purely data-driven and self-supervised, and that learns a oow-dimensional representation of the 3D shape of a neuron. By using self-supervised learning, we do not pro-specify which cell types to learn and which fcatures to use, thereby reducing bias in the classification process and opening up the possibility to discover new cell types. A similar approach can also be useful in other domains beyond neuroscience, where samples of the dataset are spatial graphs and graph-level embeddings are desired, such as tree classification in forestry.
Acknowledgments
We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS), Tübingen, for supporting Marissa. A. Weis. This project has received funding from the European Rescarch Council (ERC) under the Europcan Union's Horizon Europe research and innovation program (Grant agreement No. 101041669).
12

---

Published in Transactions on Machine Learning Research (06/2023)
References
Allen Institute. Allen cell types database technical white paper: Cell morphology and histology. 2016. URL Morph__verviee
http://help.brant ooa
Rubén Armañanzas and Giorgio Ascoli. Towards the automatic classification of neurons. Trends in Newrosciences, 38(5):307-315, 2015.
Giorgio Ascoli, Lidia Alonso-Nanclares, Stewart Anderson, Germán Barrionuevo, Ruth Benavides- Andreas Burkhalter, Gyorgy Buzsáki, Bruno Cauli, Javier Defelipe, and Alfonso Fairen. Petilla termi- nology: nomenclature of features of GABAergic interncurons of the cerebral cortex. Nature reviews. Neuroscience, 9:557 568, 2008.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. BEit: BERT pro-training of image transformers. In Proc. of the International Conf. on Learning Representations (ICLR), 2022.
Cathryn Cadwell, Athanasia Palasaniza, Xiaolong Jiang, Philipp Berens, Qiaolin Deng, Marlene Yilmaz, Jacob Reimer, Shan Shen, Matthias Bethge, Kimberley Tolias, Rickard Sandberg, and Andreas Tolias. Electrophyssilogical, transcriptomic and morphologie profiling of single neurons using patch-seq. Nature Biotechnology, 34, 2015.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in superviee vision transformers. In Proceedings of the IEEE/CVP International Conference on Computer Vision (ICCV) pp. 9650-9660, 2021.
Dexiong Chen, Leslic O' O'Bray, and Karsten Borgwardt. Structuro-aware transformer for graph representation learning, 2022a.
Hanbo Chen, Jiawei Yang, Danicl Maxim lascone, Lijuan Liu, Lei Hc, Hanchuan Peng, and Jianhua Yao. Treemoco: Contrastive neuron morphology representation learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022b.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proc. of the International Conf. on Machine learning (ICML), 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese reprosentation learning. In Proc. IEEE Conf. Computer Vision and Pattern Recognilion (CVPR), 2021.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In Proc. of the International Conf. Learning Representations (ICLR), 2021.
Javier Defelipe, Pedro López-Cruz, Ruth Benavides-Piccione Concha Bielza, Pedro Larranaga, Stewart An- derson, Andreas Burkhalter, Bruno Cauli, Alfonso Fairen, Dirk Feldmeyer, Gord Fishell, David Fitzpatrick, Tamás Freund, Guillermo Gonzalez Burgos, Shaul Hestrin, Sean Hill, Patrick Hof, Josh Huang, Edward Jones, and Giorgio Ascoli. New insights into the classification and nomenclature of cortical gabaergic interneurons. Nature reviews. Neuroscience, 14, 2013.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. CVPR, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa. Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv.org, 2020.
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafacl Bombarell, Timothy Hirzel, Alan Aspuru- Guzik, and Ryan P Adams. Convolutional networks on graphs for learning molecular fingerprints. In Advances in Neural Information Processing Systems (NeurIPS), volume 28, 2015.
13

---

Published in Transactions on Machine Learning Research (06/2023)
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. arxivoorg. 2012.09699, 2021.
Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Bench- marking graph neural networks. arXic.org, 2003.00982, 2020.
Leila Elabbady, Sharmishtaa Seshamani, Shang Mu, Gayathri Mahalingam, Casey M Schneider- Agnes Bodor, J Alexander Bae, Derrick Brittain, JoAnn Buchanan, Daniel J Bumbarger, et al. Quantitative census of local somatic features in mouse visual cortex. bioRovn 2022.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se(3)- ((l3)-transformers: 3d roto-translation equivariant attention networks. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Newral Information Processing Systems (NeurIPS), volume 33, pp. 1970 1981, 2020.
Rohan Gala, Agata Budzillo, Fahimeh Baftizadeh, Jeremy Miller, Nathan Gouwens, Anton Arkhipov, Gabe Murphy, Bosiljka Tasic, Hongkui Zeng, Michael Hawrylyez, et al. Consistent cross-modal identification of cortical neurons with coupled autoencoders. Nature Computational Science, 1(2):120-13 127, 2021.
Justin Gilmer, Samuel S. Schocnholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In Proc. of the International Conf. on Machine learning (ICML), pp. 1263-1273 2017.
Nathan Gouwens, Staci Sorensen, Jim Berg, Changkyu Lee, Tim Jarsky, Jonathan Ting, Susan Sunkin, David Feng, Costas Anastassiou, Eliza Barkan, Kris Bickley, Nicole Blesie, Thomas Braun, Krissy Brouner, Agata Budzillo, Shiella Caldejon, Tamara Casper, Dan Castelli, Peter Chong, and Christof Koch. Classification of clectrophysioloicaal and morphological neuron types in the mouse visual cortex. Nature Newroscience, 22, 2019.
Nathan W. Gouwens, Staci A. Sorensen, Fahimeh Baftizadeh, Agata Budzillo, Brian R. Lee, Tim Jarsky, Lauren Alfiler, Katherine Baker, Eliza Barkan, Kyla Berry, Darren Bertagnolli, Kris Bickley, Jasmine Bomben, Thomas Braun, Krissy Brouner, Tamara Casper, Kirsten Crichton, Tanya L. Daigle, Rachel Dalley, Rebecca A. de Frates, Nick Dec, Tsoga Desta, Samuel Dingman Lee, Nadezhda Dotson, Tom Egdorf, Lauren Ellingwood, Rachel Enstrom, Luke Esposito, Colin Farrell, David Feng, Olivia Fong, Rohan Gala, Clare Gamlin, Amanda Gary, Alexandra Glandon, Jeff Goldy, Melissa. Gorham, Lucas Graybuck, Hong Gu, Kristen Hadley, Michael J. Hawrylycz, Alex M. Henry, DiJon Hill, Madie Hupp, Sara Kebede, Tae Kyung Kim, Lisa Kim, Matthew Kroll, Changkyu Lee, Katherine E. Link, Matthew Mallory, Rusty Mann, Michelle Maxwell, Medea McGraw, Delissa McMillen, Alice Mukora, Lindsay Ng, Lydia Ng, Kiet Ngo, Philip R. Nicovich, Aaron Oldre, Daniel Park, Hanchuan Peng, Osnat Penn, Thanh Pham, Alice Pom, Zoran Popovié, Lydia Potekhina, Ramkumar Rajanbabu, Shea Ransford, David Reid, Christine Rimorin, Miranda Robertson, Kara Ronellenfitch, Augustin Ruiz, David Sandman, Kimberly Smith, Josef Sule, Susan M. Sunkin, Aaron Szafer, Michael Tieu, Amy Torkelson, Jessica Trinh, Herman Tung, Wayne Wakeraan, Katelyn Ward, Grace Williams, Zhi Zhou, Jonathan T. Ting, Anton Arkhipov, Uygar Sümbül, Ed S. Lein, Christof Koch, Zizhen Yao, Bosiljka Tasic, Jim Berg, Gabe J. Murphy, and Hongkui Zeng, Integrated morphoelectric and transcriptomic classification of cortical gabaergic cells. Cell, 183(4):935 953.c19, 2020.
William I. Hamilton, Rex Ying, and Jure Leskovee. Inductive representation learning on large graphs. Tn Advances in Newral Information Processing Systems (NewrIPS), pp. 1025-1035, 2017.
Kaveh Hassani and Amir Hosein Khasahmadi. Contrastive multi-view representation learning on graphs. Tn Proc, of the International Conf. on Machine learning (ICML), volume 119, pp. 4116 4126, 2020.
Lida Kanari, Pawel Dlotko, Martina Scolamiero, Ran Levi, Julian C. Shillcock, Kathryn Hess, and Henry Markram. topological representation of branching neuronal morphologies. Neuroinformatics, 16:3 13, 2017.
Lida Kanari, Srikanth Ramaswamy, Ying Shi, Sebastien Morand, Julie Mcystre, Rodrigo Perin, Marwan Abdellah, Yun Wang, Kathryn Hess, and Henry Markram. Objective morphological classification of neocortical pyramidal cells. Cerebral Corter, 29(4):1719-1773) 2019.
14

---

Published in Transactions on Machine Learning Research (06/2023)
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Proc. of the Interna- tional Conf- on Learning Represenlations (ICL.), 2015.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In Proc. of the International Conf. on Learning Representations (ICLR), 2017.
Nikita Kilaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In Proc. of the International Conf. Learning Representations (ICLR), 2020.
Johannes Klicpera, Stefan Weiß enberger, and Stephan Günneeaann. Diffusion improves graph learning. In Advances in Newral Information Processing Systems (NeurIPS), volume 32, 2019.
Sophie C. Laturnus and Philipp Berens. Morphvae: Generating neural morphologies from 3d-walks using a variational autoencoder with spherical latent space. In Proc. of the International Conf. on Machine learning (ICML), volume 139, pp. 6021-6031, 2021.
Yujia Li, Daniel Tarlow, Marc Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. Proc. of the International Conf. on Learning Representations (ICLR), 2016.
lyy Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with restarts. arXiv.org, 1608.03983, 2016.
Yanbin Lu, Lawrence Carin, Ronald Coifman, William Shain, and Badrinath Roysam. Quantitative arbor analytics: unsupervised harmonic co-clustering of populations of brain cell arbors based on I-measure. Newroinformatics, 13(1):47-63 63, 2015.
Henry Markram, Eilif Muller, Srikanth Ramaswamy, Michael Reimann, Marwan Abdellah, Carlos Aguado, Anastasia Ailamaki, Lidia Nlncsa-Nanclares, Nicolas Antille, Selim Arsever, Atenekeng Kahou Guy An- toine, Thomas K Berger, Ahmet Bilgili, Nenad Buncic, Athanassia Chalimourda, Giuseppe Chindemi, Jean-Denis Courcol, Fabien Delalondre, Vincent Delattre, and Felix Schürmann. Reconstruction and simulation of neocortical microcircuitry. Cell, 163:456-490 2015.
Grégoire Mialon, Dexiong Chen, Margot Selosse, and Julien Mairal. Graphit: Encoding graph structure in transformers, 2021.
The MICrONS Consortium, J. Alexander Bae, Mahaly Baptiste, Caitlyn A. Bishop, Agnes L. Bodor, Der- rick Brittain, JoAnn Buchanan, Daniel J. Bumbarger, Manuel A. Castro, Brendan Celii, Erick Cobos, Forrest Collman, Nuno Maçarico da Costa, Sven Dorkenwald, Leila Elabbady, Paul G. Fahey, Tim Fliss, Emmanouil Froudarakis, Jay Gager, Clare Gamlin, William CrnyaR,ncal, Akhilesh Halageri, James Heb- ditch, Zhen Jia, Emily Joyce, Justin Joyce, Chris Jordan, Daniel Kapner, Nico Kemnitz, Sam Kinn, Lindsey M. Kitchell, Selden Koolman, Kai Kuchner, Kisuk Lee, Kai Li, Ran Lu, Thomas Macrina, Gay- athri Mahalingam, Jordan Matelsky, Sarah McReynolds, Elanine Miranda, Eric Mitchell, Shanka Subhra Mondal, Merlin Moore, Shang Mu, Taliah Muhammad, Barak Nehoran, Oluwaseun Ogedengbe, Chris- tos Papadopoulos, Stelios Papadopoulos, Saumil Patel, Xaq Pitkow, Sergiy Popovych, Anthony Ramos, R. Clay Reid, Jacob Reimer, Patricia K. Rivlin, Victoria Rose, Casey M. Schneider-Mizell, H. Scbastian Seung, Ben Silverman, William Silversmith, Amy Sterling, Fabian H. Sinz, Cameron L. Smith, Shelby Suckow, Mare Takeno, Zheng H. Tan, Andreas S. Tolias, Russel Torres, Nicholas L. Turner, Edgar Y. Walker, Tianyu Wang, Adrian Wanner, Brock A. Wester, Grace Williams, Sarah Williams, Kyle Willic, Ryan Willie, William Wong, Jingpeng Wu, Chris Xu, Runzhe Yang, Dimitri Yatsenko, Fei Ye, Wenjing Yin, Rob Young, Szi chieh Yu, Daniel Xenes, and Chi Zhang. Functional connectomics spanning multiple areas of mouse visual cortex. bioRaiv, 2023. doi:
Annamalai Narayanan, Mahinthan Chandramohan, Rajasekar Venkatesan, Lihui Chen, Yang Liu, and Shan- tamu Jaiswal. graph2vec: Learning distributed reprosentations of graphs. arXiv.org, 1707.05005, 2017.
Marcel Oberlaender, Christiaan P. J. de Kock, Randy M. Bruno, Alejandro Ramirez, Hanno S. Meyer, Vincent Dercksen, Moritz Helmstacdter, and Bert Sakmann. Cell Type-Specific Three-Dimensions Structure of Thalamocortical Circuits in a Column of Rat Vibrissal Cortex. Cerebral Cortez, 22(10): 2375-2391, 2391, 2012.
15

---

Published in Transactions on Machine Learning Research (06/2023)
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv.org, 1807.031 0807.03748, 2018.
Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Danicl Haziza, Francisco Massa, Alaacldin El-Nouby, Mahmoud Assran, Nicolas Ballas, Woj- ciech Galuba, Russell Howes, Po-Yao Huang, Shang- Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jegou, Julien Mairal, Patrick Tabatut, Armand Joulin, and Piotr Bo- janowski. Dinov2: Learning robust visual features without supervision, 2023.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, peghhperformance deep learning library. In Advonces in Newral Information Processing Systems (NeurIPS), 2019.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay, Scikit-learn: Machine learning in Python. Journal of Machine Learning Research (JMLR), 12:2825-2830, 2011.
Hanchuan Peng, Peng Xie, Lijuan Liu, Xiuli Kuang, Yimin Wang, Lei Qu, Hui Gong, Shengdian Jiang, Anan Li, Zongcai Ruan, Liya Ding, Zizhen Yao, Chao Chen, Mengya Chen, Tanya Daigle, Rachel Dalley, Zhangean Ding, Yanjun Duan, Aaron Feiner, and Hongkui Zeng. Morphological diversity of single ncurons in molecularly defined cell types. Nature, 598:174 181, 10 2021. doi: 00.1035/441556-021-02
Sridevi Polavaram, Todd A Gillette, Ruchi Parekh, and Giorgio Ascoli. Statistical analysis and data mining of digital reconstructions of dendritic morphologies. Frontiers in neuronnatony, 8:138, 2014.
Jiczhong Qiu, Qibin Chen, Yuxiao Dong, Jing Zhang, Hongxia Yang, Ming Ding, Kuansan Wang, and Jie Tang. Goc: Graph contrastive coding for graph neural network pre-training. In Proc. of Conf. Knowledge Discovery and Data Mining (KDD), pp. 1150-116,, 2020.
Srikanth Ramaswamy, Jean-Denis Courcol, Marwan Abdellah, Stanislaw R. Adaszewski, Nicolas Antille, Se- lim Arsever, Guy Atenekeng, Ahmet Bilgili, Yury Brukau, Athanassia Chalimourda, Giuseppe Chindemi, Fabien Delalondre, Raphael Dumuse, Stefan Eilemann, Michael Emiil Gevaert, Padraig Gleeson, Joe W. Graham, Juan B. Hernando, Lida Kanari, Yury Katkov, Daniel Keller, James G. King, Rajnish Ranjan, Michael W. Reimann, Christian Rössert, Ying Shi, Julian C. Shillcock, Martin Telefont, Werner Van Geit, Jafet Villafranca Diaz, Richard Walker, Yun Wang, Stefano M. Zaninetta, Javier DeFelipe, Sean L. Hill, Jeffrey Muller, Idan Segev, Felix Schürmann, Filif B. Muller, and Henry Markram. The neocortical mi- crocircuit collaboration portal: a resource for rat somatosensory cortex. Frontiers in Neural Circuits, 9: 44, 2015.
Ladislav Rampášek, Mikhail Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for general, powerful, scalable graph transformer, 2022.
Santiago Ramón y Cajal. Histologie du système nerveur de l'homme et des vertébrés. 1911.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. arXiv.org, 2102.09844, 2021.
Federico Scala, Dmitry Kobak, Mattco Bernabucci, Yves Bernaerts, Cathryn Cadwell, Jesus Castro, Leonard Hartmanis, Xiaolong Jiang, Sophie Laturnus, Elanine Miranda, Shalaka Mulherkar, Zheng Tan, Zizhen Yao, Hongkui Zeng, Rickard Sandberg, Philipp Berens, and Andreas Tolias. Phenotypic variation of transcriptomic cell types in mouse motor cortex. Nature, 598:1 7, 2021.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEFF Transactions on Neural Networks, 20(1):61-88 180, 2009.
16

---

Published in Transactions on Machine Learning Research (06/2023)
Philiip Schubert, Sven Dorkenwald, Michal Januszewski, Viren Jain, and Joergen Kornfeld. Learning cellular morphology with neural networks. Nature Communications, 10:2736, 2019.
Ruggero Scorcioni, Sridevi Polavaram, and Giorgio A Ascoli. L-measure: a web-accessible tool for the analysis, comparison and search of digital reconstructions neuronal morphologies. Nature protocols, 3 (5):866-876, 2008.
Dominik Seidel, Yonten Dorji, Bernhard Schuldt, Emilic Isasa, and Klans Körber. Dataset: New insights into tree architecture from mobile laser scanning and geometry analysis. Dryad, 2021. doi: https:/ /doi. /1rp:0/1//dydd..///////.
Sharmishtaa Seshamani, Leila Elabbady, Casey Schneider-Mizell, Gayathri Mahalingam, Sven Dorkenwald, Agnes Bodor, Thomas Macrina, Daniel Bumbarger, JoAnn Buchanan, Mare Takeno, Wenjing Yin, Derrick Brittain, Russel Torres, Daniel Kapner, Kisuk Lee, Ran Lu, Jingpeng Wu, Nuno daCosta, R. Clay Reid, and Forrest Collman. Automated neuron shape analysis from electron mieroscopy. arXiv.org, 2006.00100, 2020.
D. A. Sholl. Dendritic organization in the neurons of the visual and motor cortices of the cat. Journal of Anatomy, 87:387-406, 1953.
Pan-Yun Sun, Jordan Hoffmann, and Jian Tang. Infograph: Unsupervised and semi-supervised graph-level representation learning via mutual information maximization. Tn Proc. of the International Conf. Learning Representations (ICLR), 2020.
Shantanu Thakoor, Corentin Tallec, Mohammad Gheshlaghi Azar, Mehdi Azabou, Eva L Dyer, Remi Munos, Petar Veličković, and Michal Valko. Large-scale representation learning on graphs via bootstrapping. Proc. of the International Conf. on Learning Representations (ICLR), 2022.
Harry BM Uylings and Jaap Van Pelt. Measures for quantifying dendritic arborizations. Network: compu- tation in neural systems, 13(3):397, 2002.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Machine Learning Research (JMLR), 9 (86):2557-22)) 2605, 2008.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Mllia Polosukhin. Attention is All you Need. Advances in Newral Information Processing Systems (NeurIPS), 2017.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph Atlention Networks. In Proc. of the International Conf. on Learning Representalions (ICLR), 2018.
Petar Veličkovié, William Fedus, William L. Hamilton, Pictro Liò, Yoshua Bengio, and R. Devon Hjelm. Deep graph infomax. In Proc. of the International Conf. on Learning Representations (ICLR), 2019.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv.org, 2020.
Yun Wang. simplified morphological classification scheme for pyramidal cells in six layers of primary somalosensory cortex of juvenile rats. IBRO Reports, 5, 2018.
Johan Winmubst, Erhan Bas, Tiago A. Ferreira, Zhuhao WI, Michael N. Economo, Patrick Edson, Ben J. Arthur, Christopher Bruns, Konrad Rokicki, David Schauder, Donald J. Olbris, Sean D. Murphy, David G. Ackerman, Cameron Arshadi, Perry Baldwin, Regina Blake, Ahmad Elsayed, Mashtura Hasan, Danicl Ramirez, Bruno Dos Santos, Monet Weldon, Amina Zafar, Joshua T. Dudman, Charles R. Gerfen, Adam W. Hantman, Wyatt Korff, Scott M. Sternson, Nelson Spruston, Karel Svoboda, and Jayaram Chandrashckar. Reconstruction of 1,000 projection neurons reveals new cell types and organization of long-range connectivity in the mouse brain. Cell, 179(1):268-28) 2012)..111.....13, 2019. ISSN 0092-8674. doi: 20..090000....222000......22
17

---

Published in Transactions on Machine Learning Research (06/2023)
Minghao Xu, Hang Wang, Bingbing Ni, Hongyu Guo, and Jian Tang. Self-supervised graph-level represen- tation learning with local and global structure. In Proc. of the International Conf. on Machine learning (ICML), volume 139, pp. 11548 11558, 2021.
Chengxuan Ying, Tianle Cai, Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? In Advonces in Neural Information Processing Systems (NewrIPS), 2021.
Yuning You, Tianlong Chen, Yongduo Sui, Ting Chen, Zhangyang Wang, and Yang Shen. Graph contrastive learning with augmentations. In Advances in Newral Information Processing Systems (NeurIPS), 2020.
Jiawei Zhang, Haopeng Zhang, Congring Xia, and Li Sun. Graph-bert: Only attention needed for learning graph representations. arXiv.org, 2001.05140, 2020.
Jie Zhao, Xuejin Chen, Zhiwei Xiong, Zheng-Jun Zha, and Feng Wu. Graph representation learning for large- scale neuronal morphological analysis. IEEE Transactions on Newral Networks and Learning Systems, pp. 1 12, 2022. doi: 10.1109/TTNNE 202.0TTT
Yanqiao Zhu, Yichen Xu, Fong Yu, Qiang Liu, Shu Wu, and Liang Wang. Graph contrastive learning with adaplive augmenlation. In Proceedings of the Web Conference 2021, pp. 2069 2080, 2021.
18

---

Published in Transactions on Machine Learning Research (06/2023)
Appendices
A
Synthetic graph dataset
To test whether model is able to use information encoded in the connectivity of the graphs, we gener- ale a synthetic graph datasel. with five classes that differ in connectivity while having similar node loca- tions. We create this synthetic graph datasct by uni- formly sampling 20 mean node positions in 3D space shared between the five classes to ensure that the presence of specific node does not encode class member- ship. For each class, we construct a distinct graph connectivity as follows: We first randomly sample a root node and two children, then we recursively sam- ple one or two children per child (with a branching probability of 50%) until all 20 nodes are connected. Using this method, we generate 100,000 graphs for the training set and 10,000 graphs for validation and test set each (with & class probability) by sampling node positions from ((,, with μ equal to the above drawn means and o = 10.
A
B
60
20
20
-0
-29
-10
-00
in [-50, 50] The mean node locations
60
40
440 20 6
-42 .0
20
20
20
Figure A.1: Example classes (A) and (B) of syn- thetic graph dataset. Samples within one class share graph connectivity. Samples between classes share mean node locations. Node locations are drawn from o)
Tab.
Tab. C.2 and Tab.
C.3 list. the hyperparamers used for experiments on the synthetic graphs.
t-SNE of the learned la- tent spaces: To visualize the learned latent space we per- form t-SNE with a perplexity of 30 to reduce the embedding to two dimensions (Fig. A.2
Linear classifier: We train a supervised linear classifier on the extracted embeddings of GRAPHDINO for 100 epochs and a learning rate of 0.01. To train the classifier, we use the test set that has not boen used in training GRAPHDINO, and split it in 8,000 samples for training the classifier and 2,000 samples for evaluating held-out test set accuracy.
Figure A.2: t-SNE embedding of syuthetic graph dataset colored by class membership for A five runs of GRAPHDINO with GRAPHArTENTION, five runs of GRAPIDINO with regular transformer attention, and C five runs of GRAPIDINO with transformer attention and without positional encoding.
19

---

Published in Transactions on Machine Learning Research (06/2023)
A.
Cluster
10
Figure B.3: A. t-SNE embedding of Trees dataset colored by cluster membership based on GMM clustering with 15 clusters. B. Three example tree morphologies are shown for different. clusters.
B Application to different domain: Tree Morphologies
We developed model that able to learn graph-level embeddings of spatially-ombeddeed graphs. So far, have shown that it yields meaningful cell types clusterings of neuronal morphologies. To show that GRAPIDINO is applicable to data domains beyond neuronal morphologies, we train our model 3D skeletons of individual trees (from forest).
The Trees dataset (Seidel et al., 2021 a highly divers dataset comprised of 391 skeletons of trees stemming from 39 different genuses and 152 species or breedings. The skeletons were extracted from LIDAR scans of the trees. Nodes of the skeletons have a 3D coordinate. We normalize the data such that the lowest point (start of the tree trunk) normalized (0, 0, 0).
GRAPHDINO learns latent space that orders tree morphologies with respect to their size, crown size and crown shape (Fig. B.3
Fig. D.7
C Extended Methods
C.1 Background: DINO
DINO (Fig. C.4 (Caron et al. 2021 is method for self-supervised image representation learning, Similar to previous approaches, it consists of two image encoders which process different views of an image. These views are obtained by image augmentation. The training objective is enforce both encoders to generate the same output distribution when the same input image shown. This can be implemented by the cross entropy loss function: Σ; logp. Both encoders transformers that share the architecture but differ in their weights: One of the encoders is the student encoder which receives weight updates through gradients of the training objective while the other encoder's (teacher) weights are an exponential moving average of the student's 's weights. In contrast to some other self-supervised methods, DINO does not require contrastive (negative) samples. To prevent collapse, i.e. predicting the same distribution independent, of the input image, two additional operations on the teacher's predictions are crucial: sharpening by adjusting the softmax temperature, and centering using batch statistics. Besides competitive performance on downstream image classification lasks, another key finding of the paper is that object segmentations emerge in the self-attention when applying DINO training visual transformer image encoders.
C.2 Data preprocessing.
To speed up data loading during training, we reduce the number nodes the graph of each neuron to 1000 nodes in the same way as when subsampling and ensure that contains only one connected compo- nent. there are unconnected components, we connect them by adding an edge between nodes of two unconnected components that have the least distance between their spatial coordinates.
20

---

Published in Transactions on Machine Learning Research (06/2023)
augmentation
student encoder
view X
softmax
p
imput image
exp. moving avg.
crnss entropy loss
View
teacher encoder
center
softmax
q
Figure C.4: The DINO method for sclf-supervised image representation learning (figure adapted from Caron
et al. 2021
C.3
Training details and hyperparameters
To select hyperparameters we run three grid searches and pick the best hyperparameters according to the lowest average loss over the BBP and M1 PatchSeq dataset.
For the optimization, we run a hyperparameter search over batch size E (32, 64, 128), learning rate e 100-,110- 10-50000101111.11111111100101110 10-5), and number of training iterations E {20, 000,50, 000, 100, 000).
For the augmentation strength, we run a hyperparameter search over jitter variance 01 E (1.000.1...001). number of deleted branchesn (1,5, 10), and graph position variance d2 (0.1, 10.0).
For the architecture, run hyperparameter search over latent dimension E (1,,32, (2, 64), number of GRAPHATTENTION blocks (depth) {5, 7, 10), and number of attention heads per block (2,4.8)
C.3.1 Architecture Hyperparameters
Table C.1: Hyperparameters used for the architecture for the different datascts. D1: Dimensionality of latent embedding z. D2: Dimensionality of probability distribution PE: Positional encoding. Softmax temperature of teacher network.
T temp:
Dataset
D1
D2
# layers
# heads
MLP dims
PE dims
T temp
Synthetic Graphs
16
300
4
4
16
16
0.04
BBP
1000
10
8
64
0.06
M1 PatchSeq
64
1000
7
8
64
32
0.06
Joint dataset (BIL, JML, ACT)
32
1000
64
0.06
Trees
32
1000
7
8
64
32
0.06
Tab. C.1
lists the hyperparameters used for the architecture for the different datasets. For the synthetic graph dataset, we clwnscale the network as it is a simpler dataset. DINO (Caron et al. 2021 uses an output
dimensionality of 65,536 for p when training on ImageNet (Deng et al. 2009) (1,000 classes). The number of classes in the neuronal datasets is unknown, but previous literature described 14 19 cell types (Gouwens) et al., 2019: Markram al.) 22115). Hence, we decrease the number of dimensions D2 of p proportionally 1,000, approximately retaining the ratio between classes and number of dimensions.
C.3.2 Optimization Hyperparameters
The learning rate is linearly increased to the value given in Tab. C.2 during the first 2,000 iterations and then decayed using exponential decay with rate 0.5 (Loshchilov & Hutter 2016)
21

---

Published in Transactions on Machine Learning Research (06/2023)
Table C.2: Hyperparameters used for optimization for the different datasets.
Dataset
Iterations
Batch size
Learning rate
Synthetic Graphs
100,000
512
10
BBP
100,000
64
10
M1 PatchSeq
50,000
128
10
Joint dataset (BIL, JML, ACT)
100,000
128
10
Trees
100,000
64
10-1
C.3.3 Augmentation Hyperparameters
Table C.3: Augmentation hyperparameters for the different. datasets. # nodes: Number of nodes to
subsample to. o1: Variance node jittering. # DB: Number of deleted branches. 02: Variance of graph translation.
Dataset
# nodes
01
# DB
0
Synthetic Graphs
15
0.1
0
0
BBP
100
0.001
10
10.0
M1 PatchSeq
100
0.1
10
10.0
Joint dataset (BIL, JML, ACT)
200
1.0
10.0
Trees
200
0.1
10.0
C.3.4 Computation
All trainings were performed on a NVIDIA Quadro RTX 5000 single GPU. Training on the neuronal BBP dataset ran for approximately 10 hours for 100,000 training iterations.
C.4 Inference
To extract the latent representation per sample, we encode the unaugmented graphs subsampled to 200 nodes using the student encoder and extract the latent representation a using the weights of the last iteration of training (no carly-stopping is used).
C.5
Evaluation
C.5.1
Evaluation on BBP
For visualization of the latent space, t-distributed stochastic neighbor embodding (t-SNE) (van der Maaten & Hinton 2008 with PCA-initialization, Euclidean distance and perplexity of 30.
For quantitative evaluation we use the subset of labeled excitatory neurons (n = 286) with the following 14 expert labels: L23-PC, L4-PC, L4-SP, L4-SS, L5-STPC, L5-TTPC1, L5-TTPC2, L5-UTPC, L6-BPC, L6-IPC, 16-TPC-L1, L6 L6--PC, L6-UTPC, L6-HPC (Markram et al. 2015).
For the ablation experiments and the comparison to INFOGRAPH Sun al. (2020) we perform k-nearest neighbor (kNN) classification with k 5 in leave-one-out setting predicting the above listed expert labels with two exceptions: We remoev the L6-HPC cells, since there are only three samples in the dataset, and we group the L5-TTPCI and L5-TTPC) into one class L5-TTPC following previous work that found that they rather form a continuum then two separate classes
(Gowwens et al. 2019 Kanari et al. 2019).
For the clustering analysis and the comparison to Gouwens et al. 20199 we follow Gouwens et al. (2019) and compute the adjusted rand index between our found clusters and the 14 export labels. To determine the optimal number of clusters, we use voos-validation to compute the lig-likelihood of held-out data of the Gaussian Mixture model and choose the number of clusters with the highest log-likelihood. The optimal
22

---

Published in Transactions on Machine Learning Research (06/2023)
number of clusters is 15 for the BBP dataset. To perform clustering within cortical layers, chose the number of clusters per layer based on the number of clusters with the majority of cells from the cortex-wide clustering (Fig. 3 four for layer 2/3, layer 5 and layer 6 and three for layer 4.
C.5.2 Comparison to InfoGraph (Sun et al., 2020)
We use the official implementation to train INFOGRAPH on the BBP dataset. We perform a hyperparameter search for INFOGRAPH as detailed in the original publication (Sun al. 2020 and extend it to include more training epochs to train it for approximately the same number of iterations as GRAPHDINO. In detail, we run grid search over learning rate (Ir) {10-2 10-3, 100111000$111 10-*), number of training epochs e {10, 20, 100, 200, 1, 000, 2, 000} and GNN layers E {4, 8, 12]. We select the hyperparameters based on the lowest unsupervised loss. The chosen hyperparameters are: Ir = 0.001, epochs = 1, 000 and four GNN layers with a hidden dimensionality of 32.
We evaluate the performance of INFOGRAPH (Sun et al. 2020 using a kNN classifier analogous to the ablation experiments (Appendix. C.5.1
C.5.3 Comparison to MorphVAE (Laturnus & Berens
2021
We follow the evaluation protocol of Laturnus & Berens 202 and perform k-nearest neighbor (kNN) classification with A = 5 on the learned latent embeddings of the excitatory neurons to predict whether they are untufted, tufted or "other" on the test set (n = 60) and report the balanced accuracy. The "other" class only contains six examples the test set. To get an estimate of the variance that is due to the chosen data split, we additionally evaluate three further data splits and report the average test set performance over the three splits. We report the performance of MORPIVAE as given in Tab. of Laturnus & Berens (2021)
C.5.4 Comparison to TreeMoCo (Chen et al., 2022b)
A fair comparison to TrecMoCo proved difficult. We tried to replicate their setting as best as possible from the information given in the paper as well as by inferring it from their code base while trying to set up a more fair benchmark for future works.
We downloaded the three datasots BIL, JML and ATC using the official code base of TREEMoCo and used it assign the eleven class labels: L1, L2/3, L4, L5, L6, VPM, CP, VPL, SUB, PRE, MG and Others as used by Chenn et al. 2022b). However, our cell counts slightly differ from those given in Chen et al. (002bb) More specifically, the JML dataset contained 1,200 neurons instead of 1,107.
Chen et al. (2022b) removed a substantial amount of the neurons (995 of 3,358 neurons) from the datasets duc to reconstruction errors. Since we did not have access to the identities of these ncurons, we trained GRAPHDINO unsupervised on all cells with more than 200 nodes (Mrotel = 3, 138; nass. = 1, 739, nJMI = 889, BACT = 510) and evaluated the proposed classes as assigned by the TREEMoCo code base. We replicated the proposed data preprocessing by centering the somata at (0, 0, 0) and aligning the neurons' first principal component to the y-axis.
Chen et al. (2022b) performs the quantitative evaluation on a 80-20% training-test data split. Since we did not have access to the exact split, we performed five cross-validations instead and report the average accuracy over folds.
According to the paper, Chen et al. (202bb) perform k-nearest neighbor classification (A 5 or = 20 depending on the dataset). We unify the evaluation and report the kNN accuracy with k 5 for all list the k = 20 performance in Tab. C.4 In their code base,
experiments in this paper. For reference, the implementation of kNN is weighted, where the neighbors vote is weighted by the cosine similarity of the embeddings. We follow the description in the paper (Chen et al.] 2022bb and use the standard kNN classification without weighing the neighbors' voles.
com/ www.s/uunnnnnnnnnaammmm IIfoorrpp
We additionally tried reach out to the authors but did not get reply.
23

---

Published in Transactions on Machine Learning Research (06/2023)
Table C.4: Cell-type classification on the TreeMoCo dataset. Performance of our model (GRAPHDINO) averaged over three random soeds and given mcan + standard deviation when using 20 for the kNN classifer.
Figure C.5: Example neurons labeled Isocortex 4 of the ACT dataset.
Model
BIL-6
ACT
Ours
78 =
54=
The performances reported by Chen et. al. (2022b) overfitted the test set: They picked the best test set performance over epochs (for the three datasets separately) (see Fig. C1 in Chen et al. (2022b) ) Additionally, they picked whether to use the latent embedding the test set performance per dataset. To give an estimate of the less overfitted performance of TREEMoCo (Chen et al. 2022b) (at least with respect to which epoch to evaluate), we report the averaged performance over the last five epochs given by Fig. C1 (Chen et al. 2022b
the projection head's output p based on
Similarly, the performance of GRAPHCL (You al. 2020) as reported by Chen et al. (2022b) is picked the best test set performance per dataset over training epochs. We therefore report the average accuracy over the last five epochs with the given by Fig. C1 (Chen et al. 2022b
D Complete cluster visualizations
In the Fig. D.6 and Fig. D.7| show the cluster assignments of all samples of the excitatory BBP dataset (n 286) and the Trees dataset (n 391), respectively.
24

---

Published in Transactions on Machine Learning Research (06/2023)
Spiny 1
Figure D.6: Clusters of spiny neurons of BBP dataset as identified by GMM based on our learned feature space. Apical dendrites are colored lighter, while basal dendrites are shown in darker color. Soma is marked by black circle.
25

---

Published in Transactions on Machine Learning Research (06/2023)
13
Cluster
12
Figure D.7: Clusters of trees as identified by GMM based on our learned feature space.
26


Reviews:
[{"review_id": "cmkoJ7Jb01", "reviewer": "Reviewer_Tde6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: GraphDINO is generic to any graph-like data structure with similar node features, and appears to not require heavy tailoring for each problem. The transformer attention mechanism is extended feasibly to graph-like structures with a novel and simple adjacency matrix bias. Experimental analyses are performed well in depth and addresses all related aspects of the analyses and proposed model.\n\nWeaknesses: Some minor changes are still necessary for clarification, especially in terms of notation for the architecture.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a self-supervised representation learning approach (based on DINO [Caron et al., 2021]), that obtains graph-level embeddings using a transformer-based architecture. Using graph data augmentation, e.g., with node subsampling, random rotation, node jittering, etc., the model is trained to maximize similarity between the embeddings for different augmented views of an input. The transformer-based architecture uses a novel, modified \"adjacency-conditioned\" attention mechanism, which imposes a learned bias on the attention values towards the adjacency matrix of the input graph. Experimental validations are thoroughly performed (e.g., on downstream classification tasks) by learning representations from graph-like neuronal morphologies on five publicly available datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the paper's claims are well supported in the manuscript, and there are no major concerns on my side. Only minor questions below:\n\n- It appears like there is no constraint on the bias factors \\lambda and \\gamma, which are learned per-node. Would the approach still reasonably work if one reduces it to a single learned parameter vector by setting e.g., \\lambda=1-\\gamma, and hence reducing the number of parameters in W by D?\n- How sensitive is the learning algorithm to the softmax temperature of the teacher encoder?\n- Variable p is not defined in the main manuscript or the supplementary until Table C.1 where the chosen values are listed? It is also confusing that it might be referring to the DINO notation in Figure C.4. A clarification on how p relates/differs from z should be added for the reader.\n- In Figure 1.B caption, the embeddings are stated to lie on a 32-dim vector space, which is not necessarily fixed and can vary. Since this becomes confusing later, perhaps this should be discussed parameterically in text and clarified in 4.2.2 with the chosen values.\n- In Table C.1, the embeddings are chosen to be 64 dimensional for the M1 PatchSeq dataset, although in Section 4.2.2 it was stated that for neuronal datasets it was fixed to 32. This should be corrected. Also since N denotes something else, in Table C.1, N layers -> \\# layers, and N heads -> \\# heads.\n- Typo on page 9, second line of Section 5.5: \"1'000\" -> \"1,000\".\n\n", "broader_impact_concerns": "No specific concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "BOZh7HykWC", "reviewer": "Reviewer_wWM4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- This work makes good use of the existing knowledge on SSL and graph transformers:  proposed strategies for domain specific data augmentation and attention on neuronal graphs make sense.\n- The authors did a good job at demonstrating the usefulness of their representation. SSL and graph transformers seem a promising direction for studying neuronal morphologies and perhaps making discoveries from large, unlabelled datasets.\n- The paper is overall well written.\n\nWeaknesses:\n- The authors could have done a slightly better job at discussing the relevant literature. For example:\n  - In SSL for graphs, [1] adapted BYOL to graph representation learning.\n  - The authors mention various works proposing graph transformers. These works contain schemes that interpolate between local and global attention, as well as different position encodingss since Laplacian eigenvectors suffer from issues: why are they not suited to the problem of learning representation of neuronal morphologies?\n\n\n- In terms of employed methods, there seems to be some room for improvement that could also be discussed in the limitations.\n  - Taken separately, transformers and SSL are often data hungry. Here, the data only consists of a few thousands of examples: although the need for data is alleviated here since the authors inject some inductive biases in the architecture (namely the adjacency matrix in attention and a non-learned positional encoding), have you considered looking for more data? This may lead to substantial improvements. This could also improve the instability of the results (error bars seem a bit high).\n  - In computer vision, having a balanced hence curated dataset (in terms of underlying class) is crucial to achieve state of the art performance with SSL. Have you considered this issue here?\n  - The authors seek a representation that is invariant to translations and 3D rotations. Instead of enforcing this property through SSL data augmentation, have you considered using equivariant architectures such as [2][3]? ([2] can be adapted to transformers). This added inductive bias could also alleviate the need for data.\n\n- Could you clarify why masking node features was not a good augmentation strategy in your setup? In SSL, masking is generally simple and beneficial.\n\n[1] Thakoor et al. Large-Scale Representation Learning on Graphs via Bootstrapping.\n\n[2] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks.\n\n[3] Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes a representation method for 3D neuronal morphologies. Seeing neuron skeletons as spatial trees, the paper proposes (i) modification to the transformer architecture suited to graph processing (ii) domain specific data augmentations to be used in the context of DINO (iii) extensive study of the resulting representation.\n\n(i) Self-attention is modified so that it interpolates between global attention and attention limited to neighboring nodes only. This interpolation is learned and depends on the node features.\n\n(ii) Data augmentations are central in non-contrastive SSL since these methods enforce similar embeddings for two different augmented versions of the same sample. Therefore, data augmentation determines to some extent to what modification the final representation will be invariant. Domain specific data augmentations are therefore crucial. Here, the authors argue that their embedding should be invariant to (a) tree specific subsampling (b) 3D rotation (c) Gaussian noise added to node coordinates (d) tree specific subgraph deletion (e) random translations of the node coordinates.\n\n(iii) The authors demonstrate that the learned representation retains many characteristics of interest for practitioners, e.g., recovering known excitatory cell types, or forming clusters that are consistent with expert labels. The learned representations seem to outperform existing baselines. They also perform ablations on their proposed augmentation and attention strategies.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above (all equally important): \n- Discussing some missing work, or cited work that haven't been used here. \n- Discussing limitations and potential improvements.\n- Clarification of node feature masking.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YCcurLCeg6", "reviewer": "Reviewer_ooUU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1. This paper is well written, with clear motivations and organizations.\n2. The proposed method is novel to some extent, with tailored data augmentation methods for the neuronal skeleton graph (and other similar graphs).\n3. This paper offers a lot of ablation studies and visualizations. They are helpful in understanding the benefits of the proposed method, and demonstrate the advantages of this method compared to the prior works.\n\nWeakness:\n1. The novelty of the proposed method is not very significant. The overall idea is built upon DINO. The positional encoding strategy is borrowed from previous work. The major novelty lies in the data augmentation, as well as in the hybrid combination of attention and adjacency matrix. Overall I think such technical novelties are not significant.\n2. Usually for the DINO method on image domain, it needs a lot of training data. In this paper I find the proposed method was only trained on a few hundred of graphs (e.g., 646). It is surprising that the trained model on such a small dataset still works quite well. Is this due to some kind of overfitting? Or is it possible that the task is not challenging enough? Can we pre-train the model on the synthetic dataset and then fine-tune it on this small dataset (in a self-supervised manner)?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a new self-supervised method for graph representation learning. This learning method is specifically designed for neuronal morphologies representation learning, but can be generalized to other similar graphs to some extent. The contributions of this paper mainly include: (1) a new self-supervised graph-level feature learning method; (2) a new attention module is introduced; (3) this paper offers a bunch of ablation studies and visualizations, which demonstrate the effectiveness of the proposed method. Overall, the findings of this paper are interesting.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "One interesting question is that: can we pre-train the model on the synthetic dataset and then fine-tune it on another small dataset? Is it possible that this pre-training will further boost the performance?", "broader_impact_concerns": "I do not see any ethical issues of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers agree that the manuscript contains interesting results, in particular for the field of computational neuroscience. It also tackles an interesting input modality that has not been considered in this context. The technical novelty is limited since similar graph transformers exist. Nevertheless, the novel application scenario and strategies to adapt the method to the considered domain are worth publication.
Decision: Accept (poster)