# Task-Specific Sparse Feature Masks for Molecular Toxicity Prediction with Chemical Language Models

arXiv: https://arxiv.org/abs/2512.11412 · HTML: https://mlsanity.com/paper/2512.11412/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 22 references that are in the map.
- First posted on arXiv 12 Dec 2025.

## Builds on (most cited first, 9 of 9)

- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [MoleculeNet: A Benchmark for Molecular Machine Learning](https://mlsanity.com/paper/1703.00564/) (2017)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [End-to-End Multi-Task Learning with Attention](https://mlsanity.com/paper/1803.10704/) (2018)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [Large-Scale Chemical Language Representations Capture Molecular Structure and Properties](https://mlsanity.com/paper/2106.09553/) (2021)
- [Massively Multitask Networks for Drug Discovery](https://mlsanity.com/paper/1502.02072/) (2015)
- [An Empirical Study of Multi-Task Learning on BERT for Biomedical Text Mining](https://mlsanity.com/paper/2005.02799/) (2020)

## Abstract (arXiv)

Reliable in silico molecular toxicity prediction is a cornerstone of modern drug discovery, offering a scalable alternative to experimental screening. However, the black-box nature of state-of-the-art models remains a significant barrier to adoption, as high-stakes safety decisions demand verifiable structural insights alongside predictive performance. To address this, we propose a novel multi-task learning (MTL) framework designed to jointly enhance accuracy and interpretability. Our architecture integrates a shared chemical language model with task-specific attention modules. By imposing an L1 sparsity penalty on these modules, the framework is constrained to focus on a minimal set of salient molecular fragments for each distinct toxicity endpoint. The resulting framework is trained end-to-end and is readily adaptable to various transformer-based backbones. Evaluated on the ClinTox, SIDER, and Tox21 benchmark datasets, our approach consistently outperforms both single-task and standard MTL baselines. Crucially, the sparse attention weights provide chemically intuitive visualizations that reveal the specific fragments influencing predictions, thereby enhancing insight into the model's decision-making process.