Source-linked AI summary

An LLM-in-the-loop RL Framework for Bioinformatics Feature Selection

Xinyuan Wang, Deepti Agrawal, Yanjie Fu

arXiv:2610.05600v1cs.LGcs.AI

TL;DR

High-dimensional bioinformatics data create feature-selection challenges, while existing RL approaches can explore inefficiently and rely on rewards that miss domain relevance. This framework uses LLM advice to guide feature exploration and hybrid rewards to combine domain knowledge with data-driven performance. Experiments report stronger benchmark performance, stable results across downstream models, and faster convergence.

  • Problem

    High-dimensional bioinformatics feature selection challenges traditional methods, while RL-based selection can explore inefficiently and rely on rewards that miss domain relevance.

    Source: §I · p. 1
  • Method

    The RL agent selects features sequentially, using LLM advice to guide exploration and LLM review to combine knowledge-driven evaluation with data-driven performance in hybrid rewards.

    Source: §I · pp. 1–2
  • Results

    Across diverse bioinformatics datasets, the framework outperforms baselines, achieves stable performance across downstream models, and converges faster; accuracy improvements reach up to 4%.

    Source: §I · p. 2, §IV · p. 5
  • Takeaways & Limitations

    The LLM’s decision-maker and reviewer roles address inefficient exploration and biased reward signals, while its explanations support transparency and human interpretability.

    Source: §IV · p. 5

    LLM inference adds cost in high-dimensional feature selection.

    Source: §IV · p. 5

Abstract

from arXiv · show

High-dimensional bioinformatics data, characterized by a large number of features relative to the number of samples, pose major challenges such as the ``curse of dimensionality,'' leading to overfitting, high computational cost, and poor generalization. Traditional feature selection methods often suffer from limited scalability and adaptability in such domains. We propose an LLM-in-the-loop reinforcement learning (RL) framework for bioinformatics feature selection, where the RL agent formulates feature selection as a sequential decision-making task, while the large language model (LLM) enhances the process in two ways: (1) guiding exploration through domain-informed advice, and (2) providing hybrid rewards that integrate data-driven performance with knowledge-driven evaluation. The LLM also produces explanations to improve interpretability for human experts without altering the RL policy update. Experiments on diverse bioinformatics datasets show that the LLM-in-the-loop framework outperforms baselines, achieves stable performance across downstream models, and converges faster than pure RL.

I. INTRODUCTION

High-dimensional bioinformatics data make feature selection important, but traditional approaches and RL face trade-offs in scalability, exploration, and reward design. The proposed LLM-in-the-loop framework adds domain-informed guidance and hybrid rewards, and experiments report improved benchmark performance and faster convergence.

I. INTRODUCTION

  • High-dimensional bioinformatics datasets may contain tens of thousands of features but only hundreds of samples, creating overfitting and computational-cost challenges.Feature selection removes irrelevant or redundant features and can support predictive performance and interpretability.Source: §I · p. 1
  • Traditional filter, wrapper, and embedded feature-selection methods respectively face limits in feature interactions, computational expense, or model dependence.Filters evaluate features individually; wrappers repeatedly train models for subsets; embedded methods integrate selection into model training.Source: §I · p. 1
  • RL formulates feature selection as sequential decisions over features, enabling a global search that can adapt to changing data environments.The introduction contrasts this formulation with traditional feature-selection approaches.Source: §I · p. 1
  • RL-based feature selection faces inefficient exploration in large feature spaces and rewards that may miss domain relevance or real-world utility.The introduction notes that feature selection can reduce downstream processing time while boosting performance on a myocardial infarction dataset.Source: §I · p. 1
  • The framework treats the LLM as a domain expert that narrows exploration by suggesting promising features and advising whether to select each candidate.The LLM also reviews decisions to produce hybrid rewards combining knowledge-driven evaluation and data-driven performance, and generates explanations outside policy updates.Source: §I · pp. 1–2
  • The framework achieves best- or second-best performance in most tasks, with accuracy improvements of up to 4% and lower regression error than baselines, while converging comparably faster.Ablations also support the value of LLM guidance in decision-making and reward shaping.Source: §I · p. 2

II. LLM-IN-THE-LOOP RL FRAMEWORK

Feature selection seeks a compact subset for downstream prediction while reducing redundancy and complexity. The framework combines RL optimization with LLM knowledge, using prioritization, decision advice, hybrid rewards, and explanations without letting the LLM determine the final subset.

A. Background and Task

  • The task selects a compact feature subset and trains a downstream model on that subset to improve prediction while reducing redundancy and complexity.The dataset is represented by samples and features with labels or regression targets.Source: A. Background and Task · p. 2
  • LLM knowledge about feature relevance, redundancy, and semantic consistency is integrated with RL optimization to address inefficient exploration and task-driven rewards.Feature selection is combinatorial, with possible subsets growing exponentially with the number of features.Source: A. Background and Task · p. 2

B. Proposed Method

  • An LLM filter prioritizes candidates from feature names, statistical summaries, and domain descriptions, but does not permanently discard low-ranked features.The RL agent can still discover useful but non-obvious candidates.Source: B. Proposed Method · p. 2
  • At each step, the RL agent combines LLM action advice with its policy, then combines downstream validation feedback and LLM review into a hybrid reward.The LLM observes the selected subset and candidate feature before the decision, and reviews after action execution.Source: B. Proposed Method · p. 2
  • The LLM softly guides exploration and reward shaping, while the RL agent learns the final feature subset from task-driven validation feedback.Natural-language explanations support human interpretability but are not used in policy updates.Source: B. Proposed Method · p. 2

C. Reinforcement Learning Framework

The framework formalizes feature selection as an MDP in which the RL agent selects or skips candidate features and maximizes expected cumulative hybrid reward. Its implementation allows repeated exploration and masks already selected features.

C. Reinforcement Learning Framework

  • The agent’s state contains the selected-feature binary vector and current candidate, and its action selects or skips that feature.Selecting adds the feature to the subset; skipping leaves the subset unchanged.Source: C. Reinforcement Learning Framework · p. 2
  • The reward combines downstream validation performance with LLM-based feedback, and the discount factor balances short- and long-term rewards.The framework is formalized as an MDP with state, action, transition, reward, and discount factor.Source: C. Reinforcement Learning Framework · p. 2
  • The objective is to maximize expected cumulative reward over the RL agent’s policy and a maximum of T decision steps.The implementation sets T = 5d; selected features are masked, while uncertain unselected features can be revisited through a prioritized queue.Source: C. Reinforcement Learning Framework · p. 3

1) LLM as Decision Maker:

The LLM decision maker scores candidate features using domain knowledge and semantic consistency. The RL agent uses that advice to adjust action probabilities and guide exploration while retaining random and policy-based choices.

1) LLM as Decision Maker:

  • The LLM scores a candidate feature in the context of the selected subset and task information; higher scores indicate greater predicted usefulness by domain knowledge and semantic consistency.The score ranges from 0 to 1.Source: 1) LLM as Decision Maker · p. 3
  • The LLM returns a score, uncertainty flag, and explanation, but only the numerical score enters the RL agent’s action selection.The explanation is retained for human interpretation.Source: 1) LLM as Decision Maker · p. 3
  • The RL agent adjusts its action probabilities using the LLM advice score, with α controlling the strength of guidance.The advice score favors selecting or skipping according to the corresponding action-specific mapping.Source: 1) LLM as Decision Maker · p. 3
  • LLM advice is added as an exploration branch alongside random and policy-based action selection.The LLM-guided branch chooses the action with the highest advice score.Source: 1) LLM as Decision Maker · p. 3
  • The exploration strategy preserves random exploration while steering the agent toward semantically promising features.Source: 1) LLM as Decision Maker · p. 3

2) LLM as Reviewer:

The LLM reviewer evaluates each selected feature for task relevance, complementarity, and redundancy, adding knowledge-driven feedback to validation-based reward. The final reward balances these signals.

2) LLM as Reviewer:

  • The reviewer scores whether a selected feature is relevant, complementary to the subset, and not overly redundant.Its score uLLM(st, at) ∈[0, 1] is rescaled before use.Source: 2) LLM as Reviewer · p. 3
  • The extrinsic reward combines validation-score improvement with penalties for subset size and feature redundancy.For classification, validation quality is accuracy; for regression, it is negative normalized ℓ2 error, and redundancy is average pairwise feature correlation.Source: 2) LLM as Reviewer · p. 3
  • The final reward mixes extrinsic and intrinsic signals using λ, which balances validation-based and LLM-based feedback.Reward components are normalized within each dataset to avoid scale mismatch.Source: 2) LLM as Reviewer · p. 3

3) Policy Learning and Reliability Control:

The policy uses a DQN trained with experience replay and a target network, while reliability controls reduce LLM influence when its outputs are uncertain or conflict with validation.

3) Policy Learning and Reliability Control:

  • A DQN agent is trained with experience replay and a target network, and its Q-network minimizes squared temporal-difference error.The temporal-difference target is computed for each transition (st, at, Rt, st+1).Source: 3) Policy Learning and Reliability… · p. 3
  • Uncertain LLM responses or repeated conflicts with validation reduce LLM influence by lowering α or increasing λ.Repeated feature-level scores are cached to reduce redundant LLM calls.Source: 3) Policy Learning and Reliability… · p. 3
  • The LLM is an auxiliary knowledge source, while downstream validation performance drives the final feature-selection policy.Source: 3) Policy Learning and Reliability… · p. 3

III. EXPERIMENT

Across the evaluated datasets, the framework performs best or second-best on most tasks, with gains reported in high-dimensional settings. The experiments cover classification and regression datasets and compare against methods from five baseline categories.

A. Experimental Setup

  • The evaluation covers diverse bioinformatics, healthcare, and scientific datasets, using accuracy for classification and ℓ2 error for regression.Baselines span embedded, filter, evolutionary, neural, and RL methods; GPT-4o provides feature-level advice and feedback but does not choose the final subset.Source: A. Experimental Setup · p. 3

B. Performance Improvement

  • In high-dimensional RNA-Seq and Glycation tasks, the framework achieves up to 4% higher classification accuracy and lower regression error than baselines.The reported effect is stable across datasets and task settings.Source: B. Performance Improvement · pp. 3–4
  • Best or second-best on most datasets, with an average rank of 1.41 across classification and regression tasks.Table I reports accuracy (%) for classification and ℓ2 error for regression.Source: B. Performance Improvement · p. 3, Table I · p. 4

C. Ablation Study

Removing either LLM module reduces performance, and pure RL performs worst, especially on high-dimensional data. The Decision Maker supports exploration while the Reviewer adds domain-informed reward feedback.

C. Ablation Study

  • Removing the Decision Maker, Reviewer, or both reduces performance, with the larger drop tending to occur when both are removed.Source: C. Ablation Study · p. 4
  • Without the Decision Maker, exploration slows and may miss useful features; without the Reviewer, reward relies only on human-defined metrics.Those metrics can be narrow or biased, and pure RL performs worst, especially on high-dimensional data.Source: C. Ablation Study · p. 4
  • The Decision Maker focuses exploration on semantically relevant features, while the Reviewer refines reward with domain-informed feedback.Together, the modules support faster convergence and more accurate selection.Source: C. Ablation Study · p. 4

D. Stability

The framework maintains stable performance across different downstream models. Although LLM inference adds per-step runtime, it reaches comparable performance with lower total runtime than the value-based RL baseline by converging with fewer exploration steps.

D. Stability

  • Comparable accuracy across SVM, Decision Tree, and neural-network classifiers, and stable MSE across multiple regressors, indicate robustness to downstream-model choice.The authors interpret this pattern as benefits not tied to one learner and selected features that support learning across tasks.Source: Figure 4 · p. 4, D. Stability · pp. 4–5

E. Convergence Analysis

  • LLM guidance adds per-step runtime but achieves lower total runtime than the value-based RL baseline by converging with fewer exploration steps.The framework pays extra inference cost for action advice and reward feedback, while reducing repeated trials over uninformative or redundant features.Source: Figure 5 · p. 5, E. Convergence Analysis · p. 5

F. Reliability of LLM Guidance

The framework limits uncertain LLM guidance to an auxiliary role, with validation-based extrinsic rewards still optimizing the policy. The conclusion reports that its decision-maker and reviewer roles address inefficient exploration and biased rewards, alongside stable results and faster convergence across diverse bioinformatics datasets.

F. Reliability of LLM Guidance

  • Validation-based extrinsic rewards still optimize the policy, while the LLM supplies auxiliary advice and intrinsic rewards rather than acting as an oracle.Constrained prompts, fixed output formats, uncertainty flags, and consistency checks are used to reduce hallucination risk; conflicting advice leads to a lower policy-bias weight α or higher extrinsic-reward weight λ.Source: F. Reliability of LLM Guidance · p. 5

IV. CONCLUSION

  • Across diverse bioinformatics datasets, the framework outperforms baselines, achieves stable results across downstream models, and converges faster despite extra LLM inference cost.The LLM acts as both decision maker and reviewer, and its explanations support transparency and human interpretability.Source: §IV · p. 5
Loading 2610.05600v1…