STaR-DRO: Stateful Tsallis Reweighting for Group-Robust Structured Prediction

Samah Fodeh; Ganesh Puthiaraju; Elyas Irankhah; Afshan Khan; Sreeraj Ramachandran; Linhai Ma; Srivani Talakokkul; Sarah Schellhorn

cs.LGApr 9, 2026

STaR-DRO: Stateful Tsallis Reweighting for Group-Robust Structured Prediction

Authors: Samah Fodeh, Ganesh Puthiaraju, Elyas Irankhah, Afshan Khan, Sreeraj Ramachandran, Linhai Ma, Srivani Talakokkul, Sarah Schellhorn

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

Figure 1: Pipeline of EPPC Miner. The pipeline progresses from clinical communication mining and codebook creation to annotation-driven dataset construction, followed by structured prompting, supervised fine-tuning, and STaR-DRO-based robust optimization.

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	$\times$	$\times$	$\times$	$\times$
MedDG Liu et al. (2022)	Chinese medical dialogue	Diagnosis, symptom inquiry	Simulated dialogues	$\times$	$\times$	$\times$	$\times$
ReMeDi Yan et al. (2022)	English medical dialogue	Medication, treatment reasoning	Clinical conversations	$\triangle$	$\triangle$	$\times$	$\times$
MediTOD Saley et al. (2024)	English history-taking	Symptom elicitation, reasoning	Simulated dialogues	$\times$	$\times$	$\times$	$\times$
BLURB Gu et al. (2021a)	Biomedical NLP	NER, RE, QA, summarization	Biomedical literature	$\times$	$\times$	$\times$	$\times$
EPPC Miner (Ours)	U.S. secure messaging	Communication functions, relational behaviors	De-identified secure msg. from portals	$\checkmark$	$\checkmark$	$\checkmark$	$\checkmark$

Table 1: Comparison between EPPC Miner and existing biomedical NLP benchmarks.

\checkmark

indicates strong support,

\times

indicates no support, and

\triangle

indicates partial or limited support. Column abbreviations are as follows: Rel./Socio-Emo. = relational and socio-emotional communication; Bidir. Inter. = bidirectional interaction; Multi-Label = multi-label annotation; Secure Msg. = secure messaging. Additional task abbreviations are: NER = named entity recognition; RE = relation extraction; QA = question answering.

Panel of Figure 2 — Figure 2: General modular prompt template for ontology-constrained structured prediction and evidence extraction.

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 $r=16$ , scaling factor $32$ , dropout $0.05$ .
Optimization (default)	bfloat16, gradient checkpointing, AdamW, cosine learning-rate decay, warmup ratio $0.03$ , learning rate $2\times 10^{-4}$ , per-device batch size $3$ , gradient accumulation $1$ , total epochs $3$ .
Regularization	Weight decay tuned over $\{0.01,\,0.05,\,0.1\}$ .
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.

Table 2: Experimental configuration and evaluation protocol.

Shot	Category	Model	Code			Sub-code			Span
Shot	Category	Model	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

Table 3: Zero-shot, one-shot, and two-shot performance of Llama instruction variants on the EPPC Miner benchmark.

Figure 3: Confusion matrices of selected representative misclassifications in zero-shot Llama-3.2-3B-Instruct, comparing Baseline Instruction (left) and Prompt-Engineered Instruction (right). The left two panels show code-level confusions and the right two panels show subcode-level confusions, highlighting careful error reduction under prompt engineering

Category	Model	Code			Sub-code			Span
Category	Model	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

Table 4: Performance of supervised fine-tuning with baseline and prompt-engineered instructions on the EPPC Miner benchmark.

Category	Model	Code			Sub-code			Span
Category	Model	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

Table 5: Comparison of Standard DRO and STaR-DRO on the EPPC Miner benchmark.

Figure 4: STaR-DRO achieves near-uniform reductions in group-wise loss across Code-level groups and model scales, with the largest reductions emerging at the 70B scale and the strongest gains concentrated in clinically salient categories.

Variant	Accuracy Change Relative to SFT			Group Loss
Variant	$\boldsymbol{\Delta_{\mathrm{SFT}}F1_{\mathrm{Code}}}$	$\boldsymbol{\Delta_{\mathrm{SFT}}F1_{\mathrm{Sub\mbox{-}code}}}$	$\boldsymbol{\Delta_{\mathrm{SFT}}F1_{\mathrm{Span}}}$	$\boldsymbol{\bar{L}_{\mathcal{G}}}$	$\boldsymbol{\bar{L}_{\mathcal{G}_{\mathrm{rare}}}}$	$\boldsymbol{\bar{L}_{\mathcal{G}_{\mathrm{freq}}}}$
SFT	$+0.00$	$+0.00$	$+0.00$	0.3281	0.3131	0.3582
Standard DRO	$-2.11$	$-1.35$	$+0.41$	0.3014	0.2862	0.3317
STaR-DRO	$\mathbf{+1.13}$	$\mathbf{+1.13}$	$-0.46$	$\mathbf{0.2932}$	$\mathbf{0.2855}$	$\mathbf{0.3086}$
STaR-DRO w/o bounded excess-only multiplier	$-0.55$	$+0.14$	$-0.10$	0.3096	0.2912	0.3463
STaR-DRO w/o Stateful EMA	$+0.07$	$+0.11$	$-0.06$	0.3103	0.2978	0.3351

Table 6: Component Ablation of STaR-DRO on Llama-3.2-3B-Instruct. Accuracy columns report absolute changes relative to SFT,

\Delta_{\mathrm{SFT}}F1_{m}=F1_{m}(\mathrm{Variant})-F1_{m}(\mathrm{SFT})

, for

m\in\{\mathrm{Code},\mathrm{Sub\mbox{-}code},\mathrm{Span}\}

\bar{L}_{\mathcal{G}}

denotes the macro-average cross-entropy over all Code-level groups, computed on the validation set.

\bar{L}_{\mathcal{G}_{\mathrm{rare}}}

and

\bar{L}_{\mathcal{G}_{\mathrm{freq}}}

denote macro-average cross-entropy over rare and frequent groups, respectively. Rare groups contain fewer than 50 annotations, whereas frequent groups contain at least 100 annotations.

Figure 5: Effect of

\ell_{2}

regularization strength on per-group validation cross-entropy loss under STaR-DRO on Llama-3.2-3B-Instruct. Teal:

\lambda_{\mathrm{wd}}=0.1

; red:

\lambda_{\mathrm{wd}}=0.01

. Arrows indicate the direction and magnitude of loss reduction (

\Delta

loss at right). Dot size reflects group sample count. Stronger weight decay uniformly reduces loss across all nine Code-level groups.

Figure 6: F1 performance and group-wise robustness of Qwen2.5-7B-Instruct on the EPPC Miner benchmark.

Figure 7: F1 performance and group-wise robustness of Qwen2.5-14B-Instruct on the EPPC Miner benchmark.

Figure 8: F1 performance and group-wise robustness of gemma-2-2b-it on the EPPC Miner benchmark.

Figure 9: F1 performance and group-wise robustness of gemma-2-9b-it on the EPPC Miner benchmark.

Grouping Criterion	Code			Sub-code			Span
Grouping Criterion	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 $\times$ Sub-code	76.81	81.75	79.20	62.45	68.70	65.43	86.51	96.45	91.21
Code $\times$ Sub-code $\times$ Number of Annotations	76.54	81.24	78.82	62.84	68.14	65.38	86.99	95.76	91.17

Table 7: Performance of STaR-DRO under different grouping criteria on the EPPC Miner benchmark for Llama-3.2-3B-Instruct.

Grouping Criterion	Code			Sub-code			Span
Grouping Criterion	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 $\times$ Sub-code	76.13	81.24	78.60	62.83	68.93	65.74	86.54	96.01	91.03
Code $\times$ Sub-code $\times$ Number of Annotations	76.29	82.12	79.10	62.13	68.63	65.22	86.54	96.45	91.22

Table 8: Performance of STaR-DRO under different grouping criteria on the EPPC Miner benchmark for Llama-3.2-3B-Instruct. Results are reported for the sample-level multi-grouping setting.

Figure 11: Converged group-weight distributions

\mathbf{q}^{*}

across eight STaR-DRO configurations on the nine Code-group EPPC Miner task with Llama-3.2-3B-Instruct. The baseline (

\eta_{t}{=}0.003

\rho_{t}{=}0.03

\alpha{=}1.1

, Epochs 2 & 3; red bars) yields a balanced two-peak profile. Increasing

\eta_{t}

\alpha

drives concentration toward fewer groups; insufficiently strong settings produce near-uniform distributions. The dashed line denotes

1/G

Figure 12: The distribution of EPPC codes based on the annotation.

Figure 13: The distribution of EPPC subcodes based on the annotation.

Figure 14: EPPC Miner communication codes and their associated Sub-codes.

Panel of Figure 15 — Figure 15: Illustration of the code–subcode relationship. Codes 1–2 share identical subcode set A (SDOH pair), Codes 3–4 largely share set B (InfoGive/InfoSeek pair, with one unique subcode in Code 3), Codes 5–6 largely share set C (Partnership pair, with several unique subcodes each), while Codes 7–9 each correspond to fully disjoint subcode sets (D–F). Each subcode set (A–F) is a subset of the subcode inventory $L_{sub}$ , i.e., $\text{Set}\subseteq 2^{L_{\text{Sub-code}}}$ , consistent with the hierarchical mapping $\mathcal{H}:L_{\text{Code}}\to 2^{L_{\text{Sub-code}}}$ .

Figure 16: Label encoding schemes used in the sensitivity analysis, showing four alternative representations of EPPC Miner codes and Sub-codes: original multi-token labels, token-formatted labels, numeric IDs, and simplified aliases.

Category	Model	Code			Sub-code			Span
Category	Model	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

Table 9: Label Encoding Sensitivity under SFT. Results are reported as Precision, Recall, and F1. T1 = Original (Multi-token), T2 = Token-Formatted (Multi-token), T3 = Numeric IDs (Single-token), and T4 = Simplified Aliases (Single-token).

Setting	Category	Model	Code			Sub-code			Span
Setting	Category	Model	Precision	Recall	F1	Precision	Recall	F1	Precision	Recall	F1
$w_{\text{Code}}=3,\;w_{\text{Sub-code}}=3,\;w_{\text{Span}}=3$	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

Table 10: Label Encoding Sensitivity under SFT + Token Weighting. Results are reported as Precision, Recall, and F1 in percentage. Panel A uses

w_{\text{Code}}=3

w_{\text{Sub-code}}=3

, and

w_{\text{Span}}=3

. Panel B uses

w_{\text{Code}}=2

w_{\text{Sub-code}}=1.5

, and

w_{\text{Span}}=1

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 papers

Equitable System-Prompt Selection via Constrained Mixed-Strategy GroupDRO
Aug 5, 2026Mengyu Xu, Qiaoxin Yang, Zhihan Liu +4Group Distributionally Robust Optimization LLM Prompting
1The University of Chicago · 2SynAI Technologies, Inc. · 3Northwestern University +2
Structure Tax: How Structured Output affects LLMs Performance
Oct 8, 2026Vineet Kumar, Kanishka, Bhuvanesh MandoraStructured Output Generation Structured Reasoning in Language Models
PayPal Artificial Intelligence PayPal, Bengaluru India
Revisiting Semantic Role Labeling: Efficient Structured Inference with Dependency-Informed Analysis
May 4, 2026Sangpil Youm, Leah Jones, Bonnie J. DorrStructured Prediction
University of Florida, Gainesville, Florida

Aug 5, 2026J/K move · Enter open · S save

Equitable System-Prompt Selection via Constrained Mixed-Strategy GroupDRO

Mengyu Xu, Qiaoxin Yang, Zhihan Liu, Ruiyao Xu+3

1The University of Chicago · 2SynAI Technologies, Inc. · 3Northwestern University+2

Group Distributionally Robust Optimization LLM Prompting Algorithmic Fairness Distributionally Robust Optimization Prompt Optimization Prompt Engineering

Open paper AskPDF arXiv

Large language models are increasingly used for information seeking, yet semantically equivalent questions phrased in different ways can receive answers of considerably different quality. System prompts are widely employed to steer response behavior, but they are typically optimized for average-case quality, so some question phrasings may still receive incomplete or low-quality answers. To address this, we formulate a constrained mixed-strategy GroupDRO framework for system-prompt selection. Instead of optimizing the system-prompt text, the framework assigns weights to system prompts in an existing pool to minimize the worst-case information-quality loss across evaluation metrics and groups, while constraining the mean loss to stay close to that of average-based selection. Because pool generation and selection are decoupled, the method applies to any system-prompt pool and can leverage an ensemble of complementary system prompts rather than a single one. Across five LLMs on two bilingual medical and consumer-finance benchmarks, the constrained method reduces the Overall Mean, Worst 25% Mean, and Worst by 13.1%, 13.2%, and 13.7% on average relative to no mitigation while keeping overall quality close to Average selection. Its multi-prompt weights reveal complementarity across metric-group pairs. Code and data are available at https://github.com/Rainxu09/equitable-system-prompt-selection.

Abstract

Structured Prediction

Explore similar work

Equitable System-Prompt Selection via Constrained Mixed-Strategy GroupDRO

Structure Tax: How Structured Output affects LLMs Performance

Revisiting Semantic Role Labeling: Efficient Structured Inference with Dependency-Informed Analysis