LLM Sparsity Prior for Robust Feature Selection

Abstract: LLMs offer a scalable mechanism to elicit domain-informed prior information for high-dimensional variable selection. However, existing methods such as LLM-Lasso are sensitive to weight quality, with performance degrading substantially when LLM-generated weights are inaccurate. To address this challenge, we first introduce a framework for quantifying the quality of LLM-generated weights, enabling rigorous evaluation of LLM-informed methods across varying weight regimes. We then propose the LLM Sparsity Prior (LSP), which integrates LLM-generated weights into the prior inclusion probabilities of Spike-and-Slab and Spike-and-Slab Lasso models via two interpretable hyperparameters governing global sparsity and weight concentration. Hierarchical hyperpriors on these parameters allow the model to dynamically discount uninformative or misleading weights, improving robustness without sacrificing gains when weights are accurate. Finally, we develop principled prompt engineering strategies and validate the method on a private medical dataset studying Acute Kidney Injury. LSP improves prediction accuracy and identifies clinically relevant features missed by the baselines, with robustness to prompt variation and particular effectiveness in low-data regimes.

• The paper introduces the LLM Sparsity Prior framework that integrates LLM-derived weights into Bayesian variable selection, dynamically discounting inaccurate priors.
• It uses hierarchical hyperpriors and efficient sampling techniques, reducing ℓ1 error up to fourfold and achieving near-perfect F1 scores with high-quality weights.
• Empirical results in AKI prediction demonstrate robust performance in low-data settings, identifying key clinical features missed by traditional methods.

LLM Sparsity Prior for Robust Feature Selection: An Expert Analysis

Motivation and Context

Classical feature selection struggles in low-data, high-dimensional scenarios where informative signal is sparse and prior information is crucial. LLMs have recently been leveraged to encode domain expertise into statistical pipelines, particularly through the extraction of feature importance weights from natural language descriptions. However, prior approaches such as LLM-Lasso inject LLM weights directly into regularization terms, leading to strong sensitivity to weight quality and instability under misaligned or noisy weights.

This paper addresses this challenge by formalizing a framework for quantifying LLM-generated weight quality, proposing a robust mechanism—the LLM Sparsity Prior (LSP)—that integrates these weights into a hierarchy of Bayesian variable selection models. The approach emphasizes dynamic discounting of inaccurate priors, interpretable hyperparameterization, and practical robustness, validated in clinically relevant low-data settings for acute kidney injury (AKI) prediction.

Quantifying LLM Weight Quality

Given the inherent variability in LLM output and the diversity of prompting strategies, rigorous evaluation of LLM-generated weights is essential before their integration into model pipelines. The paper adapts established feature importance evaluation metrics for this purpose:

• $\ell_1$ Weight Agreement: Measures normalized $\ell_1$ distance between min-max-scaled LLM weights and ground truth inclusion vectors, yielding values in $[0,1]$ where 0.5 corresponds to random guessing and 1.0 indicates perfect alignment.
• Pairwise Weight Agreement: Adapts Kendall's $\tau$ rank correlation for binary ground truth, measuring alignment in orderings across all feature pairs.

Empirical plug-in estimators are proposed when ground truth is unavailable, guiding model and prompt selection even when the true feature set is unknown.

LLM Sparsity Prior Model

The LSP framework incorporates LLM-derived weights into prior inclusion probabilities for Spike-and-Slab (SS) and Spike-and-Slab Lasso (SSL) regression models. Each feature's prior inclusion probability is parameterized as:

$$\theta_j = s \frac{w_j^{\eta}}{\frac{1}{p} \sum_{k=1}^{p} w_k^{\eta}}$$

where $s$ controls baseline sparsity, and $\eta$ governs the concentration/contrast of weight influence. Hierarchical hyperpriors (Beta for $s$, zero-inflated discrete uniform for $\eta$) allow the model to revert to an uninformative configuration when weights are misaligned, dynamically calibrating the influence of LLM priors.

Posterior inference leverages coordinate descent and MCMC schemes, adapting acceptance ratios and initialization to heterogenous prior probabilities induced by LLM weights. Efficient sampling is maintained by discretizing $\eta$ and incorporating it into Gibbs updates.

Simulation Results: Robustness and Performance

Comprehensive simulations are conducted in high-dimensional ($\ell_1$0), low-sample ($\ell_1$1) regimes. The LSP methods and LLM-Lasso are contrasted against classical baselines under controlled weight quality variation ($\ell_1$2 from 0.5 to 1.0):

• All LLM-informed methods improve with weight quality; LSP (SS) reduces $\ell_1$3 error fourfold at $\ell_1$4 and drives $\ell_1$5 scores to near-perfection as quality increases.
• Critically, LSP methods never underperform their respective baselines, achieving near-baseline accuracy when weights are random and substantial gains when weights are informative.
• LLM-Lasso is highly sensitive: it underperforms baseline with moderate-quality weights ($\ell_1$6), susceptible to misleading cross-validation selection.
• Feature inclusion probability analyses confirm LSP’s resistance to poor priors—inactive features assigned high importance are rarely included, unlike LLM-Lasso, which exhibits excessive feature inclusion under misleading weights.

  Figure 1: LSP outperforms baselines at all weight qualities and is robust to poor weights; LLM-Lasso suffers severe degradation when weights are inaccurate.

Application: Acute Kidney Injury Prediction

The proposed framework is applied to a large (~1000 features), private medical dataset predicting AKI post-cardiac surgery in five subpopulations. Feature weights are elicited from GPT-5.2o via rigorously engineered prompts with clinical context, task definition, redundancy controls, and scoring rubrics.

Empirical results:

• LSP (SS) reduces out-of-sample MSE by an average of 5% over standard Spike-and-Slab across all cohorts, with substantial improvements in low-data regimes.
• LSP (SSL) yields modest but consistent gains, with effectiveness modulated by model sparsity mechanism.
• Naive LLM weights increase MSE, confirming the necessity for robust hyperprior design.
• LSP (SS) identifies clinically relevant features missed by baselines, notably RBC transfusions, a known AKI biomarker.

  Figure 2: MSE across five prompting strategies; LSP (SS) is robust to both prompt engineering and LLM stochasticity, outperforming baselines in nearly all cases.

Prompt engineering sensitivity studies reveal LSP’s stability: LSP (SS) excels across multiple prompt variants and LLM stochastic draws, minimizing variance and rarely performing worse than baselines. Subsampling experiments show performance gains amplify as training sample size decreases, solidifying LSP’s utility in low-data settings.

Hyperparameter and Structural Sensitivity

The impact of the concentration parameter $\ell_1$7 and its prior is closely examined. Structure recovery error is minimized at intermediate $\ell_1$8 values; the zero-inflated discrete prior achieves performance near the optimum, validating the hierarchical design.

Figure 3: Structure recovery error across fixed $\ell_1$9 values, with zero-inflated prior matching optimal performance.

Implications, Limitations, and Future Directions

The practical implications are clear: LSP enables injection of scalable, domain-aware prior information into Bayesian feature selection without sacrificing robustness. Its effectiveness depends on the role of inclusion probabilities in the target model; in SSL, where sparsification is largely governed by penalty parameter $[0,1]$0, LLM weights are less influential.

Theory calls for further exploration of continuous priors on $[0,1]$1, integration of LLM-informed weights into the penalization path, and broader application to Bayesian tree ensembles, graphical models, and symbolic regression.

Prompt engineering remains a critical bottleneck; principled evaluation and aggregation of LLM-generated weights are necessary to tame output stochasticity and guide model selection in real-world deployments.

Conclusion

The LLM Sparsity Prior framework offers substantial practical and theoretical advances for robust, scalable, and interpretable feature selection in high-dimensional, low-data settings. By dynamically discounting low-quality LLM priors and amplifying their influence when aligned with data, LSP achieves consistent gains in prediction accuracy and clinically meaningful feature discovery. Its hierarchical hyperprior design and computational tractability lay a foundation for future integration with advanced Bayesian regularization techniques and broader scientific discovery tasks.