Fine-tuning LLMs with variational Bayesian last layer for high-dimensional Bayesian optimization
Haotian Xiang, Jinwen Xu, Qin Lu
TL;DR
The paper tackles high-dimensional black-box optimization with irregular variable types by using pretrained LLMs as global surrogates. It introduces LoRA-VBLL, a parameter-efficient fine-tuning framework where LoRA adapters modify the backbone and a variational Bayesian last layer models uncertainty in a closed-form, enabling recursive Bayesian updates and feature caching. An ensemble of LoRA-VBLL models across ranks automates rank selection and further improves robustness via Bayesian model averaging with data-driven weights. Across HD BO benchmarks and molecular optimization tasks, ENS-LoRA-VBLL demonstrates strong optimization performance with favorable memory and runtime characteristics, highlighting its practical potential for domains like drug discovery and material design.
Abstract
A plethora of applications entail solving black-box optimization problems with high evaluation costs, including drug discovery, material design, as well as hyperparameter tuning. Toward finding the global optimum of such black-box optimization problems with sample efficiency, Bayesian optimization (BO) is a theoretically elegant framework that relies on a probabilistic surrogate model so as to iteratively select the query point with well-balanced exploration-exploitation tradeoffs. The Gaussian process (GP), as the de-facto choice for surrogate modeling, has achieved compelling performances for vanilla BO with low-dimensional continuous variables. However, GPs fall short in coping with high-dimensional counterparts with {\it irregular} variables (e.g., categorical, ordinal, etc.). To alleviate this, neural network-based surrogates have been explored. Inspired by the powerful capabilities of LLMs, we adopt the LLM as the surrogate to model the mapping from the high-dimensional input variables to the objective function. To adapt to the current problem, we leverage the low-rank adaptation (LoRA) to fine-tune the LLM parameters together with the posterior of a linear regression head via the variational Bayesian last layer (VBLL) framework. The resulting LoRA-VBLL is not only computationally light compared to existing alternatives, but also admits recursive updates. To automate the critical selection of the LoRA rank as well as other hyperparameters, a weighted ensemble (ENS) of LoRA-VBLL surrogates has been devised, which further accommodates continual update of the per-model weight and individual LoRA-VBLL parameters via recursive Bayes. Extensive experimental results demonstrate the compelling performance of the proposed (ENS-)LoRA-VBLL approaches on various high-dimensional benchmarks and the real-world molecular optimization tasks.