Source-linked AI summary
DRG-LLaMA : Tuning LLaMA Model to Predict Diagnosis-related Group for Hospitalized Patients
Hanyin Wang, Chufan Gao, Christopher Dantona, Bryan Hull, Jimeng Sun
TL;DR
DRG assignment is labor-intensive and typically follows discharge, motivating accurate predictions during hospitalization for resource planning. This study fine-tunes a LLaMA-derived model on discharge summaries and evaluates direct and two-label DRG prediction. In the two-label approach, top-1 accuracy is 67.8% for base DRG and 67.5% for CC/MCC.
Problem
DRG assignment is a labor-intensive manual process typically performed after discharge, despite interest in accurate in-hospital predictions for resource planning.
Source: Introduction · p. 1Method
The study fine-tunes a LLaMA-derived model on MIMIC-IV discharge summaries for end-to-end DRG prediction and separate base DRG and CC/MCC prediction.
Source: Introduction · pp. 1–2Results
In the two-label approach, top-1 accuracy was 67.8% for base DRG and 67.5% for CC/MCC.
Source: Discussion · p. 6Takeaways & Limitations
DRG prediction is intended to support hospital resource planning and allocation during a patient's hospitalization.
Source: Introduction · p. 1The study used only discharge summaries, which are available after discharge, rather than earlier HPI or ED notes.
Source: Discussion · p. 6
Abstract
from arXiv · showhide
In the U.S. inpatient payment system, the Diagnosis-Related Group (DRG) is pivotal, but its assignment process is inefficient. The study introduces DRG-LLaMA, an advanced large language model (LLM) fine-tuned on clinical notes to enhance DRGs assignment. Utilizing LLaMA as the foundational model and optimizing it through Low-Rank Adaptation (LoRA) on 236,192 MIMIC-IV discharge summaries, our DRG-LLaMA-7B model exhibited a noteworthy macro-averaged F1 score of 0.327, a top-1 prediction accuracy of 52.0%, and a macro-averaged Area Under the Curve (AUC) of 0.986, with a maximum input token length of 512. This model surpassed the performance of prior leading models in DRG prediction, showing a relative improvement of 40.3% and 35.7% in macro-averaged F1 score compared to ClinicalBERT and CAML, respectively. Applied to base DRG and complication or comorbidity (CC)/major complication or comorbidity (MCC) prediction, DRG-LLaMA achieved a top-1 prediction accuracy of 67.8% and 67.5%, respectively. Additionally, our findings indicate that DRG-LLaMA's performance correlates with increased model parameters and input context lengths.
Introduction
DRG prediction is sought to support inpatient operations, but it differs from multilabel ICD coding. This study applies a LLaMA-derived model to clinical notes and reports stronger performance than prior leading DRG-prediction models.
Introduction
- Medical LLM research remains nascent, and the paper notes safety and privacy concerns including hallucination, bias, and reasoning deficiencies.The authors call for detailed analysis of LLM performance and limitations in medicine.Source: Introduction · p. 1
- DRG assignment is labor-intensive and usually follows discharge, while accurate prediction during hospitalization is sought to support resource planning and allocation.Unlike multi-label ICD coding, DRG prediction assigns one code per visit.Source: Introduction · p. 1
- The study considers both end-to-end DRG prediction and a two-label approach that predicts base DRG and CC/MCC status separately.The two labels are then used to infer the final DRG label.Source: Introduction · pp. 1–2
- DRG-LLaMA outperformed prior leading DRG-prediction models CAML and ClinicalBERT (Alsentzer et al. 2019).The model is derived from LLaMA and trained on MIMIC-IV discharge summaries.Source: Introduction · pp. 1–2
Results
DRG-LLaMA outperformed ClinicalBERT and CAML in single-label DRG prediction and achieved similar overall accuracy through a two-label approach. Performance varied with DRG frequency and improved with larger models and longer input contexts.
Study cohort
- The cohort used the discharge summary’s “brief hospital course” as input, excluded low-quality summaries and MS-DRGs occurring fewer than 2 times, and had imbalanced counts (training-set median 124.5 cases per DRG).The stratified split allocated 90% of the data to training and 10% to testing.Source: Study cohort · p. 2
DRG prediction as a single-label classification task
- 40.3% and 35.7% relative macro-F1 improvements over ClinicalBERT and CAML, respectively, accompanied DRG-LLaMA’s 52.0% top-1 and 84.8% top-5 accuracy.Results used a maximum input token size of 512; DRG-LLaMA outperformed both baselines across evaluation metrics.Source: DRG prediction as a single-label… · p. 2, Error analysis · p. 3, Table 1 · p. 3
- Top-1 accuracy rose from 55.7% among the 300 most frequent DRGs to 69.4% among the 30 most frequent, while performance declined for less frequent DRGs.High AUC scores across models reflected many infrequent classes and resulting true-negative predictions.Source: DRG prediction as a single-label… · p. 2
- A 13B model with 1024-token input achieved 54.6% top-1 accuracy, 86.5% top-5 accuracy, and 0.361 macro-F1.Across evaluated configurations, all metrics improved with larger models and longer input contexts.Source: DRG prediction as a single-label… · p. 2, Table 2 · p. 3
DRG prediction as a two-label classification task
- 51.5% accuracy after combining base DRG and CC/MCC predictions was comparable to 52.0% in single-label classification.Base DRG and CC/MCC top-1 accuracies were 67.8% and 67.5%, respectively.Source: DRG prediction as a two-label… · p. 2
Error analysis
- DRGs with over 80% top-5 accuracy had a median 309 training cases per label, versus 17 for those below 20%.Training-case counts correlated with performance, while DRG type also affected prediction.Source: Error analysis · p. 2
- In ten manually reviewed error cases, inadequate clinical concept extraction appeared in four and difficulty selecting the base DRG in three.The authors also identified erroneous CC/MCC, missing required information, and one plausible incorrect label; discharge summaries alone may omit information needed for prediction.Source: Error analysis · p. 3
Discussion
DRG-LLaMA outperformed ClinicalBERT and CAML in DRG prediction, with performance improving as model size and input context increased. The discussion also considers how input selection, training-case frequency, discharge-summary scope, and DRG coding rules shape results.
Discussion
- 67.8% base-DRG and 67.5% CC/MCC top-1 accuracy were achieved in the two-label task with a 512-token input limit.The two-label task predicts base DRG and CC/MCC separately.Source: Discussion · p. 4, Table 3 · p. 4
- DRG-LLaMA outperformed ClinicalBERT and CAML, and both the baselines and DRG-LLaMA surpassed outcomes reported in prior research.The authors note that the baselines trained faster, while attributing the stronger outcomes as plausibly linked to input selection and the larger MIMIC-IV dataset.Source: Discussion · p. 4
- DRG-LLaMA performance improved with larger models and longer input contexts, within experiments capped at 1,024 tokens.ClinicalBERT and CAML were described as robust baselines with faster training times.Source: Discussion · p. 4
- The study used only the brief hospital-course section of discharge summaries, unlike prior work using earlier clinical notes, and the summaries were available only after discharge.The authors suggest HPI or emergency-department notes as potential sources for earlier DRG prediction.Source: Discussion · pp. 4, 6
- Manual review identified errors involving missing information, clinical-concept extraction, and base-DRG selection; some assignment rules extend beyond the discharge narrative.The error review covered 10 cases, and the authors distinguish weaknesses that might respond to larger models or more data from challenges inherent to the training-data setup.Source: Discussion · pp. 5–6
- 51.5% final DRG accuracy after mapping was comparable to the 52.0% single-label result.The mapping rules address different MS-DRG split scenarios.Source: Discussion · pp. 4, 6
- For DRG-LLaMA-7B with a maximum input size of 512 tokens, prediction accuracy declined in less frequent DRGs; higher top-5 accuracy groups generally had more training cases, with substantial variance in the best-performing group.Source: Figure 2 · p. 6
- Experiments were limited to LLaMA models up to 13 billion parameters, and practical integration into hospital coding workflows remains an area for exploration.The authors also report that an extensive hyperparameter search was not feasible.Source: Discussion · pp. 6–7
Methodology
The study prepares MIMIC-IV discharge summaries and harmonizes MS-DRG labels before fine-tuning LLaMA for either direct DRG classification or prediction of base DRG and CC/MCC labels.
Dataset and Preprocessing
- Brief hospital course sections were extracted from discharge summaries; summaries with duplicated content or fewer than 40 words were filtered, MS-DRG codes harmonized across time to version 34.0, and data split 90%/10% stratified by DRG codes.Source: Dataset and Preprocessing · p. 7
Model Development
- The model-development setup fine-tunes LLaMA on discharge summaries paired with DRG codes using a classification task.The passage introduces two strategies, elaborated in the single-label and two-label approaches.Source: Model Development · p. 7
Single label approach
- The single-label approach predicts one DRG from each discharge summary, using the last token’s LLaMA representation to produce logits for 738 DRG codes and categorical cross-entropy for training.The target DRG is represented as an integer from 0 to 737.Source: Single label approach · p. 7
Two-label approach
- The two-label approach predicts base DRG and CC/MCC status separately, then applies a mapping rule to infer the final DRG.Its loss combines base-DRG cross-entropy with a CC/MCC cross-entropy term weighted by one half.Source: Two-label approach · p. 8
Addressing Computational Constraints via LoRA Training
The study uses LoRA to adapt LLaMA under computational constraints, compares it with CAML and ClinicalBERT, and reports its evaluation procedures and data, code, and privacy considerations.
Addressing Computational Constraints via LoRA Training
- LoRA freezes pretrained weights and trains low-rank matrices added to the weights, adapting attention weights to reduce dimensionality and computational cost.The approach was used because available computational resources did not permit extensive hyperparameter search.Source: Addressing Computational Constraints… · p. 8
Training Details
- For classification, the model uses the last-token embedding to calculate DRG logits, with softmax converting logits to probabilities.Training used cross-entropy loss and Adam for three epochs, with learning rate 2 × 10^-5, weight decay 0.01, batch size 4, and LoRA rank and alpha of 8.Source: Training Details · p. 8
Baseline Models
- CAML and ClinicalBERT were selected as DRG-prediction baselines; CAML uses convolutional features pooled with attention, while ClinicalBERT was further pretrained on two million MIMIC-III clinical notes.CAML used early stopping based on micro-F1 and a maximum sequence length of 512; ClinicalBERT was fine-tuned for three epochs.Source: Baseline Models · p. 8
Statistical analysis
- Evaluation reported macro- and micro-averaged F1 and AUC, top-1, top-5, and top-10 accuracy, and standard deviations from 30 bootstrap resamples.The Figure 2a smoothing spline used generalized cross-validation with default parameters.Source: Statistical analysis · p. 8
Data availability; Code availability; Ethical Concerns
- MIMIC-IV access requires a signed safe usage agreement, and the study scripts and documentation are available in a public GitHub repository.The dataset is deidentified under the HIPAA Safe Harbor provision.Source: Data availability · p. 9, Code availability · p. 9, Ethical Concerns · p. 9
Ethical Concerns
- The authors judge privacy risks to be quite low because they used open-source models and trained on MIMIC, but note that pretrained models may have encountered sensitive data.They state that this risk should not be counted out.Source: Ethical Concerns · p. 9
1 Supplementary Note 1: Examples of DRGs split in MS-DRG version 34.0
The supplement gives examples of MS-DRG version 34.0 categories with three-way splits, two-way splits, and no splits.
Three-way split DRGs
- DRGs 56 and 57 illustrate a three-way split category for degenerative nervous system disorders, with MCC and without MCC listed.The supplied examples identify the MCC and no-MCC variants.Source: §1 · p. 12
Base DRGs with no splits
- DRGs 11, 12, and 13 list tracheostomy for face, mouth, and neck diagnoses with MCC, with CC, and without CC/MCC, respectively.Source: §1 · p. 12
Two-way split DRGs with MCC/CC and no CC
- DRGs 52 and 53 illustrate a two-way split for spinal disorders and injuries, with CC/MCC versus without CC/MCC.Source: §1 · p. 12
Base DRGs with no splits
- DRG 69 is listed as transient ischemia under base DRGs with no splits.Source: §1 · p. 12
2 Supplementary Method 1. Process to address different DRG versions
The authors harmonized historical DRG codes to MS-DRG version 34.0 using text standardization, fuzzy matching, and physician review; codes without an appropriate match were excluded.
Process to address different DRG versions
- The harmonization algorithm targeted MS-DRG version 34.0 and first standardized abbreviations and capitalization in DRG names.Examples include replacing “W/O” with “WITHOUT,” “CATH” with “CATHETERIZATION,” and “PROC” with “PROCEDURES.”Source: §2 · p. 12
- Fuzzy string matching identified historical DRGs that did not match a version 34.0 code, after which an internal medicine physician manually assigned appropriate matches.Source: §2 · p. 12
- Hospitalizations with historical DRGs lacking an appropriate version 34.0 code were excluded from the cohort.Examples included “URINARY STONES W MCC” and “NASAL TRAUMA AND DEFORMITY WITH CC”.Source: §2 · p. 12
3 Supplementary Method 2. Process to dissect and derive DRGs to/from base DRGs and CC/MCC
The method separates DRGs into base DRG and CC/MCC components, then uses defined labels and clinical-knowledge rules to infer a final DRG.
3 Supplementary Method 2. Process to dissect and derive DRGs to/from base DRGs and CC/MCC
- Using regular expressions, the study extracted principal diagnosis/procedure strings before CC/MCC descriptions in MS-DRG v.34.0, identifying 340 distinct principal diagnosis/procedures as base DRGs.Source: §3 · p. 13
- CC/MCC status uses five labels, including “not applicable” when the DRG code does not specify CC/MCC status.The authors note this choice may introduce noisy signals but consider it less erroneous than labeling such cases “without CC/MCC.”Source: §3 · p. 13
- Final DRG inference first checks whether the predicted principal diagnosis/procedure matches the target base DRG, then compares CC/MCC labels within its set, and maps labels outside it according to MS-DRG splits.For a DRG with no split, any CC/MCC prediction maps to the correct DRG when the base DRG matches. For DRGs 56 and 57, predictions of without CC/MCC, with CC, or not applicable map to without MCC.Source: §3 · p. 13
4 Supplementary Results
The supplementary results describe a long-tailed distribution of training cases, computational costs that rise with model size and input length, and DRGs at both performance extremes.
4.1 Distribution of training cases per DRG
- The median was 124.5 training cases per DRG, and only 12 DRG codes exceeded 2,000 cases.Supplementary Figure 1 displays the distribution as a density plot with a dashed median line.Source: §4.1 · p. 14, Figure 1 · p. 14
4.2 Computational cost
- Doubling DRG-LLaMA’s maximum input tokens increased training time by approximately 90%, while expanding from 7 billion to 13 billion parameters increased it by roughly 70%.ClinicalBERT and CAML took considerably less time to train than DRG-LLaMA; Supplementary Table 1 presents training and evaluation times.Source: §4.2 · p. 14, Table 1 · p. 15
4.3 Analysis of DRGs with best and worst prediction performance
- Nine DRGs reached 100% top-1 accuracy, mostly (8 of 9) infrequent surgical DRGs, while zero top-5 accuracy occurred among medical and surgical DRGs with median training counts of 17.5 and 16.Many difficult cases were clinically and semantically hard to distinguish; the supplementary tables list the DRGs at the performance extremes.Source: §4.3 · p. 15, Table 2 · p. 15, Table 3 · p. 16
Related papers
This paper cites
- BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding2018
- Language Models are Few-Shot Learners2020
- Training language models to follow instructions with human feedback2022
- LLaMA: Open and Efficient Foundation Language Models2023
- HuggingFace's Transformers: State-of-the-art Natural Language Processing2019
- Towards Expert-Level Medical Question Answering with Large Language Models2023