#  ABKD: Graph Neural Network Compression with Attention-Based Knowledge
Distillation

Anshul Ahluwalia, Rohit Das, Payman Behnam, Alind Khare, Pan Li, Alexey
Tumanov
Georgia Institute of Technology
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Graph Neural Networks (GNNs) have proven to be quite versatile for a variety
of applications, including recommendation systems, fake news detection, drug
discovery, and even computer vision. Due to the expanding size of graph-
structured data, GNN models have also increased in complexity, leading to
substantial latency issues. This is primarily attributed to the irregular
structure of graph data and its access pattern into memory. The natural
solution to reduce latency is to compress large GNNs into small GNNs. One way
to do this is via knowledge distillation (KD). However, most KD approaches for
GNNs only consider the outputs of the last layers and do not consider the
outputs of the intermediate layers of the GNNs; these layers may contain
important inductive biases indicated by the graph structure. To address this
shortcoming, we propose a novel KD approach to GNN compression that we call
Attention-Based Knowledge Distillation (ABKD). ABKD is a KD approach that uses
attention to identify important intermediate teacher-student layer pairs and
focuses on aligning their outputs. ABKD enables higher compression of GNNs
with a smaller accuracy dropoff compared to existing KD approaches. On
average, we achieve a $1.79\%$ increase in accuracy with a $32.3\times$
compression ratio on OGBN-Mag, a large graph dataset, compared to state-of-
the-art approaches.

##  1 Introduction

Graph Neural Networks (GNNs) generalize Convolutional Neural Networks (CNNs)
to non-Euclidean data. GNNs are widely used in a variety of fields, such as
web-search recommendation systems (Ying et al., 2018), fake news detection for
social networks (Han et al., 2020), modeling proteins for drug discovery
(Zitnik & Leskovec, 2017), and computer vision tasks (Chen et al., 2022). Due
to the expanding size of social networks and other graph-structured data,
graph datasets have been steadily increasing in size (Wang et al., 2020). As
datasets have expanded in size, GNN models have also increased in complexity,
leading to substantial latency issues (Zhou et al., 2021b; Que et al., 2023;
Tan et al., 2023), as shown in figure 1. This is primarily attributed to the
irregular structure of graph data and their access pattern in memory (Liu et
al., 2022b).

Figure 1: Inference latency of GNNs with varying model sizes on the Flickr
Zeng et al. (2020) dataset on a standard GCN model architecture with
increasing embedding dimension. All tests were run on a Tesla V100 GPU.

Due to this limitation, large GNNs need to be compressed into smaller GNNs for
latency-sensitive applications such as real-time recommendation (Liu et al.,
2022a), visual question answering (Senior et al., 2023), image search (Formal
et al., 2020), and real-time spam detection (Li et al., 2019).

###  1.1 Knowledge Distillation

Knowledge Distillation (KD) is a common compression technique that uses a
teacher model to supervise the training of a smaller student model (Hinton et
al., 2015). While the original KD method can be applied to GNNs, it does not
take into account any information about node connectivity. GraphAKD (He et
al., 2022), LSP (Yang et al., 2021), and G-CRD (Joshi et al., 2022) are GNN-
specific knowledge distillation methods that focus on aligning final layer
node embeddings by considering node connectivity. However, these methods all
consider only the node embeddings at the final layer and do not consider
intermediate representations. By just aligning the final node embeddings, the
model cannot learn the logic behind leveraging the connectivity of the graph
or the inductive biases (Baxter, 2000) contained in the adjacency matrix.
Therefore, the student model is effectively just learning a mapping from node
attributes to more refined node embeddings. This can lead to suboptimal test
generalization when the model encounters previously unseen data.

Objective: Our goal is to achieve better generalization on out-of-distribution
data, measured by accuracy @ various compression rates, by considering
intermediate node embeddings and taking into account more of the inductive
biases that GNNs contain.

###  1.2 Attention-Based Knowledge Distillation

To improve accuracy over SOTA, we must consider some of the inductive biases
that other methods overlook - the connectivity of the input graph. We observe
that for the vast majority of GNN architectures, the $k^{th}$ layer of a GNN
computes node embeddings by aggregating information from each node’s $k$-hop
neighborhood. Therefore, each layer contains its own inductive bias. As other
KD methods consider only the node embeddings at the last layer of the teacher
and student networks, they do not leverage all the inductive biases present in
the network. However, one obvious challenge that is present in aligning
intermediate node embeddings is that teacher and student networks will likely
have a different number of hidden layers. As a result, there is no 1-1
correspondence between teacher and student layers and no way to easily figure
out which teacher node embeddings should be aligned with which student node
embeddings.

To tackle these challenges, we propose Attention-Based Knowledge Distillation
(ABKD). We use a trainable attention mechanism to learn which teacher-student
pairs are the most important to align. We also utilize trainable projections
into a common ABKD embedding space for teacher and student hidden layers. By
aligning across intermediate layers, the student learns how to use the
adjacency matrix to construct node embeddings instead of just learning a
mapping from the node attributes to the final layer node embeddings (Figure
2).

###  1.3 Our Contributions

Our contributions are summarized as the following:

  1. 1.

We design Attention-Based Knowledge Distillation (ABKD), a novel knowledge
distillation approach for GNNs that incorporates the intermediate feature maps
from every layer in both the teacher and student networks. This approach can
be utilized to train teacher and student networks of any architectural
configuration.

  2. 2.

We create an automatic feature linking mechanism using attention to identify
which teacher-student layer pairs are the most important, which we then use to
closely align their feature maps.

  3. 3.

Our approach broadly improves the test accuracy of student networks over a
large range of compression ratios.

  4. 4.

We comprehensively test our approach on several datasets using different model
architectures including GCNs (Kipf & Welling, 2016), RGCNs (Schlichtkrull et
al., 2017) and GraphSAGE (Hamilton et al., 2017). We also test on several
large datasets that are carefully curated to evaluate out-of-distribution
generalization such as OGBN-Mag (Wang et al., 2020) and OGBN-Arxiv (Wang et
al., 2020).

##  2 Related Work

Knowledge Distillation KD for GNNs is a relatively niche field that has
expanded over the last three years with the work of LSP (Yang et al., 2021).
In this work, the authors attempt to align node embeddings between the student
and teacher networks by maximizing the similarity between embeddings that
share edges. As only node embeddings between edges are aligned, this KD method
only preserves local topology. Joshi et al. (2022) extend LSP and propose two
different KD algorithms: Global Structure Preserving Distillation (GSP) and
Global Contrastive Representation Distillation (G-CRD). GSP extends LSP by
considering all pairwise similarities among node features, not just pairwise
similarities between nodes connected by edges. The authors also propose G-CRD
which aims to implicitly preserve global topology by aligning the student and
teacher node feature vectors via contrastive learning (van den Oord et al.,
2018). Another work introduces graph adversarial knowledge distillation
(GraphAKD), which trains the student model as a generator network and a
discriminator network to distinguish between the predictions of the student
and teacher models (He et al., 2022). Another work, GeometricKD forces the
number of teacher and student hidden layers to be the same to study the impact
of a student network operating on a smaller graph than the teacher (Yang et
al., 2023).

With the exception of GeometricKD, these works all consider the node
embeddings at the final layer of the teacher and student network and aim to
align those embeddings with one another in various ways. GeometricKD
constrains the student and teacher networks to have the same number of layers,
and it aligns the node embedding of teacher layer $i$ with student layer $i$,
thus forming a 1-1 correspondence between the layers, which makes it
inflexible to all student and teacher configurations. There have been several
KD approaches that have been applied to CNNs that the GNN community has tried
to adapt to GNNs with poor results, namely Fitnets (Romero et al., 2015) and
attention transfer (Zagoruyko & Komodakis, 2016). These methods both compute a
distance metric such as mean-squared error between the last layer node
embeddings of the student and teacher network and do not take into account the
adjacency matrix.

Using attention to find similarities across student and teacher layers is a
concept explored in CNNs Ji et al. (2021). However, this work’s ideas cannot
be applied to GNNs because the operations it uses to compare student and
teacher features do not apply to GNNs. GNNs need special consideration in this
regard over CNNs due to the non-spatial and unstructured form of graph data.

Compression: Other common techniques for GNN compression include quantization
and pruning. Quantization techniques differ for GNNs when compared to other
deep neural networks because of unique sources of error, such as inaccurate
gradient estimates, that arise due to the unstructured nature of GNNs (Tailor
et al., 2021). To account for inaccurate gradient estimates, DegreeQuant,
protects some number of nodes in every training iteration from lower precision
training and utilizes full precision training for those nodes (Tailor et al.,
2021). Other approaches consider the structure of the graph in their
calculations to implicitly avoid GNN-specific sources of error (Feng et al.,
2020; Bahri et al., 2021).

Pruning techniques, which involve compressing networks by selectively deleting
learned parameters, have also been applied to GNNs. These techniques involve
applying strategies that have worked for other deep networks and using them to
identify the most important channels for output activation (Chen et al.,
2021b) (Zhou et al., 2021a). One approach also works to dynamically prune
during execution (Chen et al., 2021a).

Method | Number of Layers Considered
---|---
GraphAKD | $1$
G-CRD | $1$
Fitnets | $1$
LSP | $1$
ABKD | All
Table 1: Comparison of Attention-Based Knowledge Distillation with other
Knowledge Distillation approaches.

##  3 Proposed Approach

###  3.1 Intuition and Mathematical Foundations

In this section, we first discuss the intuition behind ABKD and introduce some
of the mathematical definitions needed to explain it thoroughly.

####  3.1.1 SoftKD Intuition

In SoftKD (Hinton et al., 2015), we compute two different losses. The first,
$H(s_{p},y)$, is a cross-entropy loss between the output student probability
distribution and the ground truth labels. The other, $H(s_{p},t_{p})$, is a
cross-entropy loss between the output student probability distribution and the
output teacher probability distribution. The total loss is defined as:

$L_{KD}=H(s_{p},y)+\alpha H(s_{p},t_{p})$ (1)

Here, $\alpha$ is a hyper-parameter controlling how much the KD loss affects
the total loss. The goal is to align the output student probability
distribution with the output teacher probability distribution. The higher
$H(s_{p},t_{p})$ is, the less aligned the student and teacher output
probability distributions are.

####  3.1.2 ABKD Intuition

Similarly, with ABKD, we want to incorporate this intuition of alignment.
However, we want to go further than just aligning the final output - we want
to align the outputs at the intermediate layers, as the intermediate layers
contain inductive biases that are not present in just the final layers. As one
of our goals with ABKD is to work with any combination of teacher and student
architectural configurations, this presents one significant challenge: Teacher
and student networks will likely have a different number of hidden layers,
which means there is no 1-1 correspondence between teacher and student layers.

ABKD solves this problem by identifying which teacher-student layer pairs are
the most important to align via an attention mechanism. This mechanism works
with any arbitrary number of teacher layers and student layers, which makes
this approach amenable to any arbitrary teacher-student configuration. ABKD
also uses a reprojection technique to account for the student and teacher
networks having different hidden dimensions. The output of each hidden layer
for both the teacher and student networks is projected into a standardized
embedding dimension, which ensures that we can work with student and teacher
networks of any embedding dimension. As each layer represents its own semantic
information, an important challenge that we faced was to ensure that each
layer’s feature map was not smoothed out by a single projection matrix. To
this end, we use separate trainable linear layers for each hidden layer in
both the teacher and student networks, in order to ensure that we don’t lose
out on any valuable semantic information in the hidden layers.

These trainable linear layers help us construct the two key components of
ABKD, which are the attention map and the dissimilarity map. At a high level,
the attention map tells us how important each teacher-student layer pair is,
while the dissimilarity map tells us how distant the feature maps of each
teacher-student layer pair are. The teacher-student layer pairs with higher
attention scores are deemed as more important, and ABKD focuses on reducing
their dissimilarity scores during training.

Figure 2: ABKD generates an attention map using a trainable attention
mechanism and a dissimilarity map using a trainable subspace projection. The
loss matrix is an element-wise multiplication of the attention matrix and the
dissimilarity matrix.

####  3.1.3 Mathematical Foundation

Without loss of generality, we will consider distilling a general Graph
Convolution Network (GCN) (Kipf & Welling, 2016), in which the output of the
$l^{th}$ layer is

$H^{l}=\sigma(\hat{A}H_{l-1}W_{l})$ (2)

Here, $\sigma$ is an activation function and $\hat{A}$ is the normalized
adjacency matrix, $\hat{A}=D^{-\frac{1}{2}}AD^{-\frac{1}{2}}$, where $D$ is
the diagonal degree matrix. $H^{l-1}\in R^{n\times d}$ represents the output
of the last hidden layer and $W^{l}\in R^{d\times d}$ represents the trainable
weights of the current layer.

Consider two different tensors $T\in R^{T_{l}\times n\times d_{t}}$ and $S\in
R^{S_{l}\times n\times d_{s}}$. $T_{l}$ and $S_{l}$ represent the number of
layers in the teacher and student networks respectively. $n$ represents the
number of nodes in the graph that is being trained on and $d_{t}$ and $d_{s}$
represents the dimensionality of the teacher and student networks,
respectively. $T$ and $S$ represent the calculated layer maps before
activation is applied at every layer for the teacher and student networks,
respectively.

###  3.2 ABKD

####  3.2.1 Attention Scores

The first step of ABKD is to generate $A\in R^{T_{l}\times S_{l}}$. $A_{ij}$
will represent an "importance" score for the layer pair consisting of teacher
layer $i$ and student layer $j$. We take the average of the feature maps along
the node dimension to compute a mean node feature for every layer in both the
teacher and student networks. Call these tensors $T_{a}\in R^{T_{l}\times
d_{t}}$ and $S_{a}\in R^{S_{l}\times d_{s}}$. Then, we pass each layer in
$T_{a}$ through its own linear layer to create $T_{p}\in R^{T_{l}\times
d_{a}}$, where $d_{a}$ is the embedding dimension of ABKD. Similarly, we
create $S_{p}\in R^{S_{l}\times d_{a}}$. We can finally generate $A$ in the
following manner:

$A=softmax(\frac{T_{p}S_{p}^{T}}{\sqrt{d_{a}}})$ (3)

####  3.2.2 Dissimilarity Scores

The next step is to compute a pairwise dissimilarity score for each teacher-
student layer pair. Again, we project the features into $d_{a}$. For
calculating the attention scores, we averaged over the node dimension before
projecting, as our goal was to identify important layers. When calculating the
pairwise dissimilarity, we want to incorporate the per-node embeddings. So, we
use a separate set of projection matrices. We use $P_{t}\in R^{d_{t}\times
d_{a}}$ and $P_{s}\in R^{d_{s}\times d_{a}}$ to represent the projections.

However, distance metrics are less semantically valuable if $d_{a}$ is high.
To alleviate this problem, we define a trainable matrix $P\in R^{d_{a}\times
d_{a}}$ to project all vectors into the subspace defined by the column space
of $P$. Since the cardinality of the subspace defined by the column space of
$P$ will be smaller than or equal to the cardinality of $R^{d_{a}}$, distance
metrics within the subspace will be more valuable on average compared to
distance metrics in $R^{d_{a}}$.

The final step is to average over the embedding dimension and then produce
$D\in R^{T_{l}\times S_{l}}$, which gives dissimilarity scores for each
teacher-student layer pair. For calculating the dissimilarity, we experiment
with Euclidean and cosine distance, but Euclidean distance generally tends to
perform better. The dissimilarity score for a layer pair $(i,j)$ can be
represented as:

$D_{ij}=||(T_{i}P_{t}-S_{j}P_{s})P\frac{\mathds{1}_{d_{a}}}{d_{a}}||^{2}_{2}$
(4)

where $\mathds{1}_{d_{a}}$ is a vector of $1$s in $R^{d_{a}}$.

####  3.2.3 Final Loss Calculation

To produce the final loss matrix, we element-wise multiply $A$ and $D$ and
then take the row-wise mean to produce a single number that represents the
ABKD loss.

$L_{abkd}=(\mathds{1}_{T_{l}})^{T}(A\odot
D)(\frac{\mathds{1}_{S_{l}}}{S_{l}})$ (5)

The final loss is calculated as

$L=H(s_{p},y)+\beta L_{abkd}$ (6)

There is one important theorem to consider that proves $L_{abkd}$ distills
valuable knowledge from the teacher network to the student network.

Theorem 1: Consider a teacher layer $i$ and a student layer $j$ such that
$j<i$. Consider the weight $W_{j}^{s}$ which is the associated weight for the
$j^{th}$ layer of the student network. The weight update for $W_{j}^{s}$ will
be affected by $W_{i}^{t}$, which are the weights for the $i^{th}$ layer of
the teacher network.

Proof: To calculate the weight update, we have to first formulate a loss for
the layer pair $(i,j)$. By adapting equation 5 to a specific element in the
loss matrix, we get:

$\displaystyle L_{ij}$
$\displaystyle\sim(\frac{\mathds{1}_{n}}{n})^{T}(T^{p}_{i}(W_{j}^{ps})^{T})(W_{j}^{s})^{T}((H^{s}_{j-1})^{T}\hat{A}^{T}))\frac{\mathds{1}_{n}}{n}*D_{ij}$
$\displaystyle T^{p}_{i}$
$\displaystyle=\hat{A}H_{i-1}^{t}W^{t}_{i}W_{i}^{pt}$

where $W_{i}^{pt}$ and $W_{j}^{ps}$ represent the projection matrices used to
calculate the attention score.

Assuming that our weight updates are performed by gradient descent, to
calculate the weight update for $W_{j}^{s}$ due to $L_{ij}$, we have to first
calculate $\frac{\partial L_{ij}}{\partial W_{j}^{s}}\in R^{d_{s}\times
d_{s}}$:

$\frac{\partial L_{ij}}{\partial W_{j}^{s}}=\frac{\partial
A_{ij}}{W_{j}^{s}}D_{ij}+\frac{\partial D_{ij}}{W_{j}^{s}}A_{ij}$ (7)

We focus on the first term and obtain the result:

$\frac{\partial A_{ij}}{\partial
W_{j}^{s}}\sim((H_{j-1}^{s})^{T}\hat{A}^{T}(\frac{\mathds{1}_{n}}{n}))((\frac{\mathds{1}_{n}}{n})^{T}T_{i}^{p}(W_{j}^{ps})^{T})$
(8)

It is apparent that as $\frac{\partial A_{ij}}{\partial W_{j}^{s}}$ contains
the term $T^{p}_{i}=\hat{A}H_{i-1}^{t}W^{t}_{i}W_{i}^{pt}$, the weight update
will reflect terms from the $i^{th}$ teacher layer. This theorem goes to show
that even though student layer $j$ is responsible for aggregating information
for every node from its $j$-hop neighborhood, its weights collect knowledge
from teacher layer $i$ \- knowledge that would be unavailable to student layer
$j$ using any of the other GNN KD approaches.

An important follow-up observation is that the loss for a teacher-student
layer pair will be higher if the pair is deemed as more important and if their
projected feature maps have a high dissimilarity score ($D_{ij}>0.5)$, which
indicates that they aren’t closely aligned.

To verify this observation, consider an arbitrary teacher layer $i$ and an
arbitrary student layer $j$. Consider $A_{ij}$. If the pair $(i,j)$ is deemed
important, $A_{ij}$ will be high. Now consider $D_{ij}$. If the pair $(i,j)$
is not closely aligned, then $D_{ij}$ will be high. As $L_{ij}=A_{ij}*D_{ij}$,
if both $A_{ij}$ and $D_{ij}$ are high, then $L_{ij}$ will also be high, hence
proving that the loss is high for important but misaligned layer pairs.

##  4 Experiments

Figure 3: Accuracy vs. Compression Ratio: A comparison of KD methods applied
to student models of different sizes trained on the OGBN-Mag dataset. The
teacher model was the same model as the one described in table 3. The student
model was a two-layer RGCN and we varied the embedding dimension from 16 - 512
to induce this Pareto frontier. The y-axis represents the test accuracy of the
final trained models.

Dataset | OGBN-Mag | OGBN-Arxiv
---|---|---
Teacher | RGCN (3L-512H-5.5M) | GAT (3L-750H-1.4M)
Student | RGCN (2L-32H-170K) | GS (2L-256H-87K)
Teacher | $49.80$ | $74.20$
Student | $44.23\pm 0.47$ | $70.87\pm 0.58$
Fitnet | $44.87\pm 0.84$ | $71.32\pm 0.32$
AT | $43.87\pm 0.67$ | $71.04\pm 0.48$
LSP | $45.21\pm 0.54$ | $71.47\pm 0.45$
GSP | $44.97\pm 0.58$ | $71.97\pm 0.64$
G-CRD | $45.42\pm 0.43$ | $71.87\pm 0.56$
ABKD | $\mathbf{47.21\pm 0.32}$ | $\mathbf{73.25\pm 0.56}$
Ratio | $32.3\times$ | $16.1\times$

Table 2: Average accuracies over five trials for the OGBN-Mag and OGBN-Arxiv
Datasets. The teacher network was kept constant while the different
distillation methods were applied to student networks that were initialized
from the same set of weights. GS means GraphSAGE.

###  4.1 Experimental Setup

For our main experiments, we test ABKD on two difficult datasets OGBN-Mag and
OGBN-Arxiv (Wang et al., 2020). These datasets employ temporal splitting to
create validation and test sets that can assess a model’s ability to
generalize on out-of-distribution data (Hu et al., 2020). For OGBN-Mag, we run
experiments using RGCN (Schlichtkrull et al., 2017) as the teacher and student
models, and for OGBN-Arxiv, we run experiments using GAT (Veličković et al.,
2018) as the teacher model and GraphSAGE (Hamilton et al., 2017) as the
student model. This allows us to evaluate the effectiveness of ABKD for
different GNN architectures. It also allows us to assess if ABKD can distill
information between different types of GNN architectures. In our experiments,
we keep the teacher model architecture and weights fixed and only modify the
size of the student network. Each distillation method starts from the same set
of weights and trains for the same number of epochs across 5 runs. For our
baselines, we consider LSP (Yang et al., 2021), GSP (Joshi et al., 2022),
G-CRD (Joshi et al., 2022), Fitnet (Romero et al., 2015), and Attention
Transfer (Zagoruyko & Komodakis, 2016). We ran all experiments on a Tesla V100
GPU.

###  4.2 Experimental Results

Out-Of-Distribution Evaluation As evidenced by our results in Table 2, we are
able to outperform SOTA in several different compression settings on OGBN-Mag
and OGBN-Arxiv. We show improvements of over $1.79\%$ and $1.38\%$ with
compression ratios of $32.3\times$ and $16.1\times$ on OGBN-Mag and OGBN-Arxiv
respectively. As these datasets are intended to evaluate out-of-distribution
generalization, our results empirically prove that ABKD can train student
models that generalize better than other KD approaches.

Figure 4: Attention and Dissimilarity maps before and after training with
ABKD. Cooler colors refer to lower scores and warmer colors correspond to
higher scores.

ABKD aligns intermediate embeddings: To empirically prove that ABKD actually
aligns intermediate embeddings based on the attention matrix, we visualize the
before and after training attention and dissimilarity maps in figure 4. We
train on OGBN-Mag and use a deeper teacher network of 5 layers and a hidden
dimension of 512. The student network has 3 layers and a hidden dimension of
32. Our results show that dissimilarity scores are low where the attention
scores are high and vice versa. This is in line with the intuition that we
constructed in 3.2.3.

Dataset | Cora | Citeseer | Pubmed | NELL
---|---|---|---|---
Teacher | GCNII (64L-64H) | GCNII (64L-64H) | GCNII (64L-64H) | GCNII (64L-64H)
Student | GCNII (4L-4H) | GCNII (4L-4H) | GCNII (4L-4H) | GCNII (4L-4H)
Teacher | $88.40$ | $77.33$ | $89.78$ | $95.55$
Student | $74.22\pm 1.32$ | $68.14\pm 1.45$ | $88.66\pm 0.87$ | $85.20\pm 0.75$
LSP | $75.18\pm 1.45$ | $70.43\pm 1.32$ | $89.01\pm 1.73$ | $85.20\pm 1.32$
GSP | $78.32\pm 1.21$ | $69.43\pm 0.87$ | $89.13\pm 1.05$ | $86.13\pm 1.47$
G-CRD | $83.35\pm 1.32$ | $71.42\pm 1.21$ | $89.73\pm 1.05$ | $88.61\pm 1.01$
ABKD | $\mathbf{84.27\pm 1.48}$ | $\mathbf{71.96\pm 0.87}$ | $\mathbf{89.87\pm 0.45}$ | $\mathbf{91.83\pm 0.74}$
# Parameters | $5835$ | $14910$ | $2083$ | $22686$
Ratio | $60.7\times$ | $33.5\times$ | $141.3\times$ | $27.7\times$

Table 3: Average accuracies over several trials for a variety of datasets,
using GCNII. The results are based on the average of five trials, with each
distillation method applied to the same set of student weights.

Deep GNNs: We also test ABKD on deep GNN architectures. While most GNN
architectures are shallow due to the problem of over-smoothing, there are some
approaches that alleviate this issue, allowing for deep GNN architectures. One
of these approaches is GCNII (Chen et al., 2020). We test on Cora (Mccallum et
al., 2000), Citeseer (Sen et al., 2008), Pubmed (Namata et al., 2012) and NELL
(Carlson et al., 2010). Table 3 shows that ABKD is able to distill these deep
GCNIIs into shallower GCNIIs with higher accuracy compared to other
distillation methods in high compression settings. While currently, most GNN
use cases require shallow GNNs, this could be useful for future applications
that require deeper GNNs.

Improved Weight Initialization for Highly Compressed Networks: We find that
for smaller datasets, information from the teacher network is mainly distilled
into one layer of the student network, as shown in Figure 5. Our hypothesis
for this occurrence is that smaller datasets are not very complex and one
layer is sufficient for learning most of the patterns. Through
experimentation, we prove that when we initialize a one-layer network from the
weights of this particular layer, we get improved accuracy compared to random
initialization, as shown in Table 4.

Figure 5: Attention maps for Cora, Citeseer, and Pubmed. Each color in the
heatmap represents the importance score associated with that teacher-student
layer pair. Warmer colors mean higher important scores. It seems apparent that
most of the knowledge from the teacher layers is distilled into one student
layer.

Dataset | Initialized | Random Init | No Training
---|---|---|---
Cora | $80.29$ | $73.58$ | $65.36$
Citeseer | $70.12$ | $68.12$ | $54.20$
Pubmed | $88.76$ | $86.03$ | $72.90$

Table 4: Results for weight initialization experiment. Experimental details
can be found in the appendix.

Dataset | OGBN-Mag | OGBN-Arxiv
---|---|---
Student | $44.46\pm 0.54$ | $71.27\pm 0.48$
ABKD | $\mathbf{47.46\pm 0.45}$ | $\mathbf{73.18\pm 0.56}$
Modified ABKD | $46.56\pm 0.60$ | $71.89\pm 0.87$

Table 5: Ablation study comparing modified ABKD that only considers aligning
the last layer node embeddings with ABKD that considers intermediate layer
node embeddings. The teacher and student models are the same as the ones in
Table 2.

###  4.3 Ablation Studies

Aligning Intermediate Layers is Important: To prove that the aligning of
intermediate layers is necessary for superior performance, we experiment with
a variant of ABKD, which we call modified ABKD, where we set $A\in
R^{T_{l}\times S_{l}}$ to all zeros, but we set the bottom right value to 1.
This indicates that we are only interested in the dissimilarity between the
last layer node embeddings of the teacher and student models. The results in
table 5, prove that we gain accuracy by considering the outputs of
intermediate layers for both teacher and student models. In this experiment,
we start from the same set of initialized weights for both the ABKD and
modified ABKD approaches.

Improvements due to Subspace Projection: In our approach, we introduced the
concept of subspace projection as an alleviation to high dimensional embedding
spaces. While it is not needed for ABKD to work, it does improve our results
as the learned subspace projection matrix tends to be of lower rank than the
embedding dimension. This indicates that we can project our feature maps into
subspaces smaller than $R^{d_{a}}$, which increases the semantic value of the
dissimilarity scores.

One Linear Layer Per Hidden Layer Is Necessary: In our approach, we mention
that each hidden teacher and student layer is assigned its own linear layer
for projection into $d_{a}$. This is because each layer represents its own
$k$-hop neighborhood, and using just one linear layer would prove inadequate
in capturing the full spectrum of essential semantic information contained
within each layer. We run an experiment in which we use only one linear layer
for the teacher and student projections and demonstrate that there is a
significant accuracy dropoff compared to using individual linear layers for
the projection.

Dataset | Subspace Projection | No Projection
---|---|---
Cora | $84.27\pm 1.32$ | $83.78\pm 1.25$
Citeseer | $71.96\pm 0.83$ | $71.46\pm 1.21$
Pubmed | $89.87\pm 0.87$ | $89.03\pm 1.04$
NELL | $91.83\pm 0.85$ | $90.98\pm 0.75$
OGBN-Mag | $47.21\pm 0.32$ | $46.83\pm 0.45$
OGBN-Arxiv | $73.25\pm 0.56$ | $72.87\pm 0.34$

Table 6: Subspace Projection Ablation Results. The teacher and student
networks are the same as the ones described in tables 2 and 3.

Dataset | Multiple Linear Layers | One Linear Layer
---|---|---
Cora | $84.27\pm 1.48$ | $75.52\pm 1.32$
Citeseer | $71.96\pm 0.87$ | $68.86\pm 1.27$
Pubmed | $89.87\pm 0.45$ | $88.75\pm 0.64$
NELL | $91.83\pm 0.74$ | $85.38\pm 0.68$
OGBN-Mag | $47.21\pm 0.32$ | $44.67\pm 0.54$
OGBN-Arxiv | $73.25\pm 0.56$ | $71.23\pm 0.47$

Table 7: Linear Layer Ablation Results. The teacher and student networks are
the same as the ones described in tables 2 and 3.

##  5 Conclusion

Graph Neural Networks (GNNs) have increased in complexity in recent years,
owing to the rapidly increasing size of graph datasets. This increase in
complexity poses a problem for latency-sensitive applications such as real-
time recommendation and real-time spam detection, as GNN model sizes do not
scale well with latency. To address this predicament, we propose an innovative
solution known as Attention-Based Knowledge Distillation (ABKD). ABKD employs
an attention-based feature linking mechanism to identify important
intermediate teacher-student layer pairs and focuses on aligning the node
embeddings of those pairs. This knowledge distillation approach broadly
outperforms SOTA GNN-specific KD approaches over a wide variety of compression
settings. It also works with both deep and shallow networks as proven by our
experiments and performs well with several different types of GNN
architectures. On average, we achieve a $1.79\%$ increase in accuracy with a
$32.3\times$ compression ratio on OGBN-Mag, a large graph dataset, compared to
state-of-the-art approaches.

## References

  * Bahri et al. (2021) Mehdi Bahri, Gaétan Bahl, and Stefanos Zafeiriou.  Binary graph neural networks, 2021.
  * Baxter (2000) J. Baxter.  A model of inductive bias learning.  _Journal of Artificial Intelligence Research_ , 12:149–198, mar 2000.  doi: 10.1613/jair.731.  URL https://doi.org/10.1613%2Fjair.731.
  * Carlson et al. (2010) Andrew Carlson, Justin Betteridge, Bryan Kisiel, Burr Settles, Estevam Hruschka, and Tom Mitchell.  Toward an architecture for never-ending language learning.  _Proceedings of the AAAI Conference on Artificial Intelligence_ , 24(1):1306–1313, Jul. 2010.  doi: 10.1609/aaai.v24i1.7519.  URL https://ojs.aaai.org/index.php/AAAI/article/view/7519.
  * Chen et al. (2021a) Cen Chen, Kenli Li, Xiaofeng Zou, and Yangfan Li.  Dygnn: Algorithm and architecture support of dynamic pruning for graph neural networks.  In _2021 58th ACM/IEEE Design Automation Conference (DAC)_ , pp. 1201–1206, 2021a.  doi: 10.1109/DAC18074.2021.9586298.
  * Chen et al. (2022) Chaoqi Chen, Yushuang Wu, Qiyuan Dai, Hong-Yu Zhou, Mutian Xu, Sibei Yang, Xiaoguang Han, and Yizhou Yu.  A survey on graph neural networks and graph transformers in computer vision: A task-oriented perspective, 2022.
  * Chen et al. (2020) Ming Chen, Zhewei Wei, Zengfeng Huang, Bolin Ding, and Yaliang Li.  Simple and deep graph convolutional networks.  _CoRR_ , abs/2007.02133, 2020.  URL https://arxiv.org/abs/2007.02133.
  * Chen et al. (2021b) Tianlong Chen, Yongduo Sui, Xuxi Chen, Aston Zhang, and Zhangyang Wang.  A unified lottery ticket hypothesis for graph neural networks, 2021b.
  * Feng et al. (2020) Boyuan Feng, Yuke Wang, Xu Li, Shu Yang, Xueqiao Peng, and Yufei Ding.  Sgquant: Squeezing the last bit on graph neural networks with specialized quantization, 2020.
  * Formal et al. (2020) Thibault Formal, Stéphane Clinchant, Jean-Michel Renders, Sooyeol Lee, and Geun Hee Cho.  Learning to rank images with cross-modal graph convolutions.  In Joemon M. Jose, Emine Yilmaz, João Magalhães, Pablo Castells, Nicola Ferro, Mário J. Silva, and Flávio Martins (eds.), _Advances in Information Retrieval_ , pp. 589–604, Cham, 2020. Springer International Publishing.  ISBN 978-3-030-45439-5.
  * Hamilton et al. (2017) William L. Hamilton, Rex Ying, and Jure Leskovec.  Inductive representation learning on large graphs.  _CoRR_ , abs/1706.02216, 2017.  URL http://arxiv.org/abs/1706.02216.
  * Han et al. (2020) Yi Han, Shanika Karunasekera, and Christopher Leckie.  Graph neural networks with continual learning for fake news detection from social media, 2020.
  * He et al. (2022) Huarui He, Jie Wang, Zhanqiu Zhang, and Feng Wu.  Compressing deep graph neural networks via adversarial knowledge distillation, 2022.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network, 2015.
  * Hu et al. (2020) Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec.  Open graph benchmark: Datasets for machine learning on graphs.  _CoRR_ , abs/2005.00687, 2020.  URL https://arxiv.org/abs/2005.00687.
  * Ji et al. (2021) Mingi Ji, Byeongho Heo, and Sungrae Park.  Show, attend and distill:knowledge distillation via attention-based feature matching, 2021.
  * Joshi et al. (2022) Chaitanya K. Joshi, Fayao Liu, Xu Xun, Jie Lin, and Chuan Sheng Foo.  On representation knowledge distillation for graph neural networks.  _IEEE Transactions on Neural Networks and Learning Systems_ , pp. 1–12, 2022.  doi: 10.1109/tnnls.2022.3223018.  URL https://doi.org/10.1109%2Ftnnls.2022.3223018.
  * Kipf & Welling (2016) Thomas N. Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  _CoRR_ , abs/1609.02907, 2016.  URL http://arxiv.org/abs/1609.02907.
  * Li et al. (2019) Ao Li, Zhou Qin, Runshi Liu, Yiqun Yang, and Dong Li.  Spam review detection with graph convolutional networks.  In _Proceedings of the 28th ACM International Conference on Information and Knowledge Management_ , CIKM ’19, pp. 2703–2711, New York, NY, USA, 2019. Association for Computing Machinery.  ISBN 9781450369763.  doi: 10.1145/3357384.3357820.  URL https://doi.org/10.1145/3357384.3357820.
  * Liu et al. (2022a) Zhuoran Liu, Leqi Zou, Xuan Zou, Caihua Wang, Biao Zhang, Da Tang, Bolin Zhu, Yijie Zhu, Peng Wu, Ke Wang, and Youlong Cheng.  Monolith: Real time recommendation system with collisionless embedding table, 2022a.
  * Liu et al. (2022b) Zirui Liu, Kaixiong Zhou, Fan Yang, Li Li, Rui Chen, and Xia Hu.  EXACT: Scalable graph neural networks training via extreme activation compression.  In _International Conference on Learning Representations_ , 2022b.  URL https://openreview.net/forum?id=vkaMaq95_rX.
  * Mccallum et al. (2000) Andrew Mccallum, Kamal Nigam, and Jason Rennie.  Automating the construction of internet portals.  03 2000.
  * Namata et al. (2012) Galileo Namata, Ben London, Lise Getoor, and Bert Huang.  Query-driven active surveying for collective classification.  2012\.
  * Que et al. (2023) Zhiqiang Que, Hongxiang Fan, Marcus Loo, He Li, Michaela Blott, Maurizio Pierini, Alexander D. Tapper, and Wayne Luk.  Ll-gnn: Low latency graph neural networks on fpgas for high energy physics, 2023.
  * Romero et al. (2015) Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets, 2015.
  * Schlichtkrull et al. (2017) Michael Schlichtkrull, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling.  Modeling relational data with graph convolutional networks, 2017.
  * Sen et al. (2008) Prithviraj Sen, Galileo Namata, Mustafa Bilgic, Lise Getoor, Brian Galligher, and Tina Eliassi-Rad.  Collective classification in network data.  _AI Magazine_ , 29(3):93, Sep. 2008.  doi: 10.1609/aimag.v29i3.2157.  URL https://ojs.aaai.org/index.php/aimagazine/article/view/2157.
  * Senior et al. (2023) Henry Senior, Gregory Slabaugh, Shanxin Yuan, and Luca Rossi.  Graph neural networks in vision-language image understanding: A survey, 2023.
  * Tailor et al. (2021) Shyam A. Tailor, Javier Fernandez-Marques, and Nicholas D. Lane.  Degree-quant: Quantization-aware training for graph neural networks, 2021.
  * Tan et al. (2023) Zeyuan Tan, Xiulong Yuan, Congjie He, Man-Kit Sit, Guo Li, Xiaoze Liu, Baole Ai, Kai Zeng, Peter Pietzuch, and Luo Mai.  Quiver: Supporting gpus for low-latency, high-throughput gnn serving with workload awareness, 2023.
  * van den Oord et al. (2018) Aäron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  _CoRR_ , abs/1807.03748, 2018.  URL http://arxiv.org/abs/1807.03748.
  * Veličković et al. (2018) Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio.  Graph attention networks, 2018.
  * Wang et al. (2020) Kuansan Wang, Zhihong Shen, Chiyuan Huang, Chieh-Han Wu, Yuxiao Dong, and Anshul Kanakia.  Microsoft Academic Graph: When experts are not enough.  _Quantitative Science Studies_ , 1(1):396–413, 02 2020.  ISSN 2641-3337.  doi: 10.1162/qss_a_00021.  URL https://doi.org/10.1162/qss_a_00021.
  * Yang et al. (2023) Chenxiao Yang, Qitian Wu, and Junchi Yan.  Geometric knowledge distillation: Topology compression for graph neural networks, 2023.
  * Yang et al. (2021) Yiding Yang, Jiayan Qiu, Mingli Song, Dacheng Tao, and Xinchao Wang.  Distilling knowledge from graph convolutional networks, 2021.
  * Ying et al. (2018) Rex Ying, Ruining He, Kaifeng Chen, Pong Eksombatchai, William L. Hamilton, and Jure Leskovec.  Graph convolutional neural networks for web-scale recommender systems.  _CoRR_ , abs/1806.01973, 2018.  URL http://arxiv.org/abs/1806.01973.
  * Zagoruyko & Komodakis (2016) Sergey Zagoruyko and Nikos Komodakis.  Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer.  _CoRR_ , abs/1612.03928, 2016.  URL http://arxiv.org/abs/1612.03928.
  * Zeng et al. (2020) Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna.  Graphsaint: Graph sampling based inductive learning method, 2020.
  * Zhou et al. (2021a) Hongkuan Zhou, Ajitesh Srivastava, Hanqing Zeng, Rajgopal Kannan, and Viktor Prasanna.  Accelerating large scale real-time GNN inference using channel pruning.  _Proceedings of the VLDB Endowment_ , 14(9):1597–1605, may 2021a.  doi: 10.14778/3461535.3461547.  URL https://doi.org/10.14778%2F3461535.3461547.
  * Zhou et al. (2021b) Hongkuan Zhou, Ajitesh Srivastava, Hanqing Zeng, Rajgopal Kannan, and Viktor Prasanna.  Accelerating large scale real-time gnn inference using channel pruning.  _arXiv preprint arXiv:2105.04528_ , 2021b.
  * Zitnik & Leskovec (2017) Marinka Zitnik and Jure Leskovec.  Predicting multicellular function through multi-layer tissue networks.  _CoRR_ , abs/1707.04638, 2017.  URL http://arxiv.org/abs/1707.04638.

## Supplementary Materials

##  Appendix A Summary of Notations

In Section 3, we mathematically describe how ABKD generates the attention
matrix, $A$, and the dissimilarity matrix, $D$, which is then used to
calculate $L_{\texttt{ABKD}}$. In Table 8, we provide a summary of all the
mathematical notation used to describe ABKD.

Term | Definition
---|---
$A_{ij}$ | Attention score for the teacher-student layer pair $(i,j)$
$D_{ij}$ | Dissimilarity score for the teacher-student layer pair $(i,j)$
$\mathds{1}_{n}\in R^{n}$ | Vector of ones in $R^{n}$
$T_{i}\in R^{n\times d_{t}}$ | The output of the $i^{th}$ layer in the teacher network.
$S_{j}\in R^{n\times d_{s}}$ | The output of the $j^{th}$ layer in the student network
$T_{i}^{p}\in R^{d_{a}}$ | The projected output of the $i^{th}$ student layer after taking the node-wise mean
$W_{i}^{pt}\in R^{d_{t}\times d_{a}}$ | The learnable $i^{th}$ projection matrix for the teacher network
$W_{j}^{ps}\in R^{d_{s}\times d_{a}}$ | The learnable $j^{th}$ projection matrix for the student network
$P_{t}\in R^{d_{t}\times d_{a}}$ | The learnable projection matrix for the teacher network; used to calculate the dissimilarity score
$P_{s}\in R^{d_{s}\times d_{a}}$ | The learnable projection matrix for the student network; used to calculate the dissimilarity score
$P\in R^{d_{a}\times d_{a}}$ | The learnable subspace projection matrix. Shared between the student and teacher networks

Table 8: Summary of mathematical notation used in Section 3.

##  Appendix B Further Experimental Discussion and Ablations

###  B.1 Dataset and Teacher Network Information

Tables 6 and 7 provide information on the datasets used in Section 4 and the
corresponding teacher networks employed to supervise the training of various
student networks during our experiments.

| # of Nodes | # of Edges | # of Features | # of Classes
---|---|---|---|---
Cora | 2,708 | 10,556 | 1,433 | 7
Citeseer | 3,327 | 9,104 | 3,703 | 6
Pubmed | 19,717 | 88,648 | 500 | 3
OGBN-Mag | 1,939,743 | 21,111,007 | 128 | 349
OGBN-Arxiv | 169,343 | 1,166,243 | 128 | 40
Table 9: Specification of evaluated datasets.

###  B.2 Weight Initialization Experiment Details

In Section 4.2, we present the results of an experiment in which we
instantiate a one-layer network from the attention maps generated by ABKD. The
first step of this experiment is to run ABKD on a student network of any
arbitrary size and then generate the attention map, $A\in R^{T_{l}\times
S_{l}}$. The next step is to use a row-wise argmax and find the student layer
that has the most information distilled down to it. For example, in Figure 5,
the selected layer for Cora would be the third student layer (index 2 in the
figure). We then proceed to instantiate a new one-layer network and copy over
the weights from the identified layer from $A$. We then evaluate this new
network on the test set and report the results in column 3 of Table 4. The
first column of Table 4 represents the accuracies that we obtain after we
train the new one-layer network for $1200$ epochs; we compare this result to
the accuracy obtained from training a one-layer network from random
initialization, which we report in the second column of Table 4.

###  B.3 Hyperparameter Ablation Studies

There are two main hyperparameters that need to be tuned when training ABKD:
the loss coefficient, $\beta$, and the ABKD embedding dimension, $d_{a}$. We
present the test accuracies across various values for $\beta$ and $d_{a}$ in
Tables 10 and 11. These results show that a lower $\beta$ and a higher $d_{a}$
tend to produce slightly better results. For all of our experiments, we used a
$\beta$ of $10$ and a $d_{a}$ of 256 for this reason.

Dataset | $\beta=1$ | $\beta=10$ | $\beta=20$ | $\beta=50$
---|---|---|---|---
Cora | $86.92$ | $84.71$ | $86.19$ | $85.64$
Citeseer | $73.20$ | $74.33$ | $71.82$ | $69.82$
Pubmed | $89.03$ | $89.97$ | $88.32$ | $88.56$
NELL | $90.86$ | $88.73$ | $90.14$ | $91.32$

Table 10: Ablation results for $\beta$.

Dataset | $d_{a}=64$ | $d_{a}=128$ | $d_{a}=256$ | $d_{a}=512$
---|---|---|---|---
Cora | $87.45$ | $87.11$ | $86.92$ | $86.37$
Citeseer | $72.07$ | $72.52$ | $73.12$ | $74.62$
Pubmed | $89.58$ | $89.58$ | $89.33$ | $89.12$
NELL | $89.73$ | $90.12$ | $90.73$ | $91.14$

Table 11: Ablation results for $d_{a}$.

Dataset | Euclidean Distance | Cosine Distance
---|---|---
Cora | $87.33$ | $82.81$
Citeseer | $73.43$ | $70.98$
Pubmed | $89.58$ | $87.32$
NELL | $91.24$ | $85.66$
Table 12: Test accuracies when using Euclidean distance vs. cosine distance
for computing $D$.

In Section 3.2.2, we mention that we use a Euclidean distance metric instead
of a cosine distance metric to generate the dissimilarity matrix, $D$. We
present the results of this ablation in Table 12. As cosine distance
disregards the magnitude of the vectors, we hypothesize that the magnitude of
the hidden layer outputs is important. Furthermore, as Euclidean distance
takes into account the magnitude of the hidden layer outputs, this is likely
why it performs significantly better than cosine distance.


# Improving style transfer in dynamic contrast enhanced MRI using a spatio-
temporal approach

Adam G. Tattersall, Keith A. Goatman, Lucy E. Kershaw,
Scott I. K. Semple and Sonia Dahdouh

###### Abstract

Style transfer in DCE-MRI is a challenging task due to large variations in
contrast enhancements across different tissues and time. Current unsupervised
methods fail due to the wide variety of contrast enhancement and motion
between the images in the series. We propose a new method that combines
autoencoders to disentangle content and style with convolutional LSTMs to
model predicted latent spaces along time and adaptive convolutions to tackle
the localised nature of contrast enhancement. To evaluate our method, we
propose a new metric that takes into account the contrast enhancement.
Qualitative and quantitative analyses show that the proposed method
outperforms the state of the art on two different datasets.

##  1 Introduction

DCE-MRI is a type of quantitative imaging that can be used to monitor
microvascular perfusion [7]. Multiple T1 weighted images are taken rapidly
over a few minutes whilst a contrast agent injection causes a rapid increase
in signal intensity in tissues of interest. The resulting data contain motion
and tissue-dependent contrast enhancement (CE) which makes obtaining manually
annotated ground truth time consuming and difficult. Fully annotated datasets
are therefore scarce, which reduces the diversity of available data to train
deep learning models. We previously found that to train a robust model with
DCE-MRI data, using images strategically from the whole sequence with varying
contrast levels led to better performing models [16]. However, in practice
this can be difficult for tasks such as segmentation when ground truth is
limited. Although common augmentation techniques allow for an increase in the
number of images, they do not help to synthesise new annotated examples with
varying CE. A method that can synthesise new images from previously annotated
data would provide models with diverse examples.

In this work, we propose a new style transfer method to augment DCE-MRI with
different contrast enhancement levels. Our method uses convolutional long
short-term memory (LSTM) [2] networks which leverage the temporal information
from DCE-MRI to better predict structures of the content image and generate
images with the noise characteristics of MRI whilst allowing addition or
removal of contrast enhancement. We also propose a new metric, contrast
weighted (CW) - structural similarity index measure (SSIM), that uses
information on the localised aspect of tissue enhancement to evaluate either
structural or style similarity between a generated and a content or style
image, respectively.

##  2 Related Work

Using neural networks for style transfer was first proposed by Gatys et al.
[4] which used a pre-trained VGG-Network [15]. Instead of training, an
iterative process of passing the content, style and output image through a
pre-trained network was used. Building on this success, improvements were made
to speed up the iterative process by using a feed-forward network [17] as well
as creating perceptual losses [9] to give more visually pleasing results.
However, these networks struggled to adapt to new styles that were introduced
during inference as the network was usually tied to a small number of style
images. In 2017, Huang et al. [5] proposed Adaptive Instance Normalisation
(AdaIN) which aligns the mean and variance of the content features with those
of the style features allowing for style transfer without the restriction of a
pre-defined set of styles.

Successful style transfer methods have since been proposed such as learning
mappings between pairs of images (Pix2Pix [8] and CycleGAN [18]) as well as
content/style disentanglement methods such as Multimodal Unsupervised Image-
to-Image Translation (MUNIT) [6]. Pix2Pix was able to generate high quality
images using a conditional GAN but required paired and registered images to
train effectively. Instead, CycleGAN can be trained using unpaired data by
using two generators to map between each domain. However, both Pix2Pix and
CycleGAN struggle to generate diverse outputs when new style images are
introduced at inference due to learning to map between domains directly.

MUNIT offers an alternative approach using content/style disentanglement to
take input images and predict latent spaces for style and content encoding.
This allows style transfer whilst preserving the content of the image, giving
better control for the translation process. However, MUNIT often requires a
large amount of data to train as well as having large computational and memory
requirements. AdaIN is also used to combine the style and content latent
spaces. It can, however, result in global rather than tissue dependent
enhancement which leads to the generation of unrealistic images. Additionally,
MUNIT does not utilise any temporal information available in the DCE-MRI
sequence.

DRIT++ [11], an extension to MUNIT, enhances diversity by introducing
disentanglement at the domain level, and improving attribute manipulation.
However, it also comes with increased computational complexity, complexity in
hyperparameter tuning, and a potential trade-off between diversity and
fidelity.

Early work on style transfer for videos naively processed frames independently
leading to flickering and false discontinuities [3]. Methods have since been
proposed to ensure temporal consistency between image frames [3]. They use
three components: a style sub-network, a flow sub-network and a mask sub-
network. The style sub-network is an auto-encoder which performs the style
transfer. The flow sub-network estimates the correspondence between
consecutive image frames and warps the features. Finally, the mask sub-network
regresses a mask to compose features adjacent in time, allowing for warped
features to be reused. This allows for a smooth style transfer between image
frames, but it is prone to errors when there is large motion between frames as
well as propagating errors over time leading to inconsistent style transfer
and blurriness.

##  3 Materials and Methods

###  3.1 Data

Two datasets were used. The first contains 2D kidney DCE-MRI [13] from 13
patients with 375 images acquired continuously at a temporal resolution of 1.6
s, a spatial resolution of 384x348 with pixel sizes 1.08x1.08 mm. The second
contains 3D prostate DCE-MRI [12] from 20 patients with 40 volumes acquired
continuously at a temporal resolution of 6 s, a spatial resolution of
256x192x16 with voxel sizes 1.12x1.12x3.5 mm. Each dataset was split at a
patient level with an 80:20 ratio, training was done using five-fold cross
validation.

###  3.2 Global Architecture

Inspired by MUNIT, we use a structure composed of encoders and decoders to
first disentangle content and style latent spaces and recombine them to
generate new images. Successful disentanglement should lead to content latent
spaces containing information representing the structures of the image and
style latent spaces containing information representing the modality (MRI) and
contrast enhancements. However, temporal information is not taken into account
in most style transfer methods. Within DCE-MRI, information such as structures
within organs become visible when there are various levels of contrast
enhancement at different time points. In our method, we use additional content
encoders to predict latent spaces for images from multiple time points. Then,
we pass the predicted content latent spaces to a convolutional LSTM [2] which
has been set up to be bi-directional. We use a convolutional LSTM as this
allows modelling of the temporal information from a sequence of images. It
leverages spatial invariance by extracting relevant features from the images
regardless of their position. Using a bi-directional LSTM allows better
contextual understanding of the time series data by accessing past and future
events. This allows content visible at various time points to be learnt. The
overall architecture is shown in figure 1. At inference the images are passed
through the encoders and decoders once to perform style transfer.

Figure 1: Our model takes an input I of 5 images/volumes to content encoders
E${}^{c}_{0...4}$ and style encoders E${}^{s}_{0...4}$. Generators GCE and
GNCE construct images from latent spaces z predicted by the content and style
encoders. We also show the losses used: cycle consistency, perceptual and
adversarial. For clarity of the figure, we have only shown one example for
each of the losses.

###  3.3 Instance Normalisation

AdaIN [5] is used to normalise the activations in a neural network based on
the statistics of the style image. This enables the transfer of style from the
style image to the content image by matching their statistical properties.
However, by only using global style information, contrast enhancement is
performed on the whole image instead of specific tissues. To overcome this, we
use AdaConv [1] which captures global and local information to predict
parameters to combine style and content latent spaces. Three convolutional
neural networks are used to predict a depth-wise spatial kernel, a point-wise
kernel and a bias. The kernels and bias are then convolved with the content
code.

###  3.4 Losses

The L1 loss between the original and reconstructed images and between the
initially predicted and reconstructed latent spaces is used to ensure that
important information such as structure is not lost during the reconstruction
and translation process. Furthermore, the L1 loss is known for its smooth
gradients which can facilitate stable and effective optimisation during the
training process. A mean squared error (MSE) loss is also used as an
adversarial loss to help generate images that are indistinguishable from the
real images.

While those two losses were used in MUNIT, we found that some important
textural characteristics of the images were lost in the generated images. To
improve this, we also used two additional perceptual loss functions proposed
by Johnson et al. [9] for style transfer and super resolution tasks. The first
is a feature reconstruction loss which is the squared, normalised Euclidean
distance (Eq. 1). This encourages the feature representations predicted by a
pretrained model P of the content image x and the generated image g to be
similar. f is the dimension of the feature map.

$L_{feature}=\frac{1}{f}||P(g)-P(x)||^{2}_{2}$ (1)

Next, we use a style reconstruction loss to penalise differences in style
between the style image y and the generated image g. This loss uses the
squared Frobenius norm of the difference between the Gram matrices of the
feature representations predicted by a pretrained network P (Eq. 2).
Minimising this loss encourages the generated image to capture style patterns
and textures that are similar to the style image. The Frobenius norm is a
measure of the magnitude of the matrix, and using it as the loss term allows
the model to focus on preserving the overall style content and structure.

$L_{Frob}=||Gram(P(g))-Gram(P(y))||^{2}_{F}$ (2)

For the 2D data, we used a VGG-16 pre-trained network on ImageNet [14]
converted to grayscale. For the 3D data we used MedicalNet which is a
pretrained multi-modal multi-tissue ResNet-18 backbone.

###  3.5 Evaluating style transfer

Quantitatively evaluating style transfer is difficult with unpaired,
unregistered data. Metrics such as peak signal to noise ratio (PSNR), SSIM and
multiscale (MS)-SSIM primarily focus on the pixel-level similarity between a
pair of images. This makes it difficult to evaluate style transfer in DCE-MRI
due to the large variety of CE in an image and the localised aspect of it. We
propose a contrast weighted (CW)-SSIM (Eq. 3) as a metric to measure both the
content structure as well as evaluate the amount of CE (addition or removal)
in the style image.

For evaluating content, we weight the SSIM with a distance map dist_map
created by calculating the shortest distance from each voxel to a contrast
enhanced voxel. The distances were normalised between 0.1 and 1. We determined
which voxels were contrast enhanced by subtracting each image from the first
non-CE image and calculated the average. Then, we used an empirical threshold
of $20$ to determine which voxels were contrast enhanced. To evaluate style,
we inverted the distance map so that a voxel has a higher weighting when it is
closer to a contrast enhanced voxel. x is the generated image and y is either
the content or style image.

$\displaystyle\mathrm{CW{\text{-}}SSIM}(\mathrm{x},\mathrm{y},\mathrm{dist\\_map})=\mathrm{SSIM(x\cdot\mathrm{dist\\_map},y\cdot\mathrm{dist\\_map})}$
(3)

###  3.6 Implementation

We used PyTorch and trained with the Adam optimiser, a learning rate of 0.001
and batch size of 8 for 2D and 5 for 3D. Early stopping was used with a
patience of 20. Our experiments were run on an RTX Titan GPU. We compared our
method with MUNIT, CycleGAN and StyleGAN3 [10]. DRIT++ was not used due to the
large computational requirements with 3D data.

##  4 Results

Tables 1 and 2 show our quantitative results: PSNR between the style (image we
want to transfer style from) and generated image, SSIM and MS-SSIM between the
content (image we want to take structure from) and generated image and
finally, our proposed weighted SSIMs. For each style transfer direction and
metric, our method consistently outperforms the other approaches. Figures 2
and 3 qualitatively highlight the good results of our method on 2D and 3D
datasets. In addition, quantitative results on both figures show that our
proposed metric preserves qualitative ordering of results for both style and
content. In comparison, metrics such as SSIM struggle to accurately correlate
to visual results.

Table 1: Quantitative results for each approach when adding or removing CE in the kidney. The mean and standard deviation is shown for each metric. Metric | Direction | MUNIT | CycleGAN | StyleGAN3 | Our Method
---|---|---|---|---|---
PSNR | Non-CE to CE | 70.6 $\pm$ 3.2 | 66.5 $\pm$ 2.7 | 52.3 $\pm$ 3.9 | 81.3 $\pm$ 2.7
CE to Non-CE | 71.5 $\pm$ 3.4 | 66.8 $\pm$ 2.8 | 51.9 $\pm$ 3.8 | 78.5 $\pm$ 2.6
SSIM | Non-CE to CE | 0.61 $\pm$ 0.05 | 0.66 $\pm$ 0.06 | 0.32 $\pm$ 0.06 | 0.89 $\pm$ 0.04
CE to Non-CE | 0.58 $\pm$ 0.04 | 0.58 $\pm$ 0.05 | 0.34 $\pm$ 0.07 | 0.88 $\pm$ 0.03
MS-SSIM | Non-CE to CE | 0.55 $\pm$ 0.04 | 0.43 $\pm$ 0.07 | 0.34 $\pm$ 0.04 | 0.82 $\pm$ 0.03
CE to Non-CE | 0.52 $\pm$ 0.04 | 0.45 $\pm$ 0.05 | 0.37 $\pm$ 0.05 | 0.79 $\pm$ 0.04
Content | Non-CE to CE | 0.73 $\pm$ 0.03 | 0.61 $\pm$ 0.05 | 0.47 $\pm$ 0.05 | 0.95 $\pm$ 0.02
CW-SSIM | CE to Non-CE | 0.69 $\pm$ 0.03 | 0.57 $\pm$ 0.05 | 0.41 $\pm$ 0.05 | 0.94 $\pm$ 0.02
Style | Non-CE to CE | 0.64 $\pm$ 0.03 | 0.41 $\pm$ 0.05 | 0.67 $\pm$ 0.05 | 0.74 $\pm$ 0.02
CW-SSIM | CE to Non-CE | 0.62 $\pm$ 0.03 | 0.49 $\pm$ 0.05 | 0.64 $\pm$ 0.05 | 0.72 $\pm$ 0.02

(a) No CE Image

(b) CE Image

SSIM: 0.59
(C) CW-SSIM: 0.71
(S) CW-SSIM: 0.63

SSIM: 0.65
(C) CW-SSIM: 0.59
(S) CW-SSIM: 0.41

SSIM: 0.32
(C) CW-SSIM: 0.44
(S) CW-SSIM: 0.69

SSIM: 0.91
(C) CW-SSIM: 0.94
(S) CW-SSIM: 0.75

SSIM: 0.57
(C) CW-SSIM: 0.68
(S) CW-SSIM: 0.61 (c) MUNIT

SSIM: 0.56
(C) CW-SSIM: 0.54
(S) CW-SSIM: 0.51 (d) CycleGAN

SSIM: 0.31
(C) CW-SSIM: 0.42
(S) CW-SSIM: 0.67 (e) StyleGAN3

SSIM: 0.89
(C) CW-SSIM: 0.93
(S) CW-SSIM: 0.73 (f) Our Method

Figure 2: Example results from different style transfer approaches with the
kidney data. The first row (a and b) shows the input images. The second row
shows the results when (a) is used as the content image and (b) as the style.
The third row shows results when (b) is used as the content image and (a) as
the style. We also show scores given by the SSIM, content (C) CW-SSIM and
style (S) CW-SSIM.

(a) No CE Image

(b) CE Image

SSIM: 0.41
(C) CW-SSIM: 0.39
(S) CW-SSIM: 0.29

SSIM: 0.44
(C) CW-SSIM: 0.53
(S) CW-SSIM: 0.31

SSIM: 0.32
(C) CW-SSIM: 0.42
(S) CW-SSIM: 0.54

SSIM: 0.69
(C) CW-SSIM: 0.93
(S) CW-SSIM: 0.62

SSIM: 0.71
(C) CW-SSIM: 0.61
(S) CW-SSIM: 0.52 (c) MUNIT

SSIM: 0.61
(C) CW-SSIM: 0.42
(S) CW-SSIM: 0.19 (d) CycleGAN

SSIM: 0.42
(C) CW-SSIM: 0.28
(S) CW-SSIM: 0.49 (e) StyleGAN3

SSIM: 0.79
(C) CW-SSIM: 0.92
(S) CW-SSIM: 0.59 (f) Our Method

Figure 3: Example results from different style transfer approaches with the prostate data. The first row (a and b) shows the input images. The second row shows the results when (a) is used as the content image and (b) as the style. The third row shows results when (b) is used as the content image and (a) as the style. We also show scores given by the SSIM, content (C) CW-SSIM and style (S) CW-SSIM. Table 2: Quantitative results for each approach when adding or removing CE in the prostate. The mean and standard deviation is shown for each metric. Metric | Direction | MUNIT | CycleGAN | StyleGAN3 | Our Method
---|---|---|---|---|---
PSNR | Non-CE to CE | 40.6 $\pm$ 4.3 | 52.2 $\pm$ 3.8 | 43.5 $\pm$ 5.2 | 72.1 $\pm$ 3.6
CE to Non-CE | 68.7 $\pm$ 4.1 | 61.2 $\pm$ 3.4 | 51.9 $\pm$ 6.1 | 68.9 $\pm$ 4.1
SSIM | Non-CE to CE | 0.41 $\pm$ 0.06 | 0.46 $\pm$ 0.05 | 0.36 $\pm$ 0.09 | 0.71 $\pm$ 0.04
CE to Non-CE | 0.69 $\pm$ 0.04 | 0.59 $\pm$ 0.05 | 0.43 $\pm$ 0.08 | 0.78 $\pm$ 0.03
MS-SSIM | Non-CE to CE | 0.49 $\pm$ 0.04 | 0.38 $\pm$ 0.06 | 0.28 $\pm$ 0.05 | 0.76 $\pm$ 0.04
CE to Non-CE | 0.47 $\pm$ 0.04 | 0.39 $\pm$ 0.05 | 0.32 $\pm$ 0.06 | 0.71 $\pm$ 0.03
Content | Non-CE to CE | 0.38 $\pm$ 0.04 | 0.52 $\pm$ 0.05 | 0.41 $\pm$ 0.05 | 0.92 $\pm$ 0.03
CW-SSIM | CE to Non-CE | 0.63 $\pm$ 0.04 | 0.46 $\pm$ 0.04 | 0.32 $\pm$ 0.06 | 0.93 $\pm$ 0.04
Style | Non-CE to CE | 0.34 $\pm$ 0.05 | 0.28 $\pm$ 0.06 | 0.52 $\pm$ 0.07 | 0.59 $\pm$ 0.05
CW-SSIM | CE to Non-CE | 0.53 $\pm$ 0.04 | 0.21 $\pm$ 0.05 | 0.47 $\pm$ 0.06 | 0.57 $\pm$ 0.04

##  5 Discussion

In this work, we proposed a method for style transfer that leverages spatio-
temporal information, applied to DCE-MRI. Using convolutional LSTMs we are
able to model changes in structures caused by CE across the temporal dimension
giving better content latent spaces. In addition, using localised style
information instead of a global one by making use of AdaConv helped to tackle
the local nature of CE in DCE-MRI.

Qualitative evaluation shows that the proposed method leads to sharper images,
better content preservation, better localised CE and realistic MRI appearance.
Quantitatively, we outperform the other algorithms with each metric, for each
style transfer direction. While we expected that adding CE to images would be
an easier task than removing CE, we found that for both tasks the method
exhibits similarly performing quantitative and qualitative results. The
results for style transfer with the kidney data show that when there is clear
CE, it is easier to perform style transfer in both directions. However, when
style transfer is performed on the prostate data, the model struggles in
comparison to the kidney data. This may be due to the enhancement of the
prostate being less defined than the kidney data. In figure 3(b), we can see
the prostate with CE. Compared to figure 2(b) showing kidneys with CE, it is
harder to determine the edges. The two original images shown in figures 3(a)
and 3(b) are clear and easy to see some edges of prostate, in other images, it
is harder to see the edges of the prostate.

Our metric, CW-SSIM, enables the evaluation of the quality of the content and
style of the generated images when there are large differences in CE. When
comparing against the qualitative results, the quality of the style transfer
is in line with each of the CW-SSIM results. By using two weightings: areas of
CE and without CE, we are able to evaluate the similarity between structures
where intensity differences are small to measure for content similarity. By
using an inverted map, we can also compare regions of high CE to measure the
differences in style between the generated and style image. An advantage of
this metric over SSIM or MS-SSIM is that we are able to separate the
evaluation of content and style to ensure that key structures in the image
remain, whilst changing the intensity values to match the style. This is
important with DCE-MRI as large intensity changes are localised in specific
tissues. When we compare the results for the prostate using MS-SSIM and style
CW-SSIM, we can see that the style CW-SSIM gave a much lower score. This
follows the qualitative results closely as the styles of each of the generated
images are not similar to the style image used. By evaluating the areas of the
image separately we remove the possibility of parts of the image with little
CE inflating the score given by the SSIM.

A limitation to our style transfer approach is that better style transfer was
achieved with the kidney data compared to the prostate data. This may be due
to the structures in the prostate data not being as clearly defined as the
kidney data, as well as the prostate not enhancing as clearly as the kidney.
Similarly to MUNIT, our method has large computational and memory
requirements. Using our method, new images can be generated, to contain a
variety of stages of CE. This type of augmentation allows a small number of
images to be manually annotated making it easier to train robust models for
image processing tasks.

##  6 Conclusion

We proposed a style transfer method which used temporal information with
disentangled representations. Our method learns information from time series
data and accurately captures key structures from content images whilst
localising the addition of CE. Our method is trained in an unsupervised
setting with unpaired data. Qualitative and quantitative analysis shows that
our method outperforms other popular style transfer techniques. We also
proposed a new metric which uses information based on the localised aspect of
tissue enhancement with contrast in DCE-MRI to evaluate similarity to the
style or content image.

##  7 Acknowledgements

This work was funded by Medical Research Scotland and CMRE.

## References

  * [1] Chandran, P., Zoss, G., Gotardo, P., Gross, M., Bradley, D.: Adaptive Convolutions for Structure-Aware Style Transfer. CVPR (2021)
  * [2] Chao, Z., Pu, F., Yin, Y., Han, B., Chen, X.: Research on real-time local rainfall prediction based on MEMS sensors. Journal of Sensors (2018)
  * [3] Chen, D., Liao, J., Yuan, L., Yu, N., Hua, G.: Coherent online video style transfer. In: IEEE ICCV (2017)
  * [4] Gatys, L., Ecker, A., Bethge, M.: A Neural Algorithm of Artistic Style. Journal of Vision (2016)
  * [5] Huang, X., Belongie, S.: Arbitrary Style Transfer in Real-Time with Adaptive Instance Normalization. IEEE ICCV (2017)
  * [6] Huang, X., Liu, M.Y., Belongie, S., Kautz, J.: Multimodal unsupervised image-to-image translation. In: ECCV (2018)
  * [7] Ingrisch, M., Sourbron, S.: Tracer-kinetic modeling of dynamic contrast-enhanced MRI and CT: A primer. Journal of Pharmacokinetics and Pharmacodynamics (2013)
  * [8] Isola, P., Zhu, J.Y., Zhou, T., Efros, A.A.: Image-to-image translation with conditional adversarial networks. CVPR (2017)
  * [9] Johnson, J., Alahi, A., Fei-Fei, L.: Perceptual losses for real-time style transfer and super-resolution. In: ECCV (2016)
  * [10] Karras, T., Aittala, M., Laine, S., Härkönen, E., Hellsten, J., Lehtinen, J., Aila, T.: Alias-free generative adversarial networks. In: NeurIPS (2021)
  * [11] Lee, H.Y., Tseng, H.Y., Mao, Q., Huang, J.B., Lu, Y.D., Singh, M.K., Yang, M.H.: Drit++: Diverse image-to-image translation via disentangled representations. International Journal of Computer Vision (2020)
  * [12] Lemaître, G., Martí, R., Freixenet, J., Vilanova, J.C., Walker, P.M., Meriaudeau, F.: Computer-aided detection and diagnosis for prostate cancer based on mono and multi-parametric mri: A review. Computers in Biology and Medicine (2015)
  * [13] Lietzmann, F., Zöllner, F.G., Attenberger, U.I., Haneder, S., Michaely, H.J., Schad, L.R.: Dce-mri of the human kidney using blade: A feasibility study in healthy volunteers. Journal of Magnetic Resonance Imaging (2012)
  * [14] Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A.C., Fei-Fei, L.: ImageNet Large Scale Visual Recognition Challenge. IJCV (2015)
  * [15] Simonyan, K., Zisserman, A.: Very deep convolutional networks for large-scale image recognition. ICLR (2015)
  * [16] Tattersall, A.G., Goatman, K.A., Kershaw, L.E., Semple, S.I.K., Dahdouh, S.: Can a single image processing algorithm work equally well across all phases of DCE-MRI. SPIE (2023)
  * [17] Ulyanov, D., Lebedev, V., Vedaldi, A., Lempitsky, V.: Texture networks: Feed-forward synthesis of textures and stylized images. ICML (2016)
  * [18] Zhu, J.Y., Zhang, R., Pathak, D., Darrell, T., Efros, A.A., Wang, O., Shechtman, E.: Toward multimodal image-to-image translation. NeurIPS (2017)