M edi S wift: Efficient Sparse Pre-trained Biomedical Language Models (2024)

Abstract: Large language models (LLMs) are typically trained on general source data forvarious domains, but a recent surge in domain-specific LLMs has shown theirpotential to outperform general-purpose models in domain-specific tasks (e.g.,biomedicine). Although domain-specific pre-training enhances efficiency andleads to smaller models, the computational costs of training these LLMs remainhigh, posing budgeting challenges. We introduce MediSwift, a suite of biomedicalLMs that leverage sparse pre-training on domain-specific biomedical text data.By inducing up to 75% weight sparsity during the pre-training phase, MediSwiftachieves a 2-2.5x reduction in training FLOPs. Notably, all sparse pre-trainingwas performed on the Cerebras CS-2 system, which is specifically designed torealize the acceleration benefits from unstructured weight sparsity, therebysignificantly enhancing the efficiency of the MediSwift models. Throughsubsequent dense fine-tuning and strategic soft prompting, MediSwift modelsoutperform existing LLMs up to 7B parameters on biomedical tasks, setting newbenchmarks w.r.t efficiency-accuracy on tasks such as PubMedQA. Our results showthat sparse pre-training, along with dense fine-tuning and soft prompting,offers an effective method for creating high-performing, computationallyefficient models in specialized domains.

## introduction the landscape of large language models (llms) has been predominantly shaped by efforts aimed at creating generalized models @xcite @xcite , trained on diverse datasets that encompass a wide array of topics and domains, such as pile @xcite , massivetext (hoffmann pubmedqa accuracy in the reasoning-required setting: our dense and sparse mediswift models noticeably outperform other fine-tuned language @xmath0, improving the efficiency-accuracy pareto frontier. in particular, mediswift-xl (1.21b) achieves new state-of-the-art 76.8% accuracy at this size (i.e., being 5.8x smaller than pmc-llama). in addition, sparse pre-trained mediswift-xl models at @xmath1, 75%} outperform other models at similar or larger size. additional details are provided in furthermore, the shift towards specialized llms in the field of medicine is particularly gaining recognition for its capacity to significantly improve the accuracy and effectiveness of these models. this is achieved by closely aligning them with the specific needs and complexities of this specialized area. for example, in biomedicine, models trained on specialized literature (e.g., med-itron @xcite , biogpt @xcite , biomedlm @xcite ) exhibit significant improvements over general ones. this approach involves either continued pre-training on domain-relevant texts or pre-training models from scratch with domain-specific data (e.g., pubmed central 1 and pubmed open-access research papers), emphasizing the accuracy benefits and ability to yield models that are more compute efficient. after pre-training models on domain-specific data, further accuracy enhancements are achieved through prompt-based fine-tuning @xcite @xcite and full fine-tuning of the weights @xcite @xcite , offering a balance between adaptability, efficiency, and task-specific optimization. this dual approach merges the model's pre-trained knowledge with the specific needs of downstream tasks, thereby facilitating a more accurate and effective application of the model's capabilities across a wide array of downstream applications; maximizing the potential of domain-specific llms. although domain-specific llm pre-training offers significant benefits, the computational demands remain a significant challenge, necessitating innovative solutions for broader accessibility and efficiency. techniques proposed to mitigate these burdens include sparse attention @xcite , quantized optimization @xcite , low-rank factorization @xcite , and sequence-level curriculum learning @xcite . sparse attention enhances transformer efficiency but mainly benefits attention layers and may struggle with long-range dependencies. quantized optimization reduces memory footprint and accelerates training with lower-precision arithmetic but introduces quantization noise and relies on hardware support. low-rank factorization saves mem-1 https://pubmed.ncbi.nlm.nih.gov/ ory and computation by approximating weight matrices, risking information loss and requiring architectural changes. sequence-level curriculum learning improves convergence and generalization by progressively increasing difficulty, primarily enhancing training efficiency. among these approaches, weight sparsity emerges as a promising method, distinct from the others in its approach by setting a subset of model parameters to zero, thus reducing the flops needed during training. unlike sparse attention, which optimizes a specific component, or quantized optimization and low-rank factorization, which may introduce trade-offs in model performance, sparse pre-training directly reduces computational costs across the entire model without inherently compromising performance. however, the adoption of sparse training is limited by (1) the challenge of finding optimal sparsity patterns which retain the accuracy of the original dense model @xcite and (2) its difficulties in accelerating on hardware optimized for dense computations (e.g., nivida gpus and google tpus) @xcite . additionally, sparse pre-training of llms typically leads to an accuracy loss from optimization challenges in sparse networks @xcite , but previous studies have shown that transitioning from sparse-to-dense training can effectively address these problems @xcite . our work aims to reduce training flops by employing unstructured weight sparsity in domain-specific llms and recovering lost representational capacity by transitioning to dense weight matrices for downstream fine-tuning. despite weight sparsity not being widely used in real-world applications, advancements in specialized software kernels @xcite @xcite have facilitated inference acceleration with unstructured sparsity. recent developments have shown that the benefits of unstructured weight sparsity can be fully harnessed on specialized hardware, such as the cerebras cs-2 @xcite , for llm training. as sparse training techniques and hardware continue to co-evolve, we anticipate that the reductions in flops will lead to realized sparse acceleration. the latest innovations in software and hardware @xcite are geared towards enabling the widespread adoption of unstructured weight sparsity, offering the potential to achieve higher compression ratios and practical speedups in terms of wall-clock time. building on this momentum, we introduce mediswift, a suite of biomedical language models (lms) available in three sizes: med (302m), large (510m), and xl (1.21b). these models are based on gpt-3 and pre-trained sparsely from scratch on biomedical texts, aimed at reducing the computational costs required for training. we explore the impact of applying 50% and 75% weight sparsity during pre-training, which results in a 2-2.5x reduction in the overall training flops needed. figure 1 summarizes the performance, where our dense and sparse mediswift models noticeably outperform other language models up to 7b parameters. we demonstrate mediswift's capabilities through finetuning on established benchmarks for biomedical natural language processing (nlp) tasks (e.g., pub-medqa @xcite for question answering and hoc @xcite for document classification) showing significant improvements in the balance between efficiency and accuracy. although previous research suggests that sparse pre-training may compromise model accuracy on downstream tasks, our approach incorporates dense fine-tuning with strategic soft prompting to effectively regain performance on specialized tasks. specifically, mediswift-xl (1.21b) sets a new state-of-the-art by reaching 76.8% accuracy, despite being 5.8x smaller than pmc-llama. moreover, mediswift-xl models, pre-trained with 50% and 75% sparsity, surpass the performance of models of similar or greater sizes (e.g., mediswift-xl at 75% sparsity outperforms biomedlm while being almost 9x smaller). our work not only highlights the potential for sparse pre-training to make lm training more economically viable but also sets a new benchmark for efficiency in domain-specific applications of llms. the key contributions are: ## methodology in this section, we formalize our two-phase training framework for mediswift models to reduce computational costs and yet retain model accuracy. initially, we pre-train these models on biomedical data, applying unstructured weight sparsity to reduce the computational training flops. following this, we enhance the model through dense finetuning, reactivating weights to improve adaptability for specific tasks, and incorporate soft prompting to refine responses for task requirements. this efficient approach, combining sparse pre-training, dense fine-tuning, and soft prompting, significantly boosts both model efficiency and performance, as our results demonstrate in section 4. ## mediswift biomedical pre-training this section describes the mediswift pre-training process, including data sources, collection, and preprocessing for biomedical data. we explain the dataset's origins, statistical analysis, and preparation for efficient training. we also compare mediswift models' performance with both dense and sparse pre-training, emphasizing training convergence differences and flops savings. ## fine-tuning on biomedical tasks this section evaluates mediswift's performance on pubmedqa and hoc benchmarks using dense fine-tuning on each variant's specific training set (e.g., fine-tuning and testing with the pubmedqa dataset). during the fine-tuning phase, the total training flops constitute only a minor fraction of the flops expended during pre-training, de-spite the fine-tuning being conducted densely. as a result, the flops consumed during sparse pretraining is proportional to the combined flops of sparse pre-training and dense fine-tuning. following luo et al. (2022), we incorporate soft prompting into our fine-tuning framework by formatting sequences as [source; prompt; target]. this format helps our models better utilize contextual information, demonstrating mediswift's effectiveness and adaptability in medical text analysis. further details on hyperparameters and dataset specifics are provided in appendix a.2 and a.3. we note that all fine-tuning results were averaged across 3 random seeds. ## conclusion in conclusion, mediswift innovates in biomedical language models by combining sparse pre-training with dense fine-tuning and soft prompting, balancing computational efficiency with accuracy. available in med, large, and xl sizes, with 50% and 75% sparsity, mediswift addresses the cost of training models and sets new standards for biomedical tasks like pubmedqa. mediswift-xl, in particular, showcases superior efficiency-accuracy tradeoffs, outperforming models up to 7b parameters. this work exemplifies the potential of sparse pretraining as a cost-effective method for developing specialized, high-performance models, establishing mediswift as a benchmark in biomedical nlp.