DGS-MAML: Sharpness & Gradient Matching

• DGS-MAML is a meta-learning algorithm that combines SAM-style perturbations with gradient matching to bias inner adaptations and meta-updates toward flatter minima, enhancing few-shot performance and robustness to domain shift.
• It employs a bi-level optimization framework where adversarial perturbations are applied in both inner and outer loops, ensuring gradient alignment across training and validation splits.
• Supported by PAC-Bayes analysis and convergence guarantees, DGS-MAML shows consistent improvements over baseline methods on benchmarks like Mini-ImageNet and Omniglot.

DGS-MAML, short for Domain-Generalization Sharpness-Aware Minimization Model-Agnostic Meta-Learning, is a meta-learning algorithm that integrates sharpness-aware training and gradient matching within a bi-level optimization framework in order to improve few-shot adaptation and robustness to domain shift (Anjum et al., 13 Aug 2025). Its central design combines SAM-style adversarial perturbations with alignment between the gradients of the empirical and perturbed objectives, so that both inner adaptation and outer meta-updates are biased toward flat minima and more consistent adaptation dynamics across train/validation splits and across perturbations. The method is introduced for settings with limited training data, and the paper supports it with PAC-Bayes analysis, convergence guarantees, benchmark experiments, and publicly available source code (Anjum et al., 13 Aug 2025).

1. Problem setting and motivation

DGS-MAML is formulated for gradient-based meta-learning under few-shot adaptation, where tasks are sampled from a task distribution $T \sim p(\mathcal{T})$, and each task $m$ has a training split $\mathcal{D}_m^{train}$ and a validation split $\mathcal{D}_m^{val}$ (Anjum et al., 13 Aug 2025). In this setting, the meta-parameters $\theta \in \mathbb{R}^d$ define an initialization for a model $f_\theta$ that should adapt well to many tasks using few steps and few samples.

The motivating problem is that empirical risk on source tasks can admit sharp local minima that adapt quickly but generalize poorly to unseen target domains. The paper identifies two related issues in MAML-like bi-level optimization. First, stochastic gradient descent in MAML can converge to sharp basins, which exacerbates overfitting and harms cross-domain generalization. Second, conventional meta-objectives do not explicitly align how gradients behave across train versus validation splits and across neighborhoods in parameter space, so adaptation dynamics can differ under minor perturbations (Anjum et al., 13 Aug 2025).

DGS-MAML addresses these issues by combining two mechanisms. SAM seeks flat minima by minimizing the worst-case loss in a small neighborhood around the current parameters, thereby discouraging sharp solutions. Gradient matching, described through the SAGM perspective, encourages alignment between the gradients of the unperturbed loss and the perturbed loss. When these gradients align, descending the empirical loss simultaneously reduces the worst-case perturbed loss and the surrogate gap, which the paper describes as strongly indicative of flatter minima and better domain generalization (Anjum et al., 13 Aug 2025).

2. Formal construction within bi-level meta-learning

The baseline inner adaptation step follows MAML. For one step,

$$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$

where $\alpha > 0$ is the inner learning rate. The second-order MAML meta-objective is

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$

with gradient

$$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$

A first-order approximation replaces $m$0 by $m$1 (Anjum et al., 13 Aug 2025).

DGS-MAML augments this with SAM perturbations in both loops. For task $m$2, the inner SAM perturbation is

$m$3

and the outer SAM perturbation is

$m$4

These are the “both-loops” perturbations associated with SharpMAML (Anjum et al., 13 Aug 2025).

The gradient-matching component is introduced through a perturbed loss

$m$5

where $m$6 is a gradient displacement hyperparameter, and through the surrogate gap

$m$7

The SAGM-style objective minimizes

$m$8

Under a first-order expansion around $m$9, the sum becomes

$\mathcal{D}_m^{train}$0

so the inner-product term encourages gradient alignment. The paper states that DGS-MAML enforces this alignment implicitly by minimizing $\mathcal{D}_m^{train}$1 with the SAM-style $\mathcal{D}_m^{train}$2 and displacement $\mathcal{D}_m^{train}$3 (Anjum et al., 13 Aug 2025).

An explicit gradient-matching regularizer is also given conceptually:

$\mathcal{D}_m^{train}$4

or equivalently one may maximize $\mathcal{D}_m^{train}$5. In the implemented method, however, gradient matching is realized through the summed losses rather than through a separate regularizer (Anjum et al., 13 Aug 2025).

The task-wise inner objective is

$\mathcal{D}_m^{train}$6

with one inner update

$\mathcal{D}_m^{train}$7

where $\mathcal{D}_m^{train}$8 is the inner-step learning rate. After $\mathcal{D}_m^{train}$9 steps, one obtains $\mathcal{D}_m^{val}$0 (Anjum et al., 13 Aug 2025).

The outer objective uses the analogous validation loss

$\mathcal{D}_m^{val}$1

and the conceptual full bi-level objective is

$\mathcal{D}_m^{val}$2

where $\mathcal{D}_m^{val}$3 denotes the $\mathcal{D}_m^{val}$4 inner steps driven by $\mathcal{D}_m^{val}$5 and inner SAM. A weighted form

$\mathcal{D}_m^{val}$6

is possible, but the paper uses $\mathcal{D}_m^{val}$7 through the summed losses (Anjum et al., 13 Aug 2025).

3. Differentiation, optimization procedure, and computational profile

The exact meta-gradient differentiates through the unrolled inner loop:

$\mathcal{D}_m^{val}$8

where $\mathcal{D}_m^{val}$9 depends on the inner-loop perturbations and on $\theta \in \mathbb{R}^d$0 through $\theta \in \mathbb{R}^d$1. A first-order approximation treats $\theta \in \mathbb{R}^d$2 and $\theta \in \mathbb{R}^d$3 as stop-gradient perturbations and ignores Hessian terms in $\theta \in \mathbb{R}^d$4, yielding

$\theta \in \mathbb{R}^d$5

The paper identifies this as the practical choice to match the computational profile of SharpMAML (Anjum et al., 13 Aug 2025).

For practical differentiation of a generic $\theta \in \mathbb{R}^d$6 term, let

$\theta \in \mathbb{R}^d$7

and define

$\theta \in \mathbb{R}^d$8

With stop-gradient on $\theta \in \mathbb{R}^d$9 and $f_\theta$0, the gradient becomes

$f_\theta$1

which requires two backward passes per evaluation, one for the unperturbed loss and one for the perturbed loss (Anjum et al., 13 Aug 2025).

Per meta-iteration, the algorithm samples $f_\theta$2 tasks and their train/validation splits. In the inner loop, it computes

$f_\theta$3

forms

$f_\theta$4

constructs

$f_\theta$5

and performs one or more updates

$f_\theta$6

In the outer loop, it aggregates validation gradients at the perturbed inner points,

$f_\theta$7

sets

$f_\theta$8

builds

$f_\theta$9

and applies the meta-update

$$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$0

Relative to MAML, the computational overhead is limited. SAM requires two forward/backward passes per update, and DGS-MAML uses two passes because $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$1 sums unperturbed and perturbed losses, while reusing the unperturbed gradient. The paper reports comparable per-epoch runtimes: on Mini-ImageNet, 60–70s for DGS-MAML versus 60–65s for SharpMAML, and on Omniglot, 60–70s for both at batch size 100. The memory footprint is near SharpMAML in the first-order implementation, whereas second-order differentiation increases memory proportional to $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$2 inner steps (Anjum et al., 13 Aug 2025).

The hyperparameters are $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$3 for both inner and outer updates, $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$4 and $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$5 as inner and outer SAM radii, $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$6 for displacement along the empirical gradient, $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$7 inner adaptation steps, meta-batch and task-split batch sizes, the episode distribution $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$8 for $$\theta_m' = \theta - \alpha \nabla_\theta \mathcal{L}(\theta; \mathcal{D}_m^{train}),$$9-way $\alpha > 0$0-shot configurations, initialization $\alpha > 0$1, and the optimizer. The experiments use Adam and learning rates $\alpha > 0$2 (Anjum et al., 13 Aug 2025).

4. Theoretical analysis

The PAC-Bayes analysis is stated for a data-independent prior $\alpha > 0$3 over initializations and a posterior $\alpha > 0$4 over $\alpha > 0$5 after meta-training. Let $\alpha > 0$6 be the number of tasks in meta-training, $\alpha > 0$7 the confidence, and assume a bounded loss $\alpha > 0$8 together with $\alpha > 0$9-uniform stability. Then, with probability at least $$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$0,

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$1

Uniform stability is defined as follows: for any datasets $$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$2 and $$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$3 differing by one element,

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$4

This is presented as Theorem 3 and Definition 1 in the paper (Anjum et al., 13 Aug 2025).

To capture SAM and gradient matching, the paper takes

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$5

where $$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$6 is a local minimizer learned by DGS-MAML. The resulting Kullback–Leibler divergence is

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$7

Ensuring $$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$8 requires

$$\min_\theta \mathbb{E}_{T \sim p(\mathcal{T})}\big[\mathcal{L}(\theta'; \mathcal{D}_T^{val})\big],$$9

The paper states that, whereas the SharpMAML bound depends on $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$0, in DGS-MAML the parameter $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$1 further regularizes the posterior dispersion, reflecting robustness to gradient displacement, and tightens the generalization bound compared to SharpMAML empirically and in this PAC-Bayes construct (Anjum et al., 13 Aug 2025).

The convergence analysis assumes: (A1) $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$2-smoothness, meaning $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$3 is $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$4-Lipschitz; (A2) bounded stochastic gradients, with $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$5 and hence $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$6; and (A3) unbiasedness and bounded variance of stochastic gradients in the bi-level setting, with variances $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$7 and $$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$8 for the unperturbed and perturbed gradients (Anjum et al., 13 Aug 2025).

Defining

$$\nabla_\theta \mathcal{L}(\theta'; \mathcal{D}_T^{val}) = \big(I - \alpha \nabla_\theta^2 \mathcal{L}(\theta; \mathcal{D}_T^{train})\big)^\top \nabla_{\theta'} \mathcal{L}(\theta'; \mathcal{D}_T^{val}).$$9

the single-level SAGM convergence result gives, for fixed step size $m$00,

$m$01

Corollary 1 states that choosing $m$02, $m$03, and $m$04 yields an average stationarity rate $m$05, improving over the $m$06 rate reported for SharpMAML (Anjum et al., 13 Aug 2025).

For bi-level DGS-MAML, under (A1)–(A3), fixed $m$07, and the same $m$08,

$m$09

$m$10

Corollary 2 states that with $m$11, $m$12, and $m$13, the average stationarity rate is again $m$14, and the inner loop does not degrade the asymptotic rate relative to single-level SAGM (Anjum et al., 13 Aug 2025).

5. Experimental evaluation

The empirical study uses a Conv-4 backbone consisting of 4 blocks of $m$15 convolution with 64 filters, batch normalization, ReLU, and $m$16 max-pooling (Anjum et al., 13 Aug 2025). The datasets are Mini-ImageNet with a 64/12/24 train/val/test split, Omniglot following the Matching Networks splits, and MultiDigitMNIST through DoubleMNIST and TripleMNIST, with DoubleMNIST split 640/200/160 and TripleMNIST split 64/16/20. Evaluation is performed in $m$17-way $m$18-shot episodic classification on held-out query sets. The paper describes domain generalization here in terms of cross-dataset diversity and robustness to perturbations through SAM and gradient matching; explicit DG benchmarks such as PACS and VLCS are not included (Anjum et al., 13 Aug 2025).

The baselines include gradient-based meta-learning methods MAML and SharpMAML, other meta-learners ANIL, REPTILE, and CAVIA, metric-based few-shot methods Matching Networks and Prototypical Networks, and SharpMAML variants with $m$19 (Anjum et al., 13 Aug 2025).

The strongest gains are reported on Mini-ImageNet. For 5-way 1-shot, the accuracies are Matching Nets 0.4033, ProtoNet 0.4390, MAML 0.4463, SharpMAML $m$20 0.4509, and DGS-MAML $m$21 0.4665, corresponding to gains of +2.02 percentage points over MAML and +1.56 over SharpMAML. For 5-way 5-shot, MAML reaches 0.5729, SharpMAML $m$22 reaches 0.5759, and DGS-MAML $m$23 reaches 0.6360, a gain of +6.31 over MAML and +6.01 over SharpMAML. For 10-way 1-shot, MAML reaches 0.2160, SharpMAML reaches 0.2703–0.2740, and DGS-MAML $m$24 reaches 0.2882, with gains of +7.22 over MAML and +1.42 over the best SharpMAML. For 10-way 5-shot, MAML reaches 0.3621, SharpMAML reaches 0.3846–0.3899, and DGS-MAML reaches 0.4213–0.4264, for gains of +6.43 over MAML and up to +3.65 over SharpMAML (Anjum et al., 13 Aug 2025).

Setting	Baseline range	Best DGS-MAML
Mini-ImageNet 5-way 1-shot	MAML 0.4463; SharpMAML 0.4509	0.4665
Mini-ImageNet 5-way 5-shot	MAML 0.5729; SharpMAML 0.5759	0.6360
Mini-ImageNet 10-way 1-shot	MAML 0.2160; SharpMAML 0.2703–0.2740	0.2882
Mini-ImageNet 10-way 5-shot	MAML 0.3621; SharpMAML 0.3846–0.3899	0.4213–0.4264

On Omniglot in the 20-way regime, MAML achieves 0.8884 for 1-shot and 0.9571 for 5-shot, while DGS-MAML reaches up to 0.8982 for 1-shot with $m$25 and up to 0.9604 for 5-shot with $m$26. On TripleMNIST and DoubleMNIST, the reported improvements are modest but consistent; one example is TripleMNIST 1-shot, where MAML attains 0.9017 and DGS-MAML $m$27 attains 0.9171 (Anjum et al., 13 Aug 2025).

The ablation results identify $m$28 as crucial. On Mini-ImageNet 5-way 5-shot, $m$29 yields 0.6360 accuracy, while other values from 0.002 to 0.008 are inferior, with results in the range 0.613–0.625. For Mini-ImageNet 5-way 1-shot, $m$30 performs best in the tuning table with 0.4651, while the reported best run reaches 0.4665. In the experiments, $m$31 and $m$32 are set equal, and the best general performance is reported at $m$33 (Anjum et al., 13 Aug 2025).

The paper states that improvements are strongest where baseline accuracy is lower, especially on Mini-ImageNet, indicating better robustness to harder episodes. At the same time, explicit OOD or domain-shift metrics such as calibration are not provided. Training uses Adam, Torchmeta for episode generation, and the learning rates $m$34. Seeds and statistical tests are not reported, and mean $m$35 standard deviation values are not provided (Anjum et al., 13 Aug 2025).

6. Practical use, relation to prior work, and limitations

The paper positions DGS-MAML for few-shot adaptation under potential domain shift or noisy episodes, especially when flat minima and robustness to perturbations matter, such as in diverse image domains (Anjum et al., 13 Aug 2025). It recommends starting from $m$36, tuning $m$37 carefully per dataset, and using 1–5 inner steps $m$38 with first-order differentiation to control memory. The suggested test-time protocol follows standard MAML adaptation: adapt $m$39 to each target task using a few inner steps on the target support set, with the DGS-MAML SAM-plus-gradient-matching inner loop carried over at test time (Anjum et al., 13 Aug 2025).

Several practical pitfalls are explicitly noted. If $m$40 is too large, updates can become misaligned; if $m$41 is too small, the surrogate gap may not be reduced, so a sweep is advised. If $m$42 is too large, the perturbation may be excessive; if it is too small, the SAM effect is reduced. The paper therefore recommends balancing $m$43 and $m$44 jointly. When second-order meta-gradients are used, memory should be monitored; otherwise, first-order differentiation is preferred (Anjum et al., 13 Aug 2025).

In relation to prior work, DGS-MAML is described as extending MAML, which performs bi-level gradient-based meta-learning without explicit flatness control, and SharpMAML, which applies SAM in both loops. It also imports the gradient-alignment perspective of SAGM and GSAM into meta-learning by using the sum of empirical and perturbed losses to implicitly match gradients in both loops. Relative to domain generalization frameworks such as IRM, VRM, and DANN, DGS-MAML does not impose invariances across source domains directly; instead, it enforces local invariance in parameter space via SAM and gradient alignment, which the paper describes as complementary to domain-invariance penalties (Anjum et al., 13 Aug 2025).

The reported limitations are equally explicit. Benchmarks such as PACS, VLCS, and Office-Home are not included; statistical significance is not reported; and gains diminish on easier datasets with high baseline accuracy. This suggests that the strongest empirical case is currently on harder few-shot image benchmarks rather than on standard domain generalization suites. A plausible implication is that further evaluation on explicit DG benchmarks would be necessary to establish the method’s behavior outside the episodic regimes considered in the paper (Anjum et al., 13 Aug 2025).

The implementation is released at https://github.com/AIResearchTopics/DGSharpMAML, with dependencies on PyTorch, Torchmeta, and a standard Python machine learning stack. The repository provides training and evaluation code for Mini-ImageNet, Omniglot, and MultiDigitMNIST, including the Conv-4 backbone and episodic sampling; the license is not specified in the paper (Anjum et al., 13 Aug 2025).