STaR-DRO: Stateful Tsallis Reweighting for Group-Robust Structured Prediction
Organizations: Yale University, New Haven, CT, USA. · Medical Oncology, Yale School of Medicine, New Haven, CT, USA.
Abstract
Structured prediction with large language models requires outputs that are label-accurate, ontology-constrained, structurally valid, and evidence-grounded under label imbalance and heterogeneous group difficulty. We present a unified framework for ontology-constrained generation. First, we introduce a modular prompt-engineering architecture combining XML-style structure, expert disambiguation rules, chain-of-thought reasoning, metadata-aware decision logic, schema contracts, and a self-validation gate. It targets recurrent in-context failures, including format drift, label ambiguity, evidence hallucination, and metadata-conditioned confusion. Second, we propose STaR-DRO, combining Tsallis mirror ascent, sparse entmax-style primal mapback, EMA-smoothed group-loss tracking, rescaled ascent signals, and bounded excess-only multipliers. Unlike conventional DRO, which relies on dense Shannon-entropy exponentiated-gradient updates, can introduce high-variance stochastic reweighting, assigns positive adversarial mass to groups that are not persistently hard, and incurs costs through simplex competition, STaR-DRO upweights only persistently hard groups without suppressing easier ones. We evaluate the framework on EPPC Miner, a clinically grounded high-stakes structured-prediction task requiring hierarchical label prediction and evidence-span extraction from patient-provider secure messages. Across 1B-70B Llama models, prompt engineering improves zero-shot extraction, yielding an average label F1 gain of +14.46 and a Span F1 gain of +17.40. Building on supervised fine-tuning, STaR-DRO further improves accuracy and robustness, increasing average label F1 by +1.08 and +2.20 while reducing mean groupwise validation cross-entropy by 21.3% and 14.8% relative to SFT and standard DRO, respectively. These results advance reliable automated communication mining for patient-centered clinical care analysis.
Figures & tables
| Bench-mark | Domain | Focus | Data Source | Rel./ Socio- Emo. | Bidir. Inter. | Multi- Label | Secure Msg. |
| CBLUE Zhang et al. (2022) | Chinese biomedical | NER, RE, classification | Clinical text, dialogues | ||||
| MedDG Liu et al. (2022) | Chinese medical dialogue | Diagnosis, symptom inquiry | Simulated dialogues | ||||
| ReMeDi Yan et al. (2022) | English medical dialogue | Medication, treatment reasoning | Clinical conversations | ||||
| MediTOD Saley et al. (2024) | English history-taking | Symptom elicitation, reasoning | Simulated dialogues | ||||
| BLURB Gu et al. (2021a) | Biomedical NLP | NER, RE, QA, summarization | Biomedical literature | ||||
| EPPC Miner (Ours) | U.S. secure messaging | Communication functions, relational behaviors | De-identified secure msg. from portals |
| Component | Configuration |
| Models | Instruction-tuned Llama family (Llama-3.2-1B, Llama-3.2-3B, Llama-3.1-8B, Llama-3.3-70B), Qwen2.5 family (Qwen2.5-7B, Qwen2.5-14B), and Gemma-2 family (Gemma-2-2B, Gemma-2-9B). |
| Parameter-efficient tuning | LoRA with rank , scaling factor , dropout . |
| Optimization (default) | bfloat16, gradient checkpointing, AdamW, cosine learning-rate decay, warmup ratio , learning rate , per-device batch size , gradient accumulation , total epochs . |
| Regularization | Weight decay tuned over . |
| Quantization strategy | QLoRA-style 4-bit adaptation for 70B only; all other models fine-tuned without 4-bit quantization. |
| Compute setup | 70B trained on two GPUs; all other models trained on a single GPU. |
| Shot | Category | Model | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | |||
| Zero-shot | Baseline Instruction | Llama-3.3-70B-Instruct | 57.87 | 72.92 | 64.53 | 34.70 | 52.57 | 41.80 | 58.63 | 90.03 | 71.02 |
| Llama-3.1-8B-Instruct | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 0.00 | 28.19 | 87.45 | 42.65 | ||
| Llama-3.2-3B-Instruct | 19.92 | 63.78 | 30.36 | 7.76 | 33.22 | 12.57 | 8.82 | 71.52 | 15.70 | ||
| Llama-3.2-1B-Instruct | 22.50 | 35.69 | 27.60 | 3.51 | 10.07 | 5.21 | 5.60 | 65.97 | 10.31 | ||
| Prompt-Engineered Instruction | Llama-3.3-70B-Instruct | 69.16 | 68.67 | 68.92 | 50.72 | 49.96 | 50.34 | 84.52 | 84.89 | 84.70 | |
| Category | Model | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | ||
| SFT with Baseline Instruction | Llama-3.3-70B-Instruct | 75.98 | 82.21 | 78.98 | 64.24 | 70.84 | 67.38 | 85.08 | 96.34 | 90.36 |
| Llama-3.1-8B-Instruct | 75.59 | 75.94 | 75.77 | 62.77 | 68.70 | 65.60 | 85.69 | 96.99 | 90.99 | |
| Llama-3.2-3B-Instruct | 73.88 | 76.32 | 75.08 | 58.61 | 62.52 | 60.50 | 86.67 | 95.90 | 91.05 | |
| Llama-3.2-1B-Instruct | 67.70 | 65.41 | 66.54 | 53.72 | 52.55 | 53.13 | 87.58 | 93.04 | 90.23 | |
| SFT with Prompt-Engineered Instruction | Llama-3.3-70B-Instruct | 77.41 | 81.16 | 79.24 | 65.92 | 69.75 | 67.78 | 89.97 | 94.93 | 92.38 |
| Category | Model | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | ||
| Standard DRO | Llama-3.3-70B-Instruct | 77.42 | 82.79 | 80.02 | 65.90 | 71.70 | 68.68 | 87.81 | 96.05 | 91.74 |
| Llama-3.1-8B-Instruct | 74.96 | 79.03 | 76.94 | 61.94 | 66.87 | 64.31 | 87.93 | 95.87 | 91.73 | |
| Llama-3.2-3B-Instruct | 75.33 | 78.20 | 76.74 | 62.20 | 64.96 | 63.55 | 89.24 | 94.67 | 91.87 | |
| Llama-3.2-1B-Instruct | 74.12 | 77.78 | 75.91 | 59.82 | 64.06 | 61.86 | 87.71 | 95.94 | 91.64 | |
| STaR-DRO | Llama-3.3-70B-Instruct | 79.08 | 84.00 | 81.47 | 66.26 | 72.64 | 69.30 | 87.57 | 96.30 | 91.73 |
| Variant | Accuracy Change Relative to SFT | Group Loss | ||||
| SFT | 0.3281 | 0.3131 | 0.3582 | |||
| Standard DRO | 0.3014 | 0.2862 | 0.3317 | |||
| STaR-DRO | ||||||
| STaR-DRO w/o bounded excess-only multiplier | 0.3096 | 0.2912 | 0.3463 | |||
| STaR-DRO w/o Stateful EMA | 0.3103 | 0.2978 | 0.3351 | |||
| Grouping Criterion | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | |
| Code | 78.13 | 81.91 | 79.98 | 63.28 | 69.04 | 66.03 | 86.46 | 96.05 | 91.00 |
| Sub-code | 76.66 | 81.79 | 79.14 | 63.49 | 68.97 | 66.12 | 87.17 | 95.80 | 91.28 |
| Number of Annotations | 76.91 | 81.66 | 79.21 | 63.01 | 68.25 | 65.53 | 87.71 | 95.69 | 91.52 |
| Code Sub-code | 76.81 | 81.75 | 79.20 | 62.45 | 68.70 | 65.43 | 86.51 | 96.45 | 91.21 |
| Code Sub-code Number of Annotations | 76.54 | 81.24 | 78.82 | 62.84 | 68.14 | 65.38 | 86.99 | 95.76 | 91.17 |
| Grouping Criterion | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | |
| Code | 76.40 | 81.66 | 78.94 | 62.46 | 68.67 | 65.42 | 86.50 | 96.34 | 91.15 |
| Sub-code | 76.48 | 81.37 | 78.85 | 62.22 | 68.14 | 65.04 | 87.06 | 96.09 | 91.35 |
| Number of Annotations | 77.27 | 81.79 | 79.46 | 62.94 | 68.18 | 65.46 | 87.80 | 96.23 | 91.82 |
| Code Sub-code | 76.13 | 81.24 | 78.60 | 62.83 | 68.93 | 65.74 | 86.54 | 96.01 | 91.03 |
| Code Sub-code Number of Annotations | 76.29 | 82.12 | 79.10 | 62.13 | 68.63 | 65.22 | 86.54 | 96.45 | 91.22 |
| Category | Model | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | ||
| T1: Original (Multi-token) | Qwen2.5-14B-Instruct | 77.86 | 83.12 | 80.40 | 64.84 | 71.33 | 67.93 | 86.75 | 96.34 | 91.29 |
| Llama-3.1-8B-Instruct | 75.96 | 81.45 | 78.61 | 63.58 | 69.75 | 66.52 | 86.86 | 96.52 | 91.43 | |
| Llama-3.2-3B-Instruct | 76.29 | 81.58 | 78.85 | 62.05 | 68.03 | 64.90 | 86.84 | 96.59 | 91.46 | |
| Llama-3.2-1B-Instruct | 75.65 | 80.20 | 77.86 | 61.41 | 65.97 | 63.21 | 86.91 | 96.01 | 91.23 | |
| T2: Token-Formatted (Multi-token) | Qwen2.5-14B-Instruct | 77.38 | 82.75 | 79.98 | 64.08 | 70.54 | 67.15 | 87.11 | 96.01 | 91.34 |
| Setting | Category | Model | Code | Sub-code | Span | ||||||
| Precision | Recall | F1 | Precision | Recall | F1 | Precision | Recall | F1 | |||
| T1: Original (Multi-token) | Qwen2.5-14B-Instruct | 77.25 | 82.41 | 79.75 | 64.62 | 70.24 | 67.31 | 87.70 | 95.65 | 91.50 | |
| Llama-3.1-8B-Instruct | 76.31 | 81.83 | 78.98 | 63.56 | 69.57 | 66.43 | 87.36 | 96.23 | 91.58 | ||
| Llama-3.2-3B-Instruct | 75.92 | 80.62 | 78.20 | 61.65 | 67.13 | 64.27 | 87.68 | 95.94 | 91.62 | ||
| Llama-3.2-1B-Instruct | 75.56 | 78.78 | 77.14 | 61.56 | 64.66 | 63.07 | 89.00 | 95.32 | 92.05 | ||
| T2: Token-Formatted (Multi-token) | Qwen2.5-14B-Instruct | 77.83 | 83.29 | 80.47 | 65.48 | 71.51 | 68.36 | 87.94 | 96.16 | 91.86 | |
Structured Prediction
Structured prediction maps inputs to outputs whose labels, parts, or decisions must fit together as sequences, sets, graphs, or hierarchies.
See all 30 papersStructure Tax: How Structured Output affects LLMs Performance
structure tax'. We re-examine this claim by evaluating a battery of models, datasets and schemas, measuring task accuracy, confidence calibration, and hidden-state geometry. The tax turns out to depend on schema design rather than on structure per se: reasoning-first field ordering matches or exceeds free-form accuracy, while answer-first ordering causes steep drops, particularly in smaller models. Format sensitivity scales inversely with a task's own structural constraints, and schemas that preserve reasoning order also improve calibration with CKA showing greater separability between correct and incorrect representations in middle transformer layers. Our findings indicate that properly designed structured formats can match or exceed free-form performance, reframing the critical question from whether to structure' to `how to structure' for optimal reasoning preservation.