Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MedFeat: Model-Aware and Explainability-Driven Feature Engineering with LLMs for Clinical Tabular Prediction

Zizheng Zhang, Yiming Li, Justin Xu, Jinyu Wang, Rui Wang, Lei Song, Jiang Bian, David W Eyre, Jingjing Fu

TL;DR

Across a broad range of clinical prediction tasks, MedFeat achieves stable improvements over various baselines and discovers clinically meaningful features that generalize under distribution shift, demonstrating robustness across years and from ICU cohorts to general hospitalized patients, thereby offering insights into real-world deployment.

Abstract

In healthcare tabular predictions, classical models with feature engineering often outperform neural approaches. Recent advances in Large Language Models enable the integration of domain knowledge into feature engineering, offering a promising direction. However, existing approaches typically rely on a broad search over predefined transformations, overlooking downstream model characteristics and feature importance signals. We present MedFeat, a feedback-driven and model-aware feature engineering framework that leverages LLM reasoning with domain knowledge and provides feature explanations based on SHAP values while tracking successful and failed proposals to guide feature discovery. By incorporating model awareness, MedFeat prioritizes informative signals that are difficult for the downstream model to learn directly due to its characteristics. Across a broad range of clinical prediction tasks, MedFeat achieves stable improvements over various baselines and discovers clinically meaningful features that generalize under distribution shift, demonstrating robustness across years and from ICU cohorts to general hospitalized patients, thereby offering insights into real-world deployment. Code required to reproduce our experiments will be released, subject to dataset agreements and institutional policies.

MedFeat: Model-Aware and Explainability-Driven Feature Engineering with LLMs for Clinical Tabular Prediction

TL;DR

Across a broad range of clinical prediction tasks, MedFeat achieves stable improvements over various baselines and discovers clinically meaningful features that generalize under distribution shift, demonstrating robustness across years and from ICU cohorts to general hospitalized patients, thereby offering insights into real-world deployment.

Abstract

In healthcare tabular predictions, classical models with feature engineering often outperform neural approaches. Recent advances in Large Language Models enable the integration of domain knowledge into feature engineering, offering a promising direction. However, existing approaches typically rely on a broad search over predefined transformations, overlooking downstream model characteristics and feature importance signals. We present MedFeat, a feedback-driven and model-aware feature engineering framework that leverages LLM reasoning with domain knowledge and provides feature explanations based on SHAP values while tracking successful and failed proposals to guide feature discovery. By incorporating model awareness, MedFeat prioritizes informative signals that are difficult for the downstream model to learn directly due to its characteristics. Across a broad range of clinical prediction tasks, MedFeat achieves stable improvements over various baselines and discovers clinically meaningful features that generalize under distribution shift, demonstrating robustness across years and from ICU cohorts to general hospitalized patients, thereby offering insights into real-world deployment. Code required to reproduce our experiments will be released, subject to dataset agreements and institutional policies.
Paper Structure (42 sections, 4 equations, 4 figures, 11 tables, 1 algorithm)

This paper contains 42 sections, 4 equations, 4 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Explainability-Driven and Model-aware Feature Engineering
  4. MedFeat Overview
  5. Explainability-Driven Feedback
  6. Importance-Weighted Island Sampling
  7. Model-Aware Feature Generation
  8. Candidate Evaluation and Selection
  9. Feedback Memory and Termination
  10. Experiments
  11. Datasets and prediction tasks
  12. Baselines and Downstream Models
  13. Implementation details
  14. Main results
  15. Ablations
  16. ...and 27 more sections

Figures (4)

  • Figure 1: MedFeat iteratively augments clinical tabular data with LLM-generated features using explainability-guided, model-aware feedback. Starting from a downstream learner trained on the original features (1), the framework computes feature importance and profiles features by type and metadata (2). It then forms importance-weighted “feature islands” by sampling a small subset of influential features (3), and prompts the LLM with the island context, model constraints, and a memory bank of past successes/failures to propose clinically meaningful transformations (4). Proposed features are validated and evaluated locally (5), accepted features are added to the dataset, and the feedback memory is updated for the next iteration. Along the process, no patient-level raw data are sent to the LLM.
  • Figure 2: Evaluations on training IORD data enhanced by MedFeat-generated features from MIMIC for five different runs.
  • Figure 3: Comparisons of the performance between model retrained with full data and inference using trained model with MedFeat augmented features.
  • Figure 4: The percentage of MedFeat-generated features in the top 10 most important features in the augmented data.