Technical Summary: Expert-guided Clinical Text Augmentation via Query-Based Model Collaboration

1. Problem Formulation

Data augmentation is a standard strategy for improving model robustness and generalization by enriching training datasets with synthetic examples. However, in high-stakes domains such as healthcare, standard augmentation techniques face a critical safety challenge: the risk of generating clinically incorrect or misleading information (hallucinations).

The core problem addressed is the inability of general-purpose Large Language Models (LLMs) to distinguish between critical domain variables ($V$) (e.g., specific symptoms, diagnoses, medication dosages) and non-critical stylistic variables ($U$) (e.g., sentence structure, phrasing, patient names).

• The Risk: Naive LLM-based augmentation often alters or removes critical medical information ($V$) while attempting to modify stylistic elements ($U$), or it introduces fabricated medical facts. This distorts the semantic integrity of the data, potentially confusing downstream models and propagating errors that could impact clinical decisions.
• The Gap: Existing methods, including those using domain-specific pre-trained LLMs or causal prompting (e.g., CATO), lack explicit, fine-grained safeguards to ensure that critical medical entities remain invariant during the augmentation process.

2. Methodology: Query-Based Model Collaboration

The authors propose a lightweight, query-based model collaboration framework that integrates expert-level domain knowledge to guide the augmentation process without requiring the retraining of the generative model. The framework consists of two primary components:

1. Weak Expert ($W$): A lightweight, domain-specific model (e.g., a BERT-based Clinical NER model) trained to identify safety-critical tokens (medical keywords, dosages, findings) within the input text. It acts as a decision-maker that flags these tokens as "unalterable."
2. Strong Generalist ($G$): A general-purpose foundation model with strong generative and rewriting capabilities (e.g., Qwen, Llama) but without explicit training in the target clinical domain.

The Augmentation Pipeline:
The process follows a three-step procedure for each input clinical note $x_i$:

1. Critical Feature Extraction: The Weak Expert $W$ analyzes $x_i$ to extract a set of critical tokens $K_i$ that must be preserved.
2. Prompt Construction: A constraint-aware prompt is constructed, passing the original text $x_i$ and the explicit constraint set $K_i$ to the Strong Generalist. The prompt instructs the model to rewrite the text while preserving the semantics and verbatim presence of all tokens in $K_i$.
3. Constrained Rewriting: The Strong Generalist $G$ generates the augmented text $\tilde{x}_i$ conditioned on the prompt. It is free to rephrase non-critical elements (style, syntax) but is constrained from altering the critical medical information.

Distillation Option:
The paper also explores distilling this collaborative behavior into a single model using Direct Preference Optimization (DPO). In this setup, the Strong Generalist is trained to prefer outputs generated via the expert-guided collaboration over naive augmentations, effectively internalizing the expert's constraints.

3. Key Contributions

• Novel Framework: Introduction of a model collaboration framework where a lightweight "weak expert" supervises a "strong generalist" to ensure safe, domain-faithful text augmentation.
• Safety and Quality Metrics: Demonstration that the proposed method significantly improves the Preservation Rate (PR) of critical medical entities and reduces the Hallucination Rate (HR) at both token and concept levels compared to existing baselines.
• Downstream Performance: Empirical evidence that augmented data generated by this method leads to consistent performance gains in downstream clinical prediction tasks (readmission, mortality, length-of-stay) and zero-shot/few-shot inference tasks (phenotyping, ICD coding).
• Efficiency: Validation that the framework achieves these safety gains with negligible computational overhead compared to standard LLM augmentation, avoiding the high costs of training-based alternatives.

4. Experimental Results

Experiments were conducted on the MIMIC-III dataset across three clinical prediction tasks and zero-shot inference benchmarks.

• Data Quality:

  • Naive LLM augmentation removed ~49% of medical keywords and introduced ~59% hallucinated keywords.
  • The proposed method achieved a Preservation Rate (PR) of 0.79 (token level) and 0.73 (concept level), while maintaining a Hallucination Rate (HR) of 0.33 (token level) and 0.26 (concept level). This significantly outperformed baselines like Naive LLM, CATO, and rule-based synonym substitution.
  • Human expert annotations confirmed that the proposed method preserved critical information more effectively than baselines.

• Downstream Task Performance:

  • Models trained on the augmented data achieved the best mean performance across all tasks (Readmission, Mortality, Length-of-Stay) compared to models trained on naive augmented data or no augmentation.
  • In zero-shot and few-shot settings (Patient Phenotyping and ICD Coding), the proposed method consistently improved F1 scores, whereas naive augmentation often degraded performance relative to the original data.

• Component Analysis:

  • Using a medical-specific Weak Expert yielded better results than a general-purpose NER model, though the general model still provided improvements over no guidance.
  • Larger Strong Generalists (e.g., 3B vs 1B parameters) showed better preservation rates, but the collaboration framework consistently outperformed baselines across all model sizes.
  • Distillation: While DPO-trained single models showed improvement over zero-shot baselines, the dual-model collaboration remained the most reliable approach, particularly for models lacking strong medical priors.

• Cost Analysis:

  • The method adds only a small inference-time cost (~18% increase over naive augmentation) due to the lightweight Weak Expert pass.
  • It avoids the substantial training costs (hundreds of GPU hours) associated with fine-tuning large LLMs for augmentation, making it a cost-efficient alternative.

5. Significance and Claims

The paper claims to bridge the gap between the generative potential of LLMs and the strict safety requirements of specialized domains like healthcare. By explicitly separating the roles of knowledge extraction (Weak Expert) and text generation (Strong Generalist), the framework ensures that synthetic data remains clinically valid.

The authors emphasize that their approach:

1. Preserves Critical Information: Prevents the distortion of medical facts that is common in unguided LLM augmentation.
2. Reduces Hallucinations: Significantly lowers the rate of fabricated medical entities.
3. Improves Generalization: The resulting synthetic data helps downstream models learn robust features, improving performance across diverse clinical tasks.
4. Offers a Practical Solution: Provides a lightweight, inference-time solution that does not require expensive retraining of foundation models, making it scalable for safety-critical applications.

The work concludes that while general-purpose LLMs are powerful, they require explicit, domain-specific guidance to be safely utilized for data augmentation in high-stakes environments. The proposed query-based collaboration offers a viable path to achieving this safety without sacrificing computational efficiency.