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, Bryan Kian Hsiang Low
Organizations: Department of Computer Science, National University of Singapore, Singapore · Agency for Science, Technology and Research (A*STAR), Singapore
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.
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.
Figure 4: Predictor error w.r.t. N
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
JoBS
Default
71.9±1.5
71.1±1.2
72.3±3.2
68.7±1.0
74.4±2.0
-
DARTS
73.1±0.9
71.7±0.7
74.7±1.4
69.3±0.5
66.8±0.8
-
AutoLoRA
72.9±1.2
75.2±0.4
70.9±0.8
68.5±0.5
74.0±0.6
-
RoBoT
71.8±0.7
72.7±1.6
74.0±1.9
73.1±1.6
70.2±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.
Figure 6: (a) Ablation on choice of N . (b) Ablation on choice of Bsmall . (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±0.4
72.1±0.4
86.0±0.7
RBF (SE2)
74.5±0.3
72.0±1.1
86.4±1.2
Matérn
73.7±0.8
71.0±0.6
85.6±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
\pagecolordarkgreen!5013.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
α
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.
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
BO
JoBS
Default
59.8±1.9
60.7±1.0
62.9±2.6
62.3±1.3
64.5±1.2
74.8±1.0
-
DARTS
61.9±1.3
62.4±0.6
66.7±1.3
63.8±0.5
64.8±1.2
77.0±0.9
-
AutoLoRA
61.7±1.0
66.9±1.2
64.5±1.1
65.3±1.2
65.9±0.6
72.5±0.5
-
RoBoT
64.2±0.6
65.8±0.7
58.9±1.3
56.9±1.0
65.8±0.6
73.9±1.3
-
BO
66.8±1.2
67.9±0.5
69.8±0.9
70.9±1.0
64.7±1.4
76.1±2.0
-
JoBS
-
-
-
-
-
-
\pagecolordarkgreen!5078.0±0.0
Appendix
Table 6: TruthfulQA [ Lin et al., 2022 ] .
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
BO
JoBS
Default
76.3±1.0
73.0±0.8
74.2±1.7
79.3±0.7
77.4±1.7
80.6±0.8
-
DARTS
79.6±1.3
76.3±1.7
76.1±1.1
73.7±1.2
80.1±1.1
79.6±0.6
-
AutoLoRA
78.9±0.9
79.8±0.4
76.1±0.5
77.9±1.2
78.0±1.0
81.5±1.0
-
RoBoT
74.9±0.8
75.5±0.9
77.1±0.9
79.4±1.5
76.3±0.9
80.2±0.2
-
BO
79.7±1.3
79.4±0.3
77.0±0.4
81.1±0.9
79.4±1.1
80.7±1.2
-
JoBS
-
-
-
-
-
-
\pagecolordarkgreen!5088.0±1.5
Appendix
Table 7: CommonsenseQA [ Talmor et al., 2019 ] .
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
BO
JoBS
Default
61.2±1.3
63.5±0.9
59.7±1.8
57.9±0.6
62.1±1.4
64.2±1.2
-
DARTS
58.3±0.7
61.0±2.1
62.9±1.0
55.7±1.6
60.1±0.5
63.4±2.0
-
AutoLoRA
62.5±1.4
64.3±0.6
60.8±2.2
58.2±1.9
63.7±0.8
61.5±1.1
-
RoBoT
59.9±0.9
60.7±1.2
63.4±1.7
61.5±1.5
58.3±2.3
62.1±0.7
-
BO
55.8±1.8
57.2±0.4
61.3±1.2
60.5±1.6
63.9±1.0
59.6±1.5
-
JoBS
-
-
-
-
-
-
\pagecolordarkgreen!5074.4±0.4
Appendix
Table 8: MMLU [ Hendrycks et al., 2021 ] .
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
BO
JoBS
Default
54.7±1.3
59.2±0.7
61.4±2.0
52.8±1.5
60.6±0.9
62.3±1.2
-
DARTS
58.1±0.8
61.0±1.6
62.8±0.5
54.3±2.1
57.9±1.7
60.5±0.6
-
AutoLoRA
60.4±1.1
63.2±0.9
58.6±1.9
55.1±0.8
62.1±2.0
59.8±1.0
-
RoBoT
56.8±1.5
58.7±1.4
61.1±1.2
60.3±0.7
55.7±2.2
61.4±1.3
-
BO
52.6±2.0
55.9±0.6
59.7±1.3
58.5±1.4
63.4±1.0
57.4±0.9
-
JoBS
-
-
-
-
-
-
\pagecolordarkgreen!5065.6±0.0
Appendix
Table 9: ARC [ Clark et al., 2018 ] .
↓ Model | Data →
Default
LESS
DoReMi
IF
Diversity
BO
JoBS
Default
55.5±1.4
57.2±0.8
53.1±0.9
55.8±0.7
58.9±0.8
65.0±0.6
-
DARTS
58.2±0.8
61.3±1.2
61.0±0.7
63.3±1.0
59.2±0.6
66.7±1.8
-
AutoLoRA
67.8±1.4
64.7±0.9
70.6±2.2
68.6±1.7
66.2±1.5
69.7±2.4
-
RoBoT
58.4±1.5
62.3±1.7
64.2±1.4
57.2±1.2
63.4±1.5
68.2±1.3
-
BO
70.7±1.4
66.7±0.8
72.5±0.8
71.7±0.9
74.7±1.0
72.7±2.3
-
JoBS
-
-
-
-
-
-
\pagecolordarkgreen!5072.1±0.4
Appendix
Table 10: TriviaQA Gen [ Joshi et al., 2017 ] .
Figure 7: JoBS performance for LLM evaluation performance with Qwen-14B .
↓ Model | Method →
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±1.6
71.5±0.9
86.4±1.2
TruthfulQA
59.8±1.5
64.2±1.4
80.6±1.1
CommonsenseQA
76.4±1.2
76.3±1.2
84.3±2.4
MMLU
66.4±0.7
63.1±1.1
69.5±0.8
ARC
65.2±1.7
64.6±0.6
70.4±1.3
TriviaQA
61.7±2.4
63.2±1.5
76.2±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 ).
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.
Ruth Wan Theng Chew, Zhiliang Chen, Apivich Hemachandra +1
Recently, Morris et al. introduced TinyLoRA, a strategy for LLM fine-tuning which introduces arbitrarily few parameters, and demonstrated accuracy improvements using less than 10 parameters rather than the ≳10,000 typically demanded by LoRA-based methods. However, that study continued to use gradient-based local search methods of the kind deployed for training multi-billion parameter LLMs in the first place. This article is predicated on the observation that, in the low dimensional setting, global optimization methods become tractable, unlocking Large language Model Bayesian Optimization (LlMBO). LlMBO is sample-efficient, gradient-free, memory-light, parallelizes across devices without synchronization, and, unlike policy-gradient methods, can optimize greedy-decoding accuracy directly. We describe numerical experiments comparing LlMBO against GRPO on fine-tuning tasks for a variety of models on a math benchmark suite which show that LlMBO can be a suitable alternative in this regime, matching or exceeding GRPO's fine tuning gains, while using a fraction of the compute and memory.
Aditya Kulkarni, Davide Sartor, Carlos Soto +1
Department of Physics, University of Massachusetts Amherst · Department of Mathematics and Statistics, University of Massachusetts Amherst
LLM post-training pipelines that combine supervised fine-tuning and reinforcement learning are difficult to configure under realistic compute budgets: the configuration space is high-dimensional and heterogeneous, stages are strongly coupled, and each end-to-end evaluation is expensive. We propose AutoPipe, a budget-aware two-phase framework for configuration selection in LLM post-training. Offline, AutoPipe learns a dataset-conditioned learning-to-rank surrogate from historical runs, capturing within-dataset preferences and providing transferable guidance toward promising regions of the configuration space. Online, for a new dataset, AutoPipe uses the offline guidance to steer Bayesian optimization and models dataset-specific deviations with a Gaussian-process residual surrogate. To reduce evaluation cost, each trial is early-stopped and scored by a learned predictor that maps early training signals to a low-cost proxy for final post-training performance. Experiments on biomedical reasoning tasks show that AutoPipe achieves performance comparable to full-fidelity HPO baselines at a fraction of their online cost, while outperforming offline-only approaches and early-stopped HPO variants.