Dynamic Minimax Regret Optimization for Robust LLM Post-Training

Chengbo Zang; Haoyu Dong; Mehmet Kerem Turkcan; Gil Zussman; Zoran Kostic; Javad Ghaderi

cs.LGOct 5, 2026

Dynamic Minimax Regret Optimization for Robust LLM Post-Training

Authors: Chengbo Zang, Haoyu Dong, Mehmet Kerem Turkcan, Gil Zussman, Zoran Kostic, Javad Ghaderi

Organizations: Department of Electrical Engineering, Columbia University · Department of Computer Science & Engineering, University of California, Santa Cruz · Department of Civil Engineering, Columbia University

Abstract

Modern LLM training increasingly relies on heterogeneous data sources spanning different domains, tasks, preference distributions, and difficulty levels. We study dynamic minimax regret for group-distributionally robust LLM post-training under instantaneous mini-batch-only bandit feedback. The framework views the training as a two-player sampler-optimizer process: a sampler adaptively selects among data sources using bandit feedback, while an optimizer updates the model parameters using stochastic gradients from the selected source. We focus on the practically restrictive setting where source losses evolve with model training but historical data are not re-evaluated, requiring the sampler to track instantaneous worst-sources from stale partial feedback. We propose DUCB-OGD, a simple and scalable algorithm that couples a Discounted Upper-Confidence-Bound sampler with an Online Gradient Descent optimizer. The sampler maintains exponential moving average loss estimates and confidence radii based on discounted effective sample sizes, avoiding costly re-evaluation of past data or intrusive changes to standard training pipelines. For $K$ data sources and $T$ training steps, we prove that DUCB-OGD achieves a dynamic minimax regret of $\tilde{O}(K^{1/4}T^{3/4})$ , which is optimal up to logarithmic factors for the undiscounted objective under our feedback model. Extensive experiments across supervised fine-tuning, preference optimization, and reinforcement learning show that DUCB-OGD integrates seamlessly into modern LLM training pipelines and improves worst-group robustness with negligible computational overhead compared with standard sampling baselines.

Figures & tables

Objective	$\bm{-\ell(\theta,\cdot)}$	WS Switches / 1k Steps
SFT	$\mathbb{E}_{o\sim\mathcal{D}^{k}(\cdot\mid q)}\left[\log\pi_{\theta}(o\mid q)\right]$	41 [3,124]
DPO	$\mathbb{E}_{(o^{+},o^{-})\sim\mathcal{D}^{k}(\cdot\mid q)}\left[\log p_{\theta}^{\mathrm{BT}}(o^{+}\succ o^{-}\mid q)\right]$	31 [13,53]
RL (GRPO)	$\mathbb{E}_{o\sim\pi_{\theta}(\cdot\mid q)}\left[r(q,o)\right]$	20 [0,32]

Table 1: GDRO formulations of LLM post-training objectives and estimated worst-source (WS) switching. The final column reports the mean number of switches per 1,000 steps, with the range across experiments in brackets. The worst source has the largest exponentially smoothed loss for SFT/DPO or the lowest exponentially smoothed reward for GRPO. Additional optimization dynamics are reported in Sections 4 and D.5 .

Objective	Dataset		Model	Sampler	Accuracy
	Train	Test			Min	Avg
SFT	Tulu-v2 (subset)	GSM8k + IFEval	Qwen2.5-7B (Base)	Uniform	27.15	39.36
				Exp3-style ∗	27.73	40.72
				DUCB (Ours)	28.52	37.89
			Llama3.1-8B (Base)	Uniform	22.07	25.98
				Exp3-style ∗	21.48	27.93

Table 2: Main results of different sampling algorithms on various LLM post-training objectives across different model scales and families. The primary metric is minimum source-wise accuracy.

Panel of Figure 1 — Figure 1: GRPO training rewards and sampling weights for uniform sampling, DUCB (guided by dynamic regret), and Exp3 (guided by pseudo regret) on two data sources ( Arc and Countdown ).

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

p_{t}\leftarrow\textsc{Sampler}(\{\mathcal{D}^{k},\hat{\mu}_{t-1}^{k},N_{t-1}^{k}(\beta)\}_{k=1}^{K})

Appendix

Algorithm 1 Adaptive Sampling for Robust Multi-Source LLM Post-Training.

Model	Sampler	Signal	Accuracy		Mean Token Accuracy
Model	Sampler	Signal	Min	Avg	Min	Avg
Qwen2.5-3B (Base)	Uniform	Loss	22.85	29.79	81.23	84.97
	Exp3	Loss	23.63	25.68	78.82	82.35
	Exp3	Loss Difference 1	23.63	27.24	81.46	84.78
	DUCB (Ours)	Loss	24.02	26.95	81.35	84.71
Qwen2.5-7B (Base)	Uniform	Loss	27.15	39.36	78.79	82.34

Appendix

Table 3: SFT results on the GSM8k + IFEval datasets with a budget of 32k samples.

Model	Sampler	Signal	Accuracy		Margin
Model	Sampler	Signal	Min	Avg	Min	Avg
Qwen2.5-3B (Instruct)	Fixed-Weight	-	75.00	80.57	0.96	1.28
	Uniform	Loss	76.76	80.53	1.08	1.31
	Exp3	Excess Loss	77.78	79.98	1.12	1.40
	AMA-R 1	Excess Loss	78.19	79.45	1.03	1.30
	AMA-S 1	Excess Loss	80.25	82.35	1.07	1.31

Appendix

Table 4: DPO results on the Helpsteer3 dataset with a budget of 16k samples.

Sampler	Runs	Min Acc.	Avg Acc.
Uniform	3	76.76 $\pm$ 1.65	80.53 $\pm$ 0.70
DUCB (Ours)	3	79.33 $\pm$ 0.62	81.59 $\pm$ 0.40
Paired gain	3	+2.57 $\pm$ 1.04	-

Appendix

Table 5: Three-seed DPO results on HelpSteer3 with Qwen2.5-3B-Instruct. We report final-checkpoint mean accuracy

\pm

standard deviation. The paired gain compares DUCB and Uniform under matched seeds.

Model	Sampler	Signal	Accuracy
Model	Sampler	Signal	Min	Avg
Qwen2.5-0.5B (Instruct)	Uniform	Reward	14.73	23.01
	Exp3	Reward	16.01	22.82
	Exp3	Reward + Difference Ramesh et al. (2026)	16.30	23.23
	SWUCB (DUMP)	Advantage Wang et al. (2025b)	13.83	22.61
	DUCB (Ours)	Reward	16.53	22.88

Appendix

Table 6: GRPO results on the Reasoning dataset with a budget of 10k samples.

Model	Sampler	Signal	Accuracy
Model	Sampler	Signal	Min	Avg
Qwen3-8B (Instruct)	Uniform	Reward	81.40	83.50
	Exp3	Reward	77.22	85.52
	Exp3	Reward + Difference Ramesh et al. (2026)	79.75	82.81
	SWUCB (DUMP)	Advantage Wang et al. (2025b)	73.26	76.17
	DUCB (Ours)	Reward	81.40	87.20

Appendix

Table 7: GRPO results on the CommonsenseQA dataset with a budget of 6.4k samples.

Sampler	Time/Step (s)	Additional Cost
Uniform	21.90 $\pm$ 1.29	None
AMA-R †	21.47 $\pm$ 1.24	Training $K$ source-specific specialists
AMA-S †	21.72 $\pm$ 1.36	Training $K$ source-specific specialists
Zang et al. ‡	21 $+$ 0.36 $\times\mathrm{BufSize}_{t}$	$O(t)$ replay evaluations at step $t$
DUCB (Ours)	21.94 $\pm$ 1.14	$O(K)$ scalar state
† The reported runtime excludes specialist training;

Appendix

Table 8: Training overhead on HelpSteer3 with Qwen2.5-3B-Instruct using the same A100 GPU, sample budget, and hyperparameters. Per-step time is reported as mean

\pm

standard deviation. AMA runtimes cover generalist training but exclude the cost of training

K

specialists. The overhead for Zang et al. (2025) is analytical rather than a wall-clock measurement.

Task	Max LR	Max Grad. Norm mean / max	Max Signal Variation mean / max	WS Switches / 1k mean [min, max]
SFT	2e-5	190.24 / 298.79	0.07 / 0.13	41 [3,124]
DPO	1e-5	3.94 / 5.14	0.15 / 0.17	31 [13,53]
GRPO	1e-5	0.22 / 0.24	0.03 / 0.06	20 [0,32]

Appendix

Table 9: Optimization dynamics across the SFT, DPO, and GRPO experiments. For gradient norms and per-source signal variation, we report the mean and maximum of the run-level maxima. Worst-source variation is measured using exponentially smoothed training loss for SFT/DPO and reward for GRPO. Estimated worst-source (WS) switches are normalized per 1,000 steps.

Explore similar work

Mar 1, 2025cs.LG

Distributionally Robust Reinforcement Learning with Human Feedback

Reinforcement learning from human feedback (RLHF) has evolved to be one of the main methods for fine-tuning large language models (LLMs). However, existing RLHF methods are non-robust, and their performance deteriorates if the downstream task differs significantly from the preference dataset used in fine-tuning. In order to mitigate this problem, we introduce a distributionally robust RLHF for fine-tuning LLMs. In particular, our goal is to ensure that a fine-tuned model retains its performance even when the distribution of prompts significantly differs from the distribution encountered during fine-tuning. We formulate distributionally robust optimization (DRO) version of two popular fine-tuning methods -- (1) reward-based RLHF and (2) reward-free DPO (direct preference optimization). We propose a minibatch gradient descent based algorithms for both of them, and theoretically prove convergence guarantees for the algorithms. Subsequently, we evaluate our algorithms on an out-of-distribution (OOD) task by first training the model on the Unified-Feedback dataset and evaluating its performance on two different datasets. The experimental results show that our robust training improves the accuracy of the learned reward models on average, and markedly on some tasks, such as reasoning. Furthermore, we show that the robust versions of policy optimization methods, similarly improve performance on OOD tasks.

Debmalya Mandal, Paulius Sasnauskas, Goran Radanovic

Dept. of Computer Science University of Warwick, UK · Department of Computing Science University of Alberta, Edmonton, Canada · Max-Planck Institute for Software Systems, Germany

Reinforcement Learning From Human Feedback Distributional Reinforcement Learning Offline Reinforcement Learning

Ask PDF

Jun 4, 2026cs.LG

Cross-Epoch Adaptive Rollout Optimization for RL Post-Training

LLM post-training often relies on reinforcement learning methods that sample multiple rollouts per prompt, yet most existing approaches use a fixed rollout budget for every prompt, despite large differences in the training signal different prompts provide. In this paper, we study adaptive rollout allocation under a fixed global budget and formulate the problem as online resource allocation with prompt-level diminishing returns. Our method, CERO, maintains a Beta posterior over each prompt's success probability and uses the posterior expected Bernoulli variance as a Bayesian estimate of the value of additional rollouts. We use this estimate to construct a concave, saturating utility over cumulative allocations, yielding an objective in which decisions across prompts and epochs are coupled by the global budget. Since the resulting objective is temporally nonseparable, we derive a Fenchel-dual reformulation and update both prompt-level and budget-level dual variables via projected online gradient descent. Under fixed prompt utilities, we prove an

O(\sqrt{K})

regret bound against the offline allocation benchmark. Experiments on mathematical-reasoning problems show that CERO consistently outperforms GRPO across multiple open-weight LLMs and benchmarks, demonstrating that adaptive rollout budgeting can improve sample efficiency.

Yiming Zong, Yige Wang, Jiashuo Jiang

Department of Industrial Engineering & Decision Analytics, Hong Kong University of Science and Technology

Topic pending

Ask PDF

Apr 30, 2026cs.LG

Wasserstein Distributionally Robust Regret Optimization for Reinforcement Learning from Human Feedback

Reinforcement learning from human feedback (RLHF) is a central post-training tool for aligning large language models, but its training reward is only a learned proxy for true human utility. This creates a decision problem under objective misspecification: the policy is optimized against an estimated reward, while deployment performance is governed by an unobserved population preference. The resulting gap leads to reward over-optimization, where proxy reward keeps improving after true quality deteriorates. We propose distributionally robust regret optimization (DRRO) for RLHF with a Wasserstein ambiguity set over reward laws, using promptwise

\ell_p

distances between reward vectors as transport costs. Unlike standard distributionally robust optimization, which pessimizes worst-case value, DRRO pessimizes worst-case regret relative to the best policy under the same plausible reward perturbation. We show that the expressive-policy problem decomposes into promptwise regret problems. For each prompt, the inner adversary has a dual-norm closed form; under the

\ell_1

transport cost used by our algorithm, the optimizer has a water-filling structure. These results lead to a practical policy-gradient algorithm that adds a simple sampled bonus to GRPO-style training. Theory and experiments both show that DRRO is less over-pessimistic than standard DRO and mitigates over-optimization more effectively than existing baselines.

Yikai Wang, Shang Liu, Jose Blanchet

Department of Statistics and Operations Research, University of North Carolina · Imperial Business School, Imperial College London · Department of Management Science and Engineering, Stanford University

Reinforcement Learning From Human Feedback Distributionally-Robust Optimization

Ask PDF