Checkpoint Merging via Bayesian Optimization in LLM Pretraining

2024-03-28

arXiv.org

AI Summary

To address the high computational and environmental costs of large language model (LLM) pretraining, this paper proposes a Bayesian optimization–based checkpoint weighting and merging method that automatically identifies optimal fusion weights across multiple checkpoints sharing the same training trajectory. Our approach is the first to apply Bayesian optimization to LLM checkpoint merging, treating pretraining loss as a black-box objective function—enabling strong cross-task and cross-domain generalization without domain-specific fine-tuning. Experiments demonstrate that the method significantly outperforms single-checkpoint baselines on multi-task evaluation benchmarks. It achieves performance gains comparable to extended training, yet incurs only minimal additional overhead. Moreover, it exhibits low dependence on held-out data and robust generalization across diverse downstream tasks and domains.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The rapid proliferation of large language models (LLMs) such as GPT-4 and Gemini underscores the intense demand for resources during their training processes, posing significant challenges due to substantial computational and environmental costs. To alleviate this issue, we propose checkpoint merging in pretraining LLM. This method utilizes LLM checkpoints with shared training trajectories, and is rooted in an extensive search space exploration for the best merging weight via Bayesian optimization. Through various experiments, we demonstrate that: (1) Our proposed methodology exhibits the capacity to augment pretraining, presenting an opportunity akin to obtaining substantial benefits at minimal cost; (2) Our proposed methodology, despite requiring a given held-out dataset, still demonstrates robust generalization capabilities across diverse domains, a pivotal aspect in pretraining.

Problem

• Reducing computational and environmental costs in LLM pretraining
• Optimizing checkpoint merging via Bayesian exploration
• Enhancing generalization across domains with minimal resources

Innovation

• Checkpoint merging in LLM pretraining
• Bayesian optimization for weight search
• Robust generalization across diverse domains