STACI: Spatio-Temporal Aleatoric Conformal Inference

We thank reviewer AwKK for taking the time to read through our work and provide valuable comments. We appreciate the praise for the clarity of our paper and originality of the methodology. We provide some rebuttals for individual comments here.

In the numerics, the metrics for assessing performance are poorly explained, i.e.: Coverage, Interval, Score, Interval Width. Could these be elaborated on and explained how they are computed as they would help the reader.

We agree that descriptions of these metrics are helpful for readers and elaborated on these metrics further in the Supplementary Material Section B.3.

SVGD might suffer from mode-collapse in high dimensions, would this have a negative effect on the proposed methodology at test time on harder problems?

You raise a great point on the dangers of mode-collapse for difficult high-dimensional problems. While we have not found this to be a problem for our tested datasets, we acknowledge this could happen in other use cases and would affect the UQ from the Bayesian variant of STACI. However, we believe the conformal addition would still provide valid intervals as the exchangeability calculations are based on the mean estimates of the covariance parameters and not their predictive variances.

How would one perform model selection in this framework; as opposed to marginal likelihood maximization in GPs? Could this be elaborated on?

Model selection would be done through cross validation in our methodology for the model fitting and conformal hyperparameters. Priors for the covariance parameters used in the kernel function can also be adjusted for a more direct connection to marginal likelihood maximization in GPs.

The proposed method is, in some parts, complex. Which parts of the model could be simplified if the inference task at hand is perhaps simpler (for example not none-stationary)?

We agree that this can be a complex framework. In a simpler stationary case, the need for the INR backbone is eliminated and this becomes a single hidden layer Bayesian neural network for the first half of our framework. Removal of the temporal component further reduces the number of priors set and model parameters needed to be learned.

What is the relation between (4) and Karhunen-Loeve expansions? If there is one, could this link be expanded on to help readers more familiar with these tools understand the proposed method?

The KL transformation is an alternative to the Fourier transformation used in the spectral approximation. Both represent a stochastic process as a linear combination of basis functions. The KL transformation is a generalization of the principle components analysis (PCA) where the basis functions are orthonormal and chosen to maximize variance explained. A Fourier transformation restricts the basis functions to be trigonometric functions. Like PCA, application of the KL decomposition typically requires replication of the process to estimate the basis functions, whereas the Fourier approximation can be applied to a single realization of the space-time process by assuming a stationary covariance (in latent space, in our case).

We appreciate reviewer w1e9 for taking the time to understand and provide feedback for our method. We appreciate the praise for soundness of our work and apologize for some areas that lack clarity. We provide some rebuttals for individual comments here.

Please provide a clear description of figure 1.

We apologize for lack of clarity and will add the following. Figure 1: STACI algorithm pipeline. Data from coordinates across space and time are first fed into our approximate spatio-temporal GP architecture and trained using SVGD. This allows for estimation of the latent space and prediction of the mean surface through the variational posterior. A conformal inference step is then fit to provide customized prediction intervals for valid uncertainty quantification.

Please ensure all concepts discussed in the paper are clearly defined before used.

We apologize for not defining the nugget variance beforehand. The nugget variance is the portion of the variance that is fully random and can't be explained by our spatio-temporal model. This can be thought of as the random/measurement error term, or in other words the amount of aleatoric uncertainty. We will add this into the paper for the camera-ready version.

Figure 2 is quite hard to interpret. Please provide more guidance.

Thank you for pointing this out! We will add a more informative description akin to:

The first row of Figure 2 represents the ground truth AOD values and the predicted means from the STACI Bayesian and conformal variants as well as the competitor models. The color scheme for this row goes from blue (lower values of opacity) to red (higher values of opacity). The second row represents the widths of the prediction intervals from each of the methods assuming a desired 95% coverage. The darker purple represents narrower widths that moves lighter into yellow which represents the wider. We can see that the conformal variant of STACI is the only one that has highly varied interval widths, while the UQ from other methods is fairly stagnant in the middle of purple and yellow. The high variance areas from conformal also correspond to areas with deep red high spikes in pollution which is commonly smoothed over in these classes of models.

Further consideration for the conclusions would help. Especially what to do next.

We thank the reviewer for highlighting this and further steps are important to extend this work. We believe that there are multiple possible extensions to non-Gaussian data. For example, adding a link function, or slightly changing the likelihood, should extend to distributions such as Bernoulli, or Poisson to model spatio-temporal correlated binary and count data. We also believe an extension to model spatial point processes is a feasible adaptation of the model. Following the work of Lloyd et al. (2015), the Poisson intensity can be modeled as the square of a GP. This would allow us to continue utilizing random Fourier features for efficient likelihood computation. These extensions would allow STACI to reach a wider audience and cover additional types of data.

Reference Lloyd, Chris, et al. "Variational inference for Gaussian process modulated Poisson processes." International Conference on Machine Learning. PMLR, 2015.

We'd like to thank the reviewer WrPZ for providing valuable comments for our work and directions for further improvement. We appreciate highlighting the novelty of our approach and clarity of the paper format. For the typos, we appreciate the reviewer for pointing these out and will fix these for the camera-ready version. We provide some rebuttals for individual comments here.

Although the authors compared the proposed model with several GP-based models in the experiments, they did not include a comparison with DNN-based models.

While we did not include DNN-based models, the backbone models for our latent space are popular Implicit Neural Representation DNN-based models currently in use for coordinate-based data. DNN-based models typically do not provide inherent uncertainty quantification, which was a major focus of our work.

While the proposed spectral approximation is theoretically justified, it would be helpful to understand the error bounds associated with the overall approximation. Specifically, how do factors such as the number of random features or SVGD particles affect the fidelity of the STGP approximation and the resulting prediction intervals?

We agree that theoretical justification is important for practical use of this method and show that the spectral approximation is centered on the Matérn covariance function. The approximation accuracy increases with $J$, the number of basis functions, or the width of the final hidden layer in STACI. In Appendix B.4, we also show that increasing the number of SVGD particles does not have much of an effect on RMSE, but slighty decreases the NLL, indicating slightly better variance estimation at the cost of higher computational time.

Have you tried to compare with the mixture model of local GPs in GPSat library?

While we did not compare with the mixture model from GPSat in our original study, we thank the reviewer for bringing it to our attention. From the Deep RF paper, we see it performs very well in this area, placing near the top in each metric across multiple datasets and showing more flexibility than traditional GPs. Unfortunately we could not implement it in time thus far, however we will make an effort to have results for the camera-ready version.

Why not using Continuous Ranked Probability Score (CRPS) to evaluate the quality of predicted uncertainties?

Great point to use CRPS as a UQ quality metric! We include the CRPS results here for both datasets with the non-conformal variant of STACI. We see that STACI with FFNP backbone is still the highest performer across these use cases.

For the MSS Data:

AOD Data:

I appreciate the authors’ thoughtful rebuttal and their efforts to incorporate the feedback. My comments are as follows:

While the authors make a reasonable point regarding the lack of inherent uncertainty quantification in DNN-based models, this does not fully justify the omission of standard DNN baselines for predictive performance comparison.

The suggested comparison to the mixture model from GPSat has not yet been addressed with empirical results. Although it is understandable that implementation constraints exist, the absence of this comparison leaves a noticeable gap in the current evaluation.

Given these outstanding concerns, I would prefer to retain my score as it currently stands.

We would like to thank the reviewer 8XoG for taking the time to read through our paper and providing much constructive feedback. We appreciate the reviewer's acknowledgment of the clarity of our work and the logic of the method. Rebuttals to specific comments are provided.

The evaluation is quite limited and missed many important baselines, including both GP and more advanced Gaussian Cox process models, such as variational Bayesian approximation, mean-field approximation, path integral approximations of GP, and change of kernel in a reproducing kernel Hilbert space. Without a complete comparison, the benefit of the proposed method is hardly convincing.

While we may have been unclear, we believe we have compared our method to many of these approaches. It is computationally infeasible to fit a GP to the full dataset, and all of the methods we compare to are approximations to a GP. We covered the variational Bayesian baseline by including SVGP. This KL-divergence based method was also shown to have better performance than the mean-field approximation in the original SVGP paper, so we did not include this in our comparisons. We also believe the Gaussian Cox process model, while advanced, is not in the scope of this work as we are not measuring count data across space and time. However, a potential area of future work is extending the method to non-Gaussian data, like the Gaussian Cox process model does. Replacing the Gaussian likelihood with other distributions such as Bernoulli or Poisson to handle geo-referenced binary or count data should be fairly straightforward. It should also be possible to extend the method to more challenging cases such as inhomogeneous Poisson processes for spatial point pattern data. For example, the model of Lloyd et al (2015) takes the Poisson process intensity function as the square of a GP. Thus the likelihood with random Fourier features approximating the GP can be computed analytically in O(n) flops, permitting the same computational approach taken here to fit the non-stationary Fourier model to spatial point pattern data.

The pros and cons of the proposed algorithm compared to these baselines isn’t clear.

We explain the drawbacks of current baselines in paragraphs 2-4 of the introduction and apologize if this wasn't clear. Our method provides greater interpretability of covariance parameters than other deep GP methods and provides greater flexibility than the traditional GP. The conformal addition of our method is able to provide guaranteed coverage properties for prediction intervals, assuming the exchangeability condition holds. The drawback to the STACI method is the computational cost required for both SVGD training and conformal inference, however we do not believe it is prohibitive for most use cases.

Can any guarantees be provided for the approximation methods used in this algorithm? Can any theoretical analysis on the improvement be provided?

We agree that theoretical guarantees need to be made for the approximation and provided a theorem on the spectral covariance approximation in supplement section A. We show that the spectral approximation is centered on the Matérn covariance function with accuracy that increases with $J$, the number of basis functions, or the width of the final hidden layer in our method.

Reference Lloyd, Chris, et al. "Variational inference for Gaussian process modulated Poisson processes." International Conference on Machine Learning. PMLR, 2015.