Neural Conditional Probability for Uncertainty Quantification

We outline our contributions and offer further context in the global response, hoping these additions will better highlight the value of our work on the inference of conditional probability and uncertainty quantification.

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses

• W1: We thank the reviewer for this comment. We have added a detailed comparison of the NCP method to the main existing strategies in the global response. In brief, while the operator approach is fundamentally different, it can be viewed as a significant improvement over direct learning strategies that rely on a pre-specified dictionary of functions. The NCP method, in contrast, learns the latent space representation (dictionary) that is best adapted to the data under investigation.

• W2: The conditional operator approach has been the core idea in the ML theory of kernel mean embeddings and conditional mean embeddings [A], and, more recently, dynamical systems [B]. There, the operator is estimated on a universal (infinite-dimensional) reproducing kernel Hilbert space, in order to transfer probability distributions of one marginal to another by encoding them into the mean of the canonical feature map. While this theory is rich, it is essentially limited to infinite-dimensional feature spaces, and, hence, suffers from scalability limitations and the need for experts to design specific kernels for application at hand. How to build finite-dimensional data representations and inference models based on them, to the best of our knowledge, was an open problem. So, we hope that our work on the operator approach to inference of conditional probabilities opens a new line of research into this promising methodology that offers several significant benefits. From a probabilistic perspective, our approach allows us to learn the true joint distribution of $(X,Y)$ without needing to specify a model, even when the relationship between $X$ and $Y$ is complex. This capability provides valuable insights into $X$ and $Y$. For instance, the independence of $X$ and $Y$ is conveniently captured by the nullity of the conditional expectation operator. Moreover, the conditional expectation operator enables us to derive essential objects such as the conditional PDF, conditional CDF, conditional quantiles, and moments, even in complex nonlinear scenarios, without prior knowledge about $(X,Y)$. From a computational standpoint, we utilize the spectral properties of compact operators along with the recent NCP self-supervised loss function to create a theoretically sound training loss. This loss function guarantees that the true target operator is the unique global minimizer (see Theorem 1), while being inexpensive to evaluate, facilitating fast and stable training.

References:

[A] Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur and Bernhard Schölkopf (2017), "Kernel Mean Embedding of Distributions: A Review and Beyond", Foundations and Trends® in Machine Learning: Vol. 10: No. 1-2, pp 1-141

[B] Kostic, V., Novelli, P., Maurer, A., Ciliberto, C., Rosasco, L., and Pontil, M. (2022). Learning dynamical systems via Koopman operator regression in reproducing kernel Hilbert spaces. In Advances in Neural Information Processing Systems.

Thank you for the detailed response. I am happy to keep my score.

We appreciate the reviewer's insightful evaluation and valuable comments. In what follows, we aim to address the highlighted weaknesses and respond to the reviewer's questions.

Weaknesses

• W1. Thank you for this remark. We added several high-dimensional experiments focused on UQ tasks, while most are synthetic to explore, as requested, scalability of NCP, one is a very complex real-world problem of protein folding, see also our global response. For the latter, we used a dataset that is considered as challenging for the task of inferring dynamics of the stochastic process governed by the conditional probability kernel. Applying NCP to this setting, we were able to infer this conditional density from learned graph representations of Chignolin mini-protein ensemble. This allowed us to predict the expected average pairwise atomic distances between heavy atoms and infer its 10th and 90th percentiles. As reported in Fig. 3, folded/unfolded protein corresponds to small/large coverage that originates from less/more variations of the ensemble. To the best of our knowledge, this is the first work that was able to perform UQ to infer metastable states of Chignolin, an important contribution in its own right.

• W2. We thank the referee for the suggestion. We will add a section in Appendix summarizing the fundamental properties of operator theory that we used to develop the NCP approach and refer to it at the beginning of Section 3. This reminder will include in particular the definition of operators on infinite-dimensional Hilbert spaces and the definition of compact operators. Finally we will state the Erchart-Young-Mirsky theorem which guarantees the existence of the SVD for compact operators.

Questions

• Q1: We thank the referee for the interesting question. Maximizing the log-likelihood function to find model parameters is a central method in statistics, offering strong theoretical guarantees for models satisfying proper regularity conditions. However, this method becomes computationally complex for deep learning models, as the log-likelihood function can be non-concave with multiple local maxima, making it challenging to find the global maximum. In contrast, our approach represents the joint distribution through an operator and uses the regularized NCP loss designed according to the best low-rank rank approximation for operators. This method has a strong theoretical guarantee, as proven in our Theorem 1, that the unique global minimizer of this loss is the true target operator. Additionally, the empirical NCP loss is smooth and can be evaluated with linear complexity in batch size and latent space dimension, making it compatible with standard deep learning training methods.

• Q2: Another great question. We've been investigating an alternative to regularization that leverages the Singular Value Decomposition (SVD) properties of compact operators to encode orthogonality directly into the model architecture, rather than enforcing it through regularization. This approach is still under investigation, and it is not yet clear if it will surpass our current regularization method. Indeed, our current regularization approach offers two main advantages:

  • Ease of Computation. The regularization term is straightforward to compute and integrates seamlessly into the loss function, ensuring scalability and full compatibility with mini-batching and existing optimization methods.
  • Theoretical Guarantees. As established in Theorem 1, the regularized NCP loss uniquely identifies the true operator, a guarantee that an alternative method encoding orthogonality directly into the model might fail to have. Namely, the orthogonality is defined via the true distributions and not the observed empirical ones. Thus, currently it is not clear how to properly encode it in the architecture and obtain the same level of theoretical guarantees.

• Q3: We added additional experiments both on synthetic and real high-dimensional data to illustrate that the NCP operator approach scales without any issue to high-dimensional, more complex data. Please see our global response for more details.

Thank you for your detailed and thoughtful feedback on our submission. We appreciate your recognition of our contribution to the field of operator learning and ML theory which was the primary objective of our work. We would like to address your concerns regarding the empirical evaluations and provide additional context.

On general note

Thank you very much for giving us the perspective to understand your review. While we followed all your suggestions for the revision, we would like to stress that theoretical guarantees developed in Sec. 3 and 5 are the integral part of this paper, and the heart of key contributions. Without being too technical, let us just briefly summarize what we feel as the most important aspects:

• NCP is based on deep representation learning by training two DNNs (architecture depending on the data modalities) via a new loss so that the learned representations (u’s and v’s) can lead to inference of diverse statistics derived from the conditional probability.
• For each of the considered statistics we prove finite sample generalisation bounds in high probability that rely on the concentration inequalities and optimality gap from DNN training.
• To the best of our knowledge, this is the first work to show how to provably infer conditional probability distribution using finitely many deep features, so that the error bounds depend only on the effective problem dimension and not the ambient dimensions.

We hope that this summary brings more light to the quality and impact of our contributions.

Weaknesses

Major

1. We thank the reviewer for raising this point. We added a discussion in the global response to situate our NCP approach in comparison to the four known strategies. For the reviewer's convenience, we focus here on a detailed comparison with the best contenders to our NCP method, as shown in Tab. 1: Normalizing Flows (NF) and FlexCode (FC) from Izbicki and Lee (2017). FlexCode (FC) is an example of the direct learning strategy for conditional density estimation (CDE). Our results in Tab. 1 demonstrate that NCP, which learns an intrinsic representation adapted to the data, consistently outperforms FlexCode, which relies on a standard dictionary of functions, in all but one example. In the exception, NCP ties with FC. This also exemplifies that FC is not efficient for capturing the data's intrinsic structure when the FC model is overly misspecified. Regarding Normalizing Flows (NF), which is used in a conditional training strategy, a major limitation is the need to retrain an NF model for each conditioning. Additionally, NF builds a density model over the whole ambient space, which may not be efficient for capturing the data's intrinsic structure. Finally, to the best of our knowledge, NF lacks theoretical guarantees when applied to UQ tasks where strong guarantees of reliability are required. For more details on this aspect for the NCP, see our response to the general note above. We will incorporate this discussion in a revised version of our paper.

2. Thank you for this remark. We added high-dimensional experiments focused on UQ tasks. While most are synthetic in order to explore, as requested, the scalability of NCP, one is a complex real-world problem of protein folding, see also our global response. For the latter, we used a dataset that is considered as challenging for the task of inferring dynamics of the stochastic process governed by the conditional probability kernel. Applying NCP to this setting, we were able to infer this conditional density from learned GNN representations of Chignolin mini-protein ensemble. This allowed us to predict the expected average pairwise atomic distances between heavy atoms and infer its 10th and 90th percentiles. As reported in Fig. 3, folded/unfolded protein corresponds to small/large coverage that originates from less/more variations of the ensemble. To the best of our knowledge, this is the first work that was able to perform UQ to infer metastable states of Chignolin, an important contribution in its own right.

Minor

• We added an experiment on how time complexity and inference quality scales w.r.t. dimensionality of $X$. In the attached pdf we included Fig. 2 presenting compute times for training and statistical accuracy for conditional density estimation (CDE) using the KS-distance.
• Thank you a nice suggestion, we will follow it.

Questions

1. As requested, we provided the AC with a link to an anonymous repository, as outlined in the NeurIPS guidelines, that includes notebooks for reproducing all our experiments.

2. We thank the referee for bringing this method to our attention. From our understanding, conditional flow matching is designed to efficiently train continuous normalizing flows in high-dimensional settings, particularly for image generation tasks. However, based on our review of the relevant literature, this method has not been explored for the purpose of UQ. In our benchmark, we used a discrete-time masked autoregressive normalizing flow, which performs well on our synthetic benchmark due to its universal approximation property [Papamakarios et al. Normalizing Flows for Probabilistic Modeling and Inference. JMLR 2021, 22(57):1−64]. We did not encounter any issues with training convergence for this model. Additionally, our primary objective is to provide a simple yet effective approach for the UQ tasks we investigated. For these reasons, we believe that incorporating this more complex model in our benchmark lies outside the scope of our study.

Limitations

As elaborated above, we added additional experiments both on synthetic and real data to show that the NCP operator approach doesn't suffer from this limitation. Indeed it scales without any issue to high-dimensional and more complex data, without degradation of inference quality. Lastly, as requested, we will collect all discussions regarding limitations in one paragraph.