A Deep Learning Model for Spatially Clustered Data via Differentiable Cluster Assignment

2026-08-14

AI Summary

This study addresses nonparametric regression with unknown spatial partitions by proposing a differentiable clustering assignment mechanism that enables end-to-end joint learning of partitions and regression functions. The method integrates positional neural networks with annealed Softmax for gradient estimation and incorporates graph Laplacian regularization to prevent fragmentation, effectively handling nonlinear effects and complex error structures. Theoretically, we establish a prediction risk decomposition framework and prove that the proposed estimator achieves oracle convergence rates. Empirical evaluations demonstrate superior performance under abrupt boundaries and complex sampling designs. Collectively, this work provides a novel paradigm for modeling spatial heterogeneity in settings where partition structures are not known a priori.

Abstract

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.

Problem

• Nonparametric Regression
• Spatially Clustered Data
• Unknown Spatial Partition
• Cluster-specific Regression

Innovation

• Differentiable Cluster Assignment
• Joint Estimation
• Annealed Softmax Relaxation
• Graph-Laplacian Penalty
• Spatial Nonparametric Regression