Neural Conditional Probability for Uncertainty Quantification

Authors: vladimir r kostic
karim  lounici
grégoire  pacreau
giacomo  turri
pietro  novelli
massimiliano  pontil
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:50:09

Abstract: We introduce Neural Conditional Probability (NCP), an operator-theoretic approach to learning conditional distributions with 
a focus on statistical inference tasks. NCP can be used to build conditional confidence regions and extract key statistics such as 
conditional quantiles, mean, and covariance. It offers streamlined learning via a single unconditional training phase, allowing 
efficient inference without the need for retraining even when conditioning changes. By leveraging the approximation 
capabilities of neural networks, NCP efficiently handles a wide variety of complex probability distributions. 
We provide theoretical guarantees that ensure both optimization consistency and statistical accuracy. 
In experiments, we show that NCP with a 2-hidden-layer network matches or outperforms leading methods. 
This demonstrates that a a minimalistic architecture with a theoretically grounded loss can achieve 
competitive results,  even in the face of more complex architectures.

Main paper: 
 introduction: 
 This paper studies the problem of estimating the conditional distribution associated with a pair of random variables, given a finite sample from their joint distribution. This problem is fundamental in machine learning, and instrumental for various purposes such as building prediction intervals, performing downstream analysis, visualizing data, and interpreting outcomes. This entails predicting the probability of an event given certain conditions or variables, which is a crucial task across various domains, ranging from finance (Markowitz, 1958) to medicine (Ray et al., 2017) , to climate modeling (Harrington, 2017) and beyond. For instance, in finance, it is essential for risk assessment to estimate the probability of default given economic indicators. Similarly, in healthcare, predicting the likelihood of a disease, given patient symptoms, aids in diagnosis. In climate modeling, estimating the conditional probability of extreme weather events such as hurricanes or droughts, given specific climate indicators, helps in disaster preparedness and mitigation efforts.According to Gao and Hastie (2022) , there exist four main strategies to learn the conditional distribution. The first one relies on the Bayes formula for densities and proposes to apply non-parametric statistics to learn the joint and marginal densities separately. However, most of non-parametric techniques face a significant challenge known as the curse of dimensionality (Scott, 1991; Nagler and Czado, 2016) . The second strategy, also known as Localization method, involves training a model unconditionally on reweighted samples, where weights are determined by their proximity to the desired conditioning point (Hall et al., 1999; Yu and Jones, 1998) . These methods require retraining the model whenever the conditioning changes and may also suffer from the curse of dimensionality if the weighting strategy treats all covariates equally. The third strategy, known as Direct Learning of the conditional distribution involves finding the best linear approximation of the conditional density on a dictionary of base functions or a kernel space (Sugiyama et al., 2010; Li et al., 2007) . The performance of these methods relies crucially on the selection of bases and kernels. Again for high-dimensional settings, approaches that assign equal importance to all covariates may be less effective. Finally, the fourth strategy, known as Conditional Training, involves training models to estimate a target variable conditioned on certain covariates. This is typically based on partitioning the covariates space X into sets, followed by training models unconditionally within each partition (see Gao and Hastie, 2022; Winkler et al., 2020; Lu and Huang, 2020; Dhariwal and Nichol, 2021, and references therein) . However, this strategy requires a large dataset to provide enough samples for each conditioning and is expensive as it requires training separate models for each conditioning input set, even though they stem from the same underlying joint distribution.The principal contribution of this work is a different conditional probability approach that does not fall into any of the four aforementioned strategies. Rather than learning the conditional density directly, our method, called Neural Conditional Probability (NCP), aims to learn the conditional expectation operator E Y |X associated to the random variables X ∈ X and Y ∈ Y based on data from their joint distribution. The operator is defined, for every measurable functionf : Y → R, as [E Y |X f ](x) := E[f (Y ) | X = x].NCP is based on a principled loss, leveraging the connection between conditional expectation operators and deepCCA (Andrew et al., 2013) established in (Kostic et al., 2024) , and can be used interchangeably to: Notably, our approach extracts statistics directly from the trained operator without retraining or resampling, and it is supported by both optimization consistency and statistical guarantees. In addition our experiments show that our approach matches or exceeds the performance of leading methods, even when using a basic a 2-hidden-layer network. This demonstrates the effectiveness of a minimalistic architecture combined with a theoretically grounded loss function.Paper organization In Section 2 we review related work. Section 3 introduces the operator theoretic approach to model conditional expectation, while Section 4 discusses its training pipeline. In Section 5, we derive learning guarantees for NCP. Finally, Section 6 presents numerical experiments. 
related works: 
 Non-parametric estimators are valuable for density and conditional density estimation as they don't rely on specific assumptions about the density being estimated. Kernel estimators, pioneered by Parzen (1962) and Rosenblatt (1956) , are a widely used non-parametric density estimation method.Much effort has been dedicated to enhancing kernel estimation, focusing on aspects like bandwidth selection (Goldenshluger and Lepski, 2011) , non-linear aggregation (Rigollet and Tsybakov, 2007) , and computational efficiency (Langrené and Warin, 2020) , as well as extending it to conditional densities (Bertin et al., 2014) . A comprehensive review of kernel estimators and their variants is provided in (Silverman, 2017) . See also (Tsybakov, 2009) for a statistical analysis of their performance. However, most of non-parametric techniques face a significant challenge known as the curse of dimensionality (Scott, 1991; Nagler and Czado, 2016) , meaning that the required sample size for accurate estimation grows exponentially with the dimensionality of the data (Silverman, 2017) . Additionally, the computational complexity also increases exponentially with dimensionality (Langrené and Warin, 2020) .Examples of localization methods include the work by Hall et al. (1999) for conditional CDF estimation using local logistic regression and locally adjusted Nadaraya-Watson estimation, as well as conditional quantiles estimation via local pinball loss minimization in (Yu and Jones, 1998) . Examples of direct learning of the conditional distribution include (Sugiyama et al., 2010) via decomposition on a dictionary of base functions. Similarly, Li et al. (2007) explores quantile regression in reproducing Hilbert kernel spaces.Conditional training is a popular approach which was adopted in numerous works, as in the recent work by Gao and Hastie (2022) where a parametric exponential model for the conditional density p θ (y|x) is trained using the Lindsey method within each bin of a partition of the space X . This strategy has also been implemented in several prominent classes of generative models, including Normalizing Flow (NF) and Diffusion Models (DM) (Tabak and Vanden-Eijnden, 2010; Dinh et al., 2014; Rezende and Mohamed, 2015a; Sohl-Dickstein et al., 2015) . These models work by mapping a simple probability distribution into a more complex one. Conditional training approaches for NF and DM have been developed in many works including (e.g. Winkler et al., 2020; Lu and Huang, 2020; Dhariwal and Nichol, 2021) . In efforts to lower the computational burden of conditional diffusion models, an alternative approach used heuristic approximations applied directly to unconditional diffusion models on computer vision related tasks (see e.g. Song et al., 2023; Zhang et al., 2023) . However, the effectiveness of these heuristics in accurately mimicking the true conditional distributions remains uncertain. Another crucial aspect of these classes of generative models is that while the probability distribution is modelled explicitly, the computation of any relevant statistic, say E[Y |X] is left as an implicit problem usually solved by sampling from p θ (y|x) and then approximating E[Y |X] via simple Monte-Carlo integration. As expected, this approach quickly becomes problematic as the dimension of the output space Y becomes large.Conformal Prediction (CP) is a popular model-agnostic framework for uncertainty quantification (Vovk et al., 1999) . Conditional Conformal Prediction (CCP) was later developed to handle conditional dependencies between variables, allowing in principle for more accurate and reliable predictions (see Lei and Wasserman, 2014; Romano et al., 2019; Chernozhukov et al., 2021; Gibbs et al., 2023 , and the references cited therein). However, (CP) and (CCP) are not without limitations. The construction of these guaranteed prediction regions need to be recomputed from scratch for each value of the confidence level parameter and of the conditioning for (CCP). In addition, the produced confidence regions tend to be conservative. 
operator approach to probability modeling: 
 Consider a pair of random variables X and Y taking values in probability spaces (X , Σ X , µ) and (Y, Σ Y , ν), respectively, where X and Y are state spaces, Σ X and Σ Y are sigma algebras, and µ and ν are probability measures. Let ρ be the joint probability measure of (X, Y ) from the product space X × Y. We assume that ρ is absolutely continuous w.r.t. to the product measure of its marginals, that is ρ ≪ µ × ν, and denote the corresponding density by p = dρ/d(µ × ν), so that ρ(dx, dy) = p(x, y)µ(dx)ν(dy).The principal goal of this paper is, given a datasetD n := (x i , y i ) i∈[n] of observations of (X, Y ), to estimate the conditional probability measure p(B | x) := P[Y ∈ B | X = x], x ∈ X , B ∈ Σ Y .(1) Our approach is based on the simple fact that p(B | x) = E[1 B (Y ) | X = x],where 1 B denotes the characteristic function of set B. More broadly we address the above problem by studying the conditional expectation operatorE Y |X : L 2 ν (Y) → L 2 µ (X ), which is defined, for every f ∈ L 2 ν (Y) and x ∈ X , as [E Y |X f ](x) := E[f (Y ) | X = x] = Y f (y)p(dy | x) = Y f (y)p(x, y)ν(dy),where L 2 µ (X ) and L 2 ν (Y) denotes the Hilbert spaces of functions that are square integrable w.r.t. to µ and ν, respectively. One readily verifies that∥E Y |X ∥ = 1 and E Y |X 1 Y = 1 X .A prominent feature of the above operator is that its rank can reveal the independence of the random variables. That is, X and Y are independent random variables if and only if E Y |X is a rank one operator, in which case we have that EY |X = 1 X ⊗ 1 Y . It is thus useful to consider the deflated operator D Y |X = E Y |X -1 X ⊗ 1 Y : L 2 ν (Y) → L 2 µ (X ), for which we have that [E Y |X f ](x) = E[f (Y )] + [D Y |X f ](x), f ∈ L 2 ν (Y).(2) For dependent random variables, the deflated operator is nonzero. In many important situations, such as when the conditional probability distribution is a.e. absolutely continuous w.r.t. to the target measure, that is p(• | x) ≪ ν for µ-a.e. x ∈ X , the operator E Y |X is compact, and, hence, we can write the SVD of E Y |X and D Y |X respectively asEQUATIONwhere the left (u ⋆ i ) i∈N and right (v ⋆ i ) i∈N singular functions form complete orthonormal systems of L 2 µ (X ) and L 2 ν (Y), respectively. Notice that the only difference in the SVD ofE Y |X and D Y |X is the extra leading singular triplet (σ ⋆ 0 , u ⋆ 0 , v ⋆ 0 ) = (1, 1 µ , 1 ν ) of E Y |X . In terms of densities, the SVD of E Y |X leads to the characterization p(x, y) = ∞ i=0 σ ⋆ i u ⋆ i (x) v ⋆ i (y) = 1 + ∞ i=1 σ ⋆ i u ⋆ i (x) v ⋆ i (y).The mild assumption that E Y |X is a compact operator allows one to approximate it arbitrarily well with a (large enough) finite rank (empirical) operator. Choosing the operator norm as the measure of approximation error and appealing to the Eckart-Young-Mirsky Theorem (see Theorem 3 in Appendix B.1) one concludes that the best approximation is given by the truncated SVD, that is for every d ∈ N,D Y |X ≈ [[D Y |X ]] d := d i=1 σ ⋆ i u ⋆ i ⊗ v ⋆ i , and [[D Y |X ]] d ∈ arg min rank(A)≤d ∥D Y |X -A∥,where the minimum is given by σ ⋆ d , and the minimizer is unique whenever σ ⋆ d+1 < σ ⋆ d . This leads to the approximation of the joint density w.r.t. marginals p(x, y) ≈ 1 +EQUATIONwhich in particular, choosing f = 1 B , givesP[Y ∈ B | X = x] ≈ P[Y ∈ B] + d i=1 σ ⋆ i u ⋆ i (x) E[v ⋆ i (Y ) 1 B (Y )].Moreover, we have thatP[Y ∈ B | X ∈ A] = ⟨1 A , E Y |X 1 B ⟩ P[X ∈ A] ≈ P[Y ∈ B] + d i=1 σ ⋆ i E[u ⋆ i (X)1 A (X)] P[X ∈ A] E[v ⋆ i (Y ) 1 B (Y )],for which the approximation error is bounded in the following lemma.Lemma 1 (Approximation bound). For anyEQUATIONNeural network model Inspired by the above observations, to build the NCP model, we will parameterize the truncated SVD of the conditional expectation operator and then learn it. Specifically, we introduce two parameterized embeddings u θ : X → R d and v θ : Y → R d , and the singular values parameterized by w θ ∈ R d , respectively given byu θ (x):=[u θ 1 (x) . . . u θ d (x)] ⊤ , v θ (y):=[v θ 1 (y) . . . v θ d (y)] ⊤ , and σ θ :=[e -(w θ 1 ) 2 , . . . , e -(w θ d ) 2 ] ⊤ ,where the parameter θ takes values in a prescribed set Θ.We then aim to learn the joint density function p(x, y) in the formp θ (x, y) := 1 + i∈[d] σ θ i u θ i (x) v θ i (y) = 1 + ⟨σ θ ⊙ u θ (x), v θ (y)⟩, where ⊙ denotes element-wise product. To that end, we consider the loss L γ (θ) := L(θ) + γR(θ) composed of two terms. The first termL(θ) := E (X ′ ,Y ′ )∼µ×ν [[p θ (X ′ , Y ′ ) -1] 2 ] -2E (X,Y )∼ρ [p θ (X, Y )] -1 (6)essentially has been considered by HaoChen et al. (2022) in the specific context of augmentation graph in self-supervised deep learning, linked to kernel embeddings (Wang et al., 2022) , and rediscovered and tested on DeepCCA tasks by Wells et al. (2024) . Indeed, this loss can be written in terms of correlations between features. Namely, denoting the covariance and variance matrices byEQUATIONand abbreviating u θ := u θ (X) and v θ := v θ (Y ) for simplicity, we can writeEQUATIONIf p=p θ for some θ∈Θ, then the optimal loss is theχ 2 -divergence L(θ)=D χ 2 (ρ | µ×ν)=-i≥1 σ ⋆ i 2and, as we show below, L(θ) measures how well p θ (x, y) -1 approximates i∈[d] σ ⋆ i u ⋆ i (x) v ⋆ i (y).However, in order to obtain a useful probability model, it is of paramount importance to align the metric in the latent spaces with the metrics in the data-spaces L 2 µ (X ) and L 2 ν (Y). For different reasons, a similar phenomenon has been observed in Kostic et al. (2024) where dynamical systems are learned via transfer operators. In our setting, this leads to the second term of the loss that measures how well features u θ and v θ span relevant subspaces in L 2 µ (X ) andEQUATION)We now state our main result on the properties of the loss L γ , which extends the result in Wells et al. (2024) to infinite-dimensional operators and guarantees the uniqueness of the optimum due to R.Theorem 1. Let E Y |X : L 2 ν (Y) → L 2 µ (X ) be a compact operator and D Y |X = ∞ i=1 σ ⋆ i u ⋆ i ⊗ v ⋆ i be the SVD of its deflated version. If u θ i ∈ L 2 µ (X ) and v θ i ∈ L 2 ν (Y), for all θ ∈ Θ and i ∈ [d], then for every θ ∈ Θ, L γ (θ) ≥ -i∈[d] σ ⋆2 i . Moreover, if γ > 0 and σ ⋆ d > σ ⋆ d+1 , then the equality holds if and only if (σ θ i , u θ i , v θ i ) equals (σ ⋆ i , u ⋆ i , v ⋆ i ) ρ-a.e., up to unitary transform of singular spaces.We provide the proof in Appendix B.3. In the following section, we show how to learn these canonical features from data and construct approximations of the conditional probability measure.Comparison to previous methods NCP does not fall into any of the four categories defined by Gao and Hastie (2022) , as it does not aim to learn conditional density of Y |X directly. Instead, NCP focuses on learning the operator mapping L 2 ν (Y) → L 2 µ (X ), from which all relevant task-specific statistics can be derived without requiring retraining. This approach effectively integrates with deep representation learning to create a latent space adapted to p(y|x). As a result, NCP efficiently captures the intrinsic dimension of the data, which is supported by our theoretical guarantees that depend solely on the latent space dimension (Theorem 2). In contrast, strategies designed for learning density often encounter significant limitations, such as the curse of dimensionality, potential substantial misrepresentation errors when the pre-specified function dictionary misaligns with the true distribution p(y|x), and high computational complexity due to the need for retraining. Experiments confirm NCP's capability to learn representations tailored to a wide range of data types-including manifolds, graphs, and high-dimensional distributions-without relying on predefined dictionaries. This flexibility allows NCP to outperform popular aforementioned methods. 
training the ncp inference method: 
 In this section, we discuss how to train the model. Given a training dataset D n = (x i , y i ) i∈[n] and networks (u θ , v θ , σ θ ), we consider the empirical loss L γ (θ) := L(θ) + γ R(θ), where we replaced (8) and (9) by their empirical versions. In order to guarantee the unbiased estimation, as we show within the proof of Theorem 1, two terms of our loss can be written using two independent samples (X, Y ) and (X ′ , Y ′ ) from ρ asL(θ)=E[L(u θ (X), u θ (X ′ ), v θ (Y ), v θ (Y ′ ), σ θ )] and R(θ)=E[R(u θ (X), u θ (X ′ ), v θ (Y ), v θ (Y ′ ))],where, the loss functionals L and R are defined for u, uEQUATIONEQUATIONTherefore, at every epoch we take two independent batches D 1 n and D 2 n of equal size from D n , leading to Algorithm 1. See Appendix A.1 for the full discussion, and Appendix A.2, where we also provide in Figure 4 : u ← u θ (X train ), u ′ ← u θ (X ′ train ), v ← v θ (Y train ), v ′ ← v θ (Y ′ train ) Compute L(θ)by averaging (10) over the batches Compute R(θ) by averaging (11) over the batches Compute NCP loss L γ (θ) := L(θ) + γ R(θ) and back-propagate end for Practical guidelines for training In the following, we briefly report a few aspects to be kept in mind when using the NCP in practice, referring the reader to Appendix A for further details. First, the computational complexity of unbiased estimation of the loss for a dataset of size n is O(nd), allowing one to seamlessly use NCP in contemporary DL settings. Second, the size of latent dimension d, as indicated by Theorem 1 relates to the problem's "difficulty" in the sense of smoothness of joint density w.r.t. its marginals. Lastly, after the training, an additional post-processing may be applied to ensure the orthogonality of features u θ and v θ and improve statistical accuracy of the learned model.Performing inference with the trained NCP model We now explain how to extract important statistical objects from the trained model ( u θ , v θ , σ θ ). To this end, define the empirical operatorEQUATIONwhereE y [ v θ i f ] := 1 n j∈[n] v θ i (y j )f (y j ).Then, without any retraining nor simulation, we can compute the following statistics:▶ Conditional Expectation: [ E θ Y |X f ](x) := E y f +[ D θ Y |X f ](x), f ∈ L 2 ν (Y), x ∈ X .▶ Conditional moments of order α ≥ 1: apply previous formula to f (u) = u α .▶ Conditional covariance: Cov θ (Y |X) := E θ Y |X [Y Y ⊤ ] -E θ Y |X [Y ] E θ Y |X [Y ⊤ ].▶ Conditional probabilities: apply the above conditional expectation formula withf (y) = 1 B (y), that is, p y (B) = E y [1 B ] and p θ (B | x) = p y (B) + i∈[d] σ θ i u θ i (x) E y [ v θ i 1 B ], B ∈ Σ Y , x ∈ X .Then, integrating over an arbitrary set A ∈ Σ X we getEQUATION▶ Conditional quantiles: for scalar output Y , the conditional CDF F Y |X∈A (t) is obtained by taking B = (-∞, t], and in Algorithm 3 in Appendix C we show how to extract quantiles from it. 
statistical guarantees: 
 We introduce some standard assumptions needed to state our theoretical learning guarantees. To that end, for any A ∈ Σ X and B ∈ Σ Y we define important constants, followed by the main assumption,φ X (A):=1 ∨ 1-P[X ∈ A] P[X ∈ A] and φ Y (B):=1∨ 1-P[Y ∈ B] P[Y ∈ B] .Assumption 1. There exists finite absolute constants c u , c v > 1 such that for any θ ∈ Θ ess supEQUATIONfor some large enough absolute constant C > 0.Remark 1. It follows easily from Assumption 1 that σ 2 θ (X)≤c 2 u d and σ 2 θ (Y )≤c 2 v d and consequentlyϵ n (δ) ≲ (c u ∨c v )[ d log(eδ -1 )/n∨(d log(eδ -1 )/n)].Finally, for a given parameter θ∈Θ and δ ∈ (0, 1), let us denoteEQUATION)In the following result, we prove that NCP model approximates well the conditional probability distribution w.h.p. whenever the empirical loss L γ (θ) is well minimized. Theorem 2. Let Assumption 1 be satisfied, and in addition assume thatEQUATION)Then for everyEQUATIONandEQUATIONhold with probability at least 1 -δ w.r.t. iid draw of the dataset D n = (x j , y j ) j∈[n] from ρ. Remark 2. In Appendix B.5, we prove a similar result under a less restrictive sub-Gaussian assumption on the singular functions u θ (X) and v θ (Y ).The rate ψ n (δ) in ( 16) is pivotal for the efficacy of our method. If we appropriately choose the latent space dimension d to ensure accurate approximation (σ ⋆ d+1 ≪ 1), achieve successful training (E θ ≪ 1), and secure a large enough sample size (ε n (δ) ≪ 1), Theorem 2 provides assurance of accurate prediction of conditional probabilities. Indeed, ( 19) guarantees (up to a logarithmic factor)P[Y ∈B | X∈A]-p θ (B | A)=O P 1 √ n + P[Y ∈B] P[X∈A] σ ⋆ d+1 +E θ + d/n+φ X (A)/ √ n ,Note the inclusion of the term P[X ∈ A] in the denominator of the last term on the right-hand side, along with φ X (A). This indicates a decrease in the accuracy of conditional probability estimates for rarely encountered event A, aligning with intuition and with a known finite-sample impossibility result Lei and Wasserman (2014, Lemma 1) for conditional confidence regions when A is reduced to any nonatomic point of the distribution (i.e. A = {x} with P[X = x] = 0). For rare events, a larger sample size n and a higher-dimensional latent space characterized by d are necessary for accurate estimation of conditional probabilities.We propose next a non-asymptotic estimation guarantee for the conditional CDF of Y |X when Y is a scalar output. This result ensures in particular that accurate estimation of the true quantiles is possible with our method. Fix t ∈ R and consider the setB t = (-∞, t] meaning that P[Y ∈B t |X∈A]=F Y |X∈A (t) and P[Y ∈B t ]=F Y (t).We define similarly for the NCP estimator of the conditional CDFF Y |X∈A (t)= p θ (B t | A).The result follows from applying (19) to the set B t . Corollary 1. Let the Assumptions of Theorem 2 be satisfied. Then for any t ∈ R and δ ∈ (0, 1), it holds with probability at least 1 -δ that An important application of Corollary 1 lies in uncertainty quantification when output Y is a scalar. Indeed, for any α ∈ (0, 1/2), we can scan the empirical conditional CDFEQUATION)EQUATIONIn App B.6, we derive statistical guarantees for the conditional expectation and covariance of Y . 
experiments: 
 Conditional density estimation We applied our NCP method to a benchmark of several conditional density models including those of Rothfuss et al. (2019) ; Gao and Hastie (2022) . See Appendix C.1 for the complete description of the data models and the complete list of compared methods in Tab. 2 with references. We also plotted several conditional CDF along with our NCP estimators in Fig. 6 . To assess the performance of each method, we use Kolmogorov-Smirnov (KS) distance between the estimated and the true conditional CDFs. We test each method on nineteen different conditional values uniformly sampled between the 5%-and 95%-percentile of p(x) and computed the averaged performance over all the used conditioning values. In Tab. 1, we report mean performance (KS distance ± std) computed over 10 repetitions, each with a different seed. NCP with whitening (NCP-W) outperforms all other methods on 4 datasets, ties with FlexCode (FC) on 1 dataset, and ranks a close second on another one behind NF. These experiments underscore NCP's consistent performance. We also refer to Tab. 3 in App C.1 for an ablation study on post-treatments for NCP.Confidence regions Our goal is to estimate conditional confidence intervals for two different data models (Laplace and Cauchy). We investigate the performance of our method in (21) and compare it to the popular conditional conformal prediction approach. We refer to App C.2 for a quick description of the principle underlying CCP. We trained an NCP model combined with an MLP architecture followed by whitening post-processing. See App C.2 for the full description. We obtained that way the NCP conditional CDE model that we used according to (21) to build the conditional 90% confidence intervals. We proceeded similarly to build another set of conditional confidence intervals based on NFs. Finally, we also implemented the CCP method of Gibbs et al. (2023) .In Fig. 1 , the marginal is X ∼ Unif([0, 5]) and Y |X = x follows either a Laplace distribution (top) with location and scale parameters (µ(x), b(x)) = (x 2 , x) or a Cauchy distribution (bottom) with location and scale parameters (x 2 , 1 + x). In this experiment, we considered a favorable situation for the CCP method of Gibbs et al. (2023) by assuming prior knowledge that the true conditional location is a polynomial function (the truth is actually the square function). Every other parameter of the method was set as prescribed in their paper.In Fig. 1 , observe first that the CCP regression achieves the best estimation of the conditional mean mse = 3.6 • 10 -3 against mse = 3.8 • 10 -2 for NFs and mse = 8.3 • 10 -3 for NCP, as expected since the CCP regression model is well-specified in this example. However, the CCP confidence intervals are unreliable for most of the considered conditioning. We also notice instability for NF and CCP when conditioning in the neighborhood of x = 0, with the NF confidence region exploding at x = 0. We suspect this is due to the fact that the conditional distribution at x = 0 is degenerate, hence violating the condition of existence of a diffeomorphism with the generating prior, a fundamental requirement for NFs models to work at all. Comparatively, NCP does not exhibit such instability around x = 0; it only tends to overestimate the confidence region for conditioning close to x = 0. The Cauchy distribution is known to be more challenging due to its heavy tail and undefined moments.In Fig 1 (bottom), we notice that NF and CCP completely collapse. This is not a surprising outcome since CCP relies on estimation of the mean which is undefined in this case, creating instability in the constructed confidence regions, while NF attempts to build a diffeomorphism between a Gaussian prior and the final Cauchy distribution. We suspect the conservative confidence region produced by NF might originate from the successive Jacobians involved in the NF mapping taking large values. In comparison, our NCP method still returns some reasonable results. Although the NCP coverage might appear underestimated for larger x, actual mean coverages computed on a test set of 200 samples are 88% for NCP, 99% for NF and 79% for CCP. Tab. 5 in Appendix C.2 provides a comparison study on real data for learning a confidence region with NCP, NF and a split conformal predictor featuring a Random Forest regressor (RFSCP). Y |X ∼ N (θ(X), sin(θ(X)/2). Right: Y ∈ {1, 2, 3, 4, 5} admits discrete distribution depending on θ(X): Y |X ∼ P1 if θ(X) ∈ [0, π/2), P2 if θ(X) ∈ [π/2, π), P3 if θ(X) ∈ [π, 3π/2), P4 if θ(X) ∈ [3π/2, 2π). We take P1 = (1/5, 1/5, 1/5, 1/5, 1/5), P2 = (1/2, 1/2, 0, 0, 0), P3 = (0, 0, 1, 0, 0), P4 = (0, 0, 0, 1/2, 1/2).High-dimensional experiment in molecular dynamics We investigate protein folding dynamics and predict conditional transition probabilities between metastable states. Figure 2 shows how, by integrating our NCP approach with a graph neural network (GNN), we achieve accurate state forecasting and strong uncertainty quantification, enabling efficient tracking of transitions. For further context and a full model description, see App C.3. 
conclusion: 
 We introduced NCP, a novel neural operator approach to learn the conditional probability distribution from complex and highly nonlinear data. NCP offers a number of benefits. Notably, it streamlines the training process by requiring just one unconditional training phase to learn the joint distribution p(x, y). Subsequently, it allows us to efficiently derive conditional probabilities and other relevant statistics from the trained model analytically, without any additional conditional training steps or Monte Carlo sampling. Additionally, our method is backed by theoretical non-asymptotic guarantees ensuring the soundness of our training method and the accuracy of the obtained conditional statistics.Our experiments on learning conditional densities and confidence regions demonstrate our approach's superiority or equivalence to leading methods, even using a simple Multi-Layer Perceptron (MLP) with two hidden layers and GELU activations. This highlights the effectiveness of a minimalistic architecture coupled with a theoretically grounded loss function. While complex architectures often dominate advanced machine learning, our results show that simplicity can achieve competitive results without compromising performance. Our numerical experiments suggest that, while our approach works well across different datasets and models, the price we pay for this generality appears to be the need for a relatively large sample size (n ≳ 10 4 ) to start outperforming other methods. Hence, a future direction is to study how to incorporate prior knowledge into our method to make it more data-efficient. Future works will also investigate the performance of NCP for multi-dimensional time series, causality and more general sensitivity analysis in uncertainty quantification. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: This is a theoretical paper. Experiments were carried out on synthetic data or publicly available data.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Paper of theoretical nature. There are no particular concerns.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: there is no such risk.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: all existing methods used in our experimental study were properly cited in the main paper and/or Appendix C.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: This is a theoretical work.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: see above.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: see above.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: The appendix is organized as follows:• Appendix A provides additional details on the post-processing for NCP.• Appendix B contains the proofs of the theoretical results and additional statistical results.• In Appendix C, comprehensive details are presented regarding the experiment benchmark utilized to evaluate the performances of NCP.A.1 Practical guidelines for training NCP• It is better to choose a larger d rather than a smaller one. Typically for the problems we considered in Section 6, we used d ∈ {100, 500}.• The regularization parameter γ was found to yield the best results for γ ∈ {10 -2 , 10 -3 }.• To ensure the positivity of the singular values, we transform the vector w θ with the Gaussian function x → exp(-x 2 ) to recover σ θ during any call of the forward method. The vector w θ is initialized at random with parameters following a normal distribution of mean 0 and standard deviation 1 /d.• With the ReLU function, we observe instabilities in the loss function during training, whereas Tanh struggles to converge. In contrast, the use of GELU solves both problems.• We can compute some statistical objects as a sanity check for the convergence of NCP training. For instance, we can ensure that the computed conditional CDFsatisfies all the conditions to be a valid CDF.• After training, an additional post-processing may be applied to ensure the orthogonality of operators u θ and v θ . This whitening step is described in Alg 2 in App A. We describe in Algorithm 2 the whitening post-processing procedure that we apply after training.Algorithm 2 Whitening procedureRequire: new data (X new ,Y new ); trained u θ , σ θ and v θ Evaluate u X = u θ (X train ) and v Y = v θ (Y train ) Centering: u X ← u X -Ê(u θ (X train )) and v Y ← v Y -Ê(v θ (Y train )) u X ← u X diag(σ θ ) 1 2 and v Y ← v Y diag(σ θ ) 1 2Compute covariance matrices :C X ← u ⊤ x ux /n C Y ← v ⊤ Y v Y /n C XY ← u ⊤ X v Y /n U, V, σ new ← SVD C -1 /2 X C XY C -1 /2 Y if (X new ,Y new ) is different than (X train , Y train ) then u X ← u θ (X new ) -Ê(u θ (X train ) diag(σ θ ) 1 2 v Y ← v θ (X new ) -Ê(v θ (Y train ) diag(σ θ ) 1 2 end if Final whitening: u new X ← u X C -1 /2 X U v new Y ← v Y C -1 /2 Y V return u new X , σ new , v new Y B Proofs of theoretical resultsB.1 A reminder on Hilbert spaces and compact operators Definition 1. Given a vector space H, we say it is a Hilbert space if there exists an inner product ⟨•, •⟩ such that:H is complete with respect to the norm ∥x∥ = ⟨x, x⟩ for all x ∈ H.An important example of an infinite-dimensional Hilbert space is L 2 µ (R), the space of squareintegrable functions w.r.t probability measure µ on R with the inner product defined as ⟨f, g⟩ = R f (x)g(x) µ(dx). Definition 2 (Bounded Operators). Let H 1 and H 2 be Hilbert spaces. A linear operator T : H 1 → H 2 is called bounded if there exists a constant C ≥ 0 such that for all x ∈ H 1 , the following inequality holds:∥T x∥ H2 ≤ C∥x∥ H1 . The smallest such constant C is called the operator norm of T , denoted by ∥T ∥, and is given by:∥T ∥ = sup x̸ =0 ∥T x∥ H2 ∥x∥ H1.Bounded operators are continuous and play a key role in functional analysis. Definition 3 (Compact Operators). Let H 1 and H 2 be Hilbert spaces. A bounded linear operatorT : H 1 → H 2 is called compact if for any bounded sequence {x n } ⊂ H 1 , there exists a subsequence {x n k } such that T x n k converges in H 2 .Compact operators can be viewed as infinite-dimensional analogues of matrices with finite rank in finite-dimensional spaces.A key result in the theory of compact operators is the existence of a singular value decomposition (SVD) for compact operators. The following is the statement of the Eckart-Young-Mirsky theorem: Theorem 3 (Eckart-Young-Mirsky). Let T : H 1 → H 2 be a compact operator between Hilbert spaces. Then T can be decomposed as:T = ∞ i=1 σ i ⟨•, u i ⟩v i ,where {u i } ⊂ H 1 and {v i } ⊂ H 2 are orthonormal sets, and σ i are the singular values of T , which satisfyσ 1 ≥ σ 2 ≥ • • • ≥ 0.Moreover, for any rank-k operatorT k = k i=1 σ i ⟨•, u i ⟩v i , we have: ∥T -T k ∥ = min rank(S)≤k ∥T -S∥,where ∥ • ∥ is the operator norm induced by the Hilbert spaces.Proof of Lemma 1. It follows from ( 3) and ( 4) thatP[Y ∈ B | X ∈ A] -P[Y ∈ B] - ⟨1 A , [[D Y |X ]] d 1 B ⟩ P[X ∈ A] = ⟨1 A , (D Y |X -[[D Y |X ]] d )1 B ⟩ P[X ∈ A] .Next, by definition of the operator norm, we have|⟨1 A , (D Y |X -[[D Y |X ]] d )1 B ⟩| ≤ ∥D Y |X -[[D Y |X ]] d ∥ L 2 ν (Y)→L 2 µ (X ) ∥1 A ∥ L 2 µ (X ) ∥1 B ∥ L 2 ν (Y) = ∥D Y |X -[[D Y |X ]] d ∥ L 2 ν (Y)→L 2 µ (X ) P[X ∈ A] P[Y ∈ B],where the operator norm∥D Y |X -[[D Y |X ]] d ∥ L 2 ν (Y)→L 2 µ (X ) is upper bounded by σ ⋆ d+1 by definition of the SVD of D Y |X .Proof of Theorem 1. In the following, to simplify notation, whenever dependency on the parameters is not crucial, recalling that (X, Y ) and (X ′ , Y ′ ) are two iid samples from the joint distribution ρ, we will denote the vector-valued random variables in the latent (embedding) space as u := u θ (X),u ′ := u θ (X ′ ), v := v θ (Y ) and v ′ := v θ (Y ′), as well as s = σ θ and S := S θ . Then, we can write the training loss simply asE [L γ (u, u ′ , v, v ′ , S)]. First, let us prove that L 0 (θ) = ∥U θ S θ V * θ ∥ 2 HS -2 tr(S θ U * θ D Y |X V θ ). Indeed, we have thatU * θ D Y |X V θ = U * θ E Y |X V θ -U * θ 1 X ⊗(V * θ 1 Y ) = E[u θ (X)E[v θ (Y ) ⊤ | X]]-(E[u θ (X)])(E[v θ (Y )]) ⊤ , that is U * θ D Y |X V θ = E[uv ⊤ ] -E[u]]E[v] ⊤ is simply centered cross-covariance in the embedding space. Recalling that U * θ U θ = E[uu ⊤ ] and V * θ V θ = E[vv ⊤] are covariance matrices in the embedding space, we have thatL 0 (θ) = -2 tr E[(S 1/2 u)(S 1/2 v) ⊤ ] + 2 tr(E[S 1/2 u] E[S 1/2 v] ⊤ ) + tr(E[(S 1/2 u)(S 1/2 u) ⊤ ]E[(S 1/2 v)(S 1/2 v) ⊤ ]) = -2E[u ⊤ Sv] + 2(E[u]) ⊤ S(E[v]) + tr(E[(S 1/2 u)(S 1/2 u) ⊤ ]E[(S 1/2 v)(S 1/2 v) ⊤ ])which, by taking (X, Y ) and (X ′ , Y ′ ) to be iid random variables drawn from ρ, gives that L 0 (θ) can be written asE -uSv -u ′ Sv ′ + u ′ Sv + uSv ′ + 1 2 tr S 1/2 uu ⊤ Sv ′ v ′⊤ S 1/2 + S 1/2 u ′ u ′⊤ Svv ⊤ S 1/2 = E 1 2 u ⊤ Sv ′ 2 + 1 2 u ′⊤ Sv 2 -(u -u ′ )S(v -v ′ ) = E [L 0 (u, u ′ , v, v ′ , s)] = E [L 0 (u θ (X), u θ (X ′ ), v θ (Y ), v θ (Y ′ ), σ θ )]. which implies that L 0 (θ) = ∥U θ S θ V * θ ∥ 2 HS -2 tr(S θ U * θ D Y |X V θ ). Moreover, to show that L γ (θ) = L 0 (θ) + γ R(θ). It suffices to note that ∥U * θ U θ -I∥ 2 F = tr((U * θ U θ -I) 2 ) = tr((U * θ U θ ) 2 -2U * θ U θ + I) = = tr(E[uu ⊤ ]E[u ′ u ′⊤ ]-E[uu ⊤ ]-E[u ′ u ′⊤ ]+I) = E[tr(uu ⊤ u ′ u ′⊤ -uu ⊤ -u ′ u ′⊤ +I)] = E (u ⊤ u ′ ) 2 -∥u∥ 2 -∥u ′ ∥ 2 + d,as well as that ∥U * θ 1 µ ∥ 2 = ∥Eu∥ 2 = (Eu) ⊤ (Eu) = Eu ⊤ u ′ , and apply the analogous reasoning for random variable Y ∼ ν.EQUATIONThen, applying the Eckhart-Young-Mirsky theorem, we obtain thatL r 0 (θ) ≥ r i=d+1 σ ⋆2 i -i∈[r] σ ⋆2 i = -i∈[d] σ ⋆2 i ,with equality holding whenever(σ θ i , u θ i , v θ i ) = (σ ⋆ i , u ⋆ i , v ⋆ i ), ρ-almost everywhere. To prove that the same holds for L 0 (θ), observe that after expanding the HS norm via trace in ( 22), we have thatL r 0 (θ) = -2 tr S 1/2 θ U * θ D r V θ S 1/2 θ + ∥U θ S θ V * θ ∥ 2 HS ,and, consequently,L r 0 (θ) = ∥U θ S θ V * θ ∥ 2 HS -2 tr(S 1/2 θ U * θ D r V θ S 1/2 θ ) = L 0 (θ) + 2 tr(S θ U * θ (D Y |X -D r )V θ). Thus, using Cauchy-Schwartz inequality, we obtain|L r 0 (θ) -L 0 (θ)| ≤ |tr(S θ U * θ (D Y |X -D r )V θ )| ≤ ∥S θ ∥∥U * θ ∥ HS ∥D Y |X -[[D Y |X ]] r ∥∥V * θ ∥ HS ,and, therefore,|L r 0 (θ) -L 0 (θ)| ≤ σ ⋆ r+1 tr(U * θ U θ ) tr(V * θ V θ ) ≤ M dσ ⋆ r+1, where the constant is given by M:= max i∈[d] {∥u θ i ∥ L 2 µ (X ) , ∥v θ i ∥ L 2 ν (Y) } < ∞. So, L r 0 (θ) -M dσ ⋆ r+1 ≤ L 0 (θ) ≤ L r 0 (θ) + M dσ ⋆ r+1, and, since r > d + 1 was arbitrary, we can take r arbitrary large to obtain σ ⋆ r → 0 and conclude that L0 (θ) ≥ -i∈[d] σ ⋆2 i , with equality holding when (σ θ i , u θ i , v θ i ) = (σ ⋆ i , u ⋆ i , v ⋆ i ), ρ-almost everywhere, since then U * θ D Y |X V θ = U * θ U θ S θ = S θ = diag(σ 1 , . . . , σ d ).Finally, we prove that γ > 0 and σ ⋆ d > σ ⋆ d+1 assure uniqueness of the global optimum. First, if the global minimum is achieved σ ⋆ d > σ ⋆ d+1 allows one to use uniqueness result in the Eckhart-Young-Mirsky theorem that states that i∈[d] σ ⋆ i u ⋆ i ⊗ v θ i = i∈[d] σ θ i u θ i ⊗ v θ i . But since, γ > 0 implies that R(θ) = 0, i.e. (u θ i ) i∈[d] ⊂ L 2 µ (X ) and (v θ i ) i∈[d] ⊂ L 2 ν (Y)are two orthonormal systems in the corresponding orthogonal complements of constant functions, using the uniqueness of SVD, the proof is completed.Proof of Theorem 2. Let us denote the operators arising from centered and empirically centered features as U θ , U θ :R d → L 2 µ (X ) and V θ , V θ : R d → L 2 ν (Y) by U θ z := z ⊤ (u θ -E[u θ (X)])1 X , V θ z := z ⊤ (v θ -E[v θ (Y )])1 Y and U θ z := z ⊤ u θ , V θ z := z ⊤ v θ , respectively, for z ∈ R d .We first bound the error of the conditional expectation model as∥D Y |X -D θ Y |X ∥ as follows. ∥D Y |X -D θ Y |X ∥ = ∥D Y |X ± [[D Y |X ]] d ± U * θ S θ V θ ± U * θ S θ V θ -U * θ S θ V θ ∥ ≤ σ ⋆ d+1 + E θ + ∥U * θ S θ V θ -UIn a similar way, we obtain∥U * θ S θ V θ -U * θ S θ V θ ∥ = ∥U * θ S θ V θ ± U * θ S θ V θ -U * θ S θ V θ ∥ ≤ ∥U θ -U θ ∥∥V θ ∥ + ∥ U θ ∥∥V θ -V θ ∥ ≤ ∥U θ -U θ ∥∥V θ ∥ + ∥V θ -V θ ∥∥U θ ∥ + ∥U θ -U θ ∥∥V θ -V θ ∥ ≤ 1 + E θ ∥ E x [u θ ] -E[u θ (X)]∥ + ∥ E y [v θ ] -E[v θ (Y )]∥ + ∥ E x [u θ ] -E[u θ (X)]∥ ∥ E y [v θ ] -E[v θ (Y )]∥ ≤ 2 1 + E θ ε n (δ) + [ε n (δ)] 2 .where∥ E x [u θ ] -E[u θ (X)]∥ ≤ ε n (δ) and ∥ E y [v θ ] -E[v θ (Y )]∥ ≤ ε n (δ) hold w.p.a.l. 1 -δ in view of Lemma 2.To summarize, it holds w.p.a.l. 1 -δEQUATIONBy definition in ( 4) and ( 13), we haveP[Y ∈ B | X ∈ A] -p θ (B | A) = E[1 B (Y )] -E y [1 B ] + ⟨1 A , D Y |X 1 B ⟩ E[1 A (X)] - ⟨1 A , D θ Y |X 1 B ⟩ E x [1 A ] ,and⟨1 A , D Y |X 1 B ⟩ E[1 A (X)] = ⟨1 A , (D Y |X -D θ Y |X )1 B ⟩ E[1 A (X)] + ⟨1 A , D θ Y |X 1 B ⟩ E x [1 A ] E x [1 A ] E[1 A (X)].A (X)∥ L 2 µ (X ) = E[1 A (X)] = P[X ∈ A], ∥1 B (Y )∥ L 2 ν (Y) = E[1 B (Y )] = P[Y ∈ B], for any A ∈ Σ X and B ∈ Σ Y and |⟨1 A , (D Y |X -D θ Y |X )1 B ⟩| ≤ ∥D Y |X -D θ Y |X ∥∥1 A (X)∥ L 2 µ (X ) ∥1 B (Y )∥ L 2 ν (Y) .Combining the previous observations, we getEQUATIONandEQUATIONwhere we have used that ∥D Y |X ∥ ≤ 1.Similarly, we haveEQUATIONNext Lemmas 3 and 4 combined with ( 17) and elementary algebra give w.p.a.l. 1 -2δ thatE x [1 A ] -E[1 A (X)] E x [1 A ] ≤ 2φ X (A)ϵ n (δ), E y [1 B ] -E[1 B (Y )] E x [1 B ] ≤ 2φ Y (B)ϵ n (δ),andE[1 A (X)] E x [1 A ] ∨ E[1 B (Y )] E y [1 B ] ≤ 2, E[1 A (X)]( E y [1 B ] -E[1 B (Y )]) E x [1 A ] E y [1 B ] ≤ 4φ Y (B)ϵ n (δ).It also holds on the same probability event as above that⟨1 A , (D Y |X -D θ Y |X )1 B ⟩ E x [1 A ] E y [1 B ] ≤ ∥D Y |X -D θ Y |X ∥ E[1 A (X)]E[1 B (Y )] E[1 A (X)] E x [1 A ] E[1 B (Y )] E y [1 B ] ≤4 ∥D Y |X -D θ Y |X ∥ E[1 A (X)]E[1 B (Y )].Combining Lemma 2 and (23), we get with probability at least1 -δ that ∥D Y |X -D θ Y |X ∥ ≤ ψ n (δ). By a union bound combining the last two displays with ( 23), ( 26), ( 24) and ( 25), we get with probability at least 1 -3δEQUATIONandEQUATIONReplacing δ by δ/3, we get the result w.p.a.l. 1 -δ.The following result will be useful to investigate the theoretical properties of the NCP method in the iid setting.Lemma 2. Let Assumption 1 be satisfied. Then there exists an absolute constant C > 0 such that, for any δ ∈ (0, 1), it holds w.p.a.l. 1 -δ∥ E x [u θ ] -E[u θ (X)]∥ ≤ C c u d log(eδ -1 ) n + σ θ (X) log(eδ -1 ) n .Similarly, w.p.a.l. 1 -δ∥ E y [v θ ] -E[v θ (Y )]∥ ≤ C c v d log(eδ -1 ) n + σ θ (Y ) log(eδ -1 ) n .Proof of Lemma 2. We note thatE x [u θ ] -E[u θ (X)] = 1 n n i=1 Z i with Z i = u θ (X i ) -Eu θ (X i ), ∀i ∈ [n].We note that∥Z i ∥ ≤ 2c u d =: U and Var(Z i ) = Var(∥u θ (X i ) -E[u θ (X i )]∥) = σ 2 θ (X) for any i ∈ [n].We apply Minsker (2017, Corollary 4.1) to get for any t ≥ 1 6 (U + U 2 + 36nσ 2 θ (X)),EQUATION)Replacing t by nt and some elementary algebra give for any t ≥ 1 6U n + U 2 n 2 + 36 σ 2 θ (X) n =: c, w.p.a.l. 1 -28 exp (-t), ∥ 1 n n i=1 Z i ∥ ≤ 4U 3 t n + 2σ θ (X) t n .Replacing t by t + c, we get for any t ≥ 0, w.p.a.l. 1 -28 exp (-t + c),∥ 1 n n i=1 Z i ∥ ≤ 4U 3 t + c n + 2σ θ (X) t + c n .Up to a rescaling of the constants, there exists a numerical constant C > 0 such that for any δ ∈ (0, 1), w.p.a.l.1 -δ ∥ 1 n n i=1 Z i ∥ ≤ C U n c + σ θ (X) c n + U t n + σ θ (X) t n .Elementary computations give the following bound, that is, there exists a numerical constant C > 0 such that for any t > 0, w.p.a.l. 1 -exp(-t)∥ 1 n n i=1 Z i ∥ ≤ C c u d n ∨ c 2 u d 2 n 2 ∨ σ 3/2 θ (X) n 3/4 ∨ σ 2 θ (X) n + c u d t n + σ θ (X) t n .Under Assumption 1 and the conditionc 2 u d n ≤ 1, it also holds that σ 2 θ (X) n ≤ 1 since σ 2 θ (X) ≤ c 2 u d.Consequently, the bound simplifies and we obtain for any t > 1, w.p.a.l. 1 -exp(-t)∥ 1 n n i=1 Z i ∥ ≤ C c u d t n ∨ σ θ (X) t n ,where C > 0 is possibly a different absolute constant from the previous bound. Taking t = log eδ -1 for any δ ∈ (0, 1) gives the first result. We proceed similarly to get the second result.Proposition 1. Caponnetto and De Vito (2007, Proposition 2) Let A i , i ∈ [n] be i.i.d copies of a random variable A in a separable Hilbert space with norm ∥•∥. If there exist constants L > 0 and σ > 0 such that for every m ≥ 2, E∥A∥ m ≤ 1 2 m!L m-2 σ 2 , then with probability at least 1 -δEQUATION)Lemma 5 ((Sub-Gaussian random variable) Lemma 5.5. in Vershynin ( 2011)). Let Z be a random variable. Then, the following assertions are equivalent with parameters K i > 0 differing from each other by at most an absolute constant factor.1. Tails: P{|Z| > t} ≤ exp(1 -t 2 /K 2 1 ) for all t ≥ 0;2. Moments:(E|Z| p ) 1/p ≤ K 2 √ p for all p ≥ 1; 3. Super-exponential moment: E exp(Z 2 /K 2 3 ) ≤ 2.A random variable Z satisfying any of the above assertions is called a sub-Gaussian random variable. We will denote by K 3 the sub-Gaussian norm.Consequently, a sub-Gaussian random variable satisfies the following equivalence of moments property. There exists an absolute constant c > 0 such that for any m ≥ 2,E|Z| m 1/m ≤ cK 3 √ m E|Z| 2 1/2 . Lemma 6. Assume that ∥u θ (X) -E[u θ (X)]∥ and ∥v θ (Y ) -E[v θ (Y )]∥ are sub-Gaussian with sub-Gaussian norm K. We set σ 2 θ (X) := Var(∥u θ (X) -E[u θ (X)]∥), σ 2 θ (Y ) := Var(∥v θ (Y ) - E[v θ (Y )]∥).Then there exists an absolute constant C > 0 such that for any δ ∈ (0, 1), it holds w.p.a.l. 1 -δ∥ E x [u θ ] -E[u θ (X)]∥ ≤ C √ n σ 2 θ (X) + K 2 n log(2δ -1 ).Similarly, w.p.a.l. 1 -δ∥ E y [v θ ] -E[v θ (Y )]∥ ≤ C √ n σ 2 θ (Y ) + K 2 n log(2δ -1 )Proof. Set Z := ∥u θ (X) -Eu θ (X)∥ and we recall that σ 2 θ (X) := Var(∥u θ (X) -E[u θ (X)]∥). We check that the moment condition,EZ m ≤ 1 2 m!L m-2 σ 2 θ (X) 2 , ∀m ≥ 2,for some constant L > 0 to be specified.The condition is obviously satisfied for m = 2. Next for any m ≥ 3, the Cauchy-Schwarz inequality and the equivalence of moment property giveEZ m ≤ EZ 2(m-2) 1/2 EZ 4 1/2 ≤ 4K 2 3 σ 2 θ (X) 2 EZ 2(m-2) 1/2 .Next, by homogeneity, rescaling Z to Z/K 1 we can assume that K 1 = 1 in Lemma 5. We recall that if Z is in addition non-negative random variable, then for every integer p ≥ 1, we haveEZ p = ∞ 0 P{Z ≥ t} pt p-1 dt ≤ ∞ 0 e 1-t 2 pt p-1 dt = ep 2 Γ p 2 .With p = 2(m -2), we get that EZ p ≤ e(m -2)Γ m -2 = e(m -2)! = em!/2. Using again Lemma 5, we can take L = cK for some large enough absolute constant c > 0. Then Proposition 1 gives the result.We now derive guarantees for the estimation of the conditional expectation and the conditional covariance for vector-valued output Y ∈ R dy .We start with a general result for arbitrary vector-valued functions of Y . We consider a vector-valued function h = (h 1 , . . . , h d ) whereh j ∈ L 2 ν (Y) for any j ∈ [d].We introduce the space of square integrable vector-valued functions[L 2 ν (Y, R d )] equipped with the norm ∥h∥ = j∈[d] ∥h j ∥ 2 L 2 ν (Y) .Next we can define the conditional expectation of h(Y ) = (h 1 (Y (1) ), . . . h d (Y (dy) )) ⊤ conditionally on X ∈ A as followsE[h(Y ) | X ∈ A] = E[h 1 (Y )] + ⟨1 A , D Y |X h 1 ⟩ P(X ∈ A) , . . . , E[h d (Y )] + ⟨1 A , D Y |X h d P(X ∈ A) ⊤ = E[h(Y )] + ⟨1 A , [1 d ⊗ D Y |X ]h⟩ P(X ∈ A) .We define similarly its empirical version asE θ [h(Y ) | X ∈ A] = E y [h 1 ] + ⟨1 A , D θ Y |X h 1 ⟩ E x [1 A ] , . . . , E y [h d ] + ⟨1 A , D θ Y |X h d ⟩ E x [1 A ] ⊤ = E y [h] + ⟨1 A , [1 d ⊗ D θ Y |X ]h⟩ E x [1 A ].Assuming that h(Y ) is sub-Gaussian, we setK := ∥∥h(Y ) -E[h(Y )]∥∥ ψ2 , σ 2 (h(Y )) := Var(∥h(Y ) -E[h(Y )]∥).Defineψ n (δ) := 1 √ n σ 2 (h(Y )) + K 2 n log(3δ -1 ) + ∥h∥ P(X ∈ A) ψ n (δ/3) + 2(1 + ψ n (δ/3))φ X (A)ϵ n (δ/3) .Theorem 4. Let the assumptions of Theorem 2 be satisfied. Assume in addition that h(Y ) is sub-Gaussian. Then we have w.p.a.l. 1 -δ thatEQUATIONProof. We haveEQUATION)Recall that ∥D Y |X ∥ ≤ 1, (23) and Lemma 3. Hence, a union bound we get with w.p.a.l. 1 -2δ thatEQUATION)We now handle the first term ∥E[h(Y )] -E y [h]∥. We recall that a similar quantity was already studied in Lemma 6. We can just replace u θ (X) by h(Y ) ∈ R d to get the result since we assumed that h(Y ) is sub-Gaussian. Hence there exists an absolute constant C > 0 such that w.p.a.l.1 -δ ∥E[h(Y )] -E y [h]∥ ≤ C √ n σ 2 (h(Y )) + K 2 n log(2δ -1 ).Actually, we can handle the conditional expectation E[Y | X ∈ A] in a more direct way. Setϵ n (δ) := log(δ -1 d y ) n log(δ -1 d y ) n .Corollary 2. Let the Assumptions of Theorem 2 be satisfied. Assume in addition that Y is a sub-Gaussian vector. Then for any δ ∈ (0, 1), it holds with probability at least 1 -δ thatEQUATION)Proof. The proof of this result is identical to that of Theorem 4 up to (33). Now if we specify h(Y ) = Y ∈ R dy . Then, applying Bernstein's inequality on each of the d y components of E[Y ] -Y n and a union bound, we get w.p.a.l.1 -δ ∥E[Y ] -Y n ∥ ≲ tr(Cov(Y )) log(δ -1 d y ) n + max j∈[dy] ∥Y (j) ∥ ψ2 log(δ -1 d y ) n .Using again Definition 4, we obtainmax j∈[dy] ∥Y (j) ∥ ψ2 ≲ ∥Cov(Y )∥ ≤ tr(Cov(Y )).It follows from the last two displays, w.p.a.l.1 -δ ∥E[Y ] -Y n ∥ ≲ tr(Cov(Y ))ϵ n (δ).(35) A union bound combining the previous display with (33) gives the first result.We focus now on the conditional covariance estimation problem. We first define the conditional covariance as follows: Cov(Y |X ∈ A) = Cov(Y ) + ⟨1 A , [(1 dy ⊗ 1 dy ) ⊗ D Y |X ]h ⊗ h⟩/P[X ∈ A] -⟨1 A , [1 dy ⊗ D Y |X ]h⟩ ⊗ ⟨1 A , [1 dy ⊗ D Y |X ]h⟩/(P[X ∈ A]) 2 . (36) Note that ⟨1 A , [(1 dy ⊗ 1 dy ) ⊗ D Y |X ]h ⊗ h⟩ = (⟨1 A , D Y |X h j h k ⟩) j,EQUATION)We define the effective of covariance matrix Cov(Y ) as follows:r(Cov(Y )) := tr(Cov(Y )) ∥Cov(Y )∥ .We set for any δ ∈ (0, 1)EQUATIONandψ (2) n (δ) = ϵ (2) n (δ) + [ψ n (δ/4) + 2(1 + ψ n (δ/4))φ X (A)ϵ n (δ/4)] (E[∥Y ∥ 2 ]) 2 P[X ∈ A] + ψ (1) n (δ/4) 2∥E[Y | X ∈ A]∥ + ψ (1) n (δ/4) .Corollary 3. Let the assumptions of Corollary 2 be satisfied. Then for any δ ∈ (0, 1), it holds with probability at least 1 -δ thatEQUATION)Proof. We use again the function h(Y ) = Y . We note in view of ( 36)-( 37) thatEQUATIONNext, we note that∥⟨1 A , (1 dy ⊗ 1 dy ) ⊗ D Y |X P[X ∈ A] - D θ Y |X E x [1 A ] h ⊗ h⟩∥ ≤ ∥⟨1 A , (1 dy ⊗ 1 dy ) ⊗ D Y |X P[X ∈ A] - D θ Y |X E x [1 A ] h ⊗ h⟩∥ HS ≤ P[X ∈ A]∥ D Y |X P[X ∈ A] - D θ Y |X E x [1 A ] ∥ j,k∈[dy] ∥Y j Y k ∥ L 2 ν (Y) ≲ P[X∈A] ∥ D Y |X -D θ Y |X P[X∈A] ∥+∥ D θ Y |X ∥ 1 P[X∈A] - 1 E x [1 A ] j,k∈[dy] ∥Y j Y k ∥ L 2 ν (Y)Remind that Y is a sub-Gaussian vector. Using the equivalence of moments property of sub-Gaussian vector, we get that∥Y j Y k ∥ L 2 ν (Y) ≤ E[Y 4 j ]E[Y 4 k ] ≲ E[Y 2 j ]E[Y 2 k ], ∀j, k ∈ [d y ].By a union bound combining the last two displays with (23) and Lemma 3, we get w.p.a.l. 1 -2δEQUATION) Next, we set u = E[h(Y ) | X ∈ A] and û = E θ [h(Y ) | X ∈ A]. Then we have ∥u ⊗ u -û ⊗ û∥ ≤ ∥u -û∥(∥u∥ + ∥û∥) ≤ ∥u -û∥(2∥u∥ + ∥û -u∥).We apply next Corollary 2 to get w.p.a.l.EQUATION)Next Koltchinskii and Lounici (2017, Theorem 4 ) guarantees that w.p.a.l 1 -δEQUATIONwhere ϵ(2) n (δ) is defined in (38). A union bound combining ( 40), ( 41), ( 42) and ( 43) gives the result.Experiments were conducted on a high-performance computing cluster equipped with an Intel(R) Xeon(R) Silver 4210 CPU @ 2.20GHz Sky Lake CPU, 377GB RAM, and an NVIDIA Tesla V100 16Gb GPU. Code is available at https://github.com/pietronvll/NCP.To evaluate our method's ability to estimate conditional densities, we tested NCP on six different data models (described in the following paragraph) and compared its performance with ten other methods (detailed in Tab. 2). We assessed the methods' performance using the KS distance between the estimated conditional CDF and the true CDF. Additionally, we explored how the performance of each method scales with the number of training samples, ranging from 10 2 to 10 5 , with a validation set of 10 3 samples. We tested each method on nineteen different conditional values uniformly sampled between the 5%-and 95%-percentile of p(x). Conditional CDFs were estimated on a grid of 1000 points uniformly distributed over the support of Y . The KS distance between each pair of CDFs was averaged over all the conditioning values. In Tab. 5, we present the mean performance (KS distance ± standard deviation), computed over 10 repetitions, each with a different random seed.Synthetic data models. We included the following synthetic datasets from Rothfuss et al. (2019) and Gao and Hastie (2022) into our benchmark:• LinearGaussian, a simple univariate linear density model defined as Y = X + N (0, 0.1)where X ∼ Unif(-1, 1).• EconDensity, an economically inspired heteroscedastic density model with a quadratic dependence on the conditional variable defined asY = X 2 + ϵ Y , ϵ Y ∼ N (0, 1 + X) where X ∼ |N (0, 1)|.• ArmaJump, a first-order autoregressive model with a jump component exhibiting negative skewness and excess kurtosis, defined asx t = [c(1 -α) + αx t-1 ] + (1 -z t )ϵ t + z t [-3c + 2ϵ t ] ,where ϵ t ∼ N (0, 0.05) and z t ∼ B(1, p) denote a Gaussian shock and a Bernoulli distributed jump indicator with probability p, respectively. The parameters were left at their default value.• GaussianMixture, a bivariate Gaussian mixture model with 5 kernels where the goal is to estimate the conditional density of one variable given the other. The mixture model is defined as p(X, Y ) =5 k=1 π k N (µ k , Σ k )where π k , µ k , and Σ k are the mixing coefficient, mean vector, and covariance matrix of the k-th distribution. All the parameters were randomly initialized.• SkewNormal, a univariate skew normal distribution defined as Y = 2ϕ(X)ψ(αX) where ϕ(•) and ψ(•) are the standard normal probability and cumulative density functions, and α is a parameter regulating the skewness. The parameters were left at their default value.• Locally Gaussian or Gaussian mixture distribution (LGGMD) (Gao and Hastie, 2022) , a regression dataset where the target y depends on the three first dimensions of x, with seventeen irrelevant features added to x. The features of x are all uniformly distributed between -1 and 1. The first dimension of x gives the mean of Y |X, the second is whether the data is Gaussian or a mixture of two Gaussians, and the third gives its asymmetry. More specifically:Y |X ∼    0.5N (0.25X (1) -0.5, 0.25(0.25X (3) + 0.5) 2 ) +0.5N (0.25X (1) + 0.5, 0.25(0.25X (3) -0.5) 2 ) if X (2) ≤ 0.2 0N (0.25X (1) -0.5, 0.3) if X (2) > 0.2 (44)To sample data from EconDensity, ArmaJump, GaussianMixture, and SkewNormal, we used the library Conditional Density Estimation (Rothfuss et al., 2019) available at https://github. com/freelunchtheorem/Conditional_Density_Estimation.Training NCP. We trained an NCP model with u θ and v θ as multi-layer perceptrons (MLPs), each having two hidden layers of 64 units using GELU activation function in between. The vector σ θ has a size of d = 100, and γ is set to 10 -3 . Optimization was performed over 10 4 epochs using the Adam optimizer with a learning rate of 10 -3 . Early stopping was applied based on the validation set with patience of 1000 epochs. To ensure the positiveness of the singular values, we transform the vector σ θ with the Gaussian function x → exp(-x 2 ) during any call of the forward method. Whitening was applied at the end of training.Compared methods. We compared our NCP network with ten different CDE methods. See Tab. 2 for the exhaustive list of models including a brief summary and key hyperparameters.In particular, the methods were set up as follows:• NF was characterized by a 1D Gaussian base distribution and two Masked Affine Autoregressive flows (Papamakarios et al., 2017) followed by a LU Linear permutation flow. To match the NCP architecture, each flow was defined by two hidden layers with 64 units each.The training procedure was the same as for the NCP model. The model was implemented using the library normflows (Stimper et al., 2023) .• DDPM was characterized by a U-Net (Ronneberger et al., 2015) , a noise schedule starting from 10 -4 to 0.02 and 400 steps of diffusion as implemented in https://github.com/ TeaPearce/Conditional_Diffusion_MNIST.• CKDE's kernels bandwidth was estimated according to Silverman's rule (Silverman, 1986) .• MDN's architecture was defined by two hidden layers with 64 units each and 20 Gaussians kernels.• KMN's architecture was defined by two hidden layers with 64 units each, 50 Gaussians kernels, and kernels bandwidth was estimated according to Silverman's rule (Silverman, 1986) .• LSCDE was defined by 500 components which bandwidths were set to 0.5 and kernels center found via a k-means procedure.• NNKDE's number of neighbors was set using the heuristics k = √ n (Devroye et al., 1996) . Kernels bandwidth was estimated according to Silverman's rule (Silverman, 1986) . We used the implementation available at https://github.com/lee-group-cmu/NNKCDE.• RFCDE was characterized by a Random Forest with 1000 trees and 31 cosine basis functions.The training was performed using the rfcde library available at https://github.com/ lee-group-cmu/RFCDE.• FC was trained using a Random Forest with 1000 trees as a regression method and had 31 cosine basis functions. The training was performed using the flexcode library available at https://github.com/lee-group-cmu/FlexCode.• LinCDE was trained with 1000 LinCDE trees using the LinCDE.boost R function from https://github.com/ZijunGao/LinCDE. CKDE, MDN, KMN, and LSCDE hyperparameters were set according to Rothfuss et al. (2019) and were trained using the library Conditional Density Estimation available at https://github.com/ freelunchtheorem/Conditional_Density_Estimation. All methods involving the training of a neural network were assigned the same number of epochs given to NCP. All other method parameters were set as prescribed in their paper.Results. See Tab. 4 for the comparison of performances for n = 10 4 . See also Fig. 5 . We also carried out an ablation study on centering and whitening post-treatment for NCP in Tab. 3The objective of this next experiment is to estimate a confidence interval at coverage level 90% for two distribution models with different properties (Laplace and Cauchy) and one real dataset in order to showcase the versatility of our NCP approach. Normalizing Flows (NF) (Rezende and Mohamed, 2015b) Generative models that transform a simple distribution into a complex one through a series of invertible and differentiable transformations • Architecture • Flow type Denoising Diffusion Probabilistic Model (DDPM) (Ho et al., 2020) Generative models that learn to generate data by reversing a gradual noising process, modeling distributions through iterative refinementConditional KDE (CKDE) (Li and Racine, 2006) Nonparametric approach modeling the joint and marginal probabilities via KDE and computes the conditional density as p(y|x) = p(x, y)/p(x).Mixture Density Network (MDN) (Bishop, 1994) Uses NeuralNets which takes conditional x as input and governs all the weights of a GMM modeling p(y|x).Kernel Mixture Network (KMN) (Ambrogioni et al., 2017) Similar to MDN with the difference that NN only controls the weights of the GMM.Least-Squares CDE (LSCDE) (Sugiyama et al., 2010) Computes the conditional density as linear combination of Gaussian kernels Compared methods. We compared our NCP procedure for building conditional confidence intervals to the state-of-the-art conditional conformal prediction method in Gibbs et al. (2023) . We also developed another method based on Normalizing Flows' estimation of the conditional CDE and we added it to the benchmark.Experiment for Laplace and Cauchy distributions. We generate a dataset where the X variable follows a uniform distribution on interval [0, 5] and Y |X = x follows either a Laplace distribution with location and scale parameters (µ(x), b(x)) = (x 2 , x) or a Cauchy distribution with location and scale parameters (µ(x), b(x)) = (x 2 , 1 + x). We create a train set of 50000 samples, a validation set of 1000 samples and a test set of 1000 samples.For the Laplace distribution, we train an NCP where u θ and v θ are multi-layer perceptrons with two hidden layers of 128 cells, σ θ is a vector of size d = 500 and γ = 10 -2 . Between each layer, we use the GELU activation function. We optimize over 5000 epochs using the Adam optimizer with a learning rate of 10 -3 . We apply early stopping with regard to the validation set with a patience of 100 epochs. Whitening is applied at the end of training. To fit the Cauchy distribution, we increase the depth of the MLPs to 5 and the width to 258.We compare this NCP network with two state-of-the-art methods. The first is a normalizing flow with base distribution a 1D Gaussian and two Autoregressive Rational Quadratic spline flows (Durkan et al., 2019) followed by a LU Linear permutation flow. All flows come from the library normflows (Stimper et al., 2023) . The spline flows have each two blocks of 128 hidden units to match the NCP architecture. The normalizing flow is allowed the same number of epochs as ours with the same optimizer. The second model is the conditional conformal predictor from Gibbs et al. (2023) . This model needs a regressor as an input. We consider a situation favorable to Gibbs et al. (2023) as we assume as prior knowledge that the true conditional expectation is a polynomial function (the truth is actually the quadratic function in this example). Therefore we chose a linear regression with polynomial features as in Gibbs et al. (2023) as this regressor should fit the data without any problem. For all other choices of parameters, we follow the prescriptions of Gibbs et al. (2023) . For the sake of fairness, we note that the validation set used for early stopping in NF and NCP was also used as a calibration set for the CCP method.By design, the Conditional Conformal Predictor (CCP) gives the confidence interval directly. However NCP and NF output the conditional distribution. To find the smallest confidence interval with desired coverage, we apply the linear search algorithm described in Algorithm 3 on the discretized conditional CDFs provided by NCP and NF. The results are provided in Fig. 1 . First, observe that although the linear regression achieves the best estimation of the conditional mean, as should be expected since the model is well-specified in this case, the confidence intervals, however, are unreliable for most of the considered conditioning. We also notice instability for NF and CCP for conditioning in the neighborhood of x = 0 with NF confidence region exploding at x = 0. We expect this behavior is due to the fact that the conditional distribution at x = 0 is degenerate. Comparatively, NCP does not exhibit such instability around x = 0. It only tends to overestimate the produced confidence region for conditioning close to x = 0. Experiment on real data. We also evaluate the performance of NCP in estimating confidence intervals using the Student Performance dataset available at https://www.kaggle.com/datasets/ nikhil7280/student-performance-multiple-linear-regression/data. This dataset comprises 10000 records, each defined by five predictors: hours studied, previous scores, extracurricular activities, sleep hours, and sample question papers practiced, with a performance index as the target variable. In this experiment, the NCP's u θ and v θ are defined by MLPs with two hidden layers, each containing 32 units and using GELU activation functions, σ θ is a vector of size d = 50 and γ = 10 -2 . Optimization was performed over 50000 epochs using the Adam optimizer with a learning rate of 10 -3 . We compare NCP with a normalizing flow defined as above in which spline flows have each two blocks of 32 hidden units to match NCP architecture. The normalizing flow is trained for the same number of epochs as our model, using the same optimizer. We further compare NCP with a split conformal predictor featuring a Random Forest regressor (RFSCP) with 100 estimators. We used the implementation of the library puncc (Mendil et al., 2023) . For NCP and the normalizing flow, early stopping is based on the validation set, while for RFSCP, the validation set serves as the calibration set. We performed 10 repetitions, randomly splitting the dataset into a training set of 8000 samples and validation and test sets of 1000 samples each. We report the results of the estimated confidence interval at a coverage level of 90% in Tab. 5. The methods provide fairly good coverage. NF did not respect the 90% coverage condition. Only NCP and RFSCP both respect the coverage condition but the width of the confidence intervals for RFSCP are larger than for NCP.Discussion on Conformal Prediction. Conformal prediction (CP) is a popular model-agnostic framework for uncertainty quantification approach Vovk et al. (1999) . CP assigns nonconformity scores to new data points. These scores reflect how well each point aligns with the model's predictions. CP then uses these scores to construct a prediction region that guarantees the true outcome will fall within it with a user-specified confidence parameter. However, CP is not without limitations. The construction of these guaranteed prediction regions can be computationally expensive especially for large datasets, and need to be recomputed from scratch for each value of the confidence level parameter.In addition, the produced CP confidence regions tend to be conservative. Another limitation of regular CP is that predictions are made based on the entire input space without considering potential dependencies between variables. Conditional conformal prediction (CCP) was later developed to handle conditional dependencies between variables, allowing in principle for more accurate and reliable predictions Gibbs et al. (2023) . CCP suffers from the typical limitations of regular CP and the theoretical guarantees.Experiment on high-dimensional synthetic data. In Fig. 3 , we trained NCP for d = 100 using the same MLP architecture and the same NF with autoregressive flow as in our initial experiments based on n = 10 5 samples {(X i , Y i )} n i=1 with values in R d × Y. We plot the conditional CDF for several conditioning w.r.t. θ(x) on 10 repetitions. NCP paired with a small MLP architecture performs comparably to the NF model for Gaussian distributions. For discrete distributions, the NCP demonstrates superior performance compared to the NF model.We repeated the experiment in Fig. 3 for d ∈ {100, 200, 500, 1000} and recorded the average Kolmogorov-Smirnov (KS) distance of the NCP conditional distribution to the truth, computation time and their standard deviations over 10 repetitions. High-dimensional experiment in molecular dynamics: Chignolin folding. We investigated the dynamics of Chignolin folding, using a molecular dynamics simulation lasting 106µs and sampled every 200ps, resulting in 524, 743 data points. Our analysis focuses on 39 heavy atoms (nodes) with a cutoff radius of 5 Angstroms. To predict the conditional transition probability between metastable states, we integrate our NCP approach with a graph neural network (GNN) model. GNNs, as demonstrated by Chanussot et al. (2021) , represent the state-of-the-art in modeling atomistic systems, adeptly incorporating the roto-translational and permutational symmetries inherent in physical systems. In particular, we employed a SchNet model Schütt et al. (2019 Schütt et al. ( , 2023) ) with three interaction blocks. Each block features a 64-dimensional latent atomic environment, and the inter-atomic distances for message passing are expanded over 20 radial basis functions. After the final interaction block, each latent atomic environment is processed through a linear layer and then aggregated by averaging. The model underwent training for 100 epochs using an Adam optimizer with a learning rate of 10 -3 . We employed a batch size of 256 and set γ to 10 -3 . In Fig. 2 , we show how our NCP approach enables the tracking transitions between metastable states, demonstrating accurate forecasting and strong uncertainty quantification.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We provided the information at the beginning of Appendix C.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 


Decision: Intuitively, the idea of this paper is very simple—essentially: Learn neural networks that transform x and y to some feature space, so that the joint density is the inner-product of the feature spaces. This has an obvious problem, namely that the result may not be normalized (or even positive). Seemingly for this reason (though the issue is never really acknowledged in the paper) a quadratic loss is used that essentially measures the squared difference between the learned "density" (not guaranteed to be normalized or positive) and the target. Then, an extra regularizer is introduced that encourages these feature spaces to be orthogonal. A high-probability bound result is given.

Reviewers were uniformly in favor of accepting this paper. The method appears to be novel and the results correct. But I do have some concerns that the presentation greatly obscures the simplicity of the underlying method. It seems perverse to not first describe the method in simple terms before giving the measure-theoretic motivation. (Which, after all, is just a motivation.) The theory makes quite strong assumptions (as is basically necessary given the presence of neural networks inside the method. As a service to the community, I urge the authors to be more clear about the simplicity of what is being proposed, and not to "hide" the simplicity of what is actually being proposed behind so much measure theory.

Reviews: review: paper_summary: paper_summary: The paper proposes Neural Conditional Probability, a novel operator-theoretic approach to learning conditional probability distributions by learning parameters of the truncated SVD of the conditional expectation operator with a neural network. The authors provide a rigorous mathematical derivation and argue for statistical guarantees of their method. The empirical evaluations require major improvements to an otherwise solid paper.

**As a general note:** I do not consider myself an expert on the theoretical aspects of learning theory. My background is in Bayesian deep learning and simulation-based Bayesian inference. As such, my confidence regarding sections 3 and 5 is rather low, and my review shall mainly consult on the remaining sections that focus on presentation, embedding into other literature, and empirical evaluations.
strength_weakness: strength_weakness: ## Major
- The Related Work section does a good job of acknowledging related works that aim to learn conditional distributions. However, it utterly fails to embed the current paper into this research landscape. I recommend the authors elaborate on the precise similarities and differences between the referenced papers and their methods in the rebuttal period.
- The empirical evaluations are limited to low-dimensional toy problems. This is a stark contrast to the introduction of the method, where the authors repeatedly list the curse of dimensionality as a drawback of other methods. While I acknowledge that the paper is situated in the area of operator learning and ML theory, the quality standard of NeurIPS is not met by the authors’ experiments. This weak evaluation does not do the remainder of the paper justice and I strongly recommend the authors overhaul the experiments to feature high-dimensional tasks that cannot be solved with other state-of-the-art methods. This constitutes a major revision, and this is the main reason why I cannot recommend acceptance to NeurIPS 2024.


## Minor

- The empirical evaluation is missing some important information for real-world applications: What are the approximate wall-clock times for (1) training and (2) inference of the competing methods? Further, the authors mention the large required training set size, which might also influence the practically expected training duration in real-world tasks.
- Please fix the citations throughout the manuscript: Most citations are ‘text citations’ even if their embedding in the sentence warrants parenthesized citations (Author, 1976).
- This is just a personal preference, no need to address it: The ‘paper organization’ paragraph at the end of the introduction does not add value and the space could be used more efficiently elsewhere in the manuscript.
- The first sentence in the conclusion is incomplete.
questions: questions: - As per your answer to checklist item 5 (Open access to data and code), I would like to request access to the full and reproducible code of the empirical evaluations.
- Since you want to compare your method with state-of-the-art conditional density estimation methods: Why don't you benchmark against conditional flow matching?
limitations: limitations: - The authors mention limitations throughout the manuscript, which I appreciate. However, I would recommend adding a dedicated **Limitations** section in the conclusion to give a compact overview for readers who don’t engage with the entire paper in detail.
- The performance of NCP in high dimensions might be a limitation, but the authors do not study this crucial setting.
paper_summary: paper_summary: The authors propose a method (Neural Conditional Probability, NCP) for learning a conditional distribution P(Y | X) from a finite sample from a distribution. The method is based on following observations: (1) it is sufficient to learn the conditional expectation operator E_{Y | X}[f](x) = E[f(Y) | X = x]; (2) the conditional expectation operator can be written as (an infinite) SVD decomposition which could be truncated at some point, so the problem reduced to learning the finite number of functions in the SVD decomposition; (3) the joint distribution density can be written using the functions from the SVD decomposition of the conditional expectation operator, which gives an optimisation objective for fitting the functions from the SVD decomposition using the sample from a joint distribution. The authors provide an extensive theoretical analysis of the proposed method as well as a simulation study on a few synthetic datasets.
strength_weakness: strength_weakness: - Limited experiments restricted to synthetic data making it difficult to judge the potential applicability of this method
- It would be nice to have a short summary on the main properties of operators, their SVD decompositions, etc. I could generally follow the presentation without major problems, but having a such an operators summary would have made it easier to read the paper
questions: questions: - I am wondering about the choice of the specific loss function in Eq. (6). Could, for example, the log-likelihood potentially be used here? If so, what are the advantages of using the Eq. (6) instead of log-likelihood?
- The loss function in Eq. (9) is roughly speaking a regularisation term enforcing that the singular functions are orthonormal, is it correct? Could it be possible to build the neural nets with such properties by construction rather than enforce it by regularisation?
- What do you think about the scalability of the model to more complex datasets than in Section 6? For example, conditional image generation. Do you expect issues applying NCP in such cases?
limitations: limitations: The limitations are sufficiently addressed (NeurIPS Paper Checklist)
paper_summary: paper_summary: I am not qualified to review this paper
strength_weakness: strength_weakness: I am not qualified to review this paper
questions: questions: I am not qualified to review this paper
limitations: limitations: I am not qualified to review this paper
paper_summary: paper_summary: This paper proposes Neural Conditional Probability (NCP), a novel operator-theoretic approach for learning conditional probability distributions. Extensive theoretical results are provided to support the optimization consistency and statistical accuracy of NCP. NCP can be used to extract conditional density and compute statistical measures such as conditional mean, variance, moments and CDF once it is trained. Experiments on a collection of conditional density estimation datasets are conducted to highlight the efficacy of NCP.
strength_weakness: strength_weakness: - The proposed NCP method is not clearly motivated or introduced. In Line 49-50, the authors mention that NCP does not belong to any of the four aforementioned approaches. But how is NCP in contrast with them and in what aspects does NCP make improvements? I believe adding some intuitive explanations accompanying theoretical analysis would help improve the readability.
- Some key concepts or methods are not clearly explained, which makes it hard to understand the contributions of this work. For example, why is learning *conditional expectation operator* considered useful? Are there any baseline methods that also learn expectation operators?
questions: questions: Please see Weaknesses.
limitations: limitations: The authors have discussed the limitations of their work.

Comments: title: Acknowledgement to the reviewer
Acknowledgement to the reviewer
Acknowledgement to the reviewer
Code now available
comment: We would like to thank the reviewer for their comments, which inspired us to make additional steps and improve our work. We are happy that our rebuttal was helpful, and we commit to incorporate it in the revised manuscript.
We would like to thank the reviewer for their comments, which inspired us to make additional steps and improve our work. We are happy that our rebuttal was helpful, and we commit to incorporate it in the revised manuscript.
We would like to thank the reviewer for their comments, which inspired us to make additional steps and improve our work. We are happy that our rebuttal was helpful, and we commit to incorporate it in the revised manuscript.

Concerning the reference, we thank reviewer for bringing this paper to our attention, we will definitely include it in the revision.
Thank you for the detailed response. I am happy to keep my score.
Thank you very much for a detailed reply! I confirm my positive view of this paper.
Hello reviewers,

There has been some discussion of wanting to examine the code for this paper. The authors (following the official guidelines) have sent me a link to a repo containing the code. I have checked, to the best of my ability, that this is anonymous, so I am passing on the link to you here: https://anonymous.4open.science/r/NCP/README.md
Thank you for the detailed rebuttal which addresses my main concerns regarding the empirical evaluations. I acknowledge that the theoretical aspects of the paper are the main contribution. I will raise my score accordingly.


## Flow matching and UQ

> From our understanding, conditional flow matching is designed to efficiently train continuous normalizing flows [...] However, based on our review of the relevant literature, this method has not been explored for the purpose of UQ.

Flow matching has been studied for simulation-based inference, which is concerned with the very question of uncertainty quantification. See reference [1] which also contains a code repository.

[1] Dax et al. (2023). Flow Matching for Scalable Simulation-Based Inference. NeurIPS 2023. https://arxiv.org/abs/2305.17161