Technical Summary: LLM Sparsity Prior for Robust Feature Selection

Problem Statement

The integration of Large Language Models (LLMs) into statistical pipelines offers a scalable mechanism to elicit domain-informed prior information for high-dimensional variable selection. However, existing methods, such as LLM-Lasso, exhibit significant sensitivity to the quality of LLM-generated weights. When these weights are inaccurate, misleading, or "poor," the performance of LLM-informed methods degrades substantially, often underperforming standard baselines. Current approaches lack a principled framework to quantify weight quality prior to integration and lack robustness mechanisms to dynamically discount uninformative weights without sacrificing gains when weights are accurate.

Methodology

1. Quantifying Weight Quality

The authors first establish a framework to quantify the quality of LLM-generated feature weights ($w$) relative to a ground truth inclusion vector ($\gamma^*$). They adapt two standard metrics:

• $\ell_1$ Weight Agreement ($\phi_{\ell_1}$): Measures the $\ell_1$ distance between the ground truth and the min-max scaled weight vector.
• Pairwise Weight Agreement ($\phi_{\text{pairwise}}$): Measures the frequency with which the relative ordering of weights aligns with the strict ordering of the ground truth (analogous to Kendall's $\tau$ adapted for binary truth).
  These metrics allow for the rigorous evaluation of LLM-informed methods across varying weight regimes. In practice, where $\gamma^*$ is unknown, a plug-in estimator using a standard statistical estimate $\hat{\gamma}$ is proposed to measure empirical alignment.

2. The LLM Sparsity Prior (LSP)

The core contribution is the LLM Sparsity Prior (LSP), a Bayesian framework that integrates LLM-generated weights into the prior inclusion probabilities of Spike-and-Slab and Spike-and-Slab Lasso models.

• Weight Integration: Instead of asking the LLM to estimate continuous probabilities (which can be error-prone), the LLM is asked to assign ordinal importance rankings $w \in \{1, \dots, K\}^p$.
• Prior Formulation: The prior inclusion probability $\theta_j$ for feature $j$ is defined as:
  $$\theta_j = s \frac{w_j^\eta}{\sum_{k=1}^p w_k^\eta}$$
  where:
  • $s \in (0, 1)$ is a baseline sparsity parameter.
  • $\eta \in (0, \eta_{\max})$ is a concentration parameter controlling the degree of contrast between weights. As $\eta \to 0$, the model recovers the traditional uniform setting; larger $\eta$ amplifies weight differences.
• Hierarchical Hyperpriors: To ensure robustness, the authors assign hierarchical hyperpriors to $s$ and $\eta$:
  • $s \sim \text{Beta}(a_s, b_s)$.
  • $\eta$ follows a zero-inflated Discrete Uniform prior: $\eta \sim \pi_0 \delta_0 + (1-\pi_0)\text{DiscreteUniform}(E)$.
  • This structure allows the model to dynamically "fall back" to an uninformative baseline ($\eta=0$) if the LLM weights conflict with the data, thereby preventing performance degradation from poor weights.

3. Posterior Estimation

The LSP is compatible with efficient posterior estimation techniques:

• Spike-and-Slab: An adaptation of the Add-Delete-Swap (ADS) algorithm is used. The Metropolis-Hastings acceptance ratio is modified to account for heterogeneous prior inclusion probabilities, and weights are used to initialize the inclusion vector $\gamma$ to accelerate convergence.
• Spike-and-Slab Lasso: The coordinate descent algorithm is run multiple times across a grid of $\eta$ values to select the value maximizing the log joint posterior.

Key Contributions

1. Weight Quality Framework: A principled diagnostic tool to quantify LLM-generated weight quality, enabling rigorous comparison of generation mechanisms.
2. LLM Sparsity Prior (LSP): A novel framework integrating LLM weights into Bayesian variable selection via interpretable hyperparameters ($s$ and $\eta$) and hierarchical priors that ensure robustness to misleading weights.
3. Prompt Engineering & Validation: Development of principled prompt strategies and empirical validation on a private medical dataset (Acute Kidney Injury), demonstrating robustness to prompt variations and effectiveness in low-data regimes.

Results

Simulation Studies

• Robustness: LSP methods (both Spike-and-Slab and Spike-and-Slab Lasso variants) remain robust across all weight qualities. Under random weights ($\phi_{\ell_1} = 0.5$), LSP performance is nearly identical to the baseline (within 0.1 units of $\ell_1$ error).
• Performance Gain: As weight quality improves, LSP significantly outperforms baselines. For instance, in high-dimensional regimes with accurate weights, LSP reduced $\ell_1$ error fourfold compared to the standard Spike-and-Slab.
• Comparison to LLM-Lasso: Unlike LSP, the LLM-Lasso method is highly volatile. Under moderate-to-poor weight regimes, its performance deteriorates severely (up to 1.64 times the error of standard Lasso) because cross-validation fails to correctly revert to the standard Lasso when weights are uninformative.

Application: Acute Kidney Injury (AKI)

Using a private dataset from Baylor College of Medicine (BCM) involving cardiac surgery patients, the authors predicted post-operation serum creatinine levels.

• Predictive Accuracy: LSP (Spike-and-Slab) reduced out-of-sample Mean Squared Error (MSE) by an average of 5% across five distinct patient subpopulations compared to the standard Spike-and-Slab.
• Feature Selection: LSP identified clinically relevant features missed by baselines. Specifically, it uniquely selected Intraoperative Red Blood Cells (RBC) transfused as a key predictor, a known biomarker for AKI, whereas the baseline method missed it.
• Low-Data Regimes: The performance gain of LSP over baselines widened as the sample size decreased ($n \in \{100, 150, 200\}$), confirming its utility in data-scarce environments.
• Prompt Sensitivity: The method remained robust across five different prompt engineering strategies, with LSP (Spike-and-Slab) outperforming the baseline in 24 of 25 cases.

Significance and Claims

The paper claims that the LLM Sparsity Prior offers a reliable and effective alternative for real-world applications where the quality of LLM-generated weights is unknown. By incorporating a hierarchical prior that allows the model to dynamically discount uninformative weights, LSP achieves two goals simultaneously:

1. It capitalizes on high-quality domain knowledge to improve prediction accuracy and feature selection when weights are accurate.
2. It prevents the substantial performance degradation seen in existing methods (like LLM-Lasso) when weights are inaccurate or misleading.

The authors note that the advantage of LSP is most pronounced in low-data regimes and highly sparse settings where standard methods struggle to identify signals. They conclude that while LLM-informed methods hold great promise, their practical utility depends on robust integration mechanisms like LSP that account for the stochasticity and potential hallucinations of LLMs.