NIAQUE: Neural Interpretable Any-Quantile Estimation - Towards Large Probabilistic Regression Models

Dear Reviewer JSss, we would like to thank you for providing insightful feedback. Our detailed response appears below.

1. The Transformer model is used as a baseline. The basic principle here is that in order to operate in the multi-task setting, where each dataset represents a different task and has a different number of available features, the model needs to be able to process variable sized input. Both the proposed NIAQUE model and the Transformer model are capable of this mode operation, with NIAQUE demonstrating improved accuracy. Additionally, some of the existing distributional models such as neural process or conditional neural process pointed out by Reviewer p8Dt simply cannot do it. The multi-task operation creates the opportunity to co-train on multiple datasets and hence have one model for multiple problems. Finally, in our updated experiments, we show the feasibility of pretraining a large probabilistic regression model that can be fine-tuned on unseen target datasets, providing clear performance lift. We have revised the manuscript and included additional experiments and motivation for our work accordingly.

2. We have provided additional clarifications in the caption of Figure 3. We are happy to provide additional clarifications if Reviewer provides additional guidance.

3. We have provided additional transfer learning results in Table 2 and paragraph Transfer Learning Experiment. The learning pipeline we proposed provides the ability to train foundational probabilistic regression models that can be fine-tuned on target datasets providing a clear accuracy lift compared to the model trained on the same data from scratch. This is akin to the case of training foundational vision or language models that can be used in various downstream tasks through transfer learning. The feasibility of such transfer learning has not been demonstrated before in the context of probabilistic regression models and we are closing this gap in this paper.

Dear Reviewer NJbn, we would like to thank you for providing insightful feedback. Our detailed response appears below.

• W1: Indeed we have outlined this as a potential risk in the discussion section of our paper. We would like to note that this is not exclusive to our model, but rather is a common risk relevant for the class of multi-task models as a whole. We would also like to emphasize that most of the prominent present day models are multi-task models. We believe that an effective strategy to combat this risk factor is fine-tuning on the specific target dataset. In the light of our most recent transfer learning results summarized in Table 2 of the revised paper, we conclude that fine-tuning of a pretrained NIAQUE model is viable and produces tangible overall accuracy lift compared to the control model trained from scratch, even when we vary the size of the training set on the target unseen dataset, as shown in Table 2. This provides evidence that the model is capable of producing improved results for target datasets of varying sizes, addressing concerns raised in this comment. Additionally, based on our preliminary analysis, it appears that the lift is consistent across datasets. The metric comparisons on individual datasets for pretrained/global models and the local model will be presented in an appendix in the camera ready version of the paper.

• W2: NIAQUE model provides the capability to learn arbitrary quantiles of the target distribution. As such, all metrics of the model remain the same irrespective of the quantiles selected to compute COVERAGE @ $\alpha$ metric. The only metric that changes is COVERAGE @ $\alpha$. This is because (i) the point estimation metrics are computed based on the median and (ii) the CRPS metric is computed based on 200 quantiles sampled uniformly at random in (0, 1). In this light, the CRPS metric already does provide a measure of model success when operating on arbitrary quantiles, not just ones used for COVERAGE @ $\alpha$ computation. Additionally, for the NIAQUE model reported in Table 1, COVERAGE @ 50 is 52.2, COVERAGE @ 90 is 89.4 and COVERAGE @ 99 is 98.1. Note that this computation is done on the same model, without retraining it in any way for the specific target quantiles, but rather querying it with the user requested quantile value of interest. Therefore, the guideline for the user is as simple as: train the model in the way described in the paper and at inference time, query the model with the quantile of your choice, be it p50 (median), p5, p90 or any other. We have included the following clarification at the beginning of Section 2.1: "Therefore, at inference time the user of the model has the choice of querying the model with any combination of target quantiles that best suit the user's downstream application."

• W3: Our model and training method do not make any parametric assumptions about the nature of the target distribution. Hence, it is able to handle distributions of different properties, naturally. The LPRM-101 benchmark is expansive and it contains datasets of very different nature. Figure 3b demonstrates what happens when features are ranked by the importance weight, calculated per dataset, and then removed from model input based on their importance rank. The low-importance feature removal degrades accuracy visibly less than the removal of high importance features, across many datasets. This clearly demonstrates that the proposed importance works on a wide variety of data. We believe that this addresses the concern raised by the reviewer.

We thank the reviewer for pointing out the Neural Process and Conditional Neural Process papers. These approaches are not directly applicable to the problem we solve, because they solve the problem in the fixed dimensional input space. For example, typical settings considered in the papers include 1d regression (curve fitting), 2d regression (image completion), classification on Omniglot. The approaches operate on the the support set of $(x,y)_i$ tuples as explained in Section of 2.4 in https://arxiv.org/pdf/1807.01622. The dimensionality of $x$ and $y$ is assumed to be fixed. The meta-learning aspect then manifests itself, for example, in the case of ominglot, when the same-dimensional 32x32 images are fed at both training and inference time, whereas the classes at train and inference time do not overlap. This is quite different from our setup, in which across different datasets the dimensionality of $x$ changes. As an example, let's take the nearest neighbor baseline used by the conditional neural process in the Omniglot experiment (Table 2 of https://proceedings.mlr.press/v80/garnelo18a/garnelo18a.pdf). Suppose we would like to implement cross-dataset knowledge transfer with variable dimensionality $x$, what is then the meaning of the nearest neighbor? None really seems to emerge naturally. Yet, our transfer learning experiment shows that positive transfer across datasets with variable-dimensional input spaces is viable. To address the reviewer's point we have included the following paragraph in related work section.

Alternative approaches, such as Neural Processes \citep{garnelo2018neuralprocesses} and Conditional Neural Processes \citep{garnelo2018conditional}, also generate conditional probabilistic solutions to regression problems. However, these methods are limited to fixed-dimensional input spaces and are not directly applicable to the cross-dataset, multi-task learning problem addressed here, where datasets vary in the number of independent variables. Moreover, unlike \citet{garnelo2018neuralprocesses} and \citet{garnelo2018conditional}, our approach demonstrates the ability to transfer knowledge to entirely new datasets, even when their dependent variable domains do not overlap with the training data.

Dear Reviewer p8Dt, we would like to thank you for providing insightful feedback. We have addressed your major points by (i) including the transfer learning experiment on held-out datasets and (ii) providing the Transformer-local results in Table 1. Our results demonstrate the value of pretraining NIAQUE in multi-task fashion and confirms that the learnings on various probabilistic regression tasks are generalizeable and can be transferred on unseen regression datasets using our approach. Please refer to the updated manuscript for detailed results. We are sorry that it took some time to code and execute the experiments. Please let us know if this addresses your major points. In the meanwhile, we will work on providing the response to your other points.

Thanks for conducting the transfer learning experiments and including the Transformer-local results. Thanks also for the clarification on how this method relates to the (Conditional) Neural Process literature. I have raised my score to a 5.

Dear Reviewer uKfu, we would like to thank you for providing insightful feedback and for positive assessment of our work. Our detailed response appears below.

• The code to download benchmark datasets, pytorch dataloaders, models, training and evaluation scripts/notebooks along with trained model binaries will be publicly released if the paper is accepted. Therefore, we do not envision any problems with building on top of this work.

• The focus of this paper is to demonstrate the feasibility of multi-task probabilistic learning on heterogenous regression datasets. Our most recent results in the revised paper also demonstrate positive transfer after fine-tuning on unseen datasets in this setting. To be able to demonstrate these results we used raw dataset features, uniformly across all datasets, to ensure uniformity and apple-to-apple comparison. Overwhelming majority of datasets used in our study are not Kaggle competition datasets. Kaggle competitors typically spend significant amount of time on feature engineering. Your point definitely outlines an interesting direction for future work, that of finding the automatic feature extraction methods that work synergistically with large-scale regression models helping them compete against human feature engineering approaches. This requires significant additional investigative, experiment design and theoretical work, which is outside of the scope of our current paper.