Engineering Learned Heuristics to Improve Clustering for Multilevel Graph Partitioning

In the following, we assume a machine learning model $ℳ$ which predicts values between 0 and 1, with 0 meaning that an edge can be freely contracted and 1 meaning that an edge should definitively not be contracted. Let $p_{ℳ}:E\to [0,1]$ be the function that maps an edge to its prediction. We extend it to any set of edges $E^{\prime }$ by defining $p_{ℳ}(E^{\prime })≔{\sum }_{e\in E^{\prime }}{p}_{ℳ}(e)$. Our basic idea is to use a threshold $t$ for each clustering round and only contract an edge $e$ if $p_{ℳ}(e)\le t$.

Our full algorithm builds upon the size-constrained label propagation used by the default configuration of Mt-KaHyPar [20, 22] (although we only consider plain graphs instead of hypergraphs), which we summarize in the following. The algorithm initially places every node in its own singleton cluster. Then, it iterates in parallel over all nodes, using a random order. For each node $u$, it greedily assigns $u$ to the neighboring cluster which has the strongest connection to $u$ by incident edges and where the total node weight of the cluster is below a predefined size constraint. The clustering terminates if either the number of clusters is smaller than $n$ by a factor of 2.5, or after all nodes have been visited.

For our modified algorithm, consider an unweighted graph $G$ and let $𝒞$ be the current clustering. Our first modification of the algorithm states that $u$ may be assigned to a cluster $C\in 𝒞$ only if $p_{ℳ}(E(\{u\},C))\le t\cdot |N(u)\cap C|$. Thus, we simply consider the average prediction for all edges leading to a cluster when deciding whether the cluster is eligible.

Apart from forbidding some contractions, we also penalize clusters which are close to the threshold but still below it, to make it more likely that another adjacent cluster is chosen. In the standard algorithm, the rating for cluster $C$ is $r(u,C)≔|N(u)\cap C|$. Let ${\delta }_C≔\frac{1}{t}\left(t-\frac{p_{ℳ}(E(u,C))}{|N(u)\cap C|}\right)$ be the relative distance to the threshold for an eligible cluster. Our second modification of the algorithm replaces the standard rating function with the penalized variant $r_{\alpha }(u,C)≔{\delta }_C^{\alpha }\cdot |N(u)\cap C|$. The parameter $\alpha$ allows to interpolate between a linear penalty, quadratic penalty, or other exponents.

Our modified clustering algorithm can lead to problems if a majority of the edges has a high prediction, i.e., should not be contracted. Depending on the choice of $t$, a large part of the nodes might not have any eligible adjacent cluster and, consequently, is not contracted. This means that the coarsened graph is almost as large as the input graph, which might substantially degrade the overall running time of the multilevel algorithm.

To prevent the described problem, we use $\mathrm{\ell }$ separate clustering rounds with increasing thresholds $t_1,\mathrm{\dots },t_{\mathrm{\ell }}$ (in practice, we use equidistant values in the interval $[0.2,0.8]$ and $\mathrm{\ell }=5$). After each round, we check whether the desired reduction in node count was achieved. If yes, we terminate the clustering. If not, we start a new clustering round with a larger threshold, hopefully allowing for additional contractions. Note that we keep all clusters and only try to find new clusters for nodes that could not be clustered in the previous round. Therefore, consecutive rounds are faster since less nodes are considered. Also, for each round after the first we only aim to achieve a factor 1.75 reduction in node count instead of 2.5 – less aggressive clustering is beneficial for solution quality if a majority of edges has a high prediction (see Section 6.2).

Moreover, we introduce a two-hop clustering step after the regular clustering rounds. Two-hop clustering is a well-known technique [21, 23, 33] which clusters nodes that have no common neighbor but are adjacent to the same cluster, thereby allowing to contract nodes even if their neighbor clusters are overweight or not eligible due to edge predictions.

Our machine learning model is only applicable to the input graph, because we train on unweighted graphs while the coarse graphs have node and edge weights. To mitigate this, we instead use the accumulated predictions from the previous level. Let $C_1$ and $C_2$ be two adjacent clusters with corresponding nodes $c_1$ and $c_2$ in the coarse graph. Then we use $p_{ℳ}(\{c_1,c_2\})≔p_{ℳ}(E(C_1,C_2))$ for the prediction of the coarse edge. Our modified clustering algorithm still works as described above if we also account for edge weights, i.e., we use the summed weight of the edges incident to a cluster instead of the number of neighbors. Thereby, we can still use the original predictions. Also note that the coarser levels are less important for solution quality, which makes a potential loss of accuracy more tolerable.

Integrating a neural network model to guide coarsening decisions adds significant complexity compared to locally driven coarsening heuristics. At inference time, computing a set of edge predictions from the trained model involves multiple dense matrix-matrix multiplications. For large graphs with few features, these matrix multiplications consist of a tall-and-skinny matrix with many rows ($n$) and few columns (features), along with a much smaller matrix of trainable weights that depends on the number of hidden channels in the model.

While highly optimized libraries for matrix multiplication are widely available, they are typically optimized for large square-like matrices, and performance can degrade significantly as the matrix sizes approach the inner kernel size used to perform the computation. These kernels compute a small segment of the resulting matrix, which is approximately two by four registers in size, or 16 by four elements when using half precision and 256-bit registers [19]. Kernel code is often written for specific CPU architectures, resulting in variations in optimal shapes across different CPUs.

The performance drop for smaller matrices is particularly noticeable in popular deep learning frameworks like PyTorch, as highlighted by a similar study for graph coloring [31]. To mitigate some of the added computational cost, we implement a custom matrix kernel tailored to the dimensions of our model. Like other high-performance matrix multiplication libraries, we also leverage the single instruction, multiple data (SIMD) extensions of the x86 instruction set. Specifically, we use the AVX2 extension and 256-bit registers. Aside from the tailored shape of our kernel, the overall implementation is otherwise very similar to other publicly available kernels, such as those described by Goto and Geijn [19].

Using this custom kernel, we implement the computation of edge predictions in two phases. The first phase computes the global graph features and the specific features of every node (see Section 5.2), storing the results in an array with one feature vector per node. Since the computations are independent, we simply do this in parallel over the nodes. The second phase computes the actual edge predictions by reading the node features of both endpoints from the feature array, writing the model input features into the input matrix and then running our custom matrix multiplication kernel. We parallelize the second phase in batches of 256 edges per parallel task, which provides a reasonable trade-off between fine-grained parallelism and large enough input matrices ($256\times 32$).

Unfortunately, there is no clear way to measure how decisions in the coarsening stage affect the solution quality of the final partitioning result.¹¹1Multilevel partitioning algorithms are highly non-linear combinatorial algorithms that have multiple stages with complex dependencies. Our approach instead aims to predict whether an edge is likely to be cut by high-quality solutions. In practice, we generate training data by repeatedly running an expensive partitioning algorithm on small training instances (using $k\in \{4,8,12,16,20,24,28,32,48,64,96,128\}$ and 10 repetitions per value of $k$). On the resulting set of 120 high-quality solutions, we simply compute the fraction of solutions that cut a given edge $e$ and train $ℳ$ to predict this fraction.

Our algorithm for generating high-quality solutions is a variant of the memetic partitioning algorithm proposed by Andre et al. [3]. The idea is to use a memetic approach as a meta-heuristic on top of a multilevel algorithm, generating an initial population with ordinary partitioning runs and then using multilevel mutation and recombination operations to find new solutions. In our approach, we replace KaHyPar with Mt-KaHyPar as the underlying multilevel algorithm. Overall, this allowed us to generate high-quality training data with a time requirement between multiple minutes and a few hours per graph in the training set.

To facilitate efficient training of $ℳ$, we limit our dataset to a subset of these labeled edges rather than processing the full edge sets of all graphs. We constructed a fixed dataset of one million edges to serve as the basis for training by randomly sampling an equal number of edges from each graph in the training set. This approach guarantees that every graph contributes equally to the dataset, preventing larger graphs from dominating the training process. The collected data was subsequently split into training (70%), validation (15%), and testing (15%) sets. This strategy challenges the model $ℳ$ to generalize in two distinct ways: it must correctly classify unseen edges within the training graphs (evaluated via the test set) and it must adapt to entirely new graph instances (evaluated by the model’s performance within the Mt-KaHyPar partitioner).

For the training objective, we utilized the mean squared error (MSE) loss function. Optimization was performed using the Adam optimizer [29] with an initial learning rate of 0.001 and a weight decay of $1\times {10}^{-5}$. To enhance model performance, we implemented a learning rate scheduler which reduces the learning rate by a factor of 10 if the validation does not improve significantly in 10 epochs. We also employed early stopping with a patience of 10 and neuron dropout with a probability of 50% to mitigate overfitting.

As shown in Table 1, we consider features that mostly consist of simple metrics which can be computed without much time overhead (e.g., statistics relating to the degree sequence of the graph). For predicting a given edge $\{u,v\}$, we include global features that describe the graph as a whole, node features for both $u$ and $v$ that relate to the node and the structure of its neighborhood, and edge features that relate to both $u$ and $v$ at once.

In addition, we exploit a preprocessing step of Mt-KaHyPar to get access to additional features. Mt-KaHyPar pre-computes a modularity clustering [20, 22] with a parallel version of the Louvain algorithm [6]. Then, contractions during the coarsening phase are only allowed within the same community (this was originally developed for KaHyPar [26], see Section 3).

We exploit this by extracting multiple features that characterize the computed Louvain clustering, describing the structure of the communities and the achieved modularity. Since the Louvain algorithm computes a hierarchical clustering [6], we extract features from multiple levels of the clustering: the final level as well as the two levels before the final level. The features are computed separately for each of the three levels (entries in Table 1 are marked accordingly). Note that these features would be comparatively expensive to compute again, but since the community detection is already a part of the regular partitioning algorithm we can access them without additional overhead. Moreover, Table 2 shows features that require more computation time, since computing them is (at least) equivalent to counting the triangles contained in the graphs – which is particularly expensive if the graph contains many nodes with high degree. The idea is that these features are more expensive but might also improve the prediction accuracy. However, they are not included in our final configuration since the improvements turned out to be rather insubstantial (see Section 6.1).

During the training, we observed substantially different model performance depending on the input graph. Some instances yield cut improvements of more than 10% over the baseline, while there are also graphs without any improvement. A property which can partially explain this is the uneven label distribution in the training data (i.e., the frequency of edge cuts in high-quality partitions, see Section 5.1) – the mean over the labels varies substantially on individual graphs. Based on this, we divided Set B into three subsets of approximately equal size: high, medium, and low label values (denoted as Set B${}_{1\text{-}20}$, Set B${}_{21\text{-}45}$, and Set B${}_{46\text{-}69}$). Figure 1 visualizes the resulting distribution of mean label values.

Our guided coarsening algorithm (see Section 4) includes multiple parameters that control how the predicted edge values are used in the clustering. The most important parameters are the number of clustering subrounds, the target shrink factor for the secondary rounds and the penalty parameter $\alpha$. We evaluated different choices for these parameters with the unconstrained refinement strategy and the Ml-Clustering model on Set A. The results are summarized in Table 6.

As shown in the table, the parameters chosen for the final configuration provide optimal solution quality with only a minimal increase in running time when compared to other choices. However, the overall impact is very small (less than 1% in the geometric mean in each case). This indicates that our method is robust with regard to the details of the clustering algorithm, while it also seems unlikely that further parameter tuning can achieve meaningful improvements.