Reweighted SAM: Delta-SAM Optimization

• The paper introduces delta-SAM, which dynamically reweights per-instance gradients to create a shared adversarial perturbation that better approximates individual perturbations.
• It demonstrates improved performance on NLP tasks such as GLUE and STS while reducing adversarial risk and incurring only modest extra computational overhead.
• The method is theoretically derived using a second-order Hessian approximation and an unbiased estimator based on shared Gaussian perturbations, making complex per-instance computations tractable.

Searching arXiv for the specified paper and closely related SAM literature to ground the article. Reweighted SAM denotes a family of sharpness-aware training procedures in which a single shared adversarial weight perturbation is constructed for a mini-batch by reweighting per-instance gradients, so that the shared perturbation better approximates the stronger per-instance adversaries without incurring their prohibitive cost. In the formulation introduced as delta-SAM, or “Sharpness-Aware Minimization with Dynamic Reweighting,” the perturbation direction is obtained from a dynamically reweighted batch gradient whose weights are theoretically motivated by a second-order analysis of per-instance adversarial risk (Zhou et al., 2021).

1. Position within sharpness-aware optimization

Sharpness-Aware Minimization (SAM) adversarially perturbs model weights and minimizes the loss at the worst-case point in a small $L_2$ neighborhood. For a batch of size $N$ with per-instance loss $l_i$, the per-batch SAM objective is

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$

In practice, the inner maximizer is approximated by a single first-order step using the batch-average loss $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$, yielding

$$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$

The outer minimization then updates $\bm{w}$ to reduce the perturbed empirical risk (Zhou et al., 2021).

Within this framework, standard per-batch SAM aggregates gradients over the mini-batch by averaging,

$$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$

and one SAM step roughly doubles compute, requiring two forward and two backward passes per batch. Reweighted SAM retains the shared-perturbation structure of SAM but replaces the uniform averaging implicit in the batch gradient with instance-dependent weights. This suggests a reinterpretation of SAM in which the central design question is not only how large the perturbation radius $\rho$ should be, but also how the shared perturbation direction should reflect heterogeneity across instances in the same mini-batch.

2. Per-batch and per-instance adversaries

The motivation for reweighting arises from the gap between a single perturbation shared by all examples in a batch and the stronger objective in which each example receives its own adversarial perturbation. The per-instance objective is

$$\min_{\bm{w}} \frac{1}{N}\sum_{i=1}^N \max_{\bm{\epsilon}_i: \left\lVert\bm{\epsilon}_i\right\rVert_2\le \rho} \, l_i(\bm{w}+\bm{\epsilon}_i).$$

In first-order practice, one would take

$N$0

for each $N$1 (Zhou et al., 2021).

The associated sharpness increases are defined as

$N$2

Since the inner maximizers are independent in $N$3, $N$4. The reweighted construction is therefore motivated by the observation that sharing one perturbation across the batch yields a weaker adversary than summing per-instance worst cases (Zhou et al., 2021).

The difficulty is computational. Per-instance perturbations are computationally expensive because the inner problem must be solved separately for each instance, requiring $N$5 backward passes and effectively distinct model copies, approximately $N$6 forward and $N$7 backward passes per batch. The data further reports that this is often unaffordable in practice. Reweighted SAM is designed precisely for this efficiency bottleneck: it attempts to retain a single batch-level perturbation while moving that perturbation toward the behavior of the per-instance adversary.

3. Theoretical derivation of dynamic reweighting

The delta-SAM derivation uses a second-order expansion

$N$8

together with the structured Hessian assumption

$N$9

Under the single-step inner maximization, the SAM-style perturbations are

$l_i$0

Under these assumptions, the gradient of the per-instance adversarial risk, for the outer minimization and without differentiating through $l_i$1, is

$l_i$2

This expression identifies the target direction that a shared perturbation should approximate if it is to mimic the gradient induced by per-instance adversarial training (Zhou et al., 2021).

For a reweighted shared perturbation $l_i$3, let

$l_i$4

An ideal choice is

$l_i$5

which guarantees

$l_i$6

because $l_i$7 is positive definite. Since $l_i$8 is not directly available, the paper motivates approximating it by a reweighted batch gradient proportional to $l_i$9 (Zhou et al., 2021).

4. Reweighted perturbation and the delta-SAM algorithm

The dynamic weights are defined as

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$0

and the reweighted batch loss is

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$1

with $$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$2 treated as constants for the batch. Then

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$3

which leads to the normalized reweighted batch gradient

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$4

and the perturbation

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$5

Writing $$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$6 and $$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$7, this becomes

$$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$8

This is the core “Reweighted SAM” perturbation in delta-SAM (Zhou et al., 2021).

The algorithmic step for a mini-batch $$\min_{\bm{w}} \max_{\bm{\epsilon}: \left\lVert\bm{\epsilon}\right\rVert_2\le \rho}\frac{1}{N}\sum_{i=1}^N l_i(\bm{w}+\bm{\epsilon}).$$9 proceeds as follows. First, one shared $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$0 is sampled and, in practice, rescaled to $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$1. In no_grad mode, the method computes per-instance losses $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$2, $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$3, and $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$4, and then computes dynamic weights $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$5. Second, it forms $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$6 and backpropagates to obtain $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$7. Third, it computes the adversarial perturbation $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$8. Fourth, with perturbed weights $l(\bm{w}) = \frac{1}{N}\sum_{i=1}^N l_i(\bm{w})$9, it computes the standard unweighted empirical risk $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$0, backpropagates, and updates $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$1 using the optimizer, such as SGD or Adam (Zhou et al., 2021).

The implementation details are integral to the method’s intended behavior. The three forward-only passes used to estimate $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$2 are performed in no_grad mode; a single shared random perturbation is used per batch to keep overhead small; the perturbation vector is normalized; and the training update is performed on the unweighted loss to avoid biasing training targets. This suggests that the reweighting is intended only to shape the adversarial direction, not to redefine the empirical objective.

5. Estimation of dynamic weights

The practical obstacle is efficient estimation of

$$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$3

Delta-SAM proposes a simple, unbiased estimator based on one shared random Gaussian weight perturbation $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$4 per batch and three forward passes, without backpropagation. Under the same Hessian assumption, the paper states

$$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$5

and

$$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$6

Using a single shared $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$7 per batch, the dynamic weight is defined as

$$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$8

where $$\bm{\epsilon}^*\approx \arg\!\max_{\left\lVert\bm{\epsilon}\right\rVert_2\le \rho} \, l(\bm{w}) + \bm{\epsilon}^\intercal \nabla l(\bm{w}) = \rho \, \frac{\nabla l(\bm{w})}{\left\lVert\nabla l(\bm{w})\right\rVert_2}.$$9 avoids division by zero (Zhou et al., 2021).

The numerator is described as proxying curvature along $\bm{w}$0 via the Hessian, while the denominator proxies gradient magnitude along $\bm{w}$1, so that the ratio tracks $\bm{w}$2. This is the mechanism by which delta-SAM obtains per-instance weights without computing per-instance Hessians or per-instance adversarial perturbations. A plausible implication is that the estimator’s quality depends not only on the validity of the structured Hessian assumption but also on the variance introduced by using only one random perturbation sample per batch.

The practical settings reported for this estimator include $\bm{w}$3 with $\bm{w}$4 rescaled to $\bm{w}$5, and denominator floors $\bm{w}$6 searched in $\bm{w}$7, with the paper often using $\bm{w}$8. Perturbation radius $\bm{w}$9 is searched in $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$0 on GLUE; for STS, $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$1; and summarization uses $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$2 (Zhou et al., 2021).

6. Empirical behavior, computational profile, and limitations

The reported empirical results cover natural language understanding, self-supervised semantic textual similarity, and summarization. On GLUE development sets, BERTbase with SAM attains an average of 83.9, while delta-SAM attains 84.7; task-wise examples include MRPC $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$3, RTE $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$4, CoLA $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$5, and SST-2 $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$6. For RoBERTa_large, SAM attains 89.6 and delta-SAM 90.1, with reported gains including MRPC $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$7, CoLA $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$8, and RTE $$\nabla l(\bm{w})=\frac{1}{N}\sum_{i=1}^N \nabla l_i(\bm{w}),$$9 (Zhou et al., 2021).

In self-supervised STS, measured by Spearman’s $\rho$0, Mirror-BERTbase reports baseline 74.85, +SAM 75.50, +delta-SAM 75.72, +SAM w/ random noise 76.44, and +delta-SAM w/ random noise 76.71. Mirror-RoBERTabase reports baseline 74.98, +SAM 75.18, +delta-SAM 75.27, +SAM w/ random noise 75.90, and +delta-SAM w/ random noise 76.31. On CNN/DailyMail summarization with BART-large, the baseline reports RG-1 44.16, RG-2 21.28, RG-L 40.90, while +delta-SAM reports RG-1 44.70, RG-2 21.54, RG-L 41.81, which the data identifies as the best RG-1 and RG-L among compared methods (Zhou et al., 2021).

The paper also reports adversarial risk analysis using the per-instance metric $\rho$1 with $\rho$2. On MRPC and RTE, delta-SAM achieves lower $\rho$3 than SAM, with examples MRPC $\rho$4 and RTE $\rho$5, together with higher accuracy, approaching per-instance perturbation performance. The reported ablations state that reweighting matters, that smaller datasets show larger gains, and that moderate values of $\rho$6 matching baseline SAM often work best, whereas very large $\rho$7 can destabilize training. The paper further notes that using random input noise in self-supervised STS improves results, and that delta-SAM with random noise outperforms baselines that also use random noise (Zhou et al., 2021).

The computational profile is central to the method’s practical claim. Standard SAM uses two forward and two backward passes per batch, with time overhead approximately $\rho$8 vanilla training. Per-instance SAM requires approximately $\rho$9 forward and $$\min_{\bm{w}} \frac{1}{N}\sum_{i=1}^N \max_{\bm{\epsilon}_i: \left\lVert\bm{\epsilon}_i\right\rVert_2\le \rho} \, l_i(\bm{w}+\bm{\epsilon}_i).$$0 backward passes, plus memory for distinct perturbations; empirically it is reported as approximately $$\min_{\bm{w}} \frac{1}{N}\sum_{i=1}^N \max_{\bm{\epsilon}_i: \left\lVert\bm{\epsilon}_i\right\rVert_2\le \rho} \, l_i(\bm{w}+\bm{\epsilon}_i).$$1 delta-SAM for batch size 16. Delta-SAM uses the same number of backward passes as standard SAM, plus three forward-only passes per batch to estimate $$\min_{\bm{w}} \frac{1}{N}\sum_{i=1}^N \max_{\bm{\epsilon}_i: \left\lVert\bm{\epsilon}_i\right\rVert_2\le \rho} \, l_i(\bm{w}+\bm{\epsilon}_i).$$2, and empirically adds only about 18% extra computational cost over SAM on average. A concrete example is BERTbase on SST-2: 118 min for SAM versus 132 min for delta-SAM, approximately 12% slower (Zhou et al., 2021).

The limitations stated in the source are correspondingly specific. Delta-SAM remains slower than vanilla training because of the extra forward passes, even though it is much cheaper than per-instance SAM. Its theoretical derivation depends on the second-order expansion and the Hessian approximation $$\min_{\bm{w}} \frac{1}{N}\sum_{i=1}^N \max_{\bm{\epsilon}_i: \left\lVert\bm{\epsilon}_i\right\rVert_2\le \rho} \, l_i(\bm{w}+\bm{\epsilon}_i).$$3; deviations from these assumptions may limit the accuracy of the reweighting. The source also identifies potential extensions, including better curvature estimation, multiple random perturbation samples to reduce estimator variance, integration with ASAM or GSAM or trust-region methods, and per-layer or blockwise reweighting on broader task families such as sequence tagging and question answering. This suggests that reweighted SAM is best understood not as a replacement for sharpness-aware optimization in general, but as a specific mechanism for approximating per-instance adversarial weight perturbations under a constrained computational budget (Zhou et al., 2021).