[Paper Review] End-to-End Meta-Bayesian Optimisation with Transformer Neural Processes

Alexandre Maraval, Matthieu Zimmer+2 more|arXiv (Cornell University)|May 25, 2023

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes the first end-to-end differentiable meta-Bayesian optimisation framework using a transformer-based neural process to jointly learn surrogate models and acquisition functions. By combining reinforcement learning with an auxiliary probabilistic inductive bias, the method achieves state-of-the-art regret performance across hyperparameter tuning, antibody design, and electronic design automation tasks.

ABSTRACT

Meta-Bayesian optimisation (meta-BO) aims to improve the sample efficiency of Bayesian optimisation by leveraging data from related tasks. While previous methods successfully meta-learn either a surrogate model or an acquisition function independently, joint training of both components remains an open challenge. This paper proposes the first end-to-end differentiable meta-BO framework that generalises neural processes to learn acquisition functions via transformer architectures. We enable this end-to-end framework with reinforcement learning (RL) to tackle the lack of labelled acquisition data. Early on, we notice that training transformer-based neural processes from scratch with RL is challenging due to insufficient supervision, especially when rewards are sparse. We formalise this claim with a combinatorial analysis showing that the widely used notion of regret as a reward signal exhibits a logarithmic sparsity pattern in trajectory lengths. To tackle this problem, we augment the RL objective with an auxiliary task that guides part of the architecture to learn a valid probabilistic model as an inductive bias. We demonstrate that our method achieves state-of-the-art regret results against various baselines in experiments on standard hyperparameter optimisation tasks and also outperforms others in the real-world problems of mixed-integer programming tuning, antibody design, and logic synthesis for electronic design automation.

Motivation & Objective

• To address the limitation of existing meta-Bayesian optimisation methods that treat surrogate modeling and acquisition function learning as independent steps.
• To enable joint, end-to-end training of both components in meta-BO using a differentiable architecture.
• To overcome the challenge of sparse rewards in reinforcement learning when training acquisition functions from scratch.
• To improve sample efficiency in new optimisation tasks by leveraging knowledge from related source tasks.
• To develop a unified, differentiable framework that directly predicts acquisition function values from observed data without relying on Gaussian processes.

Proposed method

• Proposes a transformer-based neural process architecture to model the acquisition function directly, enabling end-to-end differentiable meta-optimisation.
• Employs reinforcement learning to train the acquisition function policy, using cumulative regret as the reward signal.
• Introduces an auxiliary loss that guides a subnetwork to learn a valid probabilistic model, acting as an inductive bias to stabilize training.
• Uses combinatorial analysis to formalise the logarithmic sparsity of regret-based rewards, justifying the need for auxiliary supervision.
• Trains the full model end-to-end using policy gradient methods, with the acquisition function outputting action probabilities for next evaluation points.
• Leverages the structural inductive bias of neural processes to generalise across tasks while maintaining differentiability for joint optimisation.

Experimental results

Research questions

• RQ1 Can a unified, end-to-end differentiable framework jointly optimise surrogate modeling and acquisition function learning in meta-Bayesian optimisation?
• RQ2 How does the sparsity of regret-based rewards impact the training stability of transformer-based acquisition functions in reinforcement learning?
• RQ3 Can an auxiliary probabilistic inductive bias improve training stability and performance in end-to-end meta-BO with sparse rewards?
• RQ4 Does joint learning of surrogate and acquisition function components lead to better sample efficiency than decoupled approaches?
• RQ5 Can the proposed method generalise effectively across diverse real-world optimisation tasks such as hyperparameter tuning, antibody design, and logic synthesis?

Key findings

• The proposed end-to-end meta-BO framework achieves state-of-the-art regret performance on standard hyperparameter optimisation benchmarks.
• The method outperforms existing baselines in real-world applications, including mixed-integer programming tuning and electronic design automation.
• The inclusion of an auxiliary probabilistic inductive bias significantly stabilises training and improves sample efficiency in the presence of sparse rewards.
• Empirical results demonstrate that joint learning of surrogate and acquisition function components leads to lower cumulative regret than decoupled approaches.
• The framework generalises effectively across diverse tasks, including combinatorial sequence optimisation in antibody design and chip design.
• Combinatorial analysis confirms that regret-based rewards exhibit logarithmic sparsity, justifying the need for auxiliary supervision in RL-based meta-BO.