LLM-Lasso: A Robust Framework for Domain-Informed Feature Selection and Regularization

TL;DR

LLM-Lasso introduces a novel framework that uses large language models to incorporate domain knowledge into feature selection and regularization, improving robustness and performance in biomedical applications.

Contribution

It is the first method to integrate LLM-derived domain insights with traditional Lasso feature selection, enhancing robustness and interpretability.

Findings

Outperforms standard Lasso and baselines in biomedical case studies

Effectively integrates domain knowledge without dataset access

Addresses robustness against LLM hallucinations

Abstract

We introduce LLM-Lasso, a novel framework that leverages large language models (LLMs) to guide feature selection in Lasso $\ell_1$ 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.