The latter setting often arises in my field of research, neural operators: models that predict physical quantities in science and engineering. In this field, training samples are usually costly to obtain, but the ground truth is well understood from a theoretical perspective. This enables professionals to judge the quality of model predictions.
These judgements could be complemented by feature attributions, such as Shapley values. Shapley values attribute a model’s prediction to the input features. They satisfy a set of desirable properties and are widely established in the scientific community. However, they are usually too costly to compute and even approximation methods can be slow when the number of features is large.
When we want to explain multiple samples, amortization methods can shift the computational load from inference to training. One such method, FastSHAP (Jethani et al., 2022), trains a neural network to predict Shapley values, resulting in quick attributions for any new sample. Unfortunately, FastSHAP requires homogeneous inputs with the same set of features across the training set, and physical data usually lives on a continuous domain that is discretized on meshes that may differ from input to input.
For amortization to work, we need a shared property that we can generalize over. In this paper, we choose the underlying physical domain \(\Omega \subset \mathbb{R}^d\), which is a connected subset of the \(1\)-, \(2\)-, or \(3\)-dimensional space and is often shared by all samples of a dataset. The main obstacle is that Shapley values are defined only for finitely many players, whereas our region \(\Omega\) consists of infinitely many points.
The solution lies in Aumann-Shapley values. Introduced in 1974 (Aumann & Shapley, 1974), they generalize Shapley values to infinitely many players. They still cannot give distinct attributions for every single point in our domain, but they allow each of the infinite players to contribute via a finite set of measures to the model prediction (which we call a game). This makes the attribution problem tractable. Moreover, for a game \(\nu(S)\), which assigns a payoff to each coalition \(S\), they yield a value \(\varphi\) that distributes this payoff over the players while preserving some of the Shapley properties.
To formalize this, let \((\Omega, \mathcal{C}, \boldsymbol{\mu})\) be a measure space, where \(S \in \mathcal{C}\) denotes a region in \(\Omega\) and defines a coalition. Typically, we want to model some physical property, which at a time \(t\) assigns a value to each location \(x\in \Omega\) via \(u: \Omega\times[0, \infty)\rightarrow \mathbb{R}\). For example, in the heat equation, \(u(x, t)\) is the temperature at location \(x\in \Omega\) and time \(t\). We want an AI model \(\hat{\mathcal{F}}\) to approximate \(u(x,t)\), usually using the true state \(u(x,t)\) as input to predict \(\hat{\mathcal{F}}(u(x,t))(x) \approx u(x,t+1)\). In practice, we work with a discretized version of \(u(x,t)\), i.e., a vector of evaluations of \(u(x,t)\) at specific locations. We denote this vector by \(\mathbf{u}_{t}\).
Now, we will explain the model’s predictions. In our Aumann-Shapley case, this means we will identify areas in \(\Omega\) that contributed positively or negatively to the prediction (relative to a baseline value). Concretely, we pick a location \(\hat{x}\in \Omega\), look at the prediction at that point \(\hat{\mathcal{F}}_{\hat{x}}(u(x,t)) := \hat{\mathcal{F}}(u(x,t))(\hat{x})\) compared to a baseline \(\hat{\mathcal{F}}_{\hat{x}}(\bar{u}(x,t))\) (e.g., \(\bar{u}(x,t)\equiv 0\)) and attribute that difference to the input areas. For the heat equation, we would then end up with an attribution plot like the following.
To show that our attributions are actually meaningful (i.e., that they satisfy similar properties to those of Shapley values), we build on Aumann-Shapley theory. They show that if our game has the form \(\nu(S)=f(\mu_{1}(S), \dots, \mu_{n}(S))\), where \(f\) is continuously differentiable, \(f(0)=0\), and \(\mu_{i}\) are measures, then there exists a value with the desired properties. Hence, our task is to create such a game for neural operators. Then, we know that a value for the game exists and can proceed to learn it from data.
To design our game, we need the following:
We choose \(f\) to be the neural operator that we want to explain (minus a baseline constant). For \(f\) to be continuously differentiable, both its activations and its input need to be sufficiently smooth.1 To isolate the effect of an area \(S\), we will design the input to take the value \(u(x,t)\) on \(S\) and a baseline value \(\bar{u}(x,t)\) elsewhere. As this constructed input needs to be smooth as well, we interpolate smoothly between \(S\) and \(\Omega\setminus S\), using bump functions \(\eta_{i}\) that partition the domain \(\Omega\). Concretely, for \(m>0\), we define our game as
\[\begin{aligned} \nu(S) &= \hat{\mathcal{F}}_{\hat{x}}(\mathbf{u}^{S}_{t}) - \hat{\mathcal{F}}_{\hat{x}}(\bar{\mathbf{u}}_{t}), \\ \mathbf{u}^{S}_{t} &:= \bar{\mathbf{u}}_{t} + \sum_{i=1}^{m} \mu_{i}(S)\, \eta_{i}\, (\mathbf{u}_{t} - \bar{\mathbf{u}}_{t}), \end{aligned}\]where \(\mu_{i}(S)\) is the share of the mesh nodes under \(\eta_{i}\) that belong to \(S\). The following visualization constructs a game from start to finish.
Equipped with a game that guarantees an Aumann-Shapley value, we can proceed to learn an amortized model for the value. The training procedure will closely follow FastSHAP:
After the training, we have a model that can predict an attribution map over arbitrary input fields \(u(x,t)\), and we can calculate values for coalitions \(S\) by integrating the attribution field over the respective area. In the paper, we include additional engineering tricks that can improve predictive performance. Many other techniques can surely be adapted from the discrete Shapley literature to get even better results. We can also explain other quantities of the model, including, empirically, non-differentiable readouts. In a follow-up work, we extend OperatorSHAP to capture interactions between regions of the domain. We believe that many other results from the discrete Shapley literature can likewise be transferred to the functional setting of OperatorSHAP.
We define \(f\) in function space, as this reflects PDE theory best. There, differentiability is more involved than on finite-dimensional spaces. Details can be found in the paper. ↩
In time series, the distribution may change arbitrarily over time and a deployed model can suddenly produce arbitrarily bad predictions, not hedgeable by anything we know from the training time. While this violates the assumptions of (classic) CP, there are multiple extensions that either relax its implications, e.g., by guaranteeing coverage only on average over an infinite time frame (see Adaptive Conformal Prediction, Gibbs & Candès, 2021), or by increasing the prerequisites, e.g., through knowledge about the time series (Oliveira et al., 2024), the error of the model (Auer et al., 2023), or the distributional change (Barber et al., 2023).
Personally, I do not see CP being used in life-or-death situations where minor deviations from the coverage guarantees could be critical. Therefore, due to its simplicity, I find Adaptive Conformal Prediction and its variants to be the most reasonable choice for time series applications. Nonetheless, the problem sparked our curiosity. Our field of study, partial differential equations (PDEs), which describe physical processes, is full of time series with additional structure and offers a plethora of mathematical tools to analyze their behavior.
For this article, it suffices to view PDEs as a time series on a continuous domain. For example, consider functions \(u(x,t): [0,1] \times [0, \infty) \rightarrow \mathbb{R}\) that have values at a point in time and space. Here, space is represented by the interval from $0$ to $1$ and time by the interval from $0$ to $\infty$. In practice, we rarely work with the full continuous function. Instead, functions are discretized in time and space. We will then write \(\mathbf{u}^t:=\left( u(0,t), u(\frac{1}{10},t), \dots, u(1,t) \right)\). This dichotomy raised one of our initial questions: Are there any differences between CP on the functional and the discrete form?
We will visualize some results using the heat equation. The heat equation is one of the simplest PDEs and describes how heat spreads on an object over time. Essentially, every point gets either hotter or colder, depending on the temperature difference to its surrounding area. The visualization below shows how the time series evolves for the functional and the discrete form, depending on the initial temperature distribution.
Suppose we observe the temperature of the object at some time \(t\) and want to predict it for some future time \(t+1\). Additionally, we would like to know how certain the model is in its prediction. This (un-)certainty is usually represented by a probability distribution over the possible predictions.
Now, in Conformal Prediction, we train our model and then estimate the error distribution on a held-out calibration set. We can use a sufficiently large quantile of the errors to construct prediction bands. If the test data follows the same distribution as the calibration data, we even get probabilistic guarantees that the true test value falls into the constructed bands.
Unfortunately, in time series, the data distribution may shift between calibration and test time, and with it the error distribution of our model. Thus, our guarantees vanish. Non-exchangeable CP (Barber et al., 2023), an extension of classical CP, prices this shift into our prediction bands. Depending on the total variation (TV) distance between calibration and test distributions, they widen the bands until they have recovered the probabilistic guarantees. However, if the TV distance gets too large, the resulting prediction bands will cover the whole prediction space and therefore become uninformative. The animation below shows the TV distance for two Gaussian distributions and how we have to adjust our prediction bands depending on their shift.
In the paper, we first investigate the TV distance in a function space setting. For this, we use the heat equation and assume the initial temperature along a 1D line segment is Gaussian2. Our theorem shows that for any evolution time \(t>0\), no matter how small, the TV distance between the initial and the final temperature distribution is always \(1\). This implies that even if our calibration and test time points are just an \(\varepsilon\) apart, we cannot apply (non-exchangeable) CP on them. This is because calibration and test distributions live on an infinite-dimensional (function) space. Thus, tiny changes accumulate over the infinite dimensions to result in mutually singular distributions. While we show this only on one example with the heat equation, this is a common behavior of distributions on function spaces.
While this result undermines the success of CP on functional data, it has little effect in practice. There, we usually work with a discretized version of a function, i.e., a finite-dimensional vector. And indeed, starting with an initial temperature vector from a finite-dimensional Gaussian, the TV distance between any two time points will always be strictly smaller than \(1\) (except in degenerate cases). Further, if we take a small enough time step, our CP prediction bands will even be bounded, providing useful guarantees.
In the paper, we provide a more general proposition. We can construct prediction bands for any linear PDE when the family of the initial distribution is closed under affine transformations (e.g., Gaussian, Cauchy, Laplace, and more). We use the algorithm from Tibshirani et al. (2019), which constructs the CP bands from the likelihood ratio between test and calibration densities. If the likelihood ratio of a test sample exceeds a certain threshold, prediction bands become infinitely wide and coverage is trivial. In general, this preserves the target coverage of the test samples over time, until all test samples get infinite bands. In practice, we would estimate the mean and covariance needed for the likelihood ratio from the data.
In the infinite-dimensional setting, we found that even arbitrarily small time shifts result in a TV distance of $1$. In finite dimensions, however, the TV distance is usually strictly smaller than $1$. We wanted to investigate if there is a natural connection between higher dimensionality, meaning larger grid sizes of the data, and TV distance. Indeed, it is not dimensionality itself that is driving this behavior. For a simple PDE that shifts in mean at a constant pace, we show that the drift direction relative to the covariance of the initial distribution is the deciding factor. As dimensionality implicitly affects this, we can mitigate the CP degeneration by restricting our prediction to a subset of our dimensions. To be precise, the prediction of our model can still use the full data space as input, but we need to carefully select the dimensions for which we want prediction bands (our paper shows some examples).
In conclusion, we have shown that on a function space, non-exchangeable CP will usually fail, while in discrete space, we can get robust prediction bands for linear PDEs by simply estimating the mean and covariance of the data. However, in that restricted setting, there are usually more straightforward methods to achieve accurate predictions and uncertainty bands. Our contribution should be viewed as theoretical, and we hope to inspire future work investigating different properties of PDEs and certain subclasses of PDEs that could unlock non-exchangeable CP in that setting.
Throughout the article, I will circumvent precise notions like exchangeability to prioritize readability over exactness. ↩
In particular, we assume a Gaussian distribution on a function space, which is a bit more complicated than “normal” Gaussian distributions. For this demonstration, it suffices to imagine that we randomly pick a somewhat smooth function fluctuating around 0. ↩