Improving Mixup Calibration with Wasserstein Distributionally Robust Optimization

Jiaming Hu; Yeping Jin; Debarghya Mukherjee; Ioannis Ch. Paschalidis

stat.MLJun 22, 2025

Improving Mixup Calibration with Wasserstein Distributionally Robust Optimization

Authors: Jiaming Hu, Yeping Jin, Debarghya Mukherjee, Ioannis Ch. Paschalidis

Organizations: Department of Mathematics & Statistics Boston University · Department of Electrical and Computer Engineering, Division of Systems Engineering Department of Biomedical Engineering Faculty of Computing & Data Sciences Boston University

Abstract

In many real-world applications, ensuring the robustness and stability of deep neural networks (DNNs) is crucial, particularly for image classification tasks that encounter various input perturbations. While Mixup-based data augmentation techniques have been widely adopted to enhance the resilience of trained models against such perturbations, our experiments reveal an important corruption robustness-calibration trade-off: stronger Mixup-based augmentation can improve robustness against corrupted data while substantially increasing expected calibration error (ECE). To address this challenge, we introduce DRO-Augment, a framework that integrates Wasserstein Distributionally Robust Optimization (W-DRO) with various Mixup-based data augmentation strategies to mitigate this trade-off. Our method substantially reduces ECE under strong Mixup-based augmentation while largely preserving corruption accuracy across CIFAR-10, CIFAR-100, CIFAR-10-C, and CIFAR-100-C. On the theoretical side, we establish novel generalization error bounds for neural networks trained using a variation-regularized loss function with augmented data, closely related to the W-DRO problem. Furthermore, we introduce a refined CIFAR-C benchmark that corrects inconsistencies in corruption intensities, providing a more reliable evaluation for future robustness research.

Figures & tables

Method	CIFAR-10				CIFAR-100
	Clean		-C		Clean		-C
	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$
Baseline	94.74	3.34	75.45	17.56	77.99	10.07	49.40	27.01
Mixup	95.60	14.03	79.25	12.44	79.12	14.92	53.29	11.64
Mixup + DRO	96.04	7.54	80.75	7.25	80.25	8.95	54.24	9.80
Manifold Mixup	95.83	17.44	78.71	15.11	80.19	12.83	53.31	12.73

Table 1: Average accuracy (%) and ECE (%) of PreActResNet-18 on CIFAR and CIFAR-C benchmarks. Higher accuracy is better and lower ECE is better.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Method	CIFAR-10				CIFAR-100
	Clean		-C		Clean		-C
	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$
Baseline	94.86	3.24	74.37	17.95	75.95	11.67	44.51	33.58
Mixup	94.72	20.95	77.96	16.69	75.17	16.58	48.72	13.05
Mixup + DRO	95.77	12.12	78.72	9.67	77.36	5.78	49.71	12.62
Manifold Mixup	95.81	12.12	75.38	11.41	77.79	8.36	48.84	13.12

Appendix

Table 2: Average accuracy (%) and ECE (%) of WideResNet-28

\times

2 on CIFAR and CIFAR-C benchmarks. Higher accuracy is better and lower ECE is better.

Method	$\rho$	CIFAR - 10				CIFAR - 100
		Clean		- C		Clean		- C
		Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$
Mixup + DRO	$1.0\times 10^{-4}$	96.06	6.30	80.08	7.20	80.15	5.77	54.41	9.24
	$2.0\times 10^{-4}$	96.09	6.03	79.52	6.84	80.17	7.57	54.43	10.54
	$3.0\times 10^{-4}$	95.86	6.46	79.74	7.92	80.38	9.81	54.44	9.73
	$5.0\times 10^{-4}$	96.13	6.18	80.04	6.79	80.07	8.08	54.31	9.13

Appendix

Table 3: Sensitivity study of the DRO radius

\rho

for Mixup + DRO, NoisyMix + DRO, and Manifold Mixup + DRO on CIFAR and CIFAR-C benchmarks using PreActResNet-18. Metrics follow Table 1 : higher accuracy is better and lower ECE is better.

Method	$\alpha_{\mathrm{mix}}$	CIFAR - 10				CIFAR - 100
		Clean		- C		Clean		- C
		Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$	Acc. $\uparrow$	ECE $\downarrow$
Mixup [+DRO]	0.2	95.69 [95.75]	2.52 [1.12]	76.97 [77.99]	10.14 [10.98]	79.02 [79.33]	3.53 [1.40]	52.17 [51.96]	8.46 [10.69]
	0.4	95.84 [95.97]	11.38 [2.88]	78.45 [79.46]	10.89 [7.87]	78.66 [79.62]	12.62 [6.00]	51.71 [52.15]	10.90 [9.37]
	0.6	95.60 [96.15]	11.32 [4.12]	78.66 [80.29]	10.44 [6.56]	79.54 [80.02]	12.80 [7.28]	53.03 [53.44]	10.16 [8.85]
	1.0	95.69 [96.01]	14.70 [5.95]	79.16 [79.69]	12.57 [7.47]	79.69 [80.31]	14.96 [8.32]	52.83 [55.00]	11.69 [9.21]

Appendix

Table 4: Sensitivity study of the mixing-strength parameter

\alpha_{\mathrm{mix}}

on CIFAR and CIFAR-C benchmarks using PreActResNet-18. Higher accuracy is better and lower ECE is better. Bracketed values report the corresponding method with DRO.

\textstyle\mathcal{L}_{total}=\sum_{i}\mathcal{L}(f_{\theta}(x_{i}),y_{i})+\rho\Bigl(\mathbb{E}_{P_{n}}\left[\|\nabla_{x}\mathcal{L}(f_{\theta}(x_{i}),y_{i})\|^{q}_{*}\right]\Bigr)^{\!\!1/q};

Appendix

Algorithm 1 Training with DRO-Augment

Model	S1	S2	S3	S4	S5
ResNet-18	86.04	79.03	70.00	61.83	51.59
ResNet-50	88.10	80.18	70.18	61.69	50.68
ResNet-101	87.60	79.54	69.51	60.85	49.74

Appendix

Table 5: Performance on Refined CIFAR–10–C (% accuracy)

Model	S1	S2	S3	S4	S5
ResNet-18	59.93	50.78	40.85	30.77	20.53
ResNet-50	65.11	55.83	45.51	35.05	23.86
ResNet-101	63.80	54.58	44.65	34.53	23.87

Appendix

Table 6: Performance on Refined CIFAR–100–C (% accuracy)

Type	S1	S2	S3	S4	S5
gaussian	86.20	78.01	69.30	61.61	51.22
shot	86.51	77.16	69.00	60.36	50.88
impulse	85.01	78.65	69.19	59.73	50.82
defocus	87.03	79.53	69.27	60.50	46.63
glass	84.88	79.16	67.59	60.73	52.56
motion	85.30	78.56	70.38	60.05	52.34

Appendix

Table 7: Accuracy of ResNet-18 on Refined CIFAR-10-C under different severity levels

Type	S1	S2	S3	S4	S5
gaussian	88.40	80.30	70.74	62.44	50.43
shot	89.17	79.85	71.92	62.72	51.95
impulse	86.97	81.18	70.95	61.20	50.78
defocus	89.28	78.65	64.31	53.83	39.72
glass	86.69	80.06	67.51	61.80	52.19
motion	87.08	78.72	69.60	59.01	49.99

Appendix

Table 8: Accuracy of ResNet-50 on Refined CIFAR-10-C under different severity levels

Type	S1	S2	S3	S4	S5
gaussian	87.93	79.14	70.12	61.87	51.49
shot	88.16	78.44	69.57	61.25	51.38
impulse	86.71	80.30	70.53	60.33	51.44
defocus	88.89	78.76	65.34	54.08	38.82
glass	86.35	79.18	67.53	63.04	51.37
motion	86.86	78.19	69.17	57.63	48.23

Appendix

Table 9: Accuracy of ResNet-101 on Refined CIFAR-10-C under different severity levels

Type	S1	S2	S3	S4	S5
gaussian	61.88	49.74	39.15	31.80	19.09
shot	61.45	50.45	42.01	29.19	18.49
impulse	61.29	52.34	42.23	29.00	22.49
defocus	61.01	50.95	40.08	33.01	21.55
glass	59.79	50.35	38.33	25.36	19.48
motion	59.86	52.02	41.66	32.28	20.51

Appendix

Table 10: Accuracy of ResNet-18 on Refined CIFAR-100-C under different severity levels

Type	S1	S2	S3	S4	S5
gaussian	66.76	54.96	44.10	36.41	22.75
shot	66.23	56.02	47.55	34.17	22.47
impulse	66.29	59.35	48.46	35.49	28.00
defocus	66.02	55.38	42.41	35.78	23.08
glass	65.20	53.53	40.20	27.68	21.45
motion	64.95	55.80	44.60	34.38	20.64

Appendix

Table 11: Accuracy of ResNet-50 on Refined CIFAR-100-C under different severity levels

Type	S1	S2	S3	S4	S5
gaussian	65.58	54.04	44.05	36.71	24.64
shot	64.83	54.97	47.32	35.17	24.26
impulse	65.32	57.01	47.40	35.29	28.20
defocus	65.11	54.15	42.78	36.00	23.80
glass	64.06	54.62	42.62	30.01	23.17
motion	64.00	54.06	43.82	34.58	21.41

Appendix

Table 12: Accuracy of ResNet-101 on Refined CIFAR-100-C under different severity levels

Explore similar work

VITA: A Multi-Source Vicinal Transfer Augmentation Method for Out-of-Distribution Generalization
Apr 25, 2022Minghui Chen, Cheng Wen, Feng Zheng +2Data Augmentation Out-Of-Distribution
Department of Computer Science and Engineering, Southern University of Science and Technology · The University of Sydney · JD Explore Academy, JD.com Inc +1
On the Reliability of Generative Augmentation: A Wasserstein-Based Theoretical and Empirical Study
Sep 1, 2026Chathurika S Abeykoon, Mathias Nthiani Muia, Mallory GoldsteinAugmentation Data Augmentation
InstructMixup: Instruction-Guided Salient Patch Editing for Robust Data Augmentation
Jul 21, 2026Khawar Islam, Arif Mahmood, Xin Jin +1Data Augmentation Inpainting
The University of Melbourne, Melbourne, Australia · The Information Technology University, Lahore, Pakistan · Westlake University, China

Apr 25, 2022, cs.CVJ/K move · Enter open · S save

VITA: A Multi-Source Vicinal Transfer Augmentation Method for Out-of-Distribution Generalization

Minghui Chen, Cheng Wen, Feng Zheng, Fengxiang He+1

Department of Computer Science and Engineering, Southern University of Science and Technology · The University of Sydney · JD Explore Academy, JD.com Inc+1

Data Augmentation Out-Of-Distribution Adversarial Corruption

Open paper AskPDF arXiv

Invariance to diverse types of image corruption, such as noise, blurring, or colour shifts, is essential to establish robust models in computer vision. Data augmentation has been the major approach in improving the robustness against common corruptions. However, the samples produced by popular augmentation strategies deviate significantly from the underlying data manifold. As a result, performance is skewed toward certain types of corruption. To address this issue, we propose a multi-source vicinal transfer augmentation (VITA) method for generating diverse on-manifold samples. The proposed VITA consists of two complementary parts: tangent transfer and integration of multi-source vicinal samples. The tangent transfer creates initial augmented samples for improving corruption robustness. The integration employs a generative model to characterize the underlying manifold built by vicinal samples, facilitating the generation of on-manifold samples. Our proposed VITA significantly outperforms the current state-of-the-art augmentation methods, demonstrated in extensive experiments on corruption benchmarks. Code: https://github.com/MinghuiChen43/VITA.

Abstract

Explore similar work

VITA: A Multi-Source Vicinal Transfer Augmentation Method for Out-of-Distribution Generalization

On the Reliability of Generative Augmentation: A Wasserstein-Based Theoretical and Empirical Study

InstructMixup: Instruction-Guided Salient Patch Editing for Robust Data Augmentation