\section{Introduction}
Traditional clustering techniques partition the data directly in the input space by drawing regular boundaries to separate the clusters.
This makes them unsuitable to handle complex data structures, such as images or time series, which lie in high-dimensional spaces where the relationships between samples are highly non-linear. 
Deep learning techniques allow to transform data samples into suitable representations, which can partitioned into meaningful clusters.
Remarkably, end-to-end deep learning frameworks can directly map complex data directly into their cluster assignments~\cite{kampffmeyer2019deep}.

Of particular interest for this work, are those data characterized by relationships, or interactions, among samples that are described by a graph. 
Graph Neural Networks (GNNs) are deep learning architectures specifically designed to process and make inference on such data~\cite{hamilton2020graph}.
Recently, GNNs have been adopted to cluster the nodes of an attributed graph based on their features and the graph topology.
Inspired by algorithms such as spectral clustering, such GNNs optimize an unsupervised loss composed of two terms: the first ensures that connected nodes are assigned to the same cluster; the second is a balancing term, which prevents degenerate solutions both by encouraging the samples to be assigned to only one cluster and the clusters to have similar size~\cite{bianchi2020spectral, tsitsulin2020graph}.

This paper proposes to considerably simplify the clustering objective optimized by previous GNN models by proposing a minimalist unsupervised loss, which consists only of a balancing term.
The proposed loss is used to train a GNN composed of standard message passing layers, which operate on a suitable connectivity matrix.
The empirical evaluation shows that the proposed model can achieve competitive results in terms of clustering performance and computational time.

\section{Background}
Let a graph be represented by a tuple $G = \{ \mathcal{V}, \mathcal{E} \}$, $|\mathcal{V}|=N$, with node set $\mathcal{V}$ and edge set $\mathcal{E}$. Each node is associated with a vector attribute in $\mathbb{R}^F$.
A graph is conveniently described by its adjacency matrix $\A \in \mathbb{R}^{N \times N}$ and the node features $\X \in \R^{N \times F}$.

\subsection{Spectral Clustering}
Graph clustering aims at partitioning the nodes in $K$ subsets, so that the similarity between nodes in the same subset is maximized.
The most famous graph clustering approach is spectral clustering, which exploit the $k-$way mincut objective to find a partition that minimizes the volume of edges corssing different clusters~\cite{von2007tutorial}.
To avoid degenerate solutions, the objective function includes a balancing term that penalizes partitions where clusters have very unequal sizes.
Specifically, the balanced $K$-cut objective can be defined as a ratio of two set functions:
\begin{equation}
\label{eq:balanced_k_cut}
    \min_{C_1, \dots, C_K} \sum_{k=1}^K \frac{\text{cut}(C_k, \bar{C}_k)}{\hat{B}(C_k)},
\end{equation}
where $\hat{B}(\cdot)$ is a set function that balances the size of the clusters in the partition.
Depending on the choice of $\hat{B}(\cdot)$, one obtains different cuts, such as ratio cut and normalized cut~\cite{hein2011beyond}.
The numerator of \eqref{eq:balanced_k_cut} can be expressed in matrix form. 
To see that, first one rewrites $\text{cut}(C_k, \bar{C}_k)$ as 
\begin{equation*}
    \label{eq:cut}
     \sum \limits_{i \in C_k, j \in \bar C_k} a_{ij}(1- z_i z_j) \;\; \text{s.t.} \;\; z_i, z_j \in \{ -1, 1 \},
\end{equation*}
Then,
\begin{equation*}
    \begin{aligned}
    & \sum \limits_{i,j} a_{ij}(1- z_i z_j) = \sum \limits_{i,j} a_{ij} \left(\frac{z_i^2 + z_j^2}{2} - z_i z_j \right) \\
    & = \frac{1}{2} \sum \limits_i \Bigg[ \sum \limits_j a_{ij} \Bigg]z_i^2 + \frac{1}{2} \sum \limits_j \Bigg[ \sum \limits_i a_{ij} \Bigg] z_j^2 \dots \\
    & - \sum \limits_{i,j} a_{ij} z_i z_j \\
    & = \frac{1}{2} \sum \limits_i d_{ii} z_i^2 + \frac{1}{2} \sum \limits_j d_{jj} z_j^2 - \z^T \A \z \\
    & = \z^T \D \z - \z^T \A \z = \z^T \Lb \z,
    \end{aligned}
\end{equation*}
where $\Lb$ is the graph Laplacian.
The relaxation done in spectral clustering to handle the discrete optimization problem is:
\begin{equation}
    \label{eq:general_sc_relax}
    \min_{\z_k \in \{ -1, 1 \}^{N}} \sum \limits_{k=1}^K \frac{\z_k^T \Lb \z_k}{\hat{B}(C_k)} \rightarrow \min_{\s_k \in \R^{N}} \sum \limits_{k=1}^K \frac{\s_k^T \Lb \s_k}{B(C_k)},
\end{equation}
being $B(C_k)$ the continuous counterpart of $\hat{B}(C_k)$. 
\newline

\textit{Remark 1:}
Besides the graph Laplacian, other operators matching the sparsity pattern of the adjacency matrix can be used in problem~\eqref{eq:general_sc_relax}. One of such operators is the symmetrically normalized Laplacian, $\Lb_s = \I - \D^{-\frac{1}{2}} \A \D^{-\frac{1}{2}}$, which generally yields a different partition as the edges to be cut are weighted by the degree of their end nodes.
\newline

\textit{Remark 2:}
The term $\s^T \Lb \s$ in~\eqref{eq:general_sc_relax} measures the local quadratic variation (LQV) of $\s$ on the graph. Laplacian smoothing minimizes LQV, by making similar the elements $s_i$ and $s_j$ if nodes $i$ and $j$ are connected. When using $\Lb_s$, the LQV is:

\begin{equation}
    \label{eq:lap_smooth_symm}
    \s^T \Lb_s \s = \frac{1}{2} \sum_{(i,j) \in \mathcal{E}} a_{i,j}\left(\frac{s_i}{\sqrt{d_i}} - \frac{s_j}{\sqrt{d_j}}\right)^2.
\end{equation}

\subsection{Graph Neural Networks}
The main building block of a GNN is the message passing (MP) layer that, first, combines the node features with those of the neighbors on the graph.
Then, the aggregated features are mapped into a new representation by applying an affine transformation and a nonlinearity~\cite{gilmer2017neural}.
A basic MP layer is implemented as follows:
\begin{equation}
    \label{eq:mp}
    \X^{(l+1)} = \texttt{MP}(\X^{(l)}, \tilde \A) = \sigma(\tilde \A \X^{(l)} \boldsymbol{\Theta}_l),
\end{equation}
where $\X^{(l)}$ and $\X^{(l+1)}$ are, respectively, the input and output node features of the $l$-th MP layer, $\tilde \A$ is an operator matching the sparsity pattern of $\A$, $\sigma$ is a nonlinear activation function, and $\boldsymbol{\Theta}_l$ are trainable parameters.


\subsection{Clustering with GNNs}
MinCutPool~\cite{bianchi2020spectral} is a GNN layer that computes soft cluster assignments as:
\begin{equation}
    \label{eq:soft_assign}
    \Sb = \texttt{softmax} \left( \texttt{MLP} \left( \bar \X, \boldsymbol{\Theta}_\text{MLP} \right) \right) \in \mathbb{R}^{N \times K},
\end{equation}
where $K$ is the number of clusters, $\bar \X$ are node features generated by a stack of one or more MP layers, and $\texttt{MLP}(\cdot)$ denotes a multi-layer perceptron with trainable parameters $\boldsymbol{\Theta}_\text{MLP}$.
The \texttt{softmax} function ensures that $\Sb$ is a proper cluster assignment matrix, since $\Sb \bm{1} = \bm{1}$ and $0 \leq s_{i,j} \leq 1$.

To learn the cluster assignments, MinCutPool optimizes the following unsupervised loss:
\begin{equation}
    \label{eq:norm_cut}
    \mathcal{L}_{mc} = 
    \underbrace{- \frac{\textrm{Tr} ( \Sb^T \tilde \A \Sb )}{\textrm{Tr} ( \Sb^T \tilde \D \Sb)}}_{\mathcal{L}_q} + 
    \underbrace{\bigg{\lVert} \frac{\Sb^T\Sb}{\|\Sb^T\Sb\|_F} - \frac{\I_K}{\sqrt{K}}\bigg{\rVert}_F}_{\mathcal{L}_b}, 
\end{equation}
where $\tilde{\A} = \D^{-\frac{1}{2}}\A\D^{-\frac{1}{2}}$ and $\tilde \D$ is the degree matrix of $\tilde \A$.
The first term, $\mathcal{L}_q$, minimizes the LQV, while $\mathcal{L}_b$ is a balancing term that helps preventing degenerate solutions.
Compared to problem~\eqref{eq:general_sc_relax}, the LQV and the balancing terms are summed rather than taking their ratio.
This helps both to prevent numerical issues when $\mathcal{L}_b$ gets too small and to keep $\mathcal{L}_{mc}$ in a controlled range, which is desirable when the GNN must minimize also other losses. 

\section{Proposed approach}

The cluster assignments $\Sb$, computed as in~\eqref{eq:soft_assign}, can be optimized by minimizing:
\begin{empheq}[box={\mymath[sharp corners]}]{equation}
    \label{eq:loss}
    \mathcal{L} = - \textrm{Tr} \left( \sqrt{\Sb^T\Sb} \right)
\end{empheq}
The proposed loss simplifies $\mathcal{L}_{mc}$ considerably as it consists only of a balancing term.
Such a simplification offers the following advantages:
\begin{itemize}
    \item The computational complexity is reduced, as less operations are needed to compute $\mathcal{L}$.
    \item Fewer competing terms in the loss can ease the training and speed-up the convergence.
    \item There are no ratios in $\mathcal{L}$, which could cause numerical instability during training.
\end{itemize}

Despite its simplicity, the proposed loss can still yield an optimal clustering assignment.
The two key insights that motivated its design are presented in the following.

\subsection{Removal of the LQV term}
\label{sec:lqv}
The absence of the LQV term in the loss can be compensated by the MP layers that generate the features $\bar \X$ used in \eqref{eq:soft_assign}.
In particular, consider the following MP layer:
\begin{empheq}[box={\mymath[sharp corners]}]{equation}
    \label{eq:mp_LQV}
    \small
    \X^{(l+1)} = \sigma\left( \left[\I - \delta (\I - \D^{-\frac{1}{2}}\A \D^{-\frac{1}{2}}) \right] \X^{(l)} \boldsymbol{\Theta}_l \right)
\end{empheq}
where $\tilde \A = \I - \delta (\I - \D^{-\frac{1}{2}}\A \D^{-\frac{1}{2}})$ is the operator matching the sparsity pattern of the graph and $\delta$ is an hyperparameter. 
When $\delta=1$, Eq.~\ref{eq:mp_LQV} reduces to $\X^{(l+1)} = \sigma(\D^{-1/2} \A \D^{-1/2} \X^{(l)} \boldsymbol{\Theta}_l)$, which is very similar to the update equation used in popular MP layers~\cite{defferrard2016convolutional}.
When $\delta=0$, the node features are not aggregated with those of the neighbors and the MP layer becomes equivalent to a dense layer of an MLP.

\begin{theorem}
The MP layer in \eqref{eq:mp_LQV} minimizes the LQV of the node features $\X$.
\end{theorem}
\begin{proof}
Let
\begin{equation*}
    ||\bm{x}||_{LQV} = \frac{1}{2}\sum_{i=1}^N \sum_{j=1}^N a_{i,j} \left(\frac{x_i}{\sqrt{d_i}} - \frac{x_j}{\sqrt{d_j}}\right)^2
\end{equation*}

To minimize the LQV, we first compute the derivative with respect to the $i$-th component:
\begin{align}
\label{eq:gradient_descent_LQV_alt}
(\partial||\bm{x}||_{LQV})_i &= \frac{1}{\sqrt{d_i}} \sum_j a_{i,j} \left( \frac{x_i}{\sqrt{d_i}} - \frac{x_j}{\sqrt{d_j}} \right) \nonumber\\
                             &= \frac{x_i}{d_i} \sum_j a_{i,j} - \sum_j \frac{a_{i,j}}{\sqrt{d_i}\sqrt{d_j}} x_j \nonumber\\
                             &= (\I \x)_i - (\D^{-1/2} \A \D^{-1/2} \x)_i \nonumber
\end{align}
The whole gradient is given by:
\begin{equation*}
    \nabla(||\bm{x}||_{LQV}) = (\I - \D^{-1/2} \A \D^{-1/2}) \bm{x}
\end{equation*}

The following update minimizes the LQV of the node features with gradient descent:
\begin{align*}
    \bm{x}^{(l+1)} & = \bm{x}^{(l)} - \delta  \nabla(||\bm{x}||_{LQV})^{(l)}\\
                   & = \bm{x}^{(l)} - \delta  (\I - \D^{-1/2} \A \D^{-1/2}) \bm{x}^{(l)}
\end{align*}
where $\delta$ indicates the gradient step.
\end{proof}

Clearly, there is a difference in minimizing the LQV of $\X$ rather than the LQV of $\Sb$ directly, as done by the term $\mathcal{L}_q$ in \eqref{eq:norm_cut}.
Nevertheless, since the function in \eqref{eq:soft_assign} is smooth, if two node features $\x_i$ and $\x_j$ are similar they will be mapped into similar cluster assignments $\s_i$ and $\s_j$.

\subsection{Optimality of the proposed balancing term}
The purpose of the balancing term is to encourage a partition where:
\begin{itemize}
    \item the nodes are assigned with high confidence to only one cluster;
    \item the number of elements in each cluster is approximately the same.
\end{itemize}

Typical degenerate solutions are those that violate one these two requirements. 
In particular, the first degenerate solution is when samples are uniformly assigned to all cluster, i.e., $\s_i = [1/K, 1/K, \dots, 1/K]$, $i=1, \dots, N$. The second degenerate solution occurs when all samples are assigned to the same cluster, e.g., $\s_i = [1, 0, \dots, 0]$, $i=1, \dots, N$.

It is straightforward to see that the first degenerate solution is avoided when $\mathcal{L}$ in \eqref{eq:loss} is minimized. 
The trace is the sum of the elements along the diagonal of $\sqrt{\Sb^T\Sb}$. 
Since $s_{i,j} \leq 1$, the trace is maximized when each cluster assignment $\s_i$ has one entry equal to 1, i.e., when $s_{i,j} \in \{0,1\}$. 
On the other hand, if a sample is assigned to more than one cluster $\s_i$ will have two or more non-zero entries with values less than 1, due to the \texttt{softmax} normalization.
As a consequence, at least one value in the off-diagonal of $\sqrt{\Sb^T\Sb}$ would be non-zero and the trace would be smaller. 

The second degenerate solution is avoided when the clusters assume equal size, which is ensured by the following theorem.

\begin{theorem}
The optimum of the problem
\begin{equation}
    \label{eq:opt_constr}
    \max \limits_{\Sb \mathbf{1} = \mathbf{1}, s_{i,j} \in \{0,1\}} \rm{Tr} \left( \sqrt{\Sb^T\Sb} \right)
\end{equation}
is a balanced partition with clusters of size $\frac{N}{K}$.
\end{theorem}
\begin{proof}
Let $\Sb^T\Sb = \mathbf{C} \in \mathbb{R}^{K \times K}$.
The conditions $\Sb \mathbf{1} = \mathbf{1}$ and $s_{i,j} \in \{0,1\}$ imply the constraint $\sum_{i=1}^K c_{i,i} = N$, being $c_{i,i}$ the volume of samples assigned to cluster $i$.
The solution  of the constrained optimization problem can be found using Lagrange multipliers. Specifically, the Lagrangian is
\begin{equation*}
    L(\mathbf{C}, \lambda) = \sum_{i=1}^K \sqrt{c_{i,i}} - \lambda\left(\sum_{i=1}^K c_{i,i} - N \right).
\end{equation*}
To be an optimum, $\frac{\partial L(\mathbf{C}, \lambda)}{\partial c_{i,i}} = 0$ must hold for for each $i$. 
Hence, by taking the derivative one obtains
\begin{equation}
\label{proof_eq1}
    \frac{1}{2 \sqrt{c_{i,i}}} - \lambda = 0 \rightarrow c_{i,i} = \frac{1}{4 \lambda^2}
\end{equation}
Solving for $\lambda$ can be done by considering the constraint $\sum_{i=1}^K c_{i,i} = N$, which gives
\begin{equation}
\label{proof_eq2}
    \sum_{i=1}^K \frac{1}{4 \lambda^2} = N \rightarrow \lambda = \sqrt{\frac{C}{4N}}.
\end{equation}
Finally, substituting \eqref{proof_eq2} in \eqref{proof_eq1} gives $c_{i,i} = \frac{N}{K}$.
\end{proof}


\section{Experimental evaluation}

\bgroup
\def\arraystretch{1}
\setlength\tabcolsep{1.0em}
\begin{table}[!ht]
\footnotesize
\centering
\label{tab:cit_dataset}
\begin{tabular}{lcccc}
\cmidrule[1.5pt]{1-5}
\textbf{Dataset} & $N$ & $|\mathcal{E}|$ & $F$ & $K$ \\
\cmidrule[.5pt]{1-5}
Cora     & 2,708  & 10,556   & 1,433 & 7 \\
Citeseer & 3,327  & 9,104   & 3,703 & 6 \\
Pubmed   & 19,717 & 88,648  & 500   & 3 \\
DBLP     & 17,716 & 105,734 & 1,639 & 4 \\
\cmidrule[1.5pt]{1-5}
\end{tabular}
\vspace{-.3cm}
\caption{\small Details of the citation datasets: number of nodes ($N$), number of edges ($|\mathcal{E}|$), dimension of the feature vectors ($F$), and number of classes/clusters ($K$).} 
\end{table}
\egroup

The clustering performance are tested on four citation datasets, whose details are reported in Tab.~\ref{tab:cit_dataset}.
In each dataset, the number of clusters $K$ is set equal to the number of classes to which the nodes belong to.
The distribution of the node classes for each dataset, depicted in Fig.~\ref{fig:class_distr}, shows that the classes are rather unbalanced, which makes the clustering problem non-trivial.

\begin{figure}[!ht]
    \centering
    \includegraphics[width=.9\columnwidth]{class_distr.pdf}
    \vspace{-.8cm}
    \caption{\small Distribution of the node classes.}
    \label{fig:class_distr}
\end{figure}

The proposed architecture, referred to as \textit{Just Balance GNN} (JBGNN), consists of a stack of MP layers, which outputs $\bar \X$, followed by an MLP, which outputs $\Sb$.
The MP layers are implemented as in \eqref{eq:mp_LQV} and JBNN is optimized by minimizing only the loss in \eqref{eq:loss}.
The following hyperparameters are used for each dataset.
learning rate: 5e-5;
$\delta$: 0.85;
MP activation: ReLU;
MP channels: 64;
MP layers: 10;
MLP activation: ReLU;
MLP channels: 16;
MLP hidden layers: 1.
The software implementation of JBNN is available online\footnote{\url{https://github.com/FilippoMB/Simplifying-Clustering-with-Graph-Neural-Networks}}.

\bgroup
\def\arraystretch{1}
\setlength\tabcolsep{.3em}
\begin{table*}[!ht]
    \centering
    \footnotesize
    \begin{tabular}{lcccccccc}
        \cmidrule[1.5pt]{1-9}
        & \multicolumn{2}{c}{\textbf{Cora}} & \multicolumn{2}{c}{\textbf{Citeseer}} & \multicolumn{2}{c}{\textbf{Pubmed}} & \multicolumn{2}{c}{\textbf{DBLP}}\\
        & ACC & NMI & ACC & NMI & ACC & NMI & ACC & NMI \\
        \midrule
        SC       & 0.298             & 0.028             & 0.217             & 0.014             & 0.589             & 0.182             & 0.458             & 0.023  \\
        DeepWalk & 0.229 {\t(0.020)} & 0.064 {\t(0.024)} & 0.193 {\t(0.002)} & 0.004 {\t(0.001)} & 0.361 {\t(0.001)} & 0.001 {\t(0.000)} & 0.266 {\t(0.001)} & 0.001 {\t(0.000)} \\
        Node2vec & 0.229 {\t(0.025)} & 0.060 {\t(0.029)} & 0.194 {\t(0.003)} & 0.004 {\t(0.001)} & 0.362 {\t(0.001)} & 0.001 {\t(0.000)} & 0.272 {\t(0.001)} & 0.001 {\t(0.000)} \\
        NetMF    & 0.389             & 0.251             & 0.277             & 0.127             & 0.448             & 0.058             & 0.455             & 0.037  \\
        \midrule
        GAE      & 0.464 {\t(0.062)} & 0.327 {\t(0.051)} & 0.381 {\t(0.038)} & 0.162 {\t(0.029)} & 0.588 {\t(0.071)} & 0.235 {\t(0.044)} & 0.416 {\t(0.035)} & 0.111 {\t(0.028)} \\
        VGAE     & \textbf{0.572} {\t(0.054)} & \textbf{0.437} {\t(0.028)} & 0.360 {\t(0.038)} & 0.156 {\t(0.034)} & \textbf{0.610} {\t(0.060)} & \textbf{0.24}5 {\t(0.043)} & 0.507 {\t(0.047)} & 0.212 {\t(0.021)} \\
        \midrule
        DiffPool & 0.472 {\t(0.010)} & 0.306 {\t(0.005)} & 0.336 {\t(0.007)} & 0.180 {\t(0.008)} & 0.418 {\t(0.002)} & 0.084 {\t(0.001)} & 0.370 {\t(0.042)} & 0.045 {\t(0.043)} \\
        DMoN     & 0.488 {\t(0.063)} & 0.357 {\t(0.042)} & 0.364 {\t(0.043)} & 0.196 {\t(0.029)} & 0.559 {\t(0.042)} & 0.192 {\t(0.048)} & 0.590 {\t(0.039)} & 0.334 {\t(0.026)} \\
        MinCut   & 0.534 {\t(0.041)} & 0.406 {\t(0.029)} & \textbf{0.497} {\t(0.049)} & \textbf{0.295} {\t(0.029)} & 0.572 {\t(0.034)} & 0.208 {\t(0.014)} & 0.538 {\t(0.033)} & 0.297 {\t(0.024)} \\
        JBGNN    & 0.457 {\t(0.025)} & 0.351 {\t(0.128)} & 0.334 {\t(0.019)} & 0.140 {\t(0.024)} & 0.564 {\t(0.023)} & 0.223 {\t(0.013)} & \textbf{0.607} {\t(0.008)} & \textbf{0.359} {\t(0.008)} \\
        \cmidrule[1.5pt]{1-9}
    \end{tabular}
    \vspace{-.3cm}
    \caption{\small Clustering performance metrics for each dataset. For method with stochastic components, the mean and the standard deviation (in brackets) computed over 10 independent runs is reported.}
    \label{tab:results}
\end{table*}
\egroup

The performance of JBGNN are compared against three classes of methods. 
The first are algorithms that generate node embeddings based only on the adjacency matrix. The node embeddings are then clustered with $k$-means.
Representatives of this category are Spectral Clustering (SC), DeepWalk~\cite{perozzi2014deepwalk}, Node2vec~\cite{grover2016node2vec}, and NetMF~\cite{qiu2018network}.
The second class of methods are neural networks that generate node embeddings by accounting both for the adjacency matrix and for the node features. 
Afterwards, the learned embeddings are clustered with $k$-means.
The chosen representatives are the Graph AutoEncoder (GAE) and Variational Graph AutoEncoder (VGAE)~\cite{kipf2016variational}.
Finally, the last class of methods are end-to-end GNN models that directly generate soft cluster assignments $\Sb$ by accounting both for the graph connectivity and the node features. The hard cluster assignments are computed as $\bar \s = \texttt{argmax}(\Sb)$.
DiffPool~\cite{ying2018hierarchical}, DMoN~\cite{tsitsulin2020graph}, MinCutPool~\cite{bianchi2020spectral}, and the proposed JBGNN belong to this class.

The clustering performance is computed in terms of the normalized mutual information (NMI) between true class labels $\bm y$ and cluster assignments $\bar \s$ and the cluster accuracy (ACC), defined as
\begin{equation}
    \label{eq:clust_acc}
    \text{ACC} = \frac{1}{N} \sum \limits_{i=1}^N \theta \left(y_i, h(\bar s_i) \right),
\end{equation}
where $h(\cdot)$ maps the hard-cluster assignment $\bar s_i$ to the best matching class label according to Kuhn-Munkres algorithm, and $\theta$ is the Heaviside step function, which is 1 when $y_i$ and $h(\bar s_i)$ are equal, and zero otherwise.
The results are reported in Tab.~\ref{tab:results}.
Methods with stochastic components are trained and evaluated 10 different times.

\begin{table}[!ht]
    \centering
    \footnotesize
    \begin{tabular}{lcccc}
        \cmidrule[1.5pt]{1-5}
        & \textbf{Cora} & \textbf{Citeseer} & \textbf{Pubmed} & \textbf{DBLP} \\
        \midrule
        Diffpool   & 0.009                         & 0.030                        & 0.234                          & 0.229                         \\
           & {\color[HTML]{CB0000} \t +200\%} & {\color[HTML]{CB0000} \t +42\%} & {\color[HTML]{CB0000} \t +1376\%} & {\color[HTML]{CB0000} \t +477\%} \\
        DMoN       & 0.006                         & 0.028                        & 0.030                          & 0.061                         \\
                   & {\color[HTML]{CB0000} \t +100\%} & {\color[HTML]{CB0000} \t +33\%} & {\color[HTML]{CB0000} \t +76\%}   & {\color[HTML]{CB0000} \t +27\%}  \\
        MinCut & 0.006                         & 0.029                        & 0.030                          & 0.061                         \\
                   & {\color[HTML]{CB0000} \t +100\%} & {\color[HTML]{CB0000} \t +38\%} & {\color[HTML]{CB0000} \t +76\%}   & {\color[HTML]{CB0000} \t +27\%}  \\
        \midrule
        JBGNN      & \textbf{0.003}                         & \textbf{0.021}                        & \textbf{0.017}                          & \textbf{0.048}               \\
        \cmidrule[1.5pt]{1-5}
    \end{tabular}
    \vspace{-.3cm}
    \caption{\small Training times (seconds/step). In red, the increments in training times with respect to JBGNN.}
    \label{tab:times}
\end{table}

The computational complexity of the end-to-end GNN models for clustering is measured in terms of training times.
Tab.~\ref{tab:times} reports the seconds used by each model to perform one step of gradient descent.
To make the comparison fair, the compared architectures are configured to have the same capacity and differ only in how the cluster assignments are computed and in the loss.
Times are measured on an Nvidia RTX A6000.

\section{Discussion and conclusions}

The clustering performance of JBGNN is competitive with respect to the other end-to-end GNN clustering methods: it achieves a comparable ACC and NMI on Cora and Pubmed, and it is the best overall performing method on DBLP.
Diffpool is the worse-performing method in this category, as it does not optimize a proper clustering objective.
On Citeseer, MinCutPool outperforms every other method, including JBGNN, by a large margin. 
As discussed in Section~\ref{sec:lqv}, JBGNN minimizes the LQV of $\bar \X$ while MinCutPool directly minimizes the LQV of $\Sb$, which could make a significant difference in certain datasets.

About the performance of the other methods, SC, DeepWalk, Node2vec, and NetMF always achieve worse ACC and NMI. 
This is expected, as they build embeddings only based on the graph connectivity.
On the other hand, GAE and VGAE also account for the node features and achieve superior performance.
Compared to the end-to-end GNN clustering methods, GAE and VGAE follow a significantly different approach: they optimize a link reconstruction loss to learn node embeddings, which are then clustered in post-processing with $k$-means.
This implies few limitations, as the $k$-means is computationally expensive and must be run each time, even if the neural network is already trained.
On the other hand, JBGNN directly outputs soft cluster assignment, which are efficiently evaluated at inference time and can also be used to implement differentiable operations, such as graph pooling~\cite{grattarola2022understanding}. 

In terms of computational complexity, JBGNN outperforms all other GNN clustering methods thanks to its simplicity and efficient formulation.
In particular, JBGNN achieves a speed improvement ranging from 27\% to  1,376\% compared to the other methods across the four datasets, making it particularly suitable for large-scale applications.

The proposed clustering objective relies on the assumption that all clusters have equal size.
This is reasonable, since clustering is an unsupervised task and the actual sizes are usually unknown.
However, if information about cluster size is available, a clustering objective that puts a bias towards an unbalanced partition could be used instead.


\paragraph{Acknowledgments}
I gratefully acknowledge the support of Nvidia Corporation with the donation of the two RTX A6000 GPUs used in this work.

\bibliographystyle{abbrvnat}
