Nearly Optimal Steiner Trees using Graph Neural Network Assisted Monte
  Carlo Tree Search
Graph neural networks are useful for learning problems, as well as for
combinatorial and graph problems such as the Subgraph Isomorphism Problem and
the Traveling Salesman Problem. We describe an approach for computing Steiner
Trees by combining a graph neural network and Monte Carlo Tree Search. We first
train a graph neural network that takes as input a partial solution and
proposes a new node to be added as output. This neural network is then used in
a Monte Carlo search to compute a Steiner tree. The proposed method
consistently outperforms the standard 2-approximation algorithm on many
different types of graphs and often finds the optimal solution.
Reyan Ahmed, Mithun Ghosh, Kwang-Sung Jun, Stephen Kobourov
# Nearly Optimal Steiner Trees using Graph Neural Network Assisted Monte Carlo Tree Search

###### Abstract

Graph neural networks are useful for learning problems, as well as for combinatorial and graph problems such as the Subgraph Isomorphism Problem and the Traveling Salesman Problem. We describe an approach for computing Steiner Trees by combining a graph neural network and Monte Carlo Tree Search. We first train a graph neural network that takes as input a partial solution and proposes a new node to be added as output. This neural network is then used in a Monte Carlo search to compute a Steiner tree. The proposed method consistently outperforms the standard 2-approximation algorithm on many different types of graphs and often finds the optimal solution.

## 1 Introduction

Graphs arise in many real-world applications that deal with relational information. Classical machine learning models, such as neural networks and recurrent neural networks, do not naturally handle graphs. Graph neural networks (GNN) were introduced by Gori et al. [24] in order to better capture graph structures. A GNN is a recursive neural network where nodes are treated as state vectors and the relationships between the nodes are quantified by the edges. Scarselli et al. [42] extended the notion of unfolding equivalence that leads to the transformation of the approximation property of feed-forward networks (Scarselli and Tsoi [43]) to GNNs.

Many real-world problems are modeled by combinatorial and graph problems that are known to be NP-complete. GNNs offer an alternative to traditional heuristics and approximation algorithms; indeed the initial GNN model [42] was used to approximate solutions to two classical graph problems: subgraph isomorphism and clique detection.

Recent GNN work [37, 48] suggests that combining neural networks and tree search leads to better results than just the neural network alone. Li et al. [37] combine a convolutional neural network with tree search to compute independent sets and other NP-hard problems that are efficiently reducible to the independent set problem. AlphaGo, by Silver et al. [44] combines deep convolutional neural networks and Monte Carlo Tree Search (MCTS) [12, 34] to assess Go board positions and reduce the search space. Xing et al. [48] build on this combination to tackle the traveling salesman problem (TSP).

Since Xing et al. [48] showed that the AlphaGo framework is effective for TSP, a natural question is whether this framework can be applied to other combinatorial problems such as the Steiner tree problem. Although both TSP and the Steiner tree problem are NP-complete, they are different. First, in the Steiner tree problem we are given a subset of the nodes called _terminals_ that must be spanned, whereas in TSP all nodes are equivalent. Second, the output of the Steiner tree problem is a tree, whereas the output of TSP is a path (or a cycle). When iteratively computing a TSP solution, the next node to be added can only be connected to the previous one, rather than having to choose from a set of nodes when growing a Steiner tree. Third, TSP and Go are similar in terms of the length of the instance: both the length of the game and the number of nodes in the TSP solution are fixed and taking an action in Go is equivalent to adding a node to the tour, while the number of nodes in the Steiner tree problem varies depending on the graph instance. Finally, Xing et al. [48] only considered geometric graphs, which is a restricted class of graphs.

### Background:

The Steiner tree problem is one of Karp's 21 NP-complete problems [29]: given an edge-weighted graph \(G=(V,E)\), a set of terminals \(T\subseteq V\) and cost \(k\), determine whether there exists a tree of cost at most \(k\) that spans all terminals. For \(|T|=2\) this is equivalent to the shortest path problem, for \(|T|=|V|\) this is equivalent to the minimum spanning tree problem, while for \(2<|T|<|V|\) the problem is NP-complete [11]. Due to applications in many domains, there is a long history of heuristics, approximation algorithms and exact algorithms for the problem. The classical 2-approximation algorithm for the Steiner tree problem [22] uses the _metric closure_ of \(G\), i.e., the complete edge-weighted graph \(G^{*}\) with terminal node
set \(T\) in which, for every edge \(uv\), the cost of \(uv\) equals the length of a shortest \(u\)-\(v\) path in \(G\). A minimum spanning tree of \(G^{*}\) corresponds to a 2-approximation Steiner tree in \(G\). This algorithm is easy to implement and performs well in practice [2]. The last in a long list of improvements is the LP-based algorithm of Byrka et al. [9], with approximation ratio of \(\ln(4)+\varepsilon<1.39\). The Steiner tree problem is APX-hard [7] and NP-hard to approximate within a factor of 96/95 [10]. Geometric variants of the problem, where terminals correspond to points in the Euclidean or rectilinear plane, admit polynomial-time approximation schemes [4, 38].

### Related Work:

Despite its practical and theoretical importance, the Steiner tree problem is not as well explored with machine learning approaches as other combinatorial and graph problems. In 1985, Hopfield et al. [27] proposed a neural network to compute feasible solutions for different combinatorial problems such as TSP. Bout et al. [14] developed a TSP objective function that works well in practice and Brandt et al. [8] provided different networks for solving TSP. Kohonen's 1982 self-organizing maps [35], an architecture for artificial neural networks, can also be used for such problems as shown by Fort [17, 3] and Favata et al. [16].

Recently, graph neural networks have been an active area of research. Lei et al. [36] introduced recurrent neural operations for graphs with associated kernel spaces. Gilmer et al. [23] study graph neural models as Message Passing Neural Networks. Garg et al. [21] generalized message-passing GNNs that rely on the local graph structure, proposing GNN frameworks that rely on graph-theoretic formalisms. GNNs have been widely used in many areas including physical systems [6, 41], protein-protein interaction networks [18], social science [26, 32], and knowledge graphs [25]; The survey of Zhou et al. [50] covers GNN methods and applications in general, and the survey of Vesselinova et al. [45] provides more details on attempts to solve combinatorial and graph problems with neural networks.

### Problem Statement:

In the standard optimization version of the Steiner Tree Problem we are given a weighted graph \(G=(V,E)\) and a set of terminals \(T\subseteq V\), and the objective is to compute a minimum cost tree that spans \(T\). A Steiner tree \(H\) must contain all the terminals and non-terminal nodes in \(H\) are the Steiner nodes. Several approximation algorithms have been proposed for this problem including a classical 2-approximation algorithm that first computes the metric closure of \(G\) on \(T\) and then returns the minimum spanning tree [1]. In this paper we consider whether graph neural networks can be used to compute spanning trees with close-to-optimal costs using a variety of different graph classes.

### Summary of Contributions:

We describe an approach for computing Steiner Trees by combining a graph neural network and Monte Carlo Tree Search (MCTS). We first train a graph neural network that takes as input a partial solution and proposes a new node to be added as output. This neural network is then used in a MCTS to compute a Steiner tree. The proposed method consistently outperforms the standard 2-approximation algorithm on many different types of graphs and often finds the optimal solution. We illustrate our approach in Figure 1. Our approach builds on the work of Xing et al. [48] for TSP. Since TSP is non-trivially different from the Steiner tree problem, we needed to address challenges in both training the graph neural network and testing the MCTS. We summarize our contribution below:

* To train the neural network we generate exact solutions of Steiner tree instances. From each instance, we generate several data points. The purpose of the neural network is to predict the next Steiner node, given a partial solution. Any permutation of the set of Steiner nodes can lead to a valid sequence of predictions. Hence, we use random permutations to generate data points for the network.
* After we determine the Steiner nodes for a given instance, it is not straightforward to compute the Steiner tree. For TSP, any permutation of all nodes is a feasible tour. For the Steiner tree problem, an arbitrary permutation can have many unnecessary nodes and thus a larger weight compared to the optimal solution. Selecting a subset of nodes is not enough either, since the output needs to be connected and span the terminals. We propose heuristics to compute the tree from the nodes that provide valid result with good quality.
* We evaluate our results on many different classes of graphs, including geometric graphs, Erdos-Renyi graphs, Barabasi-Albert graphs, Watts-Strogatz graphs, and known hard instances from the Stein-Lib database [33]. Our method is fully functional and available on Github.

## 2 Our approach

Let \(G(V,E)\) be a graph, where \(V\) is the set of nodes and \(E\) is the set of edges. Let \(w(u,v)\) be the weight of edge \((u,v)\in E\) and for unweighted graphs \(w(u,v)=1\) for any edge \((u,v)\in E\). Let \(T\subseteq V\) be the set of terminals.
We use \(S=\{v_{1},v_{2},\cdots,v_{i}\}\) to represent the set of nodes that are already added in a partially computed Steiner tree. Then, \(\overline{S}=V-S\) is the set of candidate nodes to be added to \(S\).

Given the graph \(G\) our goal is to derive a Steiner tree by adding node \(v\in\overline{S}\) to \(S\) in turn. A natural approach is to train a neural network to predict which node to add to the partial Steiner tree at a particular step. That is, neural network \(f(G|S)\) takes graph \(G\) and partial solution \(S\) as input, and return probabilities for the remaining nodes, indicating the likelihood they belong to the Steiner tree. We use the GNN of [26] to represent \(f(G|S)\).

Intuitively, we can directly use the probability values, selecting all nodes with probability higher than a given threshold. We can then construct a tree from the selected nodes in different ways. For example, we can compute the induced graph of the selected nodes (add an edge if it connects to selected nodes) and extract a minimum spanning tree [11]. Note that the induced graph may be disconnected and therefore the spanning tree will be also disconnected. Even if the spanning tree is connected, it may not span all the terminals, hence it may not provide a valid solution. These issues can be addressed by reducing the given threshold until we obtain a valid solution.

However, deriving trees in this fashion might not be reliable, as a learning-based algorithm has only one chance to compute the optimal solution, and it never goes back to reverse the decision. To overcome this drawback, we leverage the MCTS. We use a variant of PUCT [40] to balance exploration (i.e., visiting a state as suggested by the prior policy) and exploitation (i.e., visiting a state that has the best value). Using the concept of prior probability, the search space of the tree could be reduced substantially, enabling the search to allocate more computing resources to the states having higher values. We could get a more reliable policy after a large number of simulations as the output of the MCTS acts as the feedback information by fusing the prior probability with the scouting exploration. The overall approach is illustrated in Figure 1.

### Graph neural network architecture:

To get a useful neural network, information about the structures of the concerned graph, terminal nodes, and contextual information, i.e., the set of added nodes \(S=\{v_{1},\ldots,v_{i}\}\) in the partial solution, is required. We tag node \(u\) with \(x_{u}^{t}=1\) if it is a terminal, otherwise \(x_{u}^{t}=0\). We also tag node \(v\) with \(x_{v}^{a}=1\) if it is already added, otherwise \(x_{v}^{a}=0\). Intuitively, \(f(G|S)\) should summarize the state of such a "tagged" graph and generate the prior probability for each node to get included in \(S\).

Some combinatorial problems like the independent set problem and minimum vertex cover problem do not consider edge weights. However, edge weight is an important feature of the Steiner tree problem as the objective is computed based on the weights. Hence, we use the static edge graph neural network (SE-GNN) [48], to efficiently extract node and edge features of the

Figure 1: GNN assisted MCTS: first, train a GNN to evaluate non-terminal nodes, then use the network and heuristics to compute a Steiner tree with MCTS.
Steiner tree problem.

A GNN model consists of a stack of \(L\) neural network layers, where each layer aggregates local neighborhood information, i.e., features of neighbors of each node, and then passes this aggregated information to the next layer. We use \(H_{u}^{l}\in\mathbb{R}^{d}\) to denote the real-valued feature vector associated with node \(u\) at layer \(l\). Specifically, the basic GNN model [26] can be implemented as follows. In layer \(l=1,2,\cdots,L\), a new feature is computed as given by 2.1.

\[H_{u}^{l+1}=\sigma\Big{(}\theta_{1}^{l}H_{u}^{l}+\sum_{v\in N(u)}\theta_{2}^{l} H_{v}^{l}\Big{)} \tag{1}\]

In 2.1, \(N(u)\) is the set of neighbors of node \(u\), \(\theta_{1}^{l}\) and \(\theta_{2}^{l}\) are the parameter matrices for the layer \(l\), and \(\sigma(\cdot)\) denotes a component-wise non-linear function such as a sigmoid or a ReLU function. For \(l=0\), \(H_{u}^{0}\) denotes the feature initialization at the input layer.

The edge information is not taken into account in 2.1. To incorporate edge features, we adapt the approach in [30, 47] to the Steiner tree problem. We integrate the edge features with node features using 2.2.

\[\mu_{u}^{l+1}=\sigma\Big{(}\theta_{1}x_{u}+\theta_{2}\sum_{v\in N(u)}\mu_{v}^{ l}+\theta_{3}\sum_{v\in N(u)}\sigma(\theta_{4}w(u,v))\Big{)} \tag{2}\]

In 2.2, \(\theta_{1}\in\mathbb{R}^{l}\), \(\theta_{2},\theta_{3}\in\mathbb{R}^{l\times l}\) and \(\theta_{4}\in\mathbb{R}^{l}\) are all model parameters. We can see in 2.1 and 2.2 that the nonlinear mapping of the aggregated information is a single-layer perceptron, which is not enough to map distinct multisets into unique embeddings. Hence, as suggested in [48, 49], we replace the single perceptron with a multi-layer perceptron. Finally, we compute a new node feature \(H\) using 2.3.

\[H_{u}^{l+1}=\text{MLP}^{l}\Big{(}\theta_{1}^{l}H_{u}^{l}+\sum_{v\in N(u)}\theta _{2}^{l}H_{v}^{l}+\sum_{v\in N(u)}\theta_{3}^{l}e_{u,v}\Big{)} \tag{3}\]

In 2.3, \(e_{u,v}\) is the edge feature, \(\theta_{1}^{l}\), \(\theta_{2}^{l}\), and \(\theta_{3}^{l}\) are parameter matrices, and \(\text{MLP}^{l}\) is the multi-layer perceptron for layer \(l\). Note that SE-GNN differs from GEN [13] in the following aspects: (1) SE-GNN replaces \(x_{u}\) in 2.2 with \(H_{u}\) so that the SE-GNN can integrate the latest feature of the node itself directly. (2) Each update process in the GEN can be treated as one update layer of the SE-GNN, i.e., each calculation is equivalent to going one layer forward, thus calculating \(L\) times for \(L\) layers. Parameters of each layer of SE-GNN are independent, while parameters in GEN are shared between different update processes which limits the neural network. (3) We replace \(\sigma\) in 2.2 with MLP as suggested by [48, 49] to map distinct multisets to unique embeddings.

We initialize the node feature \(H^{0}\) as follows. Each node has a feature tag which is a 4-dimensional vector. The first element of the vector is binary and it is equal to 1 if the partial solution \(S\) contains the node. The second element of the vector is also binary and it is equal to 1 if the node is a terminal. The third and fourth elements of the feature tag are the \(x\) and \(y\) coordinates of the node. The last two are used only for geometric graphs.

### Parameterizing \(f(G|S;\theta)\):

Once the feature for every node is computed after updating \(L\) layers, we use the new feature for the nodes to define the \(f(G|S;\theta)\) function, which returns the prior probability for each node indicating how likely the node will belong to partial solution \(S\). Specifically, we fuse all node feature \(H_{u}^{L}\) as the current state representation of the graph and parameterize \(f(G|S;\theta)\) as expressed by 2.4.

\[f(G|S;\theta)=\text{softmax}(sum(H_{1}^{L}),\cdots,sum(H_{n}^{L})) \tag{4}\]

During training, we minimize the cross-entropy loss for each training sample \((G_{i},S_{i})\) in a supervised manner as given by 2.5.

\[\ell(S_{i},f(G_{i}|S_{i};\theta))=-\sum_{j=1}^{N}y_{j}\log f(G_{i}|S_{i}(1:j-1 );\theta) \tag{5}\]

In 2.5, \(S_{i}\) is an ordered set of nodes of a partial solution which is a permutation of the nodes of graph \(G_{i}\), with \(S_{i}(1:j-1)\) the ordered subset containing the first \(j-1\) elements of \(S_{k}\), and \(y_{j}\) a vector of length \(N\) with 1 in the \(S_{i}(j)\)-th position. We provide more details in Section 3.

### GNN assisted MCTS:

Similar to the implementation in [48], the GNN-MCTS uses graph neural networks as a guide of MCTS. We denote the child edges of \(u\) in MCTS by \(A(u)\). Each node \(u\) in the search tree contains edges \((u,a)\) for all legal actions \(a\in A(u)\). Each edge of MCTS stores a set of statistics:

\[\{N(u,a),Q(u,a),P(u,a)\},\]

where node \(u\) denotes the current state of the graph including the set of nodes \(S\) and other graph information, action \(a\) denotes the selection of node \(v\) from \(\overline{S}\) to add in \(S\), \(N(u,a)\) is the visit count, \(Q(u,a)\) is the action value and \(P(u,a)\) is the prior probability of selecting edge \((u,a)\).

In the Steiner tree problem, we are interested in finding a tree with minimum cost. Hence, we track the
best action value found under the subtree of each node to determine the "exploitation value" of the tree node, as suggested in [19] in the context of the stock trading problem.

The standard MCTS takes solution values in the range \([0,1]\)[34]. However, the Steiner tree can have an arbitrary solution value that does not fall in a predefined interval. This issue could be addressed by adjusting the parameters of the tree search algorithm in such a way that it is feasible for a specified interval. Adjusting parameters requires substantial trial and error due to the change in the number of nodes. Instead, we address this issue by normalizing the action value of node \(n\), whose parent is node \(p\), in the range of \([0,1]\) using 2.6.

\[Q_{n}=\frac{\tilde{Q}_{n}-w_{p}}{b_{p}-w_{p}} \tag{2.6}\]

In 2.6, \(b_{p}\) and \(w_{p}\) are the minimum and maximum action values among the children of \(p\), and \(Q_{n}\) is the action value of \(n\). The actions under \(p\) are normalized in the range of \([0,1]\) so that the best action is 0 and the worst action is 1.

The GNN-MCTS proceeds by iterating over the four phases below and then selects a move to play.

1. **Selection Strategy.** The first in-tree phase of each simulation starts at the root node \(v_{0}\) of the search tree and finishes when the simulation reaches a leaf node \(v_{l}\) at time step \(l\). At time step \(t<l\), we use a variant of PUCT [40] to balance exploration (i.e., visiting the states suggested by the prior policy) and exploitation (i.e., visiting states which have best values) according to the statistics in the search tree as given by 2.7 and 2.8 respectively. (2.7) \[a_{t}=\text{argmax}_{a}(Q(v_{t},a)+U(v_{t},a))\] (2.8) \[U(v,a)=c_{puct}P(v,a)\frac{\sqrt{\sum_{b}N(v,b)}}{1+N(v,a)}\] where \(c_{puct}\) is a constant for trading off between exploration and exploitation. We set \(c_{puct}=1.3\) according to previous experimental results [48].
2. **Expansion Strategy.** When a leaf node \(v\) is reached, the corresponding state \(s_{v}\) is evaluated by the GNN to obtain the prior probability \(p\) of its children nodes. The leaf node is expanded and the statistic of each edge \((s_{v},a)\) is initialized to \(\{N(s_{v},a)=0,Q(s_{v},a)=-\infty,P(s_{v},a)=p_{a}\}\).
3. **Back-Propagation Strategy.** For each step \(t<l\), the edge statistics are updated in a backward process. The visit counts are increased as \(N(v_{t},a_{t})=N(v_{t},a_{t})+1\), and the action value is updated to the best value.
4. **Play.** After repeating steps 1-3 several times (800 times for smaller datasets and 1200 times for larger datasets according to the previous experimental results [48]), we select the node with the biggest \(\hat{P}(a|u_{0})=\frac{Q(u_{0},a)}{\sum_{b}Q(u_{0},b)}\) as the next move \(a\) in the root position \(u_{0}\). The selected child becomes the new root node and the statistics stored in the subtree are preserved.

### Computing Steiner tree from \(S\):

There are several ways to compute a Steiner tree from the set of nodes \(S\). We provide two effective heuristics that we use in our experiments.

1. **MST-based heuristic.** In this heuristic, we first add the terminal nodes to the solution if they are not already present, and then compute the induced graph. We iteratively add nodes from \(\overline{S}\) in order computed by the MCTS until the induced graph is connected. In the last step, we compute a minimum spanning tree (MST) of the induced graph and prune degree-1 non-terminal nodes. This heuristic is effective for geometric graphs and unweighted graphs.
2. **Metric closure-based heuristic.** In this heuristic, given an input graph \(G=(V,E)\) and a set of terminals \(T\), we first compute a metric closure graph \(G^{\prime}=(T,E^{\prime})\). Every pair of nodes in \(G^{\prime}\) is connected by an edge with weight equal to the shortest path distance between them. The minimum spanning tree of the metric closure provides a 2-approximation to the optimal Steiner tree. For

Figure 2: Example graph for the Steiner tree heuristic. Considering \(D\) as a terminal node and computing the MST on the metric closure provides a better solution than the 2-approximation.
example, in Figure 2, \(A\), \(B\) and \(C\) are terminal nodes and \(D\) is not. Note that \(D\) does not appear in any shortest path as every shortest path between pairs of terminals is 5 and none of them goes through \(D\). Without loss of generality, the 2-approximation algorithm chooses the \(A-C-B\) path with total cost of 10, while the optimal solution that uses \(D\) has cost 9.

While the 2-approximation algorithm does not consider any node that does not belong to a shortest path between two terminal nodes, here we consider such nodes. Specifically, we iteratively add nodes from \(\overline{S}\) in order computed by the MCTS, even if they don't belong to any shortest path. Note that, unlike the MST-based heuristic, the metric closure-based heuristic computes the MST on the metric closure (not on the input graph).

Both of the heuristics start by selecting all the terminals as the partial solution. In the MCTS, we gradually add nodes that are not in the set of already selected nodes. For the MST-based heuristic, we stop selecting nodes when the induced graph becomes connected. For the metric closure-based heuristic we stop selecting nodes when 10% non-terminal nodes have been selected.

## 3 Model setup and training

In order to train the models, one has to provide training data consisting of input graphs \(G=(V,E)\), edge weights \(W:E\rightarrow\mathbb{R}^{+}\), and terminals \(T\subseteq V\). Given \(G,W,T\), and partial solution \(S\), our goal is to give label 1 to the next node to be added and 0 to all others. Initially, we set \(S=T\) as all terminals must be in the Steiner tree. Consider a graph with 6 nodes \(u_{1},u_{2},\cdots,u_{6}\), \(T=\{u_{1},u_{2},u_{3}\}\), and an optimal Steiner tree contains the first five nodes \(u_{1},u_{2},\cdots,u_{5}\). For this example, initially we set \(S=T=\{u_{1},u_{2},u_{3}\}\). Since we have two Steiner nodes \(u_{4}\) and \(u_{5}\), both permutations \(u_{4},u_{5}\) and \(u_{5},u_{4}\) are valid. For the first permutation, after setting \(S=\{u_{1},u_{2},u_{3}\}\), the next node to be added in the solution is \(u_{4}\). Hence for this data point, only the label for \(u_{4}\) is 1. This permutation provides another data point where \(S=\{u_{1},u_{2},u_{3},u_{4}\}\) and only the label for \(u_{5}\) is equal to 1. Similarly, we can generate two more data points from the other permutation. This exhaustive consideration of all possible permutations does not scale to larger graphs, so we randomly select 100 permutations from each optimal solution. The model is trained with Stochastic Gradient Descent, using the ADAM optimizer [31] to minimize the cross-entropy loss between the models' prediction and the ground-truth (a vector in \(\{0,1\}^{|V|}\) indicating whether a node is the next solution node or not) for each training sample.

### Data generation:

We produce training instances using several different random graph generation models: Erdos-Renyi [15], Watts-Strogatz [46], Barabasi-Albert [5], and random geometric [39] graphs. Each of these generators needs some parameters; below we describe the values we used, aiming to have graphs of comparable density across the different generators. For Erdos-Renyi model, there is an edge selection probability \(p\), which we set to \(\frac{2\ln n}{n}\) to ensure that the generated graphs are connected with high probability. In the Watts-Strogatz model, we initially create a ring lattice of constant degree \(K\) and rewire each edge with probability \(0\leq p\leq 1\), while avoiding self-loops and duplicate edges. For our experiments we use \(K=6\) and \(p=0.2\). In the Barabasi-Albert model, the graph begins with an initially connected graph of \(m_{0}\) nodes. New nodes are added to the network one at a time. Each new node is connected to \(m\leq m_{0}\) existing nodes with a probability that is proportional to the number of edges that the existing nodes already have. We set \(m_{0}=5\). In the random geometric graph model, we uniformly select \(n\) points from the Euclidean cube, and connect nodes whose Euclidean distance is not larger than a threshold \(r_{c}\), which we choose to be \(\sqrt{\frac{2\ln n}{\pi n}}\) to ensure the graph is connected with high probability.

The Steiner tree problem is NP-complete even if the input graph is unweighted [20]. We generate both unweighted and weighted Steiner tree instances using the random generators described above. The number of nodes in these instances is equal to 20 and the number of terminals is equal to 10. For each type of instance we generate 200 instances. For weighted graphs, we assign random integer weights in the range \(\{1,2,\cdots,10\}\) to each edge. Since the weighted version of the Steiner tree problem is the more general version, and the number of terminals is an important parameter, we create a second dataset of graphs with 50 nodes. For the number of terminals, we use two distributions. In the first distribution, the percentage of the number of terminals with respect to the total number of nodes is in \(\{20\%,40\%,60\%,80\%\}\). In the second distribution the percentage is in \(\{3\%,6\%,\cdots,18\%\}\). These two cases are considered to determine the behavior of the learning models on large and small terminal sets (compared with the overall graph size). As random graphs instances can be "easy" to solve, we also evaluate our approach on graphs from the SteinLib library [33], which provides hard graph instances. Specifically, we perform experiments on two SteinLib datasets: I080 and I160. Each instance of the I080 and I160 datasets contains 80 nodes and 160 nodes respectively. Both datasets have 100 in
stances.

### Computing optimal solutions:

In order to evaluate the performance of our approach (and that of the 2-approximation) we need to compute the optimal solutions. There are different integer linear programs (ILP) for the exact Steiner tree problem. The cut-based approach considers all possible combinations of partitions of terminals and ensures that there is an edge between that partition. This ILP is simple but introduces an exponential number of constraints. A better ILP approach in practice considers an arbitrary terminal as a root and sends flow to the rest of the terminals; see [2, 28] for details about these and other ILP methods for the exact Steiner tree problem.

We generate 2,000 Steiner tree instances and compute the exact solution with the flow-based ILP. We use CPLEX 12.6.2 as the ILP solver on a high-performance computer (Lenovo NeXtScale nx360 M5 system with 400 nodes with 192 GB of memory each). We use Python to implementing the algorithms described above.

### Model architectures:

For the MLP in our GNN model, we have used two hidden layers. The first hidden layer has an embedding dimension equal to 128. The second hidden layer has a convolution dimension equal to 128. We use the ReLU activation function for both layers. We also use batch normalization in both layers to normalize the contributions to a layer for every batch of the datasets. The value of early stopping is equal to 15; hence the model will automatically stop training when the chosen metric does not improve for 15 epochs. We trained the network and evaluated our

Figure 3: Performance on simple graphs. Each data point represents one graph. The lower the cost the better the algorithm is. Our algorithm (MCTS) is nearly optimal and performs better than 2-approximation.
algorithm separately for each combination of generator and node size. Recall that the neural network predicts the next Steiner node from a partial solution. Hence, for each Steiner tree instance, we generate a set of data points. Since neural network architecture can not handle different node sizes, we have trained four independent neural networks for node sizes 20, 50, 80, and 160. The same neural network can predict solution nodes for different graph generation models if the node size is the same. In total, we have trained the networks on around 200,000 data points.

### Heuristic setup:

We used the two heuristics described in Section 2.4. Recall that the MST-based heuristic just computes the minimum spanning tree on the induced graph of the partial solution. It works well for geometric graphs, unweighted Erdos-Renyi, unweighted Watts-Strogatz, and unweighted Barabasi-Albert graphs. We use the metric closure-based heuristic for all the other experiments.

## 4 Experimental results

We evaluate the performance of the proposed approach by comparing the computed trees to those computed by the classical 2-approximation algorithm and the optimal solutions. The proposed approach never performs worse than the 2-approximation algorithm. We also report running times.

The results for geometric graphs and other unweighted graphs are shown in Figure 3. The X-axis represents the graph or instance number that does not have any significance. Traditionally bar plot is used in such a scenario. However, for each instance, we show three costs for three different algorithms. Hence scatter

Figure 4: Performance on weighted graphs. Each data point represents one graph. The lower the cost the better the algorithm is. Our algorithm (MCTS) is nearly optimal and performs better than 2-approximation.
plot provided a better visualization by saving space horizontally. One can show two costs instead of three costs by showing the difference w.r.t. the optimal algorithm. However, this approach does not provide a better visualization since many differences get closer to zero. We illustrate the performance of different algorithms on the Geometric graphs in Figure 2(a). We represent the optimal solution with green triangles, our algorithm with yellow squares, and the 2-approximation with blue circles. For the geometric graph, we have 40 instances, each of which has 20 nodes and 10 terminals. A majority of the time the 2-approximation has a larger solution value and our algorithm has a solution very close to the optimal value. The 2-approximation performs worse than our algorithm in 36 instances out of 40 instances.

Our algorithm also performs well for unweighted graphs. We illustrate the performance for random graphs generated by Erdos-Renyi, Barabasi-Albert, and Watts-Strogatz models in Figure 2(b), Figure 2(c), and Figure 2(d) respectively. We have 40 instances for each type of generator. Again, each instance has 20 nodes and 10 terminals. In all of these instances, our algorithm achieves the optimal solution. For Erdos-Renyi graphs, our algorithm performs better than the 2-approximation in two instances. For Barabasi-Albert graphs, our algorithm performs better in six instances. For Watts-Strogatz graphs, our algorithm performs better in four instances.

Results for the weighted graphs are shown in Figure 4. The weighted version of the Steiner tree problem is harder than the unweighted version. Hence, we consider a larger set of instances. For each random graph generation model, we consider one dataset that has 20 nodes for each instance and another dataset that has 50 nodes. We illustrate the performance on 20 nodes Erdos-Renyi graphs in Figure 3(a). For this dataset, both of the algorithms provide solution values similar to the optimal value. We illustrate the performance on 20 nodes Barabasi-Albert and Watts-Strogatz graphs in Figure 3(b) and Figure 3(c) respectively. For Barabasi-Albert graphs, the 2-approximation performs worse than our algorithm in 24 instances out of 40 instances. Our algorithm provides an optimal solution in 39 instances. For Watts-Strogatz graphs, the 2-approximation performs worse than our algorithm in 24 instances out of 2-approximation performs worse than our algorithm in 39 instances. For Watts-Strogatz graphs, the 2-approximation performs worse than our algorithm in 24 instances out of 2-approximation performs worse than our algorithm in 39 instances. Our algorithm provides an optimal solution in 39 instances. For Watts-Strogatz graphs, the 2-approximation performs worse than our algorithm in 39 instances.

We illustrate the performance of the algorithms on 50 nodes Watts-Strogatz graphs in Figure 4(a). Again, our algorithm provides nearly optimal solutions and the 2-approximation has a noticeable difference. The 2-approximation performs worse than our algorithm in 38 instances out of 40 instances. Our algorithm provides an optimal solution in 34 instances. We illustrate the performance of the algorithms on 50 nodes Barabasi-Albert graphs in Figure 4(b). The 2-approximation

Figure 5: Performance on more weighted graphs. Each data point represents one graph. The lower the cost the better the algorithm is. Our algorithm (MCTS) is nearly optimal and performs better than 2-approximation.
performs worse than our algorithm in 34 instances out of 40 instances. Our algorithm provides an optimal solution in 31 instances. Our algorithm provides nearly optimal solutions for the remaining instances.

The SteinLib library [33] provides hard graph instances for solving the Steiner tree problem. Results for a SteinLib dataset is shown in Figure 6. We can see that there is a relatively large difference between the optimal solution value and the 2-approximation solution value. Despite this larger difference of 2-approximation solution values, our algorithm finds nearly optimal solutions.

### Running time:

The training time of the GNN depends on the dataset. The maximum training time is around 20 hours for the I160 SteinLib dataset. The average running times of the optimal algorithm, 2-approximation, and our algorithm for different test datasets are shown in Figure 1. We denote the geometric, unweighted Erdos-Renyi, unweighted Watts-Strogatz, and unweighted Barabasi-Albert graphs by GE, ER, WS, and BA respectively. We denote the weighted 20 nodes Erdos-Renyi, Watts-Strogatz, and Barabasi-Albert graphs by ER20, WS20, and BA20 respectively. We denote the weighted 50 nodes Erdos-Renyi, Watts-Strogatz, and Barabasi-Albert graphs by ER50, WS50, and BA50 respectively. We denote the 80 nodes and 160 nodes SteinLib datasets by I080 and I160 respectively. We can see the 2-approximation algorithm is the fastest. Our algorithm is a little slower, however, the solution values are closer to the optimal values.

## 5 Conclusion

We described an approach for the Steiner tree problem based on GNNs and MCTS. An experimental evaluation shows that the proposed method computes nearly optimal solutions on a wide variety of datasets in a reasonable time. The proposed method never performs worse than the standard 2-approximation algorithm. The source code and experimental data can be found on github [https://github.com/abureyanahmed/GNN-MCTS-Steiner](https://github.com/abureyanahmed/GNN-MCTS-Steiner).

One limitation of our work is we need to retrain for different node sizes. Also, the Steiner tree problem can be seen as a network sparsification technique. In fact, it is one of the simplest sparsification methods since it only considers trees. It would be interesting to see whether our proposed approach can be adapted to graph spanner problems. Our model is unable to fit different node sizes. Hence in our experiments, we try a

\begin{table}
\begin{tabular}{|l|c|c|c|c|c|c|c|c|c|c|c|} \hline Graphs/ & GE & ER & WS & BA & ER20 & WS20 & BA20 & ER50 & WS50 & BA50 & I080 & I060 \\ Algorithms & & & & & & & & & & & \\ \hline
2-apprx & 0.16 & 0.09 & 0.10 & 0.11 & 0.16 & 0.40 & 0.14 & 1.14 & 0.79 & 0.47 & 1.29 & 7.88 \\ \hline MCTS & 0.64 & 0.40 & 0.49 & 0.49 & 0.75 & 1.73 & 0.66 & 5.06 & 3.20 & 2.17 & 5.77 & 34.52 \\ \hline OPT & 5.92 & 6.33 & 5.00 & 4.68 & 22.99 & 28.61 & 29.90 & 153.71 & 125.41 & 134.46 & 1051.51 & 6188.18 \\ \hline \end{tabular}
\end{table}
Table 1: Average running time of different algorithms in seconds.

Figure 6: Performance on SteinLib datasets. Each data point represents one graph. The lower the cost the better the algorithm is. Our algorithm (MCTS) is nearly optimal and performs better than 2-approximation.
small set of node sizes. It is an interesting future work to generalize the model to handle different node sizes. A general model will provide an opportunity to explore the effectiveness of different parameters of the model.
