[Paper Review] Self-Supervised Graph Representation Learning for Neuronal Morphologies

Marissa A. Weis, Laura Pede+2 more|arXiv (Cornell University)|Dec 23, 2021

Functional Brain Connectivity Studies4 citations

TL;DR

This paper introduces GraphDINO, a self-supervised, transformer-based method for learning low-dimensional graph-level representations of 3D neuronal morphologies without manual feature engineering or annotations. By combining data augmentation, positional encoding adaptation, and a novel AC-Attention mechanism that blends global attention with local graph convolution, GraphDINO achieves clustering performance on par with expert-labeled cell types and outperforms prior methods on predicting expert labels across multiple brain regions and species.

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.

Motivation & Objective

• To develop a data-driven, self-supervised method for learning graph-level embeddings of 3D neuronal morphologies without relying on human-defined morphological features.
• To address the limitations of manual classification and feature engineering, which introduce bias and high inter-expert variance.
• To enable the discovery of novel morphological cell types through unsupervised representation learning on large-scale unlabeled datasets.
• To improve upon existing graph neural network and autoencoder-based approaches in predicting expert-annotated cell type labels.
• To create a generalizable framework applicable beyond neuroscience for spatially-embedded graph representation learning.

Proposed method

• Proposes GraphDINO, a transformer-based model for spatially-embedded tree-structured graphs representing neuronal morphologies.
• Introduces data augmentation strategies tailored for 3D spatial graphs, including rotation, translation, and subsampling, to enable self-supervised contrastive learning.
• Adapts positional encoding to preserve geometric invariance under spatial transformations such as rotation and translation.
• Develops AC-Attention, a hybrid attention mechanism that combines global self-attention with local graph convolution via node similarity-based routing.
• Employs contrastive learning with view augmentation to train the model on unlabeled data, optimizing for invariant representations under structural perturbations.
• Uses graph-level pooling to produce a single embedding per neuron for downstream clustering and classification tasks.

Experimental results

Research questions

• RQ1 Can a self-supervised, end-to-end learned representation of 3D neuronal morphologies match the quality of expert-based morphological classification?
• RQ2 How does GraphDINO’s performance compare to hand-crafted feature engineering and autoencoder-based baselines in predicting expert-labeled cell types?
• RQ3 To what extent can GraphDINO discover biologically meaningful cell type clusters without prior knowledge of morphological features?
• RQ4 Does the integration of global attention and local message passing via AC-Attention improve representation learning on spatial graphs compared to standard GNNs or transformers?
• RQ5 Can the method generalize across species and brain regions without architectural or hyperparameter re-tuning?

Key findings

• GraphDINO achieves morphological cell type clustering performance that is on par with manual feature-based classification by neuroscience experts, without using any prior knowledge of structural features.
• The method outperforms both traditional hand-engineered feature baselines and autoencoder-based approaches in predicting expert-annotated cell type labels on benchmark datasets.
• On the ACT dataset, GraphDINO’s performance on a subset restricted to excitatory neurons was significantly more meaningful than previous evaluations that mixed excitatory and inhibitory types.
• The AC-Attention mechanism enables a smooth interpolation between global attention and local graph convolution, improving representation quality on spatial graphs.
• Data augmentation via subsampling and spatial transformations effectively preserves global morphology while enabling robust self-supervised pretraining.
• The model generalizes across species (mouse and human) and multiple brain areas, demonstrating robustness to biological and structural variation.