# Automatic Estimation of Ulcerative Colitis Severity from Endoscopy Videos
using Ordinal Multi-Instance Learning$\dagger$††thanks: $\dagger$Accepted for
publication in the Special Issue Journal of Computer Methods in Biomechanics
and Biomedical Engineering: Imaging & Visualization and presented at the
MICCAI 2021 AE-CAI $|$ CARE $|$ OR 2.0 Workshop.

Evan Schwab and Gabriela Oana Cula and Kristopher Standish and Stephen S. F.
Yip and Aleksandar Stojmirovic and Louis Ghanem* and Christel Chehoud* *Co-
Senior AuthorsCONTACT Evan Schwab. Email<EMAIL_ADDRESS>Janssen R&D
Data Science Analytics and Insights

###### Abstract

Ulcerative colitis (UC) is a chronic inflammatory bowel disease characterized
by relapsing inflammation of the large intestine. The severity of UC is often
represented by the Mayo Endoscopic Subscore (MES) which quantifies mucosal
disease activity from endoscopy videos. In clinical trials, an endoscopy video
is assigned an MES based upon the most severe disease activity observed in the
video. For this reason, severe inflammation spread throughout the colon will
receive the same MES as an otherwise healthy colon with severe inflammation
restricted to a small, localized segment. Therefore, the extent of disease
activity throughout the large intestine, and overall response to treatment,
may not be completely captured by the MES. In this work, we aim to
automatically estimate UC severity for each frame in an endoscopy video to
provide a higher resolution assessment of disease activity throughout the
colon. Because annotating severity at the frame-level is expensive, labor-
intensive, and highly subjective, we propose a novel weakly supervised,
ordinal classification method to estimate frame severity from video MES labels
alone. Using clinical trial data, we first achieved $0.92$ and $0.90$ AUC for
predicting mucosal healing and remission of UC, respectively. Then, for
severity estimation, we demonstrate that our models achieve substantial
Cohen’s Kappa agreement with ground truth MES labels, comparable to the inter-
rater agreement of expert clinicians. These findings indicate that our
framework could serve as a foundation for novel clinical endpoints, based on a
more localized scoring system, to better evaluate UC drug efficacy in clinical
trials.

###### keywords:

IBD Ulcerative Colitis endoscopy video severity estimation weak supervision
multi-instance learning ordinal classification deep learning convolutional
neural networks.

††articletype: JOURNAL ARTICLE PREPRINT

##  1 Introduction

Ulcerative colitis (UC) is a disabling and chronic inflammatory bowel disease
(IBD) characterized by relapsing inflammation and ulceration of the large
intestinal mucosa. Clinical trials in IBD use standardized scoring systems to
assess both clinical outcomes and changes in disease activity. One commonly
used disease severity score for UC is the total Mayo Score (Schroeder et al.
1987), which combines clinical disease features, physician global assessment,
and mucosal disease burden as determined by video endoscopy. An important
component of the total Mayo Score is the Mayo Endoscopic Subscore (MES) which
is used to define patient-level UC severity from endoscopic videos and is
graded on a discrete 0-3 scale: Normal (MES 0), Mild (MES 1), Moderate (MES
2), Severe (MES 3). See Fig. 1 for representative video frames.

The MES scoring system, as deployed in clinical trials and used in clinical
practice, has two notable limitations. First, gastroenterologists (GIs), often
serving as central readers in a clinical trial, are required to assign a
single video-level MES based upon the most severe disease activity observed in
the video. Therefore, as illustrated in Fig. 2, if a patient with severe and
widespread UC (MES 3) has marked improvement after treatment, except for a
highly localized severe disease activity, that patient will still receive an
MES 3. A consequence of this ambiguity is the relative insensitivity of the
MES to quantify meaningful, local therapeutic effects of drug treatments which
could negatively impact drug development and medical decision making in the
clinic.

The second notable limitation is the relatively high subjectivity of scoring
UC severity by expert GIs. A Fleiss Kappa ($\kappa$) metric is commonly
utilized to assesses inter-rater agreement between multiple raters with the
following scale: $0.01–0.20$ (slight), $0.21–0.40$ (fair), $0.41–0.60$
(moderate), $0.61–0.80$ (substantial), and $0.81–1.00$ (nearly perfect). As
reported in Daperno et al. (2014), the inter-rater agreement of video-level
MES between 14 expert raters was moderate with a Fleiss $\kappa\\!=\\!0.53$
($95\%$ confidence interval (CI) 0.47–0.56). Likewise, Principi et al. (2020)
report moderate Fleiss $\kappa\\!=\\!0.53$ ($95\%$ CI 0.39–0.66) with 13
expert raters. A Fleiss $\kappa\\!=\\!0.60$ ($95\%$ CI 0.51-0.69) was also
reported between 4 expert raters for the clinical trial data (Sands et al.
2019) that served as our ground truth labels for training and validation.

These two limitations, a lack of granularity and moderate inter-rater
agreement, motivate the development of an automatic disease activity scoring
system for UC which uses frame-level severity scores to provide a more precise
measurement of disease distribution and severity throughout the bowel. Such an
approach would allow for finer assessments of meaningful therapeutic effects
in IBD clinical trials.

(a) Normal (MES 0)

(b) Mild (MES 1)

(c) Moderate (MES 2)

(d) Severe (MES 3)

Figure 1: Example endoscopy video frames and corresponding Mayo Endoscopic
Score (MES) of UC severity. Figure 2:  Human central readers in UC clinical
trials assign a video-level MES based on the most severe inflammation observed
in an entire endoscopy video. Therefore, as shown in this illustration, a
patient will be scored an MES 3 before and after treatment, without
sensitivity to wide-spread improvements (i.e., red (MES 3) reduces to green
(MES 0), yellow (MES 1) or orange (MES 2)). This example demonstrates how the
current clinical scoring system fails to capture important treatment effects
from endoscopy videos and the need to measure severity at a higher resolution.

Prior work have examined the automatic estimation of frame-level IBD severity
scores from endoscopy videos. In particular, Stidham et al. (2019); Ozawa et
al. (2019); Yao et al. (2020) solve supervised learning problems that require
frame-level labels from trained experts. Employing a set of trained experts to
perform frame-level severity labeling, however, is a very time-consuming and
expensive task, requiring large numbers of labeled frames suitable for
training deep neural networks. In Sec. 4, we report that four highly
experienced GI clinicians required 3-4 hours per person to label only 200
frames of our dataset.

Therefore, to overcome the difficulties of frame-level labeling, we utilize
weakly supervised deep learning to learn frame-level estimations trained only
with video-level labels. In particular, we combine the weakly supervised
algorithm multi-instance learning (MIL), with the task of ordinal
classification to arrive at a novel ordinal MIL framework for frame-level
severity estimation. To the best of our knowledge, this approach has not been
previously developed. Most closely related to our work is Gutierrez Becker et
al. (2021), which utilizes a similar form of weak supervision but does not
attempt ordinal classification. Additionally, Mohammed et al. (2020) solves an
alternative multi-class disease detection problem on annotated video segments.

This paper is organized as follows: In Sec. 2.1 we introduce the standard MIL
framework and build upon this to derive our novel ordinal MIL in Sec 2.2.
Then, in Sec. 3 we describe our dataset, pre-processing, training, and
experimental setup. We present our results on clinical trial data in Sec. 4.

##  2 Methods

In this section, we combine MIL with ordinal classification to estimate
ordinal categories of UC severity at the frame-level using only video-level
labels. In Sec. 2.1 we introduce the standard MIL algorithm, traditionally
developed for binary classification, which we apply to solve two clinically
relevant binary classification tasks. We then extend MIL to the case of
ordinal classification for UC severity estimation in Sec. 2.2 and compare two
established ordinal classification frameworks: 1) an aggregation of an
ensemble of ranked binary classifiers and 2) regression with discrete ordinal
labels.

###  2.1 Multi-Instance Learning Applied to Endoscopy Videos

Multi-instance learning (MIL) is a framework for solving weakly labeled
problems in which labeled data can be viewed as a collection of smaller
unlabeled parts (instances) that we are interested in estimating. An example
medical application is the localization of disease pathology in medical
images. While ground truth disease labels (eg. disease vs. no disease) may be
relatively easy to provide at the image-level, annotating the specific
location of disease pathology in an image has proven to be a more time
consuming and difficult undertaking. MIL has been applied by dividing a
labeled image into smaller unlabeled image patches and estimating the presence
or absence of disease in those patches using weak supervision.

MIL has been traditionally formulated for binary classification, exploiting
the assumption that for images with positive disease labels, at least one
image patch will be positive, which explains the positive label of the image.
Conversely, all image patches from a negative image must be negative. This
framework has shown effective performance across image modalities like digital
histopathology (Campanella et al. 2019) and chest X-ray (Schwab et al. 2020).
Multi-class formulations have also been explored (Remedios et al. 2020; Lu et
al. 2021).

In our setting, we apply the same weakly supervised paradigm by substituting
videos for images and video frames for image patches. Our first application
can be posed as two clinically relevant binary classification tasks: 1) MES 0
vs. (1,2,3) to detect the presence of UC, and 2) MES (0,1) vs. (2,3) to detect
mucosal healing which is associated with improved long-term patient outcomes
(Turner et al. 2020).

Using a neural network, MIL first estimates a probabilistic score for each
frame, and then compares an aggregate of those scores with the video label in
the loss function. The most typical aggregation is the maximum operator. It
effectively selects the frame with the maximum score as representative of the
entire video. The max aggregation is well suited for our problem because it
mimics the established clinical framework of reporting the maximum MES
observed in the video. This framework also allows us to model frames as
independent samples, utilizing a classical MIL setting without more advanced
temporal modeling.

More specifically, let $V$ be an endoscopic video with $F$ image frames
$\\{f_{i}\\}_{i=1}^{F}$ and binary label $y=\\{0,1\\}$. Because we have image
data, we use a convolutional neural network (CNN), $\phi$, to estimate
continuous disease probability scores for each frame,
$p_{i}=\phi(f_{i})\in[0,1]$. Then, to obtain the video-level probability,
$p_{v}$, we select the frame with the maximum score as the representative for
the video, i.e. $f_{v}$ with $v=\operatorname*{argmax}_{i}p_{i}$, and present
this to the binary cross-entropy loss function during training. Applying a
threshold $t\in[0,1]$ to $p_{v}$ produces the video class prediction
$\hat{y}\in\\{0,1\\}$ and at the same time, class predictions
$\hat{y}_{i}\in\\{0,1\\}$ for each frame in the video. Fig. 3 provides an
overview of MIL for binary classification. In the next section we extend the
standard binary MIL formulation to the case of ordinal classification.

Figure 3: Binary MIL Framework Overview. Endoscopy videos are split into
frames and fed into a CNN to generate frame-level probability scores. The
frame with the maximum score is selected as the representative for the video
and a binary threshold is used to produce a class for the video and each
frame.

###  2.2 Ordinal MIL for Frame-Level Severity Estimation

Ordinal classification is appropriate for ordered data labels without an
apparent numerical distance between them (i.e., Normal $<$ Mild $<$ Moderate
$<$ Severe)111For simplicity, and by clinical convention, we choose to
represent these labels with numerical values (i.e.,
$0\\!<\\!1\\!<\\!2\\!<\\!3$), but the numerical distance between labels need
not be equal.. Unlike multi-class problems, where misclassifications are
weighted equally, misclassifying Severe as Normal should be penalized more
than misclassifying Severe as Moderate. Next, we derive the novel combination
of MIL with two established ordinal classification frameworks.

Figure 4: Overview of proposed ordinal MIL framework using an ensemble of
ranked binary classifiers (A): $\phi_{>0}$: 0 vs $(1,2,3)$, $\phi_{>1}$:
$(0,1)$ vs $(2,3)$, $\phi_{>2}$: $(0,1,2)$ vs $3$. The output of (A) is three
probability scores per frame $p_{>0}$, $p_{>1}$, $p_{>2}\in[0,1]$. We
experiment with three different aggregation methods, (B), (C), (D), of the
frame probability scores. In (B), $p_{>0},p_{>1}$ and $p_{>2}$ are converted
into multi-class probabilities $p_{0}$, $p_{1}$, $p_{2}$, $p_{3}\in[0,1]$,
each corresponding to the respective MES class. The maximum probability is
selected as the ordinal class for each frame. In (C), three distinct binary
thresholds are applied to $p_{>0}$, $p_{>1}$, and $p_{>3}$ resulting in three
binary classes per frame. These are summed to generate ordinal classes per
frame. In (D), $p_{>0}$, $p_{>1}$, and $p_{>2}$ are first summed to create a
continuous score in $[0,3]$ per frame. A set of ordinal thresholds are then
applied to generate ordinal classes per frame. In each method, the frame with
the maximum class is selected as the video class.

####  2.2.1 Ensemble of Ranked Binary Classifiers

One approach to ordinal classification is to restructure a problem with $M$
ordinal labels as an ensemble of $M\\!-\\!1$ ranked binary classification
tasks and aggregate each output to estimate ordinal classes. Within a
traditional fully supervised framework, Frank and Hall (2001) convert the
$M\\!-\\!1$ binary probabilities to $M$ multi-class probabilities and take the
maximum as the ordinal class. Alternatively, Chen et al. (2017) add the
$M\\!-\\!1$ binary class predictions in $\\{0,1\\}$ to produce an ordinal
class in $\\{0,M\\!-\\!1\\}$.

In our application, with four ordinal classes, we must train and aggregate
three binary classifiers: $\phi_{>0}$: 0 vs. $(1,2,3)$, $\phi_{>1}$: $(0,1)$
vs. $(2,3)$, $\phi_{>2}$: $(0,1,2)$ vs. $3$, where $\phi_{>m}$ is a model to
decide whether the target class is greater than value $m$. This approach is
appealing because MIL is already well-defined for binary classifications and
$\phi_{>0}$ and $\phi_{>1}$ are already relevant tasks introduced in Sec. 2.1.
Thus, it only remains to repeat the MIL method for $\phi_{>2}$. In the MIL
setting, however, the output of each binary classifier results from the frame
with the maximum score. Since each classifier in the ensemble is applied to
the data independently, the frame with the maximum score may be different for
each classifier. We therefore, must first aggregate frame-level binary outputs
to produce frame-level ordinal scores and then take the frame with the maximum
ordinal score as the representative for the video.

Let $p^{i}_{>0}=\phi_{>0}(f_{i})$, $p^{i}_{>1}=\phi_{>1}(f_{i})$, and
$p^{i}_{>2}=\phi_{>2}(f_{i})$ be the ranked binary probability scores in
$[0,1]$ for each frame and let $\hat{y}_{>0}^{i}$, $\hat{y}_{>1}^{i}$, and
$\hat{y}_{>2}^{i}$ be binary class predictions $\\{0,1\\}$ for each frame. We
experiment with three ways to aggregate these probabilities to estimate
discrete ordinal classes, $\hat{z}^{i}\in\\{0,1,2,3\\}$, for each frame:

  1. 1.

(Convert) As in Frank and Hall (2001), convert the 3 binary probabilities into
ordinal class probabilities, $p^{i}_{0}=1-p^{i}_{>0}$,
$p^{i}_{1}=p^{i}_{>0}-p^{i}_{>1}$, $p^{i}_{2}=p^{i}_{>1}-p^{i}_{>2}$,
$p^{i}_{3}=p^{i}_{>2}$, and take the max multi-class probability as the
ordinal class per frame, $\hat{z}^{i}=\max_{0\leq j\leq 3}p^{i}_{j}$.

  2. 2.

(Threshold) As in Chen et al. (2017), first threshold each binary probability
score $p^{i}_{>0},p^{i}_{>1},p^{i}_{>2}$ with three independent thresholds to
produce binary classes predictions
$\hat{y}_{>0}^{i},\hat{y}_{>1}^{i},\hat{y}_{>2}^{i}$ per frame. Then sum these
class predictions to estimate the ordinal class per frame
$\hat{z}^{i}=\hat{y}_{>0}^{i}+\hat{y}_{>1}^{i}+\hat{y}_{2}^{i}$.

  3. 3.

(Sum) First sum the three binary probabilities together to produce a
continuous severity score $q^{i}=p^{i}_{>0}+p^{i}_{>1}+p^{i}_{>2}\in[0,3]$.
Then apply three ordinal thresholds
$0\\!<\\!t_{0}\\!<\\!t_{1}\\!<\\!t_{2}\\!<\\!3$ to bin the continuous scores
into ordinal classes $\hat{z}^{i}$ per frame.

Then, to mimic the clinical scoring system, we take the max ordinal class over
all frames as the video-level MES severity class estimation,
$\hat{z}=\max_{1\leq i\leq N}\hat{z}^{i}$. While the first two methods
(Convert and Threshold) produce discrete MES severity classes per frame, the
third method (Sum) also produces a continuous severity score in [0,3] per
frame which may be useful for generating a more precise severity scoring
system. Fig. 4 illustrates the proposed ordinal MIL using an ensemble of
ranked binary classifiers and the three different aggregation methods.

####  2.2.2 Regression with Discrete Labels

Another common technique for ordinal classification is to apply regression
while maintaining discrete class labels. Here the only difference between
classification and regression using CNNs is the final output layer. Instead of
a sigmoid activation for binary classification, the final layer is linear with
one neuron. Let $\phi_{r}$ be this regression CNN. Therefore, instead of a
binary probability score $p_{i}\in[0,1]$ as above, this network produces a
continuous severity score $s_{i}\in\mathbb{R}$ per frame such that
$s_{i}\\!=\\!\phi_{r}(f_{i})$. A clipping function is used to restrict the
range of values to [0,3].

To combine ordinal regression with MIL, we simply take the frame with the
maximum severity score, $f_{v}$, as the severity score as the representative
for the video, i.e. $s_{v}\\!=\\!\phi_{r}(f_{v})$ with
$v\\!=\\!\operatorname*{argmax}_{i}s_{i}$. We then apply three ordinal
thresholds $0\\!<\\!t_{0}\\!<\\!t_{1}\\!<\\!t_{2}\\!<\\!3$ to bin the
continuous scores into ordinal classes $\hat{z}^{i}\in\\{0,1,2,3\\}$ per
frame. Fig. 5 gives an overview of the ordinal MIL regression method. In the
next section we compare each of the proposed ordinal MIL methods on clinical
trial endoscopy data and validate them on a sample of labeled video frames.

Figure 5: Ordinal MIL Regression framework overview. Endoscopy videos are
split into frames and fed into a CNN-based regression model, which generates
continuous frame level severity scores in [0,3]. The frame with the maximum
score is selected as the representative score for the video. A set of three
ordinal thresholds are chosen to bin the continuous scores into discrete MES
$\\{0,1,2,3\\}$.

##  3 Experiments

###  3.1 Data and Training

The data we use in this work originates from the UNIFI clinical trial of
Ustekinumab (STELARA[R]) in UC (Sands et al. 2019). Our dataset consists of
1881 endoscopic videos from 726 subjects with 1 to 4 videos per subject
lasting anywhere from 35 seconds to 47 minutes and a mean duration of 10
minutes. We extracted frame images from our video files at a frame rate of 1
frame per second (as in Yao et al. (2020)), resulting in a total of 1,129,188
frames with an average of 600 frames per video.

An inclusion criterion for trial enrollment is an initial screening MES of 2
or 3. During follow-up visits, many subjects had severity reduction to MES 0
or 1. But at the same time, a subset of subjects did not complete the trial.
Therefore, the distribution of classes is imbalanced towards the severe
classes. Our video class distribution is 167 (MES 0), 220 (MES 1), 492 (MES
2), and 1,002 (MES 3). This imbalance is more pronounced when data samples are
relabeled for the set of ranked binary classifiers, i.e., for $\phi_{>0}$: 167
(0), 1,714 (1), $\phi_{>1}$: 387 (0), 1,494 (1), $\phi_{>2}$: 879 (0), 1,002
(1).

To address the class imbalance that skews towards the positive class, we
propose a method that exploits the MIL assumption that all frames in a class 0
video must be class 0. Then, instead of selecting the maximum frame that
represents a class 0 video, we can select the top $K$ frames (Campanella et
al. 2019). This is akin to increasing the sample size $K$ times for the 0
class. While it may be tempting to choose a very large $K$, or all frames from
each class 0 video, this may lead to a reverse class imbalance. We
experimented with $K\\!=\\!1,5,10,20,40,100,$ and $200$ for class 0. For class
1 videos, we leave $K\\!=\\!1$ since, by construction of MIL, there is no
guarantee that more than one frame will exhibit the positive class. (If a
class 0 video has less than $K$ frames, we select all the frames in that
video.) Only the maximum frame is selected at test time.

After experimentation with various CNN architectures, ResNet34 was chosen as
the best performing. The frame image dimensions were resized from
$640\\!\times\\!510\\!\times\\!3$ to $224\\!\times\\!224\\!\times\\!3$ to fit
the input to the CNN. After exhaustive parameter testing, we chose an Adam
optimizer with learning rate=1e-5, weight decay=0.01, and binary cross entropy
loss for 100 epochs for the binary classification tasks. For the ordinal MIL
regression we experimented with multiple loss functions including Mean Square
Error, Mean Absolute Error, Smooth L1, and Log Cosh and found that Mean
Absolute Error provided the best results. We employed online image
augmentation of horizontal and vertical flips, rotations, and hue changes. We
ran 5-fold cross validation with splits grouped at the subject-level with
proportional label distributions.

(a) Red wall

(b) Bleeding

(c) Forceps

(d) Rinse Water

Figure 6: Example endoscopic imaging artifacts. Frames with (a) red wall and
(b) bleeding were removed from the dataset due to the similarity with severe
inflammation that led to false positive predictions. Frames with (c) forceps
and (d) rinse water were not removed since they did not contribute to severe
disease scores.

###  3.2 Quality Control

Due to the transparent nature of MIL, we ensured a level of quality control by
visually inspecting the top scoring frames of true and false positive
predictions for consistent image features that may have misled the algorithm.
We found that frames that were entirely red (“red wall”) (Fig. 6a), where the
camera is pressed against the colon, were consistently the highest scoring
instead of severe UC. We thus sought to remove these misclassified frames from
our dataset to improve performance.

Taking a supervised approach, we manually labeled a set of 4,200 frames from
several videos as being a red wall or not. We used a ResNet pre-trained on
ImageNet to extract a feature vector of length 512 from each frame and trained
a support vector machine (SVM), achieving $90\%$ accuracy on a validation set
of frames from videos without overlapping of subjects. We applied the
classifier to the entire dataset and removed 273,262 red wall frames. After
successfully removing these red wall frames, we then discovered that images
with bleeding (Fig. 6b) became the highest scoring frames for both true and
false positives. We repeated this approach to remove images of bleeding by
manually labeling 4,604 images and training a new SVM. With $92\%$ accuracy on
a validation set, we removed 222,906 more frames.

Following the removal of these two types of image artifacts, performance
improved allowing severe images to score highly and drive the optimization. We
found it was not necessary to remove frames with objects such as forceps222As
per clinical trial protocol, biopsies were taken with forceps at a predefined
location in the left colon for all patients and therefore did not correlate to
frames with severe disease. (Fig. 6c) or artifacts like rinse water (Fig. 6d)
because these artifacts did not contribute to severe disease scores. The two
rounds of quality control led to a total removal of 496,168 frames, a $32.6\%$
average decrease in the number of frames per video.

Figure 7: Continuous frame-level MES severity estimations for two endoscopy
videos using the ordinal MIL ensemble method 3 (Sum). Each sample frame
corresponds to the square colored markers on the plot, and the maximum frame,
indicated by a triangle marker, gives the score for the video. Each dotted
line is an optimal MES threshold found via grid search over 5-fold cross
validation.

Figure 8: Comparison of inter-rater agreement measured by quadratic weighted
Cohen’s $\kappa$ between the MES estimations of each of the proposed methods
and MES labels. Red circles (with $95\%$ confidence intervals (CI) over 5-fold
cross validation) indicate the agreement between model estimations and video-
level labels. Each method achieves substantial agreement with ground truth
labels. Blue triangles indicate the agreement between model estimations and
frame-level consensus labels for 200 frames within a validation split. Purple
triangle is the average inter-rater agreement between each human rater and the
frame-level consensus labels with $95\%$ CI. The confusion matrices show MES
estimations vs. ground truth for Ordinal MIL Regression for video and frame-
level labels. The video-level confusion matrix includes estimations over all
5-folds.

##  4 Results

###  4.1 Video-Level Analysis

Our first results are for the three ranked binary classification tasks. With
$K=\\{40,200,100\\}$ we achieved AUCs of $0.90\pm 0.03$, $0.92\pm 0.02$, and
$0.88\pm 0.02$ for $\phi_{>0}$, $\phi_{>1}$, and $\phi_{>2}$, respectively,
across 5-fold cross validation. While $K$ parameters were tuned to achieve
these AUCs for each classifier independently, we found that parameters may
also be tuned with respect to their combination for ordinal classification.
The combination of $K\\!=\\!\\{40,100,40\\}$ gave us the best average ordinal
classification performance over the 5-fold cross validation for each
aggregation method in Sec. 2.2.1.

For the ordinal ensemble model 2 (Threshold), binary thresholds are needed for
each of the three binary classifiers. Using a grid search over 5-fold cross
validation, we selected the optimal combination of binary thresholds
$\\{t_{1},t_{2},t_{3}\\}\\!=\\!\\{0.20,0.19,0.22\\}$. For the ordinal ensemble
model 3 (Sum) which produces continuous severity scores ranging from [0,3],
ordinal thresholds are needed to bin these scores into discrete MES. Using a
grid search over 5-fold cross validation, we selected ordinal thresholds
$\\{t_{1},t_{2},t_{3}\\}\\!=\\!\\{0.1,1.04,2.02\\}$ to maximize performance.
Fig. 7 illustrates our frame-level severity estimation for two example videos
from a validation set using the ensemble summation method 3 to produce
continuous severity scores. The frames with the maximum score correctly
estimate the ground truth MES. For the ordinal MIL regression method, we found
$K\\!=\\!10$ produced the best performance. Over 5-fold cross validation we
found optimal ordinal thresholds
$\\{t_{1},t_{2},t_{3}\\}\\!=\\!\\{1.03,1.97,2.79\\}$.

In evaluating the performances of our methods, we calculated a quadratic
weighted Cohen’s $\kappa$ value between the output of our models and the
ground truth labels. Because the ground truth have been shown to be subjective
and variable as indicated by a moderate Fleiss333The (unweighted) Fleiss
$\kappa$ reported by studies including our own data, are not necessarily
comparable to a weighted Cohen’s $\kappa$, so we do not make claims of
comparison to human raters at the video-level. $\kappa$, measuring accuracy
may be a misleading quantity. Instead, we aim to demonstrate that we achieve
an acceptable level of agreement with the ground truth. Our three ordinal MIL
ensemble methods achieved substantial agreement with $\kappa\\!=\\!0.64$
($95\%$ CI 0.56-0.71), $\kappa\\!=\\!0.67$ ($95\%$ CI 0.66-0.69), and
$\kappa\\!=\\!0.65$ ($95\%$ CI 0.62-0.69), respectively. The ordinal MIL
regression model also achieved a substantial agreement with
$\kappa\\!=\\!0.68$ ($95\%$ CI 0.64-0.72). Fig. 8 shows a plot of the $\kappa$
comparisons and confusion matrices for the ordinal MIL regression model. These
results show a very similar performance between all methods at the video-
level. In the next section, we evaluate performance at the frame-level.

###  4.2 Frame-Level Analysis

While the above analyses show substantial $\kappa$ values at the video level,
it is also important that we also validate our methods at the frame-level. To
this end, we collected frame-level MES labels from 4 experienced, non-central
reader GIs to compute inter-rater agreement at the frame-level. We randomly
selected 50 frames from videos of each of the 4 MES classes within a single
validation split of our dataset, resulting in 200 total frames. This task was
time consuming, requiring 3-4 hours per physician to label all 200 frames. To
compare the inter-rater agreement of the 4 raters, we calculated a Fleiss
$\kappa\\!=\\!0.42$, which is on the low end of moderate and significantly
below the reported video-level Fleiss $\kappa\\!=\\!0.60$ ($95\%$ CI
0.51-0.69) (Sands et al. 2019).

In addition, because we know the ground truth video labels from which each of
the 200 frames were extracted, we analyzed how well the frame-level labels
coincided with the video-level labels using the fact that the frame MES label
should not exceed the video MES label. In fact, $27.5\%$ of the annotated
frame labels were greater than the video scores. We adjusted the frame labels
to be consistent with the video labels by taking the minimum of the two and
then took a majority vote consensus to produce our frame labels. In addition
to a Fleiss $\kappa$, we calculated how well each rater agreed with the
adjusted consensus frame labels by calculating the quadratic weighted Cohen’s
$\kappa\\!=\\!0.80,0.78,0.77,0.62$ with an average $\kappa\\!=\\!0.74$ (CI
0.61-0.87).

In comparison, for each of our models we achieved quadratic weighted Cohen’s
$\kappa\\!=\\!0.44$ (Convert) $\kappa\\!=\\!0.63$ (Threshold),
$\kappa\\!=\\!0.63$ (Sum), and $\kappa\\!=\\!0.66$ (Ordinal MIL Regression),
respectively. We found that the Threshold, Sum and the ordinal MIL regression
models achieve a substantial $\kappa$ value that falls within the $95\%$ CI of
the human raters at the frame-level. See Fig. 8 for this $\kappa$ comparison
and accompanying confusion matrices for the ordinal MIL regression. These
results indicate that the proposed ordinal MIL methods, can not only automate
the established MES scoring system at the video level with substantial
agreement to the ground truth, but can also provide frame-level MES
estimations comparable to human raters.

##  5 Conclusion

In this work, we proposed a set of novel weakly supervised ordinal
classification methods to automatically estimate the severity of UC from
endoscopy videos. Our methods have the ability to estimate frame-level
severity on a continuous scale by training only on discrete video-level
severity labels. This is beneficial for expanding models quickly to other
disease types and clinical indications without expensive, time-consuming and
subjective labeling for each new target of interest. Our proposed ordinal MIL
framework has the potential to effectively automate the established video-
level MES scoring system with a substantial level of score agreement.
Furthermore, we validate that the resulting frame-level severity estimations
are comparable to that of human raters. With continuous severity scores at the
frame level, our method will foster the development of more localized and
granular endoscopic endpoints that may be more sensitive to meaningful
therapeutic effects and better predict patient outcomes in IBD.

## Acknowledgement

We give special thanks to Dr. Louis Ghanem, Dr. Kathleen Lomax, Dr. Laurie
Conklin, and Dr. Susana Gonzalez for providing their GI expertise in scoring a
set of endoscopy video frames for use in our method validation and analysis.

## References

  * Campanella et al. (2019) Campanella G, Hanna MG, Geneslaw L, Miraflor A, Silva VWK, Busam KJ, Brogi E, Reuter VE, Klimstra DS, Fuchs TJ. 2019. Clinical-grade computational pathology using weakly supervised deep learning on whole slide images. Nature Medicine. 25(8):1301–1309.
  * Chen et al. (2017) Chen S, Zhang C, Dong M, Le J, Rao M. 2017. Using ranking-cnn for age estimation. In: IEEE Conference on Computer Vision and Pattern Recognition. p. 5183–5192.
  * Daperno et al. (2014) Daperno M, Comberlato M, Bossa F, Biancone L, Bonanomi AG, Cassinotti A, Cosintino R, Lombardi G, Mangiarotti R, Papa A, et al. 2014. Inter-observer agreement in endoscopic scoring systems: preliminary report of an ongoing study from the Italian Group for Inflammatory Bowel Disease (IG-IBD). Digestive and Liver Disease. 46(11):969–973.
  * Frank and Hall (2001) Frank E, Hall M. 2001. A simple approach to ordinal classification. In: European Conference on Machine Learning. Springer. p. 145–156.
  * Gutierrez Becker et al. (2021) Gutierrez Becker B, Arcadu F, Thalhammer A, Gamez Serna C, Feehan O, Drawnel F, Oh YS, Prunotto M. 2021. Training and deploying a deep learning model for endoscopic severity grading in ulcerative colitis using multicenter clinical trial data. Therapeutic advances in gastrointestinal endoscopy. 14:2631774521990623.
  * Lu et al. (2021) Lu MY, Williamson DF, Chen TY, Chen RJ, Barbieri M, Mahmood F. 2021. Data-efficient and weakly supervised computational pathology on whole-slide images. Nature Biomedical Engineering.
  * Mohammed et al. (2020) Mohammed A, Farup I, Pedersen M, Yildirim S, Hovde Ø. 2020. PS-DeVCEM: Pathology-sensitive deep learning model for video capsule endoscopy based on weakly labeled data. Computer Vision and Image Understanding. 201:103062.
  * Ozawa et al. (2019) Ozawa T, Ishihara S, Fujishiro M, Saito H, Kumagai Y, Shichijo S, Aoyama K, Tada T. 2019. Novel computer-assisted diagnosis system for endoscopic disease activity in patients with ulcerative colitis. Gastrointestinal endoscopy. 89(2):416–421.
  * Principi et al. (2020) Principi M, Contaldo A, Bianchi FP, Losurdo G, Iannone A, Ierardi E, Tafuri S, Di Leo A, Work IET. 2020. Inter-observer agreement of a new endoscopic score for ulcerative colitis activity: Preliminary experience. Diagnostics. 10(4):213.
  * Remedios et al. (2020) Remedios S, Wu Z, Bermudez C, Kerley CI, Roy S, Patel MB, Butman JA, Landman BA, Pham DL. 2020. Extracting 2d weak labels from volume labels using multiple instance learning in CT hemorrhage detection. In: Medical Imaging 2020: Image Processing; vol. 11313. International Society for Optics and Photonics. p. 113130F.
  * Sands et al. (2019) Sands BE, Sandborn WJ, Panaccione R, O’Brien CD, Zhang H, Johanns J, Adedokun OJ, Li K, Peyrin-Biroulet L, Van Assche G, et al. 2019. Ustekinumab as induction and maintenance therapy for ulcerative colitis. New England Journal of Medicine. 381(13):1201–1214.
  * Schroeder et al. (1987) Schroeder KW, Tremaine WJ, Ilstrup DM. 1987. Coated oral 5-aminosalicylic acid therapy for mildly to moderately active ulcerative colitis. New England Journal of Medicine. 317(26):1625–1629.
  * Schwab et al. (2020) Schwab E, Gooßen A, Deshpande H, Saalbach A. 2020. Localization of critical findings in chest X-ray without local annotations using multi-instance learning. In: IEEE International Symposium on Biomedical Imaging (ISBI). p. 1879–1882.
  * Stidham et al. (2019) Stidham RW, Liu W, Bishu S, Rice MD, Higgins PD, Zhu J, Nallamothu BK, Waljee AK. 2019. Performance of a deep learning model vs human reviewers in grading endoscopic disease severity of patients with ulcerative colitis. JAMA network open. 2(5):e193963–e193963.
  * Turner et al. (2020) Turner D, et al. 2020. STRIDE-II: An Update on the Selecting Therapeutic Targets in Inflammatory Bowel Disease (STRIDE) Initiative of the International Organization for the Study of IBD (IOIBD): Determining Therapeutic Goals for Treat-to-Target strategies in IBD. Gastroenterology.
  * Yao et al. (2020) Yao H, Najarian K, Gryak J, Bishu S, Rice MD, Waljee AK, Wilkins HJ, Stidham RW. 2020. Fully automated endoscopic disease activity assessment in ulcerative colitis. Gastrointestinal Endoscopy.


#  : A Neural Framework for Attributed Graph Clustering via Modularity
Maximization

Aritra Bhowmick1, Mert Kosan2, Zexi Huang2, Ambuj Singh2, Sourav Medya3 Work
done prior to joining Visa Research.

###### Abstract

Graph clustering is a fundamental and challenging task in the field of graph
mining where the objective is to group the nodes into clusters taking into
consideration the topology of the graph. It has several applications in
diverse domains spanning social network analysis, recommender systems,
computer vision, and bioinformatics. In this work, we propose a novel method,
, which primarily optimizes the modularity objective using graph neural
networks and scales linearly with the graph size. Our method does not require
the number of clusters to be specified as a part of the input and can also
leverage the availability of auxiliary node level information. We extensively
test on several real-world datasets of varying sizes, across multiple popular
cluster quality metrics. Our approach consistently outperforms the state-of-
the-art methods, demonstrating significant performance gains in almost all
settings.

## Introduction and Related Work

Graph clustering is a fundamental problem in network analysis and plays an
important role in uncovering structures and relationships between the nodes or
entities in a graph. It has numerous applications in several domains such as
community detection in social networks (Newman 2006a), identifying functional
modules in biological systems (Wang et al. 2010), image segmentation in
computer vision (Felzenszwalb and Huttenlocher 2004), and recommender systems
(Moradi, Ahmadian, and Akhlaghian 2015). The primary goal of graph clustering
is to group nodes with similar characteristics or functions while maintaining
a clear distinction between different clusters.

##### Node attributes.

While traditional graph clustering methods primarily rely on graph topology
such as modularity maximization (Newman 2006b; Newman and Girvan 2003), recent
research (Wang et al. 2017) has recognized the importance of incorporating
node attributes in the clustering process, offering a more comprehensive
approach for grouping nodes. Node attributes provide additional information
associated with each node and provide contextual insights that can improve the
accuracy in the clustering process.

##### GNN-based clustering.

Recently, there have been several attempts that use the power of deep learning
in the form of Graph neural networks (GNNs) (Kipf and Welling 2016; Hamilton,
Ying, and Leskovec 2017; Velickovic et al. 2017) for graph clustering. GNNs
provide a powerful tool in graph-based machine learning that is successful in
many diverse prediction tasks (Zhang and Chen 2018; Ying et al. 2018b, a), by
incorporating the graph topology node features or attributes. For GNN-based
graph clustering, MGAE (Wang et al. 2017) marginalizes the corrupted node
features to learn representations via a graph encoder and applies spectral
clustering. Another graph autoencoder based approach has been proposed in
(Park et al. 2019). To improve the efficiency of clustering, contrastive
learning methods have been used recently as well (Liu et al. 2023;
Kulatilleke, Portmann, and Chandra 2022; Xia et al. 2021). For more neural
methods on deep graph clustering, we refer the readers to this recent survey
(Yue et al. 2022).

##### Neural modularity maximization.

One of the initial methods to optimize modularity via deep learning for graph
clustering is proposed by Yang et al. (2016). They design a nonlinear
reconstruction method based on graph autoencoders, which also incorporate
constraints among node pairs. Wu et al. (2020) propose a method that obtains a
spatial proximity matrix by using the adjacency matrix and the opinion leaders
in the social network. The spatial eigenvectors of the proximity matrix are
applied subsequently to optimize modularity. Mandaglio, Amelio, and Tagarelli
(2018) is another study that incorporates the modularity metric for community
detection and graph clustering. Choong, Liu, and Murata (2018); Bhatia and
Rani (2018) try to find communities without predefined community structure.
Choong, Liu, and Murata (2018) propose a generative model for community
detection using a variational autoencoder. Bhatia and Rani (2018) firstly
analyze the possible number of communities in the graph, then fine-tune it
using modularity. Later on, Sun et al. (2021) use a graph neural network that
optimizes modularity and attributes similarity objectives. Another related
work in this domain is the method DMoN by Tsitsulin et al. (2023). This method
designs an architecture to encode cluster assignments and then formulate a
modularity-based objective function for optimizing these assignments.

##### Our contributions.

A vast majority of these methods have the limitation that they require the
number of clusters to be given as an input or they do not take full advantage
of the associated node attributes along with the additional node level
information like partial availability of labels or samples of known pairwise
memberships. To overcome these challenges, we propose a framework named Deep
Graph Cluster () that eliminates the need for a predefined number of clusters
and harnesses graph representation learning methods that can leverage node
attributes along with other available auxiliary information. Our major
contributions are as follows.

  * •

: We develop a novel framework that uses pairwise (soft) memberships between
nodes to solve the graph clustering problem via modularity maximization. The
complexity of our framework scales linearly with the size of the graph.

  * •

Handling Unknown Number of Clusters and Auxiliary Information: Our proposed
methodology can generalize well to cases when the number of clusters is not
known beforehand. Our designed loss function is also flexible towards
accommodating the additional local or node-level information. These are the
major strengths of our approach.

  * •

Extensive Empirical Evaluation: We conduct extensive experiments on seven
real-world datasets of different sizes on four different objectives that
quantify the quality of clusters. Our method shows significant performance
gain against state-of-the-art methods in most of the settings.

## Problem Definition

We consider an undirected and unweighted graph $G=(V,E)$, where
$V=\\{v_{1},v_{2},v_{3},...,v_{n}\\}$ is the set of $n$ vertices/nodes, and
$E=\\{e_{ij}=(v_{i},v_{j})\\}$ is the set of $m$ edges. The adjacency matrix
of $G$ can be represented as a non-negative symmetric matrix $A=[A_{ij}]\in
R_{+}^{n\times n}$ where $A_{ij}=1$ if there is an edge between vertices $i$
and $j$, and $A_{ij}=0$. The degree of vertex $i$ is defined as
$d_{i}=\sum_{j}a_{ij}$. In addition, we have features (or attributes)
associated with each node in the graph, $X^{0}\in\mathbb{R}^{n\times r}$ where
$r$ is the size of the feature vector on nodes.

##### Graph clustering.

Our objective is to have a disjoint clustering of the nodes in the graph. More
specifically, the problem of graph clustering is to partition the set of nodes
into $k$ clusters or groups $\\{V_{i}\\}^{k}_{i=1}$ ($V_{i}\cap
V_{j}=\emptyset$ for $i\neq j$), such that the nodes within a cluster are more
densely connected than nodes belonging to different clusters. Furthermore, in
this work, we aim to incorporate node attributes in addition to the graph
topology for graph clustering.

While there exist several metrics to measure the quality of clustering such as
conductance (Yang and Leskovec 2012) and normalized cut-ratio (Shi and Malik
2000), modularity remains the most popular and widely used metric for graph
clustering in the literature (Fortunato and Hric 2016).

##### Modularity.

The approach of graph clustering based on maximizing the modularity of the
graph has been introduced by Newman (Newman 2006b). As a graph topology-based
measure, modularity (Newman 2006b) quantifies the difference between the
fraction of the edges that fall within the clusters and the expected fraction
assuming the edges have been distributed randomly. Formally, modularity ($Q$)
is defined as follows:

$Q=\frac{1}{2m}\sum_{ij}(A_{ij}-\frac{d_{i}d_{j}}{2m})\delta(c_{i},c_{j})$ (1)

where $\delta(c_{i},c_{j})$ is the Kronecker delta, i.e.,
$\delta(c_{i},c_{j})=1$ if $c_{i}=c_{j}$ and $0$ otherwise, and $c_{i}$ is the
community to which node $i$ is assigned.

The value of modularity for unweighted and undirected graphs lies in the range
$[-1/2,1]$. The $Q$ value close to $0$ implies that the fraction of edges
inside communities is no better than a random distribution, and higher values
usually correspond to a stronger cluster structure. The modularity $Q$ can
also be expressed in the matrix form as follows:

$Q=\frac{1}{2m}\sum_{ij}B\odot M=\frac{1}{2m}Tr(BM^{T})=\frac{1}{2m}Tr(BM)$
(2)

where $M$ is a $n\times n$ symmetric matrix with $M_{ij}=\delta(c_{i},c_{j})$
and $B_{ij}=(A_{ij}-\frac{d_{i}d_{j}}{2m})$ is called the modularity matrix.

Modularity maximization.  Since a larger $Q$ implies a prominent cluster
structure, optimizing the modularity is a popular way of finding good
clusters. While modularity optimization is known to be NP-Hard (Brandes et al.
2006), there exist techniques such as spectral relaxation and greedy
algorithms, which permit efficient solutions (Newman 2006a; Blondel et al.
2008).

Our goal. We achieve graph clustering via modularity maximization. The
definition of modularity brings the idea of computing pairwise memberships
allowing a natural interpretation without knowing the number of clusters. We
aim to take advantage of that and the power of graph representation learning
techniques that can exploit both structural and non-structural information
from graphs. In the experiments, we show the efficacy of our method on four
different objectives that quantify the quality of the clusters: modularity
(Newman 2006b), conductance (Yang and Leskovec 2012), Normalized mutual
information (NMI), and F1 score.

## Method:

We present , a fully differentiable method which performs deep graph
clustering based on the graph structure and node attributes, without the need
to pre-define the number of communities. The key rationale of our method is to
parameterize a relevant clustering objective (e.g., Modularity) with
similarity between nodes computed based on graph neural network (GNN)-based
embeddings. Our method consists of four steps:

  * •

Node embeddings. As the first step, obtains the node embeddings as the output
of the GNN followed by some transformations that helps to perform efficient
clustering.

  * •

Modularity via similarity. We evaluate the similarity between all node pairs
from the embeddings and treat them as soft community pairwise memberships.

  * •

Objectives. Subsequently, it builds the cluster detection objectives based on
the soft memberships and train the GNN parameters in a differentiable manner.

  * •

Final clustering.  Finally, it computes the community memberships by
clustering the GNN node embeddings.

We introduce each step of our method in the following subsections.

### Node embeddings using GNN

We begin with a brief introduction of graph neural networks (GNNs). GNNs are
powerful graph machine learning paradigms that combine the graph structure and
node attribute information into node embeddings for different downstream
tasks. A key design element of GNNs is message passing where the nodes
iteratively update their representations (embeddings) by aggregating
information from their neighbors. In the literature, several different GNN
architectures have been proposed (Scarselli et al. 2008; Kipf and Welling
2016; Hamilton, Ying, and Leskovec 2017; Velickovic et al. 2017) which
implement different schemes of message passing. A comprehensive discussion on
the methods and applications of GNNs are described here (Zhou et al. 2020).

In this paper, we leverage the widely used Graph Convolutional Network (GCN)
(Kipf and Welling 2016) to produce node embeddings, noting that our model can
be equipped with other GNNs. With the initial node features as $X^{(0)}$, the
layer-wise message passing rule for layer $l$ ($l=0,\cdots,L-1$) is as
follows:

$X^{(l+1)}=\sigma(\tilde{A}X^{(l)}W^{(l)})$ (3)

where $\tilde{A}=D^{-\frac{1}{2}}AD^{\frac{1}{2}}$ is the normalized adjacency
matrix, $D$ is the diagonal node degree matrix, $X^{l}$ is the embedding
output of the $l$-th layer, $W^{l}$ is the learnable weight matrix of the
$l$th layer, and $\sigma$ is the activation function which introduces the non-
linearity in the feature aggregation scheme. We do not use any self loop
creation in the adjacency matrix and we choose the SELU (Klambauer et al.
2017) as the activation function for better convergence. The SELU activation
is given as:

$SELU(x)=\begin{cases}\beta x,&\text{if }x\geq 0\\\
\beta\alpha(e^{x}-1),&\text{otherwise}\end{cases}$ (4)

where $\beta\approx 1.05$ and $\alpha\approx 1.67$.

Transformation of the embeddings. Let $X=X^{L}=[X_{1},\cdots,X_{n}]^{T}$ be
the output embeddings of the last layer readout from the GNN. We introduce the
following problem-specific transformations on the embeddings (where $Z^{\circ
2}$ denotes the element-wise square operation):

$\displaystyle X_{i}$ $\displaystyle\leftarrow\frac{X_{i}}{\sum_{j}X_{ij}},$
(5) $\displaystyle X_{i}$ $\displaystyle\leftarrow\tanh(X_{i}),$
$\displaystyle X_{i}$ $\displaystyle\leftarrow X_{i}^{\circ 2},$
$\displaystyle X_{i}$ $\displaystyle\leftarrow\frac{X_{i}}{\|X_{i}\|_{2}}$

Specifically, the first normalization is used so as to prevent vanishing
gradients because of the $\tanh$ activation function for large values. Next,
after the activation, doing the element-wise square ensures the output is
constrained within the positive coordinate space. The final $L_{2}$
normalization reduces the cosine similarity computation of node pairs
(introduced in next subsection) to corresponding dot products. Thus, the final
embeddings lie on the surface of the unit sphere constrained in the positive
space. More detailed intuitions behind these will be explained in details in
the following sections.

### Modularity via embedding similarity

After obtaining the transformed GNN embeddings $X$, we demonstrate how to
compute the modularity $Q$ based on them via pairwise node similarities.

To achieve this, we focus on $M$ which is defined in Eq. 2 as a binary matrix
that encodes the pairwise memberships (via $\delta$) of the nodes in a
cluster. This pairwise relationship is transitive, i.e.,
$\delta(c_{u},c_{v})=1$ and $\delta(c_{v},c_{w})=1$ implies
$\delta(c_{u},c_{w})=1$. However, as stated before, the problem of finding the
optimal $M$ which maximizes $Q$ is NP-Hard. The main idea is to replace $M$
with a soft pairwise community membership matrix. We choose to replace $M$
with a similarity matrix which is defined based on node embeddings, where the
similarity ($f_{sim}(X_{u},X_{v})\in[0,1]$) can be viewed as soft membership
between the nodes $\\{u,v\\}$. Higher values of $f_{sim}(X_{u},X_{v})$
corresponds to higher similarity or stronger relationship between the nodes.
Although there can be many choices for the similarity function, we choose
$f_{sim}$ as the cosine similarity, $f_{sim}(X_{u},X_{v})=\cos(X_{u},X_{v})$.

Here, we also emphasize our rationale for the transformations of the
embeddings earlier. Specifically, the original $M$ only takes binary values
(i.e., 0 or 1). Our embedding transformation allows
$\cos(X_{u},X_{v})\in[0,1]$ by limiting them in the positive coordinate space,
and enables its computation via dot products
$\cos(X_{u},X_{v})=X_{u}^{T}X_{v}$, which in turn leads to an efficient
computation as discussed later.

### Objective function

We introduce a novel joint objective function that performs clustering based
on both the community structure quality measure (i.e., modularity) and local
auxiliary information.

#### Modularity optimization.

We first define our primary objective function, i.e., modularity optimization
by approximating the pairwise community membership computed with embedding
similarity:

$L_{1}=-\tilde{Q}=-\frac{1}{2m}Tr(BXX^{T})$ (6)

where $B$ is the modularity matrix, $m$ is the number of edges, and $XX^{T}$
is the similarity matrix obtained from the transformed node embeddings.

#### Auxiliary information loss.

We now discuss how our algorithm can be made more flexible by accounting for
additional local information.

Specifically, let $S\subseteq V$ be the subset of nodes whose local
information is available, and $H\in\mathbb{R}^{|S|\times|S|}$ be the pairwise
information matrix. We consider different types of additional information. In
the semi-supervised setting, partial node labels (e.g. cluster labels, class
labels) are available, and we can construct $H$ based on the pairwise
membership:

$H_{ij}=\begin{cases}1,&\text{if }c_{i}=c_{j}\\\
0,&\text{otherwise}\end{cases}$ (7)

where $c_{i}$ is the label of the node $i$. Alternatively, by collecting all
available labels in one-hot matrix form $C\in\mathbb{R}^{|S|\times p}$, we can
rewrite $H$ as:

$H=CC^{T}$ (8)

When those ground-truth information is not available, we can also leverage
traditional structure-based graph partitioning heuristics, such as Louvain
(Blondel et al. 2008), and treat their generated clusters as the node labels.
In general, any pairwise node information similarity which can be approximated
as $\langle C_{i},C_{j}\rangle\in[0,1]$, can be effectively used.

With the local information matrix $H$, the secondary objective function which
minimizes the difference between $H$ and the embedding similarity matrix is
given as :

$L_{2}=\frac{1}{|S|^{2}}\|H-X_{S}X_{S}^{T}\|_{F}^{2}$ (9)

where $X_{S}$ is the submatrix of node embeddings with only nodes in $S$.

The final objective function is a weighted combination of the two objectives:

$L=L_{1}+\lambda L_{2}$ (10)

where $\lambda$ is a hyperparameter. The GNN parameters are trained in an end-
to-end based on the total loss $L$ with the stochastic gradient descent
algorithm.

When instead of individual node labels, samples of pairwise node memberships
are available, $L_{2}$ can be written as follows:

$L_{2}=\frac{1}{|S|}\sum_{p_{ij}\in S}(1-\langle X_{i},X_{j}\rangle)^{2}$ (11)

where $p_{ij}$ are the node pairs which belong to the same community and $S$
is a set of such pairs.

### Clustering node embeddings

In this section, we illustrate how to obtain the hard cluster partitions based
on the soft pairwise memberships obtained in the previous section.

One way is to directly apply clustering algorithms based on the pairwise node
similarity matrix, such as affinity propagation (Frey and Dueck 2007).
However, computing the full similarity matrix from the embeddings is
computationally prohibitive. Instead, we take advantage of the following
observation. Since our embeddings are $L_{2}$ normalized (i.e.,
$\|X_{u}\|_{2}=1$), the cosine similarity is directly related to the Euclidean
distance in the embedding space:

$\displaystyle\|X_{u}-X_{v}\|_{2}^{2}$
$\displaystyle=\|X_{u}\|_{2}^{2}+\|X_{v}\|_{2}^{2}-2\|X_{u}\|_{2}\|X_{v}\|_{2}\cos(X_{u},X_{v})$
(12) $\displaystyle=2(1-\cos(X_{u},X_{v}))$

This allows us to apply clustering algorithms based on the euclidean distance
in the embedding space without computing the full pairwise similarity matrix.

Specifically, we apply the Balanced Iterative Reducing and Clustering using
Hierarchies (BIRCH) (Zhang, Ramakrishnan, and Livny 1996). BIRCH is a
scalable, memory-efficient, online clustering algorithm that can cluster large
datasets by first generating a small and compact summary of the dataset that
retains as much information as possible. Unlike other popular choices such as
k-means (Hartigan and Wong 1979), BIRCH does not require the number of
clusters beforehand.

### Complexity analysis

In this section, we analyze the complexity of our proposed model. The forward
pass requires us to compute two objective functions. Specifically, modularity
optimization loss $L_{1}$ can be evaluated with

$\displaystyle L_{1}$
$\displaystyle=-\frac{1}{2m}Tr(BXX^{T})=-\frac{1}{2m}Tr(X^{T}BX)$ (13)
$\displaystyle=-\frac{1}{2m}(Tr(X^{T}AX)-\frac{1}{2m}Tr(X^{T}dd^{T}X))$

We can see that $L_{1}$ can be computed with sparse matrix multiplications
between $X$ and $A$ and matrix vector multiplications between $X$ and $d$.
These multiplications lead to an overall computation cost of $O(k^{2}n)$,
where $k$ is the dimension of the embeddings $X$.

For the auxiliary information loss, we have

$\displaystyle\|H-X_{S}X_{S}^{T}\|_{F}^{2}=\sum_{ij}(H_{ij}-(XX^{T})_{ij})^{2}$
(14)
$\displaystyle=\sum_{ij}H_{ij}^{2}+\sum_{ij}(XX^{T})_{ij}^{2}-2\sum_{ij}H_{ij}(XX^{T})_{ij}$
$\displaystyle=\sum_{ij}(CC^{T})_{ij}^{2}+\sum_{ij}(XX^{T})_{ij}^{2}-2\sum_{ij}(CC^{T})_{ij}(XX^{T})_{ij}$
$\displaystyle=Tr(C^{T}CC^{T}C)-Tr(X_{S}^{T}X_{S}X_{S}^{T}X_{S})-2Tr(X_{S}^{T}CC^{T}X_{S})$

Computing $C^{T}C$, $X_{S}^{T}X_{S}$, and $X_{S}^{T}C$ requires $O(p^{2}n)$,
$O(k^{2}n)$, and $O(knp)$ respectively via matrix multiplications, assuming
$|S|=n$, since $|S|$ can be atmost $n$. Here, $p$ is the dimension of the
auxiliary information (note, $C\in\mathbb{R}^{n\times p}$ ). Thus, the overall
complexity of our model is $O(k^{2}n+p^{2}n)$, where $k,p\ll n$. This shows
that our model scales linearly with the size of the graph.

## Experimental Results

In this section, we present the performance comparison of with other
clustering and neural community detection methods on $7$ well-known real-world
datasets. We evaluate the algorithms using various metrics, including graph
conductance, modularity, NMI (Normalized Mutual Information), and F1 score
where higher scores are desired for the last three. Furthermore, we
demonstrate the robustness of our method against different GNN architectures
and assess the impact of adding auxiliary information. We also introduce an
additional regularization objective and show its effect on our method. Lastly,
we illustrate the number of communities our method identifies for each
dataset. Our code and implementation of is available at
https://github.com/pyrobits/DGCluster

#### Datasets.

We use seven publicly available real-world datasets in our experiments. Table
1 presents a summary of the dataset statistics. Our experiments include
datasets from well-known citation networks, namely Cora, Citeseer, and PubMed
(Sen et al. 2008). In these networks, nodes correspond to individual papers,
edges represent citations between papers, and features are extracted using a
bag-of-words approach applied to paper abstracts. The topic of each paper is
captured through node labels. Besides citation networks, we use two co-
purchase networks, Amazon PC and Amazon Photo (Shchur et al. 2018). These
networks include products as nodes, with edges denoting co-purchase
relationships. Features are extracted from product reviews, while node labels
show product categories. Our last two datasets, Coauthor CS and Coauthor PHY
(Shchur et al. 2018; Shchur and Günnemann 2019), are co-authorship networks
for computer science and physics, respectively. Within these networks, nodes
correspond to authors, and edges indicate co-authorship between them. Node
features are keywords extracted from authors’ publications and node labels are
their fields of study.

| $|V|$ | $|E|$ | $|X|$ | $|Y|$
---|---|---|---|---
Cora | 2708 | 5278 | 1433 | 7
Citeseer | 3327 | 4552 | 3703 | 6
Pubmed | 19717 | 44324 | 500 | 3
Amazon PC | 13752 | 245861 | 767 | 10
Amazon Photo | 7650 | 119081 | 745 | 8
Coauthor CS | 18333 | 81894 | 6805 | 15
Coauthor PHY | 34493 | 247962 | 8415 | 5
Table 1: Statistics of the datasets. $|V|$, $|E|$, $|X|$, and $|Y|$ denote the
number of nodes, edges, features, and node labels.

#### Baselines.

Our baseline methods consist of a range of classical clustering algorithms,
such as k-means, as well as state-of-the-art graph community detection
algorithms. For consistent comparison, we adopted the same baseline setting as
in the recent DMoN (Tsitsulin et al. 2023): k-means based on features, SBM by
Peixoto (2014), k-means based on DeepWalk by Perozzi, Al-Rfou, and Skiena
(2014), k-means(DGI) by Veličković et al. (2018), AGC by Zhang et al. (2019),
DAEGC by Wang et al. (2019), SDCN by Bo et al. (2020), NOCD by Shchur and
Günnemann (2019), DiffPool by Ying et al. (2018a), MinCutPool by Bianchi,
Grattarola, and Alippi (2020), and Ortho by Bianchi, Grattarola, and Alippi
(2020).

#### Performance Measures.

We use our primary objective, modularity, as our main evaluation metric.
Additionally, we evaluate the performance using other important metrics,
namely NMI, conductance, and F1 score. In cases where a dataset lacks ground
truth cluster labels, we assign node labels as their respective cluster
labels. We multiply the metrics by 100 for better readability.

  * •

Modularity ($Q$) (Newman 2006b): This serves as the primary objective, aiming
to quantify the strength of the community. It achieves this by contrasting the
proportion of edges within the community with a corresponding fraction
generated from random connections between nodes. Higher values mean better
community partitions.

  * •

Conductance ($C$) (Yang and Leskovec 2012): The graph’s conductance measures
the portion of total edges that goes outside the community. Lower values
indicate better community partitions.

  * •

Normalized mutual information (NMI): We use the NMI as a label-based metric
between the cluster assignments and true labels of the nodes. Higher values
mean better community partitions.

  * •

F1 score: We calculate the pairwise F1 score based on found pairwise node
memberships and their corresponding clusters. Since it is $O(N^{2})$, we
sample 1000 nodes to calculate this. Higher values indicate better community
partitions.

##### Other Settings.

In our experiments, we employ a GCN having two hidden layers of 256 and 128
nodes and an output dimension of 64, across all datasets. We showcase outcomes
for various $\lambda$ values. This enables a comparative analysis among the
graph structure-based metrics and graph attribute-based metrics. We use Adam
optimizer with a learning rate set to 0.001, and we set the number of epochs
to 300. Our tables and figures report average scores over 10 different runs,
using different seeds, for our method. For baselines, we report the results
from Tsitsulin et al. (2023).

| Cora | Citeseer | Pubmed | Amazon PC | Amazon Photo | Coauthor CS | Coauthor PHY
---|---|---|---|---|---|---|---
method | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$ | $C\downarrow$ | $Q\uparrow$
k-m(feat) | $61.7$ | $19.8$ | $60.5$ | $30.3$ | $55.8$ | $33.4$ | $84.5$ | $5.4$ | $79.6$ | $10.5$ | $49.1$ | $23.1$ | $57.0$ | $19.4$
SBM | $15.4$ | $77.3$ | $14.2$ | $78.1$ | $39.0$ | $53.5$ | $31.0$ | $\underline{60.8}$ | $18.6$ | 72.2 | 20.3 | $\underline{72.7}$ | $25.9$ | $\underline{66.9}$
k-m(DW) | $62.1$ | $30.7$ | $68.1$ | $24.3$ | 16.6 | 75.3 | $67.6$ | $11.8$ | $60.6$ | $22.9$ | $33.1$ | $59.4$ | $44.7$ | $47.0$
AGC | $48.9$ | $43.2$ | $41.9$ | $50.0$ | $44.9$ | $46.8$ | $43.2$ | $42.8$ | $35.3$ | $33.8$ | $41.5$ | $40.1$ | N/A | N/A
SDCN | $37.5$ | $50.8$ | $20.0$ | $62.3$ | $22.4$ | $50.3$ | $25.1$ | $45.6$ | $19.7$ | $53.3$ | $33.0$ | $55.7$ | $32.1$ | $52.8$
DAEGC | $56.8$ | $33.5$ | $47.6$ | $36.4$ | $53.6$ | $37.5$ | $39.0$ | $43.3$ | $19.3$ | $58.0$ | $39.4$ | $49.1$ | N/A | N/A
k-m(DGI) | $28.0$ | $64.0$ | $17.5$ | $73.7$ | $82.9$ | $9.6$ | $61.9$ | $22.8$ | $51.5$ | $35.1$ | $35.1$ | $57.8$ | $38.6$ | $51.2$
NOCD | $14.7$ | $\underline{78.3}$ | $6.8$ | $\underline{84.4}$ | $21.7$ | $69.6$ | $26.4$ | $59.0$ | $13.7$ | $70.1$ | $20.9$ | $72.2$ | $25.7$ | $65.5$
DiffPool | $26.1$ | $66.3$ | $26.0$ | $63.4$ | $32.9$ | $56.8$ | $35.6$ | $30.4$ | $26.5$ | $46.8$ | $33.6$ | $59.3$ | N/A | N/A
MinCutPool | $23.3$ | $70.3$ | $14.1$ | $78.9$ | $29.6$ | $63.1$ | N/C | N/C | N/C | N/C | $22.7$ | $70.5$ | $27.8$ | $64.3$
Ortho | $28.0$ | $65.6$ | $18.4$ | $74.5$ | $57.8$ | $32.9$ | N/C | N/C | N/C | N/C | $27.8$ | $65.7$ | $33.0$ | $59.5$
DMoN | $12.2$ | $76.5$ | $5.1$ | $79.3$ | $17.7$ | $65.4$ | $18.0$ | $59.0$ | $12.7$ | $70.1$ | $17.5$ | $72.4$ | 18.8 | $65.8$
($\lambda=0.0$) | 8.4 | 80.7 | $5.5$ | $87.2$ | $20.6$ | $72.5$ | $17.4$ | $61.3$ | $7.5$ | $70.9$ | 14.6 | 74.3 | $\underline{20.6}$ | 67.5
($\lambda=0.2$) | $9.7$ | $80.8$ | 4.1 | 87.4 | $20.4$ | $\underline{72.8}$ | 17.7 | 61.5 | 8.6 | $\underline{71.6}$ | $15.3$ | $74.2$ | $22.3$ | $67.3$
($\lambda=0.8$) | $14.5$ | $78.6$ | $6.5$ | $86.3$ | $24.6$ | $71.2$ | $27.3$ | $60.3$ | $12.4$ | $71.6$ | $18.1$ | $73.3$ | $21.3$ | $66.0$

Table 2: Performance across datasets evaluated using graph conductance $C$ and
graph modularity $Q$, with three $\lambda$ settings (0, 0.2, 0.8) from our
method. The best and second-best methods are highlighted in bold and
underlined for each dataset-metric pair. We fixed the optimal $\lambda$ value
per dataset in our method (i.e., 0.0 or 0.2) during comparison, and our method
demonstrates the best or comparable performance to the baselines. Unavailable
results and non-convergence are labeled as N/A and N/C.

| Cora | Citeseer | Pubmed | Amazon PC | Amazon Photo | Coauthor CS | Coauthor PHY
---|---|---|---|---|---|---|---
method | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$ | NMI $\uparrow$ | F1 $\uparrow$
k-m(feat) | 18.5 | 27.0 | 24.5 | 29.2 | 19.4 | 24.4 | 21.1 | 19.2 | 28.8 | 19.5 | 35.7 | 39.4 | 30.6 | 42.9
SBM | 36.2 | 30.2 | 15.3 | 19.1 | 16.4 | 16.7 | 48.4 | 34.6 | 59.3 | 47.4 | 58.0 | 47.7 | 45.4 | 30.4
k-m(DW) | 24.3 | 24.8 | 27.6 | 24.8 | 22.9 | 17.2 | 38.2 | 22.7 | 49.4 | 33.8 | 72.7 | 61.2 | 43.5 | 24.3
AGC | 34.1 | 28.9 | 25.5 | 27.5 | 18.2 | 18.4 | 51.3 | 35.3 | 59.0 | 44.2 | 43.3 | 31.9 | N/A | N/A
SDCN | 27.9 | 29.9 | 31.4 | 41.9 | 19.5 | 29.9 | 24.9 | 45.2 | 41.7 | 45.1 | 59.3 | 54.7 | 50.4 | 39.9
DAEGC | 8.3 | 13.6 | 4.3 | 18.0 | 4.4 | 11.6 | 42.5 | 37.3 | 47.6 | 45.0 | 36.3 | 32.4 | N/A | N/A
k-m(DGI) | 52.7 | 40.1 | 40.4 | 39.4 | 22.0 | 26.4 | 22.6 | 15.0 | 33.4 | 23.6 | 64.6 | 51.9 | 51.0 | 30.6
NOCD | 46.3 | 36.7 | 20.0 | 24.1 | 25.5 | 20.8 | 44.8 | 37.8 | 62.3 | 60.2 | 70.5 | 56.4 | 51.9 | 28.7
DiffPool | 32.9 | 34.4 | 20.0 | 23.5 | 20.2 | 26.3 | 22.1 | 38.3 | 35.9 | 41.8 | 41.6 | 34.4 | N/A | N/A
MinCut | 35.8 | 25.0 | 25.9 | 20.1 | 25.4 | 15.8 | N/C | N/C | N/C | N/C | 64.6 | 47.8 | 48.3 | 24.9
Ortho | 38.4 | 26.6 | 26.1 | 20.5 | 20.3 | 13.9 | N/C | N/C | N/C | N/C | 64.6 | 46.1 | 44.7 | 23.7
DMoN | 48.8 | 48.8 | 33.7 | 43.2 | 29.8 | 33.9 | 49.3 | 45.4 | 63.3 | 61.0 | 69.1 | 59.8 | 56.7 | 42.4
($\lambda=0.0$) | 49.9 | 42.0 | 26.3 | 20.0 | 24.9 | 29.0 | 51.0 | 47.5 | 68.0 | 64.3 | 72.0 | 72.3 | 54.7 | 40.7
($\lambda=0.2$) | 53.0 | 43.5 | 30.3 | 22.2 | 27.6 | 30.1 | 53.8 | 49.5 | 73.0 | 70.7 | 76.1 | 77.3 | 59.0 | 41.9
($\lambda=0.8$) | 62.1 | 54.5 | 41.0 | 32.2 | 32.6 | 34.6 | 60.4 | 52.2 | 77.3 | 75.9 | 82.1 | 83.5 | 65.7 | 49.2

Table 3: Performance across datasets evaluated using NMI and F1-score metrics,
with three $\lambda$ settings (0, 0.2, 0.8) from our method. The best and
second-best methods are highlighted in bold and underlined for each dataset-
metric pair. We fixed the optimal $\lambda$ value per dataset in our method
(i.e., 0.8) during comparison, and our method demonstrates the best or
comparable performance across most datasets. Unavailable results and non-
convergence are labeled as N/A and N/C.

| $C$ | $Q$ | NMI | F1
---|---|---|---|---
| GCN | GAT | GIN | SAGE | GCN | GAT | GIN | SAGE | GCN | GAT | GIN | SAGE | GCN | GAT | GIN | SAGE
Cora | 9.7 | 10.2 | 7.3 | 8.5 | 80.8 | 80.8 | 80.2 | 78.4 | 53.0 | 51.9 | 53.2 | 53.4 | 43.5 | 41.3 | 46.7 | 51.9
CiteSeer | 4.1 | 4.1 | 4.5 | 4.8 | 87.4 | 87.4 | 87.3 | 84.8 | 30.3 | 29.4 | 29.5 | 32.5 | 22.2 | 21.9 | 21.9 | 34.0
PubMed | 20.4 | 19.1 | 15.6 | 10.4 | 72.8 | 73.8 | 75.0 | 72.4 | 27.6 | 26.1 | 25.7 | 27.9 | 30.1 | 26.8 | 29.2 | 40.5
Amazon PC | 17.7 | 23.4 | 21.3 | 11.1 | 61.5 | 61.4 | 54.3 | 55.1 | 53.8 | 52.9 | 35.4 | 45.0 | 49.5 | 45.3 | 42.3 | 52.2
Amazon Photo | 8.6 | 8.5 | 6.9 | 5.2 | 71.6 | 72.1 | 69.9 | 63.8 | 73.0 | 72.1 | 61.8 | 60.6 | 70.7 | 69.4 | 60.6 | 58.2
Coauthor CS | 15.3 | 15.7 | 14.7 | 13.6 | 74.2 | 74.0 | 73.8 | 73.0 | 76.1 | 75.7 | 71.8 | 72.4 | 77.3 | 79.0 | 72.6 | 74.4
Coauthor PHYSICS | 22.3 | 20.3 | 18.5 | 14.5 | 67.3 | 67.1 | 66.7 | 66.0 | 59.0 | 59.2 | 59.7 | 64.0 | 41.9 | 43.4 | 48.7 | 55.9
AVERAGE | 14.0 | 14.5 | 12.7 | 9.7 | 73.6 | 73.8 | 72.5 | 70.5 | 53.3 | 52.5 | 48.2 | 50.8 | 47.9 | 46.7 | 46.0 | 52.4

Table 4: Performance of our method with various GNN base models assessed
across different metrics for all datasets, along with the calculation of
average performance across datasets. Notably, the performance remains
consistent across different variants of GNN models, demonstrating the absence
of a clear winner. This shows the robustness of our model in adapting to
diverse GNN architectures. For each dataset, we emphasize performance
variations: a performance increase of more than 5 is highlighted in bold,
while a decrease is marked with an underline, both in comparison to GCN.
Similarly, in the average row, we identify the best-performing GNN base model
in bold and the worst-performing with an underline, considering all metrics
across datasets.

### Performance

Table 2 and 3 show that our method achieves superior or comparable results
compared to the baselines across all evaluation metrics and datasets. Notably,
our method stands out with significant improvements in unsupervised graph
evaluation topology-based metrics such as graph conductance, and modularity
(Table 2). Our auxiliary objective also provides flexibility by tuning the
hyperparameter $\lambda$ to optimize NMI and F1 score to have better results
compared to the baselines (Table 3). While our method has a superior
performance in co-purchase and co-authorship networks under all four metrics,
it has a reasonably good performance in the citation networks. This shows the
generalizability of our algorithm in different types of networks.

### Robustness on Different Base GNNs

We conduct experiments to test the robustness of our method towards different
GNNs including GAT (Velickovic et al. 2017), GIN (Xu et al. 2018), and
GraphSAGE (Hamilton, Ying, and Leskovec 2017). Table 4 indicates that the
selection of the GNN model does not significantly alter the performance of our
method in terms of the evaluation measures. The average performance across
datasets for different metrics and GNN bases reveals specific enhancements:
GAT improves $Q$, GraphSAGE enhances $C$ and F1 scores on average. Conversely,
GIN results into general performance degradation.

Figure 1: Effect of auxiliary information on different metrics and datasets
with varying $\lambda$. As expected, $Q$ tends to decrease as $\lambda$
increases since $Q$ is dependent only on the graph structure whereas $\lambda$
adds weight to the label-based loss. The NMI increases as $\lambda$ increases
since NMI is directly related to labels.

### Auxiliary Information Effect

To check the effect of our auxiliary information, we vary the $\lambda$
hyperparameter in our additional experiments. Figure 1 illustrates the impact
of the hyperparameter on different metrics and datasets. The main goal of the
$\lambda$ parameter is to find the best underlying partition leveraging both
graph structure and auxiliary information. We can see modularity $Q$ has a
consistent value over increased $\lambda$, while NMI values are increasing
with a good amount. This shows the power of our auxiliary objective which
ensures flexibility to the user aiming to optimize their desired metric.
Additionally, the standard deviation of our results is small showing the
stability of our algorithm.

| Cora | Citeseer | Pubmed
---|---|---|---
| $Q\uparrow$ | NMI $\uparrow$ | $Q\uparrow$ | NMI $\uparrow$ | $Q\uparrow$ | NMI $\uparrow$
$\alpha=0.0$ | 80.8 | 43.5 | 87.4 | 22.2 | 72.8 | 30.1
$\alpha=0.5$ | 80.8 | 40.3 | 87.3 | 20.3 | 72.1 | 28.1
$\alpha=1.0$ | 80.8 | 38.5 | 87.4 | 19.5 | 71.7 | 26.3

Table 5: Performance of our method with varying $\alpha$ regularization
parameter. Additional regularization has a negligible effect while avoiding
potential trivial clustering.

#### Additional Regularization

In addition to the primary and the auxiliary information objective terms, we
propose another regularization term, the squared average node similarity,
$\alpha(\frac{1}{n^{2}}\sum_{ij}\langle
X_{i}X_{j}\rangle)^{2}=\alpha\|\bar{X}\|_{2}^{4}$, where $\bar{X}$ is the mean
of the node embedding vectors and $\alpha$ is a tunable parameter. This
regularizer can be used to avoid the formation of a trivial clustering where
all the nodes form a single cluster. Table 5 shows that adding this
regularizer have a negligible effect on our main objectives while avoiding
potential trivial clustering.

Figure 2: The number of communities varying $\lambda$ hyperparameter for
different datasets. We observe that the number of communities found is similar
across different $\lambda$ values except in some cases.

### Number of Communities

Our method does not assume the number of communities in the data, which is one
of its main strengths. The number of communities in real-world data is
generally unknown, and it is nontrivial to estimate it. To share some insight,
we provide the number of communities our method finds across different
datasets and $\lambda$ values. Figure 2 demonstrates the number of communities
varies for different datasets, and the results are consistent across different
$\lambda$ values.

## Conclusion

In this paper, we have studied the problem of graph clustering via maximizing
modularity. The existing techniques often require the number of clusters
beforehand or they fail to capitalize on the potential benefits of associated
node attributes and availability of supplementary information. To address
these, we have introduced , a novel neural framework that works well without a
predefined number of clusters. Moreover, our framework uses graph neural
networks (GNNs) to leverage the node attributes and additional information
present within the graph. The computational complexity of scales linearly with
the graph size. Through extensive experimentation across seven real-world
datasets of varying sizes and employing multiple distinct cluster quality
evaluation metrics, we have showcased the superior performance of : it
consistently has outperformed state-of-the-art approaches across most of the
settings.

## Acknowledgments

This material is based upon work partially supported by the National Science
Foundation under grant no. 2229876 and in part by funds provided by the
National Science Foundation, Department of Homeland Security, and IBM.

## References

  * Bhatia and Rani (2018) Bhatia, V.; and Rani, R. 2018.  Dfuzzy: a deep learning-based fuzzy clustering model for large graphs.  _Knowledge and Information Systems_ , 57: 159–181.
  * Bianchi, Grattarola, and Alippi (2020) Bianchi, F. M.; Grattarola, D.; and Alippi, C. 2020.  Spectral clustering with graph neural networks for graph pooling.  In _International Conference on Machine Learning_ , 874–883. PMLR.
  * Blondel et al. (2008) Blondel, V. D.; Guillaume, J.-L.; Lambiotte, R.; and Lefebvre, E. 2008.  Fast unfolding of communities in large networks.  _Journal of statistical mechanics: theory and experiment_ , 2008(10): P10008.
  * Bo et al. (2020) Bo, D.; Wang, X.; Shi, C.; Zhu, M.; Lu, E.; and Cui, P. 2020.  Structural deep clustering network.  In _Proceedings of the web conference 2020_ , 1400–1410.
  * Brandes et al. (2006) Brandes, U.; Delling, D.; Gaertler, M.; Görke, R.; Hoefer, M.; Nikoloski, Z.; and Wagner, D. 2006.  Maximizing modularity is hard.  _arXiv preprint physics/0608255_.
  * Choong, Liu, and Murata (2018) Choong, J. J.; Liu, X.; and Murata, T. 2018.  Learning community structure with variational autoencoder.  In _2018 IEEE international conference on data mining (ICDM)_ , 69–78. IEEE.
  * Felzenszwalb and Huttenlocher (2004) Felzenszwalb, P. F.; and Huttenlocher, D. P. 2004.  Efficient graph-based image segmentation.  _International journal of computer vision_ , 59: 167–181.
  * Fortunato and Hric (2016) Fortunato, S.; and Hric, D. 2016.  Community detection in networks: A user guide.  _Physics reports_ , 659: 1–44.
  * Frey and Dueck (2007) Frey, B. J.; and Dueck, D. 2007.  Clustering by passing messages between data points.  _science_ , 315(5814): 972–976.
  * Hamilton, Ying, and Leskovec (2017) Hamilton, W.; Ying, Z.; and Leskovec, J. 2017.  Inductive representation learning on large graphs.  _Advances in neural information processing systems_ , 30.
  * Hartigan and Wong (1979) Hartigan, J. A.; and Wong, M. A. 1979.  Algorithm AS 136: A k-means clustering algorithm.  _Journal of the royal statistical society. series c (applied statistics)_ , 28(1): 100–108.
  * Kipf and Welling (2016) Kipf, T. N.; and Welling, M. 2016.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_.
  * Klambauer et al. (2017) Klambauer, G.; Unterthiner, T.; Mayr, A.; and Hochreiter, S. 2017.  Self-normalizing neural networks.  _Advances in neural information processing systems_ , 30.
  * Kulatilleke, Portmann, and Chandra (2022) Kulatilleke, G. K.; Portmann, M.; and Chandra, S. S. 2022.  SCGC: Self-supervised contrastive graph clustering.  _arXiv preprint arXiv:2204.12656_.
  * Liu et al. (2023) Liu, Y.; Yang, X.; Zhou, S.; Liu, X.; Wang, S.; Liang, K.; Tu, W.; and Li, L. 2023.  Simple contrastive graph clustering.  _IEEE Transactions on Neural Networks and Learning Systems_.
  * Mandaglio, Amelio, and Tagarelli (2018) Mandaglio, D.; Amelio, A.; and Tagarelli, A. 2018.  Consensus community detection in multilayer networks using parameter-free graph pruning.  In _Advances in Knowledge Discovery and Data Mining: 22nd Pacific-Asia Conference, PAKDD 2018, Melbourne, VIC, Australia, June 3-6, 2018, Proceedings, Part III 22_ , 193–205. Springer.
  * Moradi, Ahmadian, and Akhlaghian (2015) Moradi, P.; Ahmadian, S.; and Akhlaghian, F. 2015.  An effective trust-based recommendation method using a novel graph clustering algorithm.  _Physica A: Statistical mechanics and its applications_ , 436: 462–481.
  * Newman (2006a) Newman, M. E. 2006a.  Finding community structure in networks using the eigenvectors of matrices.  _Physical review E_ , 74(3): 036104.
  * Newman (2006b) Newman, M. E. 2006b.  Modularity and community structure in networks.  _Proceedings of the national academy of sciences_ , 103(23): 8577–8582.
  * Newman and Girvan (2003) Newman, M. E.; and Girvan, M. 2003.  Mixing patterns and community structure in networks.  In _Statistical mechanics of complex networks_ , 66–87. Springer.
  * Park et al. (2019) Park, J.; Lee, M.; Chang, H. J.; Lee, K.; and Choi, J. Y. 2019.  Symmetric graph convolutional autoencoder for unsupervised graph representation learning.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 6519–6528.
  * Peixoto (2014) Peixoto, T. P. 2014.  Efficient Monte Carlo and greedy heuristic for the inference of stochastic block models.  _Physical Review E_ , 89(1): 012804.
  * Perozzi, Al-Rfou, and Skiena (2014) Perozzi, B.; Al-Rfou, R.; and Skiena, S. 2014.  Deepwalk: Online learning of social representations.  In _Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining_ , 701–710.
  * Scarselli et al. (2008) Scarselli, F.; Gori, M.; Tsoi, A. C.; Hagenbuchner, M.; and Monfardini, G. 2008.  The graph neural network model.  _IEEE transactions on neural networks_ , 20(1): 61–80.
  * Sen et al. (2008) Sen, P.; Namata, G.; Bilgic, M.; Getoor, L.; Galligher, B.; and Eliassi-Rad, T. 2008.  Collective classification in network data.  _AI magazine_ , 29(3): 93–93.
  * Shchur and Günnemann (2019) Shchur, O.; and Günnemann, S. 2019.  Overlapping community detection with graph neural networks.  _arXiv preprint arXiv:1909.12201_.
  * Shchur et al. (2018) Shchur, O.; Mumme, M.; Bojchevski, A.; and Günnemann, S. 2018.  Pitfalls of graph neural network evaluation.  _arXiv preprint arXiv:1811.05868_.
  * Shi and Malik (2000) Shi, J.; and Malik, J. 2000.  Normalized cuts and image segmentation.  _IEEE Transactions on pattern analysis and machine intelligence_ , 22(8): 888–905.
  * Sun et al. (2021) Sun, J.; Zheng, W.; Zhang, Q.; and Xu, Z. 2021.  Graph neural network encoding for community detection in attribute networks.  _IEEE Transactions on Cybernetics_ , 52(8): 7791–7804.
  * Tsitsulin et al. (2023) Tsitsulin, A.; Palowitch, J.; Perozzi, B.; and Müller, E. 2023.  Graph clustering with graph neural networks.  _Journal of Machine Learning Research_ , 24(127): 1–21.
  * Velickovic et al. (2017) Velickovic, P.; Cucurull, G.; Casanova, A.; Romero, A.; Lio, P.; Bengio, Y.; et al. 2017.  Graph attention networks.  _stat_ , 1050(20): 10–48550.
  * Veličković et al. (2018) Veličković, P.; Fedus, W.; Hamilton, W. L.; Liò, P.; Bengio, Y.; and Hjelm, R. D. 2018.  Deep graph infomax.  _arXiv preprint arXiv:1809.10341_.
  * Wang et al. (2019) Wang, C.; Pan, S.; Hu, R.; Long, G.; Jiang, J.; and Zhang, C. 2019.  Attributed graph clustering: A deep attentional embedding approach.  _arXiv preprint arXiv:1906.06532_.
  * Wang et al. (2017) Wang, C.; Pan, S.; Long, G.; Zhu, X.; and Jiang, J. 2017.  Mgae: Marginalized graph autoencoder for graph clustering.  In _Proceedings of the 2017 ACM on Conference on Information and Knowledge Management_ , 889–898.
  * Wang et al. (2010) Wang, J.; Li, M.; Chen, J.; and Pan, Y. 2010.  A fast hierarchical clustering algorithm for functional modules discovery in protein interaction networks.  _IEEE/ACM Transactions on Computational Biology and Bioinformatics_ , 8(3): 607–620.
  * Wu et al. (2020) Wu, L.; Zhang, Q.; Chen, C.-H.; Guo, K.; and Wang, D. 2020.  Deep learning techniques for community detection in social networks.  _IEEE Access_ , 8: 96016–96026.
  * Xia et al. (2021) Xia, W.; Wang, Q.; Gao, Q.; Zhang, X.; and Gao, X. 2021.  Self-supervised graph convolutional network for multi-view clustering.  _IEEE Transactions on Multimedia_ , 24: 3182–3192.
  * Xu et al. (2018) Xu, K.; Hu, W.; Leskovec, J.; and Jegelka, S. 2018.  How powerful are graph neural networks?  _arXiv preprint arXiv:1810.00826_.
  * Yang and Leskovec (2012) Yang, J.; and Leskovec, J. 2012.  Defining and evaluating network communities based on ground-truth.  In _Proceedings of the ACM SIGKDD Workshop on Mining Data Semantics_ , 1–8.
  * Yang et al. (2016) Yang, L.; Cao, X.; He, D.; Wang, C.; Wang, X.; and Zhang, W. 2016.  Modularity based community detection with deep learning.  In _IJCAI_ , volume 16, 2252–2258.
  * Ying et al. (2018a) Ying, R.; He, R.; Chen, K.; Eksombatchai, P.; Hamilton, W. L.; and Leskovec, J. 2018a.  Graph convolutional neural networks for web-scale recommender systems.  In _Proceedings of the 24th ACM SIGKDD international conference on knowledge discovery & data mining_, 974–983.
  * Ying et al. (2018b) Ying, Z.; You, J.; Morris, C.; Ren, X.; Hamilton, W.; and Leskovec, J. 2018b.  Hierarchical graph representation learning with differentiable pooling.  _Advances in neural information processing systems_ , 31.
  * Yue et al. (2022) Yue, L.; Jun, X.; Sihang, Z.; Siwei, W.; Xifeng, G.; Xihong, Y.; Ke, L.; Wenxuan, T.; Wang, L. X.; et al. 2022.  A survey of deep graph clustering: Taxonomy, challenge, and application.  _arXiv preprint arXiv:2211.12875_.
  * Zhang and Chen (2018) Zhang, M.; and Chen, Y. 2018.  Link prediction based on graph neural networks.  _Advances in neural information processing systems_ , 31.
  * Zhang, Ramakrishnan, and Livny (1996) Zhang, T.; Ramakrishnan, R.; and Livny, M. 1996.  BIRCH: an efficient data clustering method for very large databases.  _ACM sigmod record_ , 25(2): 103–114.
  * Zhang et al. (2019) Zhang, X.; Liu, H.; Li, Q.; and Wu, X.-M. 2019.  Attributed graph clustering via adaptive graph convolution.  _arXiv preprint arXiv:1906.01210_.
  * Zhou et al. (2020) Zhou, J.; Cui, G.; Hu, S.; Zhang, Z.; Yang, C.; Liu, Z.; Wang, L.; Li, C.; and Sun, M. 2020.  Graph neural networks: A review of methods and applications.  _AI open_ , 1: 57–81.


\tilde{X}_{t_n}^{n,u} & = & \underline{\xi}^{n,u}(W_{[0,t_n]}, Z),    \;\; u\in U,
\end{array}\right.
\end{equation}

where $\tilde \alpha$ is defined by
α̃_t^u=α̃(u,t, W_·∧t,Z),
t∈[0,T], u∈U.
Using Proposition <ref> with $\tilde{W}^u=W$ and $\tilde{Z}^u=Z$ for $u\in U$, we get
(X^u,W^u,Z^u)     =     (X̃^u, W,Z)
(X^n,u,W^u,Z^u)     =     (X̃^n,u,W,Z)
for all $n\geq 1$ and $u\in U$. Now, using condIniConvL2 we get from classical estimates on diffusion processes that
]λ(ụ)          0 ,
which gives the result.

Acknowledgments.  Marco Fuhrman wishes to thank Huyên Pham for the invitation  to Université Paris Cité where this work began.

[1]
Alexander Aurell, René Carmona, and Mathieu Lauriere.
Stochastic graphon games: Ii, the linear quadratic case.
Applied Mathematics & Optimization, 85(3), 2022.

[2]
Erhan Bayraktar, Suman Chakraborty, and Ruoyu Wu.
Graphon mean field systems.
Annals of Applied Probability, 33(5):3587–3619, 2023.

[3]
Alain Bensoussan, Jacques Frehse, and Sheung Chu Phillip Yam.
Mean field games and mean field type control theory.
Springer, 2013.

[4]
Dimitri Bertsekas and Steven Shreve.
Stochastic Optimal Control: The Discrete-Time
Athena Scientific, 1996.

[5]
Charles Bertucci and Matthias Rakotomala.
Strategic geometric graphs through mean field games.
Hal-04548364, 2024.

[6]
Rainer Buckdahn, Juan Li, Shige Peng, and Catherine Rainer.
Mean-field stochastic differential equations and associated PDEs.
The Annals of Probability, 45(2):824–878, March 2017.
Publisher: Institute of Mathematical Statistics.

[7]
Peter E. Caines and Minyi Huang.
Graphon mean field games and their equations.
SIAM Journal on Control and Optimization, 59(6):4373–4399,

[8]
Pierre Cardaliaguet, Samuel Daudin, and Panaglotis Souganidis.
An algebraic convergence rate for the optimal control of
mckean-vlasov dynamics.
SIAM Journal on Control and Optimization, 61(6):3341–3369,

[9]
René Carmona and Francois Delarue.
Probabilistic Theory of Mean Field Games with Applications, vol.
Springer, 2018.

[10]
René Carmona and Francois Delarue.
Probabilistic Theory of Mean Field Games with Applications, vol
Springer, 2018.

[11]
Fabio Coppini, Anna De Crescenzo, and Huyên Pham.
Non linear graphon mean field systems.
arXiv:2402.08628, 2024.

[12]
Samuel Daudin, Francois Delarue, and Joe Jackson.
On the optimal rate for the convergence problem in mean field
arXiv:2305.08423, 2023.

[13]
Fabrice Djete.
Extended mean field control problem: a propagation of chaos result.
Electronic Journal of Probability, 27:1–53, 2022.

[14]
Pierre-Emmanuel Jabin, David Poyato, and Juan Soler.
Mean-field limit of non exchangeable systems.
arXiv:2112.15406, 2021.

[15]
Daniel Lacker.
Limit theory for controlled McKean-Vlasov dynamics.
SIAM Journal on Control and Optimization, 55(3):1641–1672,

[16]
Daniel Lacker and Agathe Soret.
A Label-State Formulation of Stochastic Graphon Games and
Approximate Equilibria on Large Networks.
Mathematics of Operations Research, 48(4):1987–2018, November
Publisher: INFORMS.

[17]
Pierre-Louis Lions.
Cours au collège de france.
In www.college-de-france.

[18]
László Lovász.
Large Networks and Graph Limits, volume 60 of Colloquium Publications.
American Mathematical Society, 2012.

[19]
Yeneng Sun.
The exact law of large numbers via fubini extension and
characterization of insurable risks.
Journal of Economic Theory, 144(1):432–443, 2006.

[20]
Jerzy Zabczyk.
Chance and Decision. Stochastic Control in Discrete
Quaderni. Scuola Normale Superiore di Pisa, 1996.

§ HETEROGENEOUS SYSTEMIC RISK

THIS SECTION IS A SKETCH

The example similar to the one by Carmona et al:

X^u_s -\int_U G(u,v)\,\E[X^{v}_s]\,dv
\right)+\alpha^u_s
\right]\,ds + \sigma^u\,dW^u_s

\begin{array}{l}\displaystyle
J=\int_U \E\bigg[
\int_t^T\bigg\{
\frac{1}{2}(\alpha^u_s)^2+
\frac{\eta^u}{2}\left(
X^u_s -\int_U G(u,v)\,\E[X^{v}_s]\,dv
\right)^2
\\\displaystyle
\left(
X^u_s -\int_U G(u,v)\,\E[X^{v}_s]\,dv
\right)
\bigg\}ds
+ \frac{r^u}{2}\left(
X^u_T -\int_U G(u,v)\,\E[X^{v}_T]\,dv
\right)^2
\bigg]\,du
\end{array}

with $\eta^u\ge0$, $r^u\ge0$ and appropriate conditions on $\eta^u,q^u$. This corresponds to our general model with

\begin{align}
b(u,x,\mu,a) &= \; k^u \Big( x - \int_U G(u,v) \bar\mu^v dv \Big) + a, \quad \quad  \sigma(u,x,\mu,a) \; = \; \sigma^u, \\
f(u,x,\mu,a) & = \; \frac{\eta^u}{2} \Big( x - \int_U G(u,v) \bar\mu^v dv \Big)^2 + \frac{1}{2}a^2
+ q^ua \Big( x - \int_U G(u,v) \bar\mu^v dv \Big),  \\
g(u,x,\mu) & = \; \frac{r^u}{2} \Big( x - \int_U G(u,v) \bar\mu^v dv \Big)^2,
\end{align}

where we set $\bar\mu^u$ $=$ $\int_\R x \mu^u(dx)$.

The Bellman equation is written as

\begin{align}
& \partial_t \upsilon(t,\mu)  \\
& +  \int_U \int_\R \inf_{a \in A}
\Big[ b(u,x,\mu,a) \partial_x \dmu \upsilon(t,\mu)(u,x)  + \frac{1}{2}|\sigma^u|^2 \partial^2_x \dmu \upsilon(t,\mu)(u,x)  + f(u,x,\mu,a) \Big]
\mu^u(dx) du
\; = 0,
\end{align}

with the terminal condition $\upsilon(t,\mu)$ $=$ $\int_U \int_\R g(u,x,\mu) \mu^u(dx) du$, and a candidate for the optimal feedback control is  $\mra$ $=$ $(\mra^u)_{u\in U}$ with

\begin{align}
\mra^u(t,x,\mu) &= \; -  q^u \Big( x - \int_U G(u,v) \bar\mu^v dv \Big) -   \partial_x \dmu \upsilon(t,\mu)(u,x).
\end{align}

We look for a solution in the form:

\begin{align}
\upsilon(t,\mu) &= \; \int_U \Big[  K_t^u \int_\R \big( x - \int_U G(u,v) \bar\mu^v dv \big)^2 \mu^u(dx)  \\
& \quad \quad \quad + \;   \Gamma_t^u \int_\R \big( x - \int_U G(u,v) \bar\mu^v dv \big) \mu^u(dx)
+ R_t^u \Big] du
\end{align}

for some deterministic functions $t$ $\mapsto$ $(K_t^u)_{u\in U}$, $(\Gamma_t^u)_{u\in U}$ and $(R_t^u)_{u\in U}$.

We write $\upsilon(t,\mu)$ as

\begin{align}
\upsilon(t,\mu) &= \; \int_U\int_\R  K_t^u \varphi(u,x,\mu) \mu^u(dx) du  + \int_U\int_\R  \Gamma_t^u \vartheta(u,x,\mu) \mu^u(dx) du
+ \int_U R_t^u du     \\
\mbox{ with } \;\;\;  \vartheta(u,x,\mu) &= \;  x - \int_U G(u,v) \bar\mu^v dv, \quad \varphi(u,x,\mu)
\; = \; \vartheta(u,x,\mu)^2.
\end{align}

By the chain rule, we have

\begin{align}
\dmu \upsilon(t,\mu)(u,x) & = \; K_t^u \varphi(u,x,\mu) + \int_U \int_\R K_t^v \dmu \varphi(v,y,\mu)(u,x) \mu^v(dy) dv \\
& \quad + \Gamma_t^u \vartheta(u,x,\mu) + \int_U \int_\R \Gamma_t^v \dmu \vartheta(v,y,\mu)(u,x) \mu^v(dy) dv  \\
\mbox{ with } \;\;\; \dmu \varphi(u,x,\mu)(v,y) &= \; 2 \vartheta(u,x,\mu) \dmu \vartheta(u,x,\mu)(v,y), \quad
\dmu \vartheta(u,x,\mu)(v,y)  =  -2  G(u,v) y,
\end{align}

and thus

\begin{align}
\dmu \upsilon(t,\mu)(u,x) & = \; K_t^u \vartheta(u,x,\mu)^2  - 2 x\int_U \int_\R K_t^v \vartheta(v,y,\mu) G(v,u) \mu^v(dy) dv \\
&  \quad  + \; \Gamma_t^u \vartheta(u,x,\mu) - 2 x\int_U  \Gamma_t^v G(v,u)  dv  \\
& = \; K_t^u \vartheta(u,x,\mu)^2 +  \Gamma_t^u \vartheta(u,x,\mu)  \\
& \quad - \;  2 x \int_U \Big( K_t^v \big[  \bar\mu^v - \int_U G(v,w) \bar\mu^w dw \big] + \Gamma_t^v \Big) G(v,u) dv \\
\partial_x \dmu \upsilon(t,\mu)(u,x) & = \; 2 K_t^u \vartheta(u,x,\mu) + \Gamma_t^u
- 2 \int_U  \Big( K_t^v \big[  \bar\mu^v - \int_U G(v,w) \bar\mu^w dw \big] + \Gamma_t^v \Big) G(v,u) dv \\
\partial^2_x \dmu \upsilon(t,\mu)(u,x) & = \; 2 K_t^u.
\end{align}

This also yields

\begin{align}
\mra^u(t,x,\mu) &= \; -  (q^u + 2 K_t^u) \Big( x - \int_U G(u,v) \bar\mu^v dv \Big)  \\
& \quad - \Gamma_t^u + 2 \int_U  \Big( K_t^v \big[  \bar\mu^v - \int_U G(v,w) \bar\mu^w dw \big] + \Gamma_t^v \Big) G(v,u) dv.
\end{align}

We look for a solution of the Bellman equation in the form:

\[
\upsilon(t,\mu) = \int_U\int_\R K^u_t\left(x-\int_U G(u,v)\bar\mu^v\d v\right)^2\mu^u(\d x)\d u + \int_U\int_U\Gamma_t^{u,v}\bar\mu^u\bar\mu^v\d u\d v + \int_UR_t^u\d u,
\]

for some deterministic functions $t\mapsto(K^u_t)_{u\in I},(\Gamma^{u,v}_t)_{u,v\in U}, (R^u_t)_{u\in U}$, with $\Gamma_t^{u,v}=\Gamma_t^{v,u}$.
We compute the following derivatives:

\begin{align*}
\dmu \upsilon(t,\mu)(u,x) &= K^u_tx^2 -2x\int_U G(u,v)(K^u_t+K^v_t)\bar\mu^v\d v\\
& + 2x\int_U\int_U K^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v\\
& + 2x\int_U\Gamma_t^{u,v}\bar\mu^v\d v\\
\partial_x\dmu \upsilon(t,\mu)(u,x) &= 2K^u_tx - 2\int_U G(u,v)(K^u_t+K^v_t)\bar\mu^v\d v\\
&+ 2\int_U\int_U K^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v\\
&+ 2\int_U\Gamma_t^{u,v}\bar\mu^v\d v\\
\partial_x^2\dmu \upsilon(t,\mu)(u,x) &= 2K^u_t.
\end{align*}

Substituting in the Bellman equation, we obtain:

\begin{align*}
&\int_U\int_\R \dot K^u_t x^2\mu^u(\d x)\d u + \int_U\int_U\int_U \dot K^u_tG(u,v)G(u,w)\bar\mu^v\bar\mu^w\d v\d w\d u -2\int_U\int_U\dot K^u_tG(u,v)\bar\mu^v\bar\mu^u\d v\d u\\
&+\int_U\int_U\dot\Gamma_t^{u,v}\bar\mu^v\bar\mu^u\d v\d u + \int_U\dot R^u_t\d u\\
&+\int_U\int_\R \inf_a\bigg\{ \bigg( 2K^u_tx - 2\int_UG(u,v)(K^u_t+K^v_t)\bar\mu^v\d v + 2\int_U\int_UK^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v \\
&\qquad+ 2\int_U\Gamma^{u,v}_t\bar\mu^v\d v \bigg)\big( k^u(x-\int_UG(u,v)\bar\mu^v\d v)+a\big) + K^u_t(\sigma^u_t)^2\\
&\qquad+ \frac{1}{2}a^2 + \frac{\eta^u}{2}\bigg( x-\int_UG(u,v)\bar\mu^v\d v \bigg)^2 + q^ua\bigg(x-\int_UG(u,v)\bar\mu^v\d v\bigg)\bigg\}\mu^u(\d x)\d u\\
=& \int_U\int_\R \dot K^u_t x^2\mu^u(\d x)\d u + \int_U\int_U\int_U \dot K^u_tG(u,v)G(u,w)\bar\mu^v\bar\mu^w\d v\d w\d u -2\int_U\int_U\dot K^u_tG(u,v)\bar\mu^v\bar\mu^u\d v\d u\\
&+\int_U\int_U\dot\Gamma_t^{u,v}\bar\mu^v\bar\mu^u\d v\d u + \int_U\dot R^u_t\d u\\
&+ \int_U\int_\R \bigg( 2K^u_tk^ux^2-2k^ux\int_UG(u,v)(K^u_t+K^v_t)\bar\mu^v\d v + 2k^ux\int_U\int_U K^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v\\
& +2k^ux\int_U\Gamma_t^{u,v}\bar\mu^v\d v -2K^u_txk^u\int_UG(u,v)\bar\mu^v\d v + 2k^u\int_U\int_UG(u,v)(K^u_t + K^v_t)G(u,w)\bar\mu^v\bar\mu^w\d v\d w\\
&-2k^u\int_u\int_U\int_U K^w_tG(w,v)G(w,u)G(u,z)\bar\mu^v\bar\mu^z\d w\d v\d z\\
&-2k^u\int_U\int_U\Gamma_t^{u,v}G(u,w)\bar\mu^v\bar\mu^w\d v\d w + K^u_t(\sigma_t^u)^2 +\frac{\eta^u}{2}x^2 + \frac{\eta^u}{2}\int_U\int_UG(u,v)G(u,w)\bar\mu^v\bar\mu^w\d v\d w\\
&-\eta^ux\int_UG(u,v)\bar\mu^v\d v\bigg)\mu^u(\d x)\d u\\
&+\int_U\int_\R \inf_a\bigg\{ \bigg(2K^u_tx - 2\int_U G(u,v)(K^u_t+K^v_t)\bar\mu^v\d v + 2\int_U\int_U K^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v\\
&\qquad + 2\int_U\Gamma_t^{u,v}\bar\mu^v\d v + q^ux - q^u\int_UG(u,v)\bar\mu^v\d v\bigg)a + \frac{1}{2}a^2\bigg\}\mu^u(\d x)\d u\\
=& \int_U\int_\R\big( \dot K^u_t + 2K^u_tk^u + \frac{\eta^u}{2} + 2(K^u_t)^2 + \frac{1}{2}(q^u)^2 + K^u_tq^u \big)x^2\mu^u(\d x)\d u\\
&\int_U\int_U \bigg( \dot\Gamma^{u,v}_t + 2k^u\Gamma^{u,v}_t - 2\int_UK^wG(u,w)\Gamma_t^{w,v}\d w - 2\int_U\Gamma_t^{w,v}\Gamma_t^{u,w}\d w \\
&+ 4\int_U\int_UK^w_tG(w,v)G(w,z)\Gamma_t^{u,z}\d w\d z -4\int_UG(w,v)(K^w_t+K^v_t)\Gamma_t^{u,w}\d w\\
& -2\int_Uq^wG(u,w)\Gamma_t^{w,v}\d w + \Psi(K_t,\dot K_t)\bigg)\bar\mu^v\bar\mu^u\d v\d u\\
&+ \int_U\big(\dot R^u_t + K^u_t(\sigma^u_t)^2\big)\d u =0,
\end{align*}

where $\Psi(K_t,\dot K_t)$ represents terms depending on $K_t,\dot K_t$. The optimal control is

\begin{align}
\hat a(u,t,x,\mu) =& -\partial_x\dmu \upsilon(t,\mu)(u,x) -q^u\left(x-\int_UG(u,v)\bar\mu^v\d v\right)\\
=& -2K^u_tx + 2\int_U G(u,v)(K^u_t+K^v_t)\bar\mu^v\d v\\
&- 2\int_U\int_U K^w_tG(w,v)G(w,u)\bar\mu^v\d w\d v\\
&- 2\int_U\Gamma_t^{u,v}\bar\mu^v\d v - q^u\left(x-\int_UG(u,v)\bar\mu^v\d v\right)\\
=& -(q^u+2K^u_t)\left( x-\int_UG(u,v)\bar\mu^v\d v \right) \\
&+ 2\int_UG(u,v)K^v_t\left( \bar\mu^v - \int_UG(u,v)\bar\mu^w\d w\right)\d v -2\int_U\Gamma_t^{u,v}\bar\mu^v\d v.
\end{align}

We want to prove existence and uniqueness of the following Riccati equation:

\begin{align*}
&\dot\Gamma_t^{u,v} + k^u\Gamma_t^{u,v} + k^v\Gamma_t^{v,u} -\int_UK^wG(u,w)\Gamma_t^{w,v}\d w-\int_UK^wG(v,w)\Gamma_t^{w,u}\d w \\
&-2\int_U\Gamma_t^{w,v}\Gamma_t^{u,w} + 2\int_U\int_U K^w_tG(w,v)G(w,z)\Gamma_t^{u,z}\d w\d z\\
&+2\int_U\int_UK^w_tG(w,u)G(w,z)\Gamma_t^{v,z}\d w\d z -2\int_UG(w,v)(K_t^w+K_t^v)\Gamma_t^{u,w}\d w\\
&-2\int_UG(w,u)(K_t^w+K_t^v)\Gamma_t^{u,w}\d w -\int_Uq^wG(u,w)\Gamma_t^{w,v}\d w\\
&-\int_Uq^wG(v,w)\Gamma_t^{w,u}\d w + \Psi(K_t,\dot K_t)=0,\\
&\Gamma_T^{u,v} = \int_U\frac{r^w}{2}G(u,w)G(w,v)\d w - \frac{1}{2}r^uG(u,v) - \frac{1}{2}r^vG(u,v).
\end{align*}

Let us observe that, in the case with constant coefficients, we would obtain a matrix Riccati equation. Indeed, let us suppose that there exists a finite collection of intervals $\{U_i : i=1,\dots,N\}$, such that:

\begin{align}
& k^u = \sum_{i=1}^N k^i\I_{U_i}(u),\qquad \eta^u=\sum_{i=1}^N\eta^i\I_{U_i}(u),\qquad q^u=\sum_{i=1}^Nq^i\I_{U_i}(u),\\
&r^u=\sum_{i=1}^Nr^i\I_{U_i}(u) \qquad G(u,v)=\sum_{i,j=1}^N\xi_{ij}\I_{U_i}(u)\I_{U_j}(v),
\end{align}

then $\Gamma$ would also be piecewise constant and the above equation reduces to a matrix Riccati equations for $\Gamma_t =(\Gamma^{ij}_t)_{ij}$, namely:

\begin{align*}
&\Gamma_t +\Gamma_tk + k\Gamma_t - GK_tM\Gamma_t - \Gamma_tK_tMG -2\Gamma_tM\Gamma_t + 2\Gamma_tMGMK_tG + 2GMK_tG\Gamma_t\\
&\quad -2\Gamma_tMK_tG -2\Gamma_tMGK_t -2GK_tM\Gamma_t -2K_tGM\Gamma_t - GqM\Gamma_t-\Gamma_tMqG +\Psi =0,\\
& \Gamma_T = \frac{1}{2}GMrG + \frac{1}{2}Gr + \frac{1}{2}rG,
\end{align*}

where, with some abuse of notation,

\begin{align}
& k =\text{diag}([k^1,\dots,k^N]),\qquad G = (\xi_{ij})_{ij}, \qquad K_t =\text{diag}([K_t^1,\dots,K_t^N]),\\
&M=\text{diag}([|U_1|,\dots,|U_N|]),\qquad q=\text{diag}([q^1,\dots,q^N]),\qquad r=\text{diag}([r^1,\dots,r^N]).
\end{align}