Wasserstein Neural Processes
This addresses a limitation in probabilistic modeling for function learning, offering a more robust method for researchers and practitioners in machine learning, though it appears incremental as it builds directly on existing NP frameworks.
The paper tackles the problem of Neural Processes (NPs) failing to learn reasonable distributions for certain classes of problems when trained with maximum likelihood and KL divergence regularization, and shows that using Wasserstein distance approximations solves this issue, enabling approximation of a new class of function mappings while maintaining all benefits of traditional NPs.
Neural Processes (NPs) are a class of models that learn a mapping from a context set of input-output pairs to a distribution over functions. They are traditionally trained using maximum likelihood with a KL divergence regularization term. We show that there are desirable classes of problems where NPs, with this loss, fail to learn any reasonable distribution. We also show that this drawback is solved by using approximations of Wasserstein distance which calculates optimal transport distances even for distributions of disjoint support. We give experimental justification for our method and demonstrate performance. These Wasserstein Neural Processes (WNPs) maintain all of the benefits of traditional NPs while being able to approximate a new class of function mappings.