Dementia-R1: Reinforced Pretraining and Reasoning from Unstructured Clinical Notes for Real-World Dementia Prognosis

TL;DR

Dementia-R1 tackles the challenge of prognostic reasoning over long, non-monotonic dementia trajectories described in unstructured clinical notes. It introduces a two-stage reinforcement learning framework with Cold-Start pretraining that learns intermediate clinical indices (e.g., MMSE, GDS, CDR) before fine-tuning on the final dementia prognosis, using verifiable rewards and Group Relative Policy Optimization. The approach achieves a 7B-parameter model that rivals larger baselines on real-world AMC data (77.03% F1) and generalizes to the ADNI benchmark (74.91% F1), with neurologist-aligned reasoning and robust performance across time horizons. This work demonstrates that RL with verifiable rewards can yield grounded, longitudinal clinical reasoning in NLP models, enabling scalable, interpretable prognosis from unstructured medical narratives and structured records alike.

Abstract

While Large Language Models (LLMs) have shown strong performance on clinical text understanding, they struggle with longitudinal prediction tasks such as dementia prognosis, which require reasoning over complex, non-monotonic symptom trajectories across multiple visits. Standard supervised training lacks explicit annotations for symptom evolution, while direct Reinforcement Learning (RL) is hindered by sparse binary rewards. To address this challenge, we introduce Dementia-R1, an RL-based framework for longitudinal dementia prognosis from unstructured clinical notes. Our approach adopts a Cold-Start RL strategy that pre-trains the model to predict verifiable clinical indices extracted from patient histories, enhancing the capability to reason about disease progression before determining the final clinical status. Extensive experiments demonstrate that Dementia-R1 achieves an F1 score of 77.03% on real-world unstructured clinical datasets. Notably, on the ADNI benchmark, our 7B model rivals GPT-4o, effectively capturing fluctuating cognitive trajectories. Code is available at https://anonymous.4open.science/r/dementiar1-CDB5

Figures (15)

• Figure 1: Multi-dimensional Performance Profile. Dementia-R1 demonstrates a consistent and balanced performance gain across all dimensions, including intermediate clinical reasoning tasks (e.g., MMSE, CDR-SB, ADAS-Cog) and the final dementia prognosis (F1-score)
• Figure 2: Overview of the Dementia-R1 Framework. The pipeline consists of two phases: Stage 1: Cold-Start Pre-training, where the base model learns longitudinal reasoning via GRPO on forecasting tasks; and Stage 2: Task Fine-tuning, where the reasoning-aligned model is adapted for the final dementia prediction task.
• Figure 3: Examples of Longitudinal Sample Construction. Patient history is retrospectively sliced relative to a Target Anchor, applying the unified protocol across both unstructured (AMC) and structured (ADNI) data.
• Figure 4: Dataset Overview. Visualization of sample and patient counts. Training sets are balanced to prevent bias, while test sets retain natural patient prevalence.
• Figure 5: Neurologist Blind Pairwise Evaluation. Comparison between Dementia-R1 and the baseline model across six clinical dimensions.