We examine the utility of graph neural networks for the purpose of learning cortical segmentations. We show that attention-based transformer networks significantly outperform conventional GCN and linear feed-forward variants for the purpose of generating accurate reproducible cortical maps.
Jumping knowledge networks let each node combine embeddings from every layer, so the effective neighborhood size can vary across the graph. I explain the influence distribution argument and the LSTM attention aggregator. Then I apply a jumping knowledge GAT to cortical segmentation, where it holds up better than a plain GAT as depth increases.
Wang et al. argue that graph attention networks overfit their attention weights and oversmooth signals across class boundaries. Their fix adds two margin-based losses and aggregates only over the top-k attended neighbors. I walk through the method and my implementation of the layer in DGL.
Implementing a graph convolution layer whose filters are learned Gaussian kernels over node features, written as a custom DGL message-passing layer for segmenting the cortical surface.