# Learning Spatio-Temporal Model of Disease Progression with NeuralODEs from
Longitudinal Volumetric Data

Dmitrii Lachinov  Arunava Chakravarty  Christoph Grechenig  Ursula Schmidt-
Erfurth  Hrvoje Bogunović Manuscript submitted Oct, 14, 2022. This work was
supported in part by the Christian Doppler Research Association, Austrian
Federal Ministry for Digital and Economic Affairs, the National Foundation for
Research, Technology and Development, and Heidelberg Engineering.Dmitrii
Lachinov and Hrvoje Bogunović are with the Christian Doppler Lab for
Artificial Intelligence in Retina, Department of Ophthalmology and Optometry,
Medical University of Vienna, Austria (e-mail:
<EMAIL_ADDRESS>hrvoje.bogunovic@meduniwien.ac.at).Arunava
Chakravarthy, Christoph Grechenig, and Ursula Schmidt-Erfurth are with the
Department of Ophthalmology and Optometry, Medical University of Vienna,
Austria.

###### Abstract

Robust forecasting of the future anatomical changes inflicted by an ongoing
disease is an extremely challenging task that is out of grasp even for
experienced healthcare professionals. Such a capability, however, is of great
importance since it can improve patient management by providing information on
the speed of disease progression already at the admission stage, or it can
enrich the clinical trials with fast progressors and avoid the need for
control arms by the means of digital twins. In this work, we develop a deep
learning method that models the evolution of age-related disease by processing
a single medical scan and providing a segmentation of the target anatomy at a
requested future point in time. Our method represents a time-invariant
physical process and solves a large-scale problem of modeling temporal pixel-
level changes utilizing NeuralODEs. In addition, we demonstrate the approaches
to incorporate the prior domain-specific constraints into our method and
define temporal Dice loss for learning temporal objectives. To evaluate the
applicability of our approach across different age-related diseases and
imaging modalities, we developed and tested the proposed method on the
datasets with 967 retinal OCT volumes of 100 patients with Geographic Atrophy,
and 2823 brain MRI volumes of 633 patients with Alzheimer’s Disease. For
Geographic Atrophy, the proposed method outperformed the related baseline
models in the atrophy growth prediction. For Alzheimer’s Disease, the proposed
method demonstrated remarkable performance predicting the brain ventricle
changes induced by the disease, achieving the state-of-the-art result on
TADPOLE challenge.

Disease Progression, Deep Learning, Longitudinal Imaging, Retina, Geographic
Atrophy, Alzheimer’s Disease

##  1 Introduction

The progress in the medical imaging techniques brought vast possibilities for
the diagnosis and monitoring of patients with age-related diseases. Tools such
as Optical Coherence Tomography (OCT) or Magnetic Resonance Imaging (MRI) are
the current clinical gold standard and contribute to an overwhelming number of
diagnostic findings. The amount of information within these scans, however, is
so large, that it requires many years of training and practice to interpret
them. Yet, we believe, a lot stays hidden from expert eyes. For instance, we
hypothesize that the track of the natural, age-related disease progression is
determined by the subtle physiological changes which are captured and encoded
by medical scans. The class of computational methods trying to extract such
information is called Disease Progression Models (DPMs), which can, e.g.
calculate the expected future change in the target quantitative biomarkers. A
more sophisticated class of methods can forecast morphological change with
respect to the shape and size of an anatomical structure or lesion that is
expected to occur in the future, from an image acquired at the time of
screening. The task of modeling future changes is extremely challenging as the
early biomarkers that indicate either a fast or slow speed of progression are
often very subtle and subclinical, i.e., not well-understood in clinical
practice.

DPMs have already been employed as decision support and as tools complementing
phase II and III clinical trials [1, 2, 3]. They typically model the evolution
of a set of predefined biomarkers over the course of the disease and map it to
the global timeline of disease evolution, where the biomarkers are typically
extracted from clinical reports and imaging data. However, such biomarkers
relevant to model the disease progression have to be predetermined by the
human experts, which is often a non-trivial task, and likely does not cover
all the revelant ones. Thus, the search for predictive biomarkers is currently
a complicated and tedious procedure.

The state-of-the-art methods that automatically extract features from the
input volume also have some weaknesses, as some are designed for forecasting
of fixed intervals only; others are time-variant systems, as the same input
volume can produce different outcomes at different times. In this work, we
propose a novel disease progression model capable of learning directly from
longitudinal imaging data to predict the target area or volume change from a
single scan, acquired at baseline visit. The proposed model not only predicts
continuous-in-time changes in the target size, but it is time-invariant and
inherently interpretable, as it simultaneously produces a detailed
segmentation map of the current and the future target anatomy.

###  1.1 Clinical background

We focus on developing a DPM method for two clinically-relevant scenarios: (1)
modeling the growth of Geographic Atrophy (GA) from longitudinal 3D OCT
volumes of the retina to monitor Age-Related Macular Degeneration (AMD), and
(2) modeling Alzheimer’s Disease (AD) progression by capturing the
morphological changes in the brain ventricles from longitudinal 3D MRI volumes
of the brain.

####  1.1.1 Geographic Atrophy

AMD is a leading cause of blindness, affecting 8.7% of the world population
[4]. In the 80+ age group, this number grows to 25%, with the prevalence of
advanced AMD of 3.3%. AMD in its late stages can either be dry characterized
by the presence of Geographic Atrophy (GA) or wet characterized by the
presence of Macular Neovascularization. OCT has become a gold standard
modality for imaging the retina in the clinical setting due to its non-
invasive 3D acquisition and widespread availability in clinics. It acquires a
detailed volumetric scan of the retina made of a series of cross-sectional
slices or BScans (Fig. 1). On a BScan, GA is represented as an atrophy of the
outer retinal tissues [5]. Thus, each column (AScan) in a BScan can be denoted
as atrophic or not. This makes the GA lesion to be 2D in nature (Fig. 1), with
its size typically measured as an en-face area in mm2. In contrast to the
late-wet stage of AMD, the late-dry stage is more prevalent, affecting around
5 Mio worldwide, and has no approved treatments yet. However, very recently
multiple drugs have passed phase 3 trials, most prominently pagcetacoplan [6]
and avacincaptad pegol [7]. Thus, modeling the progression of GA plays a
critical role in clinical research to understand the pathophysiology, and the
simulation of the future atrophy progression is expected to be fundamental to
patient management in order to identify fast progressors that will need the
treatment the most.

####  1.1.2 Alzheimer’s Disease

AD is a gradually progressing, irreversible neurodegenerative disorder which
is the primary of Dementia. Around 50 million patients suffer from this
impairment worldwide, with the projected number of patients effectively
doubling each 5 years and reaching 152 million by 2050 [8]. Currently, only
symptomatic treatment is available. Volumetric MRI of the brain is commonly
employed in the clinical examination of AD to detect the early pathological
changes in the cortex, hippocampus and ventricles. Recent studies by [9, 10]
have correlated the ventricle volume extracted from MRI images to the
progression of AD, thereby making it a potential biomarker for future
interventions and disease progression modelling. Recent challenge, TADPOLE
[11] aimed at predicting future evolution of the AD patients, where the
ventricles’ volume was used as one of the main outcome measures. Thus, we
focus on the developmment of ventricles captured by brain MRI, and predicting
their 3D segmentation evolution over a three-year period (Fig. 1).

Figure 1: The proposed method predicts the future target anatomy segmentation
from a single volumetric scan acquired at baseline (day 0). The top row
represents the en-face projections of the volumetric 3D OCT scan with the
atrophy segmentation predicted at different future visits, denoted by the
number of days since the baseline visit. The raw baseline OCT volume is shown
as an en-face projection and with a central BScan (cross-section). Note that
new islands of atrophy form in the central right and top left locations of the
scan. The bottom row represents an axial slice of the brain T1 MRI cropped
around the ventricles, with their predicted segmentations at future time
points. The growth in a 2D slice is marginal, but it is substantial in 3D. In
both rows, true positives (blue), false positives (orange), and false
negatives (red) are depicted.

###  1.2 Related work

Addressing DPM with Deep Learning is an active area of research. Existing
methods can be broadly categorized into two approaches. In the first category
are the methods that operate on tabular data summarizing volume-level
quantitative biomarkers, such as the volume of relevant structures, which have
been derived from the imaging data through manual or automated segmentation.
Such image-level methods typically employ Recurrent Neural Networks (RNNs)
[12, 13] to fit the temporal variation of the derived biomarker values.
Methods in the second category utilize the raw imaging data directly for
spatio-temporal modelling of the disease, typically in the form of respective
future lesion segmentation maps. We will refer to these methods as pixel-level
progression methods, and focus on them as this is the approach of our proposed
method.

For modeling the GA progression, Zhang et al. [14] proposed a pixel-level
multi-scale deep Convolutional Neural Network (MS-DCNN) for joint segmentation
and prediction of GA. The proposed method takes preprocessed baseline en-face
OCT projection and produces segmentation at the specified future visits.
However, the method lacks the ability to interpolate the progression between
the visits. In addition, the dataset used in experiments was small, consisting
of only 29 patients. In a follow-up work, Zhang et al. [15] proposed a Bi-
LSTM-based prediction model and a UNet based refinement module. However, it
requires two visits to make a prediction and the dataset used was still small,
consisting of only 25 patients. Recently, Gigon et al. [16] proposed a pixel-
level model that uses a Taylor series with learnable time derivatives to
approximate patient-specific continuous GA growth. The method takes a set of
retinal layer thickness maps as an input, which makes the method application-
specific and limits the amount of information the model can learn from. As in
previously mentioned works, the method was tested on a low amount of cases,
precisely 20.

For the AD progression, Wang et al. [12] proposed to utilize a image-level
method with LSTM[17] to forecast the AD stage at the upcoming visits using 78
features derived from demographics, patient history and other sources. In this
method, the prediction interval is provided as part of the input. Nguyen et
al. [13] employed a similar approach utilizing MinimalRNN[18], where the
authors linked each recursive update with one-month interval. The authors used
the set of features provided as a part of TADPOLE[11] challenge. Among teams
participating in TADPOLE[11], Hill et al. [19] as a part of team GlassFrog
employed high-dimensional regression and disease state specific slope models
and ensembling for ventricle volume forecasting. Ismail et al. [19] as a part
of BravoLab team proposed to use LSTM for making forecasts. Venkatraghavan et
al.[19] fitted a linear mixed-effects model to estimate ventricle volume
expansion rate. The mentioned methods, however, don’t exploit the information
from the entire volume, since they rely on a set of input features, which were
pre-defined by the human experts.

Several recent works utilized Recurrent Neural Networks (RNNs) [12, 13]. Wang
et al. [12] proposed to use RNN trained on a set of biomarker values to
predict the Alzheimer’s Disease progression at the next visit. The proposed
method supports irregular visit intervals by introducing an additional feature
$\Delta t$ \- denoting the time to the next visit. Nguyen et al. [13] trained
a RNN on a set of quantitative values to predict a diagnosis and a set of
Alzheimer’s Disease biomarkers for every month up to six months in the future.
The authors interpolated the missing visits in the data by generating them
with the proposed model. This, however, doesn’t handle variable time
intervals. A notable pixel-level method, was proposed by Ezhov et al. [20].
The paper discussed the method to fit the parameters of the partial
differential equation (PDE), that defines the brain tumor growth, to the real
imaging data. Unfortunately, the solution is hardly generalizable outside the
tumor growth tasks, as it would require knowing a form of PDE beforehand. In
addition, the method was only tested on 2 real rat MRI scans. Another pixel-
level method was proposed by Petersen et al. [21]. The authors trained a deep
learning model to learn tumor growth dynamics directly from the imaging data.
The work focuses on modelling the distribution of the possible tumor
appearances. However, as the authors mentioned, the method is neither suitable
for predicting a single correct growth trajectory nor designed for it.
Expanding the work on probabilistic disease modelling, Petersen et al.[22]
introduced a continuous progression model, based on UNet[23] and attention
mechanism [24]. As before, the model wasn’t explicitly designed for predicting
a single trajectory, but rather a distribution of possible outcomes.
Additionally, the model was trained to utilize two or more images as prior
information, which reduces the method’s applicability in certain scenarios,
such as screening for fast progressors from a single scan for enriching the
clinical trials.

In this paper, we base our proposed pixel-level DPM method on NeuralODEs [25].
The distinctive features of the NeuralODEs that make them especially suitable
for DPM is that they are inherently continuous, which allows forecasting at
any future time-point; they represent some underlying physical process, and,
at the simplest level, the neural network learned by NeuralODEs correspond to
the instantaneous change in the target variable. In addition, NeuralODEs allow
for controlling the solution’s properties and imposing domain-specific
constraints on the learned function. This builds an ideal foundation for
NeuralODEs as a disease progression model. With the mentioned advantages,
NeuralODEs are steadily finding an application in learning disease
progression. Hao et al. [26] proposed a NeuralODE based diffusion model,
conditioned on individual’s connectivity, for Amyloid pathology progression,
which is a primary pathological event in the Alzheimer’s disease. De Brouwer
et al. [27] used their previously proposed GRU-ODE-Bayes model [28] for
prediction of disability progression in Multiple Sclerosis patients in a two-
year interval using a set of biomarkers.

###  1.3 Contribution

All the mentioned methods for DPM share at least one of the drawbacks: they
rely on hand-crafted features; they are inherently non-interpretable; they are
designed for fixed forecasting intervals; they are time-variant systems, as
they can produce different results for the same forecasting interval. In
contrast, we introduce a framework for pixel-level DPM based on NeuralODEs[25]
that addresses the above-mentioned drawbacks. To the best of our knowledge,
this is the first work that incorporated NeuralODEs [25] for the disease
progression and corresponding target anatomy segmentation. The proposed
method: (1) extracts the relevant features directly from the image and
provides progression estimate with corresponding segmentation masks at a
future time-point for reliability and feasibility assurance; (2) learns from
data with missing visits and variable visit intervals; (3) interprets disease
progression as autonomous ODE and introduces methods to incorporate prior
domain-knowledge. (4) introduces temporal Dice loss and highlights a
successful application of NeuralODE as shown by our large-scale evaluation

We trained the proposed method using two large-scale datasets and showed that
the method learns accurate segmentation maps and produces meaningful
progression estimates. We demonstrated the efficacy of the proposed method on
a private longitudinal GA dataset with 967 scans of 100 patients, and two
public AD datasets with 2537 and 286 scans of 633 and 143 patients each, far
larger than in related works. We used one of the public dataset, TADPOLE
D3-D4, as a hold-out test set and compared our results with those of the
participants in the TADPOLE challenge[19], achieving state-of-the-art
prediction performance from the single baseline scan.

##  2 Method

We formulate the disease progression as an Initial Value Problem comprising an
ordinary differential equation with an initial condition which specifies the
value of the unknown function at baseline. The initial conditions take the
form of the embedding vectors which are generated for each pixel of the
baseline image using any segmentation architecture such as UNet[23]. The ODE
system defined on the embedding and logit space is numerically solved for the
given time points in the future. The obtained logits corresponding to the
future time point are then converted to the segmentation maps (Fig. 2) of the
target anatomy. We utilize the method introduced in [25] for training which
requires a constant amount of memory independent of the solver’s step size and
integration time.

Figure 2: Illustration of the proposed method. The input image at the baseline
visit is first processed with a CNN to extract the feature embeddings
$\xi_{0}$ for each pixel. The feature embedding $\xi_{0}$ acts as the initial
condition for an ODE which evolves it over time to form a trajectory for each
pixel in the phase space. The solution trajectory is shown on a phase plot of
Eq. 2 projected on 2D plane using PCA, where blue represents completely
healthy pixels, and red represents pixels with GA. The computed solutions
$\xi_{T}$ are projected with a linear layer and sigmoid activation to obtain
the output segmentation map $f_{T}$. The abstract procedure of numerical
integration is presented at the bottom of the figure. The learnable function
$p(\xi)$ representing time derivative of $xi$ takes form of a small CNN with 3
convolutional blocks. The computed embeddings $\xi_{t}$ at intermediate
timepoint $t$ and function $p(\xi)$ forms an input to the black-box solver,
which iteratively updates $\xi_{t}$.

###  2.1 Basic ODE model

We aim to learn a continuous spatio-temporal model of disease progression from
a longitudinal dataset comprising a series of $n$-dimensional images (3D in
our case) for each subject, which are acquired across multiple visits at
irregular time-intervals. The time-point $t\in\mathbb{R}^{+}$ of each visit is
normalized to a scale of 1 year. Thus, $t=0.5$ and $t=1.5$ corresponds to the
6 and 18 month time-points respectively from the baseline visit at $t=0$. We
represent the entire stack of the longitudinal volumetric data for each
patient $i$ by a continuous function
$x_{i}:\Omega\times\mathbb{R}^{+}\rightarrow\mathbb{R}$ such that
$x_{i}(\omega;t$) is the intensity value at the spatial coordinate
$\omega\in\Omega$ in the image acquired at time $t$. The entire image of the
patient acquired at time $t$ is denoted by $x(\cdot,t)$. The ground-truth
segmentation masks for each image in the training set can be similarly
represented by a continuous ground-truth function
$g:\Omega\times\mathbb{R}^{+}\rightarrow[0,1]$, that corresponds to the state
of the target biomarker. The continuous functions $x$ and $g$ are not directly
observable but sampled at specific time-points $t$ corresponding to the visits
during which the medical images were actually acquired. Our objective is to
learn a continuous function $f:\Omega\times\mathbb{R}^{+}\rightarrow[0,1]$ to
approximate $g$ using a Deep Learning (DL) framework inspired from the
NeuralODE [25]. The proposed model attempts to approximate $f$ using only a
single image $x(\cdot,0)$ of the baseline visit as input. The estimated
$f(\cdot,t)$ can then be sampled at arbitrary future time-points to obtain an
estimate of the target anatomy segmentation in the future.

First, the image voxels of $x_{i}(.,0)$ are projected onto the latent
embedding in the form of convolutional feature maps $\xi_{0}$ using a CNN with
learnable parameters $\Theta_{\mathrm{CNN}}$. In this work, the 3D-UNet [29]
was used for MRI and the 3D$\rightarrow$2D PSC-UNet [30] was employed for OCT.
Subsequently, the evolution of $\xi_{t}$ over time is modeled as an Initial
Value Problem (IVP) with the differential equation
$\frac{\partial\xi}{\partial t}=p(\xi)$ and the initial condition
$\xi=\xi_{0}$. The $p(\xi)$ is modeled using another CNN with three
convolutional layers (see Fig 2 for details) with parameters
$\Theta_{\mathrm{ODE}}$. This IVP can be numerically solved to obtain the
latent feature embedding for any arbitrary future target time-point $t$ and
the corresponding evolving segmentation map is modeled as a linear
transformation at each pixel (or voxel) of the 2D (or 3D) feature embedding.
Thus, $f(\cdot;t)=\sigma(\phi_{t})$, where $\phi_{t}=W^{T}\xi_{t}+b$
represents the unnormalized logits before applying the sigmoid activation to
the linear transformation with the learnable parameters $W$ and $b$. During
training, the backpropagation is performed as standard except for
$\frac{\partial\xi_{t}}{\partial\Theta_{\mathrm{ODE}}}$ and
$\frac{\partial\xi_{t}}{\partial\xi_{0}}$, whose computations are described in
[25]. The training and inference are summarized in the Algorithms 1 and 2,
where ODE is one of the systems discussed below, ODESolve is a black box ODE
solver, and $\Theta_{\mathrm{ODE}}$ are the learnable weights of the ODE.

Algorithm 1 Forward pass, $f(\cdot,t,x(\cdot,0),\Theta_{CNN},\Theta_{ODE},W)$

0: Baseline image $g(\cdot,0)$, target time $t$, parameters $\Theta$ and $W$

  * •

Compute embeddings
$\xi_{0}~{}=\mathrm{CNN}(x(\cdot,0),\Theta_{\mathrm{CNN}})$,

  * •

Solve ODE,
$\xi_{t}=\text{ODESolve}(\xi_{0},\text{ODE},t,\Theta_{\mathrm{ODE}})$,
$\phi_{t}=W^{T}\xi_{t}+b$

  * •

Compute the final segmentation maps with sigmoid, $\sigma(\phi_{t})$

Algorithm 2 Backward pass

0: Baseline image $x(\cdot,0)$, annotations $g(\cdot,0)$ and $g(\cdot,t)$,
target time $t$, parameters $\Theta_{CNN}$, $\Theta_{ODE}$ and $W$

  * •

Compute $f(\cdot,t,x(\cdot,0),\Theta_{CNN},\Theta_{ODE},W)$ with forward pass,

  * •

Compute loss $L(f,g)$,

  * •

Compute $\frac{\partial L}{\partial\xi_{t}}$ and $\frac{\partial L}{\partial
W}$ as regular, compute $\frac{\partial\xi_{t}}{\partial\xi_{0}}$ and
$\frac{\partial\xi_{t}}{\partial\Theta_{\mathrm{ODE}}}$ as in [25], compute
$\frac{\partial\xi_{0}}{\partial\Theta_{\mathrm{CNN}}}$ as regular,

  * •

perform an update of $\Theta_{\mathrm{ODE}}$, $\Theta_{\mathrm{CNN}}$ and $W$.

###  2.2 Variations of the proposed ODE

The proposed IVP based formulation for DPM gives us a higher degree of control
over the evolution of the disease. It allows the incorporation of additional
domain-specific anatomical constraints by employing different systems of ODE
in the IVP. In certain medical conditions, once a pathology occurs at a
spatial location, the damaged tissue can never be recovered. In such cases,
for example in GA, the target biomarker can only increase in size over time.
We first summarize the basic IVP formulation without constraints as described
previously. Next, we propose two modifications to the ODE system to enforce
the non-decreasing constraints on the size of the target anatomy progression.

  * •

No constraints: This simple approach does not impose any constraints on the
way the biomarker segmentation evolves over time. In this case $\phi_{t}$ is
obtained as a linear transformation of the embedding feature space $\xi$ that
evolves over time as defined by the ODE

$\displaystyle\dot{\xi}=p(\xi),$ (1)

where the function $p(\xi)$ is modeled using a learnable CNN to predict
$\dot{\xi}$, which is a shorthand notation for the time derivative
$\frac{\partial\xi}{\partial t}$. The ODE is assumed to be autonomous or time-
invariant implying that the function $p(\xi)$ is independent of time, ie., the
CNN neither takes $t$ as input, nor does its network parameters change with
time.

  * •

Constraints on the embeddings:  The non-decreasing target anatomy constraint
can be incorporated into our ODE system by enforcing a non-negativity
constraint on the derivative of $\phi$ with respect to time, ie.,
$\dot{\phi}\geq 0$. Given the linear dependency of the logits
$\phi=W^{T}\xi+b$, we need $\dot{\phi}=W^{T}\dot{\xi}=W^{T}p(\xi)$ to be non-
negative. Considering the unit vector $\hat{W}=\frac{W}{||W||}$, we can
decompose $p(\xi)$ into orthogonal and collinear components
$p(\xi)=\hat{W^{T}}p(\xi)\hat{W}+(p(\xi)-\hat{W^{T}}p(\xi)\hat{W}$). We
explicitly constrain the projection of $p(\xi)$ on $\hat{W}$ in the collinear
term to be non-negative by clipping its negative values to 0 using a rectified
unit function denoted by $|\cdot|^{+}$. Thus, the proposed ODE becomes

$\displaystyle\dot{\xi}=|\hat{W}^{T}p(\xi)|^{+}\hat{W}+p(\xi)-\hat{W}^{T}p(\xi)\hat{W}.$
(2)

  * •

Constraints on the logits: In contrast to the above, where $\dot{\phi}$ and
$\dot{\xi}$ had a linear dependency, we want to decouple these two variables
and introduce non-linear dependency between them. Thus, we introduce the time
derivative of $\phi$ as a learnable function $q(\xi,\phi)$, and apply explicit
non-negative constraint by the means of rectified unit function. In this case,
we introduce the following system:

$\displaystyle\begin{cases}\dot{\xi}=p(\xi,\phi),\\\
\dot{\phi}=|q(\xi,\phi)|^{+},\end{cases}$ (3)

where we used softplus as rectified unit function $|\cdot|^{+}$.

In the implementation of $p(\xi)$, we used two consecutive residual blocks
(Fig. 2) consisting of instance normalization [31], convolution with filter
size $3\times 3$ and softplus activation followed by $1\times 1$ convolution
(for 2D, with natural extension to 3D case). For implementing Eq. 3, we
concatenate $\xi$ and $\phi$ and implement $q$ identically to $p$.

###  2.3 Loss function

The proposed method is trained by obtaining $f(\cdot,t)$ from the ODE at
specific discrete time points $\\{{0\leq t\leq T\\}}$, corresponding to the
actual patient visits for which the GT $g(\cdot,t)$ are available. The total
loss at a subject-level is then be defined as the mean of the segmentation
losses across all visits. The binary cross-entropy could be used at each time
step:

$\displaystyle\bar{D}(f,g)=-\sum_{t=0}^{T}\sum_{\omega\in\Omega}{g(\omega,t)\log(f(\omega,t))}+$
$\displaystyle(1-g(\omega,t))\log(1-f(\omega,t))$ (4)

Another alternative is the soft Dice loss [32]:

$\displaystyle\bar{D}(f,g)=1-\frac{1}{T}\sum_{t=0}^{T}\frac{2\sum_{\omega\in\Omega}{f(\omega,t)g(\omega,t)}}{\sum_{\omega\in\Omega}{f^{2}(\omega,t)}+\sum_{\omega\in\Omega}{g^{2}(\omega,t)}}$
(5)

Below, we approach the problem of the temporal loss using functional analysis
and derive the temporal Dice loss function. We define
$f,g:\Omega\times\mathbb{R}^{+}\rightarrow[0;1]$, depending on time, and
assume that $f$ and $g$ are continuous almost everywhere on
$\Omega\times\mathbb{R}^{+}$, and that both integrals of $f$, $g$ and their
squares are finite. In turn, the space of such square integrable functions
$F=\\{f:\Omega\times\mathbb{R}^{+}\rightarrow[0;1]\\}$ with inner product
defined as
$(f,g)=\int_{\Omega\times\mathbb{R}^{+}}{f(\omega,t)g(\omega,t)d\omega dt}$ is
a Hilbert space.

We want our method to learn the function $f$ by minimizing the respective
squared distance $\lVert f-g\rVert^{2}$.

$\lVert f-g\rVert^{2}=\lVert f\rVert^{2}+\lVert g\rVert^{2}-2(f,g)\leq\lVert
f\rVert^{2}+\lVert g\rVert^{2},$

due to non-negativity of $(f,g)$.

The minimum of $\lVert f-g\rVert^{2}$ is achieved at $f\sim g$: $f$ and $g$
being from the same equivalence class. Introducing a surrogate function
$D(f,g)$, we limit it’s the upper bound to 1: $D(f,g)=\frac{\lVert
f-g\rVert^{2}}{\lVert f\rVert^{2}+\lVert g\rVert^{2}}$. We can show that
$D(f,g)$ has a minimum at $f\sim g$. As a result, by finding the minimum of
$D(f,g)$ we also minimize the distance $\lVert f-g\rVert^{2}$. Next, we
discretize the space $\Omega$ and time $T$, thus getting the definition of
Soft Dice loss [32] from $D(f,g)$ with temporal objective:

$\displaystyle\bar{D}(f,g)=$ $\displaystyle
1-\frac{2\sum_{t=0}^{T}\sum_{\omega\in\Omega}{f(\omega,t)g(\omega,t)}}{\sum_{t=0}^{T}\sum_{\omega\in\Omega}{f^{2}(\omega,t)}+\sum_{t=0}^{T}\sum_{\omega\in\Omega}{g^{2}(\omega,t)}}.$
(6)

Effectively, this implies that we should compute Dice loss not for individual
visits independently, but for the entire stack of visits. In the experiments,
as part of ablation, we compare the temporal Dice formulation with the other
two loss functions and study their influence on the performance.

##  3 Evaluation

###  3.1 Data

Table 1 provides an overview of the three datasets used to develop and
evaluate the proposed method. The private MUV-GA dataset is used to evaluate
the performance on the task of predicting the segmentation maps for GA at
future time-points from a retinal OCT of the first baseline visit. The two
public datasets TADPOLE [11] and ADNI1 are employed for the task of predicting
the changes in the ventricle volume over time, which is an important biomarker
to track the AD progression.

Table 1: An overview of the longitudinal Datasets.

Dataset |  Scan |  #Patients |  Total #scans |  #Scans per patient |  Avg. time span (days)
---|---|---|---|---|---
MUV-GA |  OCT | 100 | 967 | 5.95 | 970
ADNI | MRI | 633 | 2537 | 4 | 704
TADPOLE D3-D4 | MRI | 143 | 286 | 2 | 1082

####  3.1.1 MUV-GA

is a longitudinal in-house dataset from the Department of Ophthalmology at
Medical University of Vienna [33] which originated from an observational study
of natural GA progression. It consists of 3D OCT scans of eyes of patients
older than 50 years, diagnosed with GA, imaged with Spectralis (Heidelberg
Engineering) device. Study visits occurred every 3 months. The patients were
diagnosed with GA secondary to AMD, and were excluded from the study when
other confounding pathologies were found. It consists of 3D retinal OCT scans
of subjects older than 50 years who have been diagnosed with GA. Longitudinal
scans were acquired for each subject at 3 month visit intervals using the
Spectralis (Heidelberg Engineering) device. The Ground Truth pixel-level
annotations of GA lesion were performed manually by a clinical expert at all
available time-points. The dataset features both missing visits and irregular
time intervals between visits. While the OCT scans are 3D, the reference
standard for the GA segmentation is available in 2D, corresponding to the en-
face lesion as depicted in Fig. 1. Thus, the CNN architectures that can be
employed for this task are limited to 3D$\rightarrow$2D networks, such as
[30], which reduce the dimensionality of the output predictions.

##### Preprocesing

First, we resample the scans to have 5.671$\mu\text{m}$ spacing across the
BScans. Then, we register the follow-up visits. To control and fix the
alignment of the scans and corresponding annotations between visits, we first
perform retinal blood vessel segmentation [30], and then use phase correlation
algorithm [34] on the blood vessels to estimate the shift between the given
image and the baseline scan. After the registration of the follow-up visits is
done, we run layer segmentation algorithm[35] and flatten the retinal
appearance according to the outer RPE boundary. Flattening helps the methods
to focus on retinal layer appearance and structure, rather than spatial
position and inclination.

####  3.1.2 ADNI

ADNI1111adni.loni.usc.edu is an open-access dataset - part of the initiative
launched in 2003 as a public-private partnership, led by Principal
Investigator Michael W. Weiner, MD. The dataset was created for measuring the
progression of mild cognitive impairment and early Alzheimer’s Disease. The
primary goal of ADNI has been to test whether serial MRI, positron emission
tomography (PET), other biological markers, and clinical and
neuropsychological assessment can be combined to measure the progression of
mild cognitive impairment (MCI) and early Alzheimer’s disease.

TADPOLE[11] is a challenge guided towards the accurate prediction of AD
progression. It aims to predict AD progression through estimation of a number
of biomarker values. In this work, we were specifically interested in the
imaging biomarkers, thus we focused on the prediction of the ventricle volume.
The TADPOLE challenge has a series of datasets, derived from ADNI database. In
this study, we use only D3 and D4 datasets, designed for making predictions
from just a single baseline scan, where D3 contains a single visit for a
patient, and D4 plays a role of the test set with the follow-up measurements.
We use TADPOLE data exclusively as a hold-out test-set and train on ADNI1
data.

As the proposed method requires imaging data to train, we query T1 MRI scans
of the patients participating in ADNI1 and having at least 2 visits. Due to
the absence of the manual ground truth ventricle annotation, we obtain pseudo-
label segmentations automatically using FastSurfer [36] at all available time
points following the protocols used in the TADPOLE challenge. We then trained
the baseline methods and the proposed method using such pseudo-label
segmentations, and test using TADPOLE D3-D4 data. Both the input MRI images
and the corresponding segmentation labels are in 3D. Similar to MUV-GA,
TADPOLE has missing visits and irregular time intervals between visits.

##### Preprocessing

The MRI scans were first imported to the FreeSurfer, normalized[37] and skull-
stripped[38]. Then, FastSurfer [36] was used to process the scans and segment
the brain structures. Follow-up visits of the same patients were aligned using
longitudinal stream of FreeSurfer [39]. Since we explicitly focus on the
evolution of the ventricle volume, we discard all other labels. The
preprocessed scans have the size of $256\times 256\times 256$. To reduce the
memory burden for the storage and data processing, we additionally crop the
cube of size $128\times 128\times 128$. We center the cropped region on the
ventricles of the baseline visit. The follow-up scans are cropped with the
region from the baseline.

###  3.2 Experiments

####  3.2.1 Data split

To assess the performance of the proposed method, we run a series of
experiment for each dataset. In our experiments, we utilize 5-Fold Cross
Validation (CV) and split VSC GA and ADNI datasets into five groups by patient
id, where four folds are used as a train-val dataset, and one fold is used as
a test dataset. We use TADPOLE D3-D4 exclusively as a hold-out test set, and
those scans don’t overlap with ADNI1.

####  3.2.2 Baselines

Since the datasets represent different tasks, we compare the proposed method
with different baselines depending on the application.

##### Oracle

Each of the datasets has a corresponding Oracle model as an estimate of the
upper bound that a predictive model can achieve. To obtain an Oracle model, we
directly segmented both the baseline and the endpoint scan. For this, we
trained PSC-UNet [30] for automated GA segmentation, and UNet3D [29] to
perform an automated segmentation of the brain ventricles.

##### Geographic atrophy

We compare the proposed method against the approach introduced in Zhang et al.
referred as MS-DCNN[14], Nguyen et al. [13], Gigon et al.[16] and the Oracle
model. The nature of the task of predicting growth on 2D en-face masks from 3D
volumetric data (3D $\longrightarrow$ 2D) poses an additional challenge and
limits a set of the baseline methods that can be utilized. The performance of
progression prediction is evaluated at the follow-up visit, one year from the
baseline scan.

##### ADNI & TADPOLE

For ADNI data, we compare the proposed method with Nguyen et al. [13],
Petersen et al. Deep Glioma Growth (DGG)[22] and Oracle model. The prediction
performance is evaluated at the 3-year follow-up visit. For TADPOLE D3-D4, we
additionally include the methods from TADPOLE challenge in comparison, namely
GlassFrog, BravoLab, EMC1 [19].

####  3.2.3 Ablation experiments

We ran a series of ablations experiments, to investigate:

##### Constraints in the ODE system

we train the proposed method with three different ODE systems, one introducing
no constraints (Eq. 1), one constraining the embedding space (Eq. 2), and one
constraining the logits space (Eq. 3). For this ablation experiment, we
employed only MUV-GA dataset, because GA can only grow in size as the atrophic
tissue cannot recover, which is not necessarily the case with brain
ventricles. As a loss, we employed the function described in Eq. 2.3.

##### Loss functions

for learning temporal segmentation objective we investigate the performance of
the loss function derived from functional distance (Eq. 2.3), compared to the
alternative standard loss functions (Eq. 2.3, and 5). We perform a 5-fold
cross validation experiments using the proposed method on both MUV-GA and ADNI
datasets, and further apply the models trained on ADNI to TADPOLE hold-out
test set. For MUV-GA dataset, we employed the ODE system with constrains
defined by Eq. 2. For ADNI data, we used ODE system described by Eq. 1.

####  3.2.4 Evaluation metrics

For cross-validation experiments on MUV-GA and ADNI datasets, we report: Dice
index of baseline and predicted visit, Hausdorff distance of the baseline and
the last visit, coefficient of determination $R^{2}$ and coefficient of
correlation Pearson’s $r$ of the change in area or volume. For MUV-GA, we
additionally report the mean squared error (MSE) of the change in square root
transformed area in mm. For ADNI and TADPOLE, we report the metrics utilized
in TADPOLE[11] challenge. Namely, MAE, WES, calculated as percentage of
intracranial volume, and CPA:

$MAE=\frac{1}{N}\sum_{i=1}^{N}{\frac{1}{ICV_{i}}|V_{i}-\bar{V}_{i}|},$
$WES=\frac{\sum_{i=1}^{N}{\frac{1}{\bar{V}_{i}^{0.75}-\bar{V}_{i}^{0.25}}|V_{i}-\bar{V}_{i}|}}{\sum_{i=1}^{N}{\frac{ICV_{i}}{\bar{V}_{i}^{0.75}-\bar{V}_{i}^{0.25}}}},$
$CPA=|0.5-\frac{1}{N}\sum_{i=1}^{N}{1_{[\bar{V}_{i}^{0.25};\bar{V}_{i}^{0.75}]}(V_{i})}|,$

where $V$ and $\bar{V}$ are predicted ventricle volumes, $\bar{V}_{i}^{0.25}$
and $\bar{V}_{i}^{0.75}$ are 25% and 75% percentile of the predicted volume of
an individual patient, together forming 50% CI, $ICV$ is the intracranial
volume, and 1 is an indicator function.

Table 2: Cross-validation experiments on MUV-GA dataset. The metrics were computed with the respect to the baseline and the 1-year visits. In contrast to the other methods, the Oracle model does not perform a future prediction but employs the Scan at the 1 year visit to highlight an upper limit of the performance that can be achieved.  Method | Dice0 | Dicet | HD${}^{95}_{0}$ | HD${}^{95}_{t}$ | $R^{2}$ | $r$ | MAE (in mm)
---|---|---|---|---|---|---|---
Oracle: PSC-UNet[30] | $0.82\pm 0.15$ | $0.87\pm 0.13$ | $0.16\pm 0.21$ | $0.10\pm 28$ | 0.73 | 0.86 | $0.46\pm 0.59$
Zhang et al. [14] | $0.62\pm 0.25$ | $0.70\pm 0.21$ | $1.50\pm 1.00$ | $1.05\pm 0.74$ | -0.18 | 0.52 | $1.19\pm 0.96$
Zhang et al. 3D [14] | $0.76\pm 0.21$ | $0.81\pm 0.14$ | $0.78\pm 1.12$ | $0.47\pm 0.69$ | -0.29 | 0.38 | $1.15\pm 1.10$
Nguyen et al. [13] | $\mathbf{0.80\pm 0.17}$ | $\mathbf{0.84\pm 0.11}$ | $\mathbf{0.30\pm 0.54}$ | $\mathbf{0.19\pm 0.29}$ | 0.28 | 0.66 | $0.89\pm 0.80$
Gigon et al. [16] | $0.76\pm 0.22$ | $0.79\pm 0.18$ | $0.57\pm 0.89$ | $0.78\pm 1.12$ | 0.36 | 0.61 | $0.85\pm 0.75$
Proposed | $\mathbf{0.80\pm 0.19}$ | $\mathbf{0.84\pm 0.13}$ | $\mathbf{0.30\pm 0.54}$ | $0.21\pm 0.30$ | 0.59 | 0.78 | $\mathbf{0.64\pm 0.63}$
Table 3: Cross-validation experiments on ADNI dataset. The metrics were computed with the respect to the baseline and the 3-year visits. The Oracle model highlights the predictive ceiling that can be achieved. Method | Dice0 | Dicet | HD${}^{95}_{0}$ | HD${}^{95}_{t}$ | $R^{2}$ | $r$ | MAE, % ICV | WES | CPA
---|---|---|---|---|---|---|---|---|---
Oracle: 3D-UNet [29] | $0.97\pm 0.02$ | $0.97\pm 0.2$ | $1.08\pm 4.20$ | $1.28\pm 4.32$ | 0.98 | 0.99 | $0.06\pm 0.10$ | 0.05 | 0.46
Nguyen et al.[13] | $0.96\pm 0.02$ | $0.93\pm 0.04$ | $1.46\pm 4.31$ | $1.90\pm 5.07$ | 0.32 | 0.58 | $0.24\pm 0.20$ | 0.23 | 0.04
Petersen et al.[22], mean sample | $0.95\pm 0.02$ | $0.91\pm 0.04$ | $\mathbf{1.10\pm 4.20}$ | $\mathbf{1.38\pm 4.40}$ | -0.11 | 0.34 | $0.24\pm 0.26$ | 0.19 | 0.25
Petersen et al.[22], best sample | $0.95\pm 0.02$ | $0.91\pm 0.04$ | $\mathbf{1.10\pm 4.20}$ | $\mathbf{1.38\pm 4.30}$ | -0.09 | 0.35 | $0.24\pm 0.26$ | 0.18 | 0.25
Proposed | $\mathbf{0.97\pm 0.02}$ | $\mathbf{0.94\pm 0.03}$ | $1.39\pm 4.21$ | $1.63\pm 4.63$ | 0.36 | 0.61 | $\mathbf{0.21\pm 0.19}$ | 0.20 | 0.08

###  3.3 Training details

Each dataset was split into five folds. For each method, we performed a brief
hyperparameter optimization based on the additional validation subsplit of the
training folds. We chose between ADAM and SGD optimizers and a number of
learning rates. For the development We used PyTorch framework, and we trained
all the models using NVidia 2080Ti GPU. Training of each model took less than
24h.

##### Proposed method

We used the function defined by Eq. 2.3 as a loss function. For solving the
ODE system, we used RK4 scheme with step size $h=\frac{1}{12}$ equivalent to a
one-month interval. For optimization, we used ADAM optimizer with default
parameters for all 3 variants of the proposed method. For NeuralODE part, we
used the code provided by Chen et
al.222https://github.com/rtqichen/torchdiffeq During training the proposed
method takes a constant amount of GPU memory, independent of the time-interval
over which the ODE is evolved by using the adjoint method [25] for
backpropagation. This allowed our method to be trained over large time
intervals with limited GPU memory.

##### Baseline methods

For MS-DCNN [14], we removed dropout and added InstanceNorm [31], which
greatly improved training stability and performance. We also extended the
method of Zhang et al. to handle 3D OCT volumes instead of en-face
projections, it is referred to as MS-DCNN 3D. We extended the approach of
Nguyen et al. [13] to handle the imaging data instead of tabular features,
introducing convolutional filters into the recurrent blocks. The work of Gigon
et al. [16] was implemented as in the paper, however, we used [35] for layer
segmentation. For implementation of the Deep Glioma Growth[22], we used code
provided by Petersen et al.333https://github.com/MIC-DKFZ/image-time-series,
with the exception that only baseline image and its segmentation was used for
generating the prior.

##  4 Results

Table 4: Results for TADPOLE D4 test set, evaluated using a single baseline scan. Method | Dice0 | HD${}^{95}_{0}$ |  MAE, % ICV | WES | CPA
---|---|---|---|---|---
Nguyen [13] | $\mathbf{0.96\pm 0.01}$ | $1.07\pm 0.46$ | $0.46$ | 0.39 | 0.02
Petersen [22]. mean sample | $0.95\pm 0.01$ | $\mathbf{1.0\pm 0.4}$ | $0.56$ | 0.4 | 0.2
Proposed | $\mathbf{0.96\pm 0.01}$ | $1.05\pm 0.44$ | $\mathbf{0.44}$ | 0.35 | 0.08
GlassFrog- LCMEM -HDR[19] | - | - | 0.48 | 0.38 | 0.24
BravoLab[19] | - | - | 0.64 | 0.64 | 0.42
GlassFrog- Average[19] | - | - | 0.68 | 0.55 | 0.24
EMC1- Std[19] | - | - | 0.80 | 0.62 | 0.48
Table 5: Ablation experiments of explicit domain-specific constraints. Eq. 1 corresponds to the system without constraints. To encode that the lesion only grows over time, Eq. 2 uses explicit constrains in the embedding space, while Eq. 3 applies the constraints on the logits. ODE | Dice0 | Dicet | R2 | r | MAE, mm
---|---|---|---|---|---
Eq. (1) | $0.79\pm 0.19$ | $0.83\pm 0.15$ | 0.44 | 0.72 | $0.78\pm 0.69$
Eq. (2) | $\mathbf{0.8\pm 0.19}$ | $\mathbf{0.84\pm 0.13}$ | $\mathbf{0.59}$ | 0.78 | $\mathbf{0.64\pm 0.63}$
Eq. (3) | $\mathbf{0.8\pm 0.19}$ | $\mathbf{0.84\pm 0.14}$ | $0.56$ | 0.76 | $0.68\pm 0.63$
Table 6: Ablation experiments of the proposed method trained with different loss functions. Eq. 2.3 represents binary cross entropy loss, Eq. 5 reflects a soft Dice loss applied on each individual scans and then averaged, Eq. 2.3 represents the Dice loss described in the paper, considering all scans as a single 4D volume. Dataset | MUV-GA, 5-Fold CV | ADNI, 5-Fold CV | TADPOLE D4
---|---|---|---
Loss | Dice0 | Dicet | R2 | r | MAE, mm | Dice0 | Dicet | R2 | r | MAE, % ICV | WES | CPA
Eq. (2.3) | $0.73\pm 0.22$ | $0.77\pm 0.18$ | $0.46$ | 0.68 | $0.71\pm 0.74$ | $0.97\pm 0.02$ | $0.94\pm 0.03$ | $0.36$ | $\mathbf{0.6}$ | $0.46\pm 0.8$ | $0.36$ | $0.12$
Eq. (5) | $0.78\pm 0.18$ | $0.83\pm 0.13$ | $0.58$ | 0.76 | $0.67\pm 0.61$ | $0.97\pm 0.02$ | $0.94\pm 0.03$ | $0.36$ | $\mathbf{0.6}$ | $0.45\pm 0.82$ | $0.36$ | $\mathbf{0.05}$
Eq. (2.3) | $\mathbf{0.8\pm 0.19}$ | $\mathbf{0.84\pm 0.13}$ | $\mathbf{0.59}$ | 0.78 | $\mathbf{0.64\pm 0.63}$ | $0.97\pm 0.02$ | $0.94\pm 0.03$ | $0.36$ | $0.61$ | $\mathbf{0.44\pm 0.83}$ | $\mathbf{0.35}$ | $0.08$

###  4.1 Progression of Geographic Atrophy

A qualitative result of our method is shown in Fig. 1, and the quantitative
cross-validation results are provided in Table 2. The top row represents the
segmentation performed with Oracle PSC-UNet[30] using baseline and first year
images.The proposed method clearly outperformed the baseline methods on all
the metrics with the exception of Nguyen et al. [13], where they both achieved
a mean Dice score of 0.8 and 0.84 for predicting the atrophy segmentation at
baseline and 1-year visits, respectively. However, the proposed method
performed better in estimating the atrophy growth, achieving $R^{2}=0.59$,
$r=0.78$, and MAE $=0.64$ mm, which is very close to the performance of the
Oracle model. The similar performance of the proposed method and the RNN based
solution of [13] may be attributed to the fact that both belong to the same
family with the RNN model being a discrete approximation of the NeuralODE. The
RNN can be represented as NeuralODE with a less sophisticated explicit Euler
method for numerical integration at discrete time-steps.

The recently introduced method of Gigon et al. [16], exploring the idea of
Taylor series decomposition, performs a bit worse both in segmentation and
prediction. The method of Zhang et al. [14] failed in quantitative prediction
of the lesion change, at the same time showing moderate performance in terms
of segmentation.

###  4.2 Progression of Alzheimer’s Disease

A qualitative result of our method for predicting the ventricle segmentation
in brain MRI at future time-points is depicted in Fig. 1. All the models
achieved a high segmentation performance of the baseline visits, with Dice
score ranging from 0.95 of Petersen et al. [22] to 0.97 of the proposed method
(Table 3). It’s important to note that we used automatically generated
annotations by FastSurfer, which potentially might not fully reflect the
complexity of the ventricle anatomy and might be easy to learn by automated
methods. The proposed method achieved the best performance in predicting
segmentation at a 3-year visit in terms of Dice, $R^{2}$, $r$ and MAE. The
other methods demonstrated similarly good results in baseline and 3-year HD95,
WES and CPA. As Petersen et al. claim, the Continuous-Time Deep Glioma Growth
model wasn’t designed for precise prediction of the changes, but rather for
their probabilistic description. However, the Deep Glioma Growth model
achieved good results in predicting the changes, with mean prediction being
close to the best sample from the predicted segmentation masks distribution
w.r.t 3-year Dice score.

The TADPOLE hold-out test set consists of on average longer forecasting
intervals, testing the extrapolation capabilities of the methods.
Quantitatively (Table 4), the proposed and baseline models demonstrated good
generalization in baseline visit segmentation, achieving Dice score of 0.95 or
above and showing a mean HD95 smaller than 1.07 mm. Importantly, the proposed
method, as in the cross-validation experiments, excelled in accurate disease
evolution prediction, surpassing the challenge contenders in all metrics. It
clearly outperformed the TADPOLE participants, in the prediction of the future
ventricle size, achieving 0.44 MAE, 0.35 WES, and demonstrating a competitive
0.02 CPA. The method of Petersen et al. [22] showed competitive predictive
performance, where we sampled the mean prediction, since we were interested in
a single prediction output. Overall, all of the imaging-based method
demonstrated competitive results, with the proposed method surpassing the
challenge contenders in all metrics.

###  4.3 Ablation experiments.

To investigate the effect of the explicit constraint enforcement as a part of
the ODE system, we compared three different ODE systems defined by Eq. 1-3
(Table 5) on the MUV-GA dataset. The effect of the constrained ODE system was
slightly pronounced on the segmentation performance of the baseline and 1-year
scans as both ODE systems with constraints (Eq. 2 and 3) performed slightly
better than the unconstrained system. However, for the atrophy growth
estimation, the effect of the constraints was more pronounced, where the
system with the constraints in the embedding space (Eq. 2) performed clearly
the best.

The effect of different loss functions on the predictive performance is
reported in Table 6. In all experiments, binary cross-entropy function
performed the worst, except in ADNI cross-validation, where all loss functions
performed equally well. This can be attributed to the nature of the dataset’s
reference standard, which was produced by an automated segmentation method.
Both Dice-based losses performed well, however, the loss we derived from
functional distance (Eq. 2.3), performed better in GA.

##  5 Conclusion

In this work, we introduced a DPM to predict the future progression of the
morphological changes in pathologies (such as GA in retinal OCT) and
clinically relevant anatomical structures (such as the ventricles in brain
MRI) from a single scan of the baseline visit. Our method produces the
segmentation maps for the desired object of interest for varying future time-
points which is clinically useful in the early detection of the fast
progressors for personalized treatment and the recruitment of patients in
clinical trials.

In contrast to the existing methods, the proposed method employs a continuous
time-invariant Neural-ODE formulation to constrain the solution space and can
be trained on datasets with missing data and variable time intervals between
the visits. We demonstrated two approaches to constrain the solution to
physiologically plausible ones and introduced temporal Dice loss for training
with longitudinal data. The state of the art performance of our method on the
two distinct tasks of predicting the progression of GA in retinal OCT and the
ventricles in brain MRI demonstrates its efficacy and generalizability across
different anatomical structures and imaging modalities.

We focused on the task of providing predictions from a single scan, targeting
the scenario of patient recruitment for clinical trials, where typically only
a baseline scan is available. In addition to enriching for fast progressors,
our method could serve as a control digital twin of a patient receiving a
treatment to predict the expected outcomes without the treatment, thereby
reducing the need for costly sham arms in clinical trials. Further extension
of this work to utilize multiple prior scans and update the predictions as the
new data becomes available offers a promising direction for future work.

## References

  * [1] J. S. Barrett, T. Nicholas, K. Azer, and B. W. Corrigan, “Role of Disease Progression Models in Drug Development,” _Pharmaceutical Research_ , 2022\.
  * [2] N. Carrillo et al., “Safety and efficacy of N-acetylmannosamine (ManNAc) in patients with GNE myopathy: an open-label phase 2 study,” _Genetics in Medicine_ , vol. 23, no. 11, pp. 2067–2075, 2021.
  * [3] C. Hu, “Exposure-response modeling of clinical end points using latent variable indirect response models.” _CPT: pharmacometrics & systems pharmacology_, vol. 3, no. 6, p. e117, jun 2014.
  * [4] W. L. Wong, X. Su, X. Li, C. M. G. Cheung, R. Klein, C.-Y. Cheng, and T. Y. Wong, “Global prevalence of age-related macular degeneration and disease burden projection for 2020 and 2040: a systematic review and meta-analysis,” _The Lancet Global Health_ , vol. 2, no. 2, pp. e106–e116, 2014.
  * [5] R. H. Guymer et al., “Incomplete retinal pigment epithelial and outer retinal atrophy in age-related macular degeneration: Classification of atrophy meeting report 4,” _Ophthalmology_ , vol. 127, no. 3, pp. 394–409, 2020\.
  * [6] D. S. Liao et al., “Complement C3 Inhibitor Pegcetacoplan for Geographic Atrophy Secondary to Age-Related Macular Degeneration: A Randomized Phase 2 Trial.” _Ophthalmology_ , vol. 127, no. 2, pp. 186–195, feb 2020.
  * [7] G. J. Jaffe et al., “C5 Inhibitor Avacincaptad Pegol for Geographic Atrophy Due to Age-Related Macular Degeneration: A Randomized Pivotal Phase 2/3 Trial,” _Ophthalmology_ , vol. 128, no. 4, pp. 576–586, apr 2021.
  * [8] Z. Breijyeh and R. Karaman, “Comprehensive Review on Alzheimer’s Disease: Causes and Treatment.” _Molecules (Basel, Switzerland)_ , vol. 25, no. 24, dec 2020.
  * [9] B. R. Ott et al., “Brain ventricular volume and cerebrospinal fluid biomarkers of Alzheimer’s disease,” _Journal of Alzheimer’s disease : JAD_ , vol. 20, no. 2, pp. 647–657, 2010.
  * [10] S. M. Nestor et al., “Ventricular enlargement as a possible measure of Alzheimer’s disease progression validated using the Alzheimer’s disease neuroimaging initiative database,” _Brain_ , vol. 131, no. 9, pp. 2443–2454, sep 2008.
  * [11] R. V. Marinescu et al., “TADPOLE Challenge: Accurate Alzheimer’s disease prediction through crowdsourced forecasting of future data.” _PRedictive Intelligence in MEdicine. PRIME (Workshop)_ , vol. 11843, pp. 1–10, oct 2019.
  * [12] T. Wang, R. G. Qiu, and M. Yu, “Predictive Modeling of the Progression of Alzheimer’s Disease with Recurrent Neural Networks,” _Scientific Reports_ , vol. 8, no. 1, p. 9161, 2018.
  * [13] M. Nguyen, T. He, L. An, D. C. Alexander, J. Feng, and B. T. T. Yeo, “Predicting Alzheimer’s disease progression using deep recurrent neural networks,” _NeuroImage_ , vol. 222, p. 117203, 2020.
  * [14] Y. Zhang, Z. Ji, S. Niu, T. Leng, D. L. Rubin, and Q. Chen, “A Multi-Scale Deep Convolutional Neural Network For Joint Segmentation And Prediction Of Geographic Atrophy In SD-OCT Images,” in _ISBI 2019_ , 2019, pp. 565–568.
  * [15] Y. Zhang et al., “An integrated time adaptive geographic atrophy prediction model for sd-oct images,” _Medical Image Analysis_ , vol. 68, p. 101893, 2021\.
  * [16] A. Gigon et. al, “Personalized Atrophy Risk Mapping in Age-Related Macular Degeneration,” _Translational Vision Science & Technology_, vol. 10, no. 13, p. 18, nov 2021.
  * [17] S. Hochreiter and J. Schmidhuber, “Long short-term memory,” _Neural Comput._ , vol. 9, no. 8, p. 1735–1780, Nov. 1997.
  * [18] M. Chen, “Minimalrnn: Toward more interpretable and trainable recurrent neural networks,” _ArXiv_ , vol. abs/1711.06788, 2017.
  * [19] R. V. Marinescu et al., “The alzheimer’s disease prediction of longitudinal evolution (tadpole) challenge: Results after 1 year follow-up,” _Machine Learning for Biomedical Imaging_ , vol. 1, 2021.
  * [20] I. Ezhov et al., “Neural parameters estimation for brain tumor growth modeling,” in _MICCAI_ , 2019.
  * [21] J. Petersen et al., “Deep Probabilistic Modeling of Glioma Growth,” in _MICCAI_ , 2019.
  * [22] J. Petersen et al., “Continuous-time deep glioma growth models,” in _MICCAI 2021_. Springer International Publishing, 2021, pp. 83–92.
  * [23] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in _MICCAI 2015_ , Cham, 2015, pp. 234–241.
  * [24] A. Vaswani et al., “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [25] T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. Duvenaud, “Neural Ordinary Differential Equations,” in _NeurIPS_ , 2018.
  * [26] W. Hao et al., “Learning Amyloid Pathology Progression from Longitudinal PIB-PET Images in Preclinical Alzheimer’s Disease,” in _ISBI 2020_ , apr 2020, pp. 572–576.
  * [27] E. De Brouwer et al., “Longitudinal modeling of MS patient trajectories improves predictions of disability progression,” _Computer Methods and Programs in Biomedicine_ , vol. 208, p. 106180, 2021.
  * [28] E. D. Brouwer, J. Simm, A. Arany, and Y. Moreau, “GRU-ODE-Bayes: Continuous modeling of sporadically-observed time series,” in _NeurIPS_ , 2019.
  * [29] Ö. Çiçek, A. Abdulkadir, S. Lienkamp, T. Brox, and O. Ronneberger, “3d u-net: Learning dense volumetric segmentation from sparse annotation,” in _MICCAI_ , 2016, pp. 424–432.
  * [30] D. Lachinov, P. Seeböck, J. Mai, F. Goldbach, U. Schmidt-Erfurth, and H. Bogunovic, “Projective skip-connections for segmentation along a subset of dimensions in retinal oct,” in _MICCAI 2021_. Springer International Publishing, 2021, pp. 431–441.
  * [31] D. Ulyanov, A. Vedaldi, and V. Lempitsky, “Instance normalization: The missing ingredient for fast stylization,” _ArXiv_ , vol. abs/1607.08022, 2016.
  * [32] F. Milletari, N. Navab, and S.-A. Ahmadi, “V-net: Fully convolutional neural networks for volumetric medical image segmentation,” _2016 Fourth International Conference on 3D Vision (3DV)_ , pp. 565–571, 2016.
  * [33] P. T. A. Bui et al., “Fundus autofluorescence and optical coherence tomography biomarkers associated with the progression of geographic atrophy secondary to age-related macular degeneration,” _Eye_ , pp. 1–7, aug 2021\.
  * [34] E. De Castro and C. Morandi, “Registration of translated and rotated images using finite fourier transforms,” _IEEE Transactions on Pattern Analysis and Machine Intelligence_ , vol. PAMI-9, no. 5, pp. 700–703, 1987.
  * [35] M. K. Garvin, M. D. Abramoff, X. Wu, S. R. Russell, T. L. Burns, and M. Sonka, “Automated 3-d intraretinal layer segmentation of macular spectral-domain optical coherence tomography images,” _IEEE Transactions on Medical Imaging_ , vol. 28, no. 9, pp. 1436–1447, 2009.
  * [36] L. Henschel, S. Conjeti, S. Estrada, K. Diers, B. Fischl, and M. Reuter, “FastSurfer - A fast and accurate deep learning based neuroimaging pipeline,” _NeuroImage_ , vol. 219, p. 117012, 2020.
  * [37] J. G. Sled, A. P. Zijdenbos, and A. C. Evans, “A nonparametric method for automatic correction of intensity nonuniformity in MRI data,” _IEEE Trans Med Imaging_ , vol. 17, pp. 87–97, 1998.
  * [38] F. Segonne, A. M. Dale, E. Busa, M. Glessner, D. Salat, H. K. Hahn, and B. Fischl, “A hybrid approach to the skull stripping problem in MRI,” _NeuroImage_ , vol. 22, no. 3, pp. 1060–1075, 2004.
  * [39] M. Reuter, H. D. Rosas, and B. Fischl, “Highly Accurate Inverse Consistent Registration: A Robust Approach,” _NeuroImage_ , vol. 53, no. 4, pp. 1181–1196, 2010.


# TLRM: Task-level Relation Module for GNN-based Few-Shot Learning

Yurong Guo1, Zhanyu Ma1∗, Xiaoxu Li2, and Yuan Dong1 ∗ Corresponding author 1
Pattern Recognition and Intelligent System Lab.,
Beijing University of Posts and Telecommunications, Beijing, China  2 Lanzhou
University of Technology, Lanzhou, China

###### Abstract

Recently, graph neural networks (GNNs) have shown powerful ability to handle
few-shot classification problem, which aims at classifying unseen samples when
trained with limited labeled samples per class. GNN-based few-shot learning
architectures mostly replace traditional metric with a learnable GNN. In the
GNN, the nodes are set as the samples’ embedding, and the relationship between
two connected nodes can be obtained by a network, the input of which is the
difference of their embedding features. We consider this method of measuring
relation of samples only models the sample-to-sample relation, while neglects
the specificity of different tasks. That is, this method of measuring relation
does not take the task-level information into account. To this end, we propose
a new relation measure method, namely the task-level relation module (TLRM),
to explicitly model the task-level relation of one sample to all the others.
The proposed module captures the relation representations between nodes by
considering the sample-to-task instead of sample-to-sample embedding features.
We conducted extensive experiments on four benchmark datasets: mini-ImageNet,
tiered-ImageNet, CUB-$200$-$2011$, and CIFAR-FS. Experimental results
demonstrate that the proposed module is effective for GNN-based few-shot
learning.

###### Index Terms:

Few-shot learning, Graph Neural Networks, Task-level Relation

##  I Introduction

Deep Learning has been achieved great success in visual recognition tasks [1,
2, 3, 4], which depends on powerful model and amounts of labelled samples [5].
However, humans can learn new concepts with little examples, or none at all.
The gap motivated researchers to study few-shot learning and zero-shot
learning.

The goal of few-shot learning is to classify unseen samples, given just a
small number of labeled samples in each class. It has attracted considerable
attention [6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]. One promising study is
metric-based few-shot learning [6, 7, 8, 9, 10, 11, 12, 13]. Given just a
query sample and a few labeled support samples, the embedding function
extracts feature for all samples, and then a metric module measures distance
between the query embedding and class embedding to give a recognition result.
Recently, there have some studies of utilizing Graph Neural Networks (GNNs)
[8, 10, 11, 18] to handle the few-shot classification task, which can be seen
as a kind of metric learning method. In GNN-based few-shot learning model, all
embedding features are connected to construct a graph. And each node is
represented by the embedding feature of a sample. Then the graph classifies
the unlabeled query by measuring the similarity between two samples.

Figure 1: The left panel shows a general framework of previous metric
approaches based on GNN. The right panel briefly illustrates our approach.
Figure 2: (a) shows the GNN-based few-shot model. And in (b), the left panel
shows a general framework of our approach used for calculating the similarity
scores, and the right panel shows the task-level relation module.

Even though GNN-based model have made significant advance in few-shot
classification, they do suffer from distinct limitation. In the metric module
of GNN-based methods, relation representation for a pair samples is obtained
by calculating the absolute difference [8, 10, 11, 18]. It only considers the
corresponding embedding features of the samples. Intuitively, pair-wise
relationship is not only dependent on the distance between corresponding
embedding features, but also related to all embedding features in a task. As
shown in the left panel of Figure 1, there is no significant difference
between the target sample and all other samples in the task. The distance
representation between two samples neglects the specificity of the task and
lacks discrimination. This will cause the problem that the similarity scores
are not significantly different, so that the category of the target sample is
not clear. To deal with the key challenge of how to learn relation
representations with distinctive information, we propose a sample-to-task
metric module, as shown in the right panel of Figure 1, which adopts a meta
learning strategy to learn the relation representations. The main
contributions of this paper are summarized as follows:

  * •

We propose an task-level relation module (TLRM). The proposed TLRM utilizes
the attention mechanism to learn task-specific relation representations for
each task.

  * •

The comprehensive experimental results on four benchmark datasets show that
our proposed module is effective for GNN-based few-shot model. In addition,
the results of semi-supervised few-shot classification and visualization of
similarity scores are provided to further evaluate our module.

##  II Related work

Meta Learning in Few-shot Learning: Meta Learning framework is an effective
study for few-shot learning, which mainly focuses on how to learn and utilize
meta-level knowledge to adapt to new tasks quickly and well. One of the
excellent studies is model-agnostic meta-learning (MAML) [19]. The MAML
learned initialization parameters by cross-task training strategy such that
the base learner can rapidly generalize new tasks using a few support samples.
Subsequently, many MAML variants [20, 21, 22, 23, 24, 25] have been developed.

Metric Learning in Few-shot Learning: On the metric learning side, most of
algorithms consist of embedding function extracting features for instances and
metric function for measuring sample between the query embedding and class
embedding. Koch et al. [9] used siamese network to compute the pair-wise
distance between samples. Prototypical networks [6] firstly built a prototype
representation of each class and measured the samples between the query
embedding and class’s prototype by using euclidean distance. Matching network
[26] used a neural network with external memories to map samples to embedding
features, which considers full context in a task. TADAM [27] introduced a
metric scaling factor to optimize the similarity metric of prototypical nets.
Zheng et al. [28] believed that the average prototype ignores the different
importance of different support samples and proposed principal characteristic
nets.

Fixed metric methods will restrict the embedding function to produce
discriminative representations. Sung et al. [7] introduced relation network
(RN) for few-shot learning. The relation network learns to learn a deep
distance metric by a neural network. However, due to the inherent local
connectivity of CNN, the RN can be sensitive to the spatial position
relationship of semantic objects in two compared images. To address this
problem, Wu et al. [29] introduced a deformable feature extractor (DFE) to
extract more efficient features, and designed a dual correlation attention
mechanism (DCA) to deal with its inherent local connectivity. Hou et al. [30]
proposed a cross attention network for few-shot classification, which is
designed to model the semantic relevance between class and query features.

GNN-based methods in Few-shot Learning: Recently, most approaches are proposed
to exploit GNN in the field of few-shot learning task. Specifically, Garcia et
al. [8] first utilized GNN to solve few-shot learning problem, where all
embedding features extracting by a convolutional neural network are densely
connected. Liu et al. [11] proposed a transductive propagation network (TPN).
The TPN utilizes the entire query set for transductive inference. To further
exploit intra-cluster similarity and inter-cluster dissimilarity, kim et al.
[10] proposed an edge-labeling graph neural network. Then in order to
explicitly model the distribution-level relation, Yang et al. [18] proposed
distribution propagation graph network (DPGN).

In the existing GNN-based few-shot learning methods, pair-wise distance
representations are absolute difference of the embedding features. However,
when the classes in the task are similar, it will lead to the problem of
insufficient discrimination in metric representations. So, in this paper, we
focus on learning distinctive relation information through an task-level
relation module.

##  III The Proposed Method

###  III-A GNN-based few-shot learning

TABLE I: $5$-way $1$-shot classification accuracy on four benchmark datasets: mini-ImageNet, tiered-ImageNet, CUB-$200$-$2011$, and CIFAR-FS Model | Trans. | mini-ImageNet | tiered-ImageNet | CUB-200-2011 | CIFAR-FS
---|---|---|---|---|---
EGNN (CVPR 19) | No | $52.86\pm 0.42$ | $57.09\pm 0.42$ | $64.82\pm 0.41$ | $65.51\pm 0.43$
EGNN + TLRM | No | $53.65\pm 0.43$ | $57.40\pm 0.42$ | $65.07\pm 0.41$ | $65.00\pm 0.42$
TPN (ICLR 18) | Yes | $59.46$ | $-$ | $-$ | $-$
EGNN | Yes | $58.94\pm 0.51$ | $62.37\pm 0.51$ | $73.18\pm 0.51$ | $72.20\pm 0.49$
DPGN (CVPR 20) | Yes | $66.41\pm 0.51$ | $71.86\pm 0.50$ | $75.25\pm 0.46$ | $75.83\pm 0.47$
EGNN + TLRM | Yes | $60.79\pm 0.51$ | $64.52\pm 0.51$ | $75.02\pm 0.46$ | $73.42\pm 0.50$
DPGN + TLRM | Yes | $66.97\pm 0.53$ | $72.24\pm 0.50$ | $77.53\pm 0.46$ | $77.05\pm 0.46$
TABLE II: $5$-way $5$-shot classification accuracy on four benchmark datasets:
mini-ImageNet, tiered-ImageNet, CUB-$200$-$2011$, and CIFAR-FS

Model | Trans. | mini-ImageNet | tiered-ImageNet | CUB-200-2011 | CIFAR-FS
---|---|---|---|---|---
EGNN | No | $68.20\pm 0.41$ | $71.13\pm 0.39$ | $80.05\pm 0.36$ | $76.95\pm 0.37$
EGNN + TLRM | No | $68.72\pm 0.40$ | $72.39\pm 0.39$ | $81.03\pm 0.36$ | $77.78\pm 0.37$
TPN | Yes | $75.65$ | $-$ | $-$ | $-$
EGNN | Yes | $75.71\pm 0.46$ | $81.04\pm 0.43$ | $87.68\pm 0.38$ | $86.13\pm 0.41$
DPGN | Yes | $82.04\pm 0.45$ | $82.70\pm 0.43$ | $87.72\pm 0.36$ | $87.85\pm 0.38$
EGNN + TLRM | Yes | $76.18\pm 0.45$ | $81.47\pm 0.43$ | $88.00\pm 0.36$ | $85.70\pm 0.39$
DPGN + TLRM | Yes | $82.58\pm 0.45$ | $83.31\pm 0.44$ | $90.39\pm 0.34$ | $89.15\pm 0.37$

  * *

“No” means non-transductive method, and “Yes” means transductive method.

As shown in Figure 2 (a), GNN-based few-shot model usually consists of a CNN
for extracting features and a GNN for propagating labels from labeled nodes to
unlabeled according to similarity scores between nodes. In the training and
testing process, GNN-based few-shot model usually adopts the episodic
mechanism, in which each episode (task) consists of the support set $S$ and
the query set $Q$. And the support set contains $N*K$ labeled support samples
and the query set contains $T$ unseen samples in a $N$-way $K$-shot problem.

Generally, the CNN $g(\cdot)$ as backbone of extracting features has two
different types $1)$ the $4$-layer convolution network (ConvNet) [11, 10] and
$2)$ the $12$-layer residual network (ResNet-$12$) used in [18]. The GNN
consists of $L$ layers to process the graph. Let
$V=\left\\{V_{1},V_{2},...,V_{N\times K+T}\right\\}$ be embedding features for
all nodes extracted by the CNN, $R_{ij}$ be relation representations between
nodes, and $s_{ij}$ be similarity score between node $i$ and $j$. Given
$V^{L-1}$ and $s^{L-1}$ from the layer $L-1$, node feature update is firstly
conducted by a neighborhood aggregation procedure. And node $i$ is updated as

$V_{i}^{L}=f_{v}\left(\sum_{j=1}^{N\times K+T}V_{j}^{L-1}s_{ij}^{L-1}\right),$
(1)

where $f_{v}$ is the feature (node) transformation network. Then, the relation
representation is obtained by calculating the absolute difference between two
vector nodes. It can be denoted as

$\begin{split}R_{ij}&=\left|V_{i}^{L}-V_{j}^{L}\right|=\begin{matrix}\sum_{k=1}^{C}\left|V_{ik}^{L}-V_{jk}^{L}\right|\end{matrix}.\end{split}$
(2)

Finally, the relation representation $R_{ij}$ is input into a Multilayer
Perceptron (MLP) to capture the similarity scores between nodes

$\begin{split}s_{ij}&=f_{s}\left(R_{ij}\right)\\\
&=f_{s}\left(\begin{matrix}\sum_{k=1}^{C}\left|V_{ik}^{L}-V_{jk}^{L}\right|\end{matrix}\right)\\\
&=\sigma\left(\begin{matrix}\sum_{k=1}^{C}\omega_{k}\left|V_{ik}^{L}-V_{jk}^{L}\right|\end{matrix}\right).\end{split}$
(3)

Where $f_{s}$ is transformation network. The goal of GNN-based few-shot
learning is to learn function $g$, $f_{v}$ and $f_{s}$ to classify query
sample $x_{query}$ by
$\hat{y}_{query}=f_{s}(f_{v}(g(x_{query};D_{support})))\in{(0,1)}^{N}$. Note
that the relationship is obtained by measuring the distance between two
corresponding node, which is node-to-node and task-agnostic.

###  III-B Task-level Relation Module

In this paper, attention mechanism is employed to transform sample embedding
to relation representations with consideration to task-specific embedding.
Note that the relation representation is task-specific and not only the
distance between nodes. We denote it as Task-level Relation Module (TLRM). The
proposed TLRM can avoid direct comparison relative relationship irrelevant
local representations. As shown in Figure 2 (b), given the feature
representations $V\in R^{(N\times K+T)\times C}$, the relation representations
can be obtained. The implementation details are performed as follows.

For node $i$, the attention value between the target embedding and all other
samples in the task can be obtained by adopting method commonly used in the
attention mechanism. The attention value is performed as follows

$a\left(V_{i},V_{j}\right)=\cfrac{\text{exp}{(e_{ij}})}{\begin{matrix}\sum_{k=1}^{N\times
K+T}\end{matrix}\text{exp}{(e_{ik}})},$ (4)

Where $a\in R^{\left(N\times K+T\right)\times\left(N\times K+T\right)}$, which
represents the similarity between nodes comparing to all other embedding in
the task. $e_{ij}$ reflects the matching degree of node $i$ to node $j$. When
the degree is higher, $a_{ij}$ is bigger. The matching degree $e_{ij}$ is
performed as follows

$e_{ij}=s\left(V_{j},V^{T}_{i}\right)/{\sqrt{C}},$ (5)

Where, first, the feature representation $V_{i}\in R^{1\times C}$ of target
sample is reshaped to $V_{i}\in R^{C\times 1}$ through a transpose operation
and $s\left(V_{i},V^{T}_{j}\right)$ is the vector multiplication operation.
And then, $a_{ij}$ is used to encode $V_{j}$ and the relation representations
can be obtained, which can be denoted as

$\widehat{R}_{ij}=a\left(V_{i},V_{j}\right)\odot V_{j},$ (6)

The relation representation $\widehat{R}_{ij}$ models the relation
representation between node $i$ and $j$, which is a task-level relation
representation of sample $i$ to $j$ comparing to all the other samples.
Afterwards, $\widehat{R}_{ij}$ is fed to an MLP to capture the relation score
for performing further classification

$s_{ij}=\text{MLP}(\widehat{R}_{ij}).$ (7)

##  IV Experiments and Discussions

###  IV-A Datasets and setups

To evaluate our module, we select two GNN-based few-shot models: EGNN and
DPGN, and four standard few-shot learning benchmarks: mini-ImageNet [26],
tiered-ImageNet [31], CUB-$200$-$2011$ [32] and CIFAR-FS [33].

For the sake of fairness, all experiments employed the same setups as EGNN and
DPGN. EGNN used ConvNet, and DPGN used ResNet-$12$ for extracting features. In
training process, the Adam optimizer was used in all experiments with the
initial learning rate $10^{-3}$. And the learning rate was decayed by $0.1$
per $15000$ iterations. The weight decay was set as $10^{-5}$. For all
datasets, $5$-way $1$-shot and $5$-way $5$-shot experiments were conducted. We
randomly sampled $10,000$ tasks and then reported the mean accuracy along with
its $95\%$ confidence interval.

TABLE III: Semi-supervised few-shot classification accuracy on mini-ImageNet.
The results are tested in transductive learning.

Model | $20\%$ | $40\%$ | $60\%$
---|---|---|---
EGNN | $63.91\pm 0.42$ | $65.84\pm 0.41$ | $68.06\pm 0.44$
EGNN + TLRM | $\bf{66.37\pm 0.41}$ | $\bf{66.82\pm 0.42}$ | $\bf{69.08\pm 0.42}$
DPGN | $74.16\pm 0.44$ | $81.23\pm 0.44$ | $80.84\pm 0.46$
DPGN + TLRM | $\bf{80.99\pm 0.44}$ | $\bf{81.70\pm 0.45}$ | $\bf{80.94\pm 0.45}$

TABLE IV: Effect of adding the proposed TLRM to different layer of EGNN. The results on $5$-way $1$-shot are reported. Model | mini-ImageNet
---|---
EGNN | $52.86\pm 0.42$
EGNN+TLRM | $L=1$ | $53.11\pm 0.43$
$L=2$ | $53.29\pm 0.44$
$L=3$ | $53.42\pm 0.42$
$\bf{L=1~{}\&~{}2~{}\&~{}3}$ | $\bf{53.65\pm 0.43}$
Figure 3: Visualization of similarity scores obtained by EGNN (top) and EGNN
with our module (bottom).

###  IV-B Results and discussions for few-shot classification

Experimental results for $5$-way $1$-shot and $5$-way $5$-shot classification
are shown in Table I and Table II. We can see that EGNN or DPGN with our TLRM
have higher accuracy than the ones without TLRM on mini-ImageNet, tiered-
ImageNet, and CUB-$200$-$2011$. Meanwhile, partial experimental results on the
CIFAR-FS dataset dropped slightly, the reason of which might lie in the
categories in the CIFAR-FS dataset are highly distinguishable. In addition,
the CUB-$200$-$2011$ dataset is the most widely used benchmark for fine-
grained image classification, which has significant intra-class variance and
inter-class similarity. Fine-grained image task is more challenging in few-
shot learning. Clearly, the improvement on the CUB-$200$-$2011$ dataset is
significant in Table I and Table II, which shows that the relationship
representation obtained by our module is more discriminative than previous
method for tasks with high similarity. And overall, our method is simple and
effective.

Semi-supervised experiments were conducted in $5$-way $5$-shot setting on
mini-ImageNet with two backbones, in which the support samples are only
partially labeled. The results are presented in Table III. Notably, the EGNN
and DPGN with our TLRM outperforms the previous backbones especially when the
labeled samples portion was decreased.

###  IV-C Ablation studies

In order to investigate the effect of our proposed TLRM on different layer of
GNN, ablation studies were conducted with $L=1$, $L=2$, and $L=3$ on mini-
ImageNet with EGNN backbone. It can be observed from Table IV that the
proposed TLRM plays a significant role in each layer of EGNN.

###  IV-D Visualization of similarity scores

For further analysis, Figure 3 shows similarity scores in the last layer of
EGNN. The similarity scores are the average of 10000 tasks in setting of 5-way
5-shot and 5 queries for each class. The 25 samples in vertical axis are
support set, and 25 samples in horizontal axis are query set. Notably, EGNN
with our module not only contributes to predicting more accurately but also
reduces the similarity score between samples in different classes and
increases the similarity score between samples in the same classes.

##  V Conclusions

In this paper, we propose an task-level relation module to capture the
relation representations by employing all the embedding features in a single
task. By considering all the samples in the task, our method can hold
discriminative relation features for each node pair. Experimental results
demonstrate that it improves the performance of recently proposed GNN-based
methods on four benchmark datasets: mini-ImageNet, tiered-ImageNet,
CUB-$200$-$2011$, and CIFAR-FS.

## References

  * [1] K. Simonyan and A. Zisserman,  “Very deep convolutional networks for large-scale image recognition,”  in CVPR, 2014.
  * [2] K. He, X. Zhang, S. Ren, and J. Sun,  “Deep residual learning for image recognition,”  in CVPR, 2016.
  * [3] D. Chang, Y. Ding, J. Xie, A. Bhunia, X. Li, Z. Ma, M. Wu, J. Guo, and Y. Song,  “The devil is in the channels: Mutual-channel loss for fine-grained image classification,”  IEEE Transactions on Image Processing, vol. 29, pp. 4683–4695, 2020\.
  * [4] Y. Ding, Z. Ma, S. Wen, J. Xie, D. Chang, Z. Si, M. Wu, and H. Ling,  “Ap-cnn: Weakly supervised attention pyramid convolutional neural network for fine-grained visual classification,”  IEEE Transactions on Image Processing, vol. 30, pp. 2826–2836, 2021\.
  * [5] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. Berg, and F. Li,  “ImageNet Large Scale Visual Recognition Challenge,”  International Journal of Computer Vision, vol. 115, no. 3, pp. 211–252, 2015.
  * [6] J. Snell, K. Swersky, and R. Zemel,  “Prototypical networks for few-shot learning,”  in NIPS, 2017.
  * [7] F. Sung, Y. Yang, L. Zhang, T. Xiang, P. Torr, and T. Hospedales,  “Learning to compare: Relation network for few-shot learning,”  in CVPR, 2018.
  * [8] V. Garcia and J. Bruna,  “Few-shot learning with graph neural networks,”  in ICLR, 2018.
  * [9] G. Koch, R. Zemel, and R. Salakhutdinov,  “Siamese neural networks for one-shot image recognition,”  in ICML, 2015.
  * [10] J. Kim, T. Kim, S. Kim, and C. Yoo,  “Edge-labeling graph neural network for few-shot learning,”  in CVPR, 2019.
  * [11] Y. Liu, J. Lee, M. Park, S. Kim, E. Yang, S. Hwang, and Y. Yang,  “Learning to propagate labels: Transductive propagation network for few-shot learning,”  in ICLR, 2018.
  * [12] W. Li, J. Xu, J. Huo, L. Wang, Y. Gao, and J. Luo,  “Distribution consistency based covariance metric networks for few-shot learning,”  in AAAI, 2019.
  * [13] W. Li, L. Wang, J. Xu, J. Huo, Y. Gao, and J. Luo,  “Revisiting local descriptor based image-to-class measure for few-shot learning,”  in CVPR, 2019.
  * [14] X. Li, J. Wu, Z. Sun, Z. Ma, J. Cao, and J. Xue,  “Bsnet: Bi-similarity network for few-shot fine-grained image classification,”  IEEE Transactions on Image Processing, vol. 30, pp. 1318–1331, 2021\.
  * [15] C. Xu, Y. Fu, C. Liu, C. Wang, J. Li, F. Huang, L. Zhang, and X. Xue,  “Learning dynamic alignment via meta-filter for few-shot learning,”  in CVPR, 2021.
  * [16] H. Zhang, P. Koniusz, S. Jian, H. Li, and P. Torr,  “Rethinking class relations: Absolute-relative supervised and unsupervised few-shot learning,”  in CVPR, 2021.
  * [17] B. Zhang, X. Li, Y. Ye, Z. Huang, and L. Zhang,  “Prototype completion with primitive knowledge for few-shot learning,”  in CVPR, 2021.
  * [18] L. Yang, L. Li, Z. Zhang, X. Zhou, E. Zhou, and Y. Liu,  “DPGN: Distribution propagation graph network for few-shot learning,”  in CVPR, 2020.
  * [19] C. Finn, P. Abbeel, and S. Levine,  “Model-agnostic meta-learning for fast adaptation of deep networks,”  in ICML, 2017.
  * [20] M. Jamal and G. Qi,  “Task agnostic meta-learning for fewshot learning,”  in CVPR, 2019.
  * [21] A. Rusu, D. Rao, J. Sygnowski, O. Vinyals, R. Pascanua, S. Osindero, and R. Hadsell,  “Meta-learning with latent embedding optimization,”  in ICLR, 2019.
  * [22] X. Jiang, M. Havaei, F. Varno, G. Chartrand, N. Chapados, and S. Matwin,  “Learning to learn with conditional class dependencies,”  in ICLR, 2019.
  * [23] S. Ravi and H. Larochelle,  “Optimization as a model for few-shot learning,”  in ICLR, 2017.
  * [24] T. Munkhdalai and H. Yu,  “Meta networks,”  in ICML, 2017.
  * [25] H. Li, W. Dong, X. Mei, C. Ma, F. Huang, and B. Hu,  “LGM-Net: learning to generate matching networks for few-shot learning,”  in ICML, 2019.
  * [26] O. Vinyals, C. Blundell, T. Lillicrap, K. Kavukcuoglu, and D. Wierstra,  “Matching networks for one shot learning,”  in NIPS, 2016.
  * [27] B. Oreshkin, A. Lacoste, and P. Rodriguez,  “TADAM: task dependent adaptive metric for improved few-shot learning,”  in NIPS, 2018.
  * [28] Y. Zheng, R. Wang, J. Yang, L. Xue, and M. Hu,  “Principal characteristic networks for few-shot learning,”  Visual communication and Image Representation, vol. 59, pp. 563–573, 2019.
  * [29] Z. Wu, Y. Li, L. Guo, and K. Jia,  “PARN: position-aware relation networks for few-shot learning,”  in ICCV, 2019.
  * [30] R. Hou, H. Chang, B. Ma, S. Shan, and X. Chen,  “Cross attention network for few-shot classification,”  in NIPS, 2019.
  * [31] M. Ren, E. Triantafillou, S. Ravi, J. Snell, K. Swersky, J. Tenenbaum, H. Larochelle, and R. Zemel,  “Meta-learning for semi-supervised few-shot classification,”  in ICLR, 2018.
  * [32] C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie,  “The Caltech-UCSD Birds-200-2011 Dataset,”  Tech. Rep. CNS-TR-2011-001, California Institute of Technology, 2011\.
  * [33] L. Bertinetto, J. Henriques, P. Torr, and A. Vedaldi,  “Meta-learning with differentiable closedform solvers,”  in ICML, 2019.