Learning Robust Neural Processes with Risk-Averse Stochastic Optimization

Thanks for your positive and valuable feedback. We have made efforts to address your concerns. If there are any further questions, please let us know and we will reply promptly.

Q1. The paper uses simple regret as the primary performance measure, but a more robust metric like the CVaR of simple regret may be more suitable given the risk-averse theme.

Reply: We agree that CVaR-based metrics align directly with our emphasis on robustness. We initially chose simple regret for consistency with prior NP-based Bayesian optimization work, which commonly measures average performance. In the final version, we plan to include CVaR of simple regret in our Bayesian optimization experiments. Preliminary results suggest that our risk-averse training indeed lowers the worst-case tail of the regret distribution, indicative of stronger robustness. We will highlight these findings to reinforce that the proposed method not only reduces average regret but also mitigates high-regret outliers.

Q2. “Perhaps give some reference on the equivalence between Eq. (3) and Eq. (4).” “… equivalence between Eq. (3) and Eq. (6).”**

Reply: Thank you for noting this. We do rely on standard transformations from the CVaR’s probability-constrained problem to its slack-variable-based and DRO-based forms. We will add more explicit citations (e.g., Rockafellar et al. (2000) for CVaR expansions and Shapiro et al. (2009) for standard references on DRO formulations) and a short annotated derivation in the supplement clarifying each step.

[Rockafellar et al. (2000)]: Rockafellar, R. T., Uryasev, S., et al. Optimization of conditional value-at-risk. Journal of risk, 2:21–42, 2000.

[Shapiro et al. (2009)]: A. Shapiro, D. Dentcheva, and A. Ruszczy´ nski. Lectures on stochastic programming: modeling and theory. SIAM, 2009.

Q3. Typographical Error in “fundom” (Line 115, right column)There is a typo: “fundom.”

Reply: We will correct “fundom” to “random” or “function domain” (whichever was intended in context). Thank you for pointing it out.

Q4. Is there a theoretical way to show how our proposed approach improves the robustness of neural processes?

Reply: As shown in Equations (3)–(6), our training reweights high-loss (or high-regret) tasks so they carry larger gradients, effectively controlling the model’s tail behavior. This is theoretically grounded in distributionally robust optimization. A formal proof of how and by how much tail performance improves remains challenging, especially in a non-convex NP setting. Nonetheless, the link between CVaR and tail distribution improvements is well-known in risk-averse optimization theory (e.g., Rockafellar et al., 2000). We will add references and clarifications that highlight how increasing the confidence level α in CVaR training can reduce the worst-case outcomes.

Q5. The robust training process also improves average performance (Figure 3). Please discuss.

Reply: Indeed, empirically, we observe that focusing on high-risk tasks can prevent overfitting to easy samples and help the model learn more stable, generalizable parameters. That can yield broader improvements in average performance as well. A moderate amount of noise or risk weighting seems beneficial for both extremes (the tail and the average) by flattening observed losses. We will stress this empirical phenomenon more in Section 5.

Lastly, thank you once again for your valuable comments.

Thanks for your positive and valuable feedback. We have made efforts to address your concerns. If there are any further questions, please let us know and we will reply promptly.

Q1. The contribution seems incremental.

Reply: While CVaR optimization is well-established (e.g., in [Curi et al., 2020]), we believe our contribution lies in adapting it specifically to the Neural Process (NP) framework. This is non-trivial because tasks in meta-learning differ greatly in difficulty and distribution, making CVaR-based reweighting particularly relevant to “worst-case adaptation.”

We provide a finite-sum DRO formulation for task-level NP optimization and find that standard SGD-based solutions for CVaR can suffer from exploding/vanishing gradients. The double-loop mirror prox with variance reduction addresses these problems, stabilizing updates while achieving robust generalization.

As with many deep-learning contexts, providing rigorous global convergence guarantees in non-convex scenarios is challenging. We plan to investigate approximate convergence or local convergence guarantees in a future revision, building upon the convex analyses presented in Appendix B.

[Curi et al., 2020] Curi, S., Levy, K. Y., Jegelka, S., and Krause, A. Adaptive sampling for stochastic risk-averse learning. Advances in Neural Information Processing Systems, 33:1036–1047, 2020.

Q2. (a) What is the sample complexity and minimum number of iterations required to achieve a certain level of accuracy?(b) In Algorithm 1 (Line 5), how is the maximum number of inner loop iterations Ks determined?

Reply: We follow standard distributionally robust optimization settings, where each task can be large, and the meta-dataset comprises many tasks. In each iteration, we typically sample at least one data point per task (one shot per task). Detailed formal sample complexity results exist for convex mirror-prox setups; translating them to non-convex NPs remains an open question.

In the paper, $S$ denotes the number of outer loops (epochs), and $K_s$ represents the number of inner-loop iterations within each outer-loop $s$. In experiments, we set $K_s$ as a constant across all outer loops, i.e., $K_s = K$ for all $s$. $K$ is chosen to scale with the average number of samples per task. Empirically, we find that $K_s = 100-500$ often suffice for stable training in NP tasks; deeper loops help slightly but yield diminishing returns relative to the added cost. For outer loop epochs $S$, usually between 50 to 500 epochs. For simple tasks (e.g., low-dimensional regression, small datasets), $S$ is 50–200 epochs. For complex tasks (e.g., image completion), $S$ often requires 200–500+ epochs. We will expand them to also show the effect of varying the number of inner/outer loop iterations in a revised version.

Q3. Theoretical Questions (Appendix B)

Reply: For Proposition B.1, we will clarify that H refers to a set of finite cardinality or finite VC dimension to avoid ambiguity.

For Theorem B.5, the $a$ in the norm $|\cdot|_{\theta a,*}$ is a typo, it should be $|\cdot|_{\theta,*}$.

For Line 740 (Eq.(26)), the second term should be $2(\sum_{m=1}^M|q_m^+ - q_m^-|G)^2$. For the second term in line 737, we also omitted the square. Thank you again for pointing out the error.

For the left-hand side of Equation (27), adding or omitting the absolute value or norm is acceptable.

For Equation (28), we will correct $||\theta^+ - \theta^-||_{\theta}$ to be squared, i.e., $||\theta^+ - \theta^-||_{\theta}^2$

Q4. Clarifications and Notation Issues

Reply: Line 127 (Left Column): Garnelo et al. (2018a) assume the target set $T$ is disjoint from the context set $C$ in standard CNP. In our exposition, $T$ could overlap with $C$ or not. The choice is optional and widely used in BNP, DIVNP, TNP, SNP, and we will clarify that it differs slightly from the original CNP assumptions. Lines 220–221: In the expression $max_{\theta} \psi_{\theta}(\theta)−min_{\theta} \psi_{\theta}(\theta) ≤ D_{\theta}^2$, we will revise the notation to separate the parameters inside $\psi(\cdot)$ from the outer $\theta$ symbol, ensuring clarity.

Q5. Lines 436–437 suggest robust solutions shift the risk distribution left (Figure 4), but in Figure 4(b), RANP still has a rightmost bar. Could you clarify?

Reply: RANP does not eliminate all high-risk tasks but significantly reduces their proportion. We will refine the text to specify “we shift the distribution, decreasing the fraction of tasks with high risk” rather than implying complete removal. We can also adjust Figure 4’s bin widths or annotate it more explicitly to highlight the reduction percentage.

Q6. Typographical Errors

Reply: We have thoroughly revised the spelling errors, typographical mistakes, and symbol ambiguities in the text.

Lastly, thank you once again for your valuable comments.

Thanks for your positive and valuable feedback. We have made efforts to address your concerns. If there are any further questions, please let us know, and we will reply promptly.

Q1. Ablation experiments on CVaR α, inner/outer loop iteration lengths, and learning rate.

Reply: We appreciate the need for more thorough ablations. Due to space limits, we included a subset of these in Appendix E.3, focusing on α. Empirically, we have observed that α values in a moderate range (e.g., 0.4–0.6) often strike a good balance between focusing on worst-case tasks and not overly biasing the model toward rare outliers. We show an ablation (Figure 8) demonstrating that too small an α (e.g., 0.1) under-treats high-risk tasks, while an extreme α (e.g., >0.8) can degrade average performance. These observations align intuitively with the idea that CVaR reweights the “tail” portion of tasks.

In the paper, $S$ denotes the number of outer loops (epochs), and $K_s$ represents the number of inner-loop iterations within each outer-loop $s$. In experiments, we set $K_s$ as a constant across all outer loops, i.e., $K_s = K$ for all $s$. $K$ is chosen to scale with the average number of samples per task. Empirically, we find that $K_s = 100-500$ often suffice for stable training in NP tasks; deeper loops help slightly but yield diminishing returns relative to the added cost. For outer loop epochs $S$, usually between 50 to 500 epochs. For simple tasks (e.g., low-dimensional regression, small datasets), $S$ is 50–200 epochs. For complex tasks (e.g., image completion), $S$ often requires 200–500+ epochs. We will expand them to also show the effect of varying the number of inner/outer loop iterations in a revised version.

We used Adam optimizer with an initial learning rate $5·10^{-4}$ and decayed the learning rate using a cosine annealing scheme.

Q2. How does the double-loop optimization compare to simpler baselines in computational overhead?

Reply: The additional overhead mainly comes from (i) maintaining snapshots of gradients and (ii) sampling from each task in each inner iteration. However, because we adopt a task-aware sampling scheme (i.e., one sample per task), the minibatch size remains comparable to that of standard NP training. In practice, we find the extra overhead to be modest—roughly 1.1–1.3× the per-epoch training time compared to standard NPs, depending on the number of tasks and the length of inner loops. We will add more quantitative details in the final version, including a time-per-epoch comparison and possible strategies (e.g., smaller inner loops) to reduce total training time.

Q3. Why does improving tail performance not degrade average accuracy?

Reply: Elevated attention to difficult tasks prevents overfitting to “easy” samples and encourages more robust parameter updates. This can reduce overall variance in predictions and in practice, often benefits average-case accuracy. Our experiments (Tables 1, 2, 3) demonstrate that while we significantly reduce high-risk failure rates, the average performance (e.g., log-likelihood) either remains on par or improves slightly compared to standard ERM-based NPs. CVaR reweights tasks in proportion to their risk. Importantly, for moderately chosen α, we are not only training on the worst tasks but rather adjusting the objective so that the upper tail influences the gradient more strongly. The result is a more balanced training signal across all tasks.

Q4. The paper largely focuses on convex analysis, but deep NPs are inherently non-convex.

Reply: We agree that in non-convex regimes, strong theoretical convergence guarantees are elusive (an issue for many deep learning algorithms). While mirror-prox methods do enjoy strong convergence rates for convex-concave saddle-point problems, we currently rely on their practical effectiveness for non-convex NPs. We have followed prior work in robust optimization, which often leverages the same theoretical tools in non-convex settings with the understanding that local minima or stable solutions can often suffice in practice. In the final version, we will highlight this limitation more explicitly.

Lastly, thank you once again for your valuable comments.