Traj-CoA: Patient Trajectory Modeling via Chain-of-Agents for Lung Cancer Risk Prediction

AI summary

Overview

Research area: Clinical NLP / LLM-based temporal reasoning on electronic health records (EHR), applied to oncology risk prediction.

Technical level: Intermediate — assumes familiarity with LLMs, RAG, and multi-agent systems, but the core ideas are explained accessibly.

Scope: Introduces Traj-CoA, a chain-of-agents framework with a structured long-term memory (EHRMem) that enables zero-shot lung cancer risk prediction from five years of long, noisy, multimodal EHR data.

What This Paper Is About

Longitudinal EHR records are too long (often over 100k tokens) and too noisy for standard LLMs to reason over reliably — models hit the "lost-in-the-middle" problem and forget events scattered across years of history. The authors build Traj-CoA, a multi-agent system that splits a patient's record into time-aware chunks, processes them sequentially through a chain of worker agents, and stores salient clinical events in a persistent memory so a final manager agent can make a well-informed risk prediction. The showcased task is predicting whether a patient will be diagnosed with lung cancer within one year of a chest-related radiology exam, using only five years of prior history.

Key Contributions

1. Traj-CoA framework — a chain-of-agents architecture specifically designed for temporal reasoning over long, noisy longitudinal EHRs, using a unified XML input to minimize task-specific feature engineering.
2. Time-aware chunking plus sequential worker agents — heterogeneous EHR data is decomposed into temporally coherent chunks that workers process one at a time, extracting task-relevant signals and filtering out local noise.
3. EHRMem, a structured long-term memory module — stores distilled clinical events and timestamps, with a deduplication mechanism to counter EHR "copy-forwarding," giving the manager agent a global, unabstracted context that complements the worker summaries.
4. Zero-shot state-of-the-art performance — Traj-CoA outperforms machine learning, deep learning, fine-tuned BERT, vanilla LLM, and RAG baselines on lung cancer risk prediction, and its accuracy keeps improving as context scales up to 160k tokens, where standard LLMs plateau or degrade.

Main Findings

• Traj-CoA outperforms all zero-shot baselines and most supervised ones. It reaches AUROC 0.766 ± 0.019 and F1 0.380 ± 0.018, beating logistic regression (0.741), XGBoost (0.763), RETAIN (0.757), PatientTM (0.730), fine-tuned Clinical ModernBERT (0.749), vanilla MedGemma (0.743), and RAG (0.753).

• Scaling beats vanilla LLMs on long context. Vanilla MedGemma actually degrades from AUROC 0.743 at 32k tokens to 0.714 at 64k tokens, while Traj-CoA's AUROC steadily improves as the context window grows from 40k to 160k tokens.

• EHRMem is essential. Removing it drops AUROC by 1.8 points and F1 by 8.1 points, confirming that keeping raw salient events is critical and that the evolving summary alone loses early signals.

• An 8k chunk size is the sweet spot. With total context fixed at 80k tokens, performance peaks at 8k-token chunks; smaller chunks (2k) cause catastrophic forgetting across a long chain, while larger chunks (16k) reintroduce the lost-in-the-middle problem.

• Reasoning is clinically aligned. Traj-CoA's identified salient events span seven clinical categories (diagnosis, procedures, labs, vitals, medications, behaviors/symptoms, demographics). Top themes include advanced age, anemia, COPD, cough, inflammatory markers, lung nodules, pneumonia, pulmonary function, smoking, and weight loss — matching established lung cancer risk knowledge.

• It uses the entire time horizon. Event distributions concentrate in the final year before prediction, as clinically expected, but significant events from earlier years survive into the final output, showing the model does not discard valuable history.

• Lower encoding complexity than vanilla prompting. Traj-CoA's complexity is O(L·L_C) for encoding versus O(L²) for vanilla full-context prompting, though it is more compute-intensive than RAG.

Methodology in Plain English

1. Unify the record. Convert all of a patient's multimodal history — demographics, diagnosis codes, labs, vitals, clinical notes, radiology reports — into a single chronological XML document. This preserves the heterogeneity of the data while putting it in a format LLMs handle well.

2. Chunk by time, not by token. Instead of cutting the record at fixed token counts (which can strip timestamps away from their events), the authors group entries by timestamp and pack them into chunks of up to 8k tokens. If one timestamp's records exceed the limit, it is split while keeping its timestamp attached to each piece.

3. Chain of worker agents. Each worker agent receives one chunk, the previous agent's running summary, and the last few entries from memory. It extracts task-relevant events, updates the summary, and writes new events with timestamps into EHRMem — skipping anything already recorded.

4. Manager agent makes the prediction. The manager receives the final worker summary plus the entire EHRMem timeline and outputs a risk score from 1 to 10. The authors deliberately instruct workers to record a slightly broader set of potentially relevant events, delegating final judgment to the manager, who can see the whole picture.

5. Evaluation. The test set contains 300 held-out instances (28 lung cancer cases, 272 matched controls) drawn from a proprietary case-control dataset. Every instance is anchored to a chest-related radiology exam, and the target is a primary lung cancer diagnosis within one year. The base model is MedGemma-27B, run fully zero-shot.

Why This Matters

Impact on research. The paper demonstrates that agentic, memory-augmented LLM architectures can handle ultra-long (160k-token) clinical sequences where monolithic long-context LLMs fail — an open problem the authors explicitly identify as under-explored. It also bridges two research threads, general-purpose multi-agent reasoning and clinical temporal modeling, with a task-agnostic design that could be repurposed for other longitudinal predictions without retraining.

Real-world applications:

• Lung cancer screening triage — flag patients at elevated one-year risk around the time of a chest X-ray or CT, helping prioritize follow-up imaging and specialist referral.
• Clinical decision support in EHR systems — deploy as a reasoning layer over existing records to surface the specific historical events driving a risk score, giving clinicians an interpretable audit trail.
• Longitudinal risk models for other diseases — the same architecture could be adapted to predict sepsis, heart failure decompensation, or progression of chronic conditions from years of history.
• Retrospective cohort research — mine long patient histories to identify temporal patterns associated with outcomes, with the extracted event timelines serving as structured intermediate data.

Industry relevance. EHR vendors, hospital systems, and health AI companies all face the same engineering problem: how to feed multi-year, multimodal patient records into LLMs without losing middle-context information or paying quadratic attention costs. Traj-CoA offers a training-free pattern — chunk, summarize, remember, then decide — that is compatible with off-the-shelf models and cheaper than full-context prompting. The open-source release (GitHub: zengsihang/Traj-CoA) lowers the barrier to adoption, and the complexity analysis gives engineering teams a concrete latency-versus-completeness trade-off to reason about versus RAG.

Future Directions

• Larger, multi-institution validation. The current evaluation uses a single-institution cohort of 300 test instances and one prediction task; the authors plan broader-scale validation across diverse clinical settings and additional prediction targets.
• More powerful base models and light training. Replacing MedGemma-27B with stronger backbones, adding external knowledge retrieval, or applying multi-agent fine-tuning (e.g., rejection-sampling fine-tuning, which the authors only explored preliminarily) could push performance higher.
• Understanding how, not just what. Current analysis shows which events are salient but not how the model weighs and synthesizes them for different patient subpopulations — a mechanistic question with implications for trust and fairness.
• Reducing prompt dependence. Traj-CoA is task-agnostic in architecture but still requires carefully hand-crafted prompts per task; automated prompt optimization or data-driven hypothesis generation could make the framework easier to generalize.

Target Audience

Clinical NLP and health AI researchers working on longitudinal patient modeling; machine learning engineers building LLM-based clinical decision support; oncologists and screening program designers interested in risk stratification around radiology exams; and multi-agent systems researchers looking for a domain where chain-of-agents with structured memory makes a measurable difference over RAG and vanilla prompting. Readers without a clinical background should still follow the architecture and results, though familiarity with EHR data structures and AUROC/AUPRC trade-offs will help interpret the evaluation sections.

Authors’ abstract

Large language models (LLMs) offer a generalizable approach for modeling patient trajectories, but suffer from the long and noisy nature of electronic health records (EHR) data in temporal reasoning. To address these challenges, we introduce Traj-CoA, a multi-agent system involving chain-of-agents for patient trajectory modeling. Traj-CoA employs a chain of worker agents to process EHR data in manageable chunks sequentially, distilling critical events into a shared long-term memory module, EHRMem, to reduce noise and preserve a comprehensive timeline. A final manager agent synthesizes the worker agents' summary and the extracted timeline in EHRMem to make predictions. In a zero-shot one-year lung cancer risk prediction task based on five-year EHR data, Traj-CoA outperforms baselines of four categories. Analysis reveals that Traj-CoA exhibits clinically aligned temporal reasoning, establishing it as a promisingly robust and generalizable approach for modeling complex patient trajectories. Implementation of Traj-CoA is available on https://github.com/zengsihang/Traj-CoA.