BAST: Bayesian Additive Regression Spanning Trees for Complex Constrained Domain

1 Introduction . Over the past few decades , data collected from complex constrained domains have attracted much attention in machine learning and spatial statistics . Domains with non-trivial geometries , such as irregular boundaries , sharp concavities , and/or interior holes due to geographic constraints ( e.g. , lakes and coasts ) , impose challenges on statistical modeling , as the Euclidean assumption underpinning many traditional statistical and machine learning methods no longer holds for data with intrinsic geometries . In this paper , we consider nonparametric regression problems with features lying on constrained domains or , more generally , compact Riemannian manifolds . To be more specific , we model the response variable Y ( s ) ∈ R at a location s on a compact Riemannian manifoldM as Y ( s ) = f ( s ) + ( s ) , ( s ) iid∼ N ( 0 , σ2 ) , ( 1 ) for some unknown function f : M → R and noise variance σ2 . In many applications , the true function f ( · ) may not be globally smooth but has discontinuities/abrupt changes across some narrow boundary regions in the domain . For example , housing prices can be substantially different in two neighboring communities , and ocean chlorophyll data that are separated by a narrow peninsula can exhibit distinct spatial patterns . It is of great need to develop new methodologies that can both respect intrinsic geometries of the domain and capture complicated local discontinuity patterns in the true function . There is growing literature on nonparametric regression and smoothing for complex domains . Spline smoothing methods [ 33 , 22 , 41 , 43 , 37 ] have been developed for data on constrained domains , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . but most of them focus on domains in R2 . Sangalli et al . [ 36 ] generalized the spline methods to constrained regions in R3 . For more general Riemannian manifolds , kernel based smoothing models , including kernel regressions [ 32 , 19 ] and local regressions [ 1 , 7 , 12 ] , were developed . Gaussian process ( GP ) regression is another popular tool for nonparametric regression problems , and many works focus on developing valid covariance kernels on spheres [ see 20 , 16 , 15 , among others ] . More recently , a few practical GP models for constrained domains and Riemannian manifolds were studied in the literature [ 24 , 31 , 4 , 13 ] . However , most of the aforementioned approaches assume globally smooth true functions and thus may not fully adapt to the ones with local discontinuities . Ensemble tree models [ 5 , 6 ] have been widely used in traditional nonparametric regression problems . One prominent example is the Bayesian additive regression trees ( BART ) model [ 10 ] , due to its versatility and capability of producing uncertainty measures . However , to our knowledge , ensemble methods have not been used for nonparametric regression on complex constrained domains , and almost all ensemble tree methods rely on binary decision tree partition models as their ensemble members ( weak learners ) . Nevertheless , binary trees may not be ideal to capture possibly highly irregular partitions on complex domains as they can only make splits parallel to Euclidean coordinate axes . For instance , in the U-shape domain shown in Figure 1 ( c ) , a complicated and over-clustered binary treed partition is needed to approximate a simple partition with three clusters ( marked by different colors ) . Moreover , the rectangular partitions do not comply to irregular domain constraints , which may cause the so called “ leakage ” problems on complex domains . With a similar partitioning idea , Menafoglio et al . [ 28 ] proposed a Voronoi tessellation based model to account for domain constraints , but their method imposes convexity restrictions on partitions . Most recently , spanning treed partition models have been demonstrated as an effective tool for characterizing partitions with flexible shapes [ 23 , 39 , 26 ] , but their focus is on traditional two-dimensional Euclidean spaces and linear regression settings . Our contribution in this paper is to propose a novel Bayesian additive regression spanning trees ( BAST ) model for nonparametric regressions on complex constrained domains with efficient Bayesian inference algorithms . The backbone of BAST is a new random spanning tree ( RST ) manifold partition model , which replaces binary decision trees in each weak learner . RST is capable of capturing irregularly shaped partitions with a small number of spanning tree edge cuts while respecting intrinsic geometries and domain boundary constraints . Equipped with a soft prediction scheme , we show that BAST achieves a superior prediction performance over other competing methods on various tasks , thanks to its strong local adaptivity to different levels of smoothness . The rest of the paper proceeds as follows . In Section 2 , we present the RST partition models on manifolds and develop a new Bayesian nonparametric regression model with RST ensembles . Section 3 discusses algorithms for Bayesian inference . In Section 4 , we illustrate the model performance by simulation experiments and a real chlorophyll data set in Aral Sea . Section 5 concludes the paper with some discussions . Additional details on Bayesian inference , hyperparameter selection , and sensitivity analysis are provided in Supplementary Materials . The code of BAST is available at https : //github.com/ztluostat/BAST . 2 Bayesian Nonparametric Regressions with Additive Spanning Trees . 2.1 A Spanning Treed Partition Model on Manifolds . Our novel nonparametric regression model is built upon an ensemble of partitions of observations . In this subsection , we introduce a stochastic partition model for data on a compact Riemannian manifold via random spanning trees ( RSTs ) , which will serve as a building block to develop the sum-of-spanning-trees model in Section 2.2 . Let M be a d-dimensional compact Riemannian manifold that is known a priori , and S = { s1 , . . . , sn } ⊆ M be a finite set of locations onM where the data are observed . We are interested in partitioning S into several disjoint subsets such that each subset consists of nearby locations where the data are relatively homogeneous that can be modeled separately . For spatial data , it is often desired to impose contiguity constraints on partitions . Below , we introduce the notion of spatially contiguous partitions on a manifold based on a spatial graph whose edges encode the relationship of spatial adjacency or neighborhood . Let G = ( S , E ) be a connected undirected graph with vertex set S and edge set E that connects s ∈ S to its “ close neighbors ” . The construction of spatial graphs on a manifold will be discussed later in this subsection . We say π ( S ) = { S1 , . . . , Sk } , where Sj ⊆ S for j = 1 , . . . , k , is a spatially contiguous partition of S with respect to G if ∪kj=1Sj = S , Sj ∩ Sj′ = ∅ for all j 6= j′ , and there exists a connected subgraph Gj = ( Sj , Ej ) of G for each j . We call each Sj a cluster , which consists of locations that are connected to each other and thus is spatially contiguous with respect to the spatial graph . Henceforth , when there is no risk of confusion , we will refer to spatially contiguous partitions simply as partitions . Figure 1 ( a ) shows an example of a partition with three clusters . For constrained domains in R2 such as the U-shaped domain in Figure 1 , spatial graphs can be constructed via constrained Delaunay triangulations [ CDTs ; 8 ] . Specifically , let G0 be a CDT mesh on S ∪ SB , where SB is a set of locations on the domain boundaries . Then the induced subgraph of G0 on S can be chosen as a spatial graph G. Edges longer than a certain threshold can be removed if desired . See Figure 1 ( a ) for an example of G constructed via CDT . For a general manifold , motivated by the nice adaptive properties of nonparametric regressions based on K nearest neighbor ( K-NN ) graphs on manifolds [ 21 , 27 ] , one may construct G by a K-NN graph that connects s ∈ S to its K nearest neighbours with respect to geodesic distance . Given G , we model partitions on manifolds in a similar spirit as the spanning treed partition models developed for two-dimensional Euclidean spaces [ 23 , 39 , 26 ] . Specifically , a connected subgraph T = ( S , ET ) of G is called a spanning tree of G if it has no cycle . A well-known property of spanning trees is that if a set of k − 1 edges in ET is removed , we obtain a disconnected subgraph of T with k connected components , which naturally defines a partition π ( S ) with k clusters by letting Sj be the vertex set of the jth component . In this case , we say π ( S ) is induced by T . See Figure 1 ( b ) for an example . This property implies that we can simplify a complicated graph partition modeling problem to modeling spanning trees as well as the number and locations of removed edges . Mathematically , conditional on T and k we assume a discrete uniform distribution on all possible partitions induced by T : p { π ( S ) | k , T } ∝ 1 { π ( S ) is induced by T and has k clusters } , ( 2 ) where 1 ( · ) is an indicator function . Next , we consider a probabilistic model on the spanning tree space . Let ωe be the weight for an edge e ∈ E and ω = { ωe : e ∈ E } . We assume an iid uniform distribution on the edge weights and let T be the resulting minimum spanning tree ( MST ) , i.e. , the spanning tree with minimum ∑ e∈ET ωe : T = MST ( ω ) , ωe iid∼ Unif ( 0 , 1 ) , ( 3 ) where MST ( ω ) denotes an MST of G based on the edge weights ω . Note that we are not assuming a discrete uniform distribution on the spanning tree space , with which it is challenging to sample spanning trees for Bayesian inference . As we will show in Section 3.1 , our model specification leads to an exact and fast spanning tree sampler , taking advantage of the Prim ’ s algorithm for MST constructions . Finally , we assume a truncated Poisson distribution with mean parameter λk on the number of clusters : k ∼ Poisson ( λk ) · 1 ( 1 ≤ k ≤ k̄ ) , ( 4 ) where k̄ is the pre-specified maximum number of clusters . The following proposition states that the support of RST is rich enough to accommodate all possible spatially contiguous partitions on a manifold with no more than k̄ clusters . The proof is postponed to Appendix C. Note that similar results do not hold for binary treed partition models . See Figure 1 ( c ) for a counterexample where there does not exist a rectangle containing all the blue points without including any green or red ones . Proposition 1 Let π ( S ) = { S1 , . . . , Sk } be an arbitrary spatially contiguous partition . Then π ( S ) is within the support of the partition model defined by ( 2 ) , ( 3 ) , and ( 4 ) if k ≤ k̄ .

Summary: This paper proposes a modification of Bayesian Additive Regression Trees tailored to complicated or constrained input space. The basic idea, like BART, is to learn a collection of piecewise-constant weak learners to approximate an unknown regression function. Unlike BART, which parametrizes each weak learner using a binary regression trees with axis-aligned splits, the proposed BAST procedure parametrizes the weak learners using a spatially-contiguous partition of the input space that is index by a random spanning tree. The proposed procedure works quite well on simulated and real world data. It represents a very interesting new direction in a growing literature around BART.