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 new Bayesian additive regression spanning trees (BAST) model for nonparametric regressions on manifolds. The model uses a new random spanning tree manifold partition model as weak learners, which allows it to capture complicated local discontinuity patterns in the true function while respecting intrinsic geometries of the domain. Building on nice properties of spanning tree structures, an efficient Bayesian inference algorithm is developed. Experiments on simulation data and the real chlorophyll data in Aral Sea show that BAST significantly outperforms other competing methods.