Hey there, fellow researchers and language enthusiasts! Today, we dive into the fascinating world of Biomedical Named Entity Recognition (NER) and explore how we can supercharge our models' inference time without sacrificing accuracy.
NER tasks have witnessed remarkable progress thanks to transformer-based models like PubMedBERT and BioBERT. These models have set new standards for accuracy, but they come with a catch – they can be painfully slow during inference due to their complexity and massive parameters.
In this exciting research paper, we tackle this issue head-on by leveraging the power of Multi-Task Learning (MTL). By jointly training our model for three different biomedical NER tasks, we aim to achieve not only a reduction in inference time but also maintain the accuracy we know and love.
The Dataset: Where It All Begins
Before we get into the nitty-gritty of our experiments, let's talk about the data. Most publicly available datasets for BioNER tasks suffer from data quality issues and lack annotations for crucial entities like tissue and strain. To address this, we rolled up our sleeves and manually curated a gold-standard corpus with cell-line, tissue, and strain entities. Trust us, it was a labor of love!
Experiments: The Magic of Multi-Task Learning
Armed with our finely curated dataset, we embarked on the journey of multi-task learning. We fine-tuned our shared model with three task heads, each dedicated to handling the cell-line, tissue, and strain entities. Our training and testing split was 70:30, and we fine-tuned the model for eight epochs with some nifty learning rate and batch size tricks.
Results and Analysis: The Moment of Truth
The moment of truth has arrived! After all the hard work, we examined the results of our experiments. The MTL model achieved F1 scores of 0.86, 0.71, and 0.87 for cell-line, tissue, and strain, respectively. Excitingly, these scores were on par with those obtained from the single-task models.
But wait, there's more! We didn't just stop at comparing our multi-task model to single-task ones. We also built a multi-class token classification model for cell-line and tissue. Unfortunately, this approach didn't fare as well and left us with some food for thought. Sometimes, one model to rule them all isn't the answer.
The Sweet Taste of Success: 3x Reduction in Inference Time
One of our main goals was to reduce inference time without compromising performance. We're thrilled to announce that our MTL model achieved this feat! When compared to running three separate BERT models, our multi-task model zipped through with a 3x reduction in inference time. Now, that's what we call a win-win!
The Tissue Challenge: Lessons Learned and Future Prospects
We couldn't ignore the elephant in the room – the tissue NER task. Despite our best efforts, tissue entity extraction proved trickier than expected. The F1 score for tissue was lower than the other entities, leaving us pondering the best approach to tackle this challenge. We're determined to explore new avenues, maybe an auxiliary classification task or rule-based methods to improve the overall accuracy.
Conclusion: Multi-Task Learning Shines Bright
As we wrap up our journey, we're excited by the results and the potential of Multi-Task Learning in the Biomedical NER domain. Our MTL model proved its mettle, delivering impressive performance and slashing inference time. But the tissue challenge remains a puzzle yet to be fully solved.
We hope this research inspires further investigations into MTL's magic and how it can revolutionize NER tasks in the biomedical domain. Until next time, happy learning and happy modeling!
References
This paper was part of my work as a Machine Learning intern at Elucidata