Source-linked AI summary
A Deep Learning Model for Spatially Clustered Data via Differentiable Cluster Assignment
Kexuan Li, Weidong Ma
TL;DR
Spatial regression relationships can vary across unknown geographic regions, while existing approaches may not capture fully nonlinear within-region effects. This paper jointly learns location-based regions and neural-network regression surfaces, showing that sufficiently accurate partition estimation preserves the oracle regression rate. In housing data, however, a flexible global spatial learner predicted better.
Problem
Existing spatial clustering methods often represent regional effects linearly or through smooth varying-coefficient expansions, limiting fully nonlinear response surfaces with covariate interactions.
Method
The model jointly estimates location-only cluster assignments and cluster-specific neural-network regressions using annealed softmax relaxation with graph-smoothness and occupancy penalties.
Results
When assignment error decays sufficiently quickly, estimating the partition does not alter the leading regression rate, which matches the oracle rate.
Takeaways & Limitations
The method is most useful for approximately region-homogeneous relationships with abrupt boundaries, where it provides an explicit and reproducible geographic submarket representation.
Takeaways & Limitations
In the housing application, methods allowing location to enter the regression surface directly predicted better than the spatially clustered model.
Abstract
from arXiv · showhide
We consider nonparametric regression when the association between a response and its covariates changes across an unknown partition of a spatial domain. The proposed estimator learns the partition and the cluster-specific regression functions jointly. A neural network depending only on location determines cluster membership, while separate neural networks describe the covariate--response relationship within the clusters. An annealed softmax relaxation permits gradient-based estimation of the otherwise discrete assignments. Graph-Laplacian and occupancy penalties are used to discourage fragmented regions and degenerate solutions. We establish identifiability up to label permutation, bound partition error under a margin condition, and decompose prediction risk into regression and assignment components. The resulting rate agrees with that of an oracle estimator when the partition is estimated sufficiently accurately. Simulations show that joint estimation is useful when regression surfaces change abruptly across spatial boundaries, including settings with nonlinear effects, unequal region sizes, preferential sampling, and spatially correlated errors. Finally, a real data analysis is provided to demonstrate the validity and effectiveness of the proposed method.
1 Introduction
The paper addresses regression relationships that are relatively homogeneous within spatial regions but change abruptly at unknown boundaries. It jointly learns location-based cluster assignments and nonlinear within-cluster regressions, with differentiable optimization, spatial regularization, and theory separating regression from partition error.
- Motivation: Spatial heterogeneity can be gradual, but the paper targets relationships that are homogeneous within regions and change at their boundaries.Existing smooth spatial methods impose continuity or local smoothness, whereas region-based approaches directly represent boundary changes.
- Proposed method: Joint estimation links spatial clustering to regression relationships rather than first clustering locations by distance or sampling density and then fitting separate regressions.This distinguishes the method from a two-stage procedure and from geographic proximity-based clustering alone.
- Proposed method: The method determines cluster membership from spatial location and models each cluster’s conditional mean with a neural network of nonspatial covariates.The location-only assignment preserves the geographic meaning of clusters, while within-region networks allow nonlinear effects.
- Estimation: Annealed continuous categorical relaxations make discrete assignments amenable to gradient-based training, while a graph-Laplacian penalty encourages neighboring locations to have similar assignment probabilities.The temperature decreases over epochs to address the computational obstacle posed by discrete partitions and hard argmax assignments.
- Theory: The theory establishes identifiability up to label permutation and separates prediction risk into regression-function and estimated-partition components.The identifiability result uses overlap of regional covariate distributions and separation of cluster effects; the analysis accounts for unknown spatial assignment beyond standard neural-network approximation results.
- Empirical evaluation: Simulations examine abrupt spatial changes under nonlinear effects, unequal region sizes, preferential sampling, and spatially correlated errors, including robustness beyond independent-sampling theory.The experiments contrast smooth and piecewise spatial heterogeneity and include two-region and three-curved-region settings.
2 Model
The model represents spatially varying covariate–response relationships through an unknown partition of the spatial domain into K0 clusters. It jointly estimates location-based cluster assignments and separate deep-neural-network regression functions for each cluster.
- Model setup: The data consist of independent observations with spatial locations s_i, covariates x_i, and responses y_i, partitioned into K0 unknown spatial clusters.Each true region is C_k = {s ∈ S : C_0(s) = k}.
- Model setup: Within cluster k, the response follows Y = f_0,C_0(S)(X) + ϵ with E(ϵ | S, X) = 0, where f_0,k maps covariates to the response.The cluster-specific regression functions may differ across the spatial partition.
- Neural-network specification: Deep neural networks estimate the complex cluster-specific regression functions, while a separate assignment network estimates regional membership probabilities.The assignment network is designed for flexible estimation of both spatial partitioning and nonlinear regression relationships.
- Neural-network specification: The assignment network depends only on spatial location, whereas the within-cluster regression functions are represented by separate neural networks.This separates spatial clustering from covariate–response modeling within each region.
3 Estimation and Optimization
The estimator replaces nondifferentiable hard assignments with an annealed differentiable relaxation for joint network training. Graph-Laplacian and occupancy penalties stabilize spatial regions and discourage cluster collapse, after which hard assignments are recovered.
- Differentiable assignment: A Gumbel–softmax or deterministic annealed-softmax relaxation makes cluster assignments differentiable for gradient-based training.The experiments use the deterministic version, setting g_k = 0 to avoid Monte Carlo variation; as τ → 0, assignments concentrate on a simplex vertex.
- Regularization: A graph-Laplacian smoothness penalty discourages fragmented cluster boundaries by coupling assignments across spatially adjacent locations.Adjacency can be defined by k-nearest neighbors or a distance-threshold graph, with weights decreasing with distance.
- Regularization: A Kullback–Leibler occupancy penalty discourages cluster collapse by favoring equal empirical occupancy across clusters.The penalty is minimized when ¯π_k = 1/K for all k.
- Optimization: The temperature is annealed across epochs, and hard assignment rules are restored after training.Training uses τ_t = max{τ_min, τ_0ρ^t}, with τ_0 > τ_min > 0 and 0 < ρ < 1.
- Optimization: Joint assignment and regression networks are trained with mini-batch gradient descent from multiple initializations because the objective is nonconvex.After the final epoch, relaxed assignments are discarded and hard rules are used.
- Model selection: Cross-validation is the default recommendation for selecting the number of clusters because neural-network degrees of freedom are difficult to approximate.The information criterion uses residual variance and an effective degrees-of-freedom approximation.
4 Theoretical properties
Under overlap, separation, regularity, and boundary assumptions, the population target is identifiable up to label permutation, while partition error is controlled by boundary localization and a margin condition. Prediction risk separates regression and assignment errors, yielding an oracle-order conclusion only when the boundary is learned sufficiently quickly.
- Identifiability: Theorem 1 identifies the population minimizers as the true assignment and regression functions, up to a permutation of cluster labels and null sets.Overlap and separation are both required; without either, one fitted regional function can represent multiple true regions without increasing population risk.
- Partition error: Theorem 2 bounds partition error through boundary localization and the margin condition, since an an-wide boundary displacement affects only a tube with probability of order an^γ.The localization error is permutation-invariant and measures distance from incorrectly assigned locations to the true boundary.
- Oracle comparison: The joint estimator matches the oracle regression order up to logarithmic factors only when assignment error is no larger than the oracle regression term.Consistent partition estimation alone is insufficient, and continuous annealed softmax optimization is not theoretically analyzed as a separate population target.
- Risk decomposition: Theorem 3 decomposes prediction risk into error when assignments are correct and the additional cost of assigning locations to incorrect regional functions.The decomposition remains applicable beyond neural-network regression estimators, while bounded fitted functions yield the constant 4M^2.
5 Simulation studies
Simulation studies evaluated prediction, spatial-partition recovery, unequal region sizes, and spatial regularization across two- and three-region settings. The proposed joint estimator generally outperformed implementable alternatives, especially when nonlinear response surfaces changed sharply across curved boundaries.
- Study design: The experiments assessed prediction accuracy, partition recovery, unequal cluster sizes, and spatial regularization in two-region and nonlinear three-region settings.Experiment 2 additionally included curved regions, spatially structured covariates, preferential sampling, and correlated errors.
- Experiment 1: At n = 1200, the proposed ISE was 0.280 with unequal cluster sizes, versus 0.776 for the global network, 0.836 for GWR, and 1.244 for two-stage clustering.Its misclassification rate declined from 0.235 at n = 300 to 0.036 at n = 1200.
- Experiment 1: At n = 1200 in the weak-separation setting, proposed ISE was 0.201, compared with 0.212 for GWR, 0.259 for the global network, and 0.330 for two-stage clustering.The proposed estimator’s misclassification rate decreased from 0.158 to 0.045 as sample size increased.
- Experiment 2: In NL-IID at n = 1200, proposed ISE was 0.344, versus 0.731 for GWR, 1.028 for the global DNN, and 2.016 for the two-stage method.Under preferential sampling and spatially correlated errors, corresponding ISE values were 0.558, 0.897, 1.147, and 2.272.
- Model selection: Prediction-based cross-validation selected the correct K0 = 3 in 0.52 of NL-IID replications and 0.29 of NL-SP replications at n = 1200.In NL-SP, K = 4 or 5 was selected in 71% of larger-sample replications.
6 King County house sales
In King County house sales, the proposed model produced a stable, geographically coherent three-region partition with heterogeneous within-region price relationships. However, global DNN and GWR methods predicted better, indicating limited value for prediction when spatial variation is continuous.
- The analysis used 21,613 King County transactions from May 2014 to May 2015, with log sale price as the response and property characteristics plus geographic coordinates as covariates.
- The proposed estimator had mean log-scale RMSE 0.248, versus 0.189 for the global DNN and 0.196 for GWR, and did not improve on two-stage K-means.Residual Moran’s I was 0.460 for the proposed estimator compared with 0.233 for the global DNN.
- The fitted three-region partition was geographically coherent and highly reproducible, with 5,644, 9,275, and 6,694 sales across the regions.Average maximum assignment probabilities were 0.920, 0.935, and 0.922, respectively.
- Regional price relationships differed substantially: fitted prices were lower in the southern region at comparable living area, while construction grade, waterfront status, and view had regionally distinct effects.
- The application illustrates a limitation: continuous spatial variation in house prices was better predicted by methods allowing location to enter the regression surface directly.
7 Discussion
The method is designed for regression relationships that are homogeneous within spatial regions and change abruptly at boundaries, rather than as a general replacement for smooth spatial regression. Its theory separates partition-assignment error from regional regression error, while the discussion identifies practical extensions.
- Scope and empirical behavior: The method performed well in piecewise nonlinear simulations, whereas the King County data favored a global spatial neural network for prediction.The three-region housing fit served a different purpose than improving prediction.
- Theoretical implications: Separation and covariate overlap identify the partition and cluster-specific functions up to label permutation.The prediction bound isolates the cost of assigning a new location to the wrong region.
- Theoretical implications: a_n = O(n^-2β/(2β+p)) is the boundary-localization rate required for partition estimation not to change the leading regression rate.This condition separates the geometric requirement on the estimated partition from the nonparametric rate of regional regressions; otherwise assignment error may dominate.
- Extensions: The current theory assumes K = K0, although K is selected by cross-validation in practice, motivating consistent data-adaptive selection of K.Other proposed extensions include cluster-specific noise variances and generalized outcomes using appropriate losses.
A Proofs of the main results
The proofs establish identifiability up to label permutation, derive the partition-error bound, and justify neural-network regression rates and prediction-risk decomposition under the stated assumptions.
- Theorem 1: Each fitted label overlaps at most one true cluster, while every true cluster is represented by at least one fitted label.The argument uses distinct cluster-specific regression functions to rule out multiple overlaps and positive cluster probabilities to ensure coverage.
- Theorem 1: With equal numbers of fitted labels and true clusters, the correspondence is one-to-one, yielding identifiability up to a permutation of cluster labels.Matched clusters also have equal regression functions almost everywhere, completing the theorem.
- Theorem 2: Theorem 2 follows by conditioning on the training sample, applying the margin assumption to the partition error, and then taking expectations under the stated assumptions.The proof handles both small and large partition-error regimes before letting the approximation parameter tend to zero.
- Proposition 1: For each cluster, Hölder-function approximation and a least-squares oracle inequality control neural-network regression error, with network complexity balanced against approximation error.The negligible low-sample event is controlled by exponential concentration and uniform boundedness.
- Theorem 3: Theorem 3 decomposes prediction error by conditioning on the training data and separately bounding its constituent terms.Test-point independence and boundedness provide the two component bounds whose combination proves the result.