Machine Learning |

Jumping-Knowledge Representation Learning With LSTMs

Sun, 14 Feb 2021 23:24:17 -0700

Background

As I mentioned in my previous post on constrained graph attention networks, graph neural networks suffer from overfitting and oversmoothing as network depth increases. These issues can ultimately be linked to the local topologies of the graph.

If we consider a 2d image as a graph (i.e. pixels become nodes), we see that images are highly regular – that is, each node has the same number of neighbors, except for those at the image periphery. When we apply convolution kernels over node signals, filters at any given layer are aggregating information from the same-sized neighborhoods irrespective of their location.

However, if we consider a graph, there is no guarantee that the graph will be regular. In fact, in many situations, graphs are highly irregular, and are characterized by unique topological neighborhood properties such as tree-like structures or expander graphs, that are sparse yet highly connected. If we compare an expander node to a node whose local topology is more regular, we would find that the number of signals that each node implicitly convolves at each network layer would vary considerably. These topological discrepancies have important implications when we consider problems like node and graph classification, as well as edge prediction. The problem ultimately boils down to one of flexibility: can we account for unique local topologies of a graph in order to dynamically aggregate local information on a node-by-node basis?

Node signal aggregation as a function of network depth. At each layer, the neural network implicitly aggregates signals over an increasingly-larger neighborhood. In this example, the network is highly regular – however, not all graphs are.

In a recent paper, the authors propose one approach to address this question, which they call “jumping knowledge representation learning”¹. Instead of utilizing the output of the last convolution layer to inform the prediction, jumping-knowledge networks aggregate the embeddings from all hidden layers to inform the final prediction. The authors develop an approach to study the “influence distribution” of nodes: for a given node $x$, the influence distribution $I_{x}$ characterizes how much the final embedding of node $x$ is influenced by the input features of every other node:

$$ \begin{align} I(x,y) &= \sum_{i=1}^{m} |\frac{\delta h_{x}^{k}}{\delta h_{y}^{0}}|{i} \\ I{x}(y) &= I(x,y) \Big/\sum_{z} I(x,z) \end{align} $$

They show that influence distribution $I_{x}$ for a $k$-layer graph convolution network is equal, in expectation, to the $k$-step random walk distribution. They point out that the random walk distribution of expander-like nodes converges quickly – the final embeddings of these nodes are representative of the whole graph and carry global information – while the random-walk distribution of nodes with tree-like topology converges slowly – these nodes carry more-local information². These two conclusions are related to the spectral gap of the graph – the smallest non-zero eigenvalue of the graph Laplacian. A large spectral gap indicates high-connectivity, while a low spectral gap indicates low connectivity. From a graph theory perspective, this local connectivity is related to the idea of centrality. Nodes with high centrality will easily saturate their random walk distribution, but will also aggregate information from large neighborhoods quickly. For graph neural networks with fixed aggregation kernels, this has important implications for representation learning, because the feature distributions of nodes with different topologies will not correspond to the same degree of locality, which may not lead to the best learned representations for all nodes. A radius that is too large may result in over-smoothing of node features, while a radius that is too small may not be robust enough to learn optimal node embeddings.

The jumping knowledge network architecture is conceptually similar to other graph neural networks, and we can, in fact, simply incorporate the jumping knowledge mechanism as an additional layer. The goal is to adaptively learn the effective neighborhood size on a node-by-node basis, rather than enforcing the same aggregation radius for every node (remember, we want to account for local topological and feature variations). The authors suggest three possible aggregation functions: concatenation, max-pooling, and an LSTM-attention mechanism ¹ ³. Each aggregator learns an optimal combination of the hidden embeddings, which is then pushed through a linear layer to generate the final network output. Concatenation determines the optimal linear combination of hidden embeddings for the entire dataset simultaneously, so it is not a node-specific aggregator. Max-pooling selects the most important hidden layer for each feature element on a node-by-node basis – however, empirically, I found that max-pooling was highly unstable in terms of model learning. The LSTM-attention aggregator treats the hidden embeddings as a sequence of elements, and assigns a unique attention score to each hidden embedding ⁴.

Schematic of a jumping-knowledge network. The neural network generates an embedding for each hidden layer. The aggregator function then optimally combines these hidden embeddings to learn the optimal abstraction of input information. Some alternative aggregation functions include max-pooling, concatenation, and an LSTM layer. In the case of an LSTM layer coupled with an attention mechanism, the aggregator computes a convex combination of hidden embeddings.

Long-Short Term Memory

Briefly, given a sequence of samples $X_{1}, X_{2}, \dots X_{t}$, the LSTM cell learns temporal dependencies between elements of a sequence by maintaining a memory of previously observed elements – in our case, the sequence elements are the embeddings learned by each consecutive hidden layer. An LSTM cell is characterized by three gates controlling information flow between elements in the sequence: input, forget, and output, as well as a cell state vector, which captures the memory and temporal dependencies between sequence elements⁵:

Schematic of an LSTM cell. The cell controls what information is remembered from previous elements in a sequence, and what information is incorporated into memory given a new element in the sequence.

$$ \begin{align} f_{t} &= \sigma(W_{f}X_{t} + U_{f}h_{t-1} + b_{f}) \\ i_{t} &= \sigma(W_{i}X_{t} + U_{i}h_{t-1} + b_{i}) \\ o_{t} &= \sigma(W_{o}X_{t} + U_{o}h_{t-1} + b_{o}) \\ \end{align} $$

where $W$, $U$, and $b$ are learnable parameters of the gates. Here, $X_{t}$ is the $t$-th sequence element, $h_{t-1}$ represents the learned LSTM cell embedding for element $t-1$, and $C_{t-1}$ represents the current memory state, given the previous $1, 2 \dots t-1$ elements. The input and forget gates determine which aspects of a sequence element are informative / uninformative, and decide what information to keep / forget, while the output gate combines the previous memory state with our new knowledge. We can roughly think of this process as updating our prior beliefs, in the Bayesian sense, with new incoming data.

$$ \begin{align} \tilde{c_{t}} &= \sigma(W_{c}X_{t} + U_{c}h_{t-1} + b_{c}) \\ c_{t} &= f_{t}\circ c_{t-1} + i_{t} \tilde{c}{t} \\ h{t} &= o_{t} \circ tanh(c_{t}) \end{align} $$

The embeddings for each element learned by the LSTM cell are represented by $h_{t}$. In the original paper¹, the authors propose to apply a bi-directional LSTM to simultaneously learn forwards and backwards embeddings, which are concatenated and pushed through a single-layer perceptron to compute layer-specific attention weights for each node:

$$ \begin{align} \alpha_{i}^{t} &= \sigma(\vec{w}{t}^{T}(h^{F}{i, t} || h^{B}{i, t})) \\ \alpha{i}^{t} &= \frac{\exp{(\alpha_{i}^{t})}}{\sum_{t=1}^{L} \exp{(\alpha_{i}^{t})}} \end{align} $$

The softmax-normalized attention weights represent a probability distribution over attention weights

$$\begin{align} \sum_{t=1}^{L} \alpha_{i}^{t} &= 1 \\ \\ \alpha_{i}^{t} &>= 0 \end{align} $$

where $\alpha_{i}^{t}$ represents how much node $i$ attends to the embedding of hidden layer $t$. The optimal embedding is then computed as the attention-weighted convex combination of hidden embeddings:

$$ \begin{align} X_{i, \mu} = \sum_{t=1}^{L} \alpha_{i}^{t}X_{i, t} \end{align} $$

An Application of Jumping Knowledge Networks to Cortical Segmentation

I’ve implemented the jumping knowledge network using DGL here. Below, I’ll demonstrate the application of jumping knowledge representation learning to a cortical segmentation task. Neuroscientifically, we have reason to believe that the influence radius will vary along the cortical manifold, even if the mesh structure is highly regular. As such, I am specifically interested in examining the importance that each node assigns to the embeddings of each hidden layer. To that end, I utilize the LSTM-attention aggregator. Similarly, as the jumping-knowledge mechanism can be incorporated as an additional layer to any general graph neural network, I will use graph attention networks (GAT) as the base network architecture, and compare vanilla GAT performance to GATs with a jumping knowledge mechanism (JKGAT).

Below, I show the prediction generated by a 9-layer JKGAT model, with 4 attention heads and 32 hidden channels per layer, with respect to the “known” or “true” cortical map. We find slight differences in the performance of our JKGAT model with respect to the ground truth map, notably in the lateral occipital cortex and the medial prefrontal cortex.

Comparison of the group-average predicted cortical segmentation produced by the JKGAT model, to the ground truth cortical segmentation. The ground truth was previously generated here. The consensus cortical map corresponds very well to the true map.

When we consider the accuracies for various parameterizations of our models, we see that the JKGAT performs quite well. Notably, it performs better than the GAT model in most cases. Likewise, as hypothesized, the JKGAT performs better than the GAT model as network depth increases, specifically because we are able to dynamically learn the optimal influence radii for each node, rather than constraining the same radius size for the entire graph. This allows us to learn more abstract representations of the input features by mitigating oversmoothing and by accounting for node topological variability, which is important for additional use-cases like graph classification.

Model accuracy comparison between GAT and JKGAT models on a node classification problem for cortical segmentation. Accuracy is represented as the fraction of correctly-labeled nodes in a graph, across 150 validation subjects. Each node in the graph has 80 features, and each graph has 30K nodes.

Similarly, we find that JKGAT networks generate segmentation predictions that are more reproducible and consistent across resampled datasets. This is important, especially in the case where we might acquire data on an individual multiple times, and want to generate a cortical map for each acquisition instance. Unless an individual suffers from an accelerating neurological disorder, experiences a traumatic neurological injury, or the time between consecutive scans is very long (on the order of years), we expect the cortical map of any given individual to remain quite static (though examining how the “map” of an individual changes over time is still an open-ended topic).

Model reproducibility comparison between GAT and JKGAT models on a node classification problem for cortical segmentation, using 150 validation subjects. Each subject has four repeated datasets. Within a given subject, the topology of each graph is the same, but the node features are re-sampled for each graph. Reproducibility is computed using the F1-score between all pairs of predicted node classifications, such that we compute 6 F1 scores for each subject.

Finally, when we consider the importance that each cortical node assigns to the unique embedding at the $k$-th layer via the LSTM-attention aggregation function, we see very interesting results. Notably, we see high spatial auto-correlation of the attention weights. Even more striking is that this spatial correlation seems to correspond to well-studied patterns of resting-state networks identified using functional MRI. Apart from the adjacency structure of our graphs, we do not encode any a priori information of brain connectivity. That the LSTM-attention aggregator of the jumping-knowledge layer identifies maps corresponding reasonably well to known functional networks of the human brain is indicative, to some extent, of how the model is learning, and more importantly, of which features are useful in distinguishing cortical areas from one another.

Let us consider the attention map for layer 4. We can interpret the maps as follows: for a given network architecture (in this case, a network with 9 layers), we find that areas in the primary motor (i.e. Brodmann areas 3a and banks of area 4) and primary auditory cortex (Brodmann areas A1 and R1) preferentially attend to the embedding of hidden layer 4, relative to the rest of the cortex – this indicates that the implicit aggregation over an influence radius of 4 layers is deemed more informative for the classification of nodes in the primary motor and auditory regions than for other cortical areas. However, whether this says anything about the implicit complexity of the cortical signals of these areas remains to be studied.

Maps of learned LSTM-attention aggregator weights. Each inset corresponds to the weights learned by every cortical node for the $k$-th layer hidden embedding (black: low, red: high). We see that most of the attention mass is distributed over layers 4-7, indicating that most nodes assign maximal importance to intermediate levels of abstraction. However, we do see spatially varying attention. Notably, within a given attention map, we find that nodes of the lateral Default Mode Network preferentially attend to the embeddings of layers 1-3, while layer 4 is preferentially attended to by the primary motor and auditory areas.

Xu et al. Representation Learning on Graphs with Jumping Knowledge Networks. 2018. ↩︎ ↩︎ ↩︎
Dinitz et al. Large Low-Diameter Graphs are Good Expanders. 2017. ↩︎
Lutzeyer et al. Comparing Graph Spectra of Adjacency and Laplacian Matrices. 2017. ↩︎
Gers, Felix. Long Short-Term Memory in Recurrent Neural Networks. 2001. ↩︎
Fan et al. Comparison of Long Short Term Memory Networks and the Hydrological Model in Runoff Simulation. 2020. ↩︎

Constrained Graph Attention Networks

Fri, 25 Dec 2020 23:24:17 -0700

In their recent paper, Wang et al. propose a few updates to the Graph Attention Network (GAT) neural network algorithm (if you want to skip the technical bit and get to the code, click here). Briefly, GATs are a recently-developed neural network architecture applied to data distributed over a graph domain. We can think of graph convolutional networks as progressively transforming and aggregating signals from within a local neighborhood of a node. At each iteration of this process, we implicitly merge signals from larger and larger neighborhoods of the node of interest, and thereby learn unique representations of nodes that are dependent on their surroundings.

GATs incorporate the seminal idea of “attention” into this learning process. In each message-passing step, rather than updating the features of a source-node via equally-weighted contributions of neighborhood nodes, GAT models learn an attention function – i.e. they learn how to differentially pay attention to various signals in the neighborhood. In this way, the algorithm can learn to focus on important signals and disregard superfluous signals. If we consider neural networks as universal function approximators, the attention mechanism improves the approximating ability by incorporating multiplicative weight factors into the learning.

Figure from Velickovic et al. For a source node $i$ and destination node $j$, vectors $\vec{h_{i}}$ and $\vec{h_{j}}$ are the input feature vectors of nodes $i$ and $j$ in layer $l$. $\mathbf{W}$ is a learned affine projection matrix. $\mathbf{\vec{a}}$ is the learned attention function. The source and destination node input features are pushed through the attention layer as $\alpha_{i,j} = \sigma\Big(\vec{a}^{T}\mathbf{W}\Big(\vec{h_{i}} || \vec{h_{j}}\Big)\Big)$ where $\sigma$ is an activation function, and $\alpha_{i,j}$ the unnormalized attention that node $i$ pays to node $j$ in layer $l$. Attention weights are then passed through a softmax layer, mapping the attentions between [0,1].

However, GATs are not without their pitfalls, as noted by Wang et al. Notably, the authors point to two important issues that GATs suffer from: overfitting of attention values and oversmoothing of signals across class boundaries. The authors propose that GATs overfit the attention function because the learning process is driven only by classification error, with complexity $O(|V|)$ i.e. the number of nodes in the graph. With regards to oversmoothing, the authors note that a single attention layer can be viewed as a form of Laplacian smoothing:

$$\begin{align} Y = AX^{l} \end{align}$$

where $A_{n \times n}$ is the attention weight matrix with $A_{i,j} = \alpha_{i,j}$ if $j \in \mathcal{N_{i}}$ and $0$ otherwise. Because $\sum_{j\in \mathcal{N_{i}}} \alpha_{i,j} = 1$, we can view $A$ as a random walk transition probability matrix. If we assume that graph $G=(V,E)$ has $K$ connected components, repeated application of $A$ to $X$ distributed over $G$ will result in a stationary distribution of node features within each connected component – that is, the feature vectors of the nodes within each connected component will converge on the component mean. However, as the authors point out, we typically have multiple layers $l_{1}\dots l_{j}$, each with their own attention matrix $A_{1} \dots A_{j}$, each representing a unique transition probability matrix. Because we generally do not have disconnected components, nodes from different classes will be connected – consequently, deep GAT networks will mix and smooth signals from different adjacent components, resulting in classification performance that worsens with network depth. Importantly, multi-head attention networks do not alleviate this convergence issue – each head can be viewed as a unique probability transition matrix, which all suffer from the same oversmoothing issue as $l \rightarrow \infty$.

Wang et al. propose to incorporate two margin-based constraints into the learning process. The first constraint, $\mathcal{L_{g}}$, addresses the overfitting issue, by enforcing that learned attentions between adjacent pairs of nodes be higher than attentions between distant pairs of nodes.

$$\begin{align} \mathcal{L_{g}} &= \sum_{i \in V} \sum_{j \in \mathcal{N_{i}} \setminus \mathcal{N_{i}^{-}}} \sum_{k \in V\setminus \mathcal{N_{i}}} max(0, \phi(v_{i},v_{k}) + \zeta_{g} - \phi(v_{i},v_{j})) \end{align}$$

The second constraint, $\mathcal{L_{b}}$, addresses the oversmoothing issue, by enforcing that learned attentions between pairs of adjacent nodes with the same label be higher than attention between pairs of adjacent nodes with different labels:

$$\begin{align} \mathcal{L_{b}} &= \sum_{i \in V}\sum_{j \in \mathcal{N_{i}^{+}}} \sum_{k \in \mathcal{N_{i}^{-}}} max(0, \phi(v_{i},v_{k}) + \zeta_{b} - \phi(v_{i},v_{j})) \end{align}$$

In both cases, $\phi(,)$ is the attention function between a pair of nodes, $\mathcal{N_{i}^{+}}$ and $\mathcal{N_{i}^{-}}$ are the nodes adjacent to node $i$ with the same (+) and different (-) labels as $i$, and $\zeta_{g}$ and $\zeta_{b}$ are slack variables controlling the margin between attention values. The first loss function, $\mathcal{L_{g}}$, can be implemented via negative sampling of nodes (the authors actually perform importance-based negative sampling based on attention-weighted node degrees, but showed that this only marginally improved classification accuracy in benchmark datasets).

The authors propose one final addition to alleviate the oversmoothing issue posed by vanilla GATs. Rather than aggregating over all adjacent nodes in a neighborhood, the authors propose to aggregate over the nodes with the $K$ greatest attention values. Because the class boundary loss $\mathcal{L_{b}}$ enforces large attentions on nodes of the same label and small attention on nodes of different labels, aggregating over the top $K$ nodes will tend to exclude nodes of different labels than the source node in the message passing step, thereby preventing oversmoothing. The authors show that this constrained aggregation approach is preferable to attention dropout proposed in the original GAT paper. Taken together, the authors deem these margin-based losses and constrained aggregation “Constrained Graph Attention Network” (C-GAT).

Implementation

I wasn’t able to find an implementation of the Constrained Graph Attention Network for my own purposes, so I’ve implemented the algorithm myself in Deep Graph Library (DGL) – the source code for this convolutional layer can be found here. This implementation makes use of the original DGL GATConv layer structure, with modifications made for the constraints and aggregations. Specifically, the API for CGATConv has the following modifications:


CGATCONV(in_feats,
 out_feats,
 num_heads,
 feat_drop=0.,
 graph_margin=0.1, # graph structure loss slack variable
 class_margin=0.1, # class boundary loss slack variable
 top_k=3, # number of messages to aggregate over
 negative_slope=0.2,
 residual=False,
 activation=None,
 allow_zero_in_degree=False)

Of note is the fact that the attn_drop parameter has been substituted by the top_k parameter in order to mitigate oversmoothing, and the two slack variables $\zeta_{g}$ and $\zeta_{b}$ are provided as graph_margin and class_margin.

With regards to the loss functions, the authors compute all-pairs differences between all edges incident on a source node, instead of summing over the positive / negative sample attentions ($\mathcal{L_{g}}$) and same / different label attentions ($\mathcal{L_{b}}$) and then differencing these summations. In this way, the C-GAT model anchors the loss values to specific nodes, promoting learning of more generalizable attention weights. The graph structure loss function $\mathcal{L_{g}}$ is implemented with the graph_loss reduction function below:

def graph_loss(nodes):

 """
 Loss function on graph structure.

 Enforces high attention to adjacent nodes and
 lower attention to distant nodes via negative sampling.
 """

 msg = nodes.mailbox['m']

 pw = msg[:, :, :, 0, :].unsqueeze(1)
 nw = msg[:, :, :, 1, :].unsqueeze(2)

 loss = (nw + self._graph_margin - pw).clamp(0)
 loss = loss.sum(1).sum(1).squeeze()

 return {'graph_loss': loss}
.
.
.
graph.srcdata.update({'ft': feat_src, 'el': el})
graph.dstdata.update({'er': er})
graph.apply_edges(fn.u_add_v('el', 'er', 'e'))
e = self.leaky_relu(graph.edata.pop('e'))

# construct the negative graph by shuffling edges
# does not assume a single graph or blocked graphs
# see cgatconv.py for ```construct_negative_graph``` function
neg_graph = [construct_negative_graph(i, k=1) for i in dgl.unbatch(graph)]
neg_graph = dgl.batch(neg_graph)

neg_graph.srcdata.update({'ft': feat_src, 'el': el})
neg_graph.dstdata.update({'er': er})
neg_graph.apply_edges(fn.u_add_v('el', 'er', 'e'))
ne = self.leaky_relu(neg_graph.edata.pop('e'))

combined = th.stack([e, ne]).transpose(0, 1).transpose(1, 2)
graph.edata['combined'] = combined
graph.update_all(fn.copy_e('combined', 'm'), graph_loss)

# compute graph structured loss
Lg = graph.ndata['graph_loss'].sum() / (graph.num_nodes() * self._num_heads)

Similarly, the class boundary loss function $\mathcal{L_{b}}$ is implemented with the following message and reduce functions:

def adjacency_message(edges):

 """
 Compute binary message on edges.

 Compares whether source and destination nodes
 have the same or different labels.
 """

 l_src = edges.src['l']
 l_dst = edges.dst['l']

 if l_src.ndim > 1:
 adj = th.all(l_src == l_dst, dim=1)
 else:
 adj = (l_src == l_dst)

 return {'adj': adj.detach()}

def class_loss(nodes):

 """
 Loss function on class boundaries.

 Enforces high attention to adjacent nodes with the same label
 and lower attention to adjacent nodes with different labels.
 """

 m = nodes.mailbox['m']

 w = m[:, :, :-1]
 adj = m[:, :, -1].unsqueeze(-1).bool()

 same_class = w.masked_fill(adj == 0, np.nan).unsqueeze(2)
 diff_class = w.masked_fill(adj == 1, np.nan).unsqueeze(1)

 difference = (diff_class + self._class_margin - same_class).clamp(0)
 loss = th.nansum(th.nansum(difference, 1), 1)

 return {'boundary_loss': loss}
.
.
.
graph.ndata['l'] = label
graph.apply_edges(adjacency_message)
adj = graph.edata.pop('adj').float()

combined = th.cat([e.squeeze(), adj.unsqueeze(-1)], dim=1)
graph.edata['combined'] = combined
graph.update_all(fn.copy_e('combined', 'm'), class_loss)
Lb = graph.ndata['boundary_loss'].sum() / (graph.num_nodes() * self._num_heads)

And finally, the constrained message aggregation is implemented using the following reduction function:

def topk_reduce_func(nodes):

 `"""
 Aggregate attention-weighted messages over the top-K
 attention-valued destination nodes
 """

 K = self._top_k

 m = nodes.mailbox['m']
 [m,_] = th.sort(m, dim=1, descending=True)
 m = m[:,:K,:,:].sum(1)

 return {'ft': m}
.
.
.
# message passing
if self._top_k is not None:
 graph.update_all(fn.u_mul_e('ft', 'a', 'm'),
 topk_reduce_func)
else:
 graph.update_all(fn.u_mul_e('ft', 'a', 'm'),
 fn.sum('m', 'ft'))

Cross-Entropy With Structure

Wed, 09 Dec 2020 01:12:32 -0700

As I mentioned in my previous post, I work with cortical surface segmentation data. Due to the biology of the human brain, there is considerable reproducible structure and function across individuals (thankfully!). One manifestation of this reproducibility is exemplified by the neocortex a.k.a. the thin (~2.5mm) gray matter layer of cell-bodies at the periphery of the brain. The neocortex is well known to have local homogeneity in terms of types of neuronal cells, protein and gene expression, and large-scale function, for example. Naturally, researchers have been trying to identify discrete delineations of the cortex for nearly 100 years, by looking for regions of local homogeneity of various features along the cortical manifold.

As in my previous post, I’m working on this problem using graph convolution networks (GCN). Given the logits output by a forward pass of a GCN, I want to classify a cortical node as belonging to some previously identified cortical area. Using categorical cross-entropy, we can calculate the loss of a given forward pass of the model $h(X; \Theta)$:

$$ \begin{align} L = -\sum_{k=1}^{K} \sum_{l \in \mathcal{L}} x_{l}^{k} \cdot log(\sigma(x^{k})_{l}) \end{align} $$

where $x^{k}$ is the output of the model for a single node, $x_{l}^{k}$ is the one-hot-encoding value of the true labels, and $\sigma$ is the softmax function. Importantly, the cross-entropy cost is high when the probability assigned to the true label of a node is small i.e. $log(0) = \infty$, while $log(1) = 0$ – as such, the cross-entropy tries to minimize the rate of false negatives.

However, we can incorporate more structure into this loss function. As I mentioned previously, we know that the brain is highly reproducible across individuals. In our case, we have years of biological evidence pointing to the fact that functional brain areas i.e. like the primary visual area (V1), will always be in the same anatomical location i.e. posterior occipital cortex – and will always be adjacent to a small subset of other functionally-defined areas, like the secondary visual area (V2), for example.

Various maps of the primate visual cortex. Tootell et al, 2003.

This leads us to the idea of assigning a high cost when nodes which should be in V1, for example, are assigned labels of regions that are not adjacent to V1. We do so by defining another cost function:

\begin{align} G = -\sum_{k=1}^{k}\sum_{l \in \mathcal{L}} \sum_{h \in \mathcal{L} \setminus \mathcal{N_{l}}} w_{l}^{k} \cdot log(1-\sigma(x^{k})_{l}) \end{align}

where $w_{l}^{k}$ is the probability weight assigned to label $h \in \mathcal{L}\setminus \mathcal{N_{l}}$ i.e. the set of labels not adjacent to label $l$. In order to follow the idea of a cross-entropy, we enforce the following constraints on weights $\mathbf{w}$:

$$ \begin{align} w_{l}^{k} >&= 0 \\ \sum_{l \in \mathcal{L}} w_{l}^{k} &= 1 \end{align} $$

such that the vector $\mathbf{w}$ is a probability distribution over labels. Importantly, if we consider more closely what this loss-function is doing, we are encouraging the predicted label of $x^{k}$ to not be in the set $\mathcal{L} \setminus \mathcal{N_{l}}$. Assume, for example, that the true label of $x^{k}$ is $t$, and that label $j$ is not adjacent to label $t$ on the cortical surface. If the softmax function assigns a probability $p(x^{k}{l} = j) = 0.05$, then $log(1-p(x^{k}{l} = j))$ will be small. However, if $p(x^{k}{l} = j) = 0.95$, then $log(1-p(x^{k}{l} = j))$ will be large. Consequently, we penalize higher probabilities assigned to labels not adjacent to our true label – i.e. ones that are not even biologically plausible. If a candidate label of $x^{k}{l} \in \mathcal{N{t}}$, we simply set $w_{l}^{k} = 0$ – that is, we do not penalize the true label (obviously), or labels adjacent to the true label, since these are the regions we really want to consider.

Below, I’ve implemented this loss function using Pytorch and Deep Graph Library. Assume that we are given the adjacency matrix of our mesh, the logits of our model, and the true label of our training data:

import numpy as np

import dgl
import dgl.function as fn
import torch.nn.functional as F

import torch as th

def structured_cross_entropy(graph, logits, target):

 """
 Compute a structured cross-entropy loss.

 Loss penalizes high logit probabilities assigned to labels
 that are not directly adjacent to the true label.

 Parameters:
 - - - - -
 graph: DGL graph
 input graph structure
 input: torch tensor
 logits from model
 target: torch tensor
 true node labeling
 Returns:
 - - - -
 loss: torch tensor
 structured cross-entropy loss
 """

 # compute one-hot encoding of true labels
 hot_encoding = F.one_hot(target).float()

 # identify adjacent labels
 weight = th.matmul(hot_encoding.t(),
 th.matmul(graph.adjacency_matrix(), hot_encoding))
 weight = (1-(weight>0).float())

 # compute inverted encoding (non-adjacent labels receive value of 1)
 inv_encoding = weight[target]
 # weight by 1/(# non adjacent)
 # all non-adjacent labels receive the same probability
 # adjacent labels and self-label receive probability of 0
 inv_encoding = inv_encoding / inv_encoding.sum(1).unsqueeze(1)
 loss = th.sum(inv_encoding*th.log(1-F.softmax(logits)), 1)

 return -loss.mean()

If we wanted to use this loss function in conjunction with another loss, like the usual cross-entropy, we could perform something like the following:


# define a regularizing parameter
gamma = 0.1
# define the usual cross-entropy loss function
loss_fcn = torch.nn.CrossEntropyLoss()
loss = loss_function(logits, target) + gamma*structured_cross_entropy(graph, logits, target)

# because our new loss functions performs computations using Pytorch
# the computation history is stored, and we can compute the gradient 
# with respect to this combined loss as

optimizer.zero_grad() # zero the gradients (no history)
loss.backward() # compute new gradients
optimizer.step() # update weights and parameters w.r.t new gradient

Because we’re optimizing two loss functions now i.e. the global accuracy of the model as defined using the conventional cross-entropy, and the desire for predicted labels to not be far away from the true label using the structured cross-entropy, this combination of loss functions will likely have the effect of slightly reducing global accuracy – however, it will have the effect of generating predictions showing fewer anatomically spurious labels i.e. we are less likely to see vertices in the frontal lobe labeled as V1, or vertices in the lateral parietal cortex labeled as Anterior Cingulate. Global predictions will be more biologically plausible. While GCNs as a whole are already better able to incorporate local spatial information than other models due to the fact that they convolve signals based on the adjacency structure of the network in question, I have found empirically that these anatomically spurious predictions are still possible – hence the need for this more-structured regularization.

Gaussian Graph Convolutional Networks

Mon, 07 Dec 2020 23:24:17 -0700

I’m using graph convolutional networks as a tool to segment the cortical surface of the brain. This research resides in the domain of node classification using inductive learning. By node classification, I mean that we wish to assign a discrete label to cortical surface locations (nodes / vertices in a graph) on the basis of some feature data and brain network topology. By inductive learning, I mean that we will train, validate, and test on datasets with possibly different graph topologies – this is in contrast to transductive learning that learns models that do not generalize to arbitrary network topology.

In conventional convolutions over regular grid domains, such as images, using approaches like ConvNet, we learn the parameters of a sliding filter that convolves the signal around a pixel of interest $p_{i,j}$, such that we aggregate the information from pixels $p_{\Delta i, \Delta j}$ for some fixed distance $\Delta$ away from $p$. Oftentimes, however, we encounter data that is distributed over a graphical domain, such as social networks, journal citations, brain connectivity, or the electrical power grid. In such cases, concepts like “up”, “down”, “left”, and “right” do not make sense – what does it mean to be “up” from something in a network? – so we need some other notion of neighborhood.

In come graph convolutional networks (GCNs). GCNs generalize the idea of neighborhood aggregation to the graph domain by utilizing the adjacency structure of a network – we can now aggregate signals near a node by using some neighborhood around it. While vanilla GCNs learn rotationally-invariant filters, recent developments in the world of Transformer networks have opened up the door for much more flexible and inductive models (see: Graph Attention Networks, GraphSAGE).

Demonstration of graph convolution network from Thomas Kipf.

I was specifically interested in applying the methodology described here, where the authors utilize Gaussian kernels as filters over the neighborhood of nodes. However, the authors did not open-source their code – as such, I needed to implement this method myself. Assume our input data to layer $l$ is $Y^{(l)} \in \mathbb{R}^{N \times q}$ for $N$ nodes in the graph. We can define the Gaussian kernel-weighted convolution as follows:

$$ \begin{align} z_{i,p}^{(l)} = \sum_{j \in \mathcal{N}{i}} \sum{q=1}^{M_{(l)}} \sum_{k=1}^{K_{(l)}} w_{p,q,k}^{(l)} \cdot y_{j,q}^{(l)} \cdot \phi(\hat{\mu}{i}, \hat{\mu}{j}; \Theta_{k}^{(l)}) + b_{p}^{(l)} \end{align} $$

Above, $y_{j,q}^{(l)}$ is the $q$-th input feature of neighboring node $j$, $w_{p,q,k}^{(l)}$ is the linear weight assigned to this feature for the $k$-th kernel, and $\phi(\hat{\mu}{i}, \hat{\mu}{j}; \Theta_{k}^{(l)})$ is the $k$-th kernel weight between node $i$ and node $j$, defined as:

$$ \begin{align} \phi(\hat{\mu_{i}}, \hat{\mu_{j}}; \sigma_{k}^{(l)}, \mu_{k}^{(l)} ) = \exp^{-\sigma_{k}^{(l)} \left\Vert (\hat{\mu_{i}} - \hat{\mu_{j}}) - \mu_{k}^{(l)} \right\Vert^{2}} \end{align} $$

Extrinsically, the kernel weights are represented by edges in a sparse affinity matrix, such that index $(i,j)$ is the Gaussian kernel weight between node $i$ and node $j$ for the $k$-th kernel in the $l$-th layer, where nodes $j$ are restricted to be within a certain neighborhood or distance of node $i$. This can be seen more clearly here:

Figure from Wu et al.. $v_{i}$ is our voxel of interest, and $v_{k}^{i}$, for demonstration purposes, is an adjacent node. Both $v_{i}$ and $v_{k}^{i}$ are characterized by embedding vectors $e_{i}, e_{k}^{i} \in \mathbb{R}^{q}$, from which we compute the kernel weight $\phi_{i,k}$ characterizing how similar the two vertices’ embedding vectors are.

I implemented a new convolutional layer called GAUSConv (available here). To implement this algorithm, I utilized the Deep Graph Library (DGL), which offers a stellar single unified API based on message passing (I’m using Pytorch as the backend). I noticed that I could formulate this problem using attention mechanisms described in the Graph Attention Network paper – however, instead of computing attention weights using a fully connected layer as described in that work, I would compute kernel weights using Gaussian filters. Similarly, just as the GAT paper describes multi-head attention for multiple attention channels, I could analogize my formulation to multi-head kernels for multiple kernel channels. To this end, I could make use of the GATConv API quite easily by replacing the attention computations with the Gaussian kernel filtrations. Likewise, I utilized the GraphConv API to incorporate linear weights from the Graph Convolution Network paper.

The GAUSConv layer is similar to both the GraphConv and GATConv layers but differs in a few places. Rather than initializing the layer with attention heads, we initialize it with the number of kernels and a kernel dropout probability.


GAUSConv(in_feats, # number of input dimensions
 out_feats, # number of output features
 num_kernels, # number of kernels for current layer
 feat_drop=0., # dropout probability of features
 kernel_drop=0., # dropout probability of kernels
 negative_slope=0.2, # leakly relu slope
 activation=None, # activation function to apply after forward pass
 random_seed=None, # for example / reproducibility purposes
 allow_zero_in_degree=False)

Importantly, in the layer instantiation, we define linear weights and kernel mean and sigma parameters, mu and sigma. We initialize both kernel parameters with the flag require_grad=True, which enables us to update these kernel parameters during the backward pass of the layer. Both parameters are initialized with values in the reset_parameters method.

# initialize feature weights and bias vector
self.weights = nn.Parameter(
 th.Tensor(num_kernels, in_feats, out_feats), requires_grad=True)
self.bias = nn.Parameter(
 th.Tensor(num_kernels, out_feats), requires_grad=True)

# initialize kernel perameters
self.mu = nn.Parameter(
 th.Tensor(1, num_kernels, in_feats), requires_grad=True)
self.sigma = nn.Parameter(
 th.Tensor(num_kernels, 1), requires_grad=True)

Now here is the clever part, and where the DGL message passing interface really shines through. DGL fuses the send and receive messages so that no messages between nodes are ever explicitly stored, using built-in message and reduce functions. To compute the kernel weights between all pairs of source and destination nodes, we use these built-in functions. The important steps are:

compute node feature differences between all source / destination node pairs
aggregate and reduce incoming messages from destination nodes scaled by the kernel weights, to update the source node features

In the forward pass of our layer, we perform the following steps:

### forward pass of GAUSConv layer ###

# compute all pairwise differences between adjacent node features
graph.ndata['h'] = feat
graph.apply_edges(fn.u_sub_v('h', 'h', 'diff'))

# compute kernel weights for each source / desintation pair
e = graph.edata['diff'].unsqueeze(1) - mu
e = -1*sigma*th.norm(e, dim=2).unsqueeze(2)
e = e.exp()
graph.edata['e'] = e

# apply kernel weights to destination node features
graph.apply_edges(fn.v_mul_e('h', 'e', 'kw'))

# apply linear projection to kernel-weighted destination node features
a = th.sum(th.matmul(graph.edata['kw'].transpose(1, 0), weights), dim=0)

# apply kernel dropout
a = self.kernel_drop(a)
graph.edata['a'] = a

# final message-passing and reduction step
# aggregate weighted destination node features to update source node features
graph.update_all(fn.copy_e('a', 'm'), fn.sum('m', 'h'))
rst = graph.ndata['h']

As an example, given a graph and features, we instantiate a GAUSConv layer and propagate our features through the network via:


# set random seed
random_seed=1

# define arbitrary input/output feature shape
n_samples = 4
in_feats=4
out_feats=2
features = th.ones(n_samples, in_feats)

# define number of kernels
num_kernels=2

# create graph structure
u, v = th.tensor([0, 0, 0, 1]), th.tensor([1, 2, 3, 3])
g = dgl.graph((u, v))
g = dgl.to_bidirected(g)
g = dgl.add_self_loop(g)

# instantiate layer
GausConv = GAUSConv(in_feats=in_feats,
 out_feats=out_feats,
 random_seed=random_seed,
 num_kernels=num_kernels,
 feat_drop=0,
 kernel_drop=0)

# forward pass of layer
logits = GausConv(g, features)

print(logits)
tensor([[0.1873, 0.7217],
 [0.1405, 0.5413],
 [0.0936, 0.3608],
 [0.1405, 0.5413]], grad_fn=<AddBackward0>)