Maximizing Intermediate Checkpoint Value in LLM Pretraining with Bayesian Optimization

Q1: Novel contributions compared to Previous Works.

(A) Our Focus: A Search Perspective for Pairwise Merging

• Key Objective. We propose to linearly merge two consecutive checkpoints $\widetilde{\Theta} _t = \lambda_t\, \Theta _t \;+\; (1 - \lambda _t)\, \Theta _{t-1}, \quad \lambda_t \in [\alpha, 1],$ at each stage of large-scale pre-training, while optimally selecting the weight $\lambda_t$ using a Bayesian optimization approach.
• Why Pairwise? As we show empirically, merging adjacent checkpoints tends to preserve or enhance performance when the models are close in training steps (rather than distant). This simplification (from merging multiple checkpoints at once to only the two most recent) drastically reduces the dimensionality of the search to one parameter, $\lambda_t$.
• Bayesian Optimization Rationale. Empirically (Figures 3--5 in our paper), the function $f(\lambda)$ mapping the merging weight $\lambda$ to downstream performance is non-monotonic. Our method:
  1. Treats $f(\lambda)$ as a black-box function (we do not assume it is strictly convex or linear).
  2. Uses Gaussian Process (GP) modeling with an acquisition function (e.g., EI, UCB) to systematically explore $\lambda\in[\alpha,1]$.
  3. Finds the near-optimal or optimal $\lambda$ in a small number of evaluations, even if the function is multi-modal.

(B) Distinctions from LAWA, SWA, and EMA

1. LAWA, SWA, and EMA:
   • Typically rely on an internal training procedure—for instance, SWA requires cyclical or stepwise learning rates and continuously accumulates model parameters from multiple epochs. LAWA uses large learning rates and tail-averaging from the final phase of training. EMA uses an exponential decay factor.
2. Ours:
   • Post-hoc approach: We do not require mid-training gradient data, specialized cyclical learning rates, or entire logs of model state. We only need:
     1. Several discrete checkpoints (e.g., from open-source releases or a standard training pipeline)
     2. A small held-out dataset to measure performance at different $\lambda$.
   • Black-box optimization: Instead of a fixed schedule (as in EMA or LAWA) or a continuous tail average (SWA), we actively search for $\lambda$ that maximizes the model’s performance metric. This is especially powerful for large language models whose training logs might not be fully available, and where advanced Hessian-based or continuous averaging is infeasible.

Hence, although both lines of work share the notion that “averaging can yield flatter minima,” our Bayesian search-based viewpoint and minimal assumptions on training steps make our method (i) more flexible, (ii) directly data-driven for each pairwise merge, and (iii) easy to apply when only discrete checkpoints are at hand.

Q2: Empirical Evidence and Comparisons

(A) Added Comparisons with LAWA, SWA, EMA

In the revised manuscript, we include additional baselines:

Our method either outperforms or closely matches these strategies, reaffirming the practical advantage.

(B) Pre-Training Perplexity

We also conducted additional experiments on pre-training perplexity. As shown in the table below:

Q3: Generalization to “Unseen” Data:

We recognize the reviewer’s concern that large LLMs might have encountered segments of certain benchmarks in their pre-training. We specifically mean “unseen” to refer to:

• A dataset or domain not used to tune $\lambda$. For instance, we tune on a Chinese domain (C-Eval) and evaluate on English tasks (MMLU, GSM8K), ensuring cross-lingual or cross-domain checks.
• Although full de-contamination is challenging, these curated benchmarks (like MMLU, GSM8K) are standard evaluation sets widely regarded as indicative of “test” performance beyond the direct training set.

Ref:

[1] Gaussian process optimization in the bandit setting: No regret and experimental design

def eval(args, subject, model, tokenizer, dev_df, test_df):
    total_loss = 0
    cors = []
    for i in range(test_df.shape[0]):
        prompt_end = format_example(test_df, i, include_answer=False)
        train_prompt = gen_prompt(dev_df, subject, args.ntrain)
        prompt = train_prompt + prompt_end

        # Tokenize prompt
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
        labels = input_ids.clone()

        # Only compute loss on the answer region (prompt_end)
        labels[:, :-len(tokenizer(prompt_end).input_ids)] = -100

        # Forward pass
        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss
        total_loss += loss.item()

        # Optional: measure accuracy by comparing predicted answer to reference
        # ...

    avg_loss = total_loss / len(test_df)
    ppl = np.exp(avg_loss)
    return ppl

Q1: Bounds on the Quadratic Approximation Error.

Our Response:

• We appreciate the request for an expanded derivation. Under the assumption that the Hessian $H(\cdot)$ is Lipschitz continuous with constant $L_H$, consider the third-order Taylor expansion remainder term when approximating $f$ around $w_t$. For $\widehat{w}_t = w_t + \Delta,\quad \text{with} \quad \Delta = \widehat{w}_t - w_t,$ the Taylor remainder $R$ can be written (via the integral form) as $R = \frac{1}{2}\int_0^1 (1-\tau)^2 \big[H(w_t+\tau \Delta) - H(w_t)\big]\, d\tau \,\Delta^{\otimes 2}.$ Taking norms and using the Lipschitz property, we have: $\|H(w_t+\tau \Delta) - H(w_t)\| \le L_H\, \tau\|\Delta\|.$ Therefore, $|R| \le \frac{1}{2}\int_0^1 (1-\tau)^2 L_H\, \tau\, d\tau \,\|\Delta\|^3.$ Evaluating the integral: $\int_0^1 (1-\tau)^2 \tau \, d\tau = \frac{1}{6},$ we obtain the bound $|R| \le \frac{L_H}{6}\|\Delta\|^3.$

• Practical Implications:

  This bound is meaningful when $\|\Delta\| = \|\widehat{w}t - w_t\|$ is small—typically a valid approximation because our analysis focuses on adjacent checkpoints in the pretraining trajectory. If $\|w_t - w{t-1}\|$ is large (and hence $\|\Delta\|$ is large, given the convex combination structure), the higher-order terms may no longer be negligible. In such cases, the quadratic approximation will incur a larger error, and our theoretical guarantees become local. Empirically, our pilot experiments demonstrate that merging distant checkpoints (which effectively produce larger $\|\Delta\|$ degrades performance, confirming that our theory is most valid in the local regime.

---

Q2: Estimating the Information Gain in GP-Based Optimization ($\gamma_T$).

Our Response:

• In one-dimensional Bayesian optimization (with a standard kernel such as the squared exponential or Matérn kernel with smoothness parameter $\nu$, it is well established (e.g., Srinivas et al. (2010)) that $\gamma_T = O(\log T).$ To be precise, for a kernel $k$ on the interval $[\alpha, 1]$ with bounded RKHS norm, one can show that the maximum information gain $\gamma_T$ grows at most logarithmically in the number of evaluations $T$.

---

Q3: Runtime Overhead of GP-Based Bayesian Optimization.

Our Response:

• Plz see the response in reviewer Znoq Q2.

---

Q4: Influence of Kernel Choice in Gaussian Process Modeling on $\gamma_T$ and Regret Bounds.

Our Response:

• The kernel function in GP regression encapsulates our assumptions about the smoothness and structure of the performance function $f(\lambda_t)$. Standard choices such as the squared exponential or Matérn kernels guarantee that $\gamma_T$ grows at most logarithmically (in one dimension), which underpins our theoretical regret bounds.

  In settings where $f(\lambda_t)$ exhibits characteristics such as periodicity or non-stationarity, alternative kernels (e.g., periodic kernels or non-stationary kernels) could more accurately model the underlying function. This improved modeling fidelity could, in turn, lead to a smaller empirical $\gamma_T$ and faster convergence in practice. While our theoretical development is stated in a general manner (i.e., assuming a kernel with bounded information gain), we acknowledge that kernel choice is a crucial hyperparameter. Our preliminary experiments with standard kernels (which are widely used and well understood) already yield superior performance. Future work will explore adaptive kernel selection strategies to further optimize convergence rates.

---

Q5: Clarification on the “Merging Weight” Label in Figure 3.

Our Response:

• The term “merging weight” in Figure 3 refers to the coefficient $\lambda_t$ in the convex combination $\widetilde{\Theta}t = \lambda_t\, \Theta_t + (1-\lambda_t)\, \Theta_{t-1}.$ It quantifies the relative contribution of the later checkpoint $\Theta_t$ versus the earlier checkpoint $\Theta_{t-1}$ during the merging process. The x-axis represents the value of $\lambda_t$ uniformly sampled over the interval $[\alpha, 1]$, where $\alpha$ is a lower bound ensuring that the more recent checkpoint retains a minimum contribution. We will add an explicit explanation in the revision to clarify this point.

Q1: Quadratic Approximation Validity (Equation (15)).

Our Response:

• Our derivation starts from the assumption that, for small perturbations, the performance function $f(\Theta)$ can be locally approximated by a quadratic form: $f(\Theta) \approx f(\Theta_t) + \nabla f(\Theta_t)^\top \Delta + \tfrac{1}{2} \Delta^\top H_t \Delta,$ where $\Delta = \widetilde{\Theta}t - \Theta_t$ and $H_t$ is the Hessian of $f$ at $\Theta_t$. This assumption leverages the smoothness $L_g$-Lipschitz continuity of $\nabla f$ and boundedness of the Hessian (i.e., $\lambda_{\min} I \preceq H_t \preceq \lambda_{\max} I$).

• Empirical Evidence:

  Empirically, Figures 1 and 3 in our paper illustrate smooth and nearly monotonic trends in performance with respect to the merging weight $\lambda_t$. These trends indicate that, within a sufficiently small region between $\Theta_{t-1}$ and $\Theta_t$, the quadratic approximation is reasonable. In addition, our ablation studies show that—despite the nonconvexity of the loss landscape—the averaged performance follows the bounds predicted by our quadratic model (cf. Eq. (17)–(25)).

---

Q2: Runtime Overhead of GP-Based Bayesian Optimization**.**

Our Response:

• Our GP-based Bayesian optimization is designed to address an expensive, derivative-free objective. In our experiments for pairwise merging the search space is one-dimensional (i.e., $\lambda_t \in [\alpha, 1]$). In a typical run we perform on the order of 10-15 evaluations. Empirical measurements show that such an optimization loop incurs an overhead of approximately 30–60 mins on a single 4090 GPU. In contrast, grid or random search methods typically require many more evaluations to reach comparable performance and may require 4–10× more time per experiment.

---

Q3: Relation to Gradient Smoothing and Sharpness-Aware Minimization.

Our Response:

• We acknowledge that recent studies on gradient smoothing and sharpness-aware minimization have focused on explicitly encouraging flat minima during training. Our approach to checkpoint merging can be viewed as complementary. By averaging weights from adjacent checkpoints, we effectively smooth over transient sharp local minima, thereby biasing the final model toward flatter regions of the loss landscape.
• In our convergence and PAC–Bayesian analyses, smoother loss landscapes (i.e., flatter minima) correspond to tighter generalization bounds. Mathematically, if one denotes model sensitivity as $\| \nabla^2 L(\Theta) \|$, then averaging (i.e., merging) reduces the effective curvature: $\widetilde{\Theta}t = \lambda_t \Theta_t + (1-\lambda_t)\Theta{t-1} \quad \Rightarrow \quad \| \widetilde{H}t \| \leq \lambda_t \| H_t \| + (1-\lambda_t)\| H{t-1} \|,$ thereby promoting flatness. We will revise the manuscript to explicitly contrast our method with these related works and add references to the most recent literature in sharpness-aware minimization.

---

Q4: Multi-Checkpoint Merging Versus Pairwise Merging.

Our Response:

• Our experiments (see Figure 7) reveal that while pairwise merging provides substantial improvements, extending the merging to three or four checkpoints often leads to performance dilution. We attribute this to the fact that the most recent checkpoints already capture the latest state of the model, whereas merging with older checkpoints may introduce bias or “stale” information.

  Furthermore, increasing the number of checkpoints increases the dimensionality of the search space exponentially—extending from a one-dimensional segment (two checkpoints) to a two-dimensional triangle (three checkpoints), and ultimately to a closed tetrahedron (four checkpoints)—which exacerbates the computational cost and complexity of the optimization process.

  We have conducted additional experiments employing Bayesian optimization for merging more than two adjacent checkpoints. The results are summarized in the table below:

  Number of Checkpoints	C-Eval
  2 (Our original method)	56.20 ± 0.52
  3	56.27 ± 0.55($\uparrow$ 0.07)
  4	55.34 ± 0.23($\downarrow$ 0.86)

  These results justify our decision to adopt pairwise merging as it achieves a favorable trade-off between performance gains and computational efficiency.

---

Q5: Generalization Across Different Architectures.

Our Response:

• Our empirical results indicate that the proposed merging method consistently improves performance across various architectures (i.e., Baichuan2, DeepSeek, and Pythia models). Although there are minor variations in the optimal merging weights—typically on the order of 0.1–0.7% accuracy—the overall trend of enhanced generalization and reduced variance is maintained.

Q1: Clarification Regarding Figure 3

Our Response:

We appreciate this observation. Figure 3 is not a result of the Bayesian Optimization procedure. Instead, it provides an empirical exploration of how the merged model’s performance changes when we uniformly sample $\lambda$ in $[0, 1]$. In other words, it is a “landscape” study (a brute-force sweep of $\lambda$ ) to:

1. Observe Variability for Large Gaps: When checkpoints differ substantially in performance, $\lambda \approx 1$ typically yields higher performance, because weighting the stronger checkpoint is beneficial.
2. Observe Smooth Behavior for Similar Checkpoints: When checkpoints have closer performance levels, a broad range of $\lambda\in[0,1]$ can give improvements, reflecting that interpolating between two similarly strong models often yields robust results.

Because Figure 3 revealed that performance often worsens if one heavily weights a weaker checkpoint, our practical merging approach restricts $\lambda\in[\alpha, 1]$ with $\alpha>0.5$. Empirically, this domain ensures the more recent (and presumably stronger) checkpoint $\Theta_t$ maintains a dominant contribution. Hence, Figure 3 is purely illustrative—an exploratory “map” of $\lambda$-versus-accuracy—while the actual search for $\lambda$ is done by Bayesian Optimization on the narrower interval $[\alpha,1]$.

Q2: Justification for Eq. 10.

Our Response:

Eq. 10 appears in our PAC-Bayesian generalization bound, where we assume $P \ = \ \mathcal{N}(\Theta_{t-1}, \sigma_P^2 I), \quad Q \ = \ \mathcal{N}(\widetilde{\Theta}t, \sigma_Q^2 I),$ and derive $D_{\mathrm{KL}}(Q \ | | \ P)$. While real neural networks often exhibit correlated parameter distributions, we rely on:

1. Analytical Tractability: A diagonal covariance structure (or isotropic $\sigma^2I$) keeps the KL divergence in closed form: $D_{\mathrm{KL}}(Q || P) \ = \ \frac{\|\widetilde{\Theta}t - \Theta{t-1}\|^2}{2\sigma^2}.$ This simplification is standard in many PAC-Bayes treatments and Bayesian NN approximations (e.g., “mean-field” approximations).
2. Boundedness: Even if true covariances are not diagonal, using a diagonal approximation typically overestimates correlations (or lumps them into the diagonal variance), making the bound somewhat looser but still valid as an upper bound.
3. Consistency with Empirical Practice: Prior works (e.g., [1]) frequently adopt similar diagonal or isotropic assumptions for theoretical clarity and to yield tractable bounds.

Q3: Part of Uncertainty Bounds in Table 1.

Our Response:

We have now run multiple-seed experiments (e.g., 5 random seeds) for each method and checkpoint merge. We report both the mean accuracy and 95% confidence intervals (CIs). Below is an illustrative example for C-Eval (5-shot):

These intervals show that differences are both statistically and practically meaningful. We will include such CIs for all key results in the revised manuscript.

Q4: Miscellaneous Comments and Reference Corrections.

Our Response:

Thank you for pointing out these reference issues. We have carefully reviewed and will correct all citation and reference errors.

Q5: Details on the Optimization Trajectories in Bayesian Optimization

Our Response:

We have added figures for model-fit plot at anonymous link: https://anonymous.4open.science/r/checkpoint-47F1/bayse.png that illustrate the Bayesian Optimization process:

• Posterior Mean & Confidence Intervals: After each iteration, we plot the Gaussian Process posterior (mean + confidence interval) over $\lambda$, and show which $\lambda$ was selected next.
• Number of Iterations: Typically, we only need about 10–15 iterations for a 1D problem to converge near an optimal $\lambda$.

Restricting $\lambda\in[\alpha,1]$ with $\alpha>0.5$ is a pragmatic choice. If users want more flexibility (e.g., to incorporate knowledge from an older checkpoint more strongly), they can lower $\alpha$. In practice, $\alpha\approx 0.5$ balanced well the merging of older vs. newer checkpoints, based on pilot experiments that showed merges with $\lambda<0.5$ rarely helped.

Ref:

[1] Dziugaite G K, Roy D M. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data[J]. arXiv preprint arXiv:1703.11008, 2017.