Learning Robust Neural Processes with Risk-Averse Stochastic Optimization

Huafeng Liu, Yiran Fu, Liping Jing, Hui Li, Shuyang Lin, Jingyue Shi, Deqiang Ouyang, Jian Yu

Abstract: Neural processes (NPs) are a promising paradigm to enable skill transfer learning across tasks with the aid of the distribution of functions. The previous NPs employ the empirical risk minimization principle in optimization. However, the fast adaption ability to different tasks can vary widely, and the worst fast adaptation can be catastrophic in risk-sensitive tasks. To achieve robust neural processes modeling, we consider the problem of training models in a risk-averse manner, which can control the worst fast adaption cases at a certain probabilistic level. By transferring the risk minimization problem to a two-level finite sum minimax optimization problem, we can easily solve it via a double-looped stochastic mirror prox algorithm with a task-aware variance reduction mechanism via sampling samples across all tasks. The mirror prox technique ensures better handling of complex constraint sets and non-Euclidean geometries, making the optimization adaptable to various tasks. The final solution, by aggregating prox points with the adaptive learning rates, enables a stable and high-quality output. The proposed learning strategy can work with various NPs flexibly and achieves less biased approximation with a theoretical guarantee. To illustrate the superiority of the proposed model, we perform experiments on both synthetic and real-world data, and the results demonstrate that our approach not only helps to achieve more accurate performance but also improves model robustness.

Lay Summary:

How can AI systems quickly learn new skills from just a few examples reliably, especially in high-stakes scenarios like healthcare or robotics? Current methods (called neural processes) are powerful but can fail unpredictably during rapid adaptation—with potentially serious consequences.

We tackle this by redesigning training to prioritize "safety nets." Instead of aiming for average performance, we focus on minimizing risks in worst-case scenarios. Our method uses advanced math (a ​mirror-prox algorithm) to balance fast adaptation and stability across diverse tasks. By smartly sharing insights from all tasks during training, we ensure AI adapts more robustly without compromising accuracy.

Tests on synthetic and real-world data prove our approach not only boosts performance but also builds resilience against unexpected challenges—theoretically guaranteed. This could help deploy trustworthy AI in risk-sensitive fields where failures are unacceptable.

Primary Area: Deep Learning->Generative Models and Autoencoders

Keywords: neural process, robust learning

Submission Number: 11914