NN-Steiner: A Mixed Neural-algorithmic Approach for the Rectilinear Steiner Minimum Tree Problem

First, we pick integers $a,b\in[0,L)$ uniformly at random and then translate the input pointset by the vector $(a,b)$. We then construct a quadtree, where the splitting of quadtree cells terminates if a cell contains 1 or 0 points. The quadtree is a tree $Q$ where each internal vertex has degree 4. The root has level $0$, and is associated with a cell of side length $L$. Any vertex $v\in Q$ at level $i$ is associated with a cell ${\mathsf{A}}_{v}$ of size (i.e., side length) $\frac{L}{2^{i}}$. Bisecting a level-$i$ cell ${\mathsf{A}}_{v}$ with a horizontal line and with a vertical line decomposes it into four level-$(i+1)$ child cells, each of size $\frac{L}{2^{i+1}}$, corresponding to the four children of $v\in Q$. As the side length of the bounding box is $L=O(n)$, and points have integral coordinates, the height (max-level) of the quadtree is at most $O(\log L)$ and the total number of vertices in $Q$ (and thus the number of cells across all levels) is $O(n\log L)\subseteq O(n\log n)$.

We now consider a special family of Steiner trees for which the crossing of quadtree cells is constrained.

Note that if a side of a cell $S$ is contained in the sides of multiple cells, then the portals on $S$ are spaced according to the cell with side containing $S$ that has the lowest level $i$. We say that the portals on $S$ have level $i$. See Fig. 2 for an example of quadtree decomposition and a $(2,1)$-light rectilinear Steiner tree.

The following theorem guarantees that the minimal $(m,r)$-light rectilinear Steiner tree is an approximate RSMT (Arora 1998). A proof is given in the supplemental (Kahng et al. 2023).

Hence, our goal is to compute a minimum $(m,r)$-light rectilinear Steiner tree. In particular, Arora proposed to use DP in a bottom-up construction. We sketch the idea here.

We process all quadtree cells in a bottom-up manner. For a fixed quadtree cell ${\mathsf{A}}$, consider an $(m,r)$-light Steiner tree $T$ restricted to ${\mathsf{A}}$; this gives rise to some Steiner forest $T_{\mathsf{A}}$ in ${\mathsf{A}}$, which can exit this cell only via portals on sides of ${\mathsf{A}}$. In particular, the portion of the Steiner tree outside ${\mathsf{A}}$ can be solved independently as long as we know the following portal configuration: (i) the set of exiting portals on the side of this cell that are used by $T$ (which connect points outside ${\mathsf{A}}$ with those inside), and (ii) how these exiting portals are connected by trees in the Steiner forest $T_{\mathsf{A}}$.

Let $\Xi_{\mathsf{A}}$ be the set of portal configurations for a cell ${\mathsf{A}}$, and let $D=|\Xi_{\mathsf{A}}|$. Each cell has $(4m+4)^{4r}$ subsets of $4r$ portals, and each subset of portals can be partitioned $\mathsf{Bell}(4r)$ ways. As the Bell number $\mathsf{Bell}(k)$ is bounded above by $k^{k}$ we have $D=(4m+4)^{4r}\mathsf{Bell}(4r)<(4m+4)^{8r}$. Our goal is to compute, for each portal configuration $\sigma\in\Xi$, the minimum cost ${\mathsf{cost}}(\sigma)$ of any rectilinear Steiner forest within ${\mathsf{A}}$ that gives rise to this boundary condition. Assuming an arbitrary but fixed order of portal configurations in $\Xi=\{\sigma_{1},\ldots,\sigma_{D}\}$, the costs of all configurations can then be represented by a vector ${\vec{\mathsf{C}}}_{\mathsf{A}}\in{\mathbb{R}}^{D}$, where ${\vec{\mathsf{C}}}_{\mathsf{A}}[i]={\mathsf{cost}}(\sigma_{i})$. We call ${\vec{\mathsf{C}}}_{\mathsf{A}}$ the cost-vector for ${\mathsf{A}}$.

We now describe the DP algorithm to compute this cost-vector for all cells in a bottom-up manner (decreasing order of levels).

Base case: ${\mathsf{A}}$ is a leaf cell. In this case, there is at most one point $p$ from $V$ contained in ${\mathsf{A}}$. We enumerate all configurations, and compute ${\vec{\mathsf{C}}}_{\mathsf{A}}$ directly, which requires solving RSMT instances of bounded size.

Inductive step: ${\mathsf{A}}$ is not a leaf cell. The four child-cells ${\mathsf{A}}_{1},\ldots,{\mathsf{A}}_{4}$ of ${\mathsf{A}}$ are from the level below ${\mathsf{A}}$’s level, and thus, by the inductive hypothesis, we have already computed the cost-vectors ${\vec{\mathsf{C}}}_{{\mathsf{A}}_{i}}$ for $i=1,\ldots,4$. Consider any portal configuration $\sigma\in\Xi$. We simply need to enumerate all choices of portal configurations $\tau_{1},\ldots,\tau_{4}$ for child-cells ${\mathsf{A}}_{1},\ldots,{\mathsf{A}}_{4}$ that are consistent with $\sigma$, meaning that the portals on common sides are the same, and the connected components of portals from the 4 child-cells do not form cycles (hence, still induce a valid Steiner forest). We have

Final construction of approximate RSMT. At the end of DP, after we compute the cost-vector for the root cell, we identify the portal configuration $\sigma^{*}\in\Xi$ with lowest cost. To obtain the corresponding rectilinear Steiner tree, we perform a top-down backtracking: (i) for the root cell, from $\sigma^{*}$ we can retrieve the set of child-cell configurations $\tau_{1}^{*},\ldots,\tau_{4}^{*}$ generating $\sigma^{*}$; (ii) we repeat this until we reach all leaf cells. At each leaf cell we once again compute the optimal Steiner forest for the chosen configuration. Combining these Steiner forests yields an optimal $(m,r)$-light RSMT.

The forward processing involves two MLPs, $\mathsf{NN_{base}}$ and $\mathsf{NN_{DP}}$, and calls $\mathsf{NN_{DP}}$ recursively in a bottom-up manner to compute an implicit encoding of the costs of possible portal configurations.

Base-case neural network $\mathsf{NN_{base}}$. At the leaf level of Arora’s PTAS, each cell contains at most $1$ point. We instead terminate the quadtree decomposition when a cell contains no more than ${\mathsf{k}_{b}}$ points, where ${\mathsf{k}_{b}}$ is a hyperparameter. Given a leaf cell ${\mathsf{B}}$ with $V_{\mathsf{B}}\subset V$ the set of points contained in ${\mathsf{B}}$, $\mathsf{NN_{base}}$ takes relative coordinates of $V_{\mathsf{B}}$ as well as the set of portals on the sides of ${\mathsf{B}}$ as input, and outputs a ${\mathsf{d}_{c}}$-dimensional vector as an implicit encoding of the cost vector ${\vec{\mathsf{C}}}_{\mathsf{B}}\in{\mathbb{R}}^{D}$; note that ${\mathsf{d}_{c}}<<D$ in practice. Coordinates are specified relative to the bottom left corner of the cell, and are normalized by the cell size. If a leaf cell has $n_{p}<{\mathsf{k}_{b}}$ input points, we pad the remaining ${\mathsf{k}_{b}}-n_{p}$ coordinates with $(-1,-1)$. Each cell can have at most $4m+8$ portals ($m$ + 2 on each side) and these portals can appear only at $4m+8$ distinct relative locations in the cell. Portals are then provided to $\mathsf{NN_{base}}$ as $4m+8$ indicators in $\{0,1\}$. Here, each corner contains two portals to simplify computation by distinguishing connections passing through from different sides. (Arora’s PTAS uses a single portal at each corner.)

Terminating with at most ${\mathsf{k}_{b}}$ points in each leaf cell is advantageous as, for a quadtree cell with very few points, it is challenging to learn a meaningful encoding of portal configurations. If ${\mathsf{k}_{b}}$ is small, the majority of cells have few points inside. (Note that a complete degree-4 tree has around 75% of its nodes at the leaf level.) Hence, training on such a collection of cells tends to provide bad supervision, harming the effectiveness of learning. On the other hand, as ${\mathsf{k}_{b}}$ increases, larger Steiner-tree instances must be computed at leaves, which may negatively affect the performance of the framework. In our experiments, ${\mathsf{k}_{b}}$ is a hyperparameter; see Sec. 2 for its effect and choice.

DP-inductive-step neural network $\mathsf{NN_{DP}}$. Next, we use another MLP, $\mathsf{NN_{DP}}$, to simulate the function $f_{DP}$ which, as mentioned in Sec. 3.1, is equivalent to the DP in Arora’s PTAS. In detail, given a cell ${\mathsf{A}}$ at level $i$, let ${\mathsf{A}}_{1},\ldots,{\mathsf{A}}_{4}$ be its 4 child-cells at level $i+1$ in a fixed order. The neural network $\mathsf{NN_{DP}}$ takes encodings of ${\vec{\mathsf{C}}}_{{\mathsf{A}}_{1}},\dots,{\vec{\mathsf{C}}}_{{\mathsf{A}}_{4}}$, and the portals of ${\mathsf{A}}$, and generates an encoding of the cost vector ${\vec{\mathsf{C}}}_{\mathsf{A}}\in{\mathbb{R}}^{D}$ of the parent cell ${\mathsf{A}}$. Note that the encodings of ${\vec{\mathsf{C}}}_{{\mathsf{A}}_{1}},\dots,{\vec{\mathsf{C}}}_{{\mathsf{A}}_{4}}$ are produced by applying either $\mathsf{NN_{DP}}$ or $\mathsf{NN_{base}}$ at ${\mathsf{A}}_{1},\ldots,{\mathsf{A}}_{4}$. Again, portals are provided via $4m+8$ indicators in $\{0,1\}$.

As described above, the forward pass applies $\mathsf{NN_{base}}$ to all leaf cells, and then $\mathsf{NN_{DP}}$ to all internal vertices, to simulate the bottom-up DP algorithm of Arora’s PTAS. Now, we use two more NNs to simulate the backtracking stage of the DP algorithm. Together, these NNs construct a portal-likelihood map $\rho:\mathcal{P}\to[0,1]$, where $\mathcal{P}$ is the set of all portals.

Root-level retrieval neural network $\mathsf{NN_{top}}$. For a cell ${\mathsf{A}}$ let $\mathrm{Portals}({\mathsf{A}})$ denote the portals in ${\mathsf{A}}$ that are one level higher than ${\mathsf{A}}$’s level, i.e., the portals on the vertical and horizontal segments bisecting ${\mathsf{A}}$. Let ${\mathsf{R}}$ be the root cell with children ${\mathsf{A}}_{1},\ldots{\mathsf{A}}_{4}$. The input to $\mathsf{NN_{top}}$ is the output of $\mathsf{NN_{DP}}$ at ${\mathsf{R}}$, and the input of $\mathsf{NN_{DP}}$ at ${\mathsf{R}}$. The output of $\mathsf{NN_{top}}$ is a vector of the likelihoods of $\mathrm{Portals}({\mathsf{R}})$, where ${\mathsf{R}}$ is the root cell.

Backward retrieval step neural network $\mathsf{NN_{retrieve}}$. We compute the rest of the portal likelihoods in a top-down manner. This is achieved by using a neural network $\mathsf{NN_{retrieve}}$ (at all non-root non-leaf cells ${\mathsf{A}}$) which computes the likelihoods of $\mathrm{Portals}({\mathsf{A}})$. The input of $\mathsf{NN_{retrieve}}$ applied at a cell ${\mathsf{A}}$ with child-cells ${\mathsf{A}}_{1},\ldots$ and ${\mathsf{A}}_{4}$, is comprised of two parts. First, $\mathsf{NN_{retrieve}}$ takes the output of four instances of either $\mathsf{NN_{base}}$ or $\mathsf{NN_{DP}}$ applied at ${\mathsf{A}}_{1},\ldots$, and ${\mathsf{A}}_{4}$ (same as $\mathsf{NN_{top}}$). Second, $\mathsf{NN_{retrieve}}$ receives likelihoods for every portal on the boundary of ${\mathsf{A}}$. At this step, these likelihoods will have been computed previously by an instance of either $\mathsf{NN_{retrieve}}$ or $\mathsf{NN_{top}}$ applied at the level above ${\mathsf{A}}$’s level.

After the backward pass, we have a portal-likelihood map $\rho$ over all portals. We select all portals with likelihood greater than a threshold $t\in(0,1)$ as the initial set of Steiner points $S$. We then compute the minimum spanning tree over $S\cup V$, which takes time $O(|S\cup V|\log|S\cup V|)$. Next, we apply three local refinement steps:

1. 1.

   First, we iterate over all leaf cells and for each cell replace every connected component with the optimal Steiner tree connecting all of the component’s Steiner points (selected portals) and input points.

2. 2.

   Next, we remove Steiner points with degree less than 3 and round the locations of the remaining Steiner points to integer coordinates.

3. 3.

   Finally, we partition²²2We partition by iterating the following procedure: (i) select a leaf, (ii) use breadth-first search to select vertices until $k$ input points or all remaining input points are selected, then (iii) remove these vertices. The subtree induced by each removed set of vertices, which contains at most $k$ input points, forms an element of the partition. the tree into subtrees with $k$ or fewer input points. We replace each subtree with the optimal Steiner tree over both its input points and any Steiner point which is adjacent to a vertex in another partition.

Optimal Steiner trees are computed using GeoSteiner 5.1 (Juhl et al. 2018). The first step, cell refinement, introduces Steiner points into the interior of cells, since initially Steiner points can only be at portals. The second step removes unnecessary Steiner points and rounds the locations of Steiner points. Rounding simplifies the tree and can be done with minimal effect on the solution since an optimal Steiner tree always lies on the Hanan grid and thus has integer Steiner points. The final step, subtree refinement, allows Steiner points at portals to be moved.

For each training pointset, we compute the optimal Steiner tree using GeoSteiner 5.1 (Juhl et al. 2018). Each time this optimal tree crosses a side of a cell, we move the crossing to the nearest portal. The resultant set of portals with crossings is used as the target in computing the loss. For the loss, we use binary cross entropy with weights of $m+1$ for portals that are Steiner points in the target, and weights of 1 for other portals. As the classification is imbalanced, with most portals not being Steiner points, this weighting scheme prevents the classifier from achieving low loss with uniformly near-zero portal likelihoods.

We train all four networks in an end-to-end manner. In training, the models are connected in a tree structure similar to that used in RNNs (recursive neural networks) in that the output of $\mathsf{NN_{DP}}$ is fed into the same $\mathsf{NN_{DP}}$ instance repeatedly in our architecture. However, RNNs are often connected linearly while NN-Steiner connects NNs in a tree structure. To accelerate training, we utilize batch-mode training. Batch mode necessitates that the same tree structure is used for each training sample. Otherwise, the model shape would be different between training samples, and these samples could not be learned in parallel.

We use pointsets sampled from a distribution tailored for compatibility with batch-mode training. Pointsets are constructed as follows: (i) Sample $n_{0}$ points from a distribution $\mathcal{D}$ on an integer grid of size $N\times N$. (ii) Construct a depth-$d$ quadtree. (iii) In each quadtree leaf cell with greater than $k_{b}$ points, remove points randomly until only $k_{b}$ points remain. With these pointsets, a depth-$d$ quadtree can always be used with no more than $k_{b}$ points in every leaf cell. Note that while we use a fixed tree structure throughout training, for testing, the tree structure is decided by the specific test pointset and can have any shape or depth.