Data-Driven Integration Kernels for Interpretable Nonlocal Operator Learning

TL;DR

Across this hierarchy, kernel-based models achieve near-baseline performance with far fewer trainable parameters, showing that much of the relevant nonlocal information can be captured through a small set of interpretable integrations when appropriate structural constraints are imposed.

Abstract

Machine learning models can represent climate processes that are nonlocal in horizontal space, height, and time, often by combining information across these dimensions in highly nonlinear ways. While this can improve predictive skill, it makes learned relationships difficult to interpret and prone to overfitting as the extent of nonlocal information grows. We address this challenge by introducing data-driven integration kernels, a framework that adds structure to nonlocal operator learning by explicitly separating nonlocal information aggregation from local nonlinear prediction. Each spatiotemporal predictor field is first integrated using learnable kernels (defined as continuous weighting functions over horizontal space, height, and/or time), after which a local nonlinear mapping is applied only to the resulting kernel-integrated features and any optional local inputs. This design confines nonlinear interactions to a small set of integrated features and makes each kernel directly interpretable as a weighting pattern that reveals which horizontal locations, vertical levels, and past timesteps contribute most to the prediction. We demonstrate the framework for South Asian monsoon precipitation using a hierarchy of neural network models with increasing structure, including baseline, nonparametric kernel, and parametric kernel models. Across this hierarchy, kernel-based models achieve near-baseline performance with far fewer trainable parameters, showing that much of the relevant nonlocal information can be captured through a small set of interpretable integrations when appropriate structural constraints are imposed.

Figures (4)

• Figure 1: Schematic of integration kernel learning. Learned kernels summarize predictor fields across horizontal space, height, and/or time into features, which are combined with local inputs and passed to a downstream nonlinear model to predict the local output
• Figure 2: Test set $R^2$ for baseline (blue), nonparametric kernel (yellow), and parametric kernel (red) models. Values are computed in the standardized, log-transformed space. Model labels indicate the dimensions treated nonlocally, with a subscript 0 denoting locality. Superscript k denotes nonparametric kernels, while TH, EXP, G, and MG denote parametric kernel families (top-hat, exponential, Gaussian, and mixture-of-Gaussians; see Appendix \ref{['app:families']}). The MIX kernel uses mixture-of-Gaussians kernels for RH and $\theta_e^*$ and an exponential kernel for $\theta_e$
• Figure 3: Learned vertical kernels for each predictor field from the two best-performing kernel-based models from Figure \ref{['fig:bars']}. Columns correspond to predictor variables. The top row shows nonparametric kernels and the bottom row shows kernels from the best parametric model, with row labels following the notation introduced in Figure \ref{['fig:bars']}. Solid lines indicate the mean kernel weight across training seeds (n = 3), with shading denoting $\pm1$ standard deviation. Larger absolute normalized kernel magnitude indicates greater contribution of a given pressure level, with sign indicating positive or negative influence. For predictors where mixture-of-Gaussians kernels were used, the dashed blue lines depict the individual components and the solid red line depicts their sum
• Figure A1: Learned vertical kernels for all kernel-based models shown in Figure \ref{['fig:bars']}, extending the results in Figure \ref{['fig:kernels']}