Composite Bayesian Optimization in function spaces using NEON—Neural Epistemic Operator Networks

Operator learning is a rising field of scientific computing where inputs or outputs of a machine learning model are functions defined in infinite-dimensional spaces. In this paper, we introduce Neon (Neural Epistemic Operator Networks), an architecture for generating predictions with uncertainty using a single operator network backbone, which presents orders of magnitude less trainable parameters than deep ensembles of comparable performance. We showcase the utility of this method for sequential decision-making by examining the problem of composite Bayesian Optimization (BO), where we aim to optimize a function <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mrow><mml:mi>f</mml:mi> <mml:mo>=</mml:mo> <mml:mi>g</mml:mi> <mml:mo>∘</mml:mo> <mml:mi>h</mml:mi></mml:mrow> </mml:math> , where <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mrow><mml:mi>h</mml:mi> <mml:mo>:</mml:mo> <mml:mi>X</mml:mi> <mml:mo>→</mml:mo> <mml:mi>C</mml:mi> <mml:mo>(</mml:mo> <mml:mi>Y</mml:mi> <mml:mo>,</mml:mo> <mml:msup><mml:mrow><mml:mi>R</mml:mi></mml:mrow> <mml:msub><mml:mi>d</mml:mi> <mml:mi>s</mml:mi></mml:msub> </mml:msup> <mml:mo>)</mml:mo></mml:mrow> </mml:math> is an unknown map which outputs elements of a function space, and <mml:math xmlns:mml="http://www.w3.org/1998/Math/MathML"><mml:mrow><mml:mi>g</mml:mi> <mml:mo>:</mml:mo> <mml:mi>C</mml:mi> <mml:mo>(</mml:mo> <mml:mi>Y</mml:mi> <mml:mo>,</mml:mo> <mml:msup><mml:mrow><mml:mi>R</mml:mi></mml:mrow> <mml:msub><mml:mi>d</mml:mi> <mml:mi>s</mml:mi></mml:msub> </mml:msup> <mml:mo>)</mml:mo> <mml:mo>→</mml:mo> <mml:mi>R</mml:mi></mml:mrow> </mml:math> is a known and cheap-to-compute functional. By comparing our approach to other state-of-the-art methods on toy and real world scenarios, we demonstrate that Neon achieves state-of-the-art performance while requiring orders of magnitude less trainable parameters.