Bi-Level Optimization for Single Domain Generalization

Marzi Heidari; Hanping Zhang; Hao Yan; Yuhong Guo

cs.LGApr 7, 2026

Bi-Level Optimization for Single Domain Generalization

Authors: Marzi Heidari, Hanping Zhang, Hao Yan, Yuhong Guo

Organizations: School of Computer Science, Carleton University, Ottawa, Canada · Canada CIFAR AI Chair, Amii, Canada

Abstract

Generalizing from a single labeled source domain to unseen target domains, without access to any target data during training, remains a fundamental challenge in robust machine learning. We address this underexplored setting, known as Single Domain Generalization (SDG), by proposing BiSDG, a bi-level optimization framework that explicitly decouples task learning from domain modeling. BiSDG simulates distribution shifts through surrogate domains constructed via label-preserving transformations of the source data. To capture domain-specific context, we propose a domain prompt encoder that generates lightweight modulation signals to produce augmenting features via feature-wise linear modulation. The learning process is formulated as a bi-level optimization problem: the inner objective optimizes task performance under fixed prompts, while the outer objective maximizes generalization across the surrogate domains by updating the domain prompt encoder. We further develop a practical gradient approximation scheme that enables efficient bi-level training without second-order derivatives. Extensive experiments on various SGD benchmarks demonstrate that BiSDG consistently outperforms prior methods, setting new state-of-the-art performance in the SDG setting.

Figures & tables

Figure 1: Overview of the BiSDG framework for Single Domain Generalization (SDG). Starting from a labeled source dataset, BiSDG constructs multiple surrogate domains via semantically coherent, label-preserving transformation pipelines to simulate unseen distribution shifts. A domain prompt encoder is introduced to capture domain-specific knowledge, which is deployed to perform FiLM-based feature modulation, supporting generalizable prediction. BiSDG formulates SGD as a bi-level optimization problem: the inner objective updates the model parameters $\Theta=(\theta,\phi)$ using the source domain, combining supervised cross-entropy loss $\mathcal{L}_{\text{in}}^{\text{cl}}$ and adversarial consistency loss $\mathcal{L}_{\text{in}}^{\text{adv}}$ . The outer objective updates the prompt encoder parameters $\omega$ to maximize generalization across surrogate domains.

Target	MixUp	CutOut	ADA	ME-ADA	AugMix	RandAug	ACVC	L2D	AdvST	BiSDG (Ours)
Art	52.8	59.8	58.0	60.7	63.9	67.8	67.8	67.6	${69.2_{(1.4)}}$	$\mathbf{70.5}_{(1.4)}$
Cartoon	17.0	21.6	25.3	28.5	27.7	28.9	30.3	42.6	$55.3_{(2.0)}$	$\mathbf{55.9}_{(1.7)}$
Sketch	23.2	28.8	30.1	29.6	30.9	37.0	46.4	47.1	$67.7_{(1.5)}$	$\mathbf{69.5}_{(1.9)}$
Avg.	31.0	36.7	37.8	39.6	40.8	44.6	48.2	52.5	64.1	$\mathbf{65.3}$

Table 1: Classification accuracy and standard deviation results (%) on the PACS dataset. Best results are in bold font.

Method	SVHN	MNIST-M	SYN	USPS	Avg.
ERM [ 15 ]	27.8	52.7	39.7	76.9	49.3
CCSA [ 24 ]	25.9	49.3	37.3	83.7	49.1
JiGen [ 2 ]	33.8	57.8	43.8	77.2	53.1
ADA [ 33 ]	35.5	60.4	45.3	77.3	54.6
ME-ADA [ 40 ]	42.6	63.3	50.4	81.0	59.3
M-ADA [ 30 ]	42.6	67.9	49.0	78.5	59.5

Table 2: Classification accuracy and standard deviation results (%) on the four target domains (SVHN, MNIST-M, SYN, and USPS) of the Digits benchmark, with MNIST as the source domain. Best results are in bold font.

Target	MixUp	CutOut	CutMix	ADA	ME-ADA	RandAug	AugMix	ACVC	AdvST	BiSDG (Ours)
Painting	38.6	38.3	38.3	38.2	39.0	41.3	40.8	41.3	$42.3_{(0.1)}$	$\mathbf{43.1}_{(0.3)}$
Infograph	13.9	13.7	13.5	13.8	14.0	13.6	13.9	12.9	${14.8}_{(0.1)}$	$\mathbf{15.9}_{(0.3)}$
Clipart	38.0	38.4	38.7	40.2	41.0	41.1	41.7	${42.8}$	$41.5_{(0.4)}$	$\mathbf{43.5}_{(0.4)}$
Sketch	26.0	26.2	26.9	24.8	25.3	30.4	29.8	30.9	$30.8_{(0.3)}$	$\mathbf{31.9}_{(0.4)}$
Quickdraw	3.7	3.7	3.6	4.3	4.3	5.3	6.3	6.6	$5.9_{(0.2)}$	$\mathbf{7.1}_{(0.2)}$
Avg.	24.0	24.1	24.2	24.3	24.7	26.3	26.5	26.9	$27.1_{(0.2)}$	$\mathbf{28.3}$

Table 3: Classification accuracy and standard deviation results (%) on the DomainNet dataset.

Method	SVHN	MNIST-M	SYN	USPS	Avg.
BiSDG (Ours)	$\mathbf{70.1}_{(0.6)}$	$\mathbf{85.2}_{(0.3)}$	$\mathbf{79.8}_{(0.5)}$	$\mathbf{96.0}_{(0.2)}$	$\mathbf{82.7}$
$-\text{w/o }\mathcal{L}_{\text{in}}^{\text{adv}}$	${68.1}_{(1.4)}$	${80.1}_{(0.6)}$	${75.3}_{(0.3)}$	${93.3}_{(0.5)}$	${79.2}$
$-\text{w/o standardization }$	${69.8}_{(0.9)}$	${83.5}_{(0.7)}$	${77.0}_{(0.4)}$	${94.7}_{(0.2)}$	${81.2}$

Table 4: Ablation study results (%) on the four target domains of the Digits benchmark, with MNIST as the source domain.

Panel of Figure 2 — Figure 2: Sensitivity analysis for four hyper-parameters $K$ , $\lambda$ , $\epsilon_{\Theta}$ and $m$ on test domain SYN of the Digits benchmark.

Explore similar work

One-Step Generalization Ratio Guided Optimization for Domain Generalization
Jun 15, 2026Sumin Cho, Dongwon Kim, Kwangsu KimDomain Generalization Adaptive Gradient Methods
Department of Computer Science and Engineering, University of Sungkyunkwan, Suwon, Korea.
Adversarial Domain Prompt Tuning and Generation for Single Domain Generalization
Jun 19, 2026Zhipeng Xu, De Cheng, Xinyang Jiang +3Synthetic Data Augmentation Adversarial Prompt Generation
Xidian University · Microsoft Research Asia · Chongqing University of Posts and Telecommunications
Generative Classifier for Domain Generalization
Apr 3, 2025Shaocong Long, Qianyu Zhou, Dexuan Zhang +1Domain Generalization Cross-Domain Generalization
Department of Computer Science and Engineering, Shanghai Jiao Tong University, Shanghai, China · College of Computer Science and Technology and the Key Laboratory of Symbolic Computation and Knowledge Engineering of the Ministry of Education, Jilin University, Changchun, China

Jun 15, 2026J/K move · Enter open · S save

One-Step Generalization Ratio Guided Optimization for Domain Generalization

Sumin Cho, Dongwon Kim, Kwangsu Kim

Department of Computer Science and Engineering, University of Sungkyunkwan, Suwon, Korea.

Domain Generalization Adaptive Gradient Methods Invariant Representation Learning Gradient Alignment

Open paper AskPDF arXiv

Domain Generalization (DG) aims to train models that generalize to unseen target domains but often overfit to domain-specific features, known as undesired correlations. Gradient-based DG methods typically guide gradients in a dominant direction but often inadvertently reinforce spurious correlations. Recent work has employed dropout to regularize overconfident parameters, but has not explicitly adjusted gradient alignment or ensured balanced parameter updates. We propose GENIE (Generalization-ENhancing Iterative Equalizer), a novel optimizer that leverages the One-Step Generalization Ratio (OSGR) to quantify each parameter's contribution to loss reduction and assess gradient alignment. By dynamically equalizing OSGR via a preconditioning factor, GENIE prevents a small subset of parameters from dominating optimization, thereby promoting domain-invariant feature learning. Theoretically, GENIE balances convergence contribution and gradient alignment among parameters, achieving higher OSGR while retaining SGD's convergence rate. Empirically, it outperforms existing optimizers and enhances performance when integrated with various DG and single-DG methods.

Abstract

Explore similar work

One-Step Generalization Ratio Guided Optimization for Domain Generalization

Adversarial Domain Prompt Tuning and Generation for Single Domain Generalization

Generative Classifier for Domain Generalization