Fine-tuning LLMs with variational Bayesian last layer for high-dimensional Bayesian optimzation

2025-10-01

AI Summary

High-dimensional Bayesian optimization suffers from prohibitive evaluation costs and degraded performance of conventional Gaussian process surrogates when handling mixed-variable search spaces (e.g., categorical, ordinal). To address this, we propose LoRA-VBLL—a novel framework that pioneers the integration of large language models (LLMs) into black-box optimization. It jointly fine-tunes LLMs via low-rank adaptation (LoRA) and variational Bayesian last-layer (VBLL) inference, enabling lightweight, sample-efficient learning. The framework supports recursive posterior updates and online adaptation, and introduces a weighted Bayesian ensemble mechanism that automatically selects optimal LoRA ranks and hyperparameters. Evaluated on multiple high-dimensional synthetic benchmarks and real-world molecular optimization tasks, LoRA-VBLL consistently outperforms state-of-the-art methods, demonstrating superior modeling capacity, robustness to variable heterogeneity, and scalability to high-dimensional, mixed-type search spaces.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Optimization
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• Optimizing high-dimensional black-box functions with irregular variables
• Improving Bayesian optimization efficiency using LLM-based surrogate models
• Automating hyperparameter selection via weighted ensemble variational Bayesian methods

Innovation

• Fine-tuning LLMs with variational Bayesian last layer
• Using LoRA for computationally efficient parameter adaptation
• Employing weighted ensemble for automated hyperparameter selection