Rethinking Bayesian Optimization for Co-Optimizing LLM Training Configurations

Zhiliang Chen; Alfred Wei Lun Leong; Shao Yong Ong; Apivich Hemachandra; Gregory Kang Ruey Lau; Chuan-Sheng Foo; Zhengyuan Liu; Nancy F. Chen

cs.LGFeb 9, 2026

Rethinking Bayesian Optimization for Co-Optimizing LLM Training Configurations

Authors: Zhiliang Chen, Alfred Wei Lun Leong, Shao Yong Ong, Apivich Hemachandra, Gregory Kang Ruey Lau, Chuan-Sheng Foo, Zhengyuan Liu, Nancy F. Chen,

Organizations: Department of Computer Science, National University of Singapore, Singapore · Agency for Science, Technology and Research (A*STAR), Singapore

Abstract

Fine-tuning an LLM to maximize performance on a downstream task requires finding the optimal data and model configuration. This is a black-box optimization problem that Bayesian optimization (BO) addresses by sequentially evaluating training configurations and using observed performance feedback to adaptively guide the search towards better configurations. However, directly applying BO to the joint data-model space suffers from two fundamental challenges: the prohibitive cost of full LLM training at each iteration, and the large number of BO iterations required for effective exploration due to the high problem dimensionality. This paper introduces JoBS, a BO-based approach that co-optimizes data and model configurations without requiring a full training run at every iteration. JoBS allocates an initial portion of the optimization budget to learn a scaling-law-inspired performance predictor that estimates fully trained LLM performance from only a small number of training steps. The remaining budget is then used to run BO with this predictor, eliminating the need for full training runs and enabling JoBS to explore significantly more data and model configurations within the same budget. We analyze JoBS's average regret and derive the optimal budget allocation between predictor learning and BO iterations. Empirically, JoBS outperforms independent data and model optimization methods and existing multi-fidelity BO baselines across a diverse set of downstream tasks. Our code is available at https://github.com/chenzhiliang94/BO-for-LLMs.

Figures & tables

Figure 1: JoBS optimally balances budget between learning a performance predictor and running Bayesian optimization iterations.

Panel of Figure 2 — Figure 2: LLM performance varies with different configurations.

Figure 3: Training a neural network to predict fully trained LLM performance from a small number of training steps.

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	JoBS
	Default	$71.9_{\pm 1.5}$	$71.1_{\pm 1.2}$	$72.3_{\pm 3.2}$	$68.7_{\pm 1.0}$	$74.4_{\pm 2.0}$	-
	DARTS	$73.1_{\pm 0.9}$	$71.7_{\pm 0.7}$	$74.7_{\pm 1.4}$	$69.3_{\pm 0.5}$	$66.8_{\pm 0.8}$	-
	AutoLoRA	$72.9_{\pm 1.2}$	$75.2_{\pm 0.4}$	$70.9_{\pm 0.8}$	$68.5_{\pm 0.5}$	$74.0_{\pm 0.6}$	-
	RoBoT	$71.8_{\pm 0.7}$	$72.7_{\pm 1.6}$	$74.0_{\pm 1.9}$	$73.1_{\pm 1.6}$	$70.2_{\pm 1.8}$	-
	JoBS	-	-	-	-	-	86.0 ±0.7

Table 1: GSM

8

K accuracy (higher is better) for all combinations of data and model optimization methods vs. JoBS .

Figure 5: Best-found LLM performance against optimization budget (training steps, in thousands) for JoBS and BO baselines across six language tasks, as well as out-of-domain (OOD) and multi-task settings. JoBS consistently achieves the highest performance across all settings. JoBS ’ curve begins later as an initial portion of the budget is allocated to predictor learning.

Panel of Figure 6 — Figure 6: (a) Ablation on choice of $N$ . (b) Ablation on choice of $B_{\text{small}}$ . (c) Wall-clock runtime of JoBS vs. data optimization methods.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Kernel	MMLU	TriviaQA	GSM8K
Hamming + RBF (SE2)	$74.4\pm 0.4$	$72.1\pm 0.4$	$86.0\pm 0.7$
RBF (SE2)	$74.5\pm 0.3$	$72.0\pm 1.1$	$86.4\pm 1.2$
Matérn	$73.7\pm 0.8$	$71.0\pm 0.6$	$85.6\pm 0.7$

Appendix

Table 2: Performance across different kernels on MMLU, TriviaQA, and GSM8K.

Method	Time (hours)
LESS	$16.3$
DoReMi	$18.5$
IF	$52$
Diversity	$122$
JoBS	$\pagecolor{darkgreen!50}\mathbf{\underline{13.3}}$

Appendix

Table 3: Wall-clock runtime comparison of data selection methods vs. JoBS .

	CQA	GSM $8$ K	PubmedQA	SciQ	TrivQA	TruthQA	Wiki	MMLU	ARC
JoBS	0	0.34	0	0.10	0.19	0	0.21	0.16	0
DoReMi	0.08	0.11	0.18	0.05	0.08	0.14	0.04	0.16	0.13

Appendix

Table 4: Optimal data mixing ratio found by JoBS vs. other baselines. The columns denote the ratio allocated to each training domain.

	Rank	NumLayers	Order	Q	K	V	Up	Down	dropout	$\alpha$
JoBS	36	25	1	1	0	1	1	0	0.112	64
DARTS	12	13	0	1	1	1	1	0	0.058	45

Appendix

Table 5: Optimal model configuration found by JoBS vs. other baselines.

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	BO	JoBS
	Default	$59.8_{\pm 1.9}$	$60.7_{\pm 1.0}$	$62.9_{\pm 2.6}$	$62.3_{\pm 1.3}$	$64.5_{\pm 1.2}$	$74.8_{\pm 1.0}$	-
	DARTS	$61.9_{\pm 1.3}$	$62.4_{\pm 0.6}$	$66.7_{\pm 1.3}$	$63.8_{\pm 0.5}$	$64.8_{\pm 1.2}$	$77.0_{\pm 0.9}$	-
	AutoLoRA	$61.7_{\pm 1.0}$	$66.9_{\pm 1.2}$	$64.5_{\pm 1.1}$	$65.3_{\pm 1.2}$	$65.9_{\pm 0.6}$	$72.5_{\pm 0.5}$	-
	RoBoT	$64.2_{\pm 0.6}$	$65.8_{\pm 0.7}$	$58.9_{\pm 1.3}$	$56.9_{\pm 1.0}$	$65.8_{\pm 0.6}$	$73.9_{\pm 1.3}$	-
	BO	$66.8_{\pm 1.2}$	$67.9_{\pm 0.5}$	$69.8_{\pm 0.9}$	$70.9_{\pm 1.0}$	$64.7_{\pm 1.4}$	$76.1_{\pm 2.0}$	-
	JoBS	-	-	-	-	-	-	$\pagecolor{darkgreen!50}\mathbf{\underline{78.0}}_{\pm 0.0}$

Appendix

Table 6: TruthfulQA [ Lin et al., 2022 ] .

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	BO	JoBS
	Default	$76.3_{\pm 1.0}$	$73.0_{\pm 0.8}$	$74.2_{\pm 1.7}$	$79.3_{\pm 0.7}$	$77.4_{\pm 1.7}$	$80.6_{\pm 0.8}$	-
	DARTS	$79.6_{\pm 1.3}$	$76.3_{\pm 1.7}$	$76.1_{\pm 1.1}$	$73.7_{\pm 1.2}$	$80.1_{\pm 1.1}$	$79.6_{\pm 0.6}$	-
	AutoLoRA	$78.9_{\pm 0.9}$	$79.8_{\pm 0.4}$	$76.1_{\pm 0.5}$	$77.9_{\pm 1.2}$	$78.0_{\pm 1.0}$	$81.5_{\pm 1.0}$	-
	RoBoT	$74.9_{\pm 0.8}$	$75.5_{\pm 0.9}$	$77.1_{\pm 0.9}$	$79.4_{\pm 1.5}$	$76.3_{\pm 0.9}$	$80.2_{\pm 0.2}$	-
	BO	$79.7_{\pm 1.3}$	$79.4_{\pm 0.3}$	$77.0_{\pm 0.4}$	$81.1_{\pm 0.9}$	$79.4_{\pm 1.1}$	$80.7_{\pm 1.2}$	-
	JoBS	-	-	-	-	-	-	$\pagecolor{darkgreen!50}\mathbf{\underline{88.0}}{\pm 1.5}$

Appendix

Table 7: CommonsenseQA [ Talmor et al., 2019 ] .

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	BO	JoBS
	Default	$61.2_{\pm 1.3}$	$63.5_{\pm 0.9}$	$59.7_{\pm 1.8}$	$57.9_{\pm 0.6}$	$62.1_{\pm 1.4}$	$64.2_{\pm 1.2}$	-
	DARTS	$58.3_{\pm 0.7}$	$61.0_{\pm 2.1}$	$62.9_{\pm 1.0}$	$55.7_{\pm 1.6}$	$60.1_{\pm 0.5}$	$63.4_{\pm 2.0}$	-
	AutoLoRA	$62.5_{\pm 1.4}$	$64.3_{\pm 0.6}$	$60.8_{\pm 2.2}$	$58.2_{\pm 1.9}$	$63.7_{\pm 0.8}$	$61.5_{\pm 1.1}$	-
	RoBoT	$59.9_{\pm 0.9}$	$60.7_{\pm 1.2}$	$63.4_{\pm 1.7}$	$61.5_{\pm 1.5}$	$58.3_{\pm 2.3}$	$62.1_{\pm 0.7}$	-
	BO	$55.8_{\pm 1.8}$	$57.2_{\pm 0.4}$	$61.3_{\pm 1.2}$	$60.5_{\pm 1.6}$	$63.9_{\pm 1.0}$	$59.6_{\pm 1.5}$	-
	JoBS	-	-	-	-	-	-	$\pagecolor{darkgreen!50}\mathbf{\underline{74.4}}_{\pm 0.4}$

Appendix

Table 8: MMLU [ Hendrycks et al., 2021 ] .

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	BO	JoBS
	Default	$54.7_{\pm 1.3}$	$59.2_{\pm 0.7}$	$61.4_{\pm 2.0}$	$52.8_{\pm 1.5}$	$60.6_{\pm 0.9}$	$62.3_{\pm 1.2}$	-
	DARTS	$58.1_{\pm 0.8}$	$61.0_{\pm 1.6}$	$62.8_{\pm 0.5}$	$54.3_{\pm 2.1}$	$57.9_{\pm 1.7}$	$60.5_{\pm 0.6}$	-
	AutoLoRA	$60.4_{\pm 1.1}$	$63.2_{\pm 0.9}$	$58.6_{\pm 1.9}$	$55.1_{\pm 0.8}$	$62.1_{\pm 2.0}$	$59.8_{\pm 1.0}$	-
	RoBoT	$56.8_{\pm 1.5}$	$58.7_{\pm 1.4}$	$61.1_{\pm 1.2}$	$60.3_{\pm 0.7}$	$55.7_{\pm 2.2}$	$61.4_{\pm 1.3}$	-
	BO	$52.6_{\pm 2.0}$	$55.9_{\pm 0.6}$	$59.7_{\pm 1.3}$	$58.5_{\pm 1.4}$	$63.4_{\pm 1.0}$	$57.4_{\pm 0.9}$	-
	JoBS	-	-	-	-	-	-	$\pagecolor{darkgreen!50}\mathbf{\underline{65.6}}_{\pm 0.0}$

Appendix

Table 9: ARC [ Clark et al., 2018 ] .

$\downarrow$ Model \| Data $\rightarrow$		Default	LESS	DoReMi	IF	Diversity	BO	JoBS
	Default	$55.5_{\pm 1.4}$	$57.2_{\pm 0.8}$	$53.1_{\pm 0.9}$	$55.8_{\pm 0.7}$	$58.9_{\pm 0.8}$	$65.0_{\pm 0.6}$	-
	DARTS	$58.2_{\pm 0.8}$	$61.3_{\pm 1.2}$	$61.0_{\pm 0.7}$	$63.3_{\pm 1.0}$	$59.2_{\pm 0.6}$	$66.7_{\pm 1.8}$	-
	AutoLoRA	$67.8_{\pm 1.4}$	$64.7_{\pm 0.9}$	$70.6_{\pm 2.2}$	$68.6_{\pm 1.7}$	$66.2_{\pm 1.5}$	$69.7_{\pm 2.4}$	-
	RoBoT	$58.4_{\pm 1.5}$	$62.3_{\pm 1.7}$	$64.2_{\pm 1.4}$	$57.2_{\pm 1.2}$	$63.4_{\pm 1.5}$	$68.2_{\pm 1.3}$	-
	BO	$70.7_{\pm 1.4}$	$66.7_{\pm 0.8}$	$72.5_{\pm 0.8}$	$71.7_{\pm 0.9}$	$74.7_{\pm 1.0}$	$72.7_{\pm 2.3}$	-
	JoBS	-	-	-	-	-	-	$\pagecolor{darkgreen!50}\mathbf{\underline{72.1}}{\pm 0.4}$

Appendix

Table 10: TriviaQA Gen [ Joshi et al., 2017 ] .

Figure 7: JoBS performance for LLM evaluation performance with Qwen-14B .

$\downarrow$ Model \| Method $\rightarrow$	LESS + AutoLoRA	DoReMi + DARTS	JoBS
Llama-3-8B-Instruct	0.73	0.76	0.81
Qwen3-14B	0.76	0.81	0.83
Qwen3-32B	0.86	0.82	0.88

Appendix

Table 11: JoBS applied to full-parameter fine-tuning of larger models.

Figure 8: JoBS performance for LLM evaluation loss (instead of performance).

Model Task		Random Data	Random Model	JoBS
Llama-3-8B-Instruct	GSM $8$ K	$67.3_{\pm 1.6}$	$71.5_{\pm 0.9}$	$86.4_{\pm{1.2}}$
	TruthfulQA	$59.8_{\pm 1.5}$	$64.2_{\pm 1.4}$	$80.6_{\pm{1.1}}$
	CommonsenseQA	$76.4_{\pm 1.2}$	$76.3_{\pm 1.2}$	$84.3_{\pm{2.4}}$
	MMLU	$66.4_{\pm 0.7}$	$63.1_{\pm 1.1}$	$69.5_{\pm{0.8}}$
	ARC	$65.2_{\pm 1.7}$	$64.6_{\pm 0.6}$	$70.4_{\pm{1.3}}$
	TriviaQA	$61.7_{\pm 2.4}$	$63.2_{\pm 1.5}$	$76.2_{\pm{1.2}}$

Appendix

Table 12: Comparison of some naive baselines with JoBS (Higher is better), averaged over

5

trials. Random Data means that we randomly selected data mixtures and applied JoBS only on the model configurations (and vice versa for Random Model ).

Asset	License	URL
Wikitext	CC-BY-SA 3.0	https://huggingface.co/datasets/mindchain/wikitext2
GSM8K	MIT	https://huggingface.co/datasets/openai/gsm8k
PubmedQA	MIT	https://huggingface.co/datasets/qiaojin/PubMedQA
SciQ	CC-BY-NC 3.0	https://huggingface.co/datasets/allenai/sciq
TriviaQA	Apache 2.0	https://huggingface.co/datasets/mandarjoshi/trivia_qa
TruthfulQA	Apache 2.0	https://huggingface.co/datasets/domenicrosati/TruthfulQA

Appendix

Table 21

Bayesian Optimization

Bayesian optimization uses probabilistic models and sequential experiments to find strong solutions when evaluating a black-box objective is costly.

See all 247 papers

BoLT: A Benchmark to Democratize Black-box Optimization Research for Expensive LLM Tasks
May 16, 2026Ruth Wan Theng Chew, Zhiliang Chen, Apivich Hemachandra +1Bayesian Optimization Benchmark Design
National University of Singapore
How Low Can You Go: Fine-Tuning LLMs with Bayesian Optimization (LlMBO)
Oct 13, 2026Aditya Kulkarni, Davide Sartor, Carlos Soto +1LLM Fine-Tuning Fine-Tuning Language Models
Department of Physics, University of Massachusetts Amherst · Department of Mathematics and Statistics, University of Massachusetts Amherst
Automatic Configuration of LLM Post-Training Pipelines
Mar 19, 2026Channe Chwa, Xinle Wu, Yao LuLLM Post-Training Hyperparameter Optimization
National University of Singapore

May 16, 2026J/K move · Enter open · S save

BoLT: A Benchmark to Democratize Black-box Optimization Research for Expensive LLM Tasks

Ruth Wan Theng Chew, Zhiliang Chen, Apivich Hemachandra, Bryan Kian Hsiang Low

National University of Singapore

Bayesian Optimization Benchmark Design Black-Box Optimization Hyperparameter Optimization

Open paper AskPDF arXiv

Optimization of LLM training and inference configurations, such as hyperparameters, data mixtures, and prompts, is critical to performance, but it is often approached heuristically in practice, leading to potentially suboptimal outcomes. By framing them as noisy, expensive, and derivative-free optimization problems, Bayesian optimization (BO) and other black-box optimization (BBO) methods offer a promising yet underexplored direction for principled, sample-efficient methods. However, LLM training and inference costs are prohibitively high for most of the BBO research community, and new methods are often only evaluated on synthetic test functions and small-scale datasets that fail to capture the challenges of modern LLM optimization problems. This impedes the development of BBO methods and makes it difficult to assess their effectiveness on modern LLM tasks. We introduce BoLT, the first LLM-centric benchmark that democratizes LLM research for the BBO community. BoLT is released at https://github.com/chewwt/bolt. BoLT covers broad and well-motivated LLM optimization problems, involving multi-fidelity, multi-objective, heteroscedastic noise, and high-dimensional search spaces. Each problem in BoLT is grounded in real experimental data and made fully reproducible and accessible through lightweight surrogate models fitted to the results of thousands of real LLM experiments. We benchmark BoLT against an extensive range of BO and BBO methods, showing that selected BO methods consistently outperform others across tasks and highlighting gaps in existing BBO methods on LLM tasks, underscoring the need to modernize benchmarks for the BBO community.

Abstract

Bayesian Optimization

Explore similar work

BoLT: A Benchmark to Democratize Black-box Optimization Research for Expensive LLM Tasks

How Low Can You Go: Fine-Tuning LLMs with Bayesian Optimization (LlMBO)

Automatic Configuration of LLM Post-Training Pipelines