LLM-Lasso: A Robust Framework for Domain-Informed Feature Selection and Regularization
Abstract
We introduce LLM-Lasso, a novel framework that leverages large language models (LLMs) to guide feature selection in Lasso regression. Unlike traditional methods that rely solely on numerical data, LLM-Lasso incorporates domain-specific knowledge extracted from natural language, enhanced through a retrieval-augmented generation (RAG) pipeline, to seamlessly integrate data-driven modeling with contextual insights. Specifically, the LLM generates penalty factors for each feature, which are converted into weights for the Lasso penalty using a simple, tunable model. Features identified as more relevant by the LLM receive lower penalties, increasing their likelihood of being retained in the final model, while less relevant features are assigned higher penalties, reducing their influence. Importantly, LLM-Lasso has an internal validation step that determines how much to trust the contextual knowledge in our prediction pipeline. Hence it addresses key challenges in robustness, making it suitable for mitigating potential inaccuracies or hallucinations from the LLM. In various biomedical case studies, LLM-Lasso outperforms standard Lasso and existing feature selection baselines, all while ensuring the LLM operates without prior access to the datasets. To our knowledge, this is the first approach to effectively integrate conventional feature selection techniques directly with LLM-based domain-specific reasoning.
Cite
@article{arxiv.2502.10648,
title = {LLM-Lasso: A Robust Framework for Domain-Informed Feature Selection and Regularization},
author = {Erica Zhang and Ryunosuke Goto and Naomi Sagan and Jurik Mutter and Nick Phillips and Ash Alizadeh and Kangwook Lee and Jose Blanchet and Mert Pilanci and Robert Tibshirani},
journal= {arXiv preprint arXiv:2502.10648},
year = {2025}
}
Comments
21 pages, 16 figures