SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexandre Maraval*
Huawei Noah's Ark Lab alexandre.maraval1@huawei.com
Matthieu Zimmer*
Huawei Noah's Ark Lab  
matthieu.zimmer@huawei.com
Antoine Grosnit
Huawei Noah's Ark Lab Technische Universität Darmstadt antoine.grosnit2@huawei.com
Haitham Bou Ammar
Huawei Noah's Ark Lab, University College London haitham.ammar@huawei.com
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.
Bayesian optimisation (BO) techniques are sample-efficient sequential model-based solvers optimising expensive-to-evaluate black-box objectives. Traditionally, BO methods operate in isolation focusing on one task at a time, a setting that led to numerous successful applications, including but not limited to hyperparameter tuning [1], drug and chip design [2], and robotics [3]. Although successful, focusing on only one objective in isolation may increase sample complexities on new tasks since standard BO algorithms work tabula-rasa as they start optimising from scratch, ignoring previously observed black-box functions.
To improve sample efficiency on new target tasks, meta-BO makes use of data collected on related source tasks [4] and attempts to transfer knowledge in between. Those methods are primarily composed of two parts: 1) a (meta) surrogate model predicting the function to optimise and 2) a (meta) acquisition function (AF) estimating how good a queried location is. Regarding surrogate
modelling, prior work relies on Gaussian processes (GPs) to perform transfer across tasks due to their data efficiency [5-9], or focuses on deep neural networks [10-12] gaining from their representation flexibility. As for acquisition functions, previous works assumed a fixed GP and trained a neural network to perform transfer [13, 14].
Although successful, those methods rely on the assumption that the two steps described above can be handled independently, potentially missing benefits from considering both steps together. Therefore, this paper advocates for an end-to-end training protocol for meta-BO where we train a model predicting AF values directly from observed data, without relying on a GP surrogate.
As shown in [10, 12, 15, 16], Neural processes (NP) are a good candidate for meta-learning due to their structural properties, thus we propose to use a new transformer-based NP to model the acquisition function directly [17].
Nevertheless, the lack of labelled AF data prohibits a supervised learning protocol. Here, we rely on reward functions to assess the goodness of AFs and formulate a new reinforcement learning (RL) problem. Our RL formulation attempts to learn an optimal policy that selects new evaluation probes by minimising per-task cumulative regrets.
Early on, we faced very unstable training with minimal learning due to the combination of RL with transformers [18]. Furthermore, we notice that this problem amplifies when the reward function is sparse, which is the case in our setting, as we formally show in Section 3.2. To mitigate this difficulty, we introduce an inductive bias in our transformer architecture via an auxiliary loss [19] that guides a part of the network to learn a valid probabilistic model, effectively specialising as a neural process with gradient updates. Having developed our framework, we demonstrate state-of-the-art regret performance in hyperparameter optimisation of real-world problems and combinatorial sequence optimisation for antibody and chip design (see Section 4). Those solid empirical results show the value of end-to-end training in meta-BO to further improve sample complexity.
Contributions We summarise our contributions as follows:  $i$  ) developing the first end-to-end transformer-based architecture for meta-BO predicting acquisition function values, ii) identifying logarithmic reward sparsity patterns hindering end-to-end learning with RL, iii) introducing NP-based inductive biases for successful model updates, and iv) demonstrating new state-of-the-art regret results on a wide range of traditional (hyperparameter tuning) and non-traditional (chip and antibody design) real-world benchmarks.
2.1 Bayesian Optimisation
In BO, we sequentially maximise an expensive-to-evaluate black-box function  $f(\pmb{x})$  for variables  $\pmb{x} \in \mathcal{X}$ . BO techniques operate in two steps. First, based on previously collected evaluation data, we fit a probabilistic surrogate model to emulate  $f(\pmb{x})$  allowing us to make probabilistic predictions of the function's behaviour on unobserved input points. Given the first step's probabilistic predictions, the second step in BO optimises an AF that trades off exploration and exploitation to find a new input query,  $\pmb{x}_{\mathrm{new}} \in \mathcal{X}$ , for evaluation. Both the model and acquisition choices play a critical role in the success of BO. Many works adopt GPs as surrogate models due to their sample efficiency and practical uncertainty estimation [20]. Moreover, on the acquisition side, the common practice is to optimise one from a set of widely adopted AFs such as Expected Improvement (EI) [21].
2.2 Transfer in Bayesian Optimisation
Transfer techniques in BO [4, 13] aim to improve the optimisation of a newly observed target black-box function by reusing knowledge from past experiences gathered in related source domains, i.e. solve  $\max_{\boldsymbol{x} \in \mathcal{X}} f_{\mathrm{Target}}(\boldsymbol{x})$  by leveraging information from  $K$  source black-boxes  $f_1(\boldsymbol{x}), \ldots, f_K(\boldsymbol{x})$ . We assume this information is available to our agent via  $\mathcal{D}_1, \ldots, \mathcal{D}_K$  such that each dataset  $\mathcal{D}_k = \{\langle \boldsymbol{x}_i^{(k)}, y_i^{(k)} \rangle\}_{i=1}^{n^{(k)}}$  consists of  $n^{(k)}$  (noisy) evaluations of  $f_k(\boldsymbol{x})$  for all  $k \in [1:K]$ . Following the meta-learning literature [4, 5, 13], we impose no additional assumptions on the process that collects  $\mathcal{D}_1, \ldots, \mathcal{D}_K$ , allowing for a diverse set of source optimisation algorithms, including but not limited to BO, genetic and evolutionary algorithms [22], and sampling-based strategies [23]. Many algorithms that use source data to improve performance on target domains exist. We categorise those based
on the part they customise within the BO pipeline, i.e., by registering if they affect initial points [24, 25], search spaces [26], surrogate models [5] or AFs [13, 14]. Since our work devises end-to-end meta-BO pipelines, we now elaborate on critical surrogate models needed for the rest of the paper and leave a detailed presentation of related work to Section 5.
Although GPs [20] are a prominent tool for single-task BO, high computational complexity [27], among other drawbacks (e.g. smoothness assumptions or limited scalability in terms of dimensions), can limit their application to transfer scenarios. Of course, one can generalise GPs to support transfer using multi-task kernels [28] or GP ensembles [6], for example. However, recent trends demonstrated superior performance to such extensions when using deep networks (e.g., neural processes) to meta-learn probabilistic surrogates [5, 13].
Neural Processes A popular method for effective meta-learning consists of directly inputting context points (i.e., available data to adjust to a target task) for a model to adapt to unseen tasks [15]. We can accomplish such a goal by relying on neural processes (NPs), a class of models that combine the flexibility of deep neural networks with the properties of stochastic processes [10]. Given an observation set of input-output pairs  $\mathcal{D}_{\mathrm{obs}}$  and a set of  $n_{\mathrm{pred}}$  locations  $\pmb{x}^{\mathrm{(pred)}}$  at which we desire to make predictions, an NP parameterised by  $\pmb{\theta}$  outputs a distribution  $p_{\pmb{\theta}}(\cdot|\pmb{x}^{\mathrm{(pred)}}, \mathcal{D}_{\mathrm{obs}})$  that approximates the true posterior over the labels  $y^{\mathrm{(pred)}}$ . Among the different parameterisations of  $p_{\pmb{\theta}}(\cdot)$ , e.g., in [29-31], transformer architectures [17] recently emerged as compelling choices of NPs [11, 12]. In this work, we also adopt a transformer architecture inspired by this architecture's broad software support and state-of-the-art supervised learning results as reported in [11].
2.3 Reinforcement Learning
In Section 3, we attempt to learn an end-to-end model that predicts acquisition values. Of course, we cannot fit the parameters of such a model with supervised learning due to the lack of labelled acquisition data. RL [32] is a viable alternative for learning from delayed and non-differentiable reward signals in those cases. In RL, we formalise problems as Markov decision processes (MDP):  $\mathcal{M} = \langle S, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle$ , where  $\mathcal{S}$  and  $\mathcal{A}$  denote the state and action spaces,  $\mathcal{P}: \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to [0,1]$  the state transition model,  $\mathcal{R}$  the reward function dictating the optimisation goal, and  $\gamma \in [0,1)$  a discount factor specifying the degree to which rewards are discounted over time. A policy  $\pi: \mathcal{S} \times \mathcal{A} \to [0,1]$  is an action-selection rule that is defined as a probability distribution over state-action pairs, where  $\pi(\boldsymbol{a}_t | s_t)$  represents the probability of selecting action  $\boldsymbol{a}_t$  in state  $s_t$ . An RL agent aims to find an optimal policy  $\pi^{\star}$  that maximises (discounted) expected returns. For determining  $\pi^{\star}$ , we use the state-of-the-art proximal policy optimisation (PPO) algorithm [33].
3 End-to-End Meta-Bayesian Optimisation
While transfer techniques in BO saw varying degrees of success in many applications, current approaches lack end-to-end differentiability, where model learning and acquisition discovery arise as two separate steps. Specifically, the surrogate model gradients hardly affect the acquisition network, and the acquisition's network gradients fail to back-propagate to the surrogate's updates.
Enabling end-to-end transfer frameworks in which we learn the surrogate and acquisition jointly hold the promise for more scalable and easier-to-deploy algorithms that are more robust to input data or task changes. Following such a framework, we also expect more accurate predictions that can lead to better regret results (see Section 4) since we optimise the entire transfer pipeline, including the intermediate probabilistic model and AF representations. Additionally, end-to-end training techniques allow us to mitigate the need for domain-specific expertise and permit stable implementations that benefit fully from GPU and computing hardware.
The most straightforward way to enable end-to-end training in Bayesian optimisation is to introduce a deep network that acquires search variables and historical evaluations of black-box functions as inputs and outputs acquisition values after a set of nonlinear transformations. Of course, it is challenging to fit the weights of such a network due, in part, to a lack of labelled acquisition data where search variables and history of evaluations are inputs and acquisition values are labels.
3.1 Reinforcement Learning for End-to-End Training
Our approach utilises RL to fit the network's parameters  $\theta$  from minimal supervision, circumventing the need for labelled acquisition data. To formalise the RL problem, we introduce an MDP where:
State:  $s_t = [\mathcal{H}_t, t, T]$  (history, BO time-step & budget) Action:  $a_t = x_t$  (choice of new probe), with  $\mathcal{H}_t = \{\langle x_1, y_1 \rangle, \dots, \langle x_{t-1}, y_{t-1} \rangle\}$  denoting the history of black-box evaluations up-to the current time-step  $t$ . Adding the current BO step  $t$  and the maximum budget  $T$  in our state variable  $s_t$  helps balance exploration versus exploitation trade-offs as noted in [34]. Our MDP's transition function is straightforward, updating  $\mathcal{H}_t$  by appending newly evaluated points, i.e.,  $\mathcal{H}_{t+1} = \mathcal{H}_t \cup \{\langle x_t, y_t \rangle\}$  and incrementing the time variable  $t$ . Regarding rewards, we follow the well-established literature [13, 14] and define  $r_t = \max_{1 \leq \ell \leq t} y_\ell$  to correspond to simple regret. Given such an MDP, our agent attempts to find a parameterised policy  $\pi_\theta$  which, when conditioned on  $s_t$ , proposes a new probe  $x_t$  that minimises cumulative regret (i.e., the sum of total discounted simple regrets).
Extensions to Multi-Task Reinforcement Learning The above MDP describes an RL configuration that learns  $\theta$  in a single BO task. We now extend this formulation to multi-task scenarios allowing for a meta-learning setup. To do so, we introduce a set of MDPs  $\mathcal{M}_1, \ldots, \mathcal{M}_K$  with  $K$  being the total number of available tasks. This paper considers same-domain multi-task learning scenarios. As such, we assume that all MDPs share the same state and action spaces and leave cross-domain extensions as an interesting avenue for future work. We define each MDP,  $\mathcal{M}_k$ , as previously introduced such that:
States:  $s_t^{(k)} = [\mathcal{H}_t^{(k)}, t^{(k)}, T^{(k)}]$   
Actions:  $\pmb{a}_t^{(k)} = \pmb{x}_t^{(k)}$   
Rewards:  $r_t^{(k)} = \max_{1 \leq \ell \leq t^{(k)}} y_t^{(k)} \forall k$ .
Moreover, for each task  $k$ , the transition model updates task-specific histories with  $\mathcal{H}_{t+1}^{(k)} = \mathcal{H}_t^{(k)} \sqcup \{\langle \pmb{x}_t^{(k)}, y_t^{(k)} \rangle\}$  and increments  $t^{(k)}$ . Contrary to the single task setup, we now seek a policy  $\pi_\theta$  which performs well on average across  $K$  tasks:
$$
\arg \max  _ {\pi_ {\boldsymbol {\theta}}} J (\pi_ {\boldsymbol {\theta}}) = \arg \max  _ {\pi_ {\boldsymbol {\theta}}} \mathbb {E} _ {k \sim p _ {\text {t a s k s}}} \left[ \mathbb {E} _ {\mathcal {H} _ {T ^ {(k)}} ^ {(k)} \sim p _ {\pi_ {\boldsymbol {\theta}}}} \left[ \sum_ {t = 1} ^ {T ^ {(k)}} \gamma^ {t - 1} r _ {t} ^ {(k)} \right] \right], \tag {1}
$$
where  $p_{\mathrm{tasks}}$  denotes the task distribution. Furthermore, the per-task history distribution  $p_{\pi_\theta}(\mathcal{H}_{T^{(k)}}^{(k)})$  is jointly parameterised by  $\theta$  and defined as:
$$
p _ {\pi_ {\boldsymbol {\theta}}} \left(\mathcal {H} _ {T ^ {(k)}} ^ {(k)}\right) = p \left(y _ {T ^ {(k)} - 1} ^ {(k)} \mid \boldsymbol {x} _ {T ^ {(k)} - 1} ^ {(k)}\right) \pi_ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {T ^ {(k)} - 1} ^ {(k)} \mid \mathcal {H} _ {T ^ {(k)} - 1} ^ {(k)}\right) \dots p \left(y _ {1} ^ {(k)} \mid \boldsymbol {x} _ {1} ^ {(k)}\right) \mu_ {0} \left(\boldsymbol {x} _ {1} ^ {(k)}\right), \tag {2}
$$
with  $\mu_0\left(\boldsymbol{x}_1^{(k)}\right)$  denoting an initial (action) distribution from which we sample the first decision  $\boldsymbol{x}_1^{(k)}$ . While it appears that off-the-shelf PPO can tackle the problem in Equation 1, Section 3.2 details difficulties associated with such an RL formulation, noting that in BO situations, the sparsity of reward signals can impede the learning of the network's weights  $\theta$ . Before presenting those arguments, we now differentiate from the closest prior art, clarifying critical MDP differences.
Connection & Differences to [13] Although we are the first to propose end-to-end multi-task MDPs for meta-BO, others have also used RL to discover AFs, for example. Our parameterisation architecture and MDP definition significantly differ from the prior art, particularly from [13], the closest method to our work. Our approach uses a transformer-based deep network model to parameterise the whole pipeline (see Section 3.4). In contrast, the work in [13] assumes a pre-trained fixed Gaussian process surrogate and uses multi-layer perceptrons only to discover acquisitions. Our state variable requires historical information from which we jointly learn a probabilistic model and an acquisition end-to-end instead of requiring posterior means and variances of a Gaussian process model. Hence, our framework is more flexible, allowing us to model non-smooth black-box functions while overcoming some drawbacks of GP surrogates, like training and inference times.
3.2 Limitations of Regret Rewards in End-to-End Training
To define Equation 1, we followed the well-established literature of meta-BO [4, 13] and utilised simple regret reward functions. Although this choice is reasonable, we face challenges in applying such rewards in end-to-end training. Apart from difficulties associated with end-to-end training
of deep architectures [18], our RL algorithm is subject to additional complexities when estimating gradients from Equation 1 due to the sparsity of the reward function. To better understand this problem, we start by noticing that for a reward component  $r_t^{(k)}$  to contribute to the cumulative summation  $\sum_{t} \gamma^{t-1} r_t^{(k)}$ , we need to observe a function value  $y_t^{(k)}$  that outperforms all values we have seen so far, i.e.,  $y_t^{(k)} > \max_{1 \leq \ell < t} y_\ell^{(k)}$ . During the early training stages of RL, we can quantify the average number of such informative events (when  $y_t^{(k)} > \max_{1 \leq \ell < t} y_\ell^{(k)}$ ) by a combinatorial argument that frames this calculation as a calculation of the number of cycles in a permutation of  $T^{(k)}$  elements, leading us to the following lemma.
Lemma 3.1. Consider a task with a horizon length (budget)  $T$ , and define  $r_t = \max_{1 \leq \ell \leq t} y_t$  the simple regret as introduced in Equation 1. For a history  $\mathcal{H}_T$ , let  $m_{\mathcal{H}}$  denote the total number of informative rewards, i.e., the number of steps  $t$  at which  $y_t > \max_{1 \leq \ell < t} y_\ell$ . Under a random policy  $\pi_{\theta}$ , the number of informative events is logarithmic in  $T$  such that:  $\mathbb{E}_{\mathcal{H} \sim p_{\pi_\theta}}[m_{\mathcal{H}}] = \mathcal{O}(\log T)$ , where  $p_{\pi_\theta}$  is induced by  $\pi_{\theta}$  as in Equation 2.
We defer the proof of Lemma 3.1 to Appendix A due to space constraints. Here, we note that this result implies that the information contained in one sampled trajectory is sparse at the beginning of RL training when the policy acts randomly. Of course, this increases the difficulty of estimating informative gradients of Equation 1 when updating  $\theta$ . One can argue that the sparsity described in Lemma 3.1 only holds under random policies during the early stages of RL training and that sparsity patterns decrease as policies improve. Interestingly, simple regret rewards do not necessarily confirm this intuition. To realise this, consider the other end of the RL training spectrum in which policies have improved to near optimality such that  $\pi_{\theta} \rightarrow \pi_{\theta^{\star}}$ . Because  $\pi_{\theta}$  has been trained to maximise regret, it will seek to suggest the optimal point of the current task, as early as possible in the BO trajectory. Consequently, the policy is encouraged to produce trajectories with even sparser rewards during later training stages, further complicating the problem of informative gradient estimates of Equation 1.
3.3 Inductive Biases and Auxiliary Tasks
Learning from sparse reward signals is a well-known difficulty in the reinforcement learning literature [32]. Many solutions, from imitation learning, [35] to exploration bonuses [36], improve reward signals to reduce agent-environment interactions and enhance gradient updates. Others [37] attempt to define more informative rewards from prior knowledge or via human interactions [38]. Unfortunately, both of those approaches are hard to use in BO. Indeed, manually engineering black-box-specific rewards is notoriously difficult and requires domain expertise and extensive knowledge of the source and target black-box functions we wish to optimise. Furthermore, learning from human feedback is data-intensive, conflicting with the goal of sample-efficient optimisation.
Another prominent direction demonstrating significant gains is the introduction of auxiliary tasks (losses) within RL that allows agents to discover relevant inductive biases leading to impressive empirical successes [19, 39, 40]. Inspired by those results, we propose introducing an inductive bias in our method via an auxiliary supervised loss. Since we have at our disposal the collected source tasks datasets on which we are training our architecture  $\mathcal{D}^{(1)},\ldots ,\mathcal{D}^{(K)}$ , we augment our objective such that our RL agent maximises not only rewards but also the likelihood of making correct predictions on these labelled datasets.
Supervised Auxiliary Loss Consider a source task  $k$  and consider that we split its corresponding dataset into an observed set and a prediction set  $\mathcal{D}^{(k)} = \mathcal{D}_{\mathrm{obs}}^{(k)} \sqcup \mathcal{D}_{\mathrm{perd}}^{(k)}$ . We define the auxiliary loss to be exactly the log-likelihood of functions values  $y^{(\mathrm{pred})}$  at predicted locations  $x^{(\mathrm{pred})}$  given observations  $\mathcal{D}_{\mathrm{obs}}$ :
$$
\mathcal {L} (\boldsymbol {\theta}) = \mathbb {E} _ {k \sim p _ {\text {t a s k}}, \mathcal {D} _ {\text {o b s}} ^ {(k)}, \mathcal {D} _ {\text {p e r d}} ^ {(k)}} \left[ \log p _ {\boldsymbol {\theta}} \left(y _ {k} ^ {(\text {p r e d})} \mid \boldsymbol {x} _ {k} ^ {(\text {p r e d})}, \mathcal {D} _ {\text {o b s}} ^ {(k)}\right) \right]. \tag {3}
$$
Interestingly, this part of our model specialises as a neural process (Section 2.2) with  $\mathcal{D}_{\mathrm{obs}}^{(k)}$  being the history  $\mathcal{H}_t^{(k)}$  and  $\pmb{x}^{(\mathrm{pred})}$  being the input points at which we wish to make predictions. To represent  $p_{\theta}(\cdot)$ , we use a head of our transformer to predict multi-modal Riemannian (or bar plot probability density function) posteriors as [11, 41]. We compute Equation 3 on random iid splits of  $\mathcal{D}^{(k)}$  and not directly on trajectories generated by the policy. This is because as the policy improves, the trajectories
it generates are composed of less diverse (non-iid.) points. Indeed as training progresses,  $\pi_{\theta}$  becomes better and therefore finds the optimal point in  $\mathcal{D}^{(k)}$  more rapidly. To fully take advantage of the labelled data at our disposal, we evaluate this auxiliary loss on iid. data, i.e. splits of  $\mathcal{D}^{(k)}$  into  $\mathcal{D}_{\mathrm{obs}}^{(k)}$  and  $\mathcal{D}_{\mathrm{pred}}^{(k)}$  sampled uniformly at random. It is sensible from a BO standpoint as well since we want the introduced inductive bias to encourage our network to maximise the likelihood not only in a neighbourhood of the optimiser but also in the rest of the dataset.
3.4 Neural Acquisition Processes