MediSwift: Efficient Sparse Pre-trained Biomedical Language Models (2024)

Abstract: Large language models (LLMs) are typically trained on general source data for various domains, but a recent surge in domain-specific LLMs has shown their potential to outperform general-purpose models in domain-specific tasks (e.g., biomedicine). Although domainspecific pre-training enhances efficiency and leads to smaller models, the computational costs of training these LLMs remain high, posing budgeting challenges. We introduce MediSwift, a suite of biomedical LMs that leverage sparse pre-training on domain-specific biomedical text data. By inducing up to 75% weight sparsity during the pre-training phase, MediSwift achieves a 2-2.5x reduction in training FLOPs. Notably, all sparse pre-training was performed on the Cerebras CS-2 system, which is specifically designed to realize the acceleration benefits from unstructured weight sparsity, thereby significantly enhancing the efficiency of the MediSwift models. Through subsequent dense fine-tuning and strategic soft prompting, MediSwift models outperform existing LLMs up to 7B parameters on biomedical tasks, setting new benchmarks w.r.t efficiency-accuracy on tasks such as PubMedQA. Our results show that sparse pre-training, along with dense finetuning and soft prompting, offers an effective method for creating high-performing, computationally efficient models in specialized domains.

heading: 1 Introduction
text: The landscape of large language models (LLMs) has been predominantly shaped by efforts aimed at creating generalized models (Achiam et al., 2023; Touvron et al., 2023a; Zhang et al., 2022; Shoeybi et al., 2020), trained on diverse datasets that encompass a wide array of topics and domains, such as Pile (Gao et al., 2020), MassiveText (Hoffmann
*Work done while at Cerebras. †Corresponding author.
et al., 2022) and RedPajama (Computer, 2023). While these comprehensive datasets have included data in specialized domains (e.g., programming code (Chen et al., 2021) and PubMed Central (Gao et al., 2020)), the overarching goal has been to forge LLMs capable of broad applicability. However, recent efforts in training models on domain-specific data are emerging, with these smaller, specialized models surpassing general-purpose ones in domain focused tasks, especially in science (Taylor et al., 2022) and medicine (Luo et al., 2022; Bolton et al., 2021). This has sparked a renewed interest in the development of LLMs tailored to specific domains,
214
suggesting a promising avenue for enhancing compute efficiency and model performance w.r.t evaluation metrics on downstream tasks.
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., MEDITRON (Chen et al., 2023), BioGPT (Luo et al., 2022), BioMedLM (Bolton et al., 2021)) 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 Central1 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 (Nori et al., 2023b; Reynolds and McDonell, 2021; Peng et al., 2023; Liu et al., 2022b) and full fine-tuning of the weights (Ziegler et al., 2019; Cohen et al., 2022; Hu et al., 2021), 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 (Dao et al., 2022b; Jaszczur et al., 2021), quantized optimization (Tang et al., 2021), low-rank factorization (Lialin et al., 2023), and sequence-level curriculum learning (Li et al., 2022). 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-
1https://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 (Frankle and Carbin, 2018; Ma et al., 2022) and (2) its difficulties in accelerating on hardware optimized for dense computations (e.g., Nivida GPUs and Google TPUs) (Hooker, 2020). Additionally, sparse pre-training of LLMs typically leads to an accuracy loss from optimization challenges in sparse networks (Evci et al., 2019), but previous studies have shown that transitioning from sparse-to-dense training can effectively address these problems (Thangarasa et al., 2023a; Dao et al., 2022a). 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 (Gale et al., 2019; Elsen et al., 2019; Ashby et al., 2019; Tang et al., 2023) 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 (Cerebras, 2023; Lie, 2022b), 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 (NeuralMagic, 2021; Gupta, 2024) 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., PubMedQA (Jin et al., 2019) for question answering and HoC (Baker et al., 2016) 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:
1. We introduce MediSwift, a family of biomedical language models in three sizes (Med, Large, and XL), and extend this by introducing both dense and sparse variants pre-trained with 50% and 75% weight sparsity. This diversification balances computational efficiency with model effectiveness in biomedical applications, offering options for different computational resource needs.
2. To our knowledge, this is the first study to highlight the benefits of sparse pre-training on biomedical texts, achieving significant computational savings. We show inducing 75% weight sparsity into MediSwift models, results
in a 2.5x reduction in training FLOPs, while improving efficiency-accuracy trade-offs in tasks like PubMedQA.
3. We demonstrate that despite the potential for sparse pre-training to reduce model accuracy, dense fine-tuning combined with soft prompting can effectively regain performance on taskspecific fine-tuning. Specifically, 50% sparse MediSwift-XL achieves a new state-of-the-art with 76.3% accuracy on PubMedQA, surpassing existing models up to 7B parameters.

heading: 2 Methodology
text: 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.

heading: 2.1 Autoregressive Language Modeling
text: Dense Pre-training Autoregressive LMs predict a series of tokens by making each token’s prediction dependent on the ones before it, similar to a Markov chain process. This method follows core principles of language modeling, aiming to understand the pattern of token sequences unsupervisedly from a corpus of text data. Consider an unsupervised corpus U consisting of tokens u1, u2, . . . , u|U|, with |U| denoting the corpus’s total token count. Our objective is to enhance the model’s ability to predict sequences by maximizing the likelihood of the observed sequences, formulated as the sum of the log probabilities of each token given its preceding context within a window of size k. The mathematical representation of this objective is as follows:
L(U) = |U|∑
i=1
log(p(ui|ui−k:i−1, θ)),
where θ denotes the neural network’s parameters, encapsulating the dense configuration of the net-
work’s architecture. The context window, k, determines the number of preceding tokens used for current token prediction. The neural network, parameterized by θ ∈ RN , where N is the total parameter count, aims to optimize these parameters across all layers L, with each layer l having its own set of parameters θl.
Sparse Pre-training Building upon our framework for dense pre-training, we introduce weight sparsity into the model, specifically to improve the computational efficiency. We achieve this by methodically reducing the number of active connections within each layer l of the model by a predefined sparsity level sl ∈ (0, 1), effectively rendering a portion of the network’s parameters inactive. This process yields a network with (1− sl)Nl active parameters per layer, where Nl denotes the original number of parameters in layer l. The overall sparsity is quantified by S, which represents the ratio of inactive parameters to the total parameter count of the initially dense model, calculated as S = ∑L
l slNl N . To apply sparsity effectively,
we employ a binary mask m ∈ {0, 1}|θ| to the model’s initial parameters θ0, resulting in a sparse parameter set m ⊙ θ0. Our approach to inducing sparsity involves random parameter pruning, a process where S percentage of the model’s weights are randomly set to zero at initialization. This mask effectively segregates the parameters into active (1) and inactive (0) states, thereby establishing a sparsity-induced version of the language model that aims to minimize a slightly modified objective:
L(U) = |U|∑
i=1
log(p(ui|ui−k:i−1,m⊙ θ)). (1)
We leverage the GPT-3 (Brown et al., 2020a) architecture for the MediSwift biomedical language model, training it with the AdamW (Loshchilov and Hutter, 2017) optimizer on a curated biomedical dataset, following the objective shown in Eq. 1 for j iterations to obtain parameters θj . This pretrained model is then fine-tuned for specific supervised tasks in the biomedical domain. GPT-3 was selected for MediSwift due to its versatility in NLP tasks and suitability for specialized domains like biomedicine. Its autoregressive nature and capability to produce contextually relevant text align with our methodology of sparse pre-training, dense fine-tuning, and soft prompting to enhance efficiency and task performance. Additionally, GPT-3
was well-supported on the Cerebras CS-2 hardware, with optimized kernels for sparse training using unstructured weight sparsity, further enhancing our model’s computational efficiency.
Our approach is fundamentally model-agnostic and adaptable to various LLMs. The core components of our methodology (i.e., sparse pre-training to reduce computational burden, followed by dense fine-tuning and soft prompting to regain or enhance task performance) can be broadly applied across different LLMs, including newer iterations of GPT.

heading: 2.2 Dense Fine-tuning and Soft Prompting
text: In this section, we detail the adaptation of our pretrained MediSwift model for tasks like biomedical question answering (QA) and document classification using dense fine-tuning and soft prompting. We align tasks with varying output formats to our pre-training format by converting task labels into natural language sequences (Li and Liang, 2021; Hu et al., 2021). This method avoids structured formats and special tokens, ensuring semantic coherence and making full use of the natural language corpus MediSwift was trained on.
Following Luo et al. (2022), each downstream fine-tuning task is represented by a training set consisting of source-target pairs defined as: Z = {(x1, y1), (x2, y2), . . . , (x|x|, y|y|)}, where both x and y are sequences of tokens. For example, in question answering (e.g., PubMedQA), x corresponds to the question and reference context description, and y the categorical answer to the question; in biomedical document classification (e.g., Hallmarks of Cancers corpus), x is the text passage and y corresponds to the hallmarks of cancer.

heading: 2.2.1 Dense Fine-tuning
text: We begin fine-tuning with parameters θj set at their pre-trained values, adjusting them by a taskspecific increment ∆θ with the same dimensionality, |∆θ| = |θ|. Unlike prior efforts that aimed at parameter efficiency for easier model deployment (Hu et al., 2021; Dettmers et al., 2023), we prioritize reducing pre-training computational costs via unstructured weight sparsity and enhance network representation by adopting dense fine-tuning (Thangarasa et al., 2023a). This approach overcomes sparse optimization challenges by reactivating previously inactive weights during the dense fine-tuning phase, thus enhancing the model’s capacity. By removing the sparsity mask m, we allow ∑L l sl·Nl weights to be reactivated
and initialize them to zero—a method proven more effective than other initialization strategies (e.g., scaled uniform or normal distribution initializations (Evci et al., 2020)). The network is then densely updated to optimize the loss function:
L(Z) = ∑
(x,y)∈Z log
|y|∏
t=1
(2)
p(yt|(x1, . . . , xt−1), θj +∆θ)

heading: 2.2.2 Soft Prompting
text: Previous research on language models in the biomedical field has mainly focused on fine-tuning for domain-specific tasks (Bolton et al., 2021). Recently, there has been a move towards improving biomedical NLP task performance through prompt engineering (Luo et al., 2022; Nori et al., 2023b; Yagnik et al., 2024). Drawing inspiration from these advancements, our approach integrates prompt-based techniques into the fine-tuning phase of MediSwift. More precisely, we adopt the soft prompting methodology as described by Luo et al. (2022), aiming to refine our model’s capability in understanding and processing biomedical text. Similar to existing work on prompt tuning (Brown et al., 2020b; Liu et al., 2021b; Lester et al., 2021), to integrate soft prompts, we insert virtual tokens between the source and target sequences, thus modifying the loss function L(Z) in Eq. 2. This adjustment accounts for the [source; prompt; target] sequence structure, impacting the model’s learning and inference. Let P denote the prompt, consisting of a sequence of virtual tokens ∈ {v1, v2, . . . , vn}, where n is the number of virtual tokens, and these tokens are represented by continuous embeddings. The modified Eq. 2, reflecting the inclusion of the prompt and its positioning, can be formalized as:
L(Z) = ∑
(x,y)∈Z log
|y|∏
t=1
(3)
p(yt|([x;P]; y<t), θj +∆θ)
Through this multi-faceted approach, our pretraining and fine-tuneing method for MediSwift not only addresses the computational efficiency challenges in LM training, but also leverages the capabilities of in-domain pre-trained LMs to improve performance on biomedical NLP tasks.

heading: 3 MediSwift Biomedical Pre-training
text: 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.

heading: 3.1 PubMed Papers and Abstracts
text: MediSwift is an in-domain biomedical language model, drawing its strength from an exclusive pretraining regimen focused solely on biomedical textual data. Its foundation lies in the vast repository of available open-access medical research papers and abstracts found in PubMed Central (PMC) (National Library of Medicine. 2003–2023), similar to the approaches used in prior models (e.g., Meditron (Chen et al., 2023), BioGPT (Luo et al., 2022), BioMedLM (Bolton et al., 2021)).
PMC consists of 4.91M full-text papers, and PubMed and PMC Abstracts comprise of 16.1 million papers (see Table 1). Moreover, we gathered the most recent PubMed entries, updated prior to 2023, directly from the official website1, utilizing the official scripts for PubMed Abstracts2 and PubMed Central3. Similar to Luo et al. (2022) and Chen et al. (2023), we filter out empty items containing solely titles without accompanying abstracts. Furthermore, prior works have show the significance of in-domain vocabulary for improving performance of specialized LMs (Gu et al., 2021; Wu et al., 2023b; Mielke et al., 2021), a critical step that is often overlooked. Therefore, inspired by Luo et al. (2022) and (Bolton et al., 2021), instead of using the standard GPT-3 vocabulary, we learned the vocabulary directly from the biomedical corpus. Employing Moses (Koehn et al., 2007) tokenization followed by byte pair encoding (BPE), we segment the corpus into word pieces and learn the vocabulary; resulting in a size of 42,384. By exclusively pre-training with biomedical texts and using a specialized vocabulary, MediSwift improves the efficiency-accuracy frontier, as empirically shown on biomedical tasks in Section 4.
2https://github.com/thoppe/The-PilePubMed?tab=readme-ov-file
3https://github.com/EleutherAI/pilepubmedcentral

heading: 3.2 MediSwift Pre-training
text: Pre-training Experimental Details We pre-train and benchmarked MediSwift in-domain biomedical language models at 3 sizes: 302M, 510M and 1.21B. All MediSwift models are pre-trained from scratch using the Cerebras CS-24, taking advantage of its ability to accelerate training with unstructured sparsity. At the time of the study, the specialized kernels of Cerebras CS-2 were designed to facilitate training with static unstructured sparsity (refer to Appendix B for additional details). In the pretraining phase of the MediSwift models, 50%, and 75% sparsity levels are explored, aside from their respective dense counterparts. The pre-training of MediSwift models are conducted on a single CS-2 for a total of 200,000 steps, with a batch size of 512 and a maximum sequence length of 1024 tokens, resulting in approximately 104.86B tokens processed in total (see Appendix A.1 for additional pre-training experimental setup details).

heading: 3.3 Sparse Pre-trained MediSwift
text: While there are several advanced sparse training techniques (Evci et al., 2020; Mocanu et al., 2018; Liu et al., 2021a), for simplicity, in this work, we adopt static sparsity, namely random pruning, for the sparse pre-training of MediSwift models. This approach mandates a uniform distribution of sparsity levels throughout each layer, irrespective of the layer’s parameter count or its FLOPs. Specifically, the scope of our sparsification process is confined to all dense linear layers within the network, including both matrices within the multilayer perceptron (MLP) module—namely, the intermediate layer and the MLP output projection, as well as the four weight matrices integral to the self-attention (Vaswani et al., 2017) mechanism:
4https://docs.cerebras.net/en/2.1.1/wsc/ how_to_guides/sparsity.html
query, key, value, and attention output projection. Notably, we ensure that the embeddings, Layer Normalization (Ba et al., 2016) components and biases are kept dense.

heading: 3.4 Analysis on Pre-trained Models
text: In Figure 2, we illustrate the training loss curves for the MediSwift-XL model for both the dense and sparse configurations at s ∈ {50%, 75%}. The 50% sparse model’s training loss closely follows that of the dense MediSwift-XL, showcasing minimal deviation throughout the training process. However, a noticeable divergence is observed at 75% sparsity, where the final training loss slightly lags behind that of the dense counterpart. Interestingly, when comparing the 75% sparse MediSwiftXL model with the dense MediSwift-Med model, the former, despite sharing the same number of nonembedding parameters, it achieves a lower training loss (refer to Table 2). This observation aligns with previous findings that larger, albeit sparser, models can surpass their smaller, densely parameterized equivalents in terms of performance (Thangarasa et al., 2023b; Liu et al., 2022a; Ramanujan et al., 2020; Golubeva et al., 2021). This superiority is further supported by the improved accuracy of the 75% sparse model on biomedical NLP tasks, when compared to the MediSwift-Med dense model (see Section 4), highlightling the benefits of training larger sparse models in comparison to smaller dense ones.
In Table 2, we provide an analysis of the com-
putational efficiency achieved through sparse pretraining in our proposed MediSwift architectures, namely MediSwift-Med and MediSwift-XL. Here, we quantify the total FLOPs required for both forward and backward propagations during the pretraining phase of these models. For the MediSwiftXL model, attention and vocabulary embeddings represent 13.3% and 6.8% of total FLOPs, respectively, hence highlighting the computational savings at the 1.21B parameter scale. Sparse pre-training with 75% sparsity reduces FLOPs by slightly over 2.5x compared to its dense counterpart. The smaller MediSwift-Med model has a higher FLOP percentage for attention and embeddings, hence achieving a 2x reduction at the same sparsity level. This indicates that FLOP savings increase with model size, demonstrating that larger models (Hoffmann et al., 2022; Kaplan et al., 2020; Hestness et al., 2017) can potentially benefit more from sparse pre-training. In addition, we emphasize that the total FLOPs required for fine-tuning these models account for a minor fraction of the overall pre-training FLOPs, reinforcing the efficiency of our approach in scaling to larger model sizes while conserving computational resources.

heading: 4 Fine-tuning on Biomedical tasks
text: 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.

heading: 4.1 Question Answering with PubmedQA
text: We assess MediSwift’s performance on the PubMedQA (Jin et al., 2019) dataset, which is derived from PubMed abstracts and includes three subsets: PQA-A, PQA-U, and PQA-L. We adhere to the original train/val/test splits, focusing on the PQAL test set for the final evaluation. Our approach utilizes multi-stage fine-tuning and soft prompting with continuous embeddings of length n = 9, which was shown to perform the best in terms of accuracy (Luo et al., 2022). Following established preprocessing methods, we format the data into [source, target] sequences, each consisting of a question, reference context, long answer, and a categorical label, [yes/no/maybe], for the answer. The performance is measured by classification accuracy, particularly under the challenging reasoning-required setting (Jin et al., 2019), where the model predicts based on the question and context without the long answer.
In Table 3, we demonstrate that across all sizes, MediSwift improves the pareto frontier in PubMedQA accuracy, notably with the dense MediSwift-XL model setting a new benchmark while being significantly smaller, at 5.8x less size than PMC-LLaMA. This trend continues with the 50% and 75% sparse variants of MediSwift, which surpass other language models of comparable or larger sizes. Specifically, the 75% sparse MediSwift-XL exceeds BioMedLM’s performance by 1.0% while being approximately 8.9x smaller. Furthermore, within the MediSwift family, the larger yet sparse 75% MediSwift-XL demonstrates superior performance over the smaller dense MediSwift-Med by 1.2%, despite both models shar-
5https://github.com/microsoft/BioGPT/tree/ main/examples/QA-PubMedQA
ing the same pre-training parameters.
Prompt Tuning Ablation Study To assess the effectiveness of soft prompting, we have conducted the requested analysis and provide the results below. Our study focused on evaluating the impact of soft prompting on the performance of our MediSwift models during the fine-tuning phase, across both dense and sparse configurations (see Table 1). The approach involves steering the pre-trained language model by appending several additional virtual tokens as prompts before the text. These continuous embeddings, distinct from the main text, are randomly initialized and learned end-to-end on down-
stream tasks, making them task-specific. Unlike prefix tuning (Li and Liang, 2021), we strategically place the virtual tokens not at the very beginning of the source input but specifically before the target sequence, resulting in a final sequence structure of [source; prompt; target] (as described in Section 2.2.2). The application of soft prompting yields a 0.7% accuracy increase in the dense MediSwift-XL model, with significant gains also seen in the 50% and 75% sparse configurations. This highlights soft prompting’s role in refining outputs through task-specific conditioning, boosting accuracy in biomedical question answering tasks.

heading: 4.2 Document Classification on HoC
text: We examine the Hallmarks of Cancers (HoC) corpus (Baker et al., 2016), comprising 1580 PubMed abstracts annotated for ten cancer hallmarks. We tackle a document classification task, assigning documents to predefined single or multi-label categories, and using MediSwift to generate label words. We follow the established train/dev/test splits of 1108/157/315 (Gu et al., 2021). Similar to Luo et al. (2022), we employ a continuous embedding of length n = 1 as the prompt, and we incorporate labels into the target sequence.
The performance is evaluated using the microF1 score, allowing direct comparison with prior models and demonstrating our method’s effectiveness. In Table 5, the dense MediSwift-XL model outperformed all similarly sized models in microF1 score, with its 50% and 75% sparse showing very competitive results, emphasizing sparse pretraining’s balance of computational efficiency and accuracy. This further showcases the potential of sparsity in optimizing language model performance for biomedical applications.

heading: 5 Related Work
text: Sparse Training for Language Models Sparse weight training for language models (LM) have emerged as a promising avenue to address the computational intensity of training large models. Recent work has explored various sparse training methods (Thangarasa et al., 2023a; Dao et al., 2022a; Chen et al., 2022), aiming to maintain or enhance model performance while significantly reducing computational requirements. Techniques such as pruning (Chen et al., 2020), sparse activations (Mirzadeh et al., 2024), along with the development of specialized software (NeuralMagic, 2021; Gupta, 2024) and hardware (Lie, 2022b,a; Dietrich et al., 2021) have been pivotal. We build on these foundations, focusing on optimizing sparse weight training strategies specifically for domainspecific LMs, pushing the boundaries of efficiency. Biomedical Language Models The evolution of language models for medical applications has progressed from adapting encoder-only architectures like BERT (Devlin et al., 2018), using biomedical corpora (Lee et al., 2019; Gu et al., 2021), to incorporating strategies like document links (Yasunaga et al., 2022b) and knowledge graphs (Yasunaga et al., 2022a). The shift towards autore-
gressive generative models, such as GPT (Brown et al., 2020b) and Llama (Touvron et al., 2023a), for pretraining on medical texts has led to significant advancements (Wu et al., 2023a; Luo et al., 2022; Bolton et al., 2021). Recent scaling efforts include GatorTronGPT with 20B parameters (Yang et al., 2022), as well as Clinical-Camel (Toma et al., 2023), MEDITRON (Chen et al., 2023) and Med42 (Christophe et al., 2023), based on Llama-270B (Touvron et al., 2023b), focusing on mixed clinical and general English texts. Our work diverges from works that scaled up medical LMs by introducing weight sparsity into the pre-training of biomedical LMs. This reduces the computational costs typically associated with large-scale models, thereby improving the balance between efficiency and accuracy in the medical domain. Prompting for Biomedical Language Models Recent research has shifted towards prompt engineering to enhance language models’ performance on biomedical tasks, such as BioGPT’s (Luo et al., 2022) use of soft prompt-tuning and Medprompt’s (Nori et al., 2023b) innovative prompting techniques for generalist foundation models. Liévin et al. (2023) and Yagnik et al. (2024) analyzed the effectiveness of prompting in the medical domain and showed that it can improve metric scores. However, combining task-specific finetuning with prompting strategies, as seen in MedPaLM 2 (Singhal et al., 2022, 2023), yields competitive results on challenging biomedical tasks. Our work extends this by integrating task-specific fine-tuning and soft prompting to address model accuracy loss during sparse pre-training, effectively achieving efficiency gains with minimal accuracy degradation on biomedical tasks.

heading: 6 Conclusion
text: 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.
Limitations
Our work on MediSwift represents a significant leap forward in developing efficient domainspecific LLMs, particularly in biomedicine, by utilizing sparse pre-training to strike a fine balance between computational efficiency and accuracy. While we have initially focused on static sparse pre-training, the emerging field of dynamic sparse training (DST) holds great promise for further improvements (Evci et al., 2020; Mocanu et al., 2018; Liu et al., 2021a).
DST offers an exciting avenue for optimizing sparsity patterns dynamically, potentially elevating model quality and training efficiency to new heights. Although the implementation of DST requires advanced software and hardware support for unstructured sparse computations—capabilities that were beyond our current scope—this innovative approach represents an interestingly opportunity for future research. As support for unstructured sparse training evolves with ML softwarehardware co-design, we anticipate these advancements will enable us to harness DST, paving the way for even more high-quality and efficient domain-specific LLMs.
Ethics Statement
While MediSwift represents a significant advancement in encoding medical knowledge from sources of high-quality evidence, it is important to note that it has not been fully adapted to deliver this knowledge in a manner that is appropriate, safe, or within the actionable constraints required by medical professionals. Therefore, we strongly recommend against deploying MediSwift directly in clinical or medical applications without thorough alignment with specific use cases.
Moreover, additional testing is crucial, including the conduct of randomized controlled trials in realworld practice settings, to ensure the model’s recommendations are reliable and beneficial in practical healthcare environments. This cautionary approach emphasizes the importance of bridging the gap between technological capabilities and the nuanced requirements of medical practice to ensure patient safety and efficacy of care.

heading: A Experimental Setup and Hyperparameter Details
text: A.1 Pre-training on Biomedical Data To train all MediSwift models, we use the AdamW optimizer (Loshchilov and Hutter, 2017) with a peak learning rate set at 2×10-4, β1 =0.9, β2 =0.95 and ϵ =10-8. A linear warmup period, amounting to 10% of the total training steps, is employed before transitioning to a cosine decay schedule, with the learning rate decreasing to a minimum of 10% of the peak value (i.e., 2×10-5). In Table 6, we provide details on the size and architecture configurations of the MediSwift models we pre-trained. Here, nparams is the total number of trainable parameters, nlayers is the number of decoder layers, and dmodel is the base size of the model. The feedforward bottleneck is four times the base size, i.e., dff = 4 × dmodel. Finally, nheads are the number of attention heads and dhead is the dimension of each attention head. The context window size is set to 1024.
Following the training FLOPs calculation described in Hoffmann et al. (2022), we compute the total pre-training FLOPs for the dense and sparse variants of MediSwift-Med, Large and XL, and report them in Table 7, along with their relative FLOPs reduction over the dense baseline. Similar to Appendix F of Hoffmann et al. (2022), we
also include the training FLOPs contributed by the embedding matrices. Additionally, in large models, the contribution of embedding matrices to the overall FLOPs and parameters is minimal.
A.2 PubMedQA Fine-tuning
As mentioned in Section 4.1, the PubMedQA dataset includes three subsets: PQA-A, PQA-U, and PQA-L. We train all of our MediSwift models on the original train/val/test splits for each of these datasets in a multi-stage manner (Jin et al., 2019), both dense and sparse, using AdamW (Loshchilov and Hutter, 2017) and a linear learning rate warmup (i.e., 10% of to the total training steps) followed by a cosine decay schedule for a maximum 5 epochs, and perform early-stopping when the models began to overfit. We perform a grid search to discover an appropriate learning rate that led to the best downstream classification accuracy on each of the three datasets for a given compute budget. More specifically, on the dense baseline and sparse variants, we select the best batch size among {8, 16, 32, 64} and select the best learning rate among {2e-4, 1e-4, 5e-5, 2.5e-5} on the validation set. After training on the final stage (i.e., PQA-L), we evaluate the model on the PQA-L test set using the official evaluation scripts6. All results were averaged across 3 random seeds.
6https://github.com/pubmedqa/pubmedqa
A.3 HoC Fine-tuning
In Section 4.2, we described the Hallmarks of Cancers (HoC) dataset which comprises of 1580 PubMed abstracts with a 1108/157/315 split for train, val and test sets following Gu et al. (2021). We train all of our MediSwift models on the original train/val/test splits for both dense and sparse, using AdamW (Loshchilov and Hutter, 2017) and a linear learning rate warmup (i.e., 10% of to the total training steps) followed by a cosine decay schedule for a total of 100 epochs, and perform early-stopping when the models began to overfit. We perform a grid search to discover an appropriate learning rate that led to the best micro-F1 score for a given compute budget. More specifically, on the dense baseline and sparse variants, we select the best batch size among {16, 32, 64} and select the best learning rate among {8e-5, 4e-5, 2e-5, 1e-5} on the validation set. All results were averaged across 3 random seeds.

heading: B Unstructured Sparsity on Specialized Hardware Accelerators
text: The Cerebras CS-2 system, designed specifically for accelerating deep learning computations, can handle unstructured sparsity efficiently due to its unique architecture (Lie, 2022b, 2021). The CS-2’s
wafer-scale engine, with its vast array of computational cores and on-wafer memory, efficiently manages unstructured sparsity’s irregular memory access, surpassing traditional architectures that often face memory bandwidth constraints. Moreover, the CS-2 has a significant amount of on-chip memory, reducing the need to access external memory. This is crucial for unstructured sparsity, as the irregular access patterns can lead to high latency if data needs to be fetched from off-chip. By keeping more data on-chip, the CS-2 minimizes these latencies. In addition, the system supports fine-grained parallelism, allowing it to execute many small, sparse operations concurrently across its thousands of cores. This is particularly advantageous for unstructured sparsity, as the workload can be distributed across many cores to maintain high utilization. In Figure 3, we highlight the potential realized gains with unstructured weight sparsity on the Cerebras CS-2.

heading: C Author Contributions
text: We provide a summary of each author’s contributions:
• Vithursan Thangarasa led the efforts on developing the method and pre-training the various dense and sparse MediSwift models. He also implemented several parts of the code for fine-tuning and proposed innovative approaches such as dense fine-tuning and soft prompting. Additionally, he wrote the manuscript, ensuring a comprehensive and cohesive presentation of the research.
• Mahmoud Salem assisted in running experiments for both single-phase and multi-phase fine-tuning on PubMedQA. His work was crucial in validating the effectiveness of the proposed methods across different fine-tuning strategies.
• Shreyas Saxena collaborated closely with Vithursan to help design the main method. His insights and feedback were vital in refining the overall approach and ensuring the robustness of the proposed techniques.
• Kevin Leong contributed to the pre-training of both dense and sparse MediSwift models on the Cerebras CS-2. His expertise in utilizing the Cerebras CS-2 hardware was essential for the efficient and effective training of the models.
• Joel Hestness was involved in discussions with Vithursan and provided valuable feedback to improve the manuscript. His contributions helped enhance the clarity of the manuscript.
• Sean Lie frequently met with Vithursan for technical discussions and provided support for bringing up the Cerebras CS-2. His involvement was critical in performing the pretraining of MediSwift models from scratch and ensuring the smooth operation of the hardware.

heading: Future Work
text: Although the implementation of DST requires advanced software and hardware support for unstructured sparse computations—capabilities that were beyond our current scope—this innovative approach represents an interestingly opportunity for future research. As support for unstructured sparse training evolves with ML softwarehardware co-design, we anticipate these advancements will enable us to harness DST, paving the way for even more high-quality and efficient domain-specific LLMs