StatLLaMA: A multi-stage training framework for building a domain-optimized statistical language model
Jing-Yi Zeng, Guan-Hua Huang
TL;DR
This work tackles the problem of building a domain-optimized statistics LLM using a resource-efficient base model. It systematically compares three multi-stage pipelines that differ in the order and type of domain knowledge injection, instruction tuning, and alignment, culminating in StatLLaMA, a stats-focused model that preserves general reasoning. The study finds that starting from a premise with instruction-following abilities (LLaMA-3.2-3B-Instruct) and applying direct preference optimization yields the best trade-off between domain expertise and general reasoning, with extremely low-intensity downstream fine-tuning to avoid catastrophic forgetting. The results offer a practical blueprint for developing compact, domain-specific LLMs and highlight the importance of data design, alignment strategy, and careful fine-tuning in resource-constrained settings.
Abstract
This study investigates how to efficiently build a domain-specialized large language model (LLM) for statistics using the lightweight LLaMA-3.2-3B family as the foundation model (FM). We systematically compare three multi-stage training pipelines, starting from a base FM with no instruction-following capability, a base FM augmented with post-hoc instruction tuning, and an instruction-tuned FM with strong general reasoning abilities across continual pretraining, supervised fine-tuning (SFT), reinforcement learning from human feedback (RLHF) preference alignment, and downstream task adaptation. Results show that pipelines beginning with a base FM fail to develop meaningful statistical reasoning, even after extensive instruction tuning, SFT, or RLHF alignment. In contrast, starting from LLaMA-3.2-3B-Instruct enables effective domain specialization. A comprehensive evaluation of SFT variants reveals clear trade-offs between domain expertise and general reasoning ability. We further demonstrate that direct preference optimization provides stable and effective RLHF preference alignment. Finally, we show that downstream fine-tuning must be performed with extremely low intensity to avoid catastrophic forgetting in highly optimized models. The final model, StatLLaMA, achieves strong and balanced performance on benchmarks of mathematical reasoning, common-sense reasoning, and statistical expertise, offering a practical blueprint for developing resource-efficient statistical LLMs. The code is available at https://github.com/HuangDLab/StatLLaMA.