MediSwift: Efficient Sparse Pre-trained Biomedical Language Models
Vithursan Thangarasa, Mahmoud Salem, Shreyas Saxena, Chen-Yu Leong, Joel Hestness, Sean Lie
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.
BibTeX
@inproceedings{thangarasa-etal-2024-mediswift,
title = "{M}edi{S}wift: Efficient Sparse Pre-trained Biomedical Language Models",
author = "Thangarasa, Vithursan and
Salem, Mahmoud and
Saxena, Shreyas and
Leong, Chen-Yu and
Hestness, Joel and
Lie, Sean",
editor = "Ku, Lun-Wei and
Martins, Andre and
Srikumar, Vivek",
booktitle = "Findings of the Association for Computational Linguistics: ACL 2024",
month = aug,
year = "2024",
address = "Bangkok, Thailand",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2024.findings-acl.14/",
doi = "10.18653/v1/2024.findings-acl.14",
pages = "214--230"
}