# Reconstructing Training Data from
Multiclass Neural Networks

Gon Buzaglo 1 , Niv Haim11footnotemark: 1 1 , Gilad Yehudai1, Gal Vardi2 &
Michal Irani1
1Weizmann Institute of Science, Rehovot, Israel
2TTI Chicago and the Hebrew University of Jerusalem Equal Contribution

###### Abstract

Reconstructing samples from the training set of trained neural networks is a
major privacy concern. Haim et al. (2022) recently showed that it is possible
to reconstruct training samples from neural network binary classifiers, based
on theoretical results about the implicit bias of gradient methods. In this
work, we present several improvements and new insights over this previous
work. As our main improvement, we show that training-data reconstruction is
possible in the multi-class setting and that the reconstruction quality is
even higher than in the case of binary classification.

Moreover, we show that using weight-decay during training increases the
vulnerability to sample reconstruction. Finally, while in the previous work
the training set was of size at most $1000$ from $10$ classes, we show
preliminary evidence of the ability to reconstruct from a model trained on
$5000$ samples from $100$ classes.

Plane |  Car |  Bird |  Cat |  Deer |  Dog |  Frog |  Horse |  Ship |  Truck
---|---|---|---|---|---|---|---|---|---
Figure 1: Reconstructed training samples from a multi-class MLP classifier
that was trained on $500$ CIFAR10 images. Each column corresponds to one class
and shows the $10$ training samples (red) that were best reconstructed from
this class, along with their reconstructed result (blue).

##  1 Introduction

Understanding memorization in data-driven machine learning models is a
fundamental question with implications on explainability, privacy, artistic
synthesis and more. Haim et al. (2022) recently demonstrated that a large
portion of the training samples are encoded in the parameters of trained
neural network binary classifiers, by explicitly reconstructing samples from
the training set of such models. The reconstruction method is based on
consequences of the implicit bias of gradient descent, as presented by Lyu &
Li (2019); Ji & Telgarsky (2020) – a homogeneous neural network trained with
gradient descent will converge (in direction) to a solution of the KKT
conditions of a maximum-margin problem. This dictates a set of equations that
relates the parameters of the trained network and the training data. The key
observation is that given a trained model (and its parameters), these
relations can be leveraged to reconstruct training samples. Thus, a loss
function is devised to show that by changing the inputs to the classifier (in
order to minimize the loss), the inputs converge to true samples from the
original training set.

The work of Haim et al. (2022) had several limitations, namely: (1) Their work
only showed reconstructions from binary classifiers; (2) For their
reconstruction method to succeed the trained network needed to be initialized
with very small weights, smaller than standard Xavier initialization (Glorot &
Bengio, 2010); and (3) Their experiments consisted of only small datasets with
at most $1000$ samples. In this work we extend their results in several
directions, and overcome some of the limitations while gaining new insights on
this reconstruction method.

##### Our contributions:

(1) We show reconstruction of large portions of actual training samples from a
trained _multi-class_ neural networks; (2) We show that the use of weight-
decay enables reconstruction of samples from models trained with standard
initialization schemes thus overcoming a major limitation of Haim et al.
(2022); (3) We show reconstruction of models trained on datasets that are
$10$x larger than shown in Haim et al. (2022); (4) We empirically analyze the
effect of weight-decay on sample reconstruction, showing that it increases the
vulnerability to such attacks.

##### Related works.

Several works have shown different privacy attacks in deep learning
architectures, which aim to leak private information from trained models. For
example, _Model inversion_ attacks aim at reconstructing class representatives
Fredrikson et al. (2015); He et al. (2015); Yang et al. (2019). Other types of
attacks target specific models, such as extracting data from _language models_
Carlini et al. (2019; 2021), which use crafted prompts; and information
leakage from collaborative deep learning (federated learning) He et al.
(2019); Melis et al. (2019); Huang et al. (2021); Hitaj et al. (2017). In
Balle et al. (2022) a reconstruction attack is shown where the attacker knows
all the training samples except for one. Recently, Carlini et al. (2023)
showed extraction of actual training images from trained diffusion models.
Their methods rely on generating many different images and using known
membership inference attacks to determine which generated image was used as a
training sample. We emphasize that their method is specific for generative
models, whereas we focus on classifiers.

##  2 Preliminaries - Implicit Bias of Gradient Methods

Neural networks are commonly trained using gradient methods, and when large
enough, they are expected to fit the training data well. However, it is
empirically known that these models converge to solutions that also generalize
well to unseen data, despite the risk of overfitting. Several works pointed to
the “ _implicit bias_ ” of gradient methods as a possible explanation. One of
the most prominent results in this area is by Soudry et al. (2018), who showed
that linear classifiers trained with gradient descent on the logistic loss
converge to the same solution as that of a hard-SVM, meaning that they
maximize the margins. This result was later extended to non-linear and
homogeneous neural networks by Lyu & Li (2019); Ji & Telgarsky (2020). Based
on these results, Haim et al. (2022) have devised a data reconstruction scheme
from trained binary classifiers (see Section 3 in Haim et al. (2022)). Below
we describe an extension of the theorem about the implicit bias of homogeneous
neural networks to a multi-class setup, based on theoretical results from
Appendix G in Lyu & Li (2019).

Formally, let
$S=\\{(\mathbf{x}_{i},y_{i})\\}_{i=1}^{n}\subseteq{\mathbb{R}}^{d}\times[C]$
be a multi-class classification training set where $C\in\mathbb{N}$ is any
number of classes, and $[C]=\\{1,\dots,C\\}$. Let
$\Phi({\bm{\theta}};\cdot):{\mathbb{R}}^{d}\to{\mathbb{R}}^{C}$ be a neural
network parameterized by ${\bm{\theta}}\in{\mathbb{R}}^{p}$. We denote the
$j$-th output of $\Phi$ on an input $\mathbf{x}$ as
$\Phi_{j}({\bm{\theta}};\mathbf{x})\in{\mathbb{R}}$. Consider a homogeneous
network, minimizing the standard cross-entropy loss and assume that after some
number of iterations the model correctly classifies all the training examples.
Then, gradient flow will converge to a KKT point of the following maximum-
margin problem:

$\min_{{\bm{\theta}}}\frac{1}{2}\left\|{\bm{\theta}}\right\|^{2}\;\;\;\;\text{s.t.
}\;\;\;\Phi_{y_{i}}({\bm{\theta}};\mathbf{x}_{i})-\Phi_{j}({\bm{\theta}};\mathbf{x}_{i})\geq
1~{}~{}~{}\forall i\in[n],\forall j\in[C]\setminus\\{y_{i}\\}\;\;~{}.$ (1)

This KKT point is characterized by the following set of conditions:

$\displaystyle{\bm{\theta}}-\sum_{i=1}^{n}\sum_{j\neq{y_{i}}}^{c}\lambda_{i,j}\nabla_{{\bm{\theta}}}(\Phi_{y_{i}}({\bm{\theta}};\mathbf{x}_{i})-\Phi_{j}({\bm{\theta}};\mathbf{x}_{i}))={\mathbf{0}}~{}$
(2) $\displaystyle\forall i\in[n],\forall
j\in[C]\setminus\\{y_{i}\\}:\;\;\Phi_{y_{i}}({\bm{\theta}};\mathbf{x}_{i})-\Phi_{j}({\bm{\theta}};\mathbf{x}_{i})\geq
1$ (3) $\displaystyle\forall i\in[n],\forall
j\in[C]\setminus\\{y_{i}\\}:\;\;\lambda_{i,j}\geq 0$ (4) $\displaystyle\forall
i\in[n],\forall
j\in[C]\setminus\\{y_{i}\\}:\;\;\lambda_{i,j}=0~{}\text{if}~{}\Phi_{y_{i}}({\bm{\theta}};\mathbf{x}_{i})-\Phi_{j}({\bm{\theta}};\mathbf{x}_{i})\neq
1$ (5)

##  3 Multi-Class Reconstruction

We use a similar reconstruction method to Haim et al. (2022). Suppose we are
given a trained classifier with parameters ${\bm{\theta}}$, our goal is to
find the set of data samples $\\{\mathbf{x}_{i}\\}_{i=1}^{n}$ that the network
trained on. A straightforward approach for such a loss would be to minimize
the norm of the L.H.S of condition eq. 2. That is, we initialize
$\\{\mathbf{x}_{i}\\}_{i=1}^{m}$ and
$\\{\lambda_{i,j}\\}_{i\in[n],j\in[C]\setminus y_{i}}$ where $m$ is a
hyperparameter.

Note that from eqs. 3 and 5, most $\lambda_{i,j}$ zero out: the distance of a
sample $\mathbf{x}_{i}$ to its nearest decision boundary,
$\Phi_{y_{i}}-\max_{j\neq{y_{i}}}\Phi_{j}$, is usually achieved for a single
class $j$ and therefore (from eq. 5) in this case at most one $\lambda_{i,j}$
will be non-zero. For some samples $\mathbf{x}_{i}$ it is also possible that
all $\lambda_{i,j}$ will vanish. We therefore turn to only optimizing on the
distance from the decision boundary. This implicitly includes eq. 5 into the
summation in eq. 2, dramatically reducing the number of summands, and
simplifying the overall optimization problem. We define the following loss:

$L_{st}(\mathbf{x}_{1},...,\mathbf{x}_{m},\lambda_{1},...,\lambda_{m})=\left\|{\bm{\theta}}-\sum_{i=1}^{m}\lambda_{i}\
\nabla_{\bm{\theta}}[\Phi_{y_{i}}(\mathbf{x}_{i};{\bm{\theta}})-\max_{j\neq{y_{i}}}\Phi_{j}(\mathbf{x}_{i};{\bm{\theta}})]\right\|_{2}^{2}$
(6)

Our reconstruction method is, given the parameters of a trained network
${\bm{\theta}}$, initialize $\mathbf{x}_{i}$ and $\lambda_{i}$ for
$i=1,\dots,m$, and minimize equation 6. In order to make sure that eq. 4 is
satisfied we optimize for $a_{i}$ and require that
$\lambda_{i}=a_{i}^{2}$111More precisely,
$\lambda_{i}=a_{i}^{2}+\lambda_{\text{min}}$, where $\lambda_{\text{min}}$ is
a hyperparameter that encourages the reconstructed samples to converge to
margin-samples.. This further simplifies the optimization problem compared to
Haim et al. (2022) that use a separate loss function.

Since $n$ is unknown we set $m\geq n$ which represents the number of samples
we want to reconstruct (thus, we only need to upper bound $n$). We can
hypothetically set $m=C\cdot n$ and with balanced labels, this way there are
enough reconstructed samples for any distribution of the labels. In practice,
we set $m$ to be slightly larger than $n$, which is enough to get good
reconstructions. The rest of the hyperparameters (including $\lambda_{\min}$)
are chosen using a hyper-parameter search.

##  4 Results

###  4.1 Multiclass Reconstruction

We compare between reconstruction from binary classifiers (as studied in Haim
et al. (2022)) and multi-class classifiers. We conduct the following
experiment: we train an MLP classifier with architecture $D$-$1000$-$1000$-$C$
on $500$ samples from the CIFAR10 (Krizhevsky et al., 2009) dataset. We use
full-batch GD, and set the learning rate as $0.5$. The model is trained to
minimize the cross-entropy loss with full-batch gradient descent, once with
two classes ($250$ samples per class) and once for the full $10$ classes ($50$
samples per class). The test set accuracy of the models is $77\%$/$32\%$
respectively, which is far from random ($50\%$/$10\%$ resp.).

To quantify the quality of our reconstructed samples, for each sample in the
original training set we search for its nearest neighbour in the reconstructed
images and measure the similarity using SSIM (Wang et al., 2004) (higher SSIM
means better reconstruction). As a rule of thumb, we say that a sample was
reconstructed well if its SSIM$>0.4$ (see discussion in appendix A). In fig. 2
we plot the quality of reconstruction (in terms of SSIM) against the distance
of the sample from the decision boundary
$\Phi_{y_{i}}(\mathbf{x}_{i};{\bm{\theta}})-\max_{j\neq{y_{i}}}\Phi_{j}(\mathbf{x}_{i};{\bm{\theta}})$.
As seen, a multi-class classifier yields much more samples that are vulnerable
to being reconstructed.

Figure 2: Multi-class classifiers are more vulnerable to training-set
reconstruction. For a training set of size $500$, a multi-class model (left)
yields much more reconstructed samples with good quality (SSIM$>0.4$), than a
binary classification model (right).

Next, we examine the dependence between the ability to reconstruct from a
model and the number of classes on which it was trained. Comparing between two
models trained on different number of classes is not immediately clear, since
we want to isolate the effect of the number of classes from the size of the
dataset (it was observed by Haim et al. (2022) that the number of
reconstructed samples decreases as the total size of the training set
increases). We therefore train models on training sets with varying number of
classes ($C\in\\{2,3,4,5,10\\}$) and varying number of samples per class
($1,5,10,50$). The results are visualized in fig. 3a. Note that for models
with same number of samples per class, the ability to reconstruct _increases_
with the number of classes, even though the total size of the training set is
larger. This further validates our hypothesis that multi-class models are more
vulnerable to reconstruction. We continue this study in appendix B.

Another way to validate this hypothesis is by showing the dependency between
the number of classes and the number of “good” reconstructions (SSIM$>0.4$) –
shown in fig. 3b. As can be seen, training on multiple classes yields more
samples that are vulnerable to reconstruction. A possible intuitive
explanation, is that multi-class classifiers have more “margin” samples. Since
margin-samples are more vulnerable to reconstruction, this results in more
samples being reconstructed from the model.

---
(a) Full results of each experiment
(b) Number of “good” reconstructions increases with number of classes and the
samples per class
Figure 3: Evaluating effect of multiple classes on the ability to reconstruct.
We show reconstructions from models trained with different number of classes
and different number of samples per class. As seen, multiple classes result in
more reconstructed samples.

###  4.2 Weight Decay Increases Reconstructability

Figure 4: Using weight-decay during training increases vulnerability to sample
reconstruction

Another drawback of Haim et al. (2022) is that reconstruction was only shown
for models whose first fully-connected layer was initialized with small (non-
standard) weights. Models with standard initialization, such as Kaiming He et
al. (2015) or Xavier Glorot & Bengio (2010) where each weight vector is
initialized with a variance of $\sim\frac{1}{d}$ (where $d$ is the input’s
layer dimension) did not yield good reconstructed samples. Haim et al. (2022)
only reconstructed from networks initialized with a variance of
$\sim\frac{1}{d^{1.5}}$. Set to better understand this drawback, we observed
an interesting phenomenon – for models with standard initialization, using
weight-decay during training enabled samples reconstruction. Moreover, for
some values of weight-decay the reconstructability is _significantly higher_
than what was observed for models with small-initialized-first-layer models.
In fig. 4 we show the number of good reconstructed samples for different
choices of the value of the weight decay ($\lambda_{\text{WD}}$). We show
results for two models trained on $C=2$ classes (fig. 4left) and $C=10$
classes (fig. 4right), both trained on $50$ samples per class. We add two
baselines trained without weight-decay: model trained with standard
initialization (black) and model with small-initialized-first-layer (red).

By looking at the exact distribution of reconstruction quality to the distance
from the margin, we observe that weight-decay (for some values) results in
more training samples being on the margin of the trained classifier, thus
being more vulnerable to reconstruction. This observation is shown in fig. 5
where we show the plots for all experiments from fig. 4left. We also provide
the train and test errors for each model. It seems that the generalization
(test error) does not change significantly. However, an interesting
observation is that reconstruction is possible even for models with non-zero
training errors (for which, the assumptions of Lyu & Li (2019) do not hold).

Figure 5: Weight-Decay “pushes” more samples to the margin, thus enabling them
to be reconstructed

###  4.3 Reconstruction From a Larger Number of Samples

One of the major limitations of Haim et al. (2022) is that they reconstruct
from models that trained on a relatively small number of samples.
Specifically, in their largest experiment, a model is trained with only $1000$
samples. Here we take a step further, and apply our reconstruction scheme for
a model trained on $5000$ data samples.

To this end, we trained a 3-layer MLP, where the number of neurons in each
hidden layer is $10,000$. Note that the size of the hidden layer is $10$ times
larger than in any other model we used. Increasing the number of neurons seems
to be one of the major reasons for which we are able to reconstruct from such
large datasets, although we believe it could be done with smaller models,
which we leave for future research. We used the CIFAR100 dataset, with 50
samples in each class, for a total of $5000$ samples.

In fig. 6a we give the best reconstructions of the model. Note that although
there is a degradation in the quality of the reconstruction w.r.t a model
trained on less samples, it is still clear that our scheme can reconstruct
some of the training samples to some extent. In fig. 6b we show a scatter plot
of the SSIM score w.r.t the distance from the boundary, similar to fig. 3a.
Although most of the samples are on or close to the margin, only a few dozens
achieve an SSIM$>0.4$. This may indicate that there is a potential for much
more images to reconstruct, and possibly with better quality.

---
(a) Full Images. Original samples from the training set (red) and
reconstructed results (blue)
(b) Scatter plot (similar to fig. 3).
Figure 6: Reconstruction from a model trained on $50$ images per class from
the CIFAR100 dataset ($100$ classes, total of $5000$ datapoints)

##  5 Conclusions and Future Work

In this paper we have shown several improvements over Haim et al. (2022). Most
notably, we have shown that it is possible to reconstruct training data in a
multi-class setting, compared to only a binary classification setting in the
previous work. Additionally, Haim et al. (2022) showed reconstructions only
from networks trained with small initialization. Here we show reconstructions
from networks trained with weight decay, which is much more standard than a
small initialization scale. Finally, we show it is possible to reconstruct
from models trained on $5000$ data samples, which is 5 times more than the
largest trained model in Haim et al. (2022)

There are a couple of future research directions that we think might be
interesting. First, to extend our reconstruction scheme to more practical
models such as CNN and ResNets. Second, to reconstruct from models trained on
more data, such as the entire CIFAR 10/100 datasets, or different types of
data such as text, time-series or tabular data. Finally, it would be
interesting to find privacy schemes which could protect from reconstruction
attacks by specifically protecting samples which lie on the margin.

#### Acknowledgments

This project received funding from the European Research Council (ERC) under
the European Union’s Horizon 2020 research and innovation programme (grant
agreement No 788535), and ERC grant 754705, and from the D. Dan and Betty Kahn
Foundation. GV acknowledges the support of the NSF and the Simons Foundation
for the Collaboration on the Theoretical Foundations of Deep Learning.

## References

  * Balle et al. (2022) Borja Balle, Giovanni Cherubin, and Jamie Hayes.  Reconstructing training data with informed adversaries.  _arXiv preprint arXiv:2201.04845_ , 2022.
  * Carlini et al. (2019) Nicholas Carlini, Chang Liu, Úlfar Erlingsson, Jernej Kos, and Dawn Song.  The secret sharer: Evaluating and testing unintended memorization in neural networks.  In _28th USENIX Security Symposium (USENIX Security 19)_ , pp. 267–284, 2019.
  * Carlini et al. (2021) Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert-Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, et al.  Extracting training data from large language models.  In _30th USENIX Security Symposium (USENIX Security 21)_ , pp. 2633–2650, 2021.
  * Carlini et al. (2023) Nicholas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramèr, Borja Balle, Daphne Ippolito, and Eric Wallace.  Extracting training data from diffusion models.  _arXiv preprint arXiv:2301.13188_ , 2023.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE conference on computer vision and pattern recognition_ , pp. 248–255. Ieee, 2009.
  * Fredrikson et al. (2015) Matt Fredrikson, Somesh Jha, and Thomas Ristenpart.  Model inversion attacks that exploit confidence information and basic countermeasures.  In _Proceedings of the 22nd ACM SIGSAC conference on computer and communications security_ , pp. 1322–1333, 2015.
  * Glorot & Bengio (2010) Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 249–256. JMLR Workshop and Conference Proceedings, 2010.
  * Haim et al. (2022) Niv Haim, Gal Vardi, Gilad Yehudai, Ohad Shamir, and Michal Irani.  Reconstructing training data from trained neural networks.  _NeurIPS_ , 2022.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Delving deep into rectifiers: Surpassing human-level performance on imagenet classification.  In _Proceedings of the IEEE international conference on computer vision_ , pp. 1026–1034, 2015.
  * He et al. (2019) Zecheng He, Tianwei Zhang, and Ruby B Lee.  Model inversion attacks against collaborative inference.  In _Proceedings of the 35th Annual Computer Security Applications Conference_ , pp. 148–162, 2019.
  * Hitaj et al. (2017) Briland Hitaj, Giuseppe Ateniese, and Fernando Perez-Cruz.  Deep models under the gan: information leakage from collaborative deep learning.  In _Proceedings of the 2017 ACM SIGSAC conference on computer and communications security_ , pp. 603–618, 2017.
  * Huang et al. (2021) Yangsibo Huang, Samyak Gupta, Zhao Song, Kai Li, and Sanjeev Arora.  Evaluating gradient inversion attacks and defenses in federated learning.  _Advances in Neural Information Processing Systems_ , 34:7232–7241, 2021.
  * Ji & Telgarsky (2020) Ziwei Ji and Matus Telgarsky.  Directional convergence and alignment in deep learning.  _Advances in Neural Information Processing Systems_ , 33:17176–17186, 2020.
  * Krizhevsky et al. (2009) Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Lyu & Li (2019) Kaifeng Lyu and Jian Li.  Gradient descent maximizes the margin of homogeneous neural networks.  _arXiv preprint arXiv:1906.05890_ , 2019.
  * Melis et al. (2019) Luca Melis, Congzheng Song, Emiliano De Cristofaro, and Vitaly Shmatikov.  Exploiting unintended feature leakage in collaborative learning.  In _2019 IEEE Symposium on Security and Privacy (SP)_ , pp. 691–706. IEEE, 2019.
  * Soudry et al. (2018) Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro.  The implicit bias of gradient descent on separable data.  _The Journal of Machine Learning Research_ , 19(1):2822–2878, 2018.
  * Wang et al. (2004) Zhou Wang, Alan C Bovik, Hamid R Sheikh, and Eero P Simoncelli.  Image quality assessment: from error visibility to structural similarity.  _IEEE transactions on image processing_ , 13(4):600–612, 2004.
  * Yang et al. (2019) Ziqi Yang, Jiyi Zhang, Ee-Chien Chang, and Zhenkai Liang.  Neural network inversion in adversarial setting via background knowledge alignment.  In _Proceedings of the 2019 ACM SIGSAC Conference on Computer and Communications Security_ , pp. 225–240, 2019.
  * Zhang et al. (2018) Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang.  The unreasonable effectiveness of deep features as a perceptual metric.  In _CVPR_ , 2018.

##  Appendix A Deciding whether a Reconstruction is “Good”

Here we justify our selection for SSIM$=0.4$ as the threshold for what we
consider as a “good” reconstruction. In general, the problem of deciding
whether a reconstruction is the correct match to a given sample, or whether a
reconstruction is a “good” reconstruction is equivalent to the problem of
comparing between images. No “synthetic” metric (like SSIM, $l2$ etc.) will be
aligned with human perception. A common metric for this purpose is LPIPS Zhang
et al. (2018) that uses a classifier trained on Imagenet Deng et al. (2009),
but since CIFAR images are much smaller than Imagenet images ($32\times 32$
vs. $224\times 224$) it is not clear that this metric will be better. As a
simple rule of thumb, we use SSIM$>0.4$ for deciding that a given
reconstruction is “good”. To justify, we plot the best reconstructions (in
terms of SSIM) in fig. 7. Note that almost all samples with SSIM$>0.4$ are
also visually similar (for a human). Also note that some of the samples with
SSIM$<0.4$ are visually similar, so in this sense we are “missing” some good
reconstructions. In general, determining whether a reconstruction is “good” is
an open question which cannot be dealt in the scope of this paper, and is a
future direction for our work.

---
Figure 7: Justifying the threshold of SSIM$=0.4$ as good rule-of-thumb for a
threshold for a “good” reconstruction. Note that samples with SSIM$>0.4$
(blue) are visually similar. Also some of the samples with SSIM$<0.4$ (red)
are similar. In general deciding whether a reconstruction is “good” is an open
question beyond the scope of this paper. The SSIM values are shown above each
train-reconstruction pair.

##  Appendix B Experiments with Different Number of Classes and Fixed Training
Set Size

To complete the experiment shown in fig. 3, we also perform experiments on
models trained on various number of classes ($C\in\\{2,3,4,5,10\\}$) but this
time with a fixed training set size of $500$ samples (distributed equally
between classes). It seems from fig. 8 that the results are not much different
from those for $50$ samples per class, and we hypothesize that the model only
needs a certain amount of “support vectors” to support its parameters, and
this number is also achieved by $50$ samples per class, and are not harmed if
more data is added for each class. Also note that for models with less
classes, not only the number of good reconstruction decreases, but also the
quality of reconstruction (lower SSIM). Since we don’t have a good heuristic
for aggregating the reconstruction quality score of a given model this
observation is hard to quantify, but evident from the plot itself.

Figure 8: Experiments of reconstruction from models trained on a a fixed
training set size ($500$ samples) for different number of classes. Number of
“good” reconstruction is shown for each model.


# DRStageNet: Deep Learning for Diabetic Retinopathy Staging from Fundus
Images

Yevgeniy Men  Jonathan Fhima  Leo Anthony Celi  Lucas Zago Ribeiro  Luis
Filipe Nakayama and Joachim A. Behar  Senior Member, IEEE YM and JB
acknowledge the support of the Technion EVPR Fund: Irving & Branna Sisenwein
Research Fund. The research was supported by a cloud computing grant from the
Israel Council of Higher Education, administered by the Israel Data Science
Initiative. LAC is funded by the National Institute of Health through NIBIB
R01 EB017205.We would like to acknowledge the assistance of ChatGPT, an AI-
based language model developed by OpenAI, for its help in editing the English
language of this manuscript.YM is with the Andrew and Erna Viterbi Faculty of
Electrical & Computer Engineering and the Faculty of Biomedical Engineering,
Technion, Israel Institute of Technology, Haifa, 3200003, Israel. JF is with
the Department of Applied Mathematics and the Biomedical Engineering Faculty,
Technion, Israel Institute of Technology, Haifa, 3200003, Israel. LAC is with
the Laboratory for Computational Physiology, Massachusetts Institute of
Technology, Cambridge, MA 02139, the Division of Pulmonary, Critical Care and
Sleep Medicine, Beth Israel Deaconess Medical Center, Boston, MA 02215 and the
Department of Biostatistics, Harvard T.H. Chan School of Public Health,
Boston, MA 02115. LZR is with the Ophthalmology department, São Paulo Federal
University, São Paulo, Brazil. LFN is with the Institute for Medical
Engineering and Science, Massachusetts Institute of Technology, Cambridge, MA,
USA and the Department of Ophthalmology, São Paulo Federal University, São
Paulo, Brazil. JAB is with the Faculty of Biomedical Engineering, Technion,
Israel Institute of Technology, Haifa, 3200003, Israel (e-mail:
jbehar@technion.ac.il).

###### Abstract

Diabetic retinopathy (DR) is a prevalent complication of diabetes associated
with a significant risk of vision loss. Timely identification is critical to
curb vision impairment. Algorithms for DR staging from digital fundus images
(DFIs) have been recently proposed. However, models often fail to generalize
due to distribution shifts between the source domain on which the model was
trained and the target domain where it is deployed. A common and particularly
challenging shift is often encountered when the source- and target-domain
supports do not fully overlap. In this research, we introduce DRStageNet, a
deep learning model designed to mitigate this challenge. We used seven
publicly available datasets, comprising a total of 93,534 DFIs that cover a
variety of patient demographics, ethnicities, geographic origins and
comorbidities. We fine-tune DINOv2, a pretrained model of self-supervised
vision transformer, and implement a multi-source domain fine-tuning strategy
to enhance generalization performance. We benchmark and demonstrate the
superiority of our method to two state-of-the-art benchmarks, including a
recently published foundation model. We adapted the grad-rollout method to our
regression task in order to provide high-resolution explainability heatmaps.
The error analysis showed that 59% of the main errors had incorrect reference
labels. DRStageNet is accessible at URL [upon acceptance of the manuscript].

diabetic retinopathy, fundus image, deep learning, self-supervised learning,
transformers.

##  1 Introduction

Diabetes mellitus (DM) is one of the largest public health concerns globally
[1]. According to estimates by the International Diabetes Federation, 536.6
million people had DM in 2021, and prevalence is projected to increase to
783.2 million by 2045 [2]. Diabetic retinopathy (DR) is a direct microvascular
end organ complication of DM. High glucose level caused by DM produces
cytokines and growth factors that lead to capillary damage of eye blood
vessels and causes increased vascular permeability and capillary occlusions.
According to a 2012 study, approximately 34.6% of DM patients suffer some
degree of DR, 10. 2% suffering from vision-threatening DR [3]. Early detection
of the disease is very important and any delay can result in rapid vision
degradation and eventual irreversible blindness. Traditionally, DR is detected
by a manual search for various lesions, including microaneurysms, hemorrhages,
hard and soft exudates and vascular abnormalities [4]. To avoid complications
related to DR, patients with DM are recommended to undergo annual examinations
[5]. The process requires highly skilled practitioners, with developing
countries suffering from an acute shortage of such experts. Recently, deep
learning (DL)-based algorithms for detection of DR from digital fundus images
(DFI) have been suggested to tackle this challenge. Some of these algorithms
focus on DR screening [6, 5, 7, 8, 9, 10], while others focus on DR staging
[11, 12, 13, 14, 15].

DR screening research and commercial devices consider the binary
classification task of referable DR (rDR) [6, 5, 7, 8, 9, 10], which defines
the positive class as a moderate or worse stage on the International Clinical
Diabetic Retinopathy [16] (ICDR) scale or the presence of diabetic macular
edema (DME). These models may be useful for nonophthalmologist professionals
for the purpose of DR screening. DR staging research often models the task as
a multiclass classification problem [11, 12, 13, 14]. Such models can support
retina specialists in diagnosing DR and in monitoring disease progression and
management. They can also be used by nonspecialists for screening. However,
the main drawback of multi-class classification algorithms is the fact that
the task is framed as a classification of multiple independent classes and
does not take into account the ordinal relations between the classes.

The ability of DR staging models to generalize across diverse datasets remains
a challenge. Models often fail to generalize [17, 7] due to distribution
shifts between the source domain on which the model was trained and the target
domain on which it is deployed. A common and particularly challenging shift
often encountered in reality is where the source and target domain supports do
not fully overlap. This commonly occurs with medical datasets that factor in
differences in demographics, ethnicities, geographic origins, and/or
comorbidities as well as technical specifications, in particular the type of
camera and field of view (FOV). Finally, the explainability of DR staging
algorithms is limited due to low causal relation between DR manifestations and
the associated explainability heatmaps; i.e., there is a substantial amount of
false positive and false negative regions that reduces their usability. This
work makes the following contributions:

Figure 1: DRStageNet consists of a DINOv2 [19] pretrained backbone that is
joined with a simple fully connected regression head. Additionally, we utilize
a multi-source domain fine-tuning approach by combining seven open source DR
datasets. DRStageNet is fine-tuned using the MSE loss.

  * •

Introduction of DRStageNet (Fig. 1), a robust, i.e., high-performing and
generalizable, algorithm for DR screening, diagnosis, staging and progression
monitoring.

  * •

Adaptation of the grad-rollout [18] method to generate high-resolution
explainability heatmaps for the regression task.

  * •

Presentation of a detailed error analysis of DRStageNet in seven independent
test sets.

##  2 Datasets

DR has multiple severity scales that are used in different countries and
different clinics. The ICDR is the most common scale in open source DFI
datasets [20]. For DR identification experiments, we selected seven
independent open datasets (Table 1) that had at least 500 DFIs available with
ICDR grading. For each dataset, DFIs labeled nongradable, DFIs with missing
labels, DFIs from children ($<18$ years) and DFIs from nondiabetes patients
were excluded. rDR was defined using the DME labels when available and with
the ICDR scale otherwise.

####  2.0.1 Kaggle EYEPACS (Eye Picture Archive Communication System) [21]

The EYEPACS dataset was provided by the Eye Picture Archive Communication
System [22] and was first introduced in the context of a Kaggle competition
[21]. It contains 88,702 macula-centered DFIs of varying resolutions, which
were captured by different cameras at different sites in the United States.
The DFIs were classified for DR by a single clinician according to the ICDR
scale using only the images as a reference. The dataset is divided into 35,126
DFIs used for training and 53,576 DFIs used for testing. Voets et al. [7]
found that approximately 20% of the dataset are ungradable DFIs and redefined
the dataset split to a training set of 28,132 DFIs from 14,404 patients
(EYEPACS-train) and a test set of 42,922 DFIs from 24,524 patients (EYEPACS-
test). We used this dataset and train-test split. For hyperparameters tuning,
the training set was divided into a 90:10 split while stratifying at the
patient level to avoid information leakage. The EYEPACS-test was used as the
unseen source domain test set.

####  2.0.2 DDR [23]

The DDR dataset includes 13,673 macula-centered DFIs obtained from 9,598
patients from 147 hospitals in China. These images were captured using
multiple cameras with a FOV of $45^{\circ}$. The dataset includes pixel-level
and bounding box annotations of microaneurysms, hemorrhages and soft and hard
exudates, as well as DR grades. Professional graders, who were trained by
ophthalmologists, evaluated the DFIs using the ICDR scale on single-image
level and also assessed their gradability. The final dataset consists of
12,519 DFIs.

####  2.0.3 The Asia Pacific Tele-Ophthalmology Society (APTOS) [24]

APTOS dataset contains a total of 5,590 macula-centered DFIs and it was made
openly accessible by the Aravind Eye Hospital in India through a Kaggle
competition. The images were captured by Aravind technicians in many rural
regions of India, under varying conditions and environments and over a long
period of time. The DFIs were later labeled by a group of doctors according to
the ICDR scale without using additional information. The only accessible
labels are of the APTOS-train split, which consists of 3,662 DFIs that we use
in this research.

####  2.0.4 Brazilian Multilabel Ophthalmological Dataset (BRSET) [25]

This dataset comprises 16,266 DFIs centered on the macula, with a FOV of
$45^{\circ}$, obtained from 8,524 patients examined at two ophthalmology
clinics in Brazil (IRB 0698/2020). It includes demographic information, as
well as anatomical parameters related to the macula, optic disc, and retinal
blood vessels. Image quality parameters such as illumination, image field, and
artifacts were also recorded to ensure quality control. The DFIs were graded
on the single-image level using the ICDR scale. The data subset of diabetic
patients consists of 1,301 individuals and 2,489 DFIs. This data subset was
used for our experiments.

####  2.0.5 Methods to Evaluate Segmentation and Indexing Techniques in the
Field of Retinal Ophthalmology (MESSIDOR2) [8]

The Messidor-2 dataset is a collection of 1,748 macula-centered DFIs obtained
using a Topcon TRC NW6 camera with $45^{\circ}$ FOV. These images are
available in one of three resolutions: 1440 × 960, 2240 × 1488, or 2304 ×
1536. While the dataset lacks official ICDR labels, alternative grading sets
have been introduced by different research groups, all of which were based on
a single DFI. Specifically, Google used a panel of 7 graders [6] who assessed
the images using the ICDR scale. Labels provided by a different panel of 3
graders [26] were publicly released by Google. Additionally, the University of
Iowa provided binary rDR grades111https://tinyurl.com/58sb5rm3. For our
research, the publicly accessible 3-grader annotations provided by Google [26]
were used.

####  2.0.6 The Indian Diabetic Retinopathy Dataset (IDRiD) [27]

IDRiD contains a total of 516 macula-centered DFIs that were acquired at an
eye clinic located in Nanded, (M.S.), India. The DFIs were acquired with a
Kowa VX-10 alpha with a $50^{\circ}$ FOV, and 4,288 × 2,848 pixels. The
dataset contains pixel-level DR lesion and optical disc annotations, as well
as image-based ICDR grade and binary classification of DME.

####  2.0.7 Diabetic Retinopathy Two-field Image Dataset (DRTiD) [28]

The DRTiD dataset contains a total of 3,100 two-field DFIs from 1,550 eyes,
i.e., one macula-centered image and optical disc-centered imaged for each eye.
Images were acquired between 2015 and 2017, using non-mydriatic retinal
cameras with FOVs ranging between $45^{\circ}$ and $50^{\circ}$. The data
acquisition was conducted as part of the Shanghai Diabetic Eye Study. A team
of three experienced ophthalmologists graded the two-field DFIs using the ICDR
scale. Intra-rater annotation discrepancies were reconciled by an expert
ophthalmologist with clinical experience of more than 10 years. We used only
the macula-centered DFIs for our experiments.

Table 1: Description of the datasets used after removing DFIs that met the
exclusion criteria. DR% is the percentage of DR DFI in the dataset defined as
an ICDR equal or superior to one. rDR% is the percentage of rDR DFI in the
dataset defined as an ICDR superior to one. NA denotes information unavailable
and VAR denotes varying devices/fields of view.

Dataset | DFI | Patients |  | Device
---
FOV | Resolution | Targets | DR% | rDR% | Origin
| Kaggle
---
EYEPACS
71,054 | 40,529 | VAR | VAR |  | 433 × 289 to
---
5184 × 3456
ICDR | 26 | 19 | US
DDR | 12,519 | 9,598$<$ | VAR | 45 |  | 702 × 706 to
---
5,184 3,456
ICDR | 50 | 45 | China
APTOS | 3,662 | NA | VAR | VAR |  | 474 × 358 to
---
4,288 2,848
ICDR | 51 | 40 | India
BRSET | 2,489 | 1,301 |  | Nikon NF505
---
Canon CR-2
45 |  | 951 × 874 to
---
2984 × 2304
| ICDR
---
DME
26 | 22 | Brazil
MESSIDOR2 | 1,744 | 874 |  | Topcon
---
TRC NW6
45 |  | 1440 × 960 to
---
2304 × 1536
| ICDR
---
DME
42 | 26 | France
DRTiD | 1,550 | NA | NA | 45-55 |  | 1,444 × 1,444 to
---
3,058 × 3,000
| ICDR
---
DME
52 | 42 | China
IDRiD | 516 | NA |  | Kowa
---
VX-10a
50 | 4288 × 2848 |  | ICDR
---
DME
67 | 62 | India

##  3 Methods

###  3.1 Deep learning for DR staging

####  3.1.1 Preprocessing

First the horizontal black regions of the images are removed and then they are
padded to a squared aspect ratio, based on the longest axis. The images were
then resized to $518\times 518$ pixels using bilinear interpolation. We found
out that this preprocessing treatment obtained better performance than
alternative techniques such as resizing without first cropping them to a
square, or using Ben Graham’s preprocessing technique [29], a preprocessing
treatment which was used by other researchers developing DR algorithms [10,
30, 31].

####  3.1.2 DRStageNet

We approached the challenge of DR detection as a regression task against
reference ICDR annotations [16]. The architecture (Fig. 1) of DRStageNet
consists of a pretrained DINOv2 [19] base vision transformer (ViT-base) [32],
having 86 million trainable parameters and a regression head that consists of
two fully connected layers with a hidden dimension of 512 and a GeLU
activation function [33]. Briefly, DINOv2 is a ViT that was trained on 142
million natural images using self-supervised learning (SSL). We chose to
utilize transfer learning from SSL because it enables learning a
representation based on a much larger dataset of natural images. This
representation can then be fine-tuned for a specific task, i.e., subsequently
training it using supervision. The output of the model is a single scalar and
the total number of trainable parameters is 86.9 million. DRStageNet is fine-
tuned using the mean squared error (MSE) loss function, where the target is
the ICDR grade. The fine-tuning step is performed over the entire network,
that is, without freezing any layer. We used a batch size of 16 DFIs with Adam
[34] optimizer with a $0.04$ weight decay and a learning rate scheduler with
an initial learning rate of $1e^{-6}$ which decreases 10-fold when the
validation loss stopped decreasing over 4 epochs. Early stopping with respect
to the validation loss was also used to reduce overfitting. Additionally, we
used the data augmentations proposed by [7], consisting of random horizontal
flips, jitter of contrast, saturation and hue. In our experiments, these
augmentations proved to be the most stable and most effective compared to
other ensembles of augmentations. We save the weights of the model with the
lowest validation loss.

####  3.1.3 Multi-source domain fine-tuning

We evaluated two methods for model training and evaluation. The first was a
single-source domain (SS) fine-tuning on EYEPACS-train and evaluation on
EYEPACS-test as well as on six external datasets (target domains). This method
is called DRStageNet-SS. The second method uses multi-source domain fine-
tuning (MST). It consists of training on a joint set of multiple source
datasets while evaluating generalization performance on a single left-out
target domain. The intuition behind this second approach is that when training
a model on a single dataset, it may overfit this specific domain distribution.
Variations in data collection equipment and inherent biases in the sample
group, such as age, ethnicity, and health conditions, can cause such a model
to fail when deployed. Furthermore, shortcut learning [35] can cause a model
to recognize misleading patterns, leading to errors in real-world
applications. An MST approach can moderate these effects by learning a wider
support set as well as prevent the model from learning shortcut features.
Training on multiple datasets should prevent shortcut features and overfitting
of a specific population sample. To implement MST, we split each of the seven
datasets Table 1 for training and validation (90:10). The EYEPACS dataset was
divided into EYEPACS-train which was included for all experiments and EYEPACS-
test. At each fine-tuning stage, we used a joint validation dataset of all
source domains aside from the left-out target domain. To evaluate performance
on a target domain, we used a leave-one-domain-out method, i.e., six out of
seven datasets were used as source domains to train the model, while the left-
out domain was used as the target domain. Therefore, in Figures 3, 2, 4 and
Tables 2 and 3, the performance measures reported are reported for EYEPACS-
test, i.e., the test set of the source domain, while the other datasets were
considered as unseen target domains. This MST approach is called DRStageNet.

####  3.1.4 Benchmarks

We benchmarked DRStageNet against two SOTA models [15, 17]. Our first
benchmark consisted of a large ImageNet [36] pretrained EfficientNet2 [37]
(117 million parameters). We added a regression head and fine-tuned it using
the same protocol as DRStageNet-SS, using a larger initial learning rate of
$1e^{-5}$. This model is similar to the work of Vijayan et al. [15], with the
one difference that we used EfficientNet2 [37], which is the most contemporary
version of EfficientNet. The second benchmark consisted of the fine-tuned
version of the recently published pretrained RETFound [17] foundation model.
RETFound is a ViT-large that consists of 303 million parameters, that was
trained on a set of 1.6 million DFIs using the masked autoencoder [38] SSL
method. Previous work showed that the use of transfer learning from a
pretrained SSL model improved the performance of medical computer vision tasks
[39, 40], specifically DR diagnosis [41, 42, 43]. We used the published source
code222https://github.com/rmaphoh/RETFound_MAE. RETFound was fine-tuned using
the same protocol as DRStageNet-SS while the initial learning rate was set to
$1e^{-5}$ and the DFIs were resized to $224\times 224$ pixels, which is the
resolution input of RETFound.

Figure 2: DRStageNet confusion matrices. For each confusion matrix but
EYEPACS, the classification is reported for a model trained on all other
datasets. For EYEPACS-test, the EYEPACS-train set and all other datasets are
used to train the model.

###  3.2 Explainability

A modified attention rollout [44] approach was used for model explainability.
This method enables the use of the inherent high-resolution attention
mechanism of our model, which is a transformer-encoder-based model. Briefly,
the rollout [44] method describes how to compute the propagation of attention
from the input to the last block of self-attention. Although this method has
modest performance when used with images of a single large object, e.g., an
image of a dog [32], it tends to emphasize irrelevant tokens and is class-
agnostic, i.e., the same heatmaps are attributed to different classes in the
image. To address these issues, we followed the GradRollout [18] method and
weighted the attention of each layer by its gradient with respect to the
input. The GradRollout method was originally developed within the context of
the classification framework in which gradient weights are taken with respect
to the desired output class. In this research, we used a regression model such
that the gradients were used with respect to the single-scalar output of the
model. We hypothesized that, similarly to the classification setting, these
gradients will emphasize the attribution of relevant information at each layer
to a larger positive output.

Let $x\in\mathbb{R}^{n\times n}$ be an input image and
$f\left(x\right)\in\mathbb{R}$ be the output of the regression transformer
model. We define $s-1$ to be the number of input tokens, which are the
flattened patches across the image, $b$ to be the transformer self-attention
block counter and $h$ be the number of heads in each transformer block. Then
$A^{b}\in\mathbb{R}^{h\times\left(s\times s\right)}$ are the attention
matrices at each block and $df/dA^{b}=W^{b}\in\mathbb{R}^{h\times\left(s\times
s\right)}$ are the gradient weights that are multiplied element-wise with
$A^{b}$. Our algorithm is described in equation (1), where
$A^{b+1}_{\text{gradrollout}}$ is the output heatmap at step $b+1$ and
$g(\cdot)$ denotes first taking the maximum over the attention block heads,
then zeroing out 10% of the pixels of the lowest intensity. In addition,the
attention matrix is normalized at the end of each step. After the algorithm
reaches the last step of multiplication, $A^{b+1}_{\text{gradrollout}}$, it
extracts the attention weights associated with the global classification token
([CLS] token) both horizontally and vertically and averages them, since the
attention matrix is not necessarily symmetric (as a result of independent key
and query matrices). Finally, we reshape the weights to a squared matrix and
use a bilinear interpolation to resize them to the input dimension of
$518\times 518$ pixels.

$\displaystyle A_{\text{gradrollout}}^{b+1}$
$\displaystyle=\begin{cases}A^{b}&,b=0\\\ \frac{1}{2}\left(g\left(A^{b}\odot
W^{b}\right)+I\right)A_{\text{gradrollout}}^{b}&,b>0\end{cases}$ (1)

###  3.3 Performance measures

To assess the models’ performance we used multiclass accuracy (MC-ACC),
linearly weighted Cohen’s kappa (LW-Kappa), MSE and mean absolute error (MAE)
measures. We estimate the confidence interval by bootstrapping 1000 times 60%
of each test set, and the lower and upper bounds represent the quantiles of
0.25 and 0.75, respectively. Furthermore, we used the Mann-Whitney U
statistical test on the performance of DRStageNet and a benchmark. We also
report the area under the curve (AUC) metric for the binary rDR task. For that
purpose, the DRStageNet output was transformed into a binary output, with the
positive class defined as higher than stage one. The F1 score and the binary
accuracy for the rDR task are also reported.

###  3.4 Error analysis

We examined instances where DRStageNet misclassified DFIs with a discrepancy
of three or more units between the reference and predicted labels (Fig. 2).
Each DFI underwent an independent and blinded review by two retinal
specialists who were unaware of the dataset origin and the assigned reference
label. The specialists assessed whether the DFIs were ungradable, identified
the presence of one or more comorbidities, and assigned a DR grade according
to the ICDR scale. Disagreements between the annotations of the two
specialists were discussed, and a consensus was reached.

(a)

(b)

(c)

(d)

Figure 3: Models’ DR staging performance across open-source datasets. (a)
linearly weighted Cohen’s kappa (LW-Kappa) performance. (b) Multiclass
accuracy (MC-ACC) performance. (c) MSE performance. (d) MAE performance.

(a)

(b)

Figure 4: Models’ rDR performance across open-source datasets. (a) rDR AUC
performance. (b) rDR F1 performance. Figure 5: A comparison between three
explainability methods. The first column represents a set of four DFIs from
the DDR dataset. They were ordered by their respective ICDR reference labels.
Image (a) is labeled mild diabetic retinopathy (stage 1) while image (m) is
labeled proliferative diabetic retinopathy (stage 4). The ground-truth
microaneurysms, hemorrhages, soft exudates and hard exudates were colored in
green, blue, red, and cyan, respectively. The ground-truths in images (a) and
(e) have bounding boxes for visual support. The second column shows the
heatmap output of Grad-CAM [45] method for the last transformer block of
DRStageNet, the third column shows the heatmaps from the rollout [44] method
and the fourth column is our explainability heatmaps. In image (p), DRStageNet
attends to some neovascularization near the optic disc that is not present in
the ground-truth segmentation. Table 2: External performance comparison of rDR
AUC and MC-ACC between DRStageNet and other methods. (*) The reported results
are estimated from Fig. 2b in the paper [17].

| Eval. dataset
---
Method | rDR AUC | MC-ACC
APTOS | DRStageNet | 0.974 | 0.739
| DRGen [46] |  | 0.703
| RETFound [17] | 0.800* |
MESSIDOR2 | DRStageNet | 0.979 | 0.796
| Papadopoulos et al. [10] | 0.976 |
| DRGen [46] |  | 0.705
| RETFound [17] | 0.820* |
IDRiD | DRStageNet | 0.984 | 0.717
| RETFound [17] | 0.820* |

##  4 Results

###  4.1 Generalization performance

Figure 3 presents the measured LW-Kappa, MC-ACC, MSE and MAE for DRStageNet,
DRStageNet-SS and two contemporary benchmarks across the seven datasets used
in this research (Table 1). DRStageNet LW-Kappa on the local EYEPACS test set
was 0.747 vs. 0.718 for the EfficientNet2 SOTA benchmark. The generalization
performance of DRStageNet was significant (p¡0.001) and non-incremental over
the EfficientNet2 SOTA benchmark for five out of six of the external datasets.
The only exception was observed in the MSE metric on the IDRID dataset, where
DRStageNet and EfficientNet2 exhibited similar performance. The confusion
matrices for DRStageNet are displayed in Fig. 2. It can be observed that most
errors were one step off the diagonal, which is similar to human inter- and
intra-rater inconsistency [47].

In addition, Table 2 compares DRStageNet to other works reporting on the
performance on one or more of the external datasets included in our study.
DRStageNet demonstrated superior performance in terms of MC-ACC compared to
DRGen [46], which used a domain generalization approach to train their model.
Furthermore, on the MESSIDOR2 dataset, DRStageNet achieved an rDR AUC score of
0.979, similar to the performance of Papadopoulos et al. [10]. Additionally,
the reported rDR AUC performance of RETFound [17] was lower than DRStageNet’s
rDR AUC scores, e.g., 0.82 vs. 0.979 AUC on MESSIDOR2.

We report the rDR AUC performance of DRStageNet versus the other two
benchmarks (Fig. 4). DRStageNet, DRStageNet-SS and EfficientNet2 showed
comparable performance on all the datasets. RETFound exhibited a lower
performance than the other three methods, but it had a higher performance than
previously reported in the author’s original article [17], e.g., 0.968 versus
approximately 0.8 AUC on APTOS. A similar performance comparison of rDR F1,
MSE and MAE and detailed numerical results for DRStageNet is presented in Fig.
4, Fig. 3 and Table 3 respectively.

Table 3: DRStageNet performance summary on the datasets used in this research.
The reported kappa is linearly weighted and ACC is the binary rDR accuracy.

| Eval. dataset
---
AUC | ACC | F1 | MC-ACC |  | Kappa
---
MSE
EYEPACS | 0.955 | 0.940 | 0.831 | 0.842 | 0.747 | 0.246
DDR | 0.970 | 0.905 | 0.885 | 0.779 | 0.779 | 0.424
APTOS | 0.974 | 0.898 | 0.888 | 0.739 | 0.789 | 0.320
BRSET | 0.970 | 0.928 | 0.852 | 0.808 | 0.769 | 0.468
MESSIDOR2 | 0.979 | 0.942 | 0.886 | 0.796 | 0.777 | 0.218
DRTiD | 0.922 | 0.847 | 0.837 | 0.618 | 0.624 | 0.633
IDRiD | 0.984 | 0.955 | 0.964 | 0.717 | 0.772 | 0.397

###  4.2 Explainability

We validated the DRStageNet heatmaps by comparing them with the ground truth
DR lesion annotations of the DDR dataset (Fig. 5) of four DFIs that correspond
to the four stages of DR. The heatmaps correctly highlight the annotated
lesions in the four DFIs. In the DFI heatmap that corresponded to a
proliferative case (p), DRStageNet highlighted some neovascularization, i.e.,
formation of new blood vessels, near the optic disc, which were not annotated
in the ground truth. A DFI is classified as proliferative DR when there exists
neovascularization or vitreous/preretinal hemorrhage. In this case, DRStageNet
attended the neovascularization near the optic disc, which may explain its
correct classification. Furthermore, Fig. 5 presents a comparative analysis of
our explainability approach with Grad-CAM and Rollout methods. Qualitatively,
we appreciate on this set of examples that these techniques generate an
important number of false positives. This makes them less valuable in
providing an explainability support tool to a prospective clinical user.

###  4.3 Error analysis

A total of 247 DFIs were reviewed. Among these, a total of 106 (43%) DFIs had
at least one comorbidity, 24 (10%) were considered ungradable, and 147 (59%)
were mislabelled. Overall, a total of 210 (85%) DFIs were either ungradable,
had a comorbidity or were mislabelled.

##  5 Discussion

(a)

(b)

(c)

Figure 6: Some examples of the extreme errors of DRStageNet. (a) Mislabelling
example from EYEPACS dataset. Original ICDR target is 4 and the reviewed
target 0. DRStageNet predicted stage 0 correctly. (b) A DFI with a comorbidity
example from EYEPACS dataset, specifically of vascular occlusion, which might
resemble multiple hemorrhages and microaneurysms similar to DR stage 3 (severe
non-proliferative DR), as predicted by DRStageNet. (c) Ungradable DFI example
from DDR dataset. This DFI has lighting and focus problems.

The single-source DRStageNet-SS method outperformed the other two single-
source benchmarks on the source domain EYEPACS-test set, and exhibited equal
or superior generalization performance on all target domains except for IDRiD.
These results demonstrate the value of using a transformer architecture that
was pretrained using SSL on a large number of natural images to create a
representation that can be fine-tuned to a specific downstream classification
task. Our MST approach, DRStageNet, exhibited even better generalization
performance except for on the DRTiD dataset, on which all the methods showed
reduced performance. The results obtained for the MST DRStageNet algorithm
highlight the value of using MST in learning a more generalizable
representation for a given task, as it avoids overfitting a specific domain or
learning shortcut features. The confusion matrices for DRStageNet (Fig. 2)
show that the majority of errors are one step off the diagonal. This is
similar to human inter- and intra-rater inconsistencies [47]. For the
secondary task rDR, performance was superior but incremental compared to the
EfficientNet2 benchmark. Indeed, the binary task is simpler than ICDR staging
and does not require the global attention mechanism that distinguishes
transformer models such as the one used in DRStageNet.

The generalization performance of all algorithms was low for the DRTiD
dataset. In DRTiD, the ground truths are based on two-field DFIs, one being
macula-centered while the other is optic disc-centered. Therefore, unlike the
other datasets used, the ICDR label is provided at the patient level as
opposed to the image level. These differences in the results obtained with a
dataset annotated at the patient level versus the single DFI level suggest the
need for integration of multiple DFIs of the retina, at least one macula-
centered and one disc-centered.

Our error analysis revealed that the majority of misclassified DFIs with a gap
of three or more were incorrectly labeled (63%). This is a recognized issue in
open-source DFI datasets, which has prompted some researchers [6, 26] to re-
annotate datasets by engaging multiple DR experts to refine the reference
labels. Additionally, a considerable number of misclassified DFIs were
associated with at least one comorbidity (40%). In fact, certain comorbidities
can be confused with DR due similarities in their patterns, e.g., vascular
occlusion (see Fig. 6), or because they are relatively rare pathologies that
are not adequately represented in the training set, preventing the network
from learning a meaningful representation.

Our explainability heatmaps exhibited high association with ground-truth
lesion segmentation of the DDR dataset. However, there were still some
discrepancies. The heatmaps sometime highlighted regions such as the macula or
the optic disc that suggest that DRStageNet identifies some important features
in these regions. Some DR manifestations were, however, not highlighted, e.g.,
in Fig. 5 (i). It is possible that DRStageNet’s decision may be driven by a
subset of DR objects while not attending the others. The main goal of the
heatmaps is to help specialists identify DR patients, which DRStageNet helps
achieve. We believe that the combination of the ICDR regression outcome and
the attention heatmaps will help reduce clinician workload and the number of
false negatives cases of DR patients.

In this research, we introduced DRStageNet, a DL model for DR, which aims to
accurately stage DR and mitigate the challenge of generalizing performance to
target domains. For this purpose, DRStageNet uses a SSL-based pretrained ViT
model and implements a multi-source domain fine-tuning strategy. We
demonstrated the superiority of this method over two SOTA models.

## References

  * [1] Xiling Lin et al.  “Global, regional, and national burden and trend of diabetes in 195 countries and territories: an analysis from 1990 to 2025”  In _Scientific Reports_ 10.1, 2020, pp. 14790
  * [2] Hong Sun et al.  “IDF Diabetes Atlas: Global, regional and country-level diabetes prevalence estimates for 2021 and projections for 2045”  In _Diabetes research and clinical practice_ 183 Diabetes Res Clin Pract, 2022
  * [3] Joanne W.Y. Yau et al.  “Global Prevalence and Major Risk Factors of Diabetic Retinopathy”  In _Diabetes Care_ 35.3, 2012, pp. 556–564
  * [4] Michael Dubow et al.  “Classification of Human Retinal Microaneurysms Using Adaptive Optics Scanning Light Ophthalmoscope Fluorescein Angiography”  In _Investigative Opthalmology & Visual Science_ 55.3, 2014, pp. 1299
  * [5] Daniel Shu Wei Ting et al.  “Development and Validation of a Deep Learning System for Diabetic Retinopathy and Related Eye Diseases Using Retinal Images From Multiethnic Populations With Diabetes”  In _JAMA_ 318.22, 2017, pp. 2211
  * [6] Varun Gulshan et al.  “Development and Validation of a Deep Learning Algorithm for Detection of Diabetic Retinopathy in Retinal Fundus Photographs”  In _JAMA_ 316.22 American Medical Association, 2016, pp. 2402–2410
  * [7] Mike Voets, Kajsa Møllersen and Lars Ailo Bongo  “Replication study: Development and validation of deep learning algorithm for detection of diabetic retinopathy in retinal fundus photographs”  In _PLoS One_ 14.6 Public Library of Science, 2018, pp. e0217541
  * [8] Michael D. Abràmoff et al.  “Automated analysis of retinal images for detection of referable diabetic retinopathy”  In _JAMA ophthalmology_ 131.3 JAMA Ophthalmol, 2013, pp. 351–357
  * [9] Michael David Abràmoff et al.  “Improved Automated Detection of Diabetic Retinopathy on a Publicly Available Dataset Through Integration of Deep Learning”  In _Investigative Ophthalmology & Visual Science_ 57.13 The Association for Research in VisionOphthalmology, 2016, pp. 5200–5206
  * [10] Alexandros Papadopoulos, Fotis Topouzis and Anastasios Delopoulos  “An interpretable multiple-instance approach for the detection of referable diabetic retinopathy in fundus images”  In _Scientific Reports 2021 11:1_ 11.1 Nature Publishing Group, 2021, pp. 1–15
  * [11] Mohaimenul Azam Khan Raiaan et al.  “A Lightweight Robust Deep Learning Model Gained High Accuracy in Classifying a Wide Range of Diabetic Retinopathy Images”  In _IEEE Access_ 11 Institute of ElectricalElectronics Engineers Inc., 2023, pp. 42361–42388
  * [12] Richa Vij and Sakshi Arora  “A novel deep transfer learning based computerized diagnostic Systems for Multi-class imbalanced diabetic retinopathy severity classification”  In _Multimedia Tools and Applications_ 82.22 Springer, 2023, pp. 34847–34884
  * [13] Francisco J. Martinez-Murcia et al.  “Deep residual transfer learning for automatic diagnosis and grading of diabetic retinopathy”  In _Neurocomputing_ 452 Elsevier, 2021, pp. 424–434
  * [14] Mohamed Shaban et al.  “A convolutional neural network for the screening and staging of diabetic retinopathy”  In _PLoS ONE_ 15.6 PLOS, 2020, pp. e0233514
  * [15] Midhula Vijayan and Venkatakrishnan S  “A Regression-Based Approach to Diabetic Retinopathy Diagnosis Using Efficientnet”  In _Diagnostics 2023, Vol. 13, Page 774_ 13.4 Multidisciplinary Digital Publishing Institute, 2023, pp. 774
  * [16] C.P Wilkinson et al.  “Proposed international clinical diabetic retinopathy and diabetic macular edema disease severity scales”  In _Ophthalmology_ 110.9, 2003, pp. 1677–1682
  * [17] Yukun Zhou et al.  “A foundation model for generalizable disease detection from retinal images”  In _Nature 2023_ Nature Publishing Group, 2023, pp. 1–8
  * [18] Jacob Gildenblat  “Exploring Explainability for Vision Transformers”  URL: https://jacobgil.github.io/deeplearning/vision-transformer-explainability
  * [19] Maxime Oquab et al.  “DINOv2: Learning Robust Visual Features without Supervision”  In _arXiv: 2304.07193_ , 2023
  * [20] Luis F Nakayama et al.  “The Challenge of Diabetic Retinopathy Standardization in an Ophthalmological Dataset”  In _Journal of Diabetes Science and Technology_ 15.6, 2021, pp. 1410–1411
  * [21] Emma Dugas, Jared, Jorge and Will Cukierski  “Diabetic Retinopathy Detection”  In _Kaggle_ , 2015  URL: https://kaggle.com/competitions/diabetic-retinopathy-detection
  * [22] Jorge Cuadros and George Bresnick  “EyePACS: An Adaptable Telemedicine System for Diabetic Retinopathy Screening”  In _Journal of diabetes science and technology (Online)_ 3.3 Diabetes Technology Society, 2009, pp. 509
  * [23] Tao Li et al.  “Diagnostic assessment of deep learning algorithms for diabetic retinopathy screening”  In _Information Sciences_ 501 Elsevier, 2019, pp. 511–522
  * [24] Maggie Karthik and Sohier Dane  “APTOS 2019 Blindness Detection”  In _Kaggle_ , 2019  URL: https://kaggle.com/competitions/aptos2019-blindness-detection
  * [25] Luis Filipe Nakayama et al.  “A Brazilian Multilabel Ophthalmological Dataset (BRSET) (version 1.0.0)”  In _PhysioNet_ , 2023  DOI: https://doi.org/10.13026/xcxw-8198
  * [26] Jonathan Krause et al.  “Grader Variability and the Importance of Reference Standards for Evaluating Machine Learning Models for Diabetic Retinopathy”  In _Ophthalmology_ 125.8 Ophthalmology, 2018, pp. 1264–1272
  * [27] Prasanna Porwal et al.  “Indian Diabetic Retinopathy Image Dataset (IDRiD)”  In _IEEE Dataport_ , 2018  DOI: https://dx.doi.org/10.21227/H25W98
  * [28] Junlin Hou et al.  “Cross-Field Transformer for Diabetic Retinopathy Grading on Two-field Fundus Images”  In _2022 IEEE International Conference on Bioinformatics and Biomedicine (BIBM)_ IEEE Computer Society, 2022, pp. 985–990
  * [29] Ben Graham  “Kaggle Diabetic Retinopathy Detection Competition Report”, 2015  URL: https://kaggle-forum-message-attachments.storage.googleapis.com/88655/2795/competitionreport.pdf
  * [30] Yi Zhou et al.  “Collaborative learning of semi-supervised segmentation and classification for medical images”  In _Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition_ 2019-June IEEE Computer Society, 2019, pp. 2074–2083
  * [31] Xianglong Zeng, Haiquan Chen, Yuan Luo and Wenbin Ye  “Automated diabetic retinopathy detection based on binocular siamese-like convolutional neural network”  In _IEEE Access_ 7 Institute of ElectricalElectronics Engineers Inc., 2019, pp. 30744–30753
  * [32] A. Dosovitskiy et al.  “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale”  In _International Conference on Learning Representations_ , 2020
  * [33] Dan Hendrycks and Kevin Gimpel  “Gaussian Error Linear Units (GELUs)”  In _arXiv:1606.08415_ , 2016
  * [34] Diederik P. Kingma and Jimmy Lei Ba  “Adam: A Method for Stochastic Optimization”  In _3rd International Conference on Learning Representations, ICLR 2015 - Conference Track Proceedings_ International Conference on Learning Representations, ICLR, 2014
  * [35] Robert Geirhos et al.  “Shortcut learning in deep neural networks”  In _Nature Machine Intelligence 2020 2:11_ 2.11 Nature Publishing Group, 2020, pp. 665–673
  * [36] Jia Deng et al.  “ImageNet: A large-scale hierarchical image database”  In _2009 IEEE Conference on Computer Vision and Pattern Recognition_ IEEE, 2009, pp. 248–255
  * [37] Mingxing Tan and Quoc V. Le  “EfficientNetV2: Smaller Models and Faster Training”  In _Proceedings of the 38th International Conference on Machine Learning_ PMLR, 2021, pp. 10096–10106
  * [38] Kaiming He et al.  “Masked Autoencoders Are Scalable Vision Learners”  In _Proceedings of the IEEE Computer Society Conference on Computer Vision and Pattern Recognition_ 2022-June IEEE Computer Society, 2021, pp. 15979–15988
  * [39] Shekoofeh Azizi et al.  “Robust and data-efficient generalization of self-supervised machine learning for diagnostic imaging”  In _Nature Biomedical Engineering 2023 7:6_ 7.6 Nature Publishing Group, 2023, pp. 756–779
  * [40] Xiaomeng Li et al.  “Self-supervised Feature Learning via Exploiting Multi-modal Data for Retinal Disease Diagnosis”  In _IEEE Transactions on Medical Imaging_ 39.12 Institute of ElectricalElectronics Engineers Inc., 2020, pp. 4023–4033
  * [41] Tuan Truong, Sadegh Mohammadi and Matthias Lenga  “How Transferable Are Self-supervised Features in Medical Image Classification Tasks?”  In _Proceedings of Machine Learning Research_ 158 ML Research Press, 2021, pp. 54–74
  * [42] Olle G. Holmberg et al.  “Self-supervised retinal thickness prediction enables deep learning from unlabelled data to boost classification of diabetic retinopathy”  In _Nature Machine Intelligence 2020 2:11_ 2.11 Nature Publishing Group, 2020, pp. 719–726
  * [43] Philippe Burlina, William Paul, T..Alvin Liu and Neil M. Bressler  “Detecting Anomalies in Retinal Diseases Using Generative, Discriminative, and Self-supervised Deep Learning”  In _JAMA ophthalmology_ 140.2 JAMA Ophthalmol, 2022, pp. 185–189
  * [44] Samira Abnar and Willem Zuidema  “Quantifying Attention Flow in Transformers”  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ Stroudsburg, PA, USA: Association for Computational Linguistics, 2020, pp. 4190–4197
  * [45] Ramprasaath R. Selvaraju et al.  “Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization”  In _International Journal of Computer Vision_ 128.2 Springer, 2016, pp. 336–359
  * [46] Mohammad Atwany and Mohammad Yaqub  “DRGen: Domain Generalization in Diabetic Retinopathy Classification”  In _Medical Image Computing and Computer Assisted Intervention – MICCAI 2022_ 13432 LNCS Springer Nature Switzerland, 2022, pp. 635–644
  * [47] Mohammad Z. Atwany, Abdulwahab H. Sahyoun and Mohammad Yaqub  “Deep Learning Techniques for Diabetic Retinopathy Classification: A Survey”  In _IEEE Access_ 10, 2022, pp. 28642–28655