11institutetext: SCALE AI Research Chair in Data-Driven Algorithms for Modern
Supply Chains
Department of Mechanical and Industrial Engineering, University of Toronto
11email<EMAIL_ADDRESS>
# Multi-Task Predict-then-Optimize

Bo Tang  0000-0002-6035-5167 Elias B. Khalil  0000-0001-5844-9642

###### Abstract

The predict-then-optimize framework arises in a wide variety of applications
where the unknown cost coefficients of an optimization problem are first
predicted based on contextual features and then used to solve the problem. In
this work, we extend the predict-then-optimize framework to a multi-task
setting: contextual features must be used to predict cost coefficients of
multiple optimization problems, possibly with different feasible regions,
simultaneously. For instance, in a vehicle dispatch/routing application,
features such as time-of-day, traffic, and weather must be used to predict
travel times on the edges of a road network for multiple traveling salesperson
problems that span different target locations and multiple $s-t$ shortest path
problems with different source-target pairs. We propose a set of methods for
this setting, with the most sophisticated one drawing on advances in multi-
task deep learning that enable information sharing between tasks for improved
learning, particularly in the small-data regime. Our experiments demonstrate
that multi-task predict-then-optimize methods provide good tradeoffs in
performance among different tasks, particularly with less training data and
more tasks.

###### Keywords:

multi-task learning predict-then-optimize data-driven optimization machine
learning

##  1 Introduction

The predict-then-optimize framework, in which the unknown coefficients for an
optimization problem are predicted and then used to solve the problem, is
emerging as a useful framework in some applications. For instance, in vehicle
routing and job scheduling, we often require optimization where the model’s
cost coefficients, e.g., travel time and execution time, are unknown but
predictable at decision time. In the conventional two-stage method, a learning
model is first trained to predict cost coefficients, after which a solver
separately optimizes accordingly. However, end-to-end approaches that learn
predictive models that minimize the decision error directly have recently
gained interest due to some improvements in experimental performance [4, 15,
10, 12, 22, 23]. Although there has been some recent work in predicting
elements of the constraint matrix in a linear programming setting [13, 17],
our focus here is on the predominant line of research that has focused on
unknown cost coefficients [10, 32, 14, 22, 1, 12, 23, 30, 5, 8].

Previous work on predict-then-optimize has focused on learning the cost
coefficients for a single optimization task. However, it is natural to
consider the setting where multiple related tasks can share information and
representations. For example, a vehicle routing application requires
predicting travel times on the edges of a road network for multiple traveling
salesperson problems (TSPs) that span different target locations and multiple
shortest path problems with different source-target pairs. These travel time
predictions should be based on the same contextual information, e.g., if the
tasks are to be executed at the same time-of-day, then the travel times that
should be predicted for the different tasks depend on the same features.
Another case is in package delivery, where distributing packages from one
depot to multiple depots results in independent delivery tasks that
nonetheless share travel time predictions since they use the same road
network. To that end, we introduce multi-task end-to-end predict-then-
optimize, which simultaneously solves multiple optimization problems with a
loss function that relates to the decision errors of all such problems.

Multi-task learning has been successfully applied to natural language
processing, computer vision, and recommendation systems. However, its
applicability to the predict-then-optimize paradigm is yet to be explored.
Predict-then-optimize with multi-task learning is attractive because of the
ability to improve model performance in the small-data regime. Machine
learning, especially with deep neural networks, is data-intensive and prone to
overfitting, which might limit applicability to the predict-then-optimize
paradigm. The need to simultaneously minimize the losses of different tasks
helps reduce overfitting and improve generalization. Multi-task learning
combines the data of all tasks, which increases the overall training data size
and alleviates task-specific noise.

To the best of our knowledge, we introduce multi-task learning for end-to-end
predict-then-optimize for the first time. We motivate and formalize this
problem before proposing a set of methods for combining the different task
losses. Our experiments show that multi-task end-to-end predict-then-optimize
provides performance benefits and good tradeoffs in performance among
different tasks, especially with less training data and more tasks, as
compared to applying standard single-task methods independently for each task.
As an additional contribution, we distinguish end-to-end predict-then-optimize
approaches that learn from observed costs (the usual setting of Elmachtoub and
Grigas [12]) and those that learn directly from (optimal) solutions without
the objective function costs themselves. This extends our framework to
applications where there are no labeled coefficients in the training data,
e.g., the Amazon Last Mile Routing Challenge [33]. The open-source code is
available 111https://github.com/khalil-research/Multi-Task_Predict-then-
Optimize.

##  2 Related Work

###  2.1 Differentiable Optimization

The key component of gradient-based end-to-end predict-then-optimize is
differentiable optimization, which allows the backpropagation algorithm to
update model parameters from the decision made by the optimizer. Based on the
KKT conditions, Amos and Kolter [3] introduced OptNet, a differentiable
optimizer to make quadratic optimization problems learnable. With OptNet,
Donti et al. [10] investigated a learning framework for quadratic programming;
Wilder et al. [32] then added a small quadratic regularization to the linear
objective function for linear programming; Ferber et al. [14] extended the
method to the discrete model with the cutting plane; Mandi and Guns [22]
adopted log-barrier regularization instead of a quadratic one. Besides OptNet,
Agrawal et al. [1] leveraged KKT conditions to differentiate conic
programming.

Except for the above approaches with KKT, an alternative methodology is to
design gradient approximations to avoid ill-defined gradients from predicted
costs to optimal solutions. Elmachtoub and Grigas [12] proposed a convex
surrogate loss. Vlastelica et al. [30] developed a differentiable optimizer
through implicit interpolation. Berthet et al. [5] demonstrated a method with
stochastic perturbation to smoothen the loss function and further constructed
the Fenchel-Young loss. Dalle et al. [8] extended the perturbation approach to
the multiplicative perturbation and the Frank-Wolfe regularization. Mulamba et
al. [24] studied a solver-free contrastive loss. Moreover, Shah et al. [26]
provided an alternate paradigm that additionally trains a model to predict
decision errors to replace the solver.

###  2.2 Multi-Task Learning

Multi-task learning, first proposed by Caruana [6], aims to learn multiple
tasks together with joint losses. In short, a model with multiple loss
functions is defined as multi-task learning. Much research has focused on the
(neural network) model architecture: the most basic model is a shared-bottom
model [6], including shared hidden layers at the bottom and task-specific
layers at the top. Besides such hard parameter sharing schemes, there is also
soft sharing so that each task keeps its own parameters. Duong et al. [11]
added $l_{2}$ norm regularization to encourage similar parameters between
tasks. Furthermore, neural networks with different experts and some gates [27,
21, 29] were designed to fuse information adaptively.

Another crucial issue is resolving the unbalanced (in magnitude) and
conflicting gradients from different tasks. There are weighting approaches,
such as UW [18], GradNorm [7], and DWA [20], that have been proposed to adjust
the weighting of different losses. Other methods, such as PCGrad [34], GradVec
[31], and CAGrad [19], were designed to alter the direction of the overall
gradient to avoid conflicts and accelerate convergence.

##  3 Building Blocks

###  3.1 Optimization Problem

For multi-task end-to-end predict-then-optimize, each task $t$ is a separate
(integer) linear optimization problem, defined as follows:

$\displaystyle\begin{aligned}
\min_{\bm{w}^{t}}\quad&{\bm{c}^{t}}^{T}\bm{w}^{t}\\\
\textrm{s.t.}\quad&\bm{A}^{t}\bm{w}^{t}\leq\bm{b}^{t}\\\
&\bm{w}^{t}\geq\bm{0}\\\ &\textrm{Some }w^{t}_{i}\textrm{ are
integer.}\end{aligned}$

The decision variables are $\bm{w}^{t}$, the constraint coefficients are
$\bm{A}^{t}$, the right-hand sides of the constraints are $\bm{b}^{t}$, and
the unknown cost coefficients are $\bm{c}^{t}$. Given a cost $\bm{c}^{t}$,
$\bm{w}^{t*}_{\bm{c}^{t}}$ is the corresponding optimal solution, and
$z^{t*}_{\bm{c}^{t}}$ is the optimal value. Additionally, we define
$S^{t}=\\{\bm{A}^{t}\bm{w}^{t}\leq\bm{b}^{t},\bm{w}^{t}\geq\bm{0},...\\}$ as
the feasible region of decision variables $\bm{w}^{t}$.

###  3.2 Gradient-based Learning

End-to-end predict-then-optimize aims to minimize a decision loss directly. As
a supervised learning problem, it requires a labeled dataset $\mathcal{D}$
consisting of features $\bm{x}$ and labels $\bm{c}$ or $\bm{w}^{*}_{c}$. We
will further discuss the difference between cost labels $\bm{c}$ and solution
labels $\bm{w}^{*}_{c}$ in Sec 4.2. As shown in Figure 1, multi-task end-to-
end predict-then-optimize predicts the unknown costs for multiple optimization
problems and then solves these optimization problems with the predicted costs.
The critical component is a differentiable optimization embedded into a
differentiable predictive model. However, the only learnable part is the
prediction model $g(\bm{x};\bm{\theta})$, since there are no parameters to
update in the solver and loss function.

Figure 1: Illustration of end-to-end multi-task end-to-end predict-then-
optimize: labeled datasets
$\mathcal{D}^{1},\mathcal{D}^{2},...,\mathcal{D}^{t}$ are used to fit a
machine learning predictor $g(\bm{x};\bm{\theta})$ that predicts costs
$\hat{\bm{c}}^{t}$ for each task $t$. The loss function
$l^{t}(\hat{\bm{c}}^{t},\cdot)$ to be minimized measures decision error
instead of prediction error.

###  3.3 Decision Losses

Loss functions for end-to-end predict-then-optimize aim to measure the error
in decision-making. For instance, regret is defined as the difference in
objective values due to an optimal solution that is based on the true costs
and another that is based on the predicted costs:

$\displaystyle
l_{\text{Regret}}(\hat{\bm{c}},\bm{c})=\bm{c}^{T}\bm{w}^{*}_{\hat{\bm{c}}}-z^{*}_{\bm{c}}$

However, with a linear objective function, regret does not provide useful
gradients for learning [12]. Besides regret, decision error can also be
defined as the difference between the true solution and its prediction, such
as using the Hamming distance [30] or the squared error of a solution to an
optimal one [5]. Because the function from costs $\bm{c}$ to optimal solutions
$\bm{w}^{*}_{\bm{c}}$ is piecewise constant, a solver with any of the
aforementioned losses has no nonzero gradients to update the model parameter
$\bm{\theta}$. Thus, the state-of-art methods, namely Smart Predict-then-
Optimize (SPO+) [12] and Perturbed Fenchel-Young Loss (PFYL), both design
surrogate decision losses which allow for a nonzero approximate gradient (or
subgradient), $\frac{\partial l(\cdot)}{\partial\hat{\bm{c}}}$.

Smart Predict-then-Optimize Loss (SPO+). SPO+ loss [12] is a differentiable
convex upper bound on the regret:

$\displaystyle\begin{aligned}
l_{\text{SPO+}}(\hat{\bm{c}},\bm{c})=-\underset{\bm{w}\in
S}{\min}\\{(2\hat{\bm{c}}-\bm{c})^{T}\bm{w}\\}+2\hat{\bm{c}}^{T}\bm{w}^{*}_{\bm{c}}-z^{*}_{\bm{c}}.\end{aligned}$

One proposed subgradient for this loss writes as follows:

$\displaystyle\begin{aligned}
2(\bm{w}^{*}_{\bm{c}}-\bm{w}^{*}_{2\hat{\bm{c}}-\bm{c}})\in\frac{\partial
l_{\text{SPO+}}(\hat{\bm{c}},\bm{c})}{\partial\hat{\bm{c}}}.\end{aligned}$

In theory, the SPO+ framework can be applied to any problem with a linear
objective. Elmachtoub and Grigas [12] have conducted experiments on shortest
path as a representative linear program and portfolio optimization as a
representative quadratically-constrained problem.

Perturbed Fenchel-Young Loss (PFYL). PFYL [5] leverages Fenchel duality, where
Berthet et al. [5] discussed only the case of linear programming. The
predicted costs are sampled with Gaussian perturbation $\bm{\xi}$, and the
expected function of the perturbed minimizer is defined as
$F(\bm{c})=\mathop{{}\mathbb{E}}_{\bm{\xi}}[\underset{\bm{w}\in
S}{\min}{\\{(\bm{c}+\sigma\bm{\xi})^{T}\bm{w}\\}}]$. With
$\Omega(\bm{w}^{*}_{\bm{c}}),$ the dual of $F(\bm{c})$, the Fenchel-Young loss
reads:

$l_{\text{FY}}(\hat{\bm{c}},\bm{w}^{*}_{\bm{c}})=\hat{\bm{c}}^{T}\bm{w}^{*}_{\bm{c}}-F(\hat{\bm{c}})-\Omega(\bm{w}^{*}_{\bm{c}}).$

Then, we can estimate the gradients by $M$ samples Monte Carlo:

$\frac{\partial
l_{\text{FY}}(\hat{\bm{c}},\bm{w}^{*}_{\bm{c}})}{\partial\hat{\bm{c}}}\approx\bm{w}^{*}_{\bm{c}}-\frac{1}{M}\sum_{m}^{M}\underset{\bm{w}\in
S}{\operatornamewithlimits{argmin}}{\\{(\hat{\bm{c}}+\sigma\bm{\xi}_{m})^{T}\bm{w}\\}}$

###  3.4 Multi-Task Loss Weighting Strategies

The general idea of multi-task learning is that multiple tasks are solved
simultaneously by the same predictive model. It is critical for a multi-task
neural network, one such flexible class of models, to balance losses among
tasks with the loss weights $u_{t}$ for task $t$. The weighting approaches we
evaluated include a uniform combination (all $u^{t}=1$) and GradNorm [7], an
adaptive loss weighting approach. The latter provides adaptive per-task
weights $u_{\text{Ada}}^{t}$ that are dynamically adjusted during training in
order to keep the scale of the gradients similar. In this work, we set the
GradNorm hyperparameters of “restoring force” to $0.1$ and the learning rate
of loss weights to $0.005$. Further tuning is possible but was not needed for
our experiments.

All the training strategies we have explored in this paper, including baseline
approaches, are summarized in Table 1. Let $T$ be the number of tasks, and
cost coefficient prediction for task $t$ be $\hat{\bm{c}}^{t}$. “mse” is the
usual two-stage baseline of training a regression model that minimizes cost
coefficients mean-squared error
$l_{\text{MSE}}=\frac{1}{n}\sum_{i}^{n}\|\hat{\bm{c}}_{i}-\bm{c}_{i}\|^{2}$
only without regard to the decision. “separated” trains one model per task,
minimizing, for each task, a decision-based loss such as SPO+ or PFYL from
Section 3.3. “comb” simply sums up the per-task decision losses, whereas
“gradnorm” does so in a weighted adaptive way. For any of these methods,
whenever “+mse” is appended to the method name, a variant of the method is
obtained that combines additional mean-squared error $l_{\text{MSE}}$ in the
cost predictions with the decision loss. Such a regularizer is known to be
useful in practice, even with the primary evaluation metric of a trained model
being its decision regret [12]. Although we refer to “separated+mse” as a
single-task method, it can also be considered as a multi-task learning method
in a broad sense because of the inclusion of two losses.

| Strategy | Losses
---|---|---
Single-Task | mse | $l_{\text{MSE}}(\bm{c},\hat{\bm{c}})$
separated | Separate $l_{\text{Decision}}(\hat{\bm{c}}^{t},\cdot)$ for each task $t$
separated+mse | Separate $l_{\text{Decision}}(\hat{\bm{c}}^{t},\cdot)+l_{\text{MSE}}(\bm{c}^{t},\hat{\bm{c}}^{t})$ for each task $t$
Multi-Task | comb | $\sum_{t}^{T}l_{\text{Decision}}^{t}(\hat{\bm{c}}^{t},\cdot)$
comb+mse | $\sum_{t}^{T}l_{\text{Decision}}^{t}(\hat{\bm{c}}^{t},\cdot)+\sum_{t}^{T}l_{\text{MSE}}^{t}(\bm{c}^{t},\hat{\bm{c}}^{t})$
gradnorm | $\sum_{t}^{T}u_{\text{Ada}}^{t}l_{\text{Decision}}^{t}(\hat{\bm{c}}^{t},\cdot)$
gradnorm+mse | $\sum_{t}^{T}u_{\text{Ada}_{1}}^{t}l_{\text{Decision}}^{t}(\hat{\bm{c}}^{t},\cdot)+\sum_{t}^{T}u_{\text{Ada}_{2}}^{t}l_{\text{MSE}}^{t}(\bm{c}^{t},\hat{\bm{c}}^{t})$
Table 1: Losses of Various Training Strategies

##  4 Learning Architectures

###  4.1 Shared Learnable Layers

The model class we will explore is deep neural networks. Besides their
capacity to represent complex functions from labeled data, neural networks
have a compositional structure that makes them particularly well-suited for
multi-task learning. A multi-task neural network shares hidden layers across
all tasks and keeps specific layers for each task. Figure 2 illustrates that
the sharing part of multi-task end-to-end predict-then-optimize depends on the
consistency of the predicted coefficients, which we will define next. At a
high level, Figure 2 distinguishes two settings. On the left, the different
tasks use the exact same predicted cost vector. On the right, each task could
have a different cost vector. In both settings, the predictions are based on
the same input feature vector.

Figure 2: Illustration of two types of multi-task end-to-end predict-then-
optimize: On the left, all optimization tasks require the same prediction as
cost coefficients. On the right, different tasks share some layers as feature
embeddings and make different coefficient predictions.

Shared Predicted Coefficients (Single-Cost). In this setting (left of Figure
2), which we will refer to as single-cost, the same cost coefficients are
shared among all tasks. For example, multiple navigation tasks on a single map
are shortest-path problems with different source-target pairs that share the
same distance matrix (i.e., costs). In this case, the cost coefficients
$\bm{c}^{t}$ for task $t$ are equal to or a subset of the shared costs
$\bm{c}$. Thus, the prediction model is defined as

$\hat{\bm{c}}=g(\bm{x};\bm{\theta}_{\text{Shared}}),$

which is the same as a single-task model, and the multiple tasks combine their
losses,

$\sum_{t}u^{t}l_{\text{Decision}}^{t}(\hat{\bm{c}},\cdot),$

based on the shared prediction, $\hat{\bm{c}}$. Therefore, as Algorithm 1
shows, all learnable layers are shared. In addition, the baseline methods we
referred to as “separated” and “separated+mse” are not practical in this same-
costs setting as they, inconsistently, produce different cost predictions for
each task, even when that is not required. Nonetheless, an experimental
assessment of their performance will be conducted in order to contrast it with
using multi-task learning.

Algorithm 1 Single-Cost Multi-Task Gradient Descent

1:coefficient matrix $\bm{A}^{1},\bm{A}^{2},...$; right-hand side
$\bm{b}^{1},\bm{b}^{2},...$; training data $\mathcal{D}$

2:Initialize predictor parameters $\bm{\theta}_{\text{Shared}}$ for predictor
$g(\bm{x};\bm{\theta}_{\text{Shared}})$

3:for epochs do

4: for each batch of training data $(\bm{x},\bm{c})$ or
$(\bm{x},\bm{w}_{\bm{c}}^{1*},\bm{w}_{\bm{c}}^{2*},...)$ do

5: Forward pass to predict cost
$\hat{\bm{c}}:=g(\bm{x};\bm{\theta}_{\text{Shared}})$

6: Forward pass to solve optimal solution
$\bm{w}^{t*}_{\hat{\bm{c}}}:={\operatornamewithlimits{argmin}}_{\bm{w}^{t}\in
S^{t}}\hat{\bm{c}}^{T}\bm{w}^{t}$ per task

7: Forward pass to sum weighted decision losses
$l(\hat{\bm{c}},\cdot):=\sum_{t}u^{t}l^{t}(\hat{\bm{c}},\cdot)$

8: Backward pass from loss $l(\hat{\bm{c}},\cdot)$ to update parameters
$\bm{\theta}_{\text{Shared}}$

9: end for

10:end for

Shared Features Embeddings (Multi-Cost). In many applications of predict-then-
optimize, the optimization problem requires cost coefficients that are
specific and heterogeneous to each task, but that can be inferred from
homogeneous contextual features. For instance, in a vehicle routing
application, features such as time of day and weather predict travel time in
different regions. Compared to the single-cost setting we just introduced, the
multi-cost predictor here has the form

$\hat{\bm{c}}^{t}=g(\bm{x};\bm{\theta}_{\text{Shared}};\bm{\theta}_{t}).$

Per-task predictions are made by leveraging the same information embedding in
the layers of the neural network that are shared across the tasks (see right
of Figure 2). Thus, the corresponding loss function is

$\sum_{t}u^{t}l_{\text{Decision}}^{t}(\hat{\bm{c}}^{t},\cdot),$

and Algorithm 2 updates the parameters of the predictor.

Algorithm 2 Multi-Cost Multi-Task Gradient Descent

1:coefficient matrix $\bm{A}^{1},\bm{A}^{2},...$; right-hand side
$\bm{b}^{1},\bm{b}^{2},...$; training data
$\mathcal{D}^{1},\mathcal{D}^{2},...$

2:Initialize predictor parameters
$\bm{\theta}_{\text{Shared}},\bm{\theta}_{1},\bm{\theta}_{2},...$ for
predictor $g(\bm{x};\bm{\theta}_{\text{Shared}};\bm{\theta}_{t})$

3:for epochs do

4: for each batch of training data
$(\bm{x}^{1},\bm{x}^{2},...,\bm{c}^{1},\bm{c}^{2},...)$ or
$(\bm{x}^{1},\bm{x}^{2},...,\bm{w}_{\bm{c}^{1}}^{1*},\bm{w}_{\bm{c}^{2}}^{2*},...)$
do

5: Forward pass to predict cost
$\hat{\bm{c}}^{t}:=g(\bm{x}^{t};\bm{\theta}_{\text{Shared}};\bm{\theta}_{t})$
per task

6: Forward pass to solve optimal solution
$\bm{w}^{t*}_{\hat{\bm{c}}^{t}}:={\operatornamewithlimits{argmin}}_{\bm{w}^{t}\in
S^{t}}\hat{\bm{c}}^{tT}\bm{w}^{t}$ per task

7: Forward pass to sum weighted decision losses
$l(\hat{\bm{c}}^{1},\hat{\bm{c}}^{2},...,\cdot):=\sum_{t}u^{t}l^{t}(\hat{\bm{c}}^{t},\cdot)$

8: Backward pass from loss $l(\hat{\bm{c}}^{1},\hat{\bm{c}}^{2},...,\cdot)$ to
update parameters
$\bm{\theta}_{\text{Shared}},\bm{\theta}_{1},\bm{\theta}_{2},...$

9: end for

10:end for

###  4.2 Label Accessibility and Learning Paradigms

We distinguish two learning paradigms that require different kinds of labels
(cost coefficients $\bm{c}$ or optimal solutions $\bm{w}^{*}_{\bm{c}}$) in the
training data: learning from costs and learning from (optimal) solutions. This
distinction is based on the availability of labeled cost coefficients
$\bm{c}$. Thus, SPO+ is learning from costs because the calculation of SPO+
loss involves true cost coefficients $\bm{c}$, whereas PFYL is learning from
solutions that do not require access to $\bm{c}$; see Section 3.3.

The need for the true cost coefficients as labels in the training data is a
key distinguishing factor because these cost coefficients provide additional
information that can be used to train the model, but they may be absent in the
data. Deriving optimal solutions from the cost coefficients is trivial, but
the opposite is intricate as it requires some form of inverse optimization
[2]. The ability to directly learn from solutions extends the applicability of
end-to-end predict-then-optimize beyond what a two-stage approach, which is
based on regressing on the cost coefficients, can do. Indeed, the recent MIT-
Amazon Last Mile Routing Challenge [33] is one such example in which good TSP
solutions are observed on historical package delivery instances, but the
corresponding edge costs are unobserved. Those good solutions are based on
experienced drivers’ tacit knowledge.

##  5 Experiments

In this section, we present experimental results for multi-task end-to-end
predict-then-optimize. In our experiments, we evaluate decision performance
using regret, and we use mean-squared error (MSE) to measure the prediction of
cost coefficients $\hat{\bm{c}}$. We use SPO+ and PFYL as typical methods for
learning from costs and learning from solutions and adopt various multi-task
learning strategies discussed in Sec 3.4, as well as two-stage and single-task
baselines. Our experiments are conducted on two datasets, including graph
routing on PyEPO TSP dataset 222https://khalil-research.github.io/PyEPO [25],
and adjusted Warcraft terrain 333
https://drive.google.com/file/d/1lYPT7dEHtH0LaIFjigLOOkxUq4yi94fy [30] to
learn single-cost decisions and multi-cost decisions. We also vary the amount
of training data size and the number of tasks.

All the numerical experiments were conducted in Python v3.7.9 with two Intel
E5-2683 v4 Broadwell CPUs, two NVIDIA P100 Pascal GPUs, and 8GB memory.
Specifically, we used PyTorch [25] v1.10.0 for the prediction model and Gurobi
[16] v9.1.2 for the optimization solver, and PyEPO [28] v0.2.0 for SPO+ and
PFYL autograd functions.

###  5.1 Benchmark Datasets and Neural Network Architecture

Graph Routing with Multiple Tasks. We used the traveling salesperson problem
dataset generated from PyEPO [28], which uses the Euclidean distance among
nodes plus polynomial function
$f(\bm{x}_{i})=(\frac{1}{\sqrt{p}}(\mathcal{B}\bm{x}_{i})_{j}+3)^{4}$ (where
$\mathcal{B}$ is a random matrix) with random noise perturbations
$f(\bm{x}_{i})\cdot\epsilon$ to map the features $\bm{x}$ into a symmetric
distance matrix of a complete graph. We discuss both learning from costs and
learning from solutions. In this experiment, the number of features $p$ is
$10$, the number of nodes $m$ is $30$, the polynomial degree of function
$f(\bm{x}_{i})$ is $4$, and the noise $\epsilon$ comes from U(0.5, 1.5). We
sample $15-22$ nodes as target locations for multiple traveling salesperson
problems (TSPs) and $54$ undirected edges for multiple shortest paths (SPs)
with different source-target pairs. Thus, all TSP and SP tasks share the same
cost coefficients.

Since the multiple routing tasks require consistent cost coefficients, the
model $g(\bm{x};\bm{\theta})$ makes one prediction of the costs that is used
for all of the tasks. The architecture of the regression network is one fully-
connected layer with a softplus activation to prevent negative cost
predictions, and all tasks share the learnable layer. For the hyperparameters,
the learning rate is $0.1$, the batch size is $32$, and the max training
iterations is $30000$ with $5$ patience early stopping. For PFYL, the number
of samples $M$ is $1$, and the perturbation temperature $\sigma$ is $1.0$. We
formulate SP as a network flow problem and use the Dantzig–Fulkerson–Johnson
(DFJ) formulation [9] to solve TSP.

Warcraft Shortest Path with Various Species. The Warcraft map shortest path
dataset [30] allows for the learning of the shortest path from RGB terrain
images, and we use $96\times 96$ RGB images for $12\times 12$ grid networks
and sample $3$ small subsets from $10000$ original training data points for
use in training. As shown in Figure 3, we modify the cost coefficients for
different species (human, naga, dwarf) and assume that the cost coefficients
are not accessible in the data. This means there are three separate datasets
of feature-solution pairs, which require learning from solutions using the
PFYL method. Similar to SP tasks in Graph Routing, the shortest path
optimization model is a linear program.

Figure 3: Multiple datasets of Warcraft terrain images for different species,
where labeled cost coefficients are unavailable.

Since the multiple Warcraft shortest paths tasks require us to predict cost
coefficients for different species, the prediction model should incorporate
task-specific layers. Following Vlastelica et al. [30], we train a truncated
ResNet18 (first five layers) for $50$ epochs with batches of size $70$, and
learning rate $0.0005$ decaying at the epochs $30$ and $40$. The first three
layers are the shared-bottom. The number of samples $M$ is $1$, and the
perturbation temperature $\sigma$ is $1.0$.

###  5.2 Performance Advantage of Multi-Task Learning

(a)  (b)

(c)  (d)

Figure 4: Performance Radar Plot for Graph Routing: Average performance for
different tasks on the test set, trained with SPO+ (Top) and PFYL (Bottom),
and $100$ training data points, including regrets and cost MSE, lower is
better. SP $i$ is the regret for shortest path task $i$, TSP $i$ is the regret
for traveling salesperson task $i$, MSE is the mean squared error of cost
coefficients.

Experimental results on multiple routing tasks on a graph, as shown in Figure
4, demonstrate that multi-task end-to-end predict-then-optimize, especially
with GradNorm, has a performance advantage over single-task approaches. Figure
4 shows the results of learning from costs with SPO+ (top), and the results of
learning from solutions with PFYL (bottom). In these “radar plots”, the per-
task regret on unseen test data is shown along each dimension (lower is
better). It can be seen that the innermost method in these figures (best seen
in color) is the red one, “gradnorm+mse”.

(a)  (b)

Figure 5: Learning from Costs for Graph Routing: Average performance for
different tasks on the test set, trained with SPO+ and $100$ training data
points for more methods, lower is better. SP $i$ is the regret for shortest
path task $i$, TSP $i$ is the regret for traveling salesperson task $i$, MSE
is the mean squared error of cost coefficients.

More experiments are shown in Fig. 5. The investigation includes the two-stage
method, single-task, and multi-task with and without cost MSE as
regularization; these two plots include a superset of the methods in Fig. 4.
Despite achieving a lower MSE, the two-stage approaches exhibit significantly
worse regret than end-to-end learning. Additionally, adding an MSE regularizer
on the cost coefficients consistently improves end-to-end learning approaches.
Thus, we always include the additional cost of MSE regularizer when learning
from costs. However, since labeled costs are absent when learning directly
from solutions, PFYL cannot add cost MSE (“+mse”) as regularization and cannot
be compared with the two-stage method (which requires cost labels.)

###  5.3 Efficiency Benefit of Multi-Task Learning

(a)  (b)

Figure 6: Training time for Graph Routing: The elapsed time of training to
convergence at different settings.

Figure 6 shows the training time for SPO+ and PFYL models when using early
stopping when five consecutive epochs exhibit non-improving loss values on
held-out validation data, a standard trick in neural network training. For
“separated” and “separated+mse”, the training time is the sum of each
individual model. We can see that the use of GradNorm to adjust the weights
dynamically allows for efficient model training as faster convergence is
achieved. The “separated+mse” baseline typically requires more time to
converge, but also converges to worse models in terms of regret, as seen in
the previous paragraph. Furthermore, “comb” and “comb+mse” usually require
more time to converge, which highlights the importance of an effective
weighting scheme for multi-tasks approaches.

###  5.4 Learning under Data Scarcity

(a)  (b)

Figure 7: More Training Data for Graph Routing: Average performance for
different graph routing tasks on the test set for SPO+ (left) and PFYL
(right), trained with PFYL and $1000$ training data points, including regrets
and cost MSE, lower is better. SP $i$ is the regret for shortest path task
$i$, TSP $i$ is the regret for traveling salesperson task $i$, MSE is the mean
squared error of cost coefficients.

In this section, we claim that the multi-task end-to-end predict-then-optimize
framework is particularly effective in the small-data regime. Compared to
Figures 4, we find that multi-task learning for graph routing loses its
advantage with more training data (Figure 7). In Warcraft shortest path
problem, Figure 8 shows that the performance of the separated single-task
model gradually improves and may even surpass multi-task learning as the
amount of training data increases. These figures show that multi-task end-to-
end predict-then-optimize can effectively leverage information from related
datasets when the size of the individual dataset is limited. Therefore, multi-
task learning is a reliable option under data scarcity.

(a)  (b)  (c)

Figure 8: Performance on Warcraft Shortest Path: Average regrets of different
strategies, trained with PFYL, decreases as the amount of training data
increases; lower is better.

###  5.5 Learning under Task Redundancy

(a)  (b)

(c)  (d)

Figure 9: Performance on Graph Routing: Average regrets of a different
strategy, trained with PFYL and $100$ training data, decreases as the amount
of tasks increases, lower is better.

Figure 9 indicates that increasing the number of related tasks improves the
model performance, especially for complicated tasks such as TSP. This
performance improvement can be attributed to the positive interaction between
the losses of the related tasks. This finding suggests the potential for using
auxiliary tasks to enhance model performance.

##  6 Conclusion

We extend the end-to-end predict-then-optimize framework to multi-task
learning, which jointly minimizes decision error for related optimization
tasks. Our results demonstrate the benefits of this approach, including an
improved performance with less training data points and the ability to handle
multiple tasks simultaneously. Future work in this area could include the
application of this method to real-world problems, as well as further
exploration of techniques for multi-task learning, such as current and novel
multi-task neural network architectures and gradient calibration methods.

## References

  * Agrawal et al. [2019] Agrawal, A., Amos, B., Barratt, S., Boyd, S., Diamond, S., Kolter, J.Z.: Differentiable convex optimization layers. In: Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., Garnett, R. (eds.) Advances in Neural Information Processing Systems, vol. 32, Curran Associates, Inc. (2019)
  * Ahuja and Orlin [2001] Ahuja, R.K., Orlin, J.B.: Inverse optimization. Operations Research 49(5), 771–783 (2001)
  * Amos and Kolter [2017] Amos, B., Kolter, J.Z.: Optnet: Differentiable optimization as a layer in neural networks. In: International Conference on Machine Learning, pp. 136–145, PMLR (2017)
  * Bengio [1997] Bengio, Y.: Using a financial training criterion rather than a prediction criterion. International Journal of Neural Systems 8(04), 433–443 (1997)
  * Berthet et al. [2020] Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.P., Bach, F.: Learning with differentiable perturbed optimizers. arXiv preprint arXiv:2002.08676 (2020)
  * Caruana [1997] Caruana, R.: Multitask learning. Machine learning 28(1), 41–75 (1997)
  * Chen et al. [2018] Chen, Z., Badrinarayanan, V., Lee, C.Y., Rabinovich, A.: Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In: International conference on machine learning, pp. 794–803, PMLR (2018)
  * Dalle et al. [2022] Dalle, G., Baty, L., Bouvier, L., Parmentier, A.: Learning with combinatorial optimization layers: a probabilistic approach. arXiv preprint arXiv:2207.13513 (2022)
  * Dantzig et al. [1954] Dantzig, G., Fulkerson, R., Johnson, S.: Solution of a large-scale traveling-salesman problem. Journal of the operations research society of America 2(4), 393–410 (1954)
  * Donti et al. [2017] Donti, P.L., Amos, B., Kolter, J.Z.: Task-based end-to-end model learning in stochastic optimization. Advances in Neural Information Processing Systems (2017)
  * Duong et al. [2015] Duong, L., Cohn, T., Bird, S., Cook, P.: Low resource dependency parsing: Cross-lingual parameter sharing in a neural network parser. In: Proceedings of the 53rd annual meeting of the Association for Computational Linguistics and the 7th international joint conference on natural language processing (volume 2: short papers), pp. 845–850 (2015)
  * Elmachtoub and Grigas [2021] Elmachtoub, A.N., Grigas, P.: Smart “predict, then optimize”. Management Science (2021)
  * Estes and Richard [2023] Estes, A.S., Richard, J.P.P.: Smart predict-then-optimize for two-stage linear programs with side information. INFORMS Journal on Optimization (2023)
  * Ferber et al. [2020] Ferber, A., Wilder, B., Dilkina, B., Tambe, M.: Mipaal: Mixed integer program as a layer. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, pp. 1504–1511 (2020)
  * Ford et al. [2015] Ford, B., Nguyen, T., Tambe, M., Sintov, N., Delle Fave, F.: Beware the soothsayer: From attack prediction accuracy to predictive reliability in security games. In: International Conference on Decision and Game Theory for Security, pp. 35–56, Springer (2015)
  * Gurobi Optimization, LLC [2021] Gurobi Optimization, LLC: Gurobi Optimizer Reference Manual (2021), URL https://www.gurobi.com
  * Hu et al. [2023] Hu, X., Lee, J.C., Lee, J.H.: Predict+ optimize for packing and covering lps with unknown parameters in constraints. In: Proceedings of the AAAI Conference on Artificial Intelligence (2023)
  * Kendall et al. [2018] Kendall, A., Gal, Y., Cipolla, R.: Multi-task learning using uncertainty to weigh losses for scene geometry and semantics. In: Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7482–7491 (2018)
  * Liu et al. [2021] Liu, B., Liu, X., Jin, X., Stone, P., Liu, Q.: Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems 34, 18878–18890 (2021)
  * Liu et al. [2019] Liu, S., Johns, E., Davison, A.J.: End-to-end multi-task learning with attention. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1871–1880 (2019)
  * Ma et al. [2018] Ma, J., Zhao, Z., Yi, X., Chen, J., Hong, L., Chi, E.H.: Modeling task relationships in multi-task learning with multi-gate mixture-of-experts. In: Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1930–1939 (2018)
  * Mandi and Guns [2020] Mandi, J., Guns, T.: Interior point solving for lp-based prediction+optimisation. In: Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M.F., Lin, H. (eds.) Advances in Neural Information Processing Systems, vol. 33, pp. 7272–7282, Curran Associates, Inc. (2020)
  * Mandi et al. [2020] Mandi, J., Stuckey, P.J., Guns, T., et al.: Smart predict-and-optimize for hard combinatorial optimization problems. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 34, pp. 1603–1610 (2020)
  * Mulamba et al. [2020] Mulamba, M., Mandi, J., Diligenti, M., Lombardi, M., Bucarey, V., Guns, T.: Contrastive losses and solution caching for predict-and-optimize. arXiv preprint arXiv:2011.05354 (2020)
  * Paszke et al. [2019] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.: Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems 32 (2019)
  * Shah et al. [2022] Shah, S., Wang, K., Wilder, B., Perrault, A., Tambe, M.: Decision-focused learning without decision-making: Learning locally optimized decision losses. In: Advances in Neural Information Processing Systems (2022)
  * Shazeer et al. [2017] Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., Dean, J.: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 (2017)
  * Tang and Khalil [sion] Tang, B., Khalil, E.B.: Pyepo: A pytorch-based end-to-end predict-then-optimize library for linear and integer programming. Mathematical Programming Computation (2022 in submission)
  * Tang et al. [2020] Tang, H., Liu, J., Zhao, M., Gong, X.: Progressive layered extraction (ple): A novel multi-task learning (mtl) model for personalized recommendations. In: Fourteenth ACM Conference on Recommender Systems, pp. 269–278 (2020)
  * Vlastelica et al. [2019] Vlastelica, M., Paulus, A., Musil, V., Martius, G., Rolínek, M.: Differentiation of blackbox combinatorial solvers. arXiv preprint arXiv:1912.02175 (2019)
  * Wang et al. [2020] Wang, Z., Tsvetkov, Y., Firat, O., Cao, Y.: Gradient vaccine: Investigating and improving multi-task optimization in massively multilingual models. arXiv preprint arXiv:2010.05874 (2020)
  * Wilder et al. [2019] Wilder, B., Dilkina, B., Tambe, M.: Melding the data-decisions pipeline: Decision-focused learning for combinatorial optimization. In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 33, pp. 1658–1665 (2019)
  * Winkenbach et al. [2021] Winkenbach, M., Parks, S., Noszek, J.: Technical proceedings of the amazon last mile routing research challenge (2021)
  * Yu et al. [2020] Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., Finn, C.: Gradient surgery for multi-task learning. Advances in Neural Information Processing Systems 33, 5824–5836 (2020)


# Medical Application of Geometric Deep Learning for the Diagnosis of Glaucoma

Alexandre H. Thiéry  Department of Statistics and Data Science, National
University of Singapore, Singapore  Corresponding authors  Fabian Braeu
Ophthalmic Engineering and Innovation Laboratory, Singapore Eye Research
Institute, Singapore National Eye Centre, Singapore  Yong Loo Lin School of
Medicine, National University of Singapore, Singapore  Tin A. Tun  Singapore
Eye Research Institute, Singapore National Eye Centre, Singapore  Duke-NUS
Graduate Medical School, Singapore  Tin Aung  Yong Loo Lin School of Medicine,
National University of Singapore, Singapore  Singapore Eye Research Institute,
Singapore National Eye Centre, Singapore  Duke-NUS Graduate Medical School,
Singapore  Michaël J.A. Girard  Ophthalmic Engineering and Innovation
Laboratory, Singapore Eye Research Institute, Singapore National Eye Centre,
Singapore  Duke-NUS Graduate Medical School, Singapore  Institute for
Molecular and Clinical Ophthalmology, Basel, Switzerland  Corresponding
authors

###### Abstract

Purpose: (1) To assess the performance of geometric deep learning (PointNet)
in diagnosing glaucoma from a single optical coherence tomography (OCT) 3D
scan of the optic nerve head (ONH); (2) To compare its performance to that
obtained with a standard 3D convolutional neural network (CNN), and with a
gold-standard glaucoma parameter, i.e. retinal nerve fiber layer (RNFL)
thickness.

Methods: 3D raster scans of the ONH were acquired with Spectralis OCT
(Heidelberg Engineering, Heidelberg, Germany) for $477$ glaucoma and $2,296$
non-glaucoma subjects at the Singapore National Eye Centre. All volumes were
automatically segmented using deep learning to identify $7$ major neural and
connective tissues including the RNFL, the prelamina, and the lamina cribrosa
(LC). Each ONH was then represented as a 3D point cloud with $1,000$ points
chosen randomly from all tissue boundaries. To simplify the problem, all ONH
point clouds were aligned with respect to the plane and center of Bruch’s
membrane opening. Geometric deep learning (PointNet) was then used to provide
a glaucoma diagnosis from a single OCT point cloud. The performance of our
approach (reported using the area under the receiver operating characteristic
curve or AUC) was compared to that obtained with a 3D CNN, and with RNFL
thickness.

Results: PointNet was able to provide a robust glaucoma diagnosis solely from
the ONH represented as a 3D point cloud (AUC=$0.95\pm 0.01$). The performance
of PointNet was superior to that obtained with a standard 3D CNN (AUC=$0.87\pm
0.02$) and with that obtained from RNFL thickness alone (AUC=$0.80\pm 0.03$).

Discussion: We provide a proof-of-principle for the application of geometric
deep learning in the field of glaucoma. Our technique requires significantly
less information as input to perform better than a 3D CNN, and with an AUC
superior to that obtained from RNFL thickness alone. Geometric deep learning
may have wide applicability in the field of Ophthalmology.

Keywords: PointNet, Glaucoma, Optical Coherence Tomography, Optic Nerve Head,
Artificial Intelligence

##  1 Introduction

With the development and progression of glaucoma, the optic nerve head (ONH)
typically exhibits complex neural- and connective-tissue structural changes
including, but not limited to: thinning of the retinal nerve fiber layer
(RNFL) and the ganglion cell complex (GCC) layer [1, 2], changes in lamina
cribrosa (LC) shape/depth/curvature [3, 4], and posterior bowing of the
peripapillary sclera [5, 6]. Clinically, optical coherence tomography (OCT) is
the mainstay of imaging to observe such changes [7], however, signal
interpretation (by humans or machines) for glaucoma diagnosis/prognosis
remains a challenge. Recently, a growing number of artificial intelligence
(AI) studies have proposed to use deep learning algorithms to provide a robust
glaucoma diagnosis from a single OCT scan of the ONH or of the macula. Such
applications could have excellent clinical value by e.g. reducing the number
of tests needed to confirm a glaucoma diagnosis. Some of these algorithms were
directly applied to the raw OCT scans [8, 9, 10], while others, first
simplified the images to only a few classes (or colors) by highlighting
relevant tissue structures [11]. Most algorithms achieved good to excellent
performance. However, such algorithms need to be able to handle a considerable
amount of information (e.g. voxel intensities distributed on 3D grids) that
could be heavily corrupted by noise, image artifacts, and 3D image orientation
issues, thus limiting their ease-of-use and deployability.

To this end, a family of algorithms, known as PointNet [12], fitting under the
category of geometric deep learning, has been proposed to solve classification
problems from structures represented as 3D point clouds (such as those in
medical imaging [13]), with excellent performance. Because such point clouds
do not need to be dense and falling onto regular grids, this intrinsically
means the amount of information needed to make e.g. a diagnosis could be
significantly reduced, thus reducing the ‘black-box’ element of AI. For our
glaucoma diagnosis application, the ONH can simply be thought of as a complex
3D structure that can be represented by a cloud of points as has been
routinely performed in 3D histomorphometric and finite element studies [14,
15, 16].

In this study, we aimed to apply a geometric deep learning solution (PointNet)
to provide a robust glaucoma diagnosis from a single OCT scan of the ONH. Each
OCT scan was first pre-processed and each ONH was represented as a 3D point
cloud, thus limiting the amount of information to be processed by several
orders of magnitude. Our approach was compared with a standard 3D
convolutional neural network (CNN), and its performance compared with that
from RNFL thickness alone (current gold standard). It may have the potential
to considerably simplify diagnosis and prognosis applications in the field of
glaucoma.

##  2 Methods

###  2.1 Patient Recruitment

A total of $2,773$ subjects ($477$ glaucoma and $2,296$ non-glaucoma) were
recruited at the Singapore National Eye Center (SNEC, Singapore; see Table 1
for further demographic details). All subjects gave written informed consent.
The study adhered to the tenets of the Declaration of Helsinki and was
approved by the institutional review board of the respective hospital.
Subjects with intraocular pressure (IOP) less than $21$ mmHg, healthy optic
discs with a vertical cup-disc ratio (VCDR) less than or equal to $0.5$, and
normal visual fields tests were considered as non-glaucoma, whereas subjects
with glaucomatous optic neuropathy, VCDR $\geq 0.7$, and/or neuroretinal rim
narrowing with repeatable glaucomatous visual field defects were considered as
glaucoma. Subjects with corneal abnormalities that have the potential to
preclude the quality of the scans were excluded from the study.

Institute |  Age |  Sex (%Male) |  Non-Glaucoma Patients |  Glaucoma Patients |  Non-Glaucoma Scans |  Glaucoma Scans
---|---|---|---|---|---|---
SNEC |  $60.5\pm 9.0$ |  $51\%$ |  $2,296$ |  $477$ |  $3,897$ |  $873$
Total | $2,773$ | $4,770$
Table 1: Summary of patient information.

###  2.2 Optical Coherence Tomography Imaging

Spectral-domain OCT imaging (Spectralis; Heidelberg Engineering, Heidelberg,
Germany) was performed on seated subjects under dark room conditions after
dilation with tropicamide $1\%$ solution. Images were acquired from either
both eyes or one eye of each subject. Each OCT volume consisted of 97 serial
horizontal Bscans ($\sim$30\text{\,}\mathrm{\SIUnitSymbolMicro m}$$ distance
between B-scans; $384$ A-scans per B-scan; $20\times$ signal averaging; axial
resolution: $\sim$3.87\text{\,}\mathrm{\SIUnitSymbolMicro m}$$) that covered a
rectangular area of $15^{\circ}\times 10^{\circ}$ centered on the ONH. The eye
tracking and enhanced depth imaging modalities of the Spectralis were used
during image acquisition. In total, we obtained $4,770$ scans ($873$ glaucoma
and $3,897$ non-glaucoma scans).

###  2.3 Automated Segmentation of OCT Images

Since PointNet requires each ONH to be described as a point cloud, it was
first necessary for us to identify (or highlight) the major neural and
connective tissues that are involved in glaucoma pathogenesis. To this end, we
used the software REFLECTIVITY (Abyss Processing Pte Ltd, Singapore) to
automatically segment the following tissue groups: (1) the retinal nerve fiber
layer (RNFL) and the prelamina; (2) the ganglion cell layer and the inner
plexiform layer (GCL+IPL); (3) all other retinal layers; (4) the retinal
pigment epithelium (RPE) with Bruch’s membrane (BM); (5) the choroid; (6) the
peripapillary sclera including the scleral flange; and (7) the lamina cribrosa
(LC). All segmented tissues can be observed in Figure 1. Note that
REFLECTIVITY was developed from advances in AI-based ONH segmentation as
described in our previous publications [17, 18]. It is also important to point
out that REFLECTIVITY cannot identify the "true" posterior boundaries of the
peripapillary sclera and of the LC, but instead provides the OCT-visible
portions of those 2 tissues. The AI-based segmentation process assigned a
label to each voxel of each 3D OCT scan to indicate the tissue class.

###  2.4 Point cloud generation

Once each voxel of each 3D OCT scan was assigned a label, the voxels situated
at the boundaries between two different tissue groups were automatically
identified. The following N=8 boundaries of interest were identified: anterior
& posterior boundaries of the RNFL+prelamina, and the posterior boundaries of
the IPL+GLC, other retina layers, RPE, choroid, sclera, and LC. On average, a
total of $\sim 50,000$ boundary voxels of interest were extracted from each 3D
OCT scan. The 3D coordinates as well as the boundary class (expressed as a
label varying between $1$ and $N=8$) of each boundary voxel of interest was
recorded for further processing. The 3D coordinates were expressed within an
[$x,y,z$] cartesian coordinate system with origin situated at the center of
the Bruch’s membrane opening (BMO) circle, and such that the BMO circle lies
in the horizontal [$x=0$, $y=0$] plane.

Figure 1: PointNet Workflow: each OCT scan of the ONH is first segmented using
deep learning to identify the following tissue structures: RNFL+prelamina,
IPL+GCL, all other retina layers, RPE, choroid, peripapillary sclera, and
lamina cribrosa. A 3D point cloud is then generated strictly from the tissue
boundaries. The 3D point cloud is ultimately passed through our PointNet
network to produce a glaucoma diagnosis.

###  2.5 PointNet for Glaucoma Diagnosis

In this work, we aimed to design a binary classifier to identify whether a
given ONH would be classified as glaucoma or non-glaucoma. For designing such
a classifier, we opted to use a simple and robust solution that produces a
probabilistic forecast based only on the geometric properties of the ONH.

Geometric Deep Learning [19] is an emerging field of AI that proposes
inductive biases and network architectures that can efficiently process data
structures such as grids, graphs and cloud of points while respecting their
intrinsic symmetries and invariances [20]. In this study, we have leveraged
the recently proposed PointNet neural architecture [12]. This deep neural
network has been especially designed to process point-clouds, i.e. an
unordered set of points. A PointNet takes a point-cloud as input and provides
a unified architecture for applications ranging from object classification,
part segmentation, to scene semantic parsing. Although simple, it has been
empirically demonstrated that the PointNet architecture exhibits a performance
on par, or even better, than the state-of-the-art [12].

For classifying a 3D OCT scan as glaucoma or non-glaucoma, we used a PointNet
that processed a point-cloud of size $S=1,000$. For this purpose, out of the
typically much larger set of points of interest extracted from each OCT scan,
a subset of S=1,000 points were randomly selected. There are two main reasons
for this approach: (1) we have empirically observed that this random sub-
sampling procedure improved the robustness of the method and helped mitigate
overfitting issues; and (2) among all the points of interests extracted from
each 3D OCT scan, there is typically a large amount of redundancy. Reducing
the point-clouds to a random subset of $S=1,000$ points helped improve the
computational efficiency of the method and led to more robust predictions. The
point location (i.e. 3 dimensional coordinates) as well as its boundary class
(represented as a one-hot-encoded vector of dimension $N=8$) were concatenated
into a vector of dimension $D=8+3=11$. In summary, the PointNet was designed
to process point-clouds of size $S=1,000$ and dimension $D=11$.

The dataset of OCT scans was split into training ($70\%$), validation
($15\%$), and test ($15\%$) sets, respectively. The split was performed in
such a way that scans from the same subject did not exist in different sets.
Furthermore, the proportion of glaucoma scans was identical in all the splits.
Our network was then trained on a Nvidia 1080Ti GPU card until optimum
performance was reached in the validation set in about $100$ epochs. For this
purpose, the standard cross-entropy loss was minimized with the ADAM optimizer
[21]. During the training process, each time an OCT scan was processed (i.e.
once during each epoch), a newly generated subset of S=1,000 points was
selected and fed into the PointNet architecture. For enhanced robustness and
to further enrich the training set, subsequent to the random subsampling
process we also used a data-augmentation scheme that applied random rigid
transformations to the subsampled point-clouds. When applying a rigid
transformation to such a point-cloud, the rigid transformation was only
applied to the spatial coordinates (i.e. the first 3 coordinates). To evaluate
the performance of our method, we reported the area under the receiver
operating characteristic curve (ROC-AUC) with uncertainty estimates obtained
from a $5$-fold cross-validation process.

Figure 2: The AUCs were found to be $0.80\pm 0.03$ for RNFL thickness,
$0.87\pm 0.02$ for the 3D CNN approach, and $0.95\pm 0.01$ for the PointNet
method.

###  2.6 Comparisons to 3D Convolutional Neural Network

To compare the performance of our PointNet with a Gold-standard approach we
trained a 3D convolutional neural network (CNN) [8]. For this purpose, each
volume was linearly resampled to isotropic dimensions $128\times 128\times
128$ voxels. The resized OCT scans were fed into a CNN composed of five
3D-convolutional layers, alternated with batch-normalization operations and
ReLU non-linearities. After a final global average pooling layer, a softmax
operation was used to produce a probabilistic binary classification estimate.
The standard cross-entropy loss was minimized with the ADAM optimizer. As
employed during the training of the PointNet , the dataset of OCT scans was
split into training ($70\%$), validation ($15\%$), and test ($15\%$) sets,
respectively. It was important to use data-augmentation during the training
process: random translations and rotations ($\pm 15$ degrees), left-right
flips, additions of Gaussian noise to the voxel intensities. We used early
stopping and selected the network with the highest validation ROC-AUC during
training. We reported the classification ROC-AUC with uncertainty estimates
obtained from a $5$-fold cross-validation process.

###  2.7 Comparisons to RNFL Thickness

To compare the performance of PointNet to that from a gold standard glaucoma
parameter [1], we extracted, from the 3D segmentation of the OCT scans, the
average RNFL thickness at a distance of $1.4$ times the BMO radius. RNFL
thickness was computed from our automated segmentations as the minimum
distance between the anterior and posterior boundaries of the RNFL. We then
reported the diagnostic power as quantified by the AUC. Since the RNFL
thickness is a scalar parameter, no classification algorithm was needed to
compute the AUC.

###  2.8 Results

We evaluated the classification performance of (1) RNFL thickness as a gold
standard glaucoma parameter; of (2) the 3D CNN approach; and of (3) our
proposed PointNet method by performing a $5$-fold cross-validation study. The
three methods were evaluated on the same splits of the data. The AUCs were
found to be $0.80\pm 0.03$ for RNFL Thickness, $0.87\pm 0.02$ for the 3D CNN
approach, and $0.95\pm 0.01$ for the PointNet method (Figure 2).

##  3 Discussion

In this study, we proposed a relatively simple AI approach based on geometric
deep learning to provide a robust glaucoma diagnosis from a single OCT scan of
the ONH. Each ONH was first pre-processed as a 3D point cloud in order to
represent all major neural- and connective-tissue boundaries. Overall, our
proposed approach required significantly less information to perform glaucoma
classification than other 3D AI approaches. It also performed better than a
standard 3D CNN (applied to raw OCT scans), and by taking into account
information about both ONH neural and connective tissue layers, it yielded an
AUC higher than that obtained from RNFL thickness alone.

In this study, we found that geometric deep learning (PointNet) was well
adapted to perform glaucoma classification with ONH tissues represented as
point clouds. In this study, a given ONH was represented with $1,000$ data
points (as a point cloud), which is significantly less than the $\sim
18,000,000$ data points (or voxels) needed to represent an 3D OCT scan. In
other words, we reduced the size of our input by $4$ orders or magnitude.
Dealing with smaller inputs may ultimately allow us to reduce the "black-box"
effect of AI, and allow us to better identify the complex 3D structural and
biomechanical signature of the glaucomatous ONH [22].

It is interesting to note that our PointNet performed better than a standard
3D CNN on the exact same dataset (AUC=$0.95\pm 0.01$ for PointNet vs
AUC=$0.87\pm 0.02$ for 3D CNN). This may not be surprising as 3D OCT scans
typically exhibit a considerable amount of noise and artifacts that would have
been otherwise eliminated through a pre-processing segmentation step, as
proposed herein. In addition, our PointNet was able to focus on the most
important ONH structural features (such as tissue boundaries, and implicitly,
tissue thicknesses), which would be considerably more challenging to fully
exploit for a 3D CNN.

In this study, we found that PointNet was able to provide a higher diagnostic
accuracy (AUC=$0.95\pm 0.01$) as compared to that obtained from a gold
standard glaucoma parameter, i.e. RNFL thickness (AUC=$0.80\pm 0.03$). This
result may not be surprising as the ONH do not only exhibit neural tissue
changes, but also connective tissue changes, such as bending of the
peripapillary sclera [5], and changes in LC morphology and pore shape/pathway
[23, 24]. PointNet has the advantage of capturing some of these features while
minimizing the total amount of information required to establish a diagnosis.
PointNet may also help us identify the contribution of each individual neural
or connective tissue for the diagnosis or prognosis of glaucoma.

Geometric deep learning may have wide applicability in the field of
Ophthalmology. It is relatively attractive for its ease of use, ease (and
speed) of training, and considerably smaller and simpler input size. While our
first application targeted glaucoma, geometric deep learning could also be
used for the diagnosis and prognosis of other optic neuropathies [25], and for
a wide range of corneal [26] and retinal disorders [27].

In this study, several limitations warrant further discussion. First, our
approach was only tested with one OCT device (Spectralis OCT), and for one
Singapore population. Second, we did not consider all structural landmarks
that could have improved the diagnosis of glaucoma such as e.g. the 3D
configuration of the central retinal vessels [28], or the presence of
peripapillary atrophy [29]; neither did we consider other optical properties
[30]. Third, our non-glaucoma population did not include other major optic
neuropathies. The inclusion of such cases could be critical for clinical
translation [31]. Fourth, the segmentation of the posterior LC and
peripapillary sclera boundaries was artificial and was solely based on the
amount of visible signal in the OCT scans, which most likely does not coincide
with the "true" anatomical landmarks [32]. This may have influenced our
diagnostic powers.

In conclusion, we provide herein a proof-of-principle for the application of
geometric deep learning in the field of Ophthalmology with a special emphasis
on glaucoma diagnosis. We found that our technique required significantly less
data as input to perform better than a 3D CNN, and with an AUC superior to
that obtained from RNFL thickness alone. Geometric deep learning may have wide
applicability in the field of Ophthalmology, and it should be explored for
other pathologies.

## Acknowledgments

We acknowledge funding from (1) the donors of the National Glaucoma Research,
a program of the BrightFocus Foundation, for support of this research
(G2021010S [MJAG]); (2) SingHealth Duke-NUS Academic Medicine Research Grant
(SRDUKAMR21A6 [MJAG]); (3) Singapore MOE Tier 1 grant (R155-000-228-114)
[AHT]; (4) the "Retinal Analytics through Machine learning aiding Physics
(RAMP)" project supported by the National Research Foundation, Prime
Minister’s Office, Singapore under its Intra-Create Thematic Grant
"Intersection Of Engineering And Health" - NRF2019-THE002-0006 [MJAG/AT].

## Disclosure

AHT and MJAG are the co-founders of the AI start-up company Abyss Processing
Pte Ltd that provides 3D AI solutions for glaucoma diagnosis and prognosis

## References

  * [1] Akiyasu Kanamori, Makoto Nakamura, Michael FT Escano, Ryu Seya, Hidetaka Maeda, and Akira Negi.  Evaluation of the glaucomatous damage on retinal nerve fiber layer thickness measured by optical coherence tomography.  American journal of ophthalmology, 135(4):513–520, 2003.
  * [2] Jean-Claude Mwanza, Mary K Durbin, Donald L Budenz, Fouad E Sayyad, Robert T Chang, Arvind Neelakantan, David G Godfrey, Randy Carter, and Alan S Crandall.  Glaucoma diagnostic accuracy of ganglion cell–inner plexiform layer thickness: comparison with nerve fiber layer and optic nerve head.  Ophthalmology, 119(6):1151–1158, 2012.
  * [3] Jeong-Ah Kim, Tae-Woo Kim, Robert N Weinreb, Eun Ji Lee, Michaël JA Girard, and Jean Martial Mari.  Lamina cribrosa morphology predicts progressive retinal nerve fiber layer loss in eyes with suspected glaucoma.  Scientific reports, 8(1):1–10, 2018.
  * [4] J Crawford Downs and Christopher A Girkin.  Lamina cribrosa in glaucoma.  Current opinion in ophthalmology, 28(2):113, 2017.
  * [5] Xiaofei Wang, Tin A Tun, Monisha Esther Nongpiur, Hla M Htoon, Yih Chung Tham, Nicholas G Strouthidis, Tin Aung, Ching-Yu Cheng, and Michael JA Girard.  Peripapillary sclera exhibits a v-shaped configuration that is more pronounced in glaucoma eyes.  British Journal of Ophthalmology, 106(4):491–496, 2022.
  * [6] Ya Xing Wang, Hongli Yang, Haomin Luo, Seung Woo Hong, Stuart K Gardiner, Jin Wook Jeoung, Christy Hardin, Glen P Sharpe, Kouros Nouri-Mahdavi, Joseph Caprioli, et al.  Peripapillary scleral bowing increases with age and is inversely associated with peripapillary choroidal thickness in healthy eyes.  American journal of ophthalmology, 217:91–103, 2020.
  * [7] Alexi Geevarghese, Gadi Wollstein, Hiroshi Ishikawa, and Joel S Schuman.  Optical coherence tomography and glaucoma.  Annual review of vision science, 7:693–726, 2021.
  * [8] Stefan Maetschke, Bhavna Antony, Hiroshi Ishikawa, Gadi Wollstein, Joel Schuman, and Rahil Garnavi.  A feature agnostic approach for glaucoma detection in oct volumes.  PloS one, 14(7):e0219126, 2019.
  * [9] An Ran Ran, Carol Y Cheung, Xi Wang, Hao Chen, Lu-yang Luo, Poemen P Chan, Mandy OM Wong, Robert T Chang, Suria S Mannil, Alvin L Young, et al.  Detection of glaucomatous optic neuropathy with spectral-domain optical coherence tomography: a retrospective training and validation deep-learning analysis.  The Lancet Digital Health, 1(4):e172–e182, 2019.
  * [10] Daniel B Russakoff, Suria S Mannil, Jonathan D Oakley, An Ran Ran, Carol Y Cheung, Srilakshmi Dasari, Mohammed Riyazzuddin, Sriharsha Nagaraj, Harsha L Rao, Dolly Chang, et al.  A 3d deep learning system for detecting referable glaucoma using full oct macular cube scans.  Translational Vision Science & Technology, 9(2):12–12, 2020.
  * [11] Satish K Panda, Haris Cheong, Tin A Tun, Sripad K Devella, Vijayalakshmi Senthil, Ramaswami Krishnadas, Martin L Buist, Shamira Perera, Ching-Yu Cheng, Tin Aung, et al.  Describing the structural phenotype of the glaucomatous optic nerve head using artificial intelligence.  American journal of ophthalmology, 236:172–182, 2022.
  * [12] Charles R Qi, Hao Su, Kaichun Mo, and Leonidas J Guibas.  Pointnet: Deep learning on point sets for 3d classification and segmentation.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 652–660, 2017.
  * [13] Benjamín Gutiérrez-Becker and Christian Wachinger.  Deep multi-structural shape analysis: application to neuroanatomy.  In International Conference on Medical Image Computing and Computer-Assisted Intervention, pages 523–531. Springer, 2018.
  * [14] Hongli Yang, Juan Reynaud, Howard Lockwood, Galen Williams, Christy Hardin, Luke Reyes, Stuart K Gardiner, and Claude F Burgoyne.  3d histomorphometric reconstruction and quantification of the optic nerve head connective tissues.  In Glaucoma, pages 207–267. Springer, 2018.
  * [15] Ian A Sigal, John G Flanagan, Inka Tertinegg, and C Ross Ethier.  Finite element modeling of optic nerve head biomechanics.  Investigative ophthalmology & visual science, 45(12):4378–4387, 2004.
  * [16] Yuejiao Jin, Xiaofei Wang, Sylvi Febriana Rachmawati Irnadiastputri, Rosmin Elsa Mohan, Tin Aung, Shamira A Perera, Craig Boote, Jost B Jonas, Leopold Schmetterer, and Michaël J A Girard.  Effect of changing heart rate on the ocular pulse and dynamic biomechanical behavior of the optic nerve head.  Investigative ophthalmology & visual science, 61(4):27–27, 2020\.
  * [17] Sripad Krishna Devalla, Prajwal K Renukanand, Bharathwaj K Sreedhar, Giridhar Subramanian, Liang Zhang, Shamira Perera, Jean-Martial Mari, Khai Sing Chin, Tin A Tun, Nicholas G Strouthidis, et al.  Drunet: a dilated-residual u-net deep learning network to segment optic nerve head tissues in optical coherence tomography images.  Biomedical optics express, 9(7):3244–3265, 2018.
  * [18] Sripad Krishna Devalla, Tan Hung Pham, Satish Kumar Panda, Liang Zhang, Giridhar Subramanian, Anirudh Swaminathan, Chin Zhi Yun, Mohan Rajan, Sujatha Mohan, Ramaswami Krishnadas, et al.  Towards label-free 3d segmentation of optical coherence tomography images of the optic nerve head using deep learning.  Biomedical optics express, 11(11):6356–6378, 2020.
  * [19] Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst.  Geometric deep learning: going beyond euclidean data.  IEEE Signal Processing Magazine, 34(4):18–42, 2017.
  * [20] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip.  A comprehensive survey on graph neural networks.  IEEE transactions on neural networks and learning systems, 32(1):4–24, 2020.
  * [21] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [22] Ian A Sigal and C Ross Ethier.  Biomechanics of the optic nerve head.  Experimental eye research, 88(4):799–807, 2009.
  * [23] Bo Wang, Katie A Lucy, Joel S Schuman, Ian A Sigal, Richard A Bilonick, Chen Lu, Jonathan Liu, Ireneusz Grulkowski, Zachary Nadler, Hiroshi Ishikawa, et al.  Tortuous pore path through the glaucomatous lamina cribrosa.  Scientific reports, 8(1):1–7, 2018.
  * [24] Takuhei Shoji, Hiroto Kuroda, Masayuki Suzuki, Hisashi Ibuki, Makoto Araie, and Shin Yoneya.  Glaucomatous changes in lamina pores shape within the lamina cribrosa using wide bandwidth, femtosecond mode-locked laser oct.  Plos one, 12(7):e0181675, 2017.
  * [25] Michaël JA Girard, Satish K Panda, Tin Aung Tun, Elisabeth A Wibroe, Raymond P Najjar, Aung Tin, Alexandre H Thiéry, Steffen Hamann, Clare Fraser, and Dan Milea.  3d structural analysis of the optic nerve head to robustly discriminate between papilledema and optic disc drusen.  arXiv preprint arXiv:2112.09970, 2021.
  * [26] Valentin Aranha Dos Santos, Leopold Schmetterer, Hannes Stegmann, Martin Pfister, Alina Messner, Gerald Schmidinger, Gerhard Garhofer, and René M Werkmeister.  Corneanet: fast segmentation of cornea oct scans of healthy and keratoconic eyes using deep learning.  Biomedical optics express, 10(2):622–641, 2019.
  * [27] Ursula Schmidt-Erfurth, Gregor S Reiter, Sophie Riedl, Philipp Seeböck, Wolf-Dieter Vogl, Barbara A Blodi, Amitha Domalpally, Amani Fawzi, Yali Jia, David Sarraf, et al.  Ai-based monitoring of retinal fluid in disease activity and under therapy.  Progress in retinal and eye research, page 100972, 2021.
  * [28] Satish K Panda, Haris Cheong, Tin A Tun, Thanadet Chuangsuwanich, Aiste Kadziauskiene, Vijayalakshmi Senthil, Ramaswami Krishnadas, Martin L Buist, Shamira Perera, Ching-Yu Cheng, et al.  The three-dimensional structural configuration of the central retinal vessel trunk and branches as a glaucoma biomarker.  American Journal of Ophthalmology, 2022.
  * [29] Ya Xing Wang, Songhomitra Panda-Jonas, and Jost B Jonas.  Optic nerve head anatomy in myopia and glaucoma, including parapapillary zones alpha, beta, gamma and delta: histology and clinical features.  Progress in retinal and eye research, 83:100933, 2021.
  * [30] Christopher Kai Shun Leung, Alexander Ka Ngai Lam, Robert Neal Weinreb, David F Garway-Heath, Marco Yu, Philip Yawen Guo, Vivian Sheung Man Chiu, Kelvin Ho Nam Wan, Mandy Wong, Ken Zhongheng Wu, et al.  Diagnostic assessment of glaucoma and non-glaucomatous optic neuropathies via optical texture analysis of the retinal nerve fibre layer.  Nature Biomedical Engineering, pages 1–12, 2022.
  * [31] Lama A Al-Aswad, Rithambara Ramachandran, Joel S Schuman, Felipe Medeiros, Malvina B Eydelman, et al.  Artificial intelligence for glaucoma: Creating and implementing ai for disease detection and progression.  Ophthalmology Glaucoma, 2022.
  * [32] Michaël JA Girard, Nicholas G Strouthidis, C Ross Ethier, and Jean Martial Mari.  Shadow removal and contrast enhancement in optical coherence tomography images of the human optic nerve head.  Investigative ophthalmology & visual science, 52(10):7738–7748, 2011.