diff --git a/TASK.md b/TASK.md index 26ecb32..43df61b 100644 --- a/TASK.md +++ b/TASK.md @@ -2,9 +2,15 @@ TASK write a new file, from the old part. ## Status: DONE +### Extension: flatness- and curvature-aware adapters +- [x] Freeze primary texts for BAR, FMLoRA/EFMLoRA, Bi-LoRA, LoRA-MGPO, and DISAM under `docs/` +- [x] Expand the catalog from 34 to 38 methods with pseudocode and calibrated evidence summaries +- [x] Separate flatness seeking, explicit low-loss curvature, and function/context curvature in Theme 8 +- [x] UAT: [catalog](adapters_as_hypotheses.md) has 38 sequential entries and 36 active pseudocode blocks; [evidence map](adapters_vargdown.argdown) has seven new verbatim quote blocks within their linked line ranges; fresh-eyes review confirmed the corrected EFMLoRA math and curvature taxonomy + ### Task 1: adapters_as_hypotheses.md - [x] Preamble with pragmatic interpretability framing -- [x] 34 catalog entries; 32 active methods have pseudocode, while deprecated Bone and out-of-scope Trainable Tokens are marked as boundary entries +- [x] 38 catalog entries; 36 active methods have pseudocode, while deprecated Bone and out-of-scope Trainable Tokens are marked as boundary entries - [x] All papers saved to docs/ (full size, no truncation) - [x] Sub-agent review completed, fixes applied diff --git a/adapters_as_hypotheses.md b/adapters_as_hypotheses.md index 2ee2304..83a1127 100644 --- a/adapters_as_hypotheses.md +++ b/adapters_as_hypotheses.md @@ -10,7 +10,7 @@ We fine tune transformers effeciently with low rank adapters - adding a new tran This is an underused source of *suggestive* evidence. Most interpretability *observes* (probing, SAEs); adapters *intervene*. If a structural constraint helps, the structure it encodes is load-bearing. The evidence is confounded by optimization dynamics (a method can win because its parameterization is optimizer-friendly, not because its structural hypothesis is correct), but the patterns are consistent enough to be worth taking seriously. -I went through 34 PEFT methods in [HuggingFace PEFT](https://github.com/huggingface/peft) and the broader literature. For each one I extracted pseudocode for the intervention, stated the hypothesis it encodes, and weighed the evidence. Three claims emerged: +I went through 38 PEFT methods in [HuggingFace PEFT](https://github.com/huggingface/peft) and the broader literature. For each one I extracted pseudocode for the intervention, stated the hypothesis it encodes, and weighed the evidence. Three claims emerged: 1. **SVD basis outperforms random-initialized and standard bases.** Methods that initialize or constrain updates in the model's own singular-vector basis (PiSSA, SVFT, SSVD, CLOVER, PSOFT) consistently outperform random-basis alternatives at comparable budgets. SVD is linear and transformers are not, and almost no papers compare against other structured bases (ICA, Fisher eigenvectors, gradient covariance), so "SVD > random" is solid but "SVD is the right basis" is stronger than the data warrants. 2. **Direction and strength decouple.** Methods that separate *which way* to move in weight space from *how far* (DoRA, DeLoRA, ROAD, AntiPaSTO) show better stability and sometimes better OOD transfer. An honest alternative: this could be an optimization benefit (giving Adam better-conditioned knobs) rather than a structural insight. @@ -940,6 +940,109 @@ def flat_lora_loss(x, y, W, A, B, σ): --- +## 35. BAR -- Balancedness-Aware Regularization + +**Paper:** [Li, Zhang, He 2024](https://arxiv.org/abs/2410.14802) (NeurIPS 2024) +**Code:** [github.com/BingcongLi/BAR](https://github.com/BingcongLi/BAR) +**Saved:** [docs/bar_balancedness_aware_regularization.md](docs/bar_balancedness_aware_regularization.md) + +**Hypothesis:** Much of SAM's benefit for factorized adapters comes from balancing the norms of LoRA's two factors, rather than directly minimizing Hessian curvature. BAR makes that implicit effect explicit. Its nBAR variant expands one factor and contracts the other according to their gradient norms, then applies the ordinary optimizer update. + +```py +# Claude: nBAR training step +def nbar_step(loss, W, A, B, α, η, optimizer): + g_A, g_B = ∇(loss(W + B @ A), (A, B)) + s = +1 if norm(g_A) >= norm(g_B) else -1 + A ← (1 + s * α * η) * A + B ← (1 - s * α * η) * B + A, B ← optimizer.step((A, B), (g_A, g_B)) + return A, B +``` + +**Evidence:** The authors report few-shot OPT-1.3B averages of 78.5 for oBAR and 79.2 for nBAR, versus 77.6 for LoRA and 78.4 for LoRA-SAM. BAR runs at about 1.03--1.05x LoRA in those experiments, while LoRA-SAM takes 3.28--4.43x. RoBERTa and GPT-2 experiments also favor BAR in most reported comparisons. These are few-shot and ordinary held-out evaluations, not controlled OOD tests; balancedness is a proposed explanation for SAM's benefit rather than a curvature estimate. + +**Grade:** PE+BL+DE=3.5 (beats LoRA in few-shot tests and retains near-LoRA training cost) + +--- + +## 36. FMLoRA / EFMLoRA -- Flat Minima LoRA + +**Paper:** [Deng et al. 2025](https://arxiv.org/abs/2508.00522) (AAAI 2026) +**Saved:** [docs/fmlora_flat_minima_lora.md](docs/fmlora_flat_minima_lora.md) + +**Hypothesis:** A full-weight SAM perturbation can be reconstructed from LoRA-factor gradients and represented by perturbing only one LoRA factor. FMLoRA performs the two-step SAM update; EFMLoRA reuses an exponential moving average of previous perturbations to recover one-forward, one-backward training. + +```py +# Claude: EFMLoRA training step +def efmlora_step(loss, W, A, B, Ê_B, ρ, β, scale, optimizer): + g_A, g_B = ∇(loss(W + scale * (B + Ê_B) @ A), (A, B)) + Ĝ_W = 0.5 / scale * (g_B @ pinv(A.T) + pinv((B + Ê_B).T) @ g_A) + E_W = ρ * Ĝ_W / norm(Ĝ_W) # Claude: full-weight SAM direction + E_B = E_W @ pinv(A) / scale # Claude: transfer into one LoRA factor + A, B ← optimizer.step((A, B), (g_A, g_B)) + Ê_B ← (1 - β) * Ê_B + β * E_B + return A, B, Ê_B +``` + +**Evidence:** The authors report RoBERTa few-shot averages of 83.1 for FMLoRA and 82.3 for EFMLoRA, versus 80.0 for LoRA and 81.3 for LoRA-SAM. On full GLUE fine-tuning, EFMLoRA averages 89.4 versus 88.4 for LoRA and 88.9 for full fine-tuning. The paper also reports gains on GPT-2, CLIP few-shot classification, and Qwen-VL-Chat. Its "distribution shift" language mostly refers to few-shot transfer or ordinary downstream test sets; it does not use a leave-one-domain-out OOD protocol. + +**Grade:** PE+BL+BF+DE=5 (beats LoRA, slightly beats full FT on reported averages, and is strongest in few-shot settings) + +--- + +## 37. Bi-LoRA -- Bi-directional Low-Rank Adaptation + +**Paper:** [Liu et al. 2025](https://arxiv.org/abs/2508.19564) (ICLR 2026) +**Code:** [github.com/CrazyElements/Bi-LoRA](https://github.com/CrazyElements/Bi-LoRA) +**Saved:** [docs/bi_lora_sharpness_aware.md](docs/bi_lora_sharpness_aware.md) + +**Hypothesis:** Task adaptation and sharpness exploration should use separate low-rank modules. A primary LoRA branch descends the task loss while an auxiliary branch ascends it inside a norm ball. Because the adversarial branch evolves independently, its perturbations need not collapse into the primary LoRA subspace. The auxiliary branch is discarded after training. + +```py +# Claude: Bi-LoRA training step +def bilora_step(loss, W, A_1, B_1, A_2, B_2, η_1, η_2, ρ): + W̃ = W + B_1 @ A_1 + B_2 @ A_2 + G_W = ∇(loss(W̃), W̃) + B_1, A_1 ← B_1 - η_1 * G_W @ A_1.T, A_1 - η_1 * B_1.T @ G_W + B_2, A_2 ← B_2 + η_2 * G_W @ A_2.T, A_2 + η_2 * B_2.T @ G_W + A_2, B_2 ← project_product_norm(A_2, B_2, ρ) + return A_1, B_1, A_2, B_2 + +def bilora_merge(W, A_1, B_1): + return W + B_1 @ A_1 # Claude: discard adversarial branch +``` + +**Evidence:** The authors fine-tune on MetaMathQA, Code-Feedback, WizardLM, and Alpaca, then evaluate on separate benchmarks including GSM8K, HumanEval, MT-Bench, MMLU, DROP, and BBH. Against LoRA, reported Llama-2 gains are +2.11 on GSM8K, +2.45 on HumanEval, and +0.34 on MT-Bench. Bi-LoRA mostly improves on Flat-LoRA in the same table, while costing one gradient step per iteration. These cross-dataset evaluations are more informative than same-dataset validation, but they are not a controlled domain-generalization study. + +**Grade:** PE+BL+BF=3.5 (beats LoRA broadly and beats full FT on some reported tasks) + +--- + +## 38. LoRA-MGPO -- Momentum-Guided Perturbation Optimization + +**Paper:** [Chang et al. 2025](https://aclanthology.org/2025.findings-emnlp.34/) (Findings EMNLP 2025) +**Code:** [github.com/llm172/LoRA-MGPO](https://github.com/llm172/LoRA-MGPO) +**Saved:** [docs/lora_mgpo_momentum_perturbation.md](docs/lora_mgpo_momentum_perturbation.md) + +**Hypothesis:** Optimizer momentum supplies a cheap, stable approximation to SAM's adversarial direction. Perturb LoRA parameters along the previous first-moment vector, normalize the radius using an EMA of gradient norms, and compute only one gradient at the perturbed point. + +```py +# Claude: LoRA-MGPO training step +def mgpo_step(loss, W, θ, m, ḡ, ρ, β, optimizer): + ε_θ = ρ * m / (norm(m) * ḡ) # Claude: θ = (A, B) + Ã, B̃ = unpack(θ + ε_θ) + g = ∇(loss(W + B̃ @ Ã), θ + ε_θ) + θ, m ← optimizer.step(θ, m, g) + ḡ ← β * ḡ + (1 - β) * norm(g) + return θ, m, ḡ +``` + +**Evidence:** The authors report a T5 GLUE average of 88.81 versus 82.08 for LoRA and 87.91 for full fine-tuning. On Llama-2, MGPO is the strongest reported PEFT method on MT-Bench, GSM8K, and HumanEval at several ranks, though it remains below full fine-tuning on GSM8K and HumanEval. The large GLUE margin is concentrated in CoLA and MRPC and comes from the authors' own setup. The experiments support optimization stability and conventional generalization, not an explicit curvature measurement or controlled OOD transfer. + +**Grade:** PE+BL+BF=3.5 (beats LoRA and slightly beats full FT on the reported GLUE average) + +--- + ## Scorecard Sorted by evidence strength (max 8). See [scoring legend](#evidence-scoring) above. @@ -947,11 +1050,15 @@ Sorted by evidence strength (max 8). See [scoring legend](#evidence-scoring) abo | # | Method | Score | Breakdown | Theme | | ---: | ------------- | ----: | ----------- | ---------------- | | 6 | PiSSA | 5.0 | PE+BL+BF+DE | SVD basis | +| 36 | FMLoRA | 5.0 | PE+BL+BF+DE | flatness | | 4 | DoRA | 4.5 | PE+BL+BF+WA | dir/strength | | 11 | AntiPaSTO* | 4.5 | PE+DE+OOD | SVD+rotation | | 34 | Flat-LoRA | 4.0 | PE+BL+OOD | flatness | | 13 | BOFT | 4.0 | PE+BF+DE | orthogonal | | 5 | DeLoRA | 3.5 | PE+BL+DE | dir/strength | +| 35 | BAR | 3.5 | PE+BL+DE | flatness | +| 37 | Bi-LoRA | 3.5 | PE+BL+BF | flatness | +| 38 | LoRA-MGPO | 3.5 | PE+BL+BF | flatness | | 8 | SSVD | 3.5 | PE+BL+DE | SVD basis | | 31 | CLOVER | 3.5 | PE+BL+BF | SVD+architecture | | 32 | PSOFT | 3.5 | PE+BL+DE | SVD+orthogonal | @@ -985,7 +1092,7 @@ Sorted by evidence strength (max 8). See [scoring legend](#evidence-scoring) abo ## Themes: What the Evidence Tells Us -Looking across all 34 methods, the successful adapters share a recipe: choose coordinates that align with pretrained structure, constrain updates to preserve that structure, and control update strength explicitly. +Looking across all 38 methods, the successful adapters share a recipe: choose coordinates that align with pretrained structure, constrain updates to preserve that structure, and control update strength explicitly. The pattern is strong enough to organize the literature by theme rather than by year. @@ -1000,13 +1107,13 @@ The *direction-versus-strength* split follows naturally. DoRA, DeLoRA, ROAD, and The *rank* debate is secondary once basis is accounted for. Full-rank updates help on harder tasks (RandLoRA, C3A), but a good low-rank subspace beats a poorly chosen full-rank update (PiSSA, SVFT). "Which subspace" matters more than "how many free directions". -*Curvature* is probably best treated as space-specific rather than as one adapter principle. Flat-LoRA's authors report that smoothing task loss around the merged weights improves LoRA, including under corruption and instruction-following shifts. CrispEdit instead constrains curvature of a capability loss, while TRAM regularizes predictive-distribution changes in function space. Counterexamples with flat non-generalizing minima make weight-space flatness weak evidence about semantic depth. Context curvature of a steering effect remains a plausible but untested diagnostic. The quote-anchored argument and counterevidence are in [the Vargdown evidence map](adapters_vargdown.argdown). +*Flatness and curvature* form one related family, but the differentiated variable matters. Flat-LoRA, FMLoRA, Bi-LoRA, and MGPO seek finite-neighborhood flatness through random, adversarial, or momentum-guided perturbations. BAR instead isolates factor balancedness as a proposed implicit effect of SAM. This flatness-seeking family does not diagonalize a Hessian. The explicit low-loss-curvature hypothesis constrains updates to low-eigenvalue directions of a named loss; CrispEdit is the closest adjacent method. Function and context curvature differentiate predictions or steering effects with respect to inputs, activations, or context paths; TRAM supports this axis, while context curvature of a steering effect remains a plausible but untested diagnostic. Counterexamples with flat non-generalizing minima make weight-space flatness weak evidence about semantic depth. The quote-anchored argument and counterevidence are in [the Vargdown evidence map](adapters_vargdown.argdown). Finally, methods that respect *functional architecture* are promising but early. CLOVER's joint Q-K and V-O treatment outperforms per-matrix updates in reported setups, and ReFT shows targeted activation interventions can be far more parameter-efficient than weight updates. Both suggest that treating transformer layers as computation graphs -- not bags of independent matrices -- is a productive direction. ### What I now believe (and didn't before) -Before writing this catalog, I thought of adapters mainly as engineering trade-offs: LoRA is cheap, full FT is better, pick your budget. After reading 34 adapter papers carefully, I updated on three things: +Before writing this catalog, I thought of adapters mainly as engineering trade-offs: LoRA is cheap, full FT is better, pick your budget. After reading 38 adapter papers carefully, I updated on three things: 1. **The SVD basis outperforms random-initialized and standard bases.** The consistent advantage of SVD-initialized methods (PiSSA > LoRA, SVFT recovering 96% of full FT with 0.006% params, CLOVER's joint SVD beating per-matrix LoRA) is hard to explain as coincidence. The model's singular vectors appear to encode meaningful computational directions that the optimizer discovers faster when given them as a starting point. But SVD is linear and transformers are not; the advantage could be a warm-start effect; and almost no papers compare against other structured bases. "SVD > random" is solid, "SVD is the right basis" remains open. Strength of evidence: moderate (multiple independent groups, multiple modalities, but all within-paper comparisons). diff --git a/adapters_vargdown.argdown b/adapters_vargdown.argdown index 09e21e7..d313606 100644 --- a/adapters_vargdown.argdown +++ b/adapters_vargdown.argdown @@ -1,6 +1,6 @@ === title: Adapters as Representational Hypotheses -- Which Geometric Priors About Transformer Internals Hold Under Intervention? -author: Compiled from 34 PEFT methods plus adjacent evidence (2021--2026) +author: Compiled from 38 PEFT methods plus adjacent evidence (2021--2026) model: mode: strict === @@ -426,25 +426,29 @@ model: +> [Natural Manifold] -// ══════════════════════════════════════════════════════════════════════ -// THEME 8: CURVATURE IS SPACE-SPECIFIC -// Adapter: Flat-LoRA. Adjacent evidence: CrispEdit, SGD subspaces, -// sharpness counterexamples, TRAM, and steering generalization. -// ══════════════════════════════════════════════════════════════════════ +// Claude: ══════════════════════════════════════════════════════════════ +// Claude: THEME 8: FLATNESS AND CURVATURE ARE SPACE-SPECIFIC +// Claude: Flatness-seeking family: Flat-LoRA, FMLoRA, Bi-LoRA, MGPO; +// Claude: BAR is the SAM-balancedness surrogate. +// Claude: Adjacent evidence: CrispEdit, DISAM, SGD subspaces, sharpness +// Claude: counterexamples, TRAM, and steering generalization. +// Claude: ══════════════════════════════════════════════════════════════ # Curvature -[Curvature Is Space-Specific]: Curvature only becomes an intervention - hypothesis after naming the scalar landscape and the space or path in - which it is measured. Weight-, capability-, and function-space curvature - make different predictions about robustness and generalization. - + +[Curvature Is Space-Specific]: Flatness and curvature are related, but the + differentiated variable must be named. Perturbation-based flatness seeking, + explicit Hessian or Gauss-Newton eigenspace constraints, and function or + context curvature make different predictions about generalization. + + + + + + + + + - + (1) [Flat-LoRA Smooths Merged Weights]: Flat-LoRA trains low-rank factors under random perturbations of the merged weight matrix so the solution @@ -453,7 +457,49 @@ model: [evidence](docs/flat_lora_full_parameter_flatness.md#L1-L35) > Despite recent progress in improving LoRA’s performance, the relationship between the LoRA optimization space and the full parameter space is often overlooked. **A solution that appears flat in the loss landscape of the LoRA space may still exhibit sharp directions in the full parameter space, potentially compromising generalization. We introduce Flat-LoRA, which aims to identify a low-rank adaptation situated in a flat region of the full parameter space.** Instead of adopting the well-established sharpness-aware minimization approach, which incurs significant computation and memory overheads, we employ a Bayesian expectation loss objective to preserve training efficiency. Further, we design a refined random perturbation generation strategy for improved performance and carefully manage memory overhead using random seeds. {reason: "ICML 2025; authors' abstract for their own method; the main perturbation scheme covers adapted linear matrices, while all-layer perturbation is reported separately in the appendix; no independent replication found", credence: 0.78} -(2) [CrispEdit Protects Capability Loss]: CrispEdit projects model-editing +(2) [BAR Makes SAM Balancedness Explicit]: BAR replaces SAM's extra + adversarial step with a factor-norm regularizer derived from the claimed + implicit balancedness dynamics of SAM. #observation + [Li, Zhang, He 2024](https://arxiv.org/abs/2410.14802) + [evidence](docs/bar_balancedness_aware_regularization.md#L15-L29) + > Sharpness-aware minimization (SAM) improves generalization of various deep learning tasks. Motivated by popular architectures such as LoRA, we explore the implicit regularization of SAM for scale-invariant problems involving two groups of variables. **Instead of focusing on commonly used sharpness, this work introduces a concept termed balancedness, defined as the difference between the squared norm of two variables.** This allows us to depict richer global behaviors of SAM. In particular, our theoretical and empirical findings reveal that i) SAM promotes balancedness; and ii) the regularization on balancedness is data-responsive – outliers have stronger impact. The latter coincides with empirical observations that SAM outperforms SGD in the presence of outliers. Leveraging the implicit regularization, we develop a resource-efficient SAM variant, balancedness-aware regularization (BAR), tailored for scale-invariant problems such as finetuning language models with LoRA. + {reason: "NeurIPS 2024; authors' abstract and theoretical framing; BAR is reused as a baseline by Flat-LoRA, FMLoRA, and Bi-LoRA, but its tests are few-shot or ordinary held-out rather than controlled OOD", credence: 0.82} +(3) [FMLoRA Transfers Full-Space Perturbations]: FMLoRA reconstructs a + full-weight SAM direction from LoRA gradients and transfers it into one + factor; EFMLoRA reuses an EMA perturbation for near-LoRA cost. #observation + [Deng et al. 2025](https://arxiv.org/abs/2508.00522) + [evidence](docs/fmlora_flat_minima_lora.md#L8-L28) + > Little research explores the correlation between the expressive ability and generalization ability of the low-rank adaptation (LoRA). Sharpness-Aware Minimization (SAM) improves model generalization for both Convolutional Neural Networks (CNNs) and Transformers by encouraging convergence to locally flat minima. However, the connection between sharpness and generalization has not been fully explored for LoRA due to the lack of tools to either empirically seek flat minima or develop theoretical methods. **In this work, we propose Flat Minima LoRA (FMLoRA) and its efficient version i.e., EFMLoRA, to seek flat minima for LoRA. Concretely, we theoretically demonstrate that perturbations in the full parameter space can be transferred to the low-rank subspace.** This approach eliminates the potential interference introduced by perturbations across multiple matrices in the low-rank subspace. Our extensive experiments on large language models and vision-language models demonstrate that EFMLoRA achieves optimize efficiency comparable to that of LoRA while simultaneously attaining comparable or even better performance. For example, on the GLUE dataset with RoBERTa-large, EFMLoRA outperforms LoRA and full fine-tuning by 1.0% and 0.5% on average, respectively. On vision-language models e.g., Qwen-VL-Chat, there are performance improvements of 1.5% and 1.0% on the SQA and VizWiz datasets, respectively. These empirical results also verify that the generalization of LoRA is closely related to sharpness, which is omitted by previous methods. + {reason: "AAAI 2026; authors' abstract and own benchmark comparisons; reported gains cover several architectures, but the claimed distribution shift is mostly few-shot or ordinary downstream evaluation", credence: 0.72} +(4) [Bi-LoRA Separates Descent and Ascent]: Bi-LoRA trains a primary LoRA + branch for task descent and an auxiliary low-rank branch for adversarial + ascent, then discards the auxiliary branch. #observation + [Liu et al. 2025](https://arxiv.org/abs/2508.19564) + [evidence](docs/bi_lora_sharpness_aware.md#L16-L20) + > Low-Rank Adaptation (LoRA) enables parameter-efficient fine-tuning of large pre-trained models. Yet LoRA can face generalization challenges. One promising way to improve the generalization is Sharpness-Aware Minimization (SAM), which has proven effective for small-scale training scenarios. **In this paper, we propose Bi-directional Lo w-R ank A daptation (Bi-LoRA), which introduces an auxiliary adversarial LoRA module. This design explicitly decouples sharpness optimization, handled by the auxiliary module, from task adaptation, performed by the primary module.** Such a separation yields two key benefits. First, it transforms SAM’s sequential computation of adversarial perturbation and gradient descent into a parallel form, which roughly halves the time and conquers the main obstacle of applying SAM in LoRA. Second, it provides perturbations from the auxiliary module that do not collapse into the restricted optimization subspace of the primary module, enabling broader sharpness exploration and flatter minima. Bi-LoRA simultaneously achieves both efficiency and effectiveness within a single framework, as validated by extensive experiments across diverse architectures and tasks. + {reason: "ICLR 2026; authors' method section; same-paper comparisons mostly favor Bi-LoRA over LoRA, LoRA-SAM, and Flat-LoRA on cross-dataset LLM tests, without a controlled domain-generalization protocol", credence: 0.76} +(5) [MGPO Reuses Optimizer Momentum]: LoRA-MGPO perturbs trainable LoRA + factors along the optimizer's previous first moment and normalizes the + radius with an EMA of gradient norms. #observation + [Chang et al. 2025](https://aclanthology.org/2025.findings-emnlp.34/) + [evidence](docs/lora_mgpo_momentum_perturbation.md#L28-L34) + > Parameter-efficient fine-tuning (PEFT), partic-ularly Low-Rank Adaptation (LoRA), adapts large language models (LLMs) by training only a small fraction of parameters. However, as the rank of the low-rank matrices used for adap-tation increases, LoRA often exhibits an un-stable "double descent" phenomenon, charac-terized by transient divergence in the training loss, which delays convergence and impairs generalization by causing instability due to the attraction to sharp local minima. **To address this, we introduce LoRA-MGPO , a framework that incorporates Momentum-Guided Pertur-bation Optimization (MGPO). MGPO stabi-lizes training dynamics by mitigating the dou-ble descent phenomenon and guiding weight perturbations using momentum vectors from the optimizer’s state, thus avoiding dual gra-dient computations.** Additionally, an adaptive normalization scheme scales the magnitude of perturbations based on an exponential mov-ing average (EMA) of gradient norms, further enhancing stability. While EMA controls the magnitude of the perturbations, MGPO guides their direction, ensuring a more stable opti-mization trajectory. + {reason: "Findings EMNLP 2025; authors' method section and released code; improvements are conventional NLU and cross-dataset NLG results, while the link to flat minima is indirect", credence: 0.64} +---- +(6) [Flatness Seeking Is Indirect Curvature Control]: Flat-LoRA, FMLoRA, + Bi-LoRA, and MGPO seek low loss in a finite weight neighborhood; BAR + instead isolates factor balancedness as a proposed implicit effect of + SAM. Near a stationary point the four neighborhood objectives are + sensitive to the largest Hessian eigenvalues, but none of the five + computes a Hessian eigenspace or restricts updates to low-eigenvalue + directions. + {reason: "local Taylor expansion links neighborhood sharpness to curvature; BAR implements a derived surrogate, and the other four implement perturbation objectives rather than spectral projection", inference: 0.90} + +> [Curvature Is Space-Specific] + + + + +(1) [CrispEdit Protects Capability Loss]: CrispEdit projects model-editing updates into the low-curvature subspace of a separate capability loss, estimated with Gauss-Newton curvature and K-FAC. #observation [Ikram et al. 2026](https://arxiv.org/abs/2602.15823) @@ -461,11 +507,11 @@ model: > We present CrispEdit, a scalable and principled second-order editing algorithm that treats capability preservation as an explicit constraint, unifying and generalizing several existing editing approaches. **CrispEdit formulates editing as constrained optimization and enforces the constraint by projecting edit updates onto the low-curvature subspace of the capability-loss landscape.** At the crux of CrispEdit is expressing capability constraint via Bregman divergence, whose quadratic form yields the Gauss–Newton Hessian exactly and even when the base model is not trained to convergence. We make this second-order procedure efficient at the LLM scale using Kronecker-factored approximate curvature (K-FAC) and a novel matrix-free projector that exploits Kronecker structure to avoid constructing massive projection matrices. {reason: "May 2026 preprint; authors' abstract for a model-editing method rather than an adapter; capability is evaluated on a designated reference set", credence: 0.66} ---- -(3) [Low Curvature Protects the Chosen Loss]: Flat-LoRA seeks task-loss - insensitivity around merged weights; CrispEdit seeks capability-loss - insensitivity along edit directions. Neither makes low curvature a +(2) [Low Curvature Protects the Chosen Loss]: CrispEdit's spectral projector + protects a designated capability loss. This is the closest prior art to + an explicit low-loss-curvature adapter, but it does not make curvature a detector of the target behavior's semantic depth. - {reason: "the methods constrain different scalar landscapes; the distinction follows from their objectives rather than their reported benchmark gains", inference: 0.82} + {reason: "CrispEdit names the protected scalar loss and reference set; projecting an update away from its stiff directions says what is preserved, not why the edited behavior generalizes", inference: 0.84} +> [Curvature Is Space-Specific] @@ -494,6 +540,33 @@ model: +> [Curvature Is Space-Specific] + + +(1) [Plain SAM Can Fail Under Domain Shift]: DISAM's authors report that + ordinary SAM can underperform ERM when source domains converge at + different rates. #observation + [Zhang et al. 2024](https://arxiv.org/abs/2405.18861) + [evidence](docs/disam_domain_shift_sharpness.md#L46-L58) + > Nonetheless, these methods cannot solve generalizability scenarios that involve training data of multiple domains with domain shifts like Domain Generalization (DG) (Ben-David et al., 2010; Li et al., 2017). **In this study, we observed that sometimes SAM even has a detrimental impact in situations where there exist domain shifts across multiple domains as shown in Figure 1(1(a)).** While a few studies have incorporated SAM-based methods in domain generalization tasks (Wang et al., 2023b; Foret et al., 2021), they cannot ensure consistent improvements in generalizability during domain shifts due to their reliance on the i.i.d assumption. Upon a thorough analysis of the behavior of SAM under domain shifts, we discovered that the degradation of the training process caused by SAM from the disparity in convergence degree among different domains as shown in Figure 1(1(a)). Given the inconsistency in the degree and direction of convergence among different domains during training (Arjovsky et al., 2019; Krueger et al., 2021), the straightforward application of SAM for perturbations may not only disrupt convergence but also generate perturbation directions that are not adequately coherent to the geometric characteristics of the entire loss landscape. + {reason: "ICLR 2024; authors' diagnosis from DomainBed experiments; this directly limits the inference from ordinary SAM or flat-adapter gains to OOD transfer", credence: 0.83} +(2) [DISAM Uses a Genuine OOD Protocol]: DISAM calibrates the SAM + perturbation using source-domain loss variance and evaluates on domains + excluded from training. #observation + [Zhang et al. 2024](https://arxiv.org/abs/2405.18861) + [evidence](docs/disam_domain_shift_sharpness.md#L1611-L1627) + > We evaluate DISAM on five datasets PACS (Li et al., 2017), VLCS (Fang et al., 2013) OfficeHome (Venkateswara et al., 2017), TerraIncognita (Beery et al., 2018) (abbreviated as TerraInc), and DomainNet (Peng et al., 2019), following the DomainBed benchmark (Gulrajani & Lopez-Paz, 2021). For fair comparison, we adhere to the training and evaluation protocol outlined in DomainBed. Evaluation. **The standard leave-one-domain-out strategy is used in evaluation. Specially, the unseen domain is used to evaluate the out-of-domain generalization, and the validation sets of source domains are used to measure the in-domain generalization, while the others are used for training.** Final accuracy is averaged across all settings, and the performance is the averaging over three trials with distinct random seeds. Detailed statistics for each case of all datasets are provided in Appendix C. + [evidence](docs/disam_domain_shift_sharpness.md#L2186-L2192) + > We propose incorporating our domain-inspired adaptive adjustment into three SAM-based methods: SAM (Foret et al., 2021), GSAM (Zhuang et al., 2022), and SAGM (Wang et al., 2023b) on five datasets of DomainBed with ResNet50 backbone. Table 1 shows that our Domain-Inspired SAM can mitigate issues arising from SAM’s training under domain shifts, by comparing averaged in-domain and out-of-domain performance of leading SAM methods, with and without DISAM. In-domain results show domain-inspired perturbations enhance convergence, especially on the TerraInc dataset with substantial domain gaps. **In Out-of-domain results, DISAM consistently improves generalization, with average improvements of 1.9% for SAM, 1.7% for GSAM, and 1.9% for SAGM.** Notably, SAM performs well when the performance gap between in-domain and out-of-domain is small but worse than ERM on datasets like TerraInc with large gaps, which proves our analysis of SAM’s shortcomings under domain shifts. This shows SAM’s inconsistent convergence for large domain shifts, which DISAM addresses by incorporating domain-inspired adaptive adjustments based on domain-level convergence degree. Incorporating CORAL constraints, a recognized effective traditional DG method on DomainBed improves SAGM with DISAM and sets new state-of-the-art results on all settings. + {reason: "ICLR 2024; standard leave-one-domain-out DomainBed protocol over five datasets, three trials, plus CLIP prompt-tuning experiments; author-reported results but materially stronger OOD evidence than downstream test accuracy", credence: 0.86} +---- +(3) [OOD Flatness Needs Shift Information]: DISAM is not a LoRA adapter, + but it is strong evidence against treating generic weight-space flatness + as domain-general by default. Its gains require domain labels and a + domain-loss variance term when constructing the perturbation. + {reason: "plain SAM sometimes loses to ERM while domain-calibrated SAM improves leave-one-domain-out accuracy; the added domain information, not flatness alone, distinguishes the methods", inference: 0.82} + +> [Curvature Is Space-Specific] + + (1) [Flatness Need Not Generalize]: Flat non-generalizing minimizers exist, diff --git a/docs/bar_balancedness_aware_regularization.md b/docs/bar_balancedness_aware_regularization.md new file mode 100644 index 0000000..ff691e4 --- /dev/null +++ b/docs/bar_balancedness_aware_regularization.md @@ -0,0 +1,9980 @@ +Title: Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems + +URL Source: https://arxiv.org/html/2410.14802 + +Markdown Content: + Abstract +1Introduction +2Preliminaries +3SAM for Non-Overparametrized Problems +4SAM for Overparametrized Problems +5Implicit Regularization Made Explicit +6Numerical Experiments +7Discussions + References +Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems +Bingcong Li &Liang Zhang &Niao He +Department of Computer Science ETH Zurich, Switzerland {bingcong.li, liang.zhang, niao.he}@inf.ethz.ch +Abstract + +Sharpness-aware minimization (SAM) improves generalization of various deep learning tasks. Motivated by popular architectures such as LoRA, we explore the implicit regularization of SAM for scale-invariant problems involving two groups of variables. Instead of focusing on commonly used sharpness, this work introduces a concept termed balancedness, defined as the difference between the squared norm of two variables. This allows us to depict richer global behaviors of SAM. In particular, our theoretical and empirical findings reveal that i) SAM promotes balancedness; and ii) the regularization on balancedness is data-responsive – outliers have stronger impact. The latter coincides with empirical observations that SAM outperforms SGD in the presence of outliers. Leveraging the implicit regularization, we develop a resource-efficient SAM variant, balancedness-aware regularization (BAR), tailored for scale-invariant problems such as finetuning language models with LoRA. BAR saves +95 +% + computational overhead of SAM, with enhanced test performance across various tasks on RoBERTa, GPT2, and OPT-1.3B. + +1Introduction + +Sharpness-aware minimization (SAM) is emerging as an appealing optimizer, because it enhances generalization performance on various downstream tasks across vision and language applications (Foret et al., 2021; Chen et al., 2022; Bahri et al., 2022). The success of SAM is typically explained using its implicit regularization (IR) toward a flat solution (Wen et al., 2023a). + +However, existing results only characterize sharpness/flatness near local minima (Wen et al., 2023a). Little is known about early convergence, despite its crucial role in SAM’s implicit regularization (Agarwala and Dauphin, 2023). In addition, theoretical understanding of SAM highly hinges upon the existence of positive eigenvalues of Hessians (Wen et al., 2023a), leaving gaps in nonconvex scenarios where the Hessian can be negative definite. The limitations above lead to our first question (Q1): can we broaden the scope of implicit regularization to depict global behaviors in SAM? + +Moreover, scenarios where SAM popularizes often involve certain form of data anomalies, such as outliers and large data variance. SAM has provable generalization benefits on sparse coding problems in the small signal-to-noise ratio (SNR) regime (Chen et al., 2023). Remarkable performance of SAM is also observed under distributional shifts, e.g., domain adaptation (Wang et al., 2023), meta-learning (Abbas et al., 2022), and transfer learning in language models (Bahri et al., 2022; Sherborne et al., 2023). Evidences above motivate our second question (Q2): can implicit regularization of SAM reflect its enhanced performance under data anomalies? + +This work answers both Q1 and Q2 within a class of scale-invariant problems. The focus on scale-invariance is motivated by its prominence in deep learning architectures. Consider variables +𝐱 +∈ +ℝ +𝑑 +1 + and +𝐲 +∈ +ℝ +𝑑 +2 +, both in high-dimensional space. The problems of interest can be categorized into non-overparametrization (NOP) and overparametrization (OP), based on whether the dimension of variables ( +𝑑 +1 ++ +𝑑 +2 +) is greater than dimension of +dom +⁢ +𝑓 +, + + + +NOP: +⁢ +min +𝐱 +, +𝐲 +⁡ +𝑓 +𝑛 +⁢ +( +𝐱𝐲 +⊤ +) += +𝔼 +𝜉 +∼ +𝒟 +⁢ +[ +𝑓 +𝑛 +𝜉 +⁢ +( +𝐱𝐲 +⊤ +) +] +, + +(1a) + + + +OP: +⁢ +min +𝐱 +, +𝐲 +⁡ +𝑓 +𝑜 +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) += +𝔼 +𝜉 +∼ +𝒟 +⁢ +[ +𝑓 +𝑜 +𝜉 +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) +] +. + +(1b) + +Here, +𝑑 +1 += +𝑑 +2 + is assumed for OP, and +𝒟 + denotes the training data. For both cases, the losses are nonconvex in +( +𝐱 +, +𝐲 +) +. Scale-invariance refers to that +( +𝛼 +⁢ +𝐱 +, +𝐲 +/ +𝛼 +) + share the same objective value +∀ +𝛼 +≠ +0 +. It naturally calls for implicit regularization from optimization algorithms to determine the value of +𝛼 +. We focus on two-variable problems in the main text for simplicity and generalize the results to multi-layer cases in the appendix. Problems (1a) and (1b) are inspired by widely-adopted modules in deep learning, where low rank adapters (LoRA) for finetuning language models is NOP, and softmax in attention falls in OP framework (Hu et al., 2022; Vaswani et al., 2017). + + + +(a) non-overparametrized (NOP) (b) overparametrized (OP) +Figure 1:Implicit regularization of SAM on balancedness. The losses for NOP and OP are +𝔼 +⁢ +[ +‖ +𝐱𝐲 +⊤ +− +( +𝐀 ++ +𝛼 +⁢ +𝐍 +) +‖ +2 +] + and +𝔼 +⁢ +[ +‖ +𝐱 +⊤ +⁢ +𝐲 +− +( +𝑎 ++ +𝛼 +⁢ +𝑛 +) +‖ +2 +] +, respectively. Here, +𝐀 + is the ground truth matrix, +𝐍 + is the Gaussian noise, and +𝛼 + controls the SNR. Left of (a) and (b): +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +| + vs. iteration. Right of (a) and (b): +| +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +| + vs. iteration, where +( +𝐠 +𝐱 +𝑡 +, +𝐠 +𝐲 +𝑡 +) + denotes stochastic gradients. + +This work studies SAM’s implicit regularization on balancedness, defined as +ℬ +𝑡 += +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) +. Balancedness is a useful alternative to sharpness for (1) because: i) it enables us to go beyond local minima and describe the behavior over SAM’s entire trajectory; ii) analyses and assumptions can be significantly simplified when working with +ℬ +𝑡 +; and, iii) it enables a data-driven perspective for understanding SAM. Building on balancedness, we answer our major questions. + +For Q1, we prove that even with imbalanced initialization, SAM drives +| +ℬ +𝑡 +| +→ +0 + for OP, while ensuring a small +| +ℬ +𝑡 +| + in NOP. In contrast, we also prove that balancedness of SGD is unchanged over iterations. This clear distinction between SAM and SGD is illustrated in Fig. 1. Thanks to the adoption of balancedness, our results on implicit regularization have no requirement on the batchsize compared to (Wen et al., 2023a) and can be extended to explain +𝑚 +-sharpness in (Foret et al., 2021). + +Regarding Q2, we present analytical and empirical evidences that data anomalies (e.g., samples with large noise) have stronger impact on balancedness for both NOP and OP. Fig. 1 showcases an example where SAM is applied on the same problem with different SNRs. Smaller SNR (i.e., larger +𝛼 +) promotes balancedness faster. Being more balanced with noisy data also aligns well with previous studies (Chen et al., 2023; Wang et al., 2023), which show that SAM performs better than SGD under data anomalies. This data-driven behavior of SAM is well depicted through balancedness. + +Our theoretical understanding on balancedness also cultivates practical tools. In particular, we explicify the implicit regularization of SAM as a data-driven regularizer. When applied on top of, e.g., SGD, it enables a computationally efficient variant of SAM, balancedness-aware regularization (BAR), suited for scale-invariant problems such as finetuning language models with LoRA (Hu et al., 2022). BAR eliminates the need to compute the second gradient in SAM, thereby significantly reducing overhead in large-scale settings. BAR improves the test performance of LoRA on three representative downstream tasks on RoBERTa, GPT2, and OPT, while saving +95 +% + computational overhead of SAM. Moreover, this is the first efficient SAM approach derived from SAM’s implicit regularization. In a nutshell, our contribution can be summarized as: + +❖ + +Theories. Balancedness is introduced as a new metric for implicit regularization in SAM. Compared to sharpness, balancedness enables us to depict richer behaviors – SAM favors balanced solutions for both NOP and OP, and data anomalies have stronger regularization on balancedness. + +❖ + +Practice. Implicit regularization of SAM is made explicit for practical merits. The resulting approach, balancedness-aware regularization (BAR), improves accuracy for finetuning language models with LoRA, while significantly saving computational overhead of SAM. + +Notation. Bold lowercase (capital) letters denote column vectors (matrices); +∥ +⋅ +∥ + stands for +ℓ +2 + (Frobenius) norm of a vector (matrix), and +( +⋅ +) +⊤ + refers to transpose. + +1.1Related Work + +Related topics are streamlined here, with comprehensive discussions deferred to Apdx. A.2. + +Scale-invariance in deep learning. Scale-invariant modules are prevalent in modern neural networks, such as LoRA, ReLU networks, and softmax in attention. However, scale-invariant problems are not yet fully understood, especially from a theoretical perspective. Neyshabur et al. (2018) develop scale-invariant PAC-Bayesian bounds for ReLU networks. A scale-invariant SGD is developed in (Neyshabur et al., 2015), and this approach becomes more practical recently in (Gonon et al., 2024). Linear neural networks entail scale-invariance and overparametrization simultaneously, and IR of (S)GD on quadratic loss is established in (Arora et al., 2018; Du et al., 2018; Gidel et al., 2019). IR of GD for softmax attention in transformers is studied in (Sheen et al., 2024) assuming linearly separable data. It is pointed out in (Dinh et al., 2017) that sharpness is sensitive to scaling, while our results indicate that when taking the training trajectory into account, SAM excludes extreme scaling. + +Mechanism behind SAM. To theoretically explain the success of SAM, Bartlett et al. (2023) analyze sharpness on quadratic losses. Wen et al. (2023a) focus on sharpness of SAM near the solution manifold on smooth loss functions, requiring batchsize to be 1 in the stochastic case. Andriushchenko and Flammarion (2022) consider sparsity of SAM on (overparametrized) diagonal linear networks on a regression problem. Chen et al. (2023) study the benign overfitting of SAM on a two-layer ReLU network. In general, existing studies on SAM’s implicit regularization focus more on sharpness and do not fully capture scale-invariance. In comparison, our results i) are Hessian-free and hence sharpness-free; ii) have no constraint on batchsize; and iii) hold for both NOP and OP. + +SAM variants. Approaches in (Kim et al., 2022; Kwon et al., 2021) modify SAM for efficiency under coordinate-wise ill-scaling, while our results suggest that SAM favors balancedness between layers. Computationally efficient SAM variants are developed through reusing or sparsifying gradients (Liu et al., 2022; Mi et al., 2022); stochastic perturbation (Du et al., 2022a); switching to SGD (Jiang et al., 2023); and connecting with distillation (Du et al., 2022b). Our BAR can be viewed as resource-efficient SAM applied specifically for scale-invariant problems such as LoRA. Different from existing works, BAR is the first to take inspiration from the implicit regularization of SAM. + +2Preliminaries + +This section briefly reviews SAM and then compares sharpness with balancedness. For a smoother presentation, our main numerical benchmark, LoRA (Hu et al., 2022), is revisited in Sec. 5. + +2.1Recap of SAM +Algorithm 1 SAM (Foret et al., 2021) +1:Initialize: +𝐰 +0 +, +𝜌 +, +𝑇 +, +𝜂 +2:for  +𝑡 += +0 +, +… +, +𝑇 +− +1 + do +3:     Sample +𝜉 + to get a minibatch +ℳ +𝑡 +4:     Define stochastic gradient on +ℳ +𝑡 + as +∇ +ℎ +𝑡 +⁢ +( +⋅ +) +5:     Find +𝜖 +𝑡 += +𝜌 +⁢ +∇ +ℎ +𝑡 +⁢ +( +𝐰 +𝑡 +) +/ +‖ +∇ +ℎ +𝑡 +⁢ +( +𝐰 +𝑡 +) +‖ +6:     Update via +𝐰 +𝑡 ++ +1 += +𝐰 +𝑡 +− +𝜂 +⁢ +∇ +ℎ +𝑡 +⁢ +( +𝐰 +𝑡 ++ +𝜖 +𝑡 +) +7:end for + +Sharpness-aware minimization (SAM) is designed originally to seek for solutions in flat basins. The idea is formalized by enforcing small loss around the entire neighborhood in parameter space, i.e., +min +𝐰 +⁡ +max +‖ +𝜖 +‖ +≤ +𝜌 +⁡ +ℎ +⁢ +( +𝐰 ++ +𝜖 +) +, where +𝜌 + is the radius of considered neighborhood, and +ℎ +⁢ +( +𝐰 +) +:= +𝔼 +𝜉 +⁢ +[ +ℎ +𝜉 +⁢ +( +𝐰 +) +] +. Practical implementation of SAM is summarized under Alg. 1. It is proved in (Wen et al., 2023a) that +‖ +∇ +ℎ +𝑡 +⁢ +( +𝐰 +) +‖ +≠ +0 + (in line 5) holds for any +𝜌 + under most initialization. Based on this result and similar to (Dai et al., 2023), we assume that SAM iterates are well-defined. + +Limitation of sharpness. Coming naturally with SAM is the so-termed sharpness, given by +𝒮 +⁢ +( +𝐰 +) +:= +max +‖ +𝜖 +‖ +≤ +𝜌 +⁡ +ℎ +⁢ +( +𝐰 ++ +𝜖 +) +− +ℎ +⁢ +( +𝐰 +) +. When +‖ +∇ +ℎ +⁢ +( +𝐰 +) +‖ +→ +0 +, +𝒮 +⁢ +( +𝐰 +) + can be approximated using (scaled) largest eigenvalue of Hessian (Zhuang et al., 2022). This approximation is widely exploited in literature to study the implicit regularization of SAM. Consequently, most results only hold locally – behaviors near +‖ +∇ +ℎ +⁢ +( +𝐰 +) +‖ +→ +0 + are studied. In addition, sharpness (the largest eigenvalue) is not always informative for scale-invariant problems (1). Consider +ℎ +⁢ +( +𝑥 +, +𝑦 +) += +𝑥 +⁢ +𝑦 + for example. The sharpness is +1 + for any +( +𝑥 +, +𝑦 +) + – these points are not distinguishable in terms of sharpness. + +2.2Prelude on Balancedness + +Balancedness +ℬ +𝑡 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + turns out to be an intriguing alternative to sharpness on the scale-invariant problem (1). Being a global metric, balancedness is capable of describing the entire trajectory of an algorithm, regardless of proximity to critical points or definiteness of Hessian. + +How does +ℬ +𝑡 + evolve in different algorithms? To set a comparing benchmark of SAM, we first borrow results from previous works on SGD. Following implicit regularization literature such as (Arora et al., 2018, 2019b; Wen et al., 2023a), we consider SGD with infinitesimally small learning rate +𝜂 +→ +0 + for the NOP problem (1a) + + +𝐱 +𝑡 ++ +1 += +𝐱 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐱 +𝑡 +, +𝐲 +𝑡 ++ +1 += +𝐲 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐲 +𝑡 +. + +(2) +Theorem 1 ((Arora et al., 2018, 2019a; Ji and Telgarsky, 2019; Ahn et al., 2023)). + +When applying SGD on the NOP (1a), the limiting flow with +𝜂 +→ +0 + satisfies +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 += +‖ +𝐱 +0 +‖ +2 +− +‖ +𝐲 +0 +‖ +2 + for all +𝑡 +> +0 +. In other words, +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +0 + holds. + +Theorem 1 shows that +ℬ +𝑡 +≡ +ℬ +0 + given +𝜂 +→ +0 +. A graphical illustration can be found in Fig. 1 (a). Another interesting observation is that given the same initialization, +ℬ +𝑡 + is fixed for SGD regardless of training datasets. This suggests that SGD is less adaptive to data. A similar result of Theorem 1 can be established for SGD on OP. The full statement is deferred to Apdx. C.1; see also Fig. 1 (b). + +Merits of being balance. Because +ℬ +0 + is preserved, SGD is sensitive to initialization. For example, +( +𝐱 +0 +, +𝐲 +0 +) + and +( +2 +⁢ +𝐱 +0 +, +0.5 +⁢ +𝐲 +0 +) + can result in extremely different trajectories, although the same objective value is shared at initialization. Most of existing works initialize +ℬ +0 +≈ +0 + to promote optimization benefits, because the variance of stochastic gradient is small and the local curvature is harmonized around a balanced solution. Take the stochastic gradient of NOP on minibatch +ℳ + for example + + +𝐠 +𝐱 += +1 +| +ℳ +| +⁢ +[ +∑ +𝜉 +∈ +ℳ +∇ +𝑓 +𝑛 +𝜉 +⁢ +( +𝐱𝐲 +⊤ +) +] +⁢ +𝐲 +, +𝐠 +𝐲 += +1 +| +ℳ +| +⁢ +[ +∑ +𝜉 +∈ +ℳ +∇ +𝑓 +𝑛 +𝜉 +⁢ +( +𝐱𝐲 +⊤ +) +] +⊤ +⁢ +𝐱 +. + +(3) + +Assuming bounded variance +𝔼 +⁢ +[ +‖ +1 +| +ℳ +| +⁢ +∑ +𝜉 +∈ +ℳ +∇ +𝑓 +𝑛 +𝜉 +⁢ +( +𝐱𝐲 +⊤ +) +− +∇ +𝑓 +𝑛 +⁢ +( +𝐱𝐲 +⊤ +) +‖ +2 +] +≤ +𝜎 +2 +, it can be seen that the variance of +[ +𝐠 +𝐱 +, +𝐠 +𝐲 +] + is bounded by +𝜎 +2 +⁢ +( +‖ +𝐱 +‖ +2 ++ +‖ +𝐲 +‖ +2 +) +. In other words, among +{ +( +𝐱 +, +𝐲 +) +| +𝐱𝐲 +⊤ += +𝐖 +} +, gradient variance is minimized if +‖ +𝐱 +‖ += +‖ +𝐲 +‖ +, i.e., being balance. Moreover, block smoothness parameters +𝐿 +𝑛 +𝐱 + and +𝐿 +𝑛 +𝐲 +1 also hint upon the difficulties for optimization, where large values typically correspond to slow convergence (Bottou et al., 2018; Nesterov, 2004). With the help of Assumption 1 (in the next subsection), it can be seen that +𝐿 +𝑛 +𝐱 += +𝐿 +𝑛 +⁢ +‖ +𝐲 +‖ +2 + and +𝐿 +𝑛 +𝐲 += +𝐿 +𝑛 +⁢ +‖ +𝐱 +‖ +2 +. In other words, a large +| +ℬ +𝑡 +| + implies difficulty for optimizing one variable than the other. For these reasons, balancedness is well-appreciated in domains such as matrix factorization/sensing – a special case of (1a) (Tu et al., 2016; Bartlett et al., 2018; Du et al., 2018; Ge et al., 2017). It is also observed that balanced neural networks are easier to optimize relative to unbalanced ones (Neyshabur et al., 2015). + +2.3Assumptions and Prerequisites + +To gain theoretical insights of scale-invariant problems in (1), we assume that the loss has Lipschitz continuous gradient on +dom +⁢ +𝑓 + following common nonconvex optimization and SAM analyses (Bottou et al., 2018; Andriushchenko and Flammarion, 2022; Wen et al., 2023a). + +Assumption 1. + +Let +𝐖 +∈ +ℝ +𝑑 +1 +× +𝑑 +2 +, and +𝑤 +∈ +ℝ +. For each +𝜉 +, +𝑓 +𝑛 +𝜉 +⁢ +( +𝐖 +) + and +𝑓 +𝑜 +𝜉 +⁢ +( +𝑤 +) + in (1) have +𝐿 +𝑛 +, and +𝐿 +𝑜 + Lipschitz continuous gradient, respectively. + +Scale-invariant problems are challenging to solve even on simple problems in Fig. 1. Even GD can diverge on some manually crafted initialization (De Sa et al., 2015; Arora et al., 2019a). With proper hyperparameters this rarely happens in practice; hence, we focus on scenarios where SGD and SAM do not diverge. This assumption is weaker than the global convergence needed in (Andriushchenko and Flammarion, 2022), and is similar to the assumption on existence (Wen et al., 2023a). + +3SAM for Non-Overparametrized Problems + +This section tackles the implicit regularization of SAM on NOP (1a). Motivated by practical scenarios such as LoRA, we focus on cases initialized with large +| +ℬ +0 +| +. + +When ambiguity is absent, the subscript in +𝑓 +𝑛 + and +𝐿 +𝑛 + is ignored in this section for convenience. Applying Alg. 1 on NOP, the update of SAM can be written as + + + +𝐱 +~ +𝑡 += +𝐱 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐠 +𝐱 +𝑡 +, + +𝐲 +~ +𝑡 += +𝐲 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐠 +𝐲 +𝑡 + +(4a) + + +𝐠 +𝐱 +~ +𝑡 += +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +~ +𝑡 +⁢ +𝐲 +~ +𝑡 +⊤ +) +⁢ +𝐲 +~ +𝑡 +, + +𝐠 +𝐲 +~ +𝑡 += +[ +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +~ +𝑡 +⁢ +𝐲 +~ +𝑡 +⊤ +) +] +⊤ +⁢ +𝐱 +~ +𝑡 + +(4b) + + +𝐱 +𝑡 ++ +1 += +𝐱 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐱 +~ +𝑡 +, + +𝐲 +𝑡 ++ +1 += +𝐲 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐲 +~ +𝑡 + +(4c) + +where +𝜌 +> +0 + is the radius of SAM perturbation; +𝑢 +𝑡 +:= +1 +/ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +; and +𝑓 +𝑡 +, +∇ +𝑓 +𝑡 + denote the loss, stochastic gradient on minibatch +ℳ +𝑡 +, respectively. + +Theorem 2. + +(Dynamics of SAM.) Suppose that Assumption 1 holds. Consider SAM for NOP in (4) with a sufficiently small +𝜌 +. Let +ℬ +𝑡 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) +. For some +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) + and +𝜂 +→ +0 +, the limiting flow of SAM guarantees that + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +𝜌 +⁢ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 ++ +𝒜 +𝑡 +. + +(5) + +Moreover, the change on +ℬ +𝑡 + depends on the difference of stochastic gradients on +𝐱 +𝑡 + and +𝐲 +𝑡 +, i.e., + + +𝜌 +⁢ +| +‖ +𝐠 +𝐱 +𝑡 +‖ +− +‖ +𝐠 +𝐲 +𝑡 +‖ +| +− +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) +≤ +| +𝑑 +⁢ +ℬ +𝑡 +𝑑 +⁢ +𝑡 +| +≤ +𝜌 +⁢ +| +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +| ++ +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) +. + +(6) + +Unlike SGD for which +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +0 +, Theorem 2 states that the balancedness for SAM is driven by gradient difference +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +. To gain some intuition, if we estimate +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +∝ +‖ +𝐲 +𝑡 +‖ +2 +− +‖ +𝐱 +𝑡 +‖ +2 + based on (3) and ignore +𝒜 +𝑡 +, it can be seen that +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 +∝ +− +𝜌 +⁢ +ℬ +𝑡 +. This indicates the contraction on +| +ℬ +𝑡 +| +. A graphical illustration on decreasing +| +ℬ +𝑡 +| +, and its relation with gradient difference can be found in Figs. 1 (a) and 2 (a). Moreover, this implicit regularization on balancedness is global as it holds for all +𝑡 + regardless of whether +( +𝐱 +𝑡 +, +𝐲 +𝑡 +) + is close to local optima. Thanks to adopting balancedness as the metric, Theorem 2 also poses no requirement on the batchsize. + +SAM promotes balancedness. As discussed in Section 2.2, unbalancedness is burdensome for optimization. SAM overcomes this by implicitly favoring relatively balanced solutions. + +Corollary 1. + +(Informal.) Under some regularity conditions, there exists +ℬ +¯ +𝑡 +𝜌 +≥ +0 + such that whenever +| +ℬ +𝑡 +| +> +ℬ +¯ +𝑡 +𝜌 +, the magnitude of +ℬ +𝑡 + shrinks, where +ℬ +¯ +𝑡 +𝜌 + can be found in (21) at appendix. + +Corollary 1 shows that SAM promotes balancedness until +| +ℬ +𝑡 +| + reaches lower bounds +ℬ +¯ +𝑡 +𝜌 +. Because +ℬ +¯ +𝑡 +𝜌 + depends on SAM’s trajectory, we plot +1 +𝑇 +⁢ +∫ +0 +𝑇 +ℬ +¯ +𝑡 +𝜌 +⁢ +𝑑 +𝑡 + using dotted lines for better visualization in Fig. 2 (a). It can be seen that our calculation on +ℬ +¯ +𝑡 +𝜌 + almost matches the balancedness of SAM after sufficient convergence. Being balance also reveals that the benefit of SAM can come from optimization, which is a perspective typically ignored in literature. + + + +(a) threshold of balancedness +ℬ +¯ +𝑡 +𝜌 + (b) relation with regularization +Figure 2:Implicit regularization of SAM on NOP +𝔼 +⁢ +[ +‖ +𝐱𝐲 +⊤ +− +( +𝐀 ++ +𝛼 +⁢ +𝐍 +) +‖ +2 +] +, where +𝛼 + controls SNR. (a) the threshold of balancedness +ℬ +¯ +𝑡 +𝜌 + in Corollary 1; (b) implicit vs. explicit regularization. + +Noisy data have stronger impact on balancedness. Although our discussions extend to more general problems, for simplicity we consider the example in Fig. 2 (a), i.e., +𝔼 +⁢ +[ +‖ +𝐱𝐲 +⊤ +− +( +𝐀 ++ +𝛼 +⁢ +𝐍 +) +‖ +2 +] +, where +𝐀 + is ground truth; +𝐍 + is data noise; and +𝛼 + determines SNR. For this problem, noisy data directly lead to noisy gradients. It can be seen in Fig. 2 (a) that smaller SNR coincides with faster decreasing of +| +ℬ +𝑡 +| +. To explain such a data-responsive behavior in implicit regularization, Theorem 2 states that balancedness changes largely when the difference of +‖ +𝐠 +𝐲 +𝑡 +‖ + and +‖ +𝐠 +𝐱 +𝑡 +‖ + is large. Since +𝔼 +⁢ +[ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +] +∝ +𝛼 +2 + if assuming elements of +𝐍 + to be iid unit Gaussian variables, it thus implies that a small SNR (large +𝛼 +) offers large regularization on balancedness. + +Extension to LoRA (multi-layer two-variable NOP). For LoRA, the objective is to minimize +𝐷 + blocks of variables simultaneously, i.e., +min +⁡ +𝔼 +𝜉 +⁢ +[ +𝑓 +𝜉 +⁢ +( +{ +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +} +𝑙 += +1 +𝐷 +) +] +. It is established in Theorem 5 in appendix that SAM cultivates balancedness in a layer-wise fashion, i.e., the magnitude of +ℬ +𝑡 +, +𝑙 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) + cannot be large for each +𝑙 +. However, the +| +d +⁢ +ℬ +𝑡 +, +𝑙 +/ +d +⁢ +𝑡 +| + can be +𝒪 +⁢ +( +𝐷 +) + times smaller than Theorem 2 in the worst case because of the additional variables. + +Validation of IR on modern architectures. Going beyond the infinitesimally small step size, we adopt +𝜂 += +0.1 + on modern language models to validate our theoretical findings. We consider finetuning a RoBERTa-large with LoRA for few-shot learning tasks. More details can be found later in Section 6.1. Balancedness of SAM and SGD on different layers in various datasets are plotted in Fig. 3. SAM has a clear trend of promoting balancedness, aligning well with our theoretical predictions. + +4SAM for Overparametrized Problems + +Next, we focus on SAM’s implicit regularization on OP (1b). Overparametrization enables SAM to have stronger regularization on balancedness. Subscripts in +𝑓 +𝑜 + and +𝐿 +𝑜 + are omitted for convenience. SAM’s per iteration update for OP can be summarized as + + + +𝐱 +~ +𝑡 += +𝐱 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐲 +𝑡 +, + +𝐲 +~ +𝑡 += +𝐲 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐱 +𝑡 + +(7a) + + +𝐠 +𝐱 +~ +𝑡 += +𝑓 +𝑡 +′ +⁢ +( +𝐱 +~ +𝑡 +⊤ +⁢ +𝐲 +~ +𝑡 +) +⁢ +𝐲 +~ +𝑡 +, + +𝐠 +𝐲 +~ +𝑡 += +𝑓 +𝑡 +′ +⁢ +( +𝐱 +~ +𝑡 +⊤ +⁢ +𝐲 +~ +𝑡 +) +⁢ +𝐱 +~ +𝑡 + +(7b) + + +𝐱 +𝑡 ++ +1 += +𝐱 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐱 +~ +𝑡 +, + +𝐲 +𝑡 ++ +1 += +𝐲 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐲 +~ +𝑡 + +(7c) + +where +𝑢 +𝑡 +:= +sgn +⁢ +( +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +) +/ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +; +𝑓 +𝑡 + and +𝑓 +𝑡 +′ + denote the loss, stochastic gradient on minibatch +ℳ +𝑡 +, respectively. Different from NOP, SAM has stronger regularization on balancedness, where +| +ℬ +𝑡 +| + decreases whenever the norm of stochastic gradient is large. To see this, it is convenient to define +𝒞 +𝑡 +:= +| +‖ +𝐱 +𝑡 +‖ +− +‖ +𝐲 +𝑡 +‖ +| +. Note that +𝒞 +𝑡 +≤ +2 +⁢ +| +ℬ +𝑡 +| +. + +Theorem 3. + +Consider +𝜂 +→ +0 + for (7). The limiting flow of SAM on OP ensures a decreasing magnitude of +ℬ +𝑡 + whenever +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +⋅ +𝒞 +𝑡 +> +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +. Moreover, the speed of decrease can be lower- and upper- bounded as + + +𝜌 +⁢ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +⋅ +𝒞 +𝑡 +− +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +≤ +| +𝑑 +⁢ +ℬ +𝑡 +𝑑 +⁢ +𝑡 +| +≤ +𝜌 +⁢ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +⁢ +2 +⁢ +| +ℬ +𝑡 +| ++ +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +. + + +Given +𝜌 +→ +0 + and sufficiently noisy data, Theorem 3 implies that +| +ℬ +𝑡 +| +→ +0 +. Moreover, Theorem 3 also states that the regularization power on balancedness is related to both gradient norm and balancedness itself. The elbow-shaped curve of +| +ℬ +𝑡 +| + in Fig. 1 (b) demonstrates that the regularization power is reducing, as both gradient norm and balancedness shrink over time. + +Figure 3:Implicit regularization of SAM on LoRA. We consider few shot learning with LoRA on a RoBERTa-large. For datasets RTE, SST-5, and MNLI, 1st, 12th and 24th query layers’ +2 +⁢ +| +ℬ +𝑡 +, +𝑙 +| + are plotted, respectively. The layers are chosen to represent early, middle, and final stages of RoBERTa. The averaged +ℬ +¯ +𝑡 +, +𝑙 +𝜌 + in Corollary 1 is +0.37 +, +0.21 +, and +0.29 +, respectively. + +Noisy data have stronger impact on balancedness. As shown in Fig. 1 (b), balancedness is promoted faster on problems with lower SNR. This data-responsive behavior can be already seen from Theorem 3, because +| +d +⁢ +ℬ +𝑡 +/ +d +⁢ +𝑡 +| + is directly related with +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +, and +𝔼 +⁢ +[ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +] + is clearly larger when data are more noisy. In other words, SAM exploits noisy data for possible optimization merits from balancedness (see discussions in Sec. 2.2). Overall, the implicit regularization on balancedness aligns well with the empirical observations in presence of data anomalies (Wang et al., 2023; Sherborne et al., 2023), where SAM outperforms SGD by a large margin. + +Extension to +𝑚 +-sharpness. +𝑚 +-sharpness is a variant of SAM suitable for distributed training. It is observed to empirically improve SAM’s performance (Foret et al., 2021). +𝑚 +-sharpness evenly divides minibatch +ℳ +𝑡 + into +𝑚 + disjoint subsets, i.e., +{ +𝑓 +𝑡 +, +𝑗 +} +𝑗 += +1 +𝑚 +, and perform SAM update independently on each subset; see (38) in appendix. It turns out that +𝑚 +-sharpness can also be explained using balancedness. With formal proofs in Apdx. C.3, the IR of +𝑚 +-sharpness amounts to substitute +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| + in Theorem 3 with +1 +𝑚 +⁢ +∑ +𝑗 += +1 +𝑚 +| +𝑓 +𝑡 +, +𝑗 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +. This means that the regularization on balancedness from +𝑚 +-sharpness is more profound than vanilla SAM, because +1 +𝑚 +⁢ +∑ +𝑗 += +1 +𝑚 +| +𝑓 +𝑡 +, +𝑗 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +≥ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +. + +Finally, we connect balancedness with sharpness on local minima of OP. + +Lemma 1. + +Let +𝒲 +∗ += +{ +( +𝐱 +, +𝐲 +) +| +𝐱 +⊤ +⁢ +𝐲 += +𝑤 +, +𝑓 +′ +⁢ +( +𝑤 +) += +0 +, +𝑓 +′′ +⁢ +( +𝑤 +) +> +0 +} + be non-empty. For the OP problem (1b), minimizing sharpness within +𝒲 +∗ + is equivalent to finding +ℬ += +0 + in +𝒲 +∗ +. + +This link showcases that by studying balancedness we can also obtain the implicit regularization on sharpness for free. A concurrent work also links balancedness with sharpness (the largest eigenvalue) for some one-hidden layer neural networks (Singh and Hofmann, 2024). Compared with (Wen et al., 2023a), this is achieved with less assumptions and simplified analyses. More importantly, balancedness enables us to cope with arbitrary batchsize, to explain SAM’s stronger regularization with noisy data, and to extend results to +𝑚 +-sharpness. + +5Implicit Regularization Made Explicit + +Next, insights from our theoretical understanding of SAM are leveraged to build practical tools. We adopt LoRA (Hu et al., 2022) as our major numerical benchmark for scale-invariant problems given its prevalence in practice. More diverse examples on both OP and NOP can be found in Apdx. A.3. Compared to full parameter-tuning, LoRA is more economical in terms of memory not only for finetuning, but also for serving multiple downstream tasks. LoRA and its variants are actively developed and well welcomed by the community; see e.g., HuggingFace’s PEFT codebase.2 + +5.1Overview of LoRA + +Given a pretrained model with frozen weight +𝐖 +𝑙 +∈ +ℝ +𝑑 +1 +× +𝑑 +2 + on a particular layer +𝑙 +, the objective of LoRA is to find low rank matrices +𝐗 +𝑙 +∈ +ℝ +𝑑 +1 +× +𝑟 +, and +𝐘 +𝑙 +∈ +ℝ +𝑑 +2 +× +𝑟 + with +𝑟 +≪ +min +⁡ +{ +𝑑 +1 +, +𝑑 +2 +} + such that the loss is minimized for a downstream task, i.e., + + +min +{ +𝐗 +𝑙 +, +𝐘 +𝑙 +} +𝑙 +⁡ +ℒ +⁢ +( +{ +𝐖 +𝑙 ++ +𝐗 +𝑙 +⁢ +𝐘 +𝑙 +⊤ +} +𝑙 +) +. + +(8) + +LoRA enjoys parameter efficiency for finetuning thanks to the low-rank matrices +𝐗 +𝑙 + and +𝐘 +𝑙 +. For instance, it only requires 0.8M trainable parameters to finetune a 355M-parameter RoBERTa-large (Hu et al., 2022). The outer product of +𝐗 +𝑙 + and +𝐘 +𝑙 + induces scale-invariance, and the number of variables renders it NOP. The downside of LoRA, on the other hand, is the drop on test performance due to the parsimony on trainable parameters. Unbalancedness is also unavoidable for LoRA, due to the need of initializing at +𝐗 +𝑙 +∼ +𝒩 +⁢ +( +0 +, +𝜎 +2 +) +, +𝐘 +𝑙 += +𝟎 +; see an example of RoBERTa-large in Fig. 3. The unbalancedness leads to instability of LoRA when finetuning RoBERTa on datasets SST-2 and MNLI; see more details in Apdx. D.4. + +Integrating SAM with LoRA is a case with mutual benefits – LoRA reduces the additional memory requirement of SAM, while SAM not only overcomes the distributional shift in finetuning (Zhou et al., 2022), but also mitigates the possible inefficiency associated with LoRA’s unbalancedness. + +5.2Balancedness-Aware Regularization (BAR) + +However, directly applying SAM variants on LoRA exhibits two concerns: i) SAM doubles computational cost due to the need of two gradients; and ii) additional efforts are required to integrate SAM with gradient accumulation and low-precision training (HuggingFace,), which are common techniques for memory and runtime efficiency in large-scale finetuning. Note that concern i) is annoying given the size of language models, especially in setups involving model parallelism. + +Our balancedness-aware regularization (BAR) is a highly efficient approach to address both concerns, and it fixes the accuracy drop of LoRA relative to full-parameter finetuning. BAR is also the first efficient SAM variant derived from implicit regularization. The key observation for our algorithm design is that SAM’s implicit regularization on balancedness can be achieved with an explicit regularizer +𝛼 +𝑡 +⁢ +| +𝐱 +⊤ +⁢ +𝐱 +− +𝐲 +⊤ +⁢ +𝐲 +| +. This regularizer originates from matrix sensing; see e.g., (Tu et al., 2016; Ge et al., 2017). For OP, choosing +𝛼 +𝑡 +:= +𝒪 +⁢ +( +| +𝑓 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +/ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +) + recovers SAM’s dynamic on +ℬ +𝑡 + up to an error of +𝒪 +⁢ +( +𝜌 +2 +) +; cf. Lemma 2 in appendix. By ignoring this error, it can be seen that +ℬ +𝑡 + decreases when +‖ +𝐱 +𝑡 +‖ +≥ +‖ +𝐲 +𝑡 +‖ +. Following this dynamic, we regulate balancedness based on whether +‖ +𝐱 +𝑡 +‖ +≥ +‖ +𝐲 +𝑡 +‖ +. The resultant approach is termed as overparamterized BAR (oBAR) to reflect its source in OP. + +On the other hand, because LoRA is NOP inherently, we take inspiration from Theorem 2 – dropping the term +𝒜 +𝑡 + and mimicking dynamics of SAM. In particular, we regulate the objective with +𝛼 +𝑡 +⁢ +( +𝐱 +⊤ +⁢ +𝐱 +− +𝐲 +⊤ +⁢ +𝐲 +) + if +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +< +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +; otherwise +𝛼 +𝑡 +⁢ +( +𝐲 +⊤ +⁢ +𝐲 +− +𝐱 +⊤ +⁢ +𝐱 +) +. The resultant approach is termed as nBAR. A graphical illustration can be found in Fig. 2 (b). It can be observed that nBAR shares similar performance as SAM on NOP. Both nBAR and oBAR can be implemented in the same manner as weight decay, and their detailed steps are summarized in Algs. 2 and 3, respectively. + +Another benefit of BAR, in additional to the lightweight computation, is that it can be applied individually on each LoRA layer. As previously discussed (cf. Theorem 5), the number of layers has a negative impact on balancedness. By overcoming this “curse of multi-layer”, BAR can induce better test performance over SAM. + +Schedule of +𝛼 +𝑡 +. In both nBAR and oBAR, one can employ a decreasing scheduler for +𝛼 +𝑡 + for algorithmic flexibility. This is motivated by the fact that for both NOP and OP problems, the implicit regularization of SAM is less powerful after sufficient balancedness or near optimal. Commonly adopted cosine and linear schedules work smoothly. + +Algorithm 2 nBAR +1:Initialize: learning rate +{ +𝜂 +𝑡 +} +, regularization coefficient +{ +𝛼 +𝑡 +} +2:for  +𝑡 += +0 +, +… +, +𝑇 +− +1 + do +3:     Get stochastic gradient +𝐠 +𝐱 +𝑡 + and +𝐠 +𝐲 +𝑡 +4:     if  +‖ +𝐠 +𝐱 +𝑡 +‖ +≥ +‖ +𝐠 +𝐲 +𝑡 +‖ + then +5:          +𝐱 +𝑡 +← +( +1 ++ +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐱 +𝑡 +6:          +𝐲 +𝑡 +← +( +1 +− +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐲 +𝑡 +7:     else +8:          +𝐱 +𝑡 +← +( +1 +− +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐱 +𝑡 +9:          +𝐲 +𝑡 +← +( +1 ++ +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐲 +𝑡 +10:     end if +11:     Optimizer update (via Adam or SGD) +12:end for + Algorithm 3 oBAR +1:Initialize: learning rate +{ +𝜂 +𝑡 +} +, regularization coefficient +{ +𝛼 +𝑡 +} +2:for  +𝑡 += +0 +, +… +, +𝑇 +− +1 + do +3:     Get stochastic gradient +𝐠 +𝐱 +𝑡 + and +𝐠 +𝐲 +𝑡 +4:     if  +‖ +𝐱 +𝑡 +‖ +≥ +‖ +𝐲 +𝑡 +‖ + then +5:          +𝐱 +𝑡 +← +( +1 +− +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐱 +𝑡 +6:          +𝐲 +𝑡 +← +( +1 ++ +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐲 +𝑡 +7:     else +8:          +𝐱 +𝑡 +← +( +1 ++ +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐱 +𝑡 +9:          +𝐲 +𝑡 +← +( +1 +− +𝛼 +𝑡 +⁢ +𝜂 +𝑡 +) +⁢ +𝐲 +𝑡 +10:     end if +11:     Optimizer update (via Adam or SGD) +12:end for +6Numerical Experiments + +To demonstrate the effectiveness of BAR, numerical experiments are conducted on various deep learning tasks using language models (LMs). Bold and underlined numbers are used to highlight the best and second best performance, respectively. More experimental details can be found in Apdx. D. Code is available at https://github.com/BingcongLi/BAR. + +6.1Few-shot Learning with RoBERTa-large and OPT-1.3B +Table 1:Few shot learning on RoBERTa (355M). +† + denotes results reported by (Malladi et al., 2023) +RoBERTa SST-2 SST-5 SNLI MNLI RTE TREC avg ( +↑ +) +LoRA 91.1±0.8 52.3±2.9 84.3 +± +0.3 + 78.1±1.3 77.5±2.3 96.6±1.0 80.0 +LoRA-SAM 92.2±0.4 54.2±2.0 85.5±0.7 78.7±1.0 80.6±4.3 96.7±0.2 81.3 +LoRA-oBAR 91.5±0.9 54.5±2.7 84.9±0.5 78.3±2.2 79.7±2.0 96.7±0.5 80.9 +LoRA-nBAR 91.4±0.5 55.0±2.0 84.9±1.4 78.1±0.2 81.0±1.0 96.7±1.0 81.2 +Zero-Shot† 79.0 35.5 50.2 48.8 51.4 32.0 49.5 +Table 2:Runtime of BAR (normalized to LoRA, 1x) on OPT-1.3B. SAM relies on FP32 for stability. LoRA and BAR adopt FP16 training since this is the default choice for large models. nBAR and oBAR share similar runtime, hence reported together. +runtime ( +↓ +) SST-2 CB RTE COPA ReCoRD SQuAD +LoRA-SAM 4.43x 3.34x 4.10x 3.28x 4.35x 3.54x +LoRA-BAR 1.05x 1.03x 1.04x 1.05x 1.04x 1.03x +Table 3:Performance of BAR for few shot learning using OPT-1.3B. +OPT-1.3B SST-2 CB RTE COPA ReCoRD SQuAD avg ( +↑ +) +Prefix 92.9±1.0 71.6±3.0 65.2±2.6 73.0±1.0 69.7±1.0 82.1±1.4 75.8 +LoRA 93.1±0.2 72.6±3.7 69.1±4.8 78.0±0.0 70.8±1.0 81.9±1.8 77.6 +LoRA-SAM 93.5±0.5 74.3±1.0 70.6±2.7 78.0±0.0 70.9±1.2 83.0±0.7 78.4 +LoRA-oBAR 93.6±0.6 75.6±4.5 70.4±4.8 78.0±0.0 70.9±0.8 82.5±0.5 78.5 +LoRA-nBAR 93.7±0.7 79.8±4.4 70.5±2.4 78.0±0.0 71.0±1.0 82.3±1.8 79.2 +Zero-Shot 53.6 39.3 53.1 75.0 70.2 27.2 53.1 +Table 4:Finetuning RoBERTa (355M) with BAR. Results marked with +† + are taken from (Hu et al., 2022), and those with +∗ + refer to Adapter +P + in (Hu et al., 2022). +RoBERTa # para STS-B RTE MRPC CoLA QQP avg ( +↑ +) +FT† 355M 92.4 86.6 90.9 68.0 90.2 85.6 +Adapter∗ 0.8M 91.9 +± +0.4 80.1 +± +2.9 89.7 +± +1.2 67.8 +± +2.5 91.7 +± +0.2 84.2 +LoRA 0.8M 92.4 +± +0.1 88.2 +± +0.6 89.6 +± +0.5 64.8 +± +1.4 91.4 +± +0.1 85.3 +LoRA-oBAR 0.8M 92.6 +± +0.1 88.7 +± +0.2 90.3 +± +0.9 65.1 +± +1.0 91.6 +± +0.1 85.7 +LoRA-nBAR 0.8M 92.6 +± +0.2 89.2 +± +1.3 90.3 +± +0.4 65.6 +± +1.2 91.6 +± +0.1 85.9 + +The first task to consider is few-shot learning with LoRA (Malladi et al., 2023), where the goal is to finetune a language model with a small training set. We follow the settings in (Malladi et al., 2023), and choose the backbones as RoBERTa-large, a masked LM with 355M parameters, and OPT-1.3B, an autoregressive LM (Liu et al., 2019; Zhang et al., 2022). + +Results of the proposed oBAR and nBAR on RoBERTa-large are summarized in Table 1. As indicated by the zero-shot performance, the distributional shift between finetuning and pretraining datasets is obvious. This is a natural setting suitable for SAM and BAR. The averaged test accuracy is improved by +0.9 + and +1.2 + via oBAR and nBAR, respectively. The performance of nBAR is close to SAM. Moreover, BAR saves 74% additional runtime of SAM; see more details in Table 7 in the appendix. + +The proposed nBAR and oBAR perform even better when scaling up to OPT-1.3B. BAR reduces the overhead of SAM by more than +95 +% + because of its compatibility with FP16 training; see Table 2. Note that applying FP16 directly with SAM leads to underflow; see more in Apdx. D. This signifies the flexibility of BAR over SAM when scaling to large problems, as FP16 is the default choice for LMs. Prefix tuning (Li and Liang, 2021) is also included as a benchmark for comparisons on test performance. We report F1 score for SQuAD and accuracy for other datasets in Table 3. The averaged improvement over LoRA is +0.9 + and +1.6 + from oBAR and nBAR, respectively, both outperforming SAM. We conjecture that the performance gap between SAM and BAR comes from their different effectiveness in regularizing balancedness. Balancedness of a particular layer is decreasing slower in SAM due to multiple layers, as shown in Theorem 5, while BAR promotes balancedness faster as it can be applied individually on each LoRA layer. Comparing the absolute improvement for RoBERTa-large (355M) and OPT-1.3B, it is conjectured that BAR has more potential for larger models, and the verification is left for future due to hardware constraints. + +6.2Finetuning with RoBERTa-large + +Having demonstrated the power of BAR in few-shot learning, we then apply it to finetune RoBERTa-large with LoRA. The results can be found in Table 4. It can be observed that nBAR and oBAR improve the performance of LoRA and prefix tuning (Li and Liang, 2021) on most of tested datasets. On average, oBAR leads to a gain of +0.4 +, and nBAR raises the test performance by +0.6 +. BAR thereby fills the gap of test performance between LoRA (0.8M) and full-parameter (355M) finetuning. + +6.3Text Generation on GPT2-medium + +Lastly, we consider BAR on a text-generation problem using GPT2-medium, a model with 345M parameters. Results on WebNLG (Gardent et al., 2017) are reported in Table 5. It can be seen that oBAR matches the performance of prefix tuning, while nBAR achieves the best BLEU score. + +Table 5:Finetuning GPT2 (345M) with BAR on WebNLG. Results of prefix tuning and full-parameter finetuning are obtained from (Hu et al., 2022). +GPT2 FT∗ Prefix∗ LoRA LoRA-oBAR LoRA-nBAR +# param 354M 0.35M 0.35M 0.35M 0.35M +BLEU ( +↑ +) 46.5 55.1 54.99 +± +0.24 55.15 +± +0.19 55.20 +± +0.16 +7Discussions + +This work provides theoretical and empirical evidence on the implicit regularization of SAM for both scale-invariant NOP and OP problems. Balancedness, as an alternative to commonly adopted sharpness, is employed as the metric to capture global and data-responsive behaviors of SAM. We find that i) SAM promotes variables to have (relatively) balanced norms; and ii) noisy data have stronger impact on balancedness. Lastly, we explicify the implicit regularization as a data-driven regularizer to foster the design of a computationally efficient SAM variant, termed BAR. The effectiveness of BAR is demonstrated using various tasks on RoBERTa-large, GPT2 and OPT. BAR saves +95 +% + overhead of SAM and enhances the accuracy of LoRA to the level of full-parameter finetuning. + +Limitation and Future directions. + +Our approach, BAR, is best applied on scale-invariant modules in neural networks. Finetuning language models with LoRA, as a popular option in practice, is a setting naturally suitable for our approach. However, our approach does not apply for linear models, e.g., logistic regression. Regarding future directions, an interesting one is whether SAM has other forms of implicit regularization beyond balancedness and sharpness. The exploration of other scale-invariant architectures beyond LoRA, e.g., the softmax function in attention, is also deferred to future work. + +Acknowledgements + +We thank anonymous reviewers for their suggestions. BL is supported by Swiss National Science Foundation (SNSF) Project Funding No. 200021-207343. LZ gratefully acknowledges funding by the Max Planck ETH Center for Learning Systems (CLS). NH is supported by ETH research grant funded through ETH Zurich Foundations and SNSF Project Funding No. 200021-207343. + +References +Abbas et al. (2022) Momin Abbas, Quan Xiao, Lisha Chen, Pin-Yu Chen, and Tianyi Chen.Sharp-MAML: Sharpness-aware model-agnostic meta learning.In Proc. Int. Conf. Machine Learning, pages 10–32. PMLR, 2022. +Agarwala and Dauphin (2023) Atish Agarwala and Yann Dauphin.SAM operates far from home: eigenvalue regularization as a dynamical phenomenon.In Proc. Int. Conf. Machine Learning, pages 152–168. PMLR, 2023. +Ahn et al. (2023) Kwangjun Ahn, Sébastien Bubeck, Sinho Chewi, Yin Tat Lee, Felipe Suarez, and Yi Zhang.Learning threshold neurons via edge of stability.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Andriushchenko and Flammarion (2022) Maksym Andriushchenko and Nicolas Flammarion.Towards understanding sharpness-aware minimization.In Proc. Int. Conf. Machine Learning, pages 639–668. PMLR, 2022. +Arora et al. (2018) Sanjeev Arora, Nadav Cohen, and Elad Hazan.On the optimization of deep networks: Implicit acceleration by overparameterization.In Proc. Int. Conf. Machine Learning, pages 244–253. PMLR, 2018. +Arora et al. (2019a) Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu.A convergence analysis of gradient descent for deep linear neural networks.In Proc. Int. Conf. Learning Represention, 2019a. +Arora et al. (2019b) Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.Implicit regularization in deep matrix factorization.In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019b. +Arora et al. (2019c) Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.In Proc. Int. Conf. Machine Learning, pages 322–332. PMLR, 2019c. +Arora et al. (2022) Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi.Understanding gradient descent on the edge of stability in deep learning.In Proc. Int. Conf. Machine Learning, pages 948–1024. PMLR, 2022. +Bahri et al. (2022) Dara Bahri, Hossein Mobahi, and Yi Tay.Sharpness-aware minimization improves language model generalization.In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 7360–7371, 2022. +Barrett and Dherin (2021) David Barrett and Benoit Dherin.Implicit gradient regularization.In Proc. Int. Conf. Learning Represention, 2021. +Bartlett et al. (2018) Peter Bartlett, Dave Helmbold, and Philip Long.Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks.In Proc. Int. Conf. Machine Learning, pages 521–530. PMLR, 2018. +Bartlett et al. (2023) Peter Bartlett, Philip Long, and Olivier Bousquet.The dynamics of sharpness-aware minimization: Bouncing across ravines and drifting towards wide minima.J. Mach. Learn. Res., 24(316):1–36, 2023. +Bottou et al. (2018) Léon Bottou, Frank E Curtis, and Jorge Nocedal.Optimization methods for large-scale machine learning.SIAM Review, 60(2):223–311, 2018. +Bowman et al. (2015) Samuel Bowman, Gabor Angeli, Christopher Potts, and Christopher D Manning.A large annotated corpus for learning natural language inference.In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 632–642, 2015. +Cer et al. (2017) Daniel Cer, Mona Diab, Eneko Agirre, Iñigo Lopez-Gazpio, and Lucia Specia.SemEval-2017 task 1: Semantic textual similarity-multilingual and cross-lingual focused evaluation.In Proc. Int. Workshop Semant. Eval., pages 1–14. ACL, 2017. +Chaudhari et al. (2017) Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.Entropy-SGD: Biasing gradient descent into wide valleys.In Proc. Int. Conf. Learning Represention, 2017. +Chen et al. (2022) Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong.When vision transformers outperform ResNets without pre-training or strong data augmentations.In Proc. Int. Conf. Learning Represention, 2022. +Chen et al. (2024) Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia.Long-LoRA: Efficient fine-tuning of long-context large language models.In Proc. Int. Conf. Learning Represention, 2024. +Chen et al. (2023) Zixiang Chen, Junkai Zhang, Yiwen Kou, Xiangning Chen, Cho-Jui Hsieh, and Quanquan Gu.Why does sharpness-aware minimization generalize better than SGD?In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Dai et al. (2023) Yan Dai, Kwangjun Ahn, and Suvrit Sra.The crucial role of normalization in sharpness-aware minimization.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +De Marneffe et al. (2019) Marie-Catherine De Marneffe, Mandy Simons, and Judith Tonhauser.The CommitmentBank: Investigating projection in naturally occurring discourse.Proc. Sinn und Bedeutung, 23(2):107–124, 2019. +De Sa et al. (2015) Christopher De Sa, Christopher Re, and Kunle Olukotun.Global convergence of stochastic gradient descent for some non-convex matrix problems.In Proc. Int. Conf. Machine Learning, pages 2332–2341. PMLR, 2015. +Dettmers et al. (2023) Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer.QLoRA: Efficient finetuning of quantized LLMs.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Dinh et al. (2017) Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.Sharp minima can generalize for deep nets.In Proc. Int. Conf. Machine Learning, pages 1019–1028. PMLR, 2017. +Dolan and Brockett (2005) Bill Dolan and Chris Brockett.Automatically constructing a corpus of sentential paraphrases.In Proc. Int. Workshop Paraphrasing, 2005. +Du et al. (2022a) Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Y. F. Tan.Efficient sharpness-aware minimization for improved training of neural networks.In Proc. Int. Conf. Learning Represention, 2022a. +Du et al. (2022b) Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Y. F. Tan, and Joey Tianyi Zhou.Sharpness-aware training for free.In Proc. Adv. Neural Info. Processing Systems, 2022b. +Du et al. (2018) Simon S Du, Wei Hu, and Jason D Lee.Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced.In Proc. Adv. Neural Info. Processing Systems, volume 31, 2018. +Dziugaite and Roy (2017) Gintare Karolina Dziugaite and Daniel M. Roy.Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.In Proc. Conf. Uncerntainty in Artif. Intel., 2017. +Foret et al. (2021) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.Sharpness-aware minimization for efficiently improving generalization.In Proc. Int. Conf. Learning Represention, 2021. +Gardent et al. (2017) Claire Gardent, Anastasia Shimorina, Shashi Narayan, and Laura Perez-Beltrachini.The WebNLG challenge: Generating text from RDF data.In Proc. Int. Conf. Nat. Lang. Gener., pages 124–133. ACL, 2017. +Ge et al. (2017) Rong Ge, Chi Jin, and Yi Zheng.No spurious local minima in nonconvex low rank problems: A unified geometric analysis.In Proc. Int. Conf. Machine Learning, pages 1233–1242. PMLR, 2017. +Gidel et al. (2019) Gauthier Gidel, Francis Bach, and Simon Lacoste-Julien.Implicit regularization of discrete gradient dynamics in linear neural networks.In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019. +Gonon et al. (2024) Antoine Gonon, Nicolas Brisebarre, Elisa Riccietti, and Rémi Gribonval.A path-norm toolkit for modern networks: consequences, promises and challenges.In Proc. Int. Conf. Learning Represention, 2024. +Houlsby et al. (2019) Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly.Parameter-efficient transfer learning for NLP.In Proc. Int. Conf. Machine Learning, pages 2790–2799. PMLR, 2019. +Hu et al. (2022) Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen.LoRA: Low-rank adaptation of large language models.In Proc. Int. Conf. Learning Represention, 2022. +(38) HuggingFace.Gradient accumulation.URL https://huggingface.co/docs/accelerate/en/usage_guides/gradient_accumulation. +Izmailov et al. (2018) Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry P. Vetrov, and Andrew Gordon Wilson.Averaging weights leads to wider optima and better generalization.In Proc. Conf. Uncerntainty in Artif. Intel., pages 876–885, 2018. +Jastrzębski et al. (2017) Stanisław Jastrzębski, Zachary Kenton, Devansh Arpit, Nicolas Ballas, Asja Fischer, Yoshua Bengio, and Amos Storkey.Three factors influencing minima in SGD.arXiv:1711.04623, 2017. +Ji and Telgarsky (2019) Ziwei Ji and Matus Telgarsky.Gradient descent aligns the layers of deep linear networks.In Proc. Int. Conf. Learning Represention, 2019. +Jiang et al. (2023) Weisen Jiang, Hansi Yang, Yu Zhang, and James Kwok.An adaptive policy to employ sharpness-aware minimization.In Proc. Int. Conf. Learning Represention, 2023. +Jiang et al. (2020) Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.Fantastic generalization measures and where to find them.In Proc. Int. Conf. Learning Represention, 2020. +Keskar et al. (2016) Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.On large-batch training for deep learning: Generalization gap and sharp minima.In Proc. Int. Conf. Learning Represention, 2016. +Kim et al. (2022) Minyoung Kim, Da Li, Shell Xu Hu, and Timothy M. Hospedales.Fisher SAM: Information geometry and sharpness aware minimisation.In Proc. Int. Conf. Machine Learning, pages 11148–11161, 2022. +Kopiczko et al. (2024) Dawid Jan Kopiczko, Tijmen Blankevoort, and Yuki M Asano.VeRA: Vector-based random matrix adaptation.In Proc. Int. Conf. Learning Represention, 2024. +Kwon et al. (2021) Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi.ASAM: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.In Proc. Int. Conf. Machine Learning, pages 5905–5914. PMLR, 2021. +Li and Giannakis (2023) Bingcong Li and Georgios B Giannakis.Enhancing sharpness-aware optimization through variance suppression.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Li and Liang (2021) Xiang Lisa Li and Percy Liang.Prefix-tuning: Optimizing continuous prompts for generation.In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 4582–4597, 2021. +Li et al. (2022) Zhiyuan Li, Tianhao Wang, and Sanjeev Arora.What happens after SGD reaches zero loss? – A mathematical framework.In Proc. Int. Conf. Learning Represention, 2022. +Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov.RoBERTa: A robustly optimized BERT pretraining approach.arXiv preprint arXiv:1907.11692, 2019. +Liu et al. (2022) Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You.Towards efficient and scalable sharpness-aware minimization.In Proc. Conf. Computer Vision and Pattern Recognition, pages 12350–12360, 2022. +Loshchilov and Hutter (2019) Ilya Loshchilov and Frank Hutter.Decoupled weight decay regularization.In Proc. Int. Conf. Learning Represention, 2019. +Lyu and Li (2020) Kaifeng Lyu and Jian Li.Gradient descent maximizes the margin of homogeneous neural networks.In Proc. Int. Conf. Learning Represention, 2020. +Malladi et al. (2023) Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora.Fine-tuning language models with just forward passes.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Mi et al. (2022) Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao.Make sharpness-aware minimization stronger: A sparsified perturbation approach.In Proc. Adv. Neural Info. Processing Systems, volume 35, 2022. +Nesterov (2004) Yurii Nesterov.Introductory lectures on convex optimization: A basic course, volume 87.Springer Science & Business Media, 2004. +Neyshabur et al. (2015) Behnam Neyshabur, Russ R Salakhutdinov, and Nati Srebro.Path-SGD: Path-normalized optimization in deep neural networks.In Proc. Adv. Neural Info. Processing Systems, volume 28, 2015. +Neyshabur et al. (2017) Behnam Neyshabur, Srinadh Bhojanapalli, David Mcallester, Nathan Srebro, and Nati Srebro.Exploring generalization in deep learning.In Proc. Adv. Neural Info. Processing Systems, volume 30, pages 5947–5956, 2017. +Neyshabur et al. (2018) Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro.A PAC-bayesian approach to spectrally-normalized margin bounds for neural networks.In Proc. Int. Conf. Learning Represention, 2018. +Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang.SQuAD: 100,000+ questions for machine comprehension of text.In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 2383–2392, 2016. +Rajpurkar et al. (2018) Pranav Rajpurkar, Robin Jia, and Percy Liang.Know what you don’t know: Unanswerable questions for SQuAD.In Proc. Conf. Assoc. Comput. Linguist. Meet., pages 784–789, 2018. +Roemmele et al. (2011) Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon.Choice of plausible alternatives: An evaluation of commonsense causal reasoning.In AAAI Spring Symposium Series, 2011. +Sheen et al. (2024) Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H Zhou.Implicit regularization of gradient flow on one-layer softmax attention.arXiv preprint arXiv:2403.08699, 2024. +Sherborne et al. (2023) Tom Sherborne, Naomi Saphra, Pradeep Dasigi, and Hao Peng.TRAM: Bridging trust regions and sharpness aware minimization.In Proc. Int. Conf. Learning Represention, 2023. +Si and Yun (2023) Dongkuk Si and Chulhee Yun.Practical sharpness-aware minimization cannot converge all the way to optima.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023. +Singh and Hofmann (2024) Sidak Pal Singh and Thomas Hofmann.Closed form of the hessian spectrum for some neural networks.In High-dimensional Learning Dynamics 2024: The Emergence of Structure and Reasoning, 2024. +Socher et al. (2013) Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Y Ng, and Christopher Potts.Recursive deep models for semantic compositionality over a sentiment treebank.In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 1631–1642, 2013. +Tahmasebi et al. (2024) Behrooz Tahmasebi, Ashkan Soleymani, Dara Bahri, Stefanie Jegelka, and Patrick Jaillet.A universal class of sharpness-aware minimization algorithms.arXiv preprint arXiv:2406.03682, 2024. +Tu et al. (2016) Stephen Tu, Ross Boczar, Max Simchowitz, Mahdi Soltanolkotabi, and Ben Recht.Low-rank solutions of linear matrix equations via procrustes flow.In Proc. Int. Conf. Machine Learning, pages 964–973. PMLR, 2016. +Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.Attention is all you need.In Proc. Adv. Neural Info. Processing Systems, volume 30, 2017. +Voorhees and Tice (2000) Ellen M Voorhees and Dawn M Tice.Building a question answering test collection.In Proc. Annu. Int. ACM SIGIR Conf. Res. Dev. Inf. Retr., pages 200–207, 2000. +Wang et al. (2019a) Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman.SuperGLUE: A stickier benchmark for general-purpose language understanding systems.In Proc. Adv. Neural Info. Processing Systems, volume 32, 2019a. +Wang et al. (2019b) Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman.GLUE: A multi-task benchmark and analysis platform for natural language understanding.In Proc. Int. Conf. Learning Represention, 2019b. +Wang et al. (2023) Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang.Sharpness-aware gradient matching for domain generalization.In Proc. Conf. Computer Vision and Pattern Recognition, pages 3769–3778, 2023. +Wang and Mao (2022) Ziqiao Wang and Yongyi Mao.On the generalization of models trained with SGD: Information-theoretic bounds and implications.In Proc. Int. Conf. Learning Represention, 2022. +Warstadt et al. (2019) Alex Warstadt, Amanpreet Singh, and Samuel R Bowman.Neural network acceptability judgments.Trans. Assoc. Comput. Linguist., 7:625–641, 2019. +Wen et al. (2023a) Kaiyue Wen, Tengyu Ma, and Z hiyuan Li.How does sharpness-aware minimization minimizes sharpness.In Proc. Int. Conf. Learning Represention, 2023a. +Wen et al. (2023b) Kaiyue Wen, Tengyu Ma, and Zhiyuan Li.Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization.In Proc. Adv. Neural Info. Processing Systems, volume 36, 2023b. +Williams et al. (2018) Adina Williams, Nikita Nangia, and Samuel R Bowman.A broad-coverage challenge corpus for sentence understanding through inference.In Proc. Conf. North Am. Chapter Assoc. Comput. Linguist., pages 1112–1122, 2018. +Woodworth et al. (2020) Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.Kernel and rich regimes in overparametrized models.In Proc. Annual Conf. Learning Theory, pages 3635–3673. PMLR, 2020. +Wu et al. (2020) Dongxian Wu, Shu-Tao Xia, and Yisen Wang.Adversarial weight perturbation helps robust generalization.In Proc. Adv. Neural Info. Processing Systems, volume 33, pages 2958–2969, 2020. +Xia et al. (2024) Wenhan Xia, Chengwei Qin, and Elad Hazan.Chain of LoRA: Efficient fine-tuning of language models via residual learning.arXiv preprint arXiv:2401.04151, 2024. +Zhang et al. (2023a) Qingru Zhang, Minshuo Chen, Alexander Bukharin, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao.Adaptive budget allocation for parameter-efficient fine-tuning.In Proc. Int. Conf. Learning Represention, 2023a. +Zhang et al. (2023b) Ruipeng Zhang, Ziqing Fan, Jiangchao Yao, Ya Zhang, and Yanfeng Wang.Domain-inspired sharpness aware minimization under domain shifts.In Proc. Int. Conf. Learning Represention, 2023b. +Zhang et al. (2018) Sheng Zhang, Xiaodong Liu, Jingjing Liu, Jianfeng Gao, Kevin Duh, and Benjamin Van Durme.ReCoRD: Bridging the gap between human and machine commonsense reading comprehension.arXiv preprint arXiv:1810.12885, 2018. +Zhang et al. (2022) Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al.OPT: Open pre-trained transformer language models.arXiv preprint arXiv:2205.01068, 2022. +Zhao et al. (2022) Yang Zhao, Hao Zhang, and Xiuyuan Hu.Penalizing gradient norm for efficiently improving generalization in deep learning.In Proc. Int. Conf. Machine Learning, pages 26982–26992, 2022. +Zhou et al. (2022) Wenxuan Zhou, Fangyu Liu, Huan Zhang, and Muhao Chen.Sharpness-aware minimization with dynamic reweighting.In Proc. Conf. Empir. Methods Nat. Lang. Process., pages 5686–5699, 2022. +Zhuang et al. (2022) Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan, and Ting Liu.Surrogate gap minimization improves sharpness-aware training.In Proc. Int. Conf. Learning Represention, 2022. + +Supplementary Document for +“Implicit Regularization of Sharpness-Aware Minimization +for Scale-Invariant Problems” + +Appendix AMissing Details +A.1Broad Impact + +The theories and approaches are applicable across various scenarios. The proposed algorithmic tool simplifies finetuning language models, improves performance of downstream tasks, and consumes less resource compared to SAM. For tasks such as sentiment classification, our approach facilitates real world systems such as recommendation by improving accuracy. However, caution is advised when the downstream tasks of language models involve generation. For these tasks, users should thoroughly review generated content and consider to implement gating methods to ensure safety and trustworthiness. + +A.2More on Related Work + +Sharpness and generalization. Sharpness is observed to relate with generalization of SGD in deep learning (Keskar et al., 2016). It is found that sharpness varies with the ratio between learning rate and batchsize in SGD (Jastrzębski et al., 2017). Large scale experiments also indicate sharpness-based measures align with generalization in practical scenarios (Jiang et al., 2020; Chen et al., 2022). Theoretical understandings on generalization error using sharpness-related metrics can be found in e.g., (Dziugaite and Roy, 2017; Neyshabur et al., 2017; Wang and Mao, 2022). There is a large body of literature exploring sharpness for improved generalization. Entropy SGD leverages local entropy in search of a flat valley (Chaudhari et al., 2017). A similar approach as SAM is also developed in (Wu et al., 2020) while putting more emphases on adversarial robustness. Stochastic weight averaging is proposed for finding flatter minima in (Izmailov et al., 2018). It is shown later in (Wen et al., 2023b) that the interplay between sharpness and generalization subtly depends on data distributions and model architectures, and there are unveiled reasons beyond sharpness for the benefit of SAM. + +SAM variants. Although SAM is successful in various deep learning tasks, it can be improved further by leveraging local geometry in a fine-grained manner. For example, results in (Zhao et al., 2022; Barrett and Dherin, 2021) link SAM with gradient norm penalization. Zhuang et al. (2022) optimize sharpness gap and training loss jointly. A more accurate manner to solve inner maximization in SAM is developed in (Li and Giannakis, 2023). SAM and its variants are also widely applied to domain generalization problems; see e.g., (Zhang et al., 2023b; Wang et al., 2023). + +Other perspectives for SAM. The convergence of SAM is comprehensively studied in (Si and Yun, 2023). Agarwala and Dauphin (2023) focus on the edge-of-stability-like behavior of unnormalized SAM on quadratic problems. Dai et al. (2023) argue that the normalization in SAM, i.e., line 5 of Alg. 1, is critical. Sharpness measure is generalized to any functions of Hessian in (Tahmasebi et al., 2024). However, even the generalized sharpness cannot provide implicit regularization for simple functions such as +ℎ +⁢ +( +𝑥 +, +𝑦 +) += +𝑥 +⁢ +𝑦 +, because the Hessian is the same for all +( +𝑥 +, +𝑦 +) +. In addition, when Hessian is negative definite, some of the generalized sharpness measures (e.g., determinate of Hessian) may not be necessarily meaningful. + +Implicit regularization. The regularization effect can come from optimization algorithms rather than directly from the regularizer in objective functions. This type of the behavior is termed as implicit regularization or implicit bias of the optimizer. The implicit regularization of (S)GD is studied from multiple perspectives, such as margin (Ji and Telgarsky, 2019; Lyu and Li, 2020), kernel (Arora et al., 2019c), and Hessian (Li et al., 2022; Arora et al., 2022). Initialization can also determine the implicit regularization (Woodworth et al., 2020). Most of these works explore the overparametrization regime. + +LoRA and parameter-efficient finetuning. LoRA (Hu et al., 2022), our major numerical benchmark, is an instance of parameter-efficient finetuning (PEFT) approaches. PEFT reduces the resource requirement for large language models on various downstream tasks, at the cost of possible accuracy drops on test performance. The latter, together with the transfer learning setup jointly motivate the adoption of SAM. Other commonly adopted PEFT methods include, e.g., adapters (Houlsby et al., 2019) and prefix tuning (Li and Liang, 2021). There are also various efforts to further improve LoRA via adaptivity (Zhang et al., 2023a), chaining (Xia et al., 2024), aggressive parameter saving (Kopiczko et al., 2024), low-bit training (Dettmers et al., 2023), and modifications for long-sequences (Chen et al., 2024). Most of these efforts are orthogonal to BAR proposed in this work. + +A.3Additional Applications of Scale-Invariant Problems in Deep Learning + +Attention in transformers. Attention is one of the backbones of modern neural networks (Vaswani et al., 2017). Given the input +𝐃 +, attention can be written as + + +min +𝐐 +, +𝐊 +, +𝐕 +⁡ +softmax +⁢ +( +1 +𝛼 +⁢ +𝐃𝐐𝐊 +⊤ +⁢ +𝐃 +⊤ +) +⁢ +𝐃𝐕 + +(9) + +where +{ +𝐐 +, +𝐊 +, +𝐕 +} + are query, key, and value matrices to be optimized. This is a scale-invariant problem because scaling +{ +𝐐 +, +𝐊 +} + does not modify the objective function. Considering the number of variables, the optimization of +{ +𝐐 +, +𝐊 +} + is considered as OP. + +Two-layer linear neural networks. This problem is a simplified version of two-layer ReLU neural nets, and its objective can be defined as + + +𝑓 +⁢ +( +𝐖 +1 +, +𝐖 +2 +) += +1 +2 +⁢ +𝔼 +( +𝐚 +, +𝐛 +) +⁢ +[ +‖ +𝐖 +1 +⁢ +𝐖 +2 +⁢ +𝐚 +− +𝐛 +‖ +2 +] +. + +(10) + +This is usually adopted as an example for overparametrization, and can be extended to deeper linear neural networks; see e.g., (Arora et al., 2019a). Moreover, it is known that the optimization for such problem is quite challenging, and GD can fail to converge if +𝐖 +1 + and +𝐖 +2 + are not initialized with balancedness (Arora et al., 2019a). An extension of (10) is two-layer ReLU networks, which are widely adopted in theoretical frameworks to understand the behavior of neural networks. ReLU networks are scale-invariant, but only when the scaling factor is positive. + +Other examples. For ResNets, two-variable scale-invariant submodules also include affine BatchNorm and the subsequent convolutional layer. For transformers, scale-invariant submodules besides attention include LayerNorm and its subsequent linear layer. + +A.4SAM Pays More Attention to Difficult Examples + +Testing example for NOP. The problem presented below is adopted in Fig. 1 (a) and Fig. 2 for visualization of SAM’s behavior on NOP. We consider a special case of problem (1a), where the goal is to fit (rank-1) matrices by minimizing + + +𝑓 +𝑛 +⁢ +( +𝐱 +, +𝐲 +) += +𝔼 +𝜉 +⁢ +[ +‖ +𝐱𝐲 +⊤ +− +( +𝐀 ++ +𝛼 +⁢ +𝐍 +𝜉 +) +‖ +2 +] + +(11) + +where +𝐀 +∈ +ℝ +3 +× +3 +:= +diag +⁢ +[ +0.5 +, +0 +, +0 +] + and +𝐍 +𝜉 +∈ +ℝ +3 +× +3 + denote the ground truth and Gaussian noise, respectively; and +𝛼 + controls the SNR. Here we choose +𝐍 +𝜉 +:= +diag +⁢ +[ +1.0 +, +0.8 +, +0.5 +] +⁢ +𝐔 +𝜉 +, where entries of +𝐔 +𝜉 + are unit Gaussian random variables. + +In our simulation of Fig. 1 (a), we set the step size to be +𝜂 += +10 +− +4 + and the total number of iterations as +𝑇 += +10 +5 + for both SGD and SAM. Parameter +𝜌 + is chosen as +0.1 + for SAM. For both algorithms, initialization is +𝐱 +0 += +[ +0.2 +, +− +0.1 +, +0.3 +] +⊤ + and +𝐲 +0 += +− +3 +⁢ +𝐱 +0 +. Note that we choose a small step size to mimic the settings of our theorems. + +Testing example for OP. The problem presented below is adopted in Fig. 1 (b) for visualization of SAM on OP. A special case of problem (1b) is considered with objective function + + +𝑓 +𝑜 +⁢ +( +𝐱 +, +𝐲 +) += +𝔼 +𝜉 +⁢ +[ +‖ +𝐱 +⊤ +⁢ +𝐲 +− +( +𝑎 ++ +𝛼 +⁢ +𝑛 +𝜉 +) +‖ +2 +] + +(12) + +where +𝑎 +∈ +ℝ + and +𝑛 +𝜉 +∈ +ℝ + denote the ground truth and Gaussian noise, respectively. We choose +𝑎 += +0.5 + and +𝑛 +𝜉 + as a unit Gaussian random variable. Here, +𝛼 + controls the SNR of this problem. + +In our simulation of Fig. 1 (b), we set +𝜂 += +10 +− +4 + and +𝑇 += +10 +5 + for both SGD and SAM. Parameter +𝜌 + is set as +0.2 + for SAM. For both algorithms, initialization is +𝐱 +0 += +[ +0.2 +, +− +0.1 +, +0.3 +] +⊤ + and +𝐲 +0 += +− +3 +⁢ +𝐱 +0 +. + +A.5Scale-Invariance in OP + +Scale-invariance also bothers OP in the same fashion as it burdens NOP. For completeness, the scale-invariance of OP can be verified by + + +𝑓 +𝑜 +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) += +𝑓 +𝑜 +⁢ +( +( +𝛼 +⁢ +𝐱 +) +⊤ +⁢ +( +1 +𝛼 +⁢ +𝐲 +) +) +, +∀ +𝛼 +≠ +0 +. + +(13) + +An optimizer has to determine +𝛼 + for OP despite it does not influence objective value. Hence, scaling is redundant for OP. + +Similar to NOP, the (stochastic) gradient of OP is not scale-invariant. In particular, given a minibatch of data +ℳ +, the stochastic gradient for OP (1b) can be written as + + +𝐠 +𝐱 += +1 +| +ℳ +| +⁢ +[ +∑ +𝜉 +∈ +ℳ +( +𝑓 +𝑜 +𝜉 +) +′ +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) +] +⁢ +𝐲 +, +𝐠 +𝐲 += +1 +| +ℳ +| +⁢ +[ +∑ +𝜉 +∈ +ℳ +( +𝑓 +𝑜 +𝜉 +) +′ +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) +] +⁢ +𝐱 +. + +(14) + +Consequently, being balance also brings optimization benefits for OP as discussed previously in Section 2.2 . + +A.6BAR in Detail +Figure 4:The value of +𝑓 +⁢ +( +𝑥 +, +𝑦 +) +. Once SGD reaches the dotted line, i.e., the hard constraint +| +𝑥 +| += +| +𝑦 +| +, it can only converge to a saddle point +( +0 +, +0 +) +. + +BAR is inspired jointly from the balancedness-promoting regularizer +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +| + and the dynamics of SAM on both NOP and OP. The implementation of BAR is similar as weight decay in AdamW (Loshchilov and Hutter, 2019). + +Here we use nBAR as an example. If ignoring +𝒜 +𝑡 + in Theorem 2, it can be seen that +ℬ +𝑡 + for NOP decreases whenever +‖ +𝐠 +𝐱 +𝑡 +‖ +< +‖ +𝐠 +𝐲 +𝑡 +‖ +. In other words, the balancedness of SAM is driven by the difference between the gradient norms at +𝐱 +𝑡 + and +𝐲 +𝑡 +. nBAR mimics this and triggers balancedness when stochastic gradients +𝐠 +𝐱 +𝑡 + and +𝐠 +𝐲 +𝑡 + are not balanced; see Alg. 2. + +Finally, we illustrate more on the reasons for employing regularization in OP rather than posing +‖ +𝐱 +𝑡 +‖ += +‖ +𝐲 +𝑡 +‖ + as a hard constraint or initializing in a balanced manner, i.e., +‖ +𝐱 +0 +‖ += +‖ +𝐲 +0 +‖ +. First, it is quite clear that +‖ +𝐱 +‖ += +‖ +𝐲 +‖ + is a nonconvex set and how to project on such a set is still debatable. Second, the ‘symmetry’ associated with the scale-invariant problems does not always favor this constraint. For the purpose of graphical illustration, we consider a +2 +-dimensional example +𝑓 +⁢ +( +𝑥 +, +𝑦 +) += +30000 +⁢ +( +𝑥 +⁢ +𝑦 +− +0.005 +) +2 +. It is quite clear that the objective is symmetric regarding the line +𝑥 += +− +𝑦 +, which satisfies +| +𝑥 +| += +| +𝑦 +| +; see Fig. 4. However, it is not hard to see that SGD can never leave +𝑥 += +− +𝑦 + once it reaches this line via a hard constraint or initialized on this line. In other words, directly adding +‖ +𝐱 +‖ += +‖ +𝐲 +‖ + as a constraint can trap the algorithm at saddle points. This symmetric pattern is even more complicated in high dimension, i.e., symmetry over multiple lines or hyperplanes. Hence, one should be extremely careful about this hard constraint, and regularization is a safer and more practical choice. + +Appendix BMissing Proofs for NOP +B.1Proof of Theorem 1 +Proof. + +For notational convenience, we let +𝐆 +𝑡 +:= +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +) +. Then, we have that + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 += +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +d +⁢ +𝐱 +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐠 +𝐱 +𝑡 += +− +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐆 +𝑡 +⁢ +𝐲 +𝑡 +. + + +Similarly, we have that + + +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 += +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +d +⁢ +𝐲 +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +𝐠 +𝐲 +𝑡 += +− +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +𝑡 +. + + +Combining these two inequalities, we arrive at + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 +− +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 += +0 +. + + +The proof is thus completed. ∎ + +B.2Extension to Stochastic Normalized Gradient Descent (SNGD) + +Next, we extend Theorem 1 to SNGD, whose updates can be written as + + +𝐱 +𝑡 ++ +1 += +𝐱 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐱 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +, + +𝐲 +𝑡 ++ +1 += +𝐲 +𝑡 +− +𝜂 +⁢ +𝐠 +𝐲 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +. + +(15) +Theorem 4. + +When applying SNGD (15) on NOP problem (1a), the limiting flow with +𝜂 +→ +0 + guarantees that +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 += +‖ +𝐱 +0 +‖ +2 +− +‖ +𝐲 +0 +‖ +2 + for all +𝑡 +> +0 +. In other words, +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +0 + holds. + +Proof. + +For notational convenience, we let +𝐆 +𝑡 +:= +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +) +. Then, we have that + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 += +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +d +⁢ +𝐱 +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐠 +𝐱 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 += +− +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐆 +𝑡 +⁢ +𝐲 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +. + + +Similarly, we have that + + +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 += +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +d +⁢ +𝐲 +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +𝐠 +𝐲 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 += +− +2 +⁢ +𝐲 +𝑡 +⊤ +⁢ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +. + + +Combining these two inequalities, we arrive at + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 +− +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 += +0 +. + + +The proof is thus completed. ∎ + +B.3Proof of Theorem 2 +Proof. + +Denote +𝐆 +𝑡 += +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +) + and +𝐆 +~ +𝑡 += +∇ +𝑓 +𝑡 +⁢ +( +𝐱 +~ +𝑡 +⁢ +𝐲 +~ +𝑡 +⊤ +) + for notational convenience. Following SAM updates in (4) and setting +𝜂 +→ +0 +, we have that + + +d +⁢ +𝐱 +𝑡 +d +⁢ +𝑡 += +− +𝐆 +~ +𝑡 +⁢ +( +𝐲 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +𝑡 +) +, +d +⁢ +𝐲 +𝑡 +d +⁢ +𝑡 += +− +𝐆 +~ +𝑡 +⊤ +⁢ +( +𝐱 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +⁢ +𝐲 +𝑡 +) +. + + +This gives that + + + +1 +2 +⁢ +d +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) +d +⁢ +𝑡 + += +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +𝐲 +𝑡 +⊤ +⁢ +𝐆 +~ +𝑡 +⊤ +⁢ +𝐆 +𝑡 +⁢ +𝐲 +𝑡 +− +𝐱 +𝑡 +⊤ +⁢ +𝐆 +~ +𝑡 +⁢ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +𝑡 +] + +(16a) + + += +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +] ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +𝐲 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +− +𝐱 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⁢ +𝐠 +𝐲 +𝑡 +] +⏟ +:= +𝒜 +𝑡 +. + +(16b) + +The second term in (16b) is +𝒜 +𝑡 + in Theorem 2. Next, we give upper bound on +| +𝒜 +𝑡 +| +. Using Assumption 1, we have that + + +‖ +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +‖ + +≤ +𝐿 +⁢ +‖ +𝐱 +~ +𝑡 +⁢ +𝐲 +~ +𝑡 +⊤ +− +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +‖ + + += +𝐿 +⁢ +‖ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +( +𝐱 +𝑡 +⁢ +𝐠 +𝐲 +𝑡 +⊤ ++ +𝐠 +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +) ++ +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐠 +𝐱 +𝑡 +⁢ +𝐠 +𝐲 +𝑡 +⊤ +‖ + + +≤ +( +𝑎 +) +𝐿 +⁢ +𝜌 +⁢ +‖ +𝐱 +𝑡 +⁢ +𝐠 +𝐲 +𝑡 +⊤ ++ +𝐠 +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ +‖ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 ++ +𝐿 +⁢ +𝜌 +2 +⁢ +‖ +𝐠 +𝐱 +𝑡 +⁢ +𝐠 +𝐲 +𝑡 +⊤ +‖ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 + + +≤ +( +𝑏 +) +𝐿 +⁢ +𝜌 +⁢ +( +‖ +𝐱 +𝑡 +‖ ++ +‖ +𝐲 +𝑡 +‖ +) ++ +𝐿 +⁢ +𝜌 +2 +2 += +𝒪 +⁢ +( +𝐿 +⁢ +𝜌 +) + + +where (a) uses the definition of +𝑢 +𝑡 +; (b) follows from +‖ +𝐚𝐛 +⊤ +‖ += +‖ +𝐚 +‖ +⁢ +‖ +𝐛 +‖ + and the finite convergence assumption. To bound +𝒜 +𝑡 +, we also have + + +𝜌 +⁢ +𝑢 +𝑡 +⁢ +| +𝐲 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +| + += +𝜌 +⁢ +| +𝐲 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +| +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +≤ +𝜌 +⁢ +| +𝐲 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +| +‖ +𝐠 +𝐱 +𝑡 +‖ + + +≤ +𝜌 +⁢ +‖ +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +‖ +⁢ +‖ +𝐲 +𝑡 +‖ += +𝒪 +⁢ +( +𝐿 +⁢ +𝜌 +2 +) + +(17) + +where the last line also uses the finite convergence. We can bound +𝜌 +⁢ +𝑢 +𝑡 +⁢ +| +𝐱 +𝑡 +⊤ +⁢ +( +𝐆 +~ +𝑡 +− +𝐆 +𝑡 +) +⁢ +𝐠 +𝐲 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) + in a similar manner. Combining (B.3) with (16b) gives the bound on +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) + . ∎ + +B.4Proof of Corollary 1 + +Here, we prove the formal version of Corollary 1. + +Corollary 2. + +Suppose that +‖ +𝐠 +𝐱 +𝑡 +‖ +> +0 + and +‖ +𝐠 +𝐲 +𝑡 +‖ +> +0 + and +𝜌 +→ +0 +, then there exists +ℬ +¯ +𝑡 + such that the magnitude of +ℬ +𝑡 + shrinks whenever +| +ℬ +𝑡 +| +> +ℬ +¯ +𝑡 +. + +Proof. + +Without loss of generality, we suppose that +ℬ +𝑡 +> +0 +, i.e., +‖ +𝐱 +𝑡 +‖ +> +‖ +𝐲 +𝑡 +‖ +> +0 +. Let +𝐱 +¯ +𝑡 + and +𝐲 +¯ +𝑡 + be the scaled version of +𝐱 +𝑡 + and +𝐲 +𝑡 + such that +‖ +𝐱 +¯ +𝑡 +‖ += +‖ +𝐲 +¯ +𝑡 +‖ + and +𝐱 +¯ +𝑡 +⁢ +𝐲 +¯ +𝑡 +⊤ += +𝐱 +𝑡 +⁢ +𝐲 +𝑡 +⊤ + are satisfied. This suggests that +𝐱 +𝑡 += +𝛼 +𝑡 +⁢ +𝐱 +¯ +𝑡 + and +𝐲 +𝑡 += +𝐲 +¯ +𝑡 +/ +𝛼 +𝑡 +, where +𝛼 +𝑡 += +‖ +𝐱 +𝑡 +‖ +/ +‖ +𝐲 +𝑡 +‖ +. Next, we show that whenever +ℬ +𝑡 + is large enough, we have that + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +𝜌 +⁢ +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 ++ +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) +< +0 +. + +(18) + +Since +𝜌 +→ +0 +, we only need to show that for some small +𝜖 += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +) +≥ +0 +, + + +‖ +𝐠 +𝐱 +𝑡 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +< +− +𝜖 +. + +(19) + +By the definition of +𝐠 +𝐱 +𝑡 +, +𝐠 +𝐲 +𝑡 + and +𝐱 +¯ +𝑡 +, +𝐲 +¯ +𝑡 +, we have that (19) can be rewritten as + + +𝛼 +𝑡 +2 +⁢ +‖ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +¯ +𝑡 +‖ +2 +− +‖ +𝐆 +𝑡 +⁢ +𝐲 +¯ +𝑡 +‖ +2 +/ +𝛼 +𝑡 +2 +𝛼 +𝑡 +2 +⁢ +‖ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +¯ +𝑡 +‖ +2 ++ +‖ +𝐆 +𝑡 +⁢ +𝐲 +¯ +𝑡 +‖ +2 +/ +𝛼 +𝑡 +2 +> +𝜖 +. + +(20) + +Note that the function +ℎ +⁢ +( +𝑧 +) +:= +( +𝑎 +⁢ +𝑧 +− +𝑏 +/ +𝑧 +) +/ +𝑎 +⁢ +𝑧 ++ +𝑏 +/ +𝑧 + is monotonically increasing in +𝑧 + when +𝑎 +, +𝑏 +> +0 + and +𝑧 +> +0 + as +ℎ +′ +⁢ +( +𝑧 +) += +( +𝑎 +2 +⁢ +𝑧 ++ +6 +⁢ +𝑎 +⁢ +𝑏 +/ +𝑧 ++ +𝑏 +2 +/ +𝑧 +3 +) +/ +( +2 +⁢ +( +𝑎 +⁢ +𝑧 ++ +𝑏 +/ +𝑧 +) +3 +/ +2 +) +> +0 +. This implies that +ℎ +⁢ +( +𝑧 +) +> +0 + when +𝑧 +> +𝑏 +/ +𝑎 +, and thus the condition in (20) can be satisfied for +𝜖 += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +) +→ +0 + when +𝛼 +𝑡 +2 +> +𝛼 +¯ +2 +, where +𝛼 +¯ +2 +:= +‖ +𝐆 +𝑡 +⁢ +𝐲 +¯ +𝑡 +‖ +/ +‖ +𝐆 +𝑡 +⊤ +⁢ +𝐱 +¯ +𝑡 +‖ +. This condition on +𝛼 +𝑡 + is equivalent to + + +ℬ +𝑡 + += +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + + += +1 +2 +⁢ +( +‖ +𝛼 +𝑡 +⁢ +𝐱 +¯ +𝑡 +‖ +2 +− +‖ +𝐲 +¯ +𝑡 +/ +𝛼 +𝑡 +‖ +2 +) + + +> +1 +2 +⁢ +( +‖ +𝛼 +¯ +⁢ +𝐱 +¯ +𝑡 +‖ +2 +− +‖ +𝐲 +¯ +𝑡 +/ +𝛼 +¯ +‖ +2 +) +. + + +Combining everything together, we have that +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 +< +0 + if + + +ℬ +𝑡 +> +ℬ +¯ +𝑡 +:= +1 +2 +⁢ +( +‖ +𝛼 +¯ +⁢ +𝐱 +¯ +𝑡 +‖ +2 +− +‖ +𝐲 +¯ +𝑡 +/ +𝛼 +¯ +‖ +2 +) +. + +(21) + +The proof is thus completed. We also note that in the case of +𝜌 +> +0 +, the same condition as (21) can be derived by obtaining the inverse function of +ℎ +⁢ +( +𝑧 +) + evaluated at +𝜖 += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +) +, and the corresponding +𝛼 +¯ +𝜌 + and +ℬ +¯ +𝑡 +𝜌 + can be defined similarly. ∎ + +B.5Extension to LoRA (layer-wise NOP problem) + +Let +𝑙 +∈ +{ +1 +, +2 +, +… +, +𝐷 +} + be the layer index. Denote +𝑓 +𝑡 + as the loss function on minibatch +ℳ +𝑡 +. To simplify the notation, we also let +𝐆 +𝑡 +, +𝑙 +:= +∇ +𝐱 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +⊤ +𝑓 +𝑡 +⁢ +( +{ +𝐱 +𝑡 +, +𝑙 +, +𝐲 +𝑡 +, +𝑙 +} +𝑙 +) +, +𝐆 +~ +𝑡 +, +𝑙 +:= +∇ +𝐱 +~ +𝑡 +, +𝑙 +⁢ +𝐲 +~ +𝑡 +, +𝑙 +⊤ +𝑓 +𝑡 +⁢ +( +{ +𝐱 +~ +𝑡 +, +𝑙 +, +𝐲 +~ +𝑡 +, +𝑙 +} +𝑙 +) +, and +𝑢 +𝑡 +:= +1 +/ +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +. The update of SAM for layer +𝑙 + can be written as + + + +𝐱 +~ +𝑡 +, +𝑙 += +𝐱 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +, + +𝐲 +~ +𝑡 +, +𝑙 += +𝐲 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +, +𝑙 +⊤ +⁢ +𝐱 +𝑡 +, +𝑙 + +(22a) + + +𝐠 +𝐱 +~ +𝑡 +, +𝑙 += +𝐆 +~ +𝑡 +, +𝑙 +⁢ +𝐲 +~ +𝑡 +, +𝑙 +, + +𝐠 +𝐲 +~ +𝑡 +, +𝑙 += +𝐆 +~ +𝑡 +, +𝑙 +⊤ +⁢ +𝐱 +~ +𝑡 +, +𝑙 + +(22b) + + +𝐱 +𝑡 ++ +1 +, +𝑙 += +𝐱 +𝑡 +, +𝑙 +− +𝜂 +⁢ +𝐠 +𝐱 +~ +𝑡 +, +𝑙 +, + +𝐲 +𝑡 ++ +1 +, +𝑙 += +𝐲 +𝑡 +, +𝑙 +− +𝜂 +⁢ +𝐠 +𝐲 +~ +𝑡 +, +𝑙 +. + +(22c) + +Refined assumption for LoRA. Direct translating Assumption 1 to our multi-layer setting gives + + +‖ +∇ +𝑓 +𝑡 +⁢ +( +{ +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +} +𝑙 +) +− +∇ +𝑓 +𝑡 +⁢ +( +{ +𝐚 +𝑙 +⁢ +𝐛 +𝑙 +⊤ +} +𝑙 +) +‖ +2 +≤ +𝐿 +2 +⁢ +∑ +𝑙 += +1 +𝐷 +‖ +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +− +𝐚 +𝑙 +⁢ +𝐛 +𝑙 +⊤ +‖ +2 +. + +(23) + +However, the above assumption is loose, and our proof only needs block-wise smoothness, i.e., + + +‖ +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +) +− +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +𝐚 +𝑙 +⁢ +𝐛 +𝑙 +⊤ +) +‖ +2 +≤ +𝐿 +^ +2 +⁢ +‖ +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +− +𝐚 +𝑙 +⁢ +𝐛 +𝑙 +⊤ +‖ +2 +, +∀ +𝑙 + +(24) + +where +∇ +𝑙 + refers to the gradient on +𝐱 +𝑙 +⁢ +𝐲 +𝑙 +⊤ +. It can be seen that +𝐷 +⁢ +𝐿 +^ +≥ +𝐿 +, but one can assume that +𝐷 +⁢ +𝐿 +^ +≈ +𝐿 + for intuitive understandings. + +Theorem 5. + +Suppose that block smoothness assumption in (24) holds. Consider the limiting flow of SAM in (22) with +𝜂 +→ +0 + and a sufficiently small +𝜌 +. Let +ℬ +𝑡 +, +𝑙 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) + and +ℬ +𝑡 += +∑ +𝑙 += +1 +𝐷 +ℬ +𝑡 +, +𝑙 +. For some +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +) +, SAM guarantees that + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +𝜌 +⁢ +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 ++ +𝒜 +𝑡 +. + +(25) + +Furthermore, for per layer balancedness it satisfies that for some +| +𝒜 +𝑡 +, +𝑙 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +) +. + + +d +⁢ +ℬ +𝑡 +, +𝑙 +d +⁢ +𝑡 += +𝜌 +⁢ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 ++ +𝒜 +𝑡 +, +𝑖 +. + +(26) + +Understanding Theorem 5. +𝒜 +𝑡 +, +𝑖 + and +𝒜 +𝑡 + are at the same order because of the possible unbalancedness among gradient norms for different layers. Comparing per layer balancedness +ℬ +𝑡 +, +𝑙 + with Theorem 2, it can be roughly estimate that the regularization power is +𝒪 +⁢ +( +𝐷 +) + times smaller in +ℬ +𝑡 +, +𝑙 +. This estimation comes from +𝐿 +^ +≈ +𝐿 +/ +𝐷 +, and the first term is also +𝒪 +⁢ +( +𝐷 +) + smaller than the same term in Theorem 2. In other words, the regularization on balancedness can be reduced by +𝒪 +⁢ +( +𝐷 +) + times in LoRA in the worst case, and the worst case comes from gradient unbalancedness among layers. + +Proof. + +Following (22) and setting +𝜂 +→ +0 +, we have that + + +d +⁢ +𝐱 +𝑡 +, +𝑙 +d +⁢ +𝑡 += +− +𝐆 +~ +𝑡 +, +𝑙 +⁢ +( +𝐲 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +, +𝑙 +⊤ +⁢ +𝐱 +𝑡 +, +𝑙 +) +, +d +⁢ +𝐲 +𝑡 +, +𝑙 +d +⁢ +𝑡 += +− +𝐆 +~ +𝑡 +, +𝑙 +⊤ +⁢ +( +𝐱 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝐆 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +) +. + + +This gives that + + + +d +⁢ +ℬ +𝑡 +, +𝑙 +d +⁢ +𝑡 + += +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +𝐲 +𝑡 +, +𝑙 +⊤ +⁢ +𝐆 +~ +𝑡 +, +𝑙 +⊤ +⁢ +𝐆 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +− +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +𝐆 +~ +𝑡 +, +𝑙 +⁢ +𝐆 +𝑡 +, +𝑙 +⊤ +⁢ +𝐱 +𝑡 +, +𝑙 +] + +(27a) + + += +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +] ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +[ +𝐲 +𝑡 +, +𝑙 +⊤ +⁢ +( +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +, +𝑙 +− +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +( +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +) +⁢ +𝐠 +𝐲 +𝑡 +, +𝑙 +] +⏟ +:= +𝒜 +𝑡 +, +𝑙 +. + +(27b) + +Proof for (25). Let +𝒜 +𝑡 +:= +∑ +𝑙 +𝒜 +𝑡 +, +𝑙 +. To start with, we have that + + +‖ +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +‖ + +≤ +𝐿 +^ +⁢ +‖ +𝐱 +~ +𝑡 +, +𝑙 +⁢ +𝐲 +~ +𝑡 +, +𝑙 +⊤ +− +𝐱 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +⊤ +‖ + + += +𝐿 +^ +⁢ +‖ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +( +𝐱 +𝑡 +, +𝑙 +⁢ +𝐠 +𝐲 +𝑡 +, +𝑙 +⊤ ++ +𝐠 +𝐱 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +⊤ +) ++ +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐠 +𝐱 +𝑡 +, +𝑙 +⁢ +𝐠 +𝐲 +𝑡 +, +𝑙 +⊤ +‖ + + +Next, based on finite convergence assumption, we have that + + +𝜌 +⁢ +𝑢 +𝑡 +⁢ +∑ +𝑙 += +1 +𝐷 +| +𝐲 +𝑡 +, +𝑙 +⊤ +⁢ +( +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +, +𝑙 +| + +(28) + + +≤ +∑ +𝑙 += +1 +𝐷 +𝒪 +⁢ +( +𝜌 +⁢ +𝑢 +𝑡 +⁢ +‖ +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +‖ +⋅ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) + + +≤ +( +𝑎 +) +∑ +𝑙 += +1 +𝐷 +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐿 +^ +⁢ +‖ +𝐱 +𝑡 +, +𝑙 +⁢ +𝐠 +𝐲 +𝑡 +, +𝑙 +⊤ ++ +𝐠 +𝐱 +𝑡 +, +𝑙 +⁢ +𝐲 +𝑡 +, +𝑙 +⊤ +‖ +⋅ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) + + +≤ +( +𝑏 +) +∑ +𝑙 += +1 +𝐷 +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐿 +^ +⁢ +( +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ ++ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) +⋅ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) + + += +𝜌 +2 +⁢ +𝐿 +^ +⋅ +𝒪 +⁢ +( +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) ++ +∑ +𝑙 += +1 +𝐷 +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +⁢ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +) + + += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +) + + +where in (a) we use the fact that +𝜌 + is chosen small; (b) uses finite convergence assumption and +‖ +𝐚𝐛 +⊤ +‖ += +‖ +𝐚 +‖ +⁢ +‖ +𝐛 +‖ +. Using similar arguments, we can bound +𝒜 +𝑡 += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +) +. + +Proof for (26). Next, we give upper bound on +| +𝒜 +𝑡 +, +𝑙 +| +. Using similar argument as (28), we have that + + +𝜌 +⁢ +𝑢 +𝑡 +⁢ +| +𝐲 +𝑡 +, +𝑙 +⊤ +⁢ +( +𝐆 +~ +𝑡 +, +𝑙 +− +𝐆 +𝑡 +, +𝑙 +) +⊤ +⁢ +𝐠 +𝐱 +𝑡 +, +𝑙 +| + +(29) + + +≤ +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐿 +^ +⁢ +( +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ ++ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) +⋅ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +) + + += +𝜌 +2 +⁢ +𝐿 +^ +⋅ +𝒪 +⁢ +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) ++ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +⁢ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +) +. + +(30) + +Using (29), we have that + + +| +𝒜 +𝑡 +, +𝑙 +| + +≤ +𝜌 +2 +⁢ +𝐿 +^ +⋅ +𝒪 +⁢ +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) ++ +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +⁢ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +∑ +𝑙 += +1 +𝐷 +( +‖ +𝐠 +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +) + + += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +) +. + + +The proof is is thus completed. ∎ + +Appendix CMissing Proofs for OP +C.1Unbalancedness of SGD in OP +Theorem 6. + +Applied SGD or SNGD on problem (1b), both of them ensure that +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 += +‖ +𝐱 +0 +‖ +2 +− +‖ +𝐲 +0 +‖ +2 + for all +𝑡 +> +0 +. In other words, +ℬ +𝑡 + keeps unchanged. + +Proof. + +We consider SGD and NSGD separately. + +SGD. It is straightforward to see that + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 += +− +2 +⁢ +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 += +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 +. + + +This completes the proof of SGD. + +NSGD. The gradient update of NSGD is + + +d +⁢ +𝐱 +𝑡 +d +⁢ +𝑡 += +− +𝐠 +𝐱 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +, +d +⁢ +𝐲 +𝑡 +d +⁢ +𝑡 += +− +𝐠 +𝐲 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 +. + +(31) + +Then we have that for NSGD, + + +d +⁢ +‖ +𝐱 +𝑡 +‖ +2 +d +⁢ +𝑡 += +− +2 +⁢ +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +‖ +𝐠 +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐠 +𝐲 +𝑡 +‖ +2 += +d +⁢ +‖ +𝐲 +𝑡 +‖ +2 +d +⁢ +𝑡 +. + + +This gives the result for SNGD. ∎ + +C.2Proof of Theorem 3 + +To prove this theorem, we first focus on the dynamic of SAM. + +Lemma 2. + +Suppose that Assumption 1 holds. Consider the limiting flow of SAM in (7) with +𝜂 +→ +0 +. Let +ℬ +𝑡 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + and +𝜌 + be small. Then, for some +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +, SAM guarantees + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝜌 +⁢ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⁢ +ℬ +𝑡 ++ +𝒜 +𝑡 +. + +(32) +Proof. + +For notational convenience, we write +𝑓 +𝑡 +′ +:= +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) + and +𝑓 +~ +𝑡 +′ +:= +𝑓 +𝑡 +′ +⁢ +( +𝐱 +~ +𝑡 +⊤ +⁢ +𝐲 +~ +𝑡 +) +. Using similar arguments as Theorem 2, we have that + + +1 +2 +⁢ +d +d +⁢ +𝑡 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + += +− +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝑓 +~ +𝑡 +′ +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + +(33) + + += +− +𝜌 +⁢ +sgn +⁢ +( +𝑓 +𝑡 +′ +) +⁢ +𝑓 +~ +𝑡 +′ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + + += +− +𝜌 +⁢ +| +𝑓 +𝑡 +′ +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + + ++ +𝜌 +⁢ +sgn +⁢ +( +𝑓 +𝑡 +′ +) +⁢ +( +𝑓 +𝑡 +′ +− +𝑓 +~ +𝑡 +′ +) +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) +⏟ +:= +𝒜 +𝑡 +. + + +Next we bound +| +𝒜 +𝑡 +| +. To start with, we have that + + +| +𝐱 +~ +𝑡 +⊤ +⁢ +𝐲 +~ +𝑡 +− +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +| + += +| +𝜌 +2 +⁢ +𝑢 +𝑡 +2 +⁢ +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +‖ +𝐱 +𝑡 +‖ +2 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +‖ +𝐲 +𝑡 +‖ +2 +| + +(34) + + +≤ +𝜌 +2 +⁢ +| +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 ++ +𝜌 +⁢ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 + + +≤ +𝜌 +2 +2 ++ +𝜌 +⁢ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +. + + +Using Assumption 1 and (34), we arrive at + + +| +𝑓 +𝑡 +′ +− +𝑓 +𝑡 +′ +~ +| +≤ +𝐿 +⁢ +| +𝐱 +~ +𝑡 +⊤ +⁢ +𝐲 +~ +𝑡 +− +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +⁢ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +) +. + +(35) + +Hence, we arrive at + + +| +𝒜 +𝑡 +| +≤ +𝜌 +⁢ +| +𝑓 +𝑡 +′ +− +𝑓 +𝑡 +′ +~ +| +⁢ +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +. + + +The proof is thus completed. ∎ + +Next, the proof of Theorem 3 is provided. + +Proof. + +Lemma 2 has already indicated the concentration of +ℬ +𝑡 + towards +0 +, if the magnitude of the first term is larger than +| +𝒜 +𝑡 +| +. To see this, notice that we can lower bound +2 +⁢ +| +ℬ +𝑡 +| +/ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 + by + + +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +| += +| +( +‖ +𝐱 +𝑡 +‖ ++ +‖ +𝐲 +𝑡 +‖ +) +⁢ +( +‖ +𝐱 +𝑡 +‖ +− +‖ +𝐲 +𝑡 +‖ +) +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +| +| +≥ +| +‖ +𝐱 +𝑡 +‖ +− +‖ +𝐲 +𝑡 +‖ +| += +𝒞 +𝑡 +. + +(36) + +Hence, long as +𝜌 +⁢ +| +𝑓 +𝑡 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +⋅ +𝒞 +𝑡 +> +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +, we have the first term dominating the dynamic of SAM, leading to contraction of +ℬ +𝑡 +. This completes the proof to the first part. + +Next we prove the second part, which is the lower- and upper- bound on +ℬ +𝑡 +. The lower bound can be seen from (36). For the upper bound, we have + + +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +| +≤ +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +| +| += +2 +⁢ +| +ℬ +𝑡 +| +. + +(37) + +Plugging (37) into (33) finishes the proof. ∎ + +C.3 +𝑚 +-sharpness for OP + +𝑚 +-sharpness is a variant of SAM that is empirically observed to improve generalization, and it is especially useful for distributed training on multiple GPUs (Foret et al., 2021). However, the reason behind the improved performance is not fully understood. (Andriushchenko and Flammarion, 2022) show that +𝑚 +-sharpness is more sparse-promoting for diagonal linear neural networks minimized via a quadratic loss. However, diagonal linear networks are not scale-invariant. + +For consistent notation with (7), we use +𝑓 +𝑡 +⁢ +( +⋅ +) + to denote the loss function on minibatch +ℳ +𝑡 +. In +𝑚 +-sharpness, the minibatch +ℳ +𝑡 + is divided into +𝑚 + disjoint subsets. Without loss of generality, we also assume that the minibatch is evenly divided. We denote the loss function on each subset as +𝑓 +𝑡 +, +𝑖 +, +𝑖 +∈ +{ +1 +, +2 +, +… +, +𝑚 +} +. Note that we have +1 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +𝑓 +𝑡 +, +𝑖 += +𝑓 +𝑡 +. With these definitions, the update of +𝑚 +-sharpness can be written as + + + +𝐱 +~ +𝑡 +, +𝑖 += +𝐱 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +, +𝑖 +⁢ +𝐲 +𝑡 +, + +𝐲 +~ +𝑡 +, +𝑖 += +𝐲 +𝑡 ++ +𝜌 +⁢ +𝑢 +𝑡 +, +𝑖 +⁢ +𝐱 +𝑡 + +(38a) + + +𝐠 +𝐱 +~ +𝑡 +, +𝑖 +𝑖 += +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +~ +𝑡 +, +𝑖 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑖 +) +⁢ +𝐲 +~ +𝑡 +, +𝑖 +, + +𝐠 +𝐲 +~ +𝑡 +, +𝑖 +𝑖 += +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +~ +𝑡 +, +𝑖 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑖 +) +⁢ +𝐱 +~ +𝑡 +, +𝑖 + +(38b) + + +𝐱 +𝑡 ++ +1 += +𝐱 +𝑡 +− +𝜂 +⁢ +1 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +𝐠 +𝐱 +~ +𝑡 +, +𝑖 +𝑖 +, + +𝐲 +𝑡 ++ +1 += +𝐲 +𝑡 +− +𝜂 +⁢ +1 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +𝐠 +𝐲 +~ +𝑡 +, +𝑖 +𝑖 +. + +(38c) + +where +𝑢 +𝑡 +, +𝑖 +:= +sgn +⁢ +( +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +) +/ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +. Comparing with the SAM update for OP in (7), the difference is that perturbed gradient is calculated on each +𝑓 +𝑡 +, +𝑖 +. Next, we analyze the dynamic of SAM with +𝑚 +-sharpness. + +Lemma 3. + +Suppose that Assumption 1 holds. Consider the limiting flow of SAM in (38) with +𝜂 +→ +0 +. Let +ℬ +𝑡 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + and +𝜌 + be small. Then, for some +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +) +, SAM guarantees that + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +− +2 +⁢ +𝜌 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +| +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⁢ +ℬ +𝑡 ++ +𝒜 +𝑡 +. + +(39) +Proof. + +For notational convenience, we write +𝑓 +𝑡 +, +𝑖 +′ +:= +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +) + and +𝑓 +~ +𝑡 +, +𝑖 +′ +:= +𝑓 +𝑡 +, +𝑖 +′ +⁢ +( +𝐱 +~ +𝑡 +, +𝑖 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑖 +) +. Then, we have that + + +1 +2 +⁢ +d +d +⁢ +𝑡 +⁢ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + += +− +𝜌 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +𝑢 +𝑡 +, +𝑖 +⁢ +𝑓 +~ +𝑡 +, +𝑖 +′ +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + +(40) + + += +− +𝜌 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +sgn +⁢ +( +𝑓 +𝑡 +, +𝑖 +′ +) +⁢ +𝑓 +~ +𝑡 +, +𝑖 +′ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + + += +− +𝜌 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +| +𝑓 +𝑡 +, +𝑖 +′ +| +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) + + ++ +𝜌 +𝑚 +⁢ +∑ +𝑖 += +1 +𝑚 +sgn +⁢ +( +𝑓 +𝑡 +, +𝑖 +′ +) +⁢ +( +𝑓 +𝑡 +, +𝑖 +′ +− +𝑓 +~ +𝑡 +, +𝑖 +′ +) +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +⋅ +( +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +) +⏟ +:= +𝒜 +𝑡 +, +𝑖 +. + + +Next, using (34) and Assumption 1, we have + + +| +𝑓 +𝑡 +, +𝑖 +′ +− +𝑓 +~ +𝑡 +, +𝑖 +′ +| +≤ +𝐿 +⁢ +| +𝐱 +~ +𝑡 +, +𝑖 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑖 +− +𝐱 +𝑡 +⊤ +⁢ +𝐲 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +⁢ +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +) +. + + +Hence, we can bound +| +𝒜 +𝑡 +, +𝑖 +| + as + + +| +𝒜 +𝑡 +, +𝑖 +| +≤ +| +𝑓 +𝑡 +, +𝑖 +′ +− +𝑓 +~ +𝑡 +, +𝑖 +′ +| +⁢ +| +‖ +𝐱 +𝑡 +‖ +2 +− +‖ +𝐲 +𝑡 +‖ +2 +‖ +𝐱 +𝑡 +‖ +2 ++ +‖ +𝐲 +𝑡 +‖ +2 +| += +𝒪 +⁢ +( +𝜌 +⁢ +𝐿 +⁢ +| +ℬ +𝑡 +| +) +. + + +The proof is thus completed by plugging +| +𝒜 +𝑡 +, +𝑖 +| + into (40). ∎ + +C.4Extension to Layer-wise OP + +We start with the notation. Let +𝑙 +∈ +{ +1 +, +2 +, +… +, +𝐷 +} + be the layer index. Denote +𝑓 +𝑡 + as the loss on minibatch +ℳ +𝑡 +. Let +𝑓 +𝑡 +, +𝑙 +′ +:= +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +{ +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +𝑡 +, +𝑙 +} +𝑙 +) +, i.e., the +𝑙 +-th entry of gradient (w.r.t. the variable +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +𝑡 +, +𝑙 +), +𝑓 +~ +𝑡 +, +𝑙 +′ +:= +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +{ +𝐱 +~ +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑙 +} +𝑙 +) +, and +𝑢 +𝑡 +:= +1 +/ +∑ +𝑙 += +1 +𝐷 +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +[ +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +] +. The update of SAM for layer +𝑙 + can be written as + + + +𝐱 +~ +𝑡 +, +𝑙 += +𝐱 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝑓 +𝑡 +, +𝑙 +′ +⁢ +𝐲 +𝑡 +, +𝑙 +, + +𝐲 +~ +𝑡 +, +𝑙 += +𝐲 +𝑡 +, +𝑙 ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝑓 +𝑡 +, +𝑙 +′ +⁢ +𝐱 +𝑡 +, +𝑙 +, + +(41a) + + +𝐠 +𝐱 +~ +𝑡 +, +𝑙 += +𝑓 +~ +𝑡 +, +𝑙 +′ +⁢ +𝐲 +~ +𝑡 +, +𝑙 +, + +𝐠 +𝐲 +~ +𝑡 +, +𝑙 += +𝑓 +~ +𝑡 +, +𝑙 +′ +⁢ +𝐱 +~ +𝑡 +, +𝑙 + +(41b) + + +𝐱 +𝑡 ++ +1 +, +𝑙 += +𝐱 +𝑡 +, +𝑙 +− +𝜂 +⁢ +𝐠 +𝐱 +~ +𝑡 +, +𝑙 +, + +𝐲 +𝑡 ++ +1 +, +𝑙 += +𝐲 +𝑡 +, +𝑙 +− +𝜂 +⁢ +𝐠 +𝐲 +~ +𝑡 +, +𝑙 +. + +(41c) + +Refined assumption for LoRA. Our proof only needs block-wise smoothness, i.e., + + +| +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +𝐱 +𝑙 +⊤ +⁢ +𝐲 +𝑙 +) +− +∇ +𝑙 +𝑓 +𝑡 +⁢ +( +𝐚 +𝑙 +⊤ +⁢ +𝐛 +𝑙 +) +| +2 +≤ +𝐿 +^ +2 +⁢ +| +𝐱 +𝑙 +⊤ +⁢ +𝐲 +𝑙 +− +𝐚 +𝑙 +⊤ +⁢ +𝐛 +𝑙 +| +2 +, +∀ +𝑙 +, + +(42) + +where +∇ +𝑙 + refers to the gradient on +𝐱 +𝑙 +⊤ +⁢ +𝐲 +𝑙 +. It can be seen that +𝐷 +⁢ +𝐿 +^ +≥ +𝐿 +, but one can assume that +𝐷 +⁢ +𝐿 +^ +≈ +𝐿 + for more clear intuition. + +Theorem 7. + +Suppose that block smoothness assumption in (42) holds. Consider the limiting flow of SAM in (41) with +𝜂 +→ +0 + and a sufficiently small +𝜌 +. Let +ℬ +𝑡 +, +𝑙 +:= +1 +2 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) + and +ℬ +𝑡 +max += +max +𝑙 +⁡ +| +ℬ +𝑡 +, +𝑙 +| +. For some +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +⁢ +ℬ +𝑡 +max +) +, SAM guarantees that + + +d +⁢ +ℬ +𝑡 +d +⁢ +𝑡 += +− +𝜌 +⁢ +∑ +𝑙 += +1 +𝐷 +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +∑ +𝑙 += +1 +𝐷 +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +[ +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +] ++ +𝒜 +𝑡 +. + +(43) + +Furthermore, for some +| +𝒜 +𝑡 +, +𝑙 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +⁢ +| +ℬ +𝑡 +, +𝑙 +| +) +, per layer balancedness satisfies that + + +d +⁢ +ℬ +𝑡 +, +𝑙 +d +⁢ +𝑡 += +− +𝜌 +⁢ +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +∑ +𝑙 += +1 +𝐷 +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +[ +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +] ++ +𝒜 +𝑡 +, +𝑖 +. + +(44) +Proof. + +Using a similar derivation as before, we have that + + +1 +2 +⁢ +d +d +⁢ +𝑡 +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) + += +− +𝜌 +⁢ +𝑢 +𝑡 +⁢ +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⋅ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) + + ++ +𝜌 +⁢ +𝑢 +𝑡 +⁢ +𝑓 +𝑡 +, +𝑙 +′ +⁢ +( +𝑓 +𝑡 +, +𝑙 +′ +− +𝑓 +~ +𝑡 +, +𝑙 +′ +) +⋅ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 +− +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) +⏟ +:= +𝒜 +𝑡 +, +𝑙 + + +Next, based on (42), we have that + + +| +𝑓 +𝑡 +, +𝑙 +′ +− +𝑓 +~ +𝑡 +, +𝑙 +′ +| +≤ +𝐿 +^ +⁢ +| +𝐱 +~ +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +~ +𝑡 +, +𝑙 +− +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +𝑡 +, +𝑙 +| +≤ +𝜌 +⁢ +𝐿 +^ +⁢ +𝑢 +𝑡 +⁢ +| +𝑓 +𝑡 +, +𝑙 +′ +| +⁢ +( +‖ +𝐱 +𝑡 +, +𝑙 +‖ +2 ++ +‖ +𝐲 +𝑡 +, +𝑙 +‖ +2 +) ++ +𝜌 +2 +⁢ +𝐿 +^ +⁢ +𝑢 +𝑡 +2 +⁢ +| +𝑓 +𝑡 +, +𝑙 +′ +| +2 +⁢ +| +𝐱 +𝑡 +, +𝑙 +⊤ +⁢ +𝐲 +𝑡 +, +𝑙 +| +. + + +Combining these two equations, and applying similar argument as Theorem 5, it is not difficult to arrive at +| +𝒜 +𝑡 +, +𝑖 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +⁢ +| +ℬ +𝑡 +, +𝑙 +| +) + and +| +𝒜 +𝑡 +| += +𝒪 +⁢ +( +𝜌 +2 +⁢ +𝐿 +^ +⁢ +ℬ +𝑡 +max +) +. ∎ + +C.5Proof of Lemma 1 +Proof. + +Within +𝒲 +∗ +, the Hessian on +( +𝐱 +, +𝐲 +) + can be calculated as +𝑓 +′′ +⁢ +( +𝐱 +⊤ +⁢ +𝐲 +) +⁢ +[ +𝐲 +⊤ +, +𝐱 +⊤ +] +⊤ +⁢ +[ +𝐲 +⊤ +, +𝐱 +⊤ +] +. The largest eigenvalue is +𝑓 +′′ +⁢ +( +𝑤 +) +⁢ +( +‖ +𝐱 +‖ +2 ++ +‖ +𝐲 +‖ +2 +) +. By the AM-GM inequality, it can be seen that the largest eigenvalue is minimized when +‖ +𝐱 +‖ += +‖ +𝐲 +‖ +, whose balancedness is +0 +. ∎ + +Appendix DMissing Experimental Details + +We mainly focus on finetuning LMs with LoRA. This setting naturally includes distributional shift – the finetuning dataset does not usually have the same distribution as the pretraining dataset as validated through zero-shot performance. All experiments are performed on a server with AMD EPYC 7742 CPUs and NVIDIA GeForce RTX 3090 GPUs each with 24GiB memory. All numerical results from Section 6 report test performance (e.g., accuracy, F1 scores, or BLEU scores) and the standard deviation across multiple runs. + +D.1Details on Datasets + +Our evaluations are carried out on commonly-used datasets in the literature. + +GLUE benchmark. GLUE is designed to provide a general-purpose evaluation of language understanding (Wang et al., 2019b). Those adopted in our work include MNLI (inference, (Williams et al., 2018)), SST-2 (sentiment analysis, (Socher et al., 2013)), MRPC (paraphrase detection, (Dolan and Brockett, 2005)), CoLA (linguistic acceptability (Warstadt et al., 2019)), QNLI (inference (Rajpurkar et al., 2018)), QQP3 (question-answering), RTE4 (inference), and STS-B (textual similarity (Cer et al., 2017)). These datasets are released under different permissive licenses. + +SuperGLUE benchmark. SuperGLUE (Wang et al., 2019a) is another commonly adopted benchmark for language understanding and is more challenging compared with GLUE. The considered datasets include CB (inference, (De Marneffe et al., 2019)), ReCoRD (multiple-choice question answering (Zhang et al., 2018)), COPA (question answering (Roemmele et al., 2011)). These datasets are released under different permissive licenses. + +WebNLG Challenge. This dataset is commonly used for data-to-text evaluation (Gardent et al., 2017). It has 22K examples in total with 14 distinct categories. Among them, 9 are seen during training, and the unseen training data are used to test the generalization performance. The dataset is released under license CC BY-NC-SA 4.0. + +Additional datasets. We also use SQuAD (question answering (Rajpurkar et al., 2016)) in our experiments, which is released under license CC BY-SA 4.0. Other datasets include TREC (topic classification (Voorhees and Tice, 2000)) and SNLI (inference (Bowman et al., 2015)). Both of them are licensed under CC BY-SA 4.0. + +D.2Details on Language Models + +We summarize the adopted language models in our evaluation. All model checkpoints are obtained from HuggingFace. + +RoBERTa-large. This is a +355 +M parameter model. The model checkpoint5 is released under the MIT license. + +OPT-1.3B. The model checkpoint6 is released under a non-commercial license. 7 + +GPT2-medium. This is a +345 +M parameter model. Its checkpoint8 is under MIT License. + +D.3Few-shot Learning with RoBERTa and OPT + +Experiments on RoBERTa-large. We follow the +𝑘 +-shot learning setup in (Malladi et al., 2023) and focus on classification tasks. The training set contains +𝑘 += +512 + samples per class while the test set has +1000 + samples. We also employ prompts for finetuning; where the adopted prompts are the same as those in (Malladi et al., 2023, Table 13). AdamW is adopted as the base optimizer, and hyperparameters are tuned from those in Table 6. Our experiments are averaged over +3 + random trials. The estimated runtime is about 5 minutes per dataset. + +Table 6:Hyperparameters used for few-shot learning with RoBERTa-large. +Hyper-parameters Values +LoRA +𝑟 + (rank) 8 +LoRA +𝛼 + 16 +# iterations 1000 +batchsize 16 +learning rate 1 +× +10 +− +4 +, 3 +× +10 +− +4 +, 5 +× +10 +− +4 + + +𝜌 + for SAM 0.05, 0.1, 0.2 + +𝜇 +0 + for BAR 0.5, 1.0, 2.0 +scheduler for BAR linear, cosine + +The per-iteration runtime on the SST-5 dataset of BAR, SAM, and the baseline optimizer are compared in Table 7. It can be seen that SAM is much more slower than the baseline approach, and BAR reduces 74% additional runtime of SAM, while achieving comparable accuracy. We believe that this runtime saving can be even larger with additional engineering efforts such as kernel fusion, which we leave for future work. This validates the computational efficiency of BAR. + +Table 7:Per-iteration runtime for finetuning RoBERTa-large on SST5. +SST5 baseline SAM BAR +time (s) 0.105 0.265 0.146 + +Experiments on OPT. For OPT-1.3B, we consider tasks from the SuperGLUE benchmark covering classification and multiple-choice. We also consider generation tasks on SQuAD. Following (Malladi et al., 2023), we randomly sample +1000 + data for training and the other +1000 + for testing. AdamW is adopted as base optimizer. The hyperparameters adopted are searched over values in Table 8. Estimated runtime is less than or around 10 minutes, depending on the dataset. + +If we directly apply FP16 training with SAM, underflow can happen if one does not take care of the gradient scaling on the two gradients calculated per iteration. This means that SAM is not flexible enough to be integrated with the codebase for large scale training, as FP16 is the default choice for finetuning LMs. We employ FP32 to bypass the issue with SAM. Consequently, the training speed is significantly slowed down; see a summary in Table 9. It further demonstrates the effectiveness of BAR for large scale-training. + +Overall, the results for few-shot learning indicate that given limited data, BAR can effectively improve generalization using significantly reduced computational resources relative to SAM. + +Table 8:Hyperparameters used for few-shot learning with OPT-1.3B. +Hyper-parameters Values +LoRA +𝑟 + (rank) 8 +LoRA +𝛼 + 16 +# iterations 1000 +batchsize 2, 4, 8 +learning rate 1 +× +10 +− +5 +, 1 +× +10 +− +4 +, 5 +× +10 +− +4 + + +𝜌 + for SAM 0.05, 0.1, 0.2 + +𝜇 +0 + for BAR 0.2, 0.5, 1.0, 2.0 +scheduler for BAR linear, cosine +Table 9:Per-iteration runtime for finetuning OPT-1.3B on RTE. +RTE baseline SAM BAR +precision FP16 FP32 FP16 +time (s) 0.1671 0.708 0.1731 +D.4Finetuning with RoBERTa-large +Table 10:Experiments on finetuning RoBERTa (355M). Results marked with +† + are taken from (Hu et al., 2022), and those with +∗ + refer to Adapter +P + in (Hu et al., 2022). +RoBERTa # para SST2 STS-B RTE QQP QNLI MRPC MNLI CoLA avg +FT† 355M 96.4 92.4 86.6 92.2 94.7 90.9 90.2 68.0 88.9 +Adapter∗ 0.8M 96.6 91.9 80.1 91.7 94.8 89.7 - 67.8 - +LoRA 0.8M 95.8 92.4 88.2 91.4 94.7 89.6 90.6 64.8 88.4 +LoRA-oBAR 0.8M 96.0 92.6 88.7 91.6 94.8 90.3 90.6 65.1 88.7 +LoRA-nBAR 0.8M 96.0 92.6 89.2 91.6 94.7 90.3 90.8 65.6 88.9 + +Our implementation is inspired from (Hu et al., 2022)9, which is under MIT License. The hyperparameters are chosen the same as provided in its GitHub Repo. AdamW is adopted as the base optimizer. However, we employ single GPU rather than multiple ones and use gradient accumulation rather than parallelism due to memory constraint. We also note that there could be failure cases for LoRA using certain seed, e.g., SST-2 with seed 1 and MNLI with seed 2. These cases are ignored when comparing. We consider the GLUE benchmark and report the mismatched accuracy for MNLI, Matthew’s correlation for CoLA, Pearson correlation for STS-B, and accuracy for other datasets. Larger values indicate better results for all datasets. For LoRA, we employ +𝑟 += +8 + and +𝛼 += +16 +. Experiments are conducted over three random trials for all datasets, with the exception of QQP, for which only two trials are performed due to its large size. The results of final test performance can be found in Table 10. Estimated runtime varies for different datasets from 2 to 15 hours, except for QQP which takes 3 days on our device. + +For the hyperparameters of oBAR and nBAR, +𝜇 +0 + is typically chosen from +{ +0.2 +, +0.5 +, +1.0 +} +; however, for QQP, a value of +0.05 + is used. The scheduler is chosen from linear and constant. We also observe that for datasets such as COLA and RTE, setting weight decay as +0 + works best for BAR. + +D.5GPT2 medium on WebNLG Challenge + +AdamW is adopted as base optimizer. The hyperparameters can be found in Table 11. Our results are obtained from three random trials. Each trial takes roughly 8 hours on our hardware. + +Table 11:Hyperparameters used for GPT2. +Hyper-parameters Values +LoRA +𝑟 + (rank) 4 +LoRA +𝛼 + 32 +# epochs 5 +batchsize 8 +learning rate 2 +× +10 +− +4 + +label Smooth 0.1 + +𝜇 +0 + for BAR 0.1, 0.15, 0.2, 0.25, 0.3 +scheduler for BAR linear, constant +beam size 10 +length penalty 0.8 +Generated on Fri Oct 18 18:16:53 2024 by LaTeXML +Report Issue +Report Issue for Selection diff --git a/docs/bi_lora_sharpness_aware.md b/docs/bi_lora_sharpness_aware.md new file mode 100644 index 0000000..8ffb45f --- /dev/null +++ b/docs/bi_lora_sharpness_aware.md @@ -0,0 +1,1180 @@ +Title: Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models + +URL Source: https://arxiv.org/html/2508.19564 + +Markdown Content: +Yuhang Liu 1,, Tao Li 1,∗, Zhehao Huang 1, Zuopeng Yang 1, Xiaolin Huang 1,2, + +1 Institute of Image Processing and Pattern Recognition, School of Automation and Intelligent Sensing, + +Shanghai Jiao Tong University + +2 MoE Key Laboratory of System Control and Information Processing (Shanghai) + +{yuhangliu,li.tao,kinght_h,yzpeng,xiaolinhuang}@sjtu.edu.cn + +###### Abstract + +Low-Rank Adaptation (LoRA) enables parameter-efficient fine-tuning of large pre-trained models. Yet LoRA can face generalization challenges. One promising way to improve the generalization is Sharpness-Aware Minimization (SAM), which has proven effective for small-scale training scenarios. In this paper, we propose Bi-directional Lo w-R ank A daptation (Bi-LoRA), which introduces an auxiliary adversarial LoRA module. This design explicitly decouples sharpness optimization, handled by the auxiliary module, from task adaptation, performed by the primary module. Such a separation yields two key benefits. First, it transforms SAM’s sequential computation of adversarial perturbation and gradient descent into a parallel form, which roughly halves the time and conquers the main obstacle of applying SAM in LoRA. Second, it provides perturbations from the auxiliary module that do not collapse into the restricted optimization subspace of the primary module, enabling broader sharpness exploration and flatter minima. Bi-LoRA simultaneously achieves both efficiency and effectiveness within a single framework, as validated by extensive experiments across diverse architectures and tasks. Code is available at [https://github.com/CrazyElements/Bi-LoRA](https://github.com/CrazyElements/Bi-LoRA). + +## 1 Introduction + +The paradigm of pretraining followed by fine-tuning has become the de facto standard in machine learning, demonstrating state-of-the-art performance across various tasks(Devlin et al., [2018](https://arxiv.org/html/2508.19564#bib.bib19 "Bert: pre-training of deep bidirectional transformers for language understanding"); Kolesnikov et al., [2020](https://arxiv.org/html/2508.19564#bib.bib8 "Big transfer (bit): general visual representation learning"); Dosovitskiy et al., [2021](https://arxiv.org/html/2508.19564#bib.bib9 "An image is worth 16x16 words: transformers for image recognition at scale"); Radford et al., [2021](https://arxiv.org/html/2508.19564#bib.bib5 "Learning transferable visual models from natural language supervision")). However, as model sizes continue to grow, full fine-tuning (Full FT) becomes memory-prohibitive in resource-constrained settings. The most successful approach for reducing memory cost in fine-tuning large-scale model is Low-Rank Adaptation (LoRA)(Hu et al., [2022](https://arxiv.org/html/2508.19564#bib.bib32 "LoRA: low-rank adaptation of large language models")), which introduces trainable, task-specific low-rank matrices to model weight updates. LoRA greatly lowers memory requirements by significantly reducing the trainable parameters and has become one of the most popular solutions for fine-tuning and deploying large models due to its simplicity and performance comparable to Full FT. + +Despite LoRA’s memory efficiency, it still faces generalization challenges when fine-tuned with limited data. In such scenarios, the risk of overfitting becomes particularly acute and can significantly compromise model performance (Li et al., [2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape"); Deng et al., [2025](https://arxiv.org/html/2508.19564#bib.bib65 "EFlat-lora: efficiently seeking flat minima for better generalization in fine-tuning large language models and beyond"); Lin et al., [2024](https://arxiv.org/html/2508.19564#bib.bib66 "LoRA dropout as a sparsity regularizer for overfitting control"); Li et al., [2024b](https://arxiv.org/html/2508.19564#bib.bib67 "LoRASC: expressive and generalizable low-rank adaptation for large models via slow cascaded learning")). + +Motivated by the connection between flatness of the loss landscape and generalization, one promising direction for improving generalization is to seek flat minima(Hochreiter and Schmidhuber, [1994](https://arxiv.org/html/2508.19564#bib.bib20 "Simplifying neural nets by discovering flat minima"); [1997](https://arxiv.org/html/2508.19564#bib.bib21 "Flat minima")). Sharpness-Aware Minimization (SAM)(Foret et al., [2021](https://arxiv.org/html/2508.19564#bib.bib27 "Sharpness-aware minimization for efficiently improving generalization")) is a widely used technique that enhances generalization by formulating optimization as a min-max problem, effectively minimizing the worst-case loss within a local neighborhood. It has achieved state-of-the-art performance in small-scale training scenarios(Chen et al., [2022](https://arxiv.org/html/2508.19564#bib.bib2 "When vision transformers outperform resnets without pre-training or strong data augmentations"); Zhuang et al., [2022](https://arxiv.org/html/2508.19564#bib.bib11 "Surrogate gap minimization improves sharpness-aware training")). However, SAM requires calculating an adversarial perturbation of model parameters at each training step, incurring additional memory overhead of a copy of the model weights and doubling the training time. When applied to large-scale models, such extra memory cost and computation becomes particularly pronounced. + +![Image 1: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/bi_lora_framework.png) + +Figure 1: Overview of our proposed Bi-LoRA. During fine-tuning (Left), Bi-LoRA utilizes two opposing LoRA modules: the primary module (B_{1}A_{1}) is optimized by conventional gradient descent for task-specific adaptation, while the auxiliary module (B_{2}A_{2}) performs gradient ascent for sharpness optimization. These two modules are decoupled and simultaneously optimized. During inference (Right), only (B_{1}A_{1}) is retained and merged with the pretrained weights. + +Thus, a natural approach to finding flat minima when fine-tuning large models, without sacrificing memory efficiency, is to apply SAM to the LoRA parameters, referred to as LoRA-SAM(Li et al., [2024a](https://arxiv.org/html/2508.19564#bib.bib53 "Implicit regularization of sharpness-aware minimization for scale-invariant problems")). However, this straightforward integration still suffers from the doubled training cost of SAM. To address this essential problem, we introduce an auxiliary adversarial LoRA module to decouple the sharpness optimization from the task adaptation.This design, Bi-directional Low-Rank Adaptation (Bi-LoRA) shown in Figure[1](https://arxiv.org/html/2508.19564#S1.F1 "Figure 1 ‣ 1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), enables updating both modules in one forward and backward pass, making SAM practical for fine-tuning large-scale models. + +Decoupling the two LoRA modules also enables broader exploration of perturbations. During inference, the LoRA parameters are merged into the pretrained weights, so we should care about the sharpness of the full parameter space. In contrast, LoRA-SAM limits adversarial perturbations to a restricted subspace (see Proposition[1](https://arxiv.org/html/2508.19564#Thmproposition1 "Proposition 1 (Perturbation Space of LoRA-SAM). ‣ 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). Worse, LoRA-SAM collapses quickly into it (see Figure[2(c)](https://arxiv.org/html/2508.19564#S2.F2.sf3 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), thereby optimizing sharpness only there. Empirically, Bi-LoRA alleviates this issue: its decoupled auxiliary perturbation space converges more slowly than the primary optimization space, enabling exploring perturbation beyond the restricted subspace to capture sharpness and promote flatter minima. Bi-LoRA maintains LoRA’s efficiency while improving generalization, making it a strong alternative for fine-tuning large-scale models. + +Our contributions can be summarized as follows: + +* • +We propose Bi-LoRA, which introduces an auxiliary LoRA module to model SAM’s adversarial weight perturbation, decoupling it from LoRA optimization. This design enables simultaneous optimization of SAM’s two steps with minor additional memory costs. + +* • +We point out that directly applying SAM to LoRA parameters can only optimize the sharpness within the restricted subspace, thereby limiting its potential to improve generalization. Bi-LoRA broadens the sharpness exploration: the auxiliary perturbation module converges more slowly than the primary one, enabling exploration beyond the restricted subspace. + +* • +Extensive experiments across a wide range of fine-tuning tasks, including natural language understanding, mathematics, code generation, chat, instruction following, and Text-to-Image generation, demonstrate that Bi-LoRA achieves superior generalization performance. + +## 2 Issues of LoRA-SAM + +In this section, we formalize LoRA-SAM and analyze its subspace collapse issue during training. Related work on LoRA and SAM is summarized in Appendix[E](https://arxiv.org/html/2508.19564#A5 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +### 2.1 Preliminaries + +![Image 2: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_layer_60.png) + +(a) Norm of different components + +![Image 3: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_ratio_layer_60.png) + +(b) Norm ratio + +![Image 4: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/lora_trajectory_comparison_cola.png) + +(c) Cosine similarity to final + +Figure 2: Training statistics for LoRA-SAM during fine-tuning. (a) and (b): Frobenius norms of different terms and the ratio of the Frobenius norms of (B\epsilon_{A}+\epsilon_{B}A) to that of (\epsilon_{B}\epsilon_{A}) in Eqn.([5](https://arxiv.org/html/2508.19564#S2.E5 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). We monitor a fixed LoRA module and more similar experiments can be found in Appendix[T](https://arxiv.org/html/2508.19564#A20 "Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). It could be observed that the norm of the third term (\epsilon_{B}\epsilon_{A}) is several orders of magnitude (>10^{4}) smaller than that of the first two terms, making it negligible. Note that \|B\epsilon_{A}\|_{F} is initially zero since B are initialized to zero by default(Hu et al., [2022](https://arxiv.org/html/2508.19564#bib.bib32 "LoRA: low-rank adaptation of large language models")). (c): Cosine similarity between the final LoRA parameters and those during fine-tuning. The LoRA parameters converge rapidly during fine-tuning. The experiments are conducted on CoLA with the T5-base model. + +Given a pre-trained weight matrix W_{0}\in\mathbb{R}^{m\times n}, LoRA utilizes low-rank matrices B and A to model the weight change \Delta W during the fine-tuning, + +\displaystyle W=W_{0}+\Delta W\approx W_{0}+BA,(1) + +where B\in\mathbb{R}^{m\times r} and A\in\mathbb{R}^{r\times n} with r\ll\min\{m,n\}. Here we omit the scaling factor s=\alpha/r for the sake of simplicity in the equation, as it can be easily incorporated into B and A. + +During gradient back-propagation, the loss gradient w.r.t. B and A is computed using the chain rule: + +\frac{\partial\mathcal{L}}{\partial B}=(\nabla_{W}\mathcal{L}){A}^{\top},\quad\frac{\partial\mathcal{L}}{\partial A}={B}^{\top}(\nabla_{W}\mathcal{L}),(2) + +where \mathcal{L} is the loss objective to be minimized. + +### 2.2 LoRA-SAM + +For optimizing the sharpness while keeping memory efficiency, a natural idea is to apply SAM over the LoRA parameters (LoRA-SAM) which formulates the following optimization target: + +\displaystyle\min_{B,A}~~\max_{\|(\mathbf{\epsilon}_{B},\mathbf{\epsilon}_{A})\|\leq\rho}~~\mathcal{L}\left({W}_{0}+({B}+\mathbf{\epsilon}_{B})({A}+\mathbf{\epsilon}_{A})\right),\,(3) + +where \epsilon_{B}\in\mathbb{R}^{m\times r},\epsilon_{A}\in\mathbb{R}^{r\times n} are the adversarial weight perturbations over low-rank matrices, \|(\epsilon_{B},\epsilon_{A})\| denotes the norm of weight perturbations (a typical setting is the \ell_{2}-norm), and \rho is the neighborhood radius. + +To efficiently solve the inner maximization problem in Eqn.([3](https://arxiv.org/html/2508.19564#S2.E3 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), Foret et al. ([2021](https://arxiv.org/html/2508.19564#bib.bib27 "Sharpness-aware minimization for efficiently improving generalization")) employ first-order Taylor expansion for approximation, and the resulting perturbations are calculated as follows: + +\displaystyle\epsilon_{B}\displaystyle=\rho\cdot{\frac{\partial\mathcal{L}}{\partial B}}/{F_{\text{total}}},\quad\epsilon_{A}=\rho\cdot{\frac{\partial\mathcal{L}}{\partial A}}/{F_{\text{total}}},\quad F_{\text{total}}=\sqrt{\left\|\frac{\partial\mathcal{L}}{\partial B}\right\|_{F}^{2}+\left\|\frac{\partial\mathcal{L}}{\partial A}\right\|_{F}^{2}}\,,(4) + +where \|\cdot\|_{F} denotes the Frobenius norm. Then we can rewrite the perturbation given by Eqn.([3](https://arxiv.org/html/2508.19564#S2.E3 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) as + +\displaystyle\epsilon_{W}\displaystyle=B\epsilon_{A}+\epsilon_{B}A+\epsilon_{B}\epsilon_{A}.\,(5) + +Since both \epsilon_{A} and \epsilon_{B} are several orders of magnitude smaller than the original matrices A and B, the cross term \epsilon_{B}\epsilon_{A} is negligible, as evidenced by Figures[2(a)](https://arxiv.org/html/2508.19564#S2.F2.sf1 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [2(b)](https://arxiv.org/html/2508.19564#S2.F2.sf2 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), which report the magnitudes of these terms. Therefore, Eqn.([5](https://arxiv.org/html/2508.19564#S2.E5 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) can be approximately simplified as: + +\displaystyle\epsilon_{W}\approx c\left[BB^{\top}(\nabla_{W}\mathcal{L})+(\nabla_{W}\mathcal{L}){A}^{\top}A\right],\,(6) + +where c=\rho/\sqrt{\|\left(\nabla_{W}\mathcal{L}\right){A}^{\top}\|^{2}_{F}+\|{B}^{\top}(\nabla_{W}\mathcal{L})\|^{2}_{F}}. Now one can observe that the perturbation of LoRA-SAM is confined to a restricted subspace, as demonstrated below: + +###### Proposition 1(Perturbation Space of LoRA-SAM). + +The effective weight perturbation in LoRA-SAM can be decomposed into two terms: BB^{\top}(\nabla_{W}\mathcal{L}) and (\nabla_{W}\mathcal{L})A^{\top}A. The column space of the first term is given by \text{Col}(B), while the row space of the second term is given by \text{Row}(A). + +The space for adversarial weight perturbation in LoRA-SAM is primarily dominated by \text{Col}(B) and \text{Row}(A), from which it follows that LoRA-SAM only cares about the sharpness within the subspace defined by B and A, failing to capture the sharpness in a broader space. + +Moreover,B and A converge rapidly during training, as shown in Figure[2(c)](https://arxiv.org/html/2508.19564#S2.F2.sf3 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). This leads to the fast convergence of the column and row spaces\text{Col}(B)and\text{Row}(A), further restricting the subspace for sharpness optimization and hindering the effectiveness of LoRA-SAM, potentially leading to suboptimal performance. Appendix[U](https://arxiv.org/html/2508.19564#A21 "Appendix U Restricting perturbation in the optimization subspace limits LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") further corroborates this subspace collapse phenomenon. + +In Figure[4](https://arxiv.org/html/2508.19564#S3.F4 "Figure 4 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), we compare the loss landscape flatness of different methods. We observe that LoRA-SAM achieves the flattest loss landscape within the LoRA parameter space as in Figure[4(a)](https://arxiv.org/html/2508.19564#S3.F4.sf1 "In Figure 4 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), reflecting its original sharpness-minimization objective. In contrast, Figure[4(b)](https://arxiv.org/html/2508.19564#S3.F4.sf2 "In Figure 4 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that while LoRA-SAM yields a flatter landscape than LoRA in the full space, the loss still rises sharply with perturbation magnitude. This indicates the limitation of applying perturbation in a restricted subspace, highlighting the importance of considering sharpness over a broader space beyond LoRA subspace. + +## 3 Bi-LoRA: Bi-directional Low-Rank Adaptation + +In this section, we introduce Bi-LoRA, a novel LoRA variant, which contains an auxiliary module beyond the primary module in regular LoRA. The aim of the auxiliary module is to decouple the sharpness optimization from task adaptation, which makes the training efficient and enhances the generalization improvement. Specifically, the proposed Bi-LoRA takes the following formulation, + +W=W_{0}+B_{1}A_{1}+B_{2}A_{2},\,(7) + +where the first LoRA module (A_{1},B_{1}) serves as the primary module responsible for task-specific adaptation, similar to standard LoRA, while the second module (A_{2},B_{2}) acts as the auxiliary LoRA module for modeling adversarial perturbation in SAM. With these two modules, the Bi-LoRA’s optimization objective is given below, + +\displaystyle\min_{B_{1},A_{1}}\max_{\|B_{2}A_{2}\|_{F}\leq\rho}\mathcal{L}\left(W_{0}+B_{1}A_{1}+B_{2}A_{2}\right),\,(8) + +where \rho>0 is the neighborhood radius that controls the magnitude of perturbations as in the original SAM. After training, we discard the auxiliary modules (A_{2},B_{2}), as they serve solely to optimize sharpness during training and guide the primary module (A_{1},B_{1}) towards a flat region. Consequently, the adapted weights are reduced to the primary module, i.e., W=W_{0}+B_{1}A_{1}, preserving the original LoRA structure and ensuring no additional computational overhead during inference. + +In LoRA-SAM, the adversarial perturbation is calculated as Eqn.([4](https://arxiv.org/html/2508.19564#S2.E4 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), where the backpropagation on (\epsilon_{A},\epsilon_{B}) and on (A,B) must be carried out sequentially. Now in Bi-LoRA, the task adaptation and the perturbation are decoupled, as shown in Eqn.([8](https://arxiv.org/html/2508.19564#S3.E8 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). Then one can simultaneously perform task adaptation and sharpness-aware minimization _in one gradient step_. Specifically, during each iteration, the primary LoRA module (A_{1},B_{1}) is updated through standard gradient descent, while the auxiliary LoRA module (A_{2},B_{2}) is updated through gradient ascent for sharpness optimization. + +The bi-directional gradient update approach embodies the core concept of our Bi-LoRA. Concretely, one update step for Bi-LoRA is formalized as follows: + +\displaystyle\left\{\begin{aligned} B_{1}^{k+1}&=B_{1}^{k}-\eta_{1}\left(\nabla_{W}\mathcal{L}\right){A^{k\top}_{1}},\quad A_{1}^{k+1}=A_{1}^{k}-\eta_{1}B_{1}^{k\top}\left(\nabla_{W}\mathcal{L}\right),\,\\ +B_{2}^{k+1}&=B_{2}^{k}+\eta_{2}\left(\nabla_{W}\mathcal{L}\right)A_{2}^{k\top},\quad A_{2}^{k+1}=A_{2}^{k}+\eta_{2}B_{2}^{k\top}\left(\nabla_{W}\mathcal{L}\right),\,\end{aligned}\right.(9) + +where \nabla_{W}\mathcal{L}=\left.\frac{\partial\mathcal{L}}{\partial W}\right|_{W=W_{0}+B^{k}_{1}A^{k}_{1}+B^{k}_{2}A^{k}_{2}} is the gradient of the loss \mathcal{L} w.r.t. the merged weight W, k denotes the iteration index, and \eta_{1},\eta_{2} are learning rates. We use the same learning rate for both LoRA modules in our experiments, though other choices are possible (see Appendix[Q](https://arxiv.org/html/2508.19564#A17 "Appendix Q Learning Rates for Tuning the Two Modules in Bi-LoRA ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). + +The independent backpropagation for the two LoRA modules (A_{1},B_{1}) and (A_{2},B_{2}) makes the updates parallelizable. This eliminates the doubled computational cost typically incurred by LoRA-SAM. The following proposition ensures that the adversarial direction induced by (A_{2},B_{2}) can still increase the inner objective, in non-negatively alignment with the SAM’s perturbation direction. The Proof and further discussions are provided in Appendix[F.1](https://arxiv.org/html/2508.19564#A6.SS1 "F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +###### Proposition 2(Alignment of Bi-LoRA’s ascent direction with previous full gradient). + +Let G_{t}=\nabla_{W}\mathcal{L}(W_{0}+B_{1,t}A_{1,t}+\tilde{\epsilon}_{t}) denote the full gradient at step t, with \tilde{\epsilon}_{t}=B_{2,t}A_{2,t}. After one Bi-LoRA update, + +\langle G_{t},\tilde{\epsilon}_{t+1}-\tilde{\epsilon}_{t}\rangle\geq 0, + +i.e., Bi-LoRA increases the inner objective along the previous SAM’s perturbation direction. + +![Image 5: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/cola_wd0.0_lr5e-4_bilora.png) + +Figure 3: Cosine similarity between the main LoRA (task adaptation, red), the auxiliary LoRA (perturbation, blue), and their trajectories (green) during training T5-base on CoLA. The auxiliary converges only in the last 20% of steps, substantially slower and independent of the main. + +Decoupling sharpness optimization from task adaptation not only makes the backpropagation parallelizable but also partially alleviates the inconsistency of LoRA-SAM, where perturbations are applied to a restricted subspace (A,B) while inference is determined by the full parameters (see Proposition[1](https://arxiv.org/html/2508.19564#Thmproposition1 "Proposition 1 (Perturbation Space of LoRA-SAM). ‣ 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). Bi-LoRA eliminates the dependence of adversarial weight perturbations on the primary LoRA module. In Eqn.[8](https://arxiv.org/html/2508.19564#S3.E8 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), we see the column space for perturbations in Bi-LoRA is spanned by \text{Col}(B_{2}), which is independent of the LoRA optimization space, i.e., \text{Col}(B_{1}). + +Although Bi-LoRA’s perturbation space \mathrm{Col}(B_{2}) (auxiliary) remains a subspace, it is strictly decoupled from the optimization space \mathrm{Col}(B_{1}) (primary). Moreover, we find that the auxiliary modules converge much more slowly than the primary ones, preserving flexibility for sharpness-aware updates (see Figure[3](https://arxiv.org/html/2508.19564#S3.F3 "Figure 3 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and Appendix[L](https://arxiv.org/html/2508.19564#A12 "Appendix L Convergence Analysis of Main and Auxiliary LoRA Modules ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). Therefore, Bi-LoRA enables a more effective capture of sharpness in the full parameter space. And Figure[4](https://arxiv.org/html/2508.19564#S3.F4 "Figure 4 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") confirms that Bi-LoRA achieves a significantly flatter loss landscape in the full parameter space, leading to improved generalization. + +![Image 6: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/loss_landscape_1d_lora_space.png) + +(a) LoRA parameters + +![Image 7: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/loss_landscape_1d_full_space.png) + +(b) Full parameters + +Figure 4: Loss landscape visualization along random “filter-normalized” directions following Li et al. ([2018](https://arxiv.org/html/2508.19564#bib.bib7 "Visualizing the loss landscape of neural nets")), focusing on (a) LoRA parameters and (b) full parameters. It can be observed that while LoRA-SAM attains the greatest flatness within the LoRA subspace, which aligns with its sharpness-minimization objective, it is not the flattest in the full parameter space. This distinction is critical for inference, as LoRA parameters are ultimately merged into the pretrained weights. In contrast, Bi-LoRA delivers a substantially greater improvement on flatness in the full parameter space. All experiments were averaged over five independent runs with T5-base fine-tuned on CoLA. + +Numerous studies have discussed SAM’s convergence and generalization(Andriushchenko and Flammarion, [2022](https://arxiv.org/html/2508.19564#bib.bib68 "Towards understanding sharpness-aware minimization"); Dai et al., [2023](https://arxiv.org/html/2508.19564#bib.bib69 "The crucial role of normalization in sharpness-aware minimization"); Khanh et al., [2024](https://arxiv.org/html/2508.19564#bib.bib70 "Fundamental convergence analysis of sharpness-aware minimization")). For Bi-LoRA, the following proposition establishes that the Bi-LoRA objective (Eqn.[8](https://arxiv.org/html/2508.19564#S3.E8 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) is essentially equivalent to the vanilla LoRA optimization with this regularization term. + +###### Proposition 3(Bi-LoRA is a Regularized LoRA). + +From the gradient-norm perspective, SAM can be expressed as + +\mathcal{L}(W)+\rho\left\|\nabla_{W}\mathcal{L}(W)\right\|_{2}. + +If the inner maximization of Bi-LoRA is solved to convergence and attains its optimum, its objective reduces to a low-rank counterpart: + +\min_{B_{1},A_{1}}\;\mathcal{L}(W_{0}+B_{1}A_{1})+\rho\left\|\nabla_{W_{0}+B_{1}A_{1}}\mathcal{L}\right\|_{(r)}, + +where \|\cdot\|_{(r)} denotes the Ky Fan r-norm (the sum of the top-r singular values). Thus, ideally, Bi-LoRA can be viewed as LoRA equipped with an explicit low-rank gradient-norm regularizer. + +Proofs and further details are deferred to the Appendix[F.2](https://arxiv.org/html/2508.19564#A6.SS2 "F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). Proposition[3](https://arxiv.org/html/2508.19564#Thmproposition3 "Proposition 3 (Bi-LoRA is a Regularized LoRA). ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") coincides with the regularization view of SAM(Yue et al., [2023](https://arxiv.org/html/2508.19564#bib.bib71 "Sharpness-aware minimization revisited: weighted sharpness as a regularization term")). Bi-LoRA can be analyzed in the same way, differing only in that its regularizer is a low-rank norm. In practice, however, we perform a single inner step to preserve high efficiency, which is crucial for fine-tuning models. This choice introduces a gap between the practical Bi-LoRA and analyses that require optimality of the inner optimization. Given our application-oriented focus, we do not delve deeply into detailed theoretical discussions. + +##### Norm Constraints. + +In Eqn.([8](https://arxiv.org/html/2508.19564#S3.E8 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), we impose a norm constraint to keep the perturbation sufficiently small to avoid disrupting normal model training. Thus, we apply global clipping to the auxiliary module (A_{2},B_{2}) after each update, by constraining their total Frobenius norm, which maintains the weight perturbation within a controlled magnitude. Specifically, suppose that there are N LoRA layers. The i-th auxiliary module is scaled as follows: + +\left\{\begin{aligned} B^{(i)}_{2}&\leftarrow\sqrt{\rho/c_{\text{norm}}}\cdot B^{(i)}_{2},\\ +A^{(i)}_{2}&\leftarrow\sqrt{{\rho}/{c_{\text{norm}}}}\cdot A^{(i)}_{2},\\ +\end{aligned}\quad\text{if}\quad c_{\text{norm}}>\rho.\,\right.(10) + +Notice the normalization is applied only if the total Frobenius norm over N auxiliary LoRA modules c_{\text{norm}} exceeds the neighborhood radius \rho, where c_{\text{norm}}=\sqrt{\sum_{j=1}^{N}\|B_{2}^{(j)}A_{2}^{(j)}\|^{2}_{F}}. This ensures the perturbation remains constrained within the preset \rho-norm ball. The overall procedure for Bi-LoRA is summarized in Algorithm [1](https://arxiv.org/html/2508.19564#alg1 "Algorithm 1 ‣ Norm Constraints. ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +Algorithm 1 Bi-LoRA + +1:Input: Initial weight + +W_{0} +, learning rates + +\eta_{1} +, + +\eta_{2} +, radius + +\rho +, number of LoRA layers + +N + +2:Output: Adapted weight + +W +for inference + +3: Initialize LoRA modules + +B^{0}_{1} +, + +A^{0}_{1} +, + +B^{0}_{2} +, + +A^{0}_{2} +; + +4: + +k\leftarrow 0 +; + +5:while _not converged_ do + +6: Sample mini-batch data + +\mathcal{B} +; + +7: Apply gradient descent to primary and ascent to auxiliary LoRA modules via Eqn.([9](https://arxiv.org/html/2508.19564#S3.E9 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")); + +8: Clip the auxiliary modules via Eqn.([10](https://arxiv.org/html/2508.19564#S3.E10 "In Norm Constraints. ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")); + +9: + +k\leftarrow k+1 +; + +10:end while + +11: Remove the auxiliary LoRA modules + +B_{2}^{k},A_{2}^{k} +; + +12:return + +W^{k}=W_{0}+B^{k}_{1}A^{k}_{1} + +## 4 Experiments + +In this section, we evaluate Bi-LoRA across diverse models and benchmark tasks. We test Bi-LoRA’s capacities on: (1) Llama 2/3.1 models(Touvron et al., [2023](https://arxiv.org/html/2508.19564#bib.bib44 "Llama 2: open foundation and fine-tuned chat models"); Dubey et al., [2024](https://arxiv.org/html/2508.19564#bib.bib43 "The llama 3 herd of models")) for mathematical reasoning, coding, dialogue, and instruction following; (2) Qwen 2.5-14B(Qwen et al., [2025](https://arxiv.org/html/2508.19564#bib.bib58 "Qwen2.5 technical report")), a larger backbone for instruction-following; and (3) SDXL for text-to-image generation. We further demonstrate that Bi-LoRA can be integrated with existing LoRA variants to provide consistent improvement, and conduct ablation studies to examine its hyperparameter sensitivity. More results on natural language understanding tasks with T5-base are provided in Appendix[I](https://arxiv.org/html/2508.19564#A9 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +### 4.1 Baselines + +We mainly compare with the following baseline methods: + +* • +Full FT fine-tunes all model parameters. + +* • +LoRA applies low-rank adaptation to all linear modules. + +* • +LoRA-SAM applies sharpness-aware minimization over LoRA parameters. + +* • +Random perturbation baselines: Random Output Perturbation (ROP) perturbs logits y with Gaussian noise \epsilon, i.e., \tilde{y}=y+\epsilon,\ \epsilon\sim\mathcal{N}(0,\rho^{2}I); Random Weight Perturbation (RWP)_full/LoRA perturb either full weights or only LoRA adapters with Gaussian noise. + +* • +SAM variants, including LoRA-oBAR/nBAR(Li et al., [2024a](https://arxiv.org/html/2508.19564#bib.bib53 "Implicit regularization of sharpness-aware minimization for scale-invariant problems")), Flat-LoRA(Li et al., [2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape")), LoRA-ESAM(Du et al., [2022a](https://arxiv.org/html/2508.19564#bib.bib59 "Efficient sharpness-aware minimization for improved training of neural networks")), LoRA-LookSAM(Liu et al., [2022b](https://arxiv.org/html/2508.19564#bib.bib60 "Towards efficient and scalable sharpness-aware minimization")), S 2 SAM(Ji et al., [2024](https://arxiv.org/html/2508.19564#bib.bib31 "A single-step, sharpness-aware minimization is all you need to achieve efficient and accurate sparse training")), and WSAM(Yue et al., [2023](https://arxiv.org/html/2508.19564#bib.bib71 "Sharpness-aware minimization revisited: weighted sharpness as a regularization term")). + +* • +LoRA variants, including LoRA-GA(Wang et al., [2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")), PiSSA(Meng et al., [2024](https://arxiv.org/html/2508.19564#bib.bib46 "PiSSA: principal singular values and singular vectors adaptation of large language models")), DoRA(Liu et al., [2024](https://arxiv.org/html/2508.19564#bib.bib54 "Dora: weight-decomposed low-rank adaptation")), HiRA(Huang et al., [2025](https://arxiv.org/html/2508.19564#bib.bib73 "HiRA: parameter-efficient hadamard high-rank adaptation for large language models")), and DeLoRA(Bini et al., [2025](https://arxiv.org/html/2508.19564#bib.bib72 "Decoupling angles and strength in low-rank adaptation")). + +To more rigorously evaluate the effectiveness in improving generalization, we adopt a stronger training protocol with larger learning rate than prior works(Wang et al., [2025](https://arxiv.org/html/2508.19564#bib.bib57 "LoRA-pro: are low-rank adapters properly optimized?"); [2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")). The same protocol is applied across all methods for fairness. Unless otherwise specified, all results are averaged over three independent runs with standard errors. See Appendix[R](https://arxiv.org/html/2508.19564#A18 "Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") for detailed hyperparameter settings and more training protocol. + +### 4.2 Results on Llama Models + +Setting. We evaluate the performance of Bi-LoRA on Llama 2-7B/3.1-8B across four tasks: mathematical reasoning, code generation, dialogue generation, and instruction following, following Wang et al. ([2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")); Li et al. ([2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape")); Ren et al. ([2024](https://arxiv.org/html/2508.19564#bib.bib39 "MELoRA: mini-ensemble low-rank adapters for parameter-efficient fine-tuning")). Llama 2-7B is fine-tuned on the first three tasks, while Llama 3.1-8B is used for instruction following. Each task focuses on a specific capability and uses well-established datasets and metrics for training and evaluation, as detailed below: + +* • +Mathematical reasoning. Our model is fine-tuned on a 100k subset of MetaMathQA and evaluated on GSM8K. Performance is measured by accuracy. + +* • +Code generation. We fine-tune our model on a 100k subset of Code-Feedback and evaluate it on the HumanEval benchmark, using the PASS@1 metric. + +* • +Dialogue generation. We train our model on the WizardLM dataset and evaluate it on MT-Bench. The response quality is assessed with GPT-4, and we report the first-turn score on a 10-point scale. + +* • +Instruction following. Fine-tuned on Cleaned Alpaca(Taori et al., [2023](https://arxiv.org/html/2508.19564#bib.bib42 "Stanford alpaca: an instruction-following llama model")) and evaluated on INSTRUCTEVAL(Chia et al., [2023](https://arxiv.org/html/2508.19564#bib.bib41 "Instructeval: towards holistic evaluation of instruction-tuned large language models")), reporting exact match for MMLU, DROP, and BBH, and PASS@1 for HumanEval. + +In our experiments, following Du et al. ([2022a](https://arxiv.org/html/2508.19564#bib.bib59 "Efficient sharpness-aware minimization for improved training of neural networks")), (Liu et al., [2022b](https://arxiv.org/html/2508.19564#bib.bib60 "Towards efficient and scalable sharpness-aware minimization")), LoRA-ESAM perturbs 50% of the parameters on the top-50% sharpness-sensitive data, while LoRA-LookSAM applies SAM perturbation every five steps, and thus their costs are denoted as “\times 1.x”. + +Table 1: Results of fine-tuning Llama 2-7B and Llama 3.1-8B on different tasks. “Cost” indicates the gradient steps per training iteration, e.g., one step (Cost \times 1) for Full FT, LoRA, and Bi-LoRA. + +Method Cost Llama 2-7B Llama 3.1-8B +GSM8K HumanEval MT-Bench MMLU DROP HEval BBH +Full FT\times 1 59.74±0.69 33.12±0.32 6.16±0.09 64.31±0.31 51.52±0.45 41.45±1.58 44.78±0.33 +LoRA\times 1 58.21±0.34 24.75±0.23 5.92±0.10 63.38±0.39 49.82±0.54 43.15±0.93 42.82±0.27 +LoRA-SAM\times 2 59.16±0.52 26.59±0.36 5.97±0.08 63.46±0.19 50.94±0.22 44.36±1.13 43.49±0.40 +ROP\times 1 59.24±0.70 25.41±0.11 6.05±0.08 63.63±0.20 49.96±0.26 42.27±0.81 43.47±0.29 +RWP_full\times 1 59.41±0.59 26.26±0.35 6.01±0.12 63.40±0.14 50.11±0.12 44.31±0.81 43.35±0.18 +RWP_LoRA\times 1 58.81±0.27 24.92±0.11 5.80±0.11 63.50±0.01 50.16±0.67 42.68±0.70 44.10±0.25 +LoRA-oBAR\times 1 59.26±0.53 26.30±0.33 5.97±0.07 63.62±0.12 49.92±0.48 43.49±0.20 43.44±0.12 +LoRA-nBAR\times 1 59.72±0.25 26.50±0.23 6.10±0.06 63.45±0.10 49.80±0.03 45.23±0.20 43.39±0.17 +Flat-LoRA\times 1 59.44±0.33 26.67±0.23 5.98±0.05 63.67±0.33 50.44±0.17 44.31±0.73 43.99±0.10 +LoRA-ESAM\times 1.x 58.33±0.22 24.84±0.04 5.80±0.13 61.79±1.65 49.31±0.30 42.48±1.33 43.40±0.15 +LoRA-LookSAM\times 1.x 58.55±0.28 25.28±0.29 5.94±0.08 63.45±0.29 50.34±0.34 42.48±0.54 43.28±0.20 +Bi-LoRA\times 1 60.32±0.30 27.20±0.42 6.26±0.06 63.67±0.15 51.53±0.33 46.12±0.89 43.45±0.31 + +Results. We begin with Llama 2-7B for math, code, and chat tasks. The results, presented in Table[1](https://arxiv.org/html/2508.19564#S4.T1 "Table 1 ‣ 4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") (left), demonstrate Bi-LoRA’s superior performance. Compared to LoRA, Bi-LoRA achieves improvements of 2.11% accuracy on GSM8K, 2.45% on HumanEval, and 0.34 on MT-Bench. Importantly, these gains are even more pronounced than those achieved with the smaller T5-base model (Appendix[I](https://arxiv.org/html/2508.19564#A9 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), and Bi-LoRA operates at a similar speed as LoRA, underscoring its scalability. Furthermore, Bi-LoRA narrows the gap between LoRA fine-tuning and full fine-tuning, and notably outperforms full fine-tuning on GSM8K and MT-Bench tasks, which are not attainable by competing methods. In contrast, we observe that LoRA-SAM does not consistently improve upon LoRA. Random-perturbation baselines provide only marginal or task-specific gains relative to LoRA-SAM, showing that trivial random noise perturbation offers limited generalization improvement. Unlike efficient-SAM variants that either incur extra cost (LoRA-ESAM, LoRA-LookSAM) or compromise performance on certain tasks (e.g., Flat-LoRA on DROP), Bi-LoRA achieves consistent gains on all benchmarks at a single-step computation. + +Next, we turn to the instruction following task with Llama 3.1-8B. As shown in Table[1](https://arxiv.org/html/2508.19564#S4.T1 "Table 1 ‣ 4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") (right), Bi-LoRA outperforms almost all baselines, surpassing LoRA by 0.29% on MMLU, 1.71% on DROP, 2.97% on HumanEval, and 0.63% on BBH. While LoRA-SAM and the two types of baselines show improvements over LoRA, Bi-LoRA achieves more substantial gains, particularly on DROP and HumanEval datasets, while requiring only half the training time. Overall, Bi-LoRA offers a better accuracy and efficiency trade-off than both random perturbation and efficient SAM baselines. + +### 4.3 Results on Qwen Model + +Table 2: Results of fine-tuning Qwen 2.5-14B on instruction-following tasks. “Cost” indicates the gradient steps per training iteration, e.g., one step (Cost \times 1) for LoRA and Bi-LoRA. + +Method Cost MMLU DROP HEval BBH Avg. +Vanilla\times 1 79.28±0.17 51.55±0.76 70.93±1.13 57.47±0.03 64.81 +LoRA-SAM\times 2 79.19±0.31 54.54±1.03 70.53±1.81 58.25±0.11 65.63 +ROP\times 1 79.22±0.28 52.74±1.42 71.34±0.35 57.84±0.49 65.28 +RWP_full\times 1 78.92±0.07 54.75±1.01 71.13±1.47 57.76±0.07 65.64 +RWP_LoRA\times 1 79.07±0.22 51.69±0.71 71.34±0.93 57.45±0.22 64.89 +LoRA-oBAR\times 1 79.37±0.28 55.70±0.61 69.92±0.20 58.53±0.46 65.88 +LoRA-nBAR\times 1 79.34±0.32 55.96±0.45 69.92±0.20 58.67±0.44 65.97 +Flat-LoRA\times 1 79.51±0.20 54.90±1.06 71.54±0.20 58.13±0.45 66.02 +LoRA-ESAM\times 1.x 79.55±0.03 54.83±0.97 65.04±0.81 58.43±0.32 64.46 +LoRA-LookSAM\times 1.x 79.63±0.23 56.35±1.55 67.88±1.08 58.41±0.07 65.57 +Bi-LoRA\times 1 79.67±0.05 56.49±0.22 71.34±0.93 58.93±0.05 66.61 + +To demonstrate Bi-LoRA’s performance across model size and architecture, we next experiment on Qwen 2.5-14B, a distinct and larger backbone compared to LlaMA 2/3.1. We focus on instruction-following tasks. Table[2](https://arxiv.org/html/2508.19564#S4.T2 "Table 2 ‣ 4.3 Results on Qwen Model ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that Bi-LoRA achieves the highest average score, improving over LoRA-SAM by 0.98%. Compared to the strongest alternative (Flat-LoRA) across both random perturbation and efficient SAM variants baselines, Bi-LoRA delivers an additional 0.59% average gain. This indicates its applicability across architectures and scales. + +### 4.4 Results on Diffusion Models + +Setting. We apply Bi-LoRA to a subject-driven generalization task and finetune SDXL(Podell et al., [2023](https://arxiv.org/html/2508.19564#bib.bib6 "Sdxl: improving latent diffusion models for high-resolution image synthesis")) via Dreambooth(Ruiz et al., [2023](https://arxiv.org/html/2508.19564#bib.bib18 "Dreambooth: fine tuning text-to-image diffusion models for subject-driven generation")) on 3D Icons dataset, which contains 23 square-icon images. + +Table 3: CLIP I2T and T2T similarity (%) for LoRA and Bi-LoRA on the 3D Icons dataset. + +Method CLIP I2T (\uparrow)CLIP T2T (\uparrow) +LoRA 32.43_{\pm 0.46}42.27_{\pm 2.08} +Bi-LoRA\mathbf{33.14}_{\pm 0.41}\mathbf{46.79}_{\pm 2.37} +![Image 8: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/sdxl_image_show.png) + +Figure 5: Images generated by SDXL fine-tuned with LoRA and Bi-LoRA on the 3D icon datasets, where each column uses the _same_ seed for fair comparisons. + +Results. Table[3](https://arxiv.org/html/2508.19564#S4.T3 "Table 3 ‣ 4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") reports the average CLIP image‐text (I2T) and text‐text (T2T) similarity scores over six prompt instances of the form “a ToK icon of a , in the style of TOK”, each evaluated across 200 runs. Compared with LoRA, Bi-LoRA improves the average I2T and T2T similarities by 0.71 and 4.52, respectively, demonstrating its stronger personalization capability. As shown in Figure[5](https://arxiv.org/html/2508.19564#S4.F5 "Figure 5 ‣ 4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), in the second row, the image generated by Bi-LoRA artfully merges the fluffy rabbit with the icon, whereas LoRA either fails to integrate the rabbit into the icon (fourth) or even does not generate an icon (third). Furthermore, Bi-LoRA better preserves the attributes of the rabbit, such as the eyes in the first column. More qualitative examples are shown in Appendix[K](https://arxiv.org/html/2508.19564#A11 "Appendix K Qualitative Results on SDXL ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +### 4.5 Integration with Other LoRA Variants + +Table 4: Results of fine-tuning T5-base on MRPC and CoLA using various LoRA variants, both standalone and combined with Bi-LoRA. + +Method MRPC CoLA Avg. +LoRA-GA 88.81±0.47 58.87±1.00 73.84 +PiSSA 88.15±0.24 58.66±0.47 73.41 +DoRA 88.81±0.29 59.89±0.72 74.35 +LoRA-GA + Bi-LoRA 89.62±0.24 60.70±0.27 75.16 +PiSSA + Bi-LoRA 89.95±0.50 59.77±0.83 74.86 +DoRA + Bi-LoRA 89.54±0.13 60.77±0.47 75.16 + +In this section, we evaluate Bi-LoRA’s effectiveness when combined with existing advanced LoRA variants. Specifically, we consider three LoRA variants, including LoRA-GA(Wang et al., [2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")), PiSSA(Meng et al., [2024](https://arxiv.org/html/2508.19564#bib.bib46 "PiSSA: principal singular values and singular vectors adaptation of large language models")) and DoRA(Liu et al., [2024](https://arxiv.org/html/2508.19564#bib.bib54 "Dora: weight-decomposed low-rank adaptation")), and fine-tune the T5-base model on the MRPC and CoLA datasets, with detailed training settings provided in Appendix[R.1](https://arxiv.org/html/2508.19564#A18.SS1 "R.1 Experiments on the GLUE and SuperGLUE datasets ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). As shown in Table[4](https://arxiv.org/html/2508.19564#S4.T4 "Table 4 ‣ 4.5 Integration with Other LoRA Variants ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), integrating Bi-LoRA improves the average score across MRPC and CoLA by 1.45%, achieving gains of up to 1.80% on MRPC and 1.83% on CoLA. These results confirm that Bi-LoRA can be seamlessly integrated with previous approaches and deliver consistent improvements. + +### 4.6 Ablations and Hyperparameter Sensitivity + +Table 5: Ablations of each component in Bi-LoRA using Llama 3.1-8B on instruction following tasks. We compare vanilla LoRA, Dual-LoRA (two descent branches with global clipping), and Bi-LoRA under different clipping schemes (no clipping, per-layer clipping, and global clipping). + +Method Clipping Adversarial MMLU DROP HEval BBH Avg. +LoRA✗✗63.38±0.39 49.82±0.54 43.15±0.93 42.82±0.27 49.79 +Dual-LoRA global✗63.49±0.12 49.97±0.17 42.71±0.47 43.25±0.15 49.86 +Bi-LoRA (no clip)✗\checkmark 61.33±0.42 47.31±3.66 40.24±0.93 41.96±0.48 47.71 +Bi-LoRA (per-layer clip)per-layer\checkmark 63.49±0.26 50.26±0.30 43.29±0.61 43.06±0.28 50.03 +Bi-LoRA (global clip)global\checkmark 63.67±0.15 51.53±0.33 46.12±0.89 43.45±0.15 51.19 + +Setting. In this section, we ablate Bi-LoRA to quantify each component’s contribution, and investigate the hyperparameter sensitivity from three factors: (1) auxiliary learning rate \eta_{2}, (2) ranks of primary (r_{1}) and auxiliary r_{2} LoRA modules, and (3) joint sensitivity to neighborhood radius \rho and auxiliary rank r_{2}. We fine-tune Llama 3.1-8B and evaluate its performance on the instruction following tasks. + +Component-wise ablations of Bi-LoRA. We examine whether the improvements of Bi-LoRA arise from the additional capacity introduced by the second LoRA branch or from the adversarial ascent step. To this end, we ablate the clipping schemes and the usage of the auxiliary LoRA module. + +Table[5](https://arxiv.org/html/2508.19564#S4.T5 "Table 5 ‣ 4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that Dual-LoRA, which shares the same two-LoRA architecture as Bi-LoRA but updates both LoRA modules with gradient descent, only slightly outperforms vanilla LoRA (+0.07 Avg.). Intuitively, the parallel descent branch in Dual-LoRA behaves similarly to merging with another LoRA update whose norm is constrained by the global radius\rho, which prevents it from fully leveraging the increased rank and representational capacity. This indicates that the performance gains of Bi-LoRA primarily stem from the adversarial ascent step rather than from merely adding an extra LoRA branch. We further compare global clipping with two alternatives: no-clipping and per-layer clipping. Removing clipping leads to noticeable degradation (-3.48 Avg.), as the perturbations become excessively large. And per-layer clipping yields only marginal improvements over LoRA (+0.24 Avg.). These observations support global clipping as the more stable and effective choice. A more detailed explanation of why we adopt global clipping is provided in Appendix[G](https://arxiv.org/html/2508.19564#A7 "Appendix G Reasons of Using Global Clipping ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +Sensitivity to auxiliary learning rate \eta_{2}. Table[6](https://arxiv.org/html/2508.19564#S4.T6 "Table 6 ‣ 4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that Bi-LoRA is insensitive to auxiliary learning rate \eta_{2}: the average score remains within a narrow band of [50.55,51.19] as \eta_{2} varies from 1e-4 to 1e-3. This indicates that Bi-LoRA is robust to \eta_{2} in a relatively wide range. Interestingly, the best average performance is obtained when \eta_{2}= 3e-4, which coincides with the optimal primary learning rate \eta_{1} given by Ren et al. ([2024](https://arxiv.org/html/2508.19564#bib.bib39 "MELoRA: mini-ensemble low-rank adapters for parameter-efficient fine-tuning")). This suggests that using a shared learning rate for the primary and auxiliary branches yields more coordinated optimization, consistent with our observation in Table[A7](https://arxiv.org/html/2508.19564#A17.T7 "Table A7 ‣ Appendix Q Learning Rates for Tuning the Two Modules in Bi-LoRA ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). Hence, in all main experiments, we simply set \eta_{2}=\eta_{1}, reducing hyperparameter tuning while retaining strong performance. + +Table 6: Sensitivity of Bi-LoRA to the auxiliary learning rate \eta_{2} with Llama 3.1-8B on instruction following tasks, with the same primary learning rate \eta_{1} as in Section[4.2](https://arxiv.org/html/2508.19564#S4.SS2 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +\eta_{2}MMLU DROP HEval BBH Avg +1e-4 63.89±0.25 51.05±0.13 45.12±0.35 42.84±0.15 50.73 +3e-4 63.67±0.15 51.53±0.33 46.12±0.89 43.45±0.15 51.19 +5e-4 63.47±0.35 51.23±0.24 45.93±0.73 43.05±0.12 50.92 +1e-3 63.40±0.17 51.32±0.22 44.51±0.93 42.96±0.23 50.55 + +Sensitivity to ranks r_{1} and r_{2}. Figure[6(a)](https://arxiv.org/html/2508.19564#S4.F6.sf1 "In Figure 6 ‣ 4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that the average point rises monotonically up to the auxiliary rank r_{2}=8, which we adopt as the default for fine-tuning scenario. Additionally, further increasing r_{2} does not yield significant improvements or even downgrades the performance. + +![Image 9: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/effect_rank2_llama_if.png) + +(a) Effect of r_{1} and r_{2}. + +![Image 10: Refer to caption](https://arxiv.org/html/2508.19564v2/x1.png) + +(b) Effect of \rho and r_{2}. + +Figure 6: Sensitivity of Bi-LoRA to primary (r_{1}) and auxiliary (r_{2}) LoRA ranks and perturbation radius \rho with Llama 3.1-8B on instruction-following tasks. + +Sensitivity to neighborhood radius \rho and auxiliary rank r_{2}. Figure[6(b)](https://arxiv.org/html/2508.19564#S4.F6.sf2 "In Figure 6 ‣ 4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") further examines the joint effect of \rho and r_{2}. The heatmap also corroborates that enlarging r_{2} from 2 to 8 is beneficial. Overall, these trends suggest that \rho and r_{2} govern complementary aspects of neighborhood size and auxiliary capacity and can be tuned largely independently for a given architecture and dataset, in line with previous work of SAM. + +### 4.7 Training Time and Memory Cost + +Table 7: Peak memory and time per optimization step (relative to LoRA in parentheses). + +Model & Dataset Method Memory (GB)Time (s) +Llama 3.1-8B Cleaned Alpaca LoRA 23.69 8.93 (100%) +LoRA-SAM 24.00 19.23 (215%) +Bi-LoRA 24.32 9.71 (109%) + +We present detailed comparisons of memory and time per optimization step for LoRA, LoRA-SAM, and Bi-LoRA on Llama 3.1-8B (Cleaned Alpaca). Table[7](https://arxiv.org/html/2508.19564#S4.T7 "Table 7 ‣ 4.7 Training Time and Memory Cost ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that Bi-LoRA significantly reduces the LoRA-SAM’s training overhead, decreasing it from over 210% to about 110% relative to vanilla LoRA. This gain comes from jointly performing optimization and perturbation, eliminating SAM’s extra gradient step. For memory, Bi-LoRA incurs only a minimal overhead (< 0.7GB), as it adds a lightweight auxiliary module. Additional comparisons on training time are provided in Appendix[S](https://arxiv.org/html/2508.19564#A19 "Appendix S Time Comparisons Across Tasks ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +## 5 Conclusion + +In this paper, we propose Bi-LoRA, a novel dual-LoRA variant that leverages an auxiliary LoRA module to enhance the generalization performance of low-rank adaptation. Bi-LoRA decouples optimization from weight perturbation for better optimizing the sharpness of the loss landscape, allowing both to be updated in a single backward pass without requiring additional gradient steps, as in SAM. Extensive experiments across diverse tasks and architectures demonstrate Bi-LoRA’s efficiency and effectiveness in improving generalization. + +#### Acknowledgments + +The research leading to these results has received funding from National Key Research Development Project (2023YFF1104202) and National Natural Science Foundation of China (62376155). The authors would like to thank the chairs and reviewers for their thoughtful comments on this paper. Yuhang Liu also gratefully acknowledges the financial support provided by Prof. Xin Yang at Shanghai Jiao Tong University. + +## References + +* M. Andriushchenko and N. Flammarion (2022)Towards understanding sharpness-aware minimization. In International Conference on Machine Learning (ICML), Cited by: [§3](https://arxiv.org/html/2508.19564#S3.p7.1 "3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* M. Bini, L. Girrbach, and Z. Akata (2025)Decoupling angles and strength in low-rank adaptation. In International Conference on Learning Representations (ICLR), Cited by: [6th item](https://arxiv.org/html/2508.19564#S4.I1.i6.p1.1.2 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* X. Chen, C. Hsieh, and B. Gong (2022)When vision transformers outperform resnets without pre-training or strong data augmentations. In International Conference on Learning Representations (ICLR), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p3.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. K. Chia, P. Hong, L. Bing, and S. Poria (2023)Instructeval: towards holistic evaluation of instruction-tuned large language models. arXiv preprint arXiv:2306.04757. Cited by: [4th item](https://arxiv.org/html/2508.19564#S4.I2.i4.p1.1 "In 4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Dai, S. P. Karimireddy, S. Stich, and M. Jaggi (2023)The crucial role of normalization in sharpness-aware minimization. In Neural Information Processing Systems (NeurIPS), Cited by: [§3](https://arxiv.org/html/2508.19564#S3.p7.1 "3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Deng, Q. Zhu, J. Pang, L. Yang, Z. Fu, and B. Zhang (2025)EFlat-lora: efficiently seeking flat minima for better generalization in fine-tuning large language models and beyond. arXiv preprint arXiv:2508.00522. Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p2.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer (2023)QLoRA: efficient finetuning of quantized llms. In Advances in Neural Information Processing System (NeurIPS), Cited by: [Appendix M](https://arxiv.org/html/2508.19564#A13.p1.1 "Appendix M Bi-LoRA in Quantized Settings ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Devlin, M. Chang, K. Lee, and K. Toutanova (2018)Bert: pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p1.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. (2021)An image is worth 16x16 words: transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p1.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Dou, E. Zhou, Y. Liu, S. Gao, W. Shen, L. Xiong, Y. Zhou, X. Wang, Z. Xi, X. Fan, S. Pu, J. Zhu, R. Zheng, T. Gui, Q. Zhang, and X. Huang (2024)LoRAMoE: alleviating world knowledge forgetting in large language models via MoE-style plugin. In Proceedings of the Association for Computational Linguistics (Volume 1: Long Papers), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Du, H. Yan, J. Feng, J. T. Zhou, L. Zhen, R. S. M. Goh, and V. Tan (2022a)Efficient sharpness-aware minimization for improved training of neural networks. In International Conference on Learning Representations (ICLR), Cited by: [§R.5](https://arxiv.org/html/2508.19564#A18.SS5.p1.1 "R.5 Efficient SAM Variants for LoRA Fine-Tuning ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.2](https://arxiv.org/html/2508.19564#S4.SS2.p1.1 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Du, D. Zhou, J. Feng, V. Y. F. Tan, and J. T. Zhou (2022b)Sharpness-aware training for free. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang, A. Fan, et al. (2024)The llama 3 herd of models. arXiv preprint arXiv:2407.21783. Cited by: [§4](https://arxiv.org/html/2508.19564#S4.p1.1 "4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur (2021)Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations (ICLR), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p3.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§2.2](https://arxiv.org/html/2508.19564#S2.SS2.p2.8 "2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Hochreiter and J. Schmidhuber (1994)Simplifying neural nets by discovering flat minima. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p3.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Hochreiter and J. Schmidhuber (1997)Flat minima. Neural computation. Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p3.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen (2022)LoRA: low-rank adaptation of large language models. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§1](https://arxiv.org/html/2508.19564#S1.p1.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Figure 2](https://arxiv.org/html/2508.19564#S2.F2 "In 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Q. Huang, T. Ko, Z. Zhuang, L. Tang, and Y. Zhang (2025)HiRA: parameter-efficient hadamard high-rank adaptation for large language models. In International Conference on Learning Representations (ICLR), Cited by: [6th item](https://arxiv.org/html/2508.19564#S4.I1.i6.p1.1.2 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Ji, G. Li, J. Fu, F. Afghah, L. Guo, X. Yuan, and X. Ma (2024)A single-step, sharpness-aware minimization is all you need to achieve efficient and accurate sparse training. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§F.1](https://arxiv.org/html/2508.19564#A6.SS1.p1.5.5 "F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* W. Jiang, H. Yang, Y. Zhang, and J. Kwok (2023)An adaptive policy to employ sharpness-aware minimization. In The Eleventh International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* P. D. Khanh, H. Luong, B. S. Mordukhovich, and D. B. Tran (2024)Fundamental convergence analysis of sharpness-aware minimization. In Neural Information Processing Systems (NeurIPS), Cited by: [§3](https://arxiv.org/html/2508.19564#S3.p7.1 "3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* M. Kim, D. Li, S. X. Hu, and T. Hospedales (2022)Fisher sam: information geometry and sharpness aware minimisation. In International Conference on Machine Learning (ICML), Cited by: [Appendix N](https://arxiv.org/html/2508.19564#A14.p1.4 "Appendix N Effect on Adversarial Robustness ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung, S. Gelly, and N. Houlsby (2020)Big transfer (bit): general visual representation learning. In European conference on computer vision (ECCV), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p1.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. A. Koohpayegani, N. K. L., P. Nooralinejad, S. Kolouri, and H. Pirsiavash (2024)NOLA: compressing lora using linear combination of random basis. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Kwon, J. Kim, H. Park, and I. K. Choi (2021)ASAM: adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks. In International Conference on Machine Learning (ICML), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* B. Li and G. B. Giannakis (2023)Enhancing sharpness-aware optimization through variance suppression. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* B. Li, L. Zhang, and N. He (2024a)Implicit regularization of sharpness-aware minimization for scale-invariant problems. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [§R.5](https://arxiv.org/html/2508.19564#A18.SS5.p1.1 "R.5 Efficient SAM Variants for LoRA Fine-Tuning ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§1](https://arxiv.org/html/2508.19564#S1.p4.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein (2018)Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Figure 4](https://arxiv.org/html/2508.19564#S3.F4 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Li, Y. Yang, Y. Shen, F. Wei, Z. Lu, L. Qiu, and Y. Yang (2024b)LoRASC: expressive and generalizable low-rank adaptation for large models via slow cascaded learning. In Findings of the Association for Computational Linguistics: EMNLP, Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p2.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* T. Li, Z. He, Y. Li, Y. Wang, L. Shang, and X. Huang (2025)Flat-lora: low-rank adaptation over a flat loss landscape. In International Conference on Machine Learning (ICML), Cited by: [§R.5](https://arxiv.org/html/2508.19564#A18.SS5.p1.1 "R.5 Efficient SAM Variants for LoRA Fine-Tuning ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Appendix I](https://arxiv.org/html/2508.19564#A9.p1.1 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§1](https://arxiv.org/html/2508.19564#S1.p2.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.2](https://arxiv.org/html/2508.19564#S4.SS2.p1.2 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* T. Li, L. Tan, Z. Huang, Q. Tao, Y. Liu, and X. Huang (2022)Low dimensional trajectory hypothesis is true: dnns can be trained in tiny subspaces. IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI). Cited by: [§F.1](https://arxiv.org/html/2508.19564#A6.SS1.p1.5.5 "F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* T. Li, P. Zhou, Z. He, X. Cheng, and X. Huang (2024c)Friendly sharpness-aware minimization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Lin, X. Ma, X. Chu, Y. Jin, Z. Yang, Y. Wang, and H. Mei (2024)LoRA dropout as a sparsity regularizer for overfitting control. arXiv preprint arXiv:2404.09610. Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p2.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Liu, C. Wang, H. Yin, P. Molchanov, Y. F. Wang, K. Cheng, and M. Chen (2024)Dora: weight-decomposed low-rank adaptation. In International Conference on Machine Learning (ICML), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [6th item](https://arxiv.org/html/2508.19564#S4.I1.i6.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.5](https://arxiv.org/html/2508.19564#S4.SS5.p1.1 "4.5 Integration with Other LoRA Variants ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Liu, S. Mai, X. Chen, C. Hsieh, and Y. You (2022a)Towards efficient and scalable sharpness-aware minimization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Liu, S. Mai, X. Chen, C. Hsieh, and Y. You (2022b)Towards efficient and scalable sharpness-aware minimization. In Conference on Computer Vision and Pattern Recognition (CVPR), Cited by: [§R.5](https://arxiv.org/html/2508.19564#A18.SS5.p1.1 "R.5 Efficient SAM Variants for LoRA Fine-Tuning ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.2](https://arxiv.org/html/2508.19564#S4.SS2.p1.1 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* F. Meng, Z. Wang, and M. Zhang (2024)PiSSA: principal singular values and singular vectors adaptation of large language models. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [6th item](https://arxiv.org/html/2508.19564#S4.I1.i6.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.5](https://arxiv.org/html/2508.19564#S4.SS5.p1.1 "4.5 Integration with Other LoRA Variants ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* D. Podell, Z. English, K. Lacey, A. Blattmann, T. Dockhorn, J. Müller, J. Penna, and R. Rombach (2023)Sdxl: improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952. Cited by: [§4.4](https://arxiv.org/html/2508.19564#S4.SS4.p1.1 "4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Qwen, :, A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, H. Lin, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Lin, K. Dang, K. Lu, K. Bao, K. Yang, L. Yu, M. Li, M. Xue, P. Zhang, Q. Zhu, R. Men, R. Lin, T. Li, T. Tang, T. Xia, X. Ren, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Wan, Y. Liu, Z. Cui, Z. Zhang, and Z. Qiu (2025)Qwen2.5 technical report. arXiv preprint arXiv:2412.15115. Cited by: [§4](https://arxiv.org/html/2508.19564#S4.p1.1 "4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, et al. (2021)Learning transferable visual models from natural language supervision. In International Conference on Machine Learning (ICML), Cited by: [§1](https://arxiv.org/html/2508.19564#S1.p1.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu (2020)Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research (JMLR). Cited by: [Appendix I](https://arxiv.org/html/2508.19564#A9.p1.1 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* P. Ren, C. Shi, S. Wu, M. Zhang, Z. Ren, M. de Rijke, Z. Chen, and J. Pei (2024)MELoRA: mini-ensemble low-rank adapters for parameter-efficient fine-tuning. In Proceedings of the Association for Computational Linguistics (Volume 1: Long Papers), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.2](https://arxiv.org/html/2508.19564#S4.SS2.p1.2 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.6](https://arxiv.org/html/2508.19564#S4.SS6.p4.8.8.7 "4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman (2023)Dreambooth: fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Cited by: [§4.4](https://arxiv.org/html/2508.19564#S4.SS4.p1.1 "4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* R. Taori, I. Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto (2023)Stanford alpaca: an instruction-following llama model. Cited by: [4th item](https://arxiv.org/html/2508.19564#S4.I2.i4.p1.1 "In 4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* C. Tian, Z. Shi, Z. Guo, L. Li, and C. Xu (2024)HydraLoRA: an asymmetric LoRA architecture for efficient fine-tuning. In Advances in Neural Information Processing Systems (NeurIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. (2023)Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: [§4](https://arxiv.org/html/2508.19564#S4.p1.1 "4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman (2019a)Superglue: a stickier benchmark for general-purpose language understanding systems. In Advances in Neural Information Processing Systems (NeurIPS), Vol. 32. Cited by: [Appendix I](https://arxiv.org/html/2508.19564#A9.p1.1 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman (2019b)GLUE: a multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations (ICLR), Cited by: [Appendix I](https://arxiv.org/html/2508.19564#A9.p1.1 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* S. Wang, L. Yu, and J. Li (2024)LoRA-ga: low-rank adaptation with gradient approximation. In Advances in Neural Information Processing Systems (NuerIPS), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [Appendix I](https://arxiv.org/html/2508.19564#A9.p1.1 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [6th item](https://arxiv.org/html/2508.19564#S4.I1.i6.p1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.1](https://arxiv.org/html/2508.19564#S4.SS1.p3.1 "4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.2](https://arxiv.org/html/2508.19564#S4.SS2.p1.2 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.5](https://arxiv.org/html/2508.19564#S4.SS5.p1.1 "4.5 Integration with Other LoRA Variants ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Z. Wang, J. Liang, R. He, Z. Wang, and T. Tan (2025)LoRA-pro: are low-rank adapters properly optimized?. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§4.1](https://arxiv.org/html/2508.19564#S4.SS1.p3.1 "4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Wu, H. Piao, L. Huang, R. Wang, W. Li, H. Pfister, D. Meng, K. Ma, and Y. Wei (2025)SD-loRA: scalable decoupled low-rank adaptation for class incremental learning. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Yen, S. Si, Z. Meng, F. Yu, S. S. Duvvuri, I. S. Dhillon, C. Hsieh, and S. Kumar (2025)LoRA done RITE: robust invariant transformation equilibration for loRA optimization. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Y. Yue, J. Jiang, Z. Ye, N. Gao, Y. Liu, and K. Zhang (2023)Sharpness-aware minimization revisited: weighted sharpness as a regularization term. In ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Cited by: [§3](https://arxiv.org/html/2508.19564#S3.p8.1 "3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5th item](https://arxiv.org/html/2508.19564#S4.I1.i5.p1.1.1 "In 4.1 Baselines ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Zhang, Y. Zhao, D. Chen, X. Tian, H. Zheng, and W. Zhu (2024)MiLoRA: efficient mixture of low-rank adaptation for large language models fine-tuning. In Findings of the Association for Computational Linguistics: EMNLP, Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* Q. Zhang, M. Chen, A. Bukharin, P. He, Y. Cheng, W. Chen, and T. Zhao (2023)Adaptive budget allocation for parameter-efficient fine-tuning. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p1.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian (2024)GaLore: memory-efficient LLM training by gradient low-rank projection. In International Conference on Machine Learning (ICML), Cited by: [§F.1](https://arxiv.org/html/2508.19564#A6.SS1.p1.5.5 "F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). +* J. Zhuang, B. Gong, L. Yuan, Y. Cui, H. Adam, N. Dvornek, S. Tatikonda, J. Duncan, and T. Liu (2022)Surrogate gap minimization improves sharpness-aware training. In International Conference on Learning Representations (ICLR), Cited by: [Appendix E](https://arxiv.org/html/2508.19564#A5.p2.1 "Appendix E Related Work ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [§1](https://arxiv.org/html/2508.19564#S1.p3.1 "1 Introduction ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +## Appendix + +## Appendix A Ethics statement + +This paper proposes Bi-LoRA, a parameter-efficient fine-tuning method that introduces adversarial perturbations to improve generalization of large language models, which can reduce memory and energy costs, and thus broaden access to LLM research. + +For dataset use, we use only publicly available datasets under their licenses, with no human subjects or personal data involved. + +Regarding safety and privacy, we recognize potential risks. By reducing hardware barriers, Bi-LoRA could also enable malicious actors to efficiently fine-tune models for harmful applications or exacerbate dataset biases if such biases are present. Furthermore, the auxiliary adversarial adapter, while improving robustness, introduces an additional parameter pathway that could hypothetically be exploited for backdoor insertion or undesired behaviors if misused. To mitigate these risks, we will anonymously release only training code and experimental scripts in the supplementary materials that follow community standards for responsible research, and we rely exclusively on open benchmarks with no sensitive or personal data. We encourage future users of Bi-LoRA to adopt similar safeguards, including dataset audits for bias, robust model evaluation, and adherence to ethical guidelines on downstream applications. + +## Appendix B reproducibility statement + +We provide detailed experimental settings and hyperparameters in Appendix[R](https://arxiv.org/html/2508.19564#A18 "Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") to ensure reproducibility. The Bi-LoRA implementation and example experiment scripts will be released anonymously in the supplementary materials. All models, benchmarks, and evaluation metrics used in this paper are either open-sourced or explicitly identified when closed-source resources are involved. + +## Appendix C LLM usage Statement + +We use large language models (LLMs) as general-purpose tools at the following stages of this work. + +* • +Writing. We employ both closed-source and open-source LLMs solely for text polishing, including grammar correction, clarity improvement, and minor stylistic refinement. Importantly, they are not used to generate research ideas, technical content, and substantive arguments. + +* • +Benchmarking. For evaluation on the MT-Bench benchmark, we use the GPT-4 API to obtain GPT-4 scores, following the standard evaluation protocol for MT-Bench. + +No part of the research design, methodology, or core scientific contributions relied on LLMs. Their role is restricted to auxiliary assistance in language refinement and standardized benchmark scoring. + +## Appendix D Details on Sharpness-Aware Minimization (SAM) + +Let \mathcal{L}(W) denote the empirical loss of model parameters W. Standard ERM can converge to sharp minima where small parameter perturbations significantly increase loss. SAM’s goal is to target the training toward flat regions by minimizing the worst-case loss inside a : + +\min_{W}\;\;\max_{\lVert\varepsilon\rVert\leq\rho}\;\mathcal{L}\!\left(W+\varepsilon\right),(A1) + +where \rho>0 is the neighborhood radius, \|\cdot\| is a kind of norm, and \varepsilon is the perturbation. Intuitively, Eqn,([A1](https://arxiv.org/html/2508.19564#A4.E1 "In Appendix D Details on Sharpness-Aware Minimization (SAM) ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) prefers parameters whose entire neighborhood has uniformly low loss (flat minima), which has been empirically linked to improved generalization. A practical single-step approximation solves the inner maximization by first-order ascent and then updates W using the gradient evaluated at the perturbed point. + +Denote by \mathcal{L}(W;\mathcal{B}) the mini-batch loss on \mathcal{B} and by \nabla_{W}\mathcal{L} its gradient. A typical SAM update computes a normalized gradient ascent perturbation and then a descent step on the same mini-batch \mathcal{B}_{t}, but at different parameter values: + +\displaystyle{\varepsilon}_{t}\displaystyle=\rho\cdot\frac{\nabla_{W}\mathcal{L}(W_{t};\mathcal{B}_{t})}{\bigl\lVert\nabla_{W}\mathcal{L}(W_{t};\mathcal{B}_{t})\bigr\rVert},(A2) +\displaystyle W_{t+1}\displaystyle\leftarrow W_{t}\;-\;\eta\,\nabla_{W}\mathcal{L}\!\left(W_{t}+{\varepsilon}_{t};\mathcal{B}_{t}\right),(A3) + +i.e., two gradient evaluations per iteration. Eqn.([A2](https://arxiv.org/html/2508.19564#A4.E2 "In Appendix D Details on Sharpness-Aware Minimization (SAM) ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) determines the perturbation direction at W_{t} (clean weight), while Eqn.([A3](https://arxiv.org/html/2508.19564#A4.E3 "In Appendix D Details on Sharpness-Aware Minimization (SAM) ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) computes the update at the perturbed parameter W_{t}+\epsilon_{t} (perturbed weight). + +Algorithm A1 Sharpness-Aware Minimization (SAM) + +1:Input: Initial weights + +W_{0} +; learning rate + +\eta +; radius + +\rho +; optimizer + +\mathsf{Opt}(\cdot) +; max steps + +K + +2:Output: Final weights + +W + +3: Initialize model weights + +W\leftarrow W_{0} + +4:for + +t=0,1,\dots,K{-}1 +do + +5: Sample mini-batch + +\mathcal{B}_{t} + +6: + +g\leftarrow\nabla_{W}\ell(W;\mathcal{B}_{t}) +{gradient at clean weights} + +7: + +\varepsilon\leftarrow\rho\cdot\dfrac{g}{\lVert g\rVert} +{first-order worst-case direction} + +8: + +W_{\mathrm{adv}}\leftarrow W+\varepsilon + +9: + +g_{\mathrm{adv}}\leftarrow\nabla_{W}\mathcal{L}(W_{\mathrm{adv}};\mathcal{B}_{t}) +{use the same batch to produce perturbed gradient} + +10: + +W\leftarrow\mathsf{Opt}\!\bigl(W,\;g_{\mathrm{adv}},\;\eta\bigr) +{e.g., SGD/AdamW step} + +11:end for + +12:return + +W + +## Appendix E Related Work + +Low-Rank Adaptation (LoRA). LoRA(Hu et al., [2022](https://arxiv.org/html/2508.19564#bib.bib32 "LoRA: low-rank adaptation of large language models")) is a widely adopted parameter-efficient fine-tuning (PEFT) method. It models weight updates via low-rank matrices without incurring additional inference costs. Many studies have been proposed to improve the performance of LoRA(Koohpayegani et al., [2024](https://arxiv.org/html/2508.19564#bib.bib40 "NOLA: compressing lora using linear combination of random basis"); Dou et al., [2024](https://arxiv.org/html/2508.19564#bib.bib48 "LoRAMoE: alleviating world knowledge forgetting in large language models via MoE-style plugin"); Tian et al., [2024](https://arxiv.org/html/2508.19564#bib.bib50 "HydraLoRA: an asymmetric LoRA architecture for efficient fine-tuning"); Wang et al., [2025](https://arxiv.org/html/2508.19564#bib.bib57 "LoRA-pro: are low-rank adapters properly optimized?"); Yen et al., [2025](https://arxiv.org/html/2508.19564#bib.bib56 "LoRA done RITE: robust invariant transformation equilibration for loRA optimization"); Wu et al., [2025](https://arxiv.org/html/2508.19564#bib.bib55 "SD-loRA: scalable decoupled low-rank adaptation for class incremental learning")). On the resource perspective, AdaLoRA(Zhang et al., [2023](https://arxiv.org/html/2508.19564#bib.bib38 "Adaptive budget allocation for parameter-efficient fine-tuning")) dynamically adjusts the rank allocation, while MELoRA(Ren et al., [2024](https://arxiv.org/html/2508.19564#bib.bib39 "MELoRA: mini-ensemble low-rank adapters for parameter-efficient fine-tuning")) trains multiple mini-LoRA modules in parallel to cut down on trainable parameters. From the optimization perspective, LoRA-GA(Wang et al., [2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")), PISSA(Meng et al., [2024](https://arxiv.org/html/2508.19564#bib.bib46 "PiSSA: principal singular values and singular vectors adaptation of large language models")) and MiLoRA(Zhang et al., [2024](https://arxiv.org/html/2508.19564#bib.bib47 "MiLoRA: efficient mixture of low-rank adaptation for large language models fine-tuning")) accelerate convergence and boost performance through enhanced initialization, and DoRA(Liu et al., [2024](https://arxiv.org/html/2508.19564#bib.bib54 "Dora: weight-decomposed low-rank adaptation")) decomposes the adaptation into magnitude and direction for better optimization. In this paper, we enhance LoRA optimization by optimizing the loss landscape sharpness of the full parameter space through an auxiliary LoRA module. Our approach is orthogonal to prior works. + +Sharpness-Aware Minimization (SAM). SAM formulates the optimization objective as a min-max problem to seek flat minima, encouraging model parameters to reside in regions with consistently low loss values, yielding state‐of‐the‐art generalization and robustness. Despite its effectiveness, SAM doubles the computational cost because its inner maximization is approximated via an additional gradient ascent step, thereby limiting its application to large-scale models. Several SAM variants have been developed to enhance its generalization performance(Kim et al., [2022](https://arxiv.org/html/2508.19564#bib.bib4 "Fisher sam: information geometry and sharpness aware minimisation"); Li and Giannakis, [2023](https://arxiv.org/html/2508.19564#bib.bib1 "Enhancing sharpness-aware optimization through variance suppression"); Li et al., [2024c](https://arxiv.org/html/2508.19564#bib.bib30 "Friendly sharpness-aware minimization")). ASAM(Kwon et al., [2021](https://arxiv.org/html/2508.19564#bib.bib26 "ASAM: adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks")) introduces adaptive sharpness to address SAM’s scale-dependency issue, while GSAM(Zhuang et al., [2022](https://arxiv.org/html/2508.19564#bib.bib11 "Surrogate gap minimization improves sharpness-aware training")) jointly minimizes surrogate and perturbed losses to locate flatter region. Meanwhile, other works aim to improve SAM’s training efficiency(Du et al., [2022b](https://arxiv.org/html/2508.19564#bib.bib29 "Sharpness-aware training for free"); Jiang et al., [2023](https://arxiv.org/html/2508.19564#bib.bib34 "An adaptive policy to employ sharpness-aware minimization"); Ji et al., [2024](https://arxiv.org/html/2508.19564#bib.bib31 "A single-step, sharpness-aware minimization is all you need to achieve efficient and accurate sparse training")). LookSAM(Liu et al., [2022a](https://arxiv.org/html/2508.19564#bib.bib28 "Towards efficient and scalable sharpness-aware minimization")) applies adversarial perturbations only periodically. Recent studies have also focused on applying SAM to LoRA fine-tuning, aiming to enhance generalization while maintaining training efficiency. For example, BAR(Li et al., [2024a](https://arxiv.org/html/2508.19564#bib.bib53 "Implicit regularization of sharpness-aware minimization for scale-invariant problems")) makes SAM’s implicit balancedness regularization explicit for scale-invariant tasks like LoRA, achieving SAM-level generalization gains with significantly less computation. Flat-LoRA(Li et al., [2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape")) replaces the costly inner maximization with random weight perturbation to tackle the coupling issue in LoRA-SAM without sacrificing efficiency. In this paper, we introduce a novel framework that integrates SAM with LoRA through dual LoRA modules. Our method employs auxiliary LoRA modules to generate adversarial perturbations decoupled from optimization, incurring minimal additional overhead. + +## Appendix F Proofs + +### F.1 Proof and Discussion of Proposition[2](https://arxiv.org/html/2508.19564#Thmproposition2 "Proposition 2 (Alignment of Bi-LoRA’s ascent direction with previous full gradient). ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"): Alignment of Bi-LoRA’s ascent direction with previous full gradient + +###### Proof. + +Let + +G_{t}=\nabla_{W}\mathcal{L}(W_{0}+B_{1,t}A_{1,t}+\tilde{\epsilon}_{t}),(A4) + +denote the full gradient at step t, with \tilde{\epsilon}_{t}=B_{2,t}A_{2,t} the auxiliary perturbation. + +One update step for Bi-LoRA is formalized as + +\begin{cases}B_{1,t+1}=B_{1,t}-\eta_{1}G_{t}A_{1,t}^{\top},&A_{1,t+1}=A_{1,t}-\eta_{1}B_{1,t}^{\top}G_{t},\\ +B_{2,t+1}=B_{2,t}+\eta_{2}G_{t}A_{2,t}^{\top},&A_{2,t+1}=A_{2,t}+\eta_{2}B_{2,t}^{\top}G_{t},\end{cases}(A5) + +where \eta_{1},\eta_{2}>0 are step sizes. + +Expanding the auxiliary perturbation Eqn.([A5](https://arxiv.org/html/2508.19564#A6.E5 "In Proof. ‣ F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) gives + +\tilde{\epsilon}_{t+1}=B_{2,t+1}A_{2,t+1}=\tilde{\epsilon}_{t}+\eta_{2}\left(G_{t}A_{2,t}^{\top}A_{2,t}+B_{2,t}B_{2,t}^{\top}G_{t}\right)+\mathcal{O}(\eta_{2}^{2}).(A6) + +Thus, the alignment between the previous gradient Eqn.Eqn.([A4](https://arxiv.org/html/2508.19564#A6.E4 "In Proof. ‣ F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) and the update of the auxiliary perturbation is + +\langle G_{t},\tilde{\epsilon}_{t+1}-\tilde{\epsilon}_{t}\rangle=\eta_{2}\Big(\|A_{2,t}G_{t}^{\top}\|_{F}^{2}+\|B_{2,t}^{\top}G_{t}\|_{F}^{2}\Big)+\mathcal{O}(\eta_{2}^{2})\;\;\geq 0,(A7) + +since both terms on the right-hand side are nonnegative (squared Frobenius norms). + +Therefore, Eqn.by Eqn.([A7](https://arxiv.org/html/2508.19564#A6.E7 "In Proof. ‣ F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), the auxiliary update (B_{2},A_{2}) always increases the inner objective in the direction of the full gradient Eqn.Eqn.([A4](https://arxiv.org/html/2508.19564#A6.E4 "In Proof. ‣ F.1 Proof and Discussion of Proposition 2: Alignment of Bi-LoRA’s ascent direction with previous full gradient ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")). This proves that the Bi-LoRA perturbation aligns with previous SAM’s perturbation direction. ∎ + +As reported in several efficient SAM variants, reusing the previous ascent direction can still yield strong generalization (e.g., S^{2} SAM(Ji et al., [2024](https://arxiv.org/html/2508.19564#bib.bib31 "A single-step, sharpness-aware minimization is all you need to achieve efficient and accurate sparse training")) in the sparse training scenario). In general, aligning with the previous, not the current direction, has benefit on efficiency. The reason that aligning with the previous gradient works may be attributed to the low-dimension/low-rank properties of the gradients(Li et al., [2022](https://arxiv.org/html/2508.19564#bib.bib74 "Low dimensional trajectory hypothesis is true: dnns can be trained in tiny subspaces"); Zhao et al., [2024](https://arxiv.org/html/2508.19564#bib.bib75 "GaLore: memory-efficient LLM training by gradient low-rank projection")), especially in later stage of training, or in fine-tuning stage. Due to the low-dimension/low-rank property, the consecutive gradient directions tend to exhibit similarity. And the intuition behind Bi-LoRA is that adversarial perturbations across different steps are "continuous": the optimal adversarial perturbation at iteration t is likely to be close to the optimal perturbation at iteration t+1, given that the weight differences between the continuous steps are small under LoRA optimization. This means we can obtain the adversarial perturbation at step t+1 with a "slight" adjustment based on the perturbation at t. As a result, this method can effectively optimize the sharpness while maintaining the same training efficiency as regular LoRA training. + +### F.2 Proof and discussions of Proposition[3](https://arxiv.org/html/2508.19564#Thmproposition3 "Proposition 3 (Bi-LoRA is a Regularized LoRA). ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"): Equivalence between SAM and Bi-LoRA + +###### Proof. + +SAM solves the min–max objective + +\min_{W}\;\max_{\|\epsilon\|_{2}\leq\rho}\;\mathcal{L}(W+\epsilon),(A8) + +where \rho>0 controls the perturbation radius. + +A first-order Taylor expansion gives + +\mathcal{L}(W+\epsilon)\approx\mathcal{L}(W)+\langle\nabla_{W}\mathcal{L}(W),\epsilon\rangle.(A9) + +The inner maximization is + +\max_{\|\epsilon\|_{2}\leq\rho}\;\langle\nabla_{W}\mathcal{L}(W),\epsilon\rangle,(A10) + +which is attained at + +\epsilon^{\star}=\rho\,\frac{\nabla_{W}\mathcal{L}(W)}{\|\nabla_{W}\mathcal{L}(W)\|_{2}}.(A11) + +Substituting Eqn.([A11](https://arxiv.org/html/2508.19564#A6.E11 "In Proof. ‣ F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) into Eqn.([A8](https://arxiv.org/html/2508.19564#A6.E8 "In Proof. ‣ F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) yields the regularized objective + +\mathcal{L}(W)+\rho\|\nabla_{W}\mathcal{L}(W)\|_{2}.(A12) + +For Bi-LoRA, the weight matrix is decomposed as W=W_{0}+B_{1}A_{1}+B_{2}A_{2}. The training objective is + +\min_{B_{1},A_{1}}\;\max_{\begin{subarray}{c}B_{2},A_{2}\\ +\mathrm{rank}(B_{2}A_{2})\leq r,\;\|B_{2}A_{2}\|_{F}\leq\rho\end{subarray}}\;\mathcal{L}(W_{0}+B_{1}A_{1}+B_{2}A_{2}),(A13) + +where the perturbation is explicitly parameterized as + +\tilde{\epsilon}=B_{2}A_{2},\quad\mathrm{rank}(\tilde{\epsilon})\leq r,\quad\|\tilde{\epsilon}\|_{F}\leq\rho.(A14) + +At \theta=W_{0}+B_{1}A_{1}, expanding gives + +\mathcal{L}(\theta+\tilde{\epsilon})\approx\mathcal{L}(\theta)+\langle\nabla_{\theta}\mathcal{L},\tilde{\epsilon}\rangle.(A15) + +Thus the inner maximization becomes + +\max_{\begin{subarray}{c}\mathrm{rank}(\tilde{\epsilon})\leq r\\ +\|\tilde{\epsilon}\|_{F}\leq\rho\end{subarray}}\;\langle\nabla_{\theta}\mathcal{L},\tilde{\epsilon}\rangle.(A16) + +Let the SVD of \nabla_{\theta}\mathcal{L} be + +\nabla_{\theta}\mathcal{L}=U\Sigma V^{\top},\qquad\sigma_{1}\geq\sigma_{2}\geq\cdots.(A17) + +The optimum of Eqn.([A16](https://arxiv.org/html/2508.19564#A6.E16 "In Proof. ‣ F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) is achieved at + +\tilde{\epsilon}^{\star}=\rho\,U_{r}V_{r}^{\top},(A18) + +where U_{r},V_{r} are the top-r singular vectors of \nabla_{\theta}\mathcal{L}, and \mathrm{rank}(\tilde{\epsilon}^{\star})=r with \|\tilde{\epsilon}^{\star}\|_{F}=\rho. + +Substituting Eqn.([A18](https://arxiv.org/html/2508.19564#A6.E18 "In Proof. ‣ F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) into Eqn.([A16](https://arxiv.org/html/2508.19564#A6.E16 "In Proof. ‣ F.2 Proof and discussions of Proposition 3: Equivalence between SAM and Bi-LoRA ‣ Appendix F Proofs ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) gives + +\langle\nabla_{\theta}\mathcal{L},\tilde{\epsilon}^{\star}\rangle=\rho\sum_{i=1}^{r}\sigma_{i}(\nabla_{\theta}\mathcal{L})=\rho\|\nabla_{\theta}\mathcal{L}\|_{(r)},(A19) + +where + +\|\nabla_{\theta}\mathcal{L}\|_{(r)}:=\sum_{i=1}^{r}\sigma_{i}(\nabla_{\theta}\mathcal{L})(A20) + +is the Ky Fan r-norm. + +Hence, Bi-LoRA reduces to the regularized objective + +\min_{A_{1},B_{1}}\;\mathcal{L}(W_{0}+B_{1}A_{1})+\rho\|\nabla_{W_{0}+B_{1}A_{1}}\mathcal{L}\|_{(r)}.(A21) + +∎ + +When r covers all singular values, the Ky Fan norm becomes the nuclear norm, and Bi-LoRA recovers the SAM objective. Conversely, SAM specializes to Bi-LoRA when only the sharpest r singular directions are penalized. + +And Proposition[3](https://arxiv.org/html/2508.19564#Thmproposition3 "Proposition 3 (Bi-LoRA is a Regularized LoRA). ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") represents the ideal objective of Bi-LoRA and serves to illustrate our motivation intuitively. In practice, however, we adopt a simultaneous optimization approach formalized in Eqn.([9](https://arxiv.org/html/2508.19564#S3.E9 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")), which does not introduce additional gradient steps and thus maintains efficiency. Experimental results demonstrate that our approximate approach can already significantly improve performance compared to LoRA. + +Furthermore, Figure[3](https://arxiv.org/html/2508.19564#S3.F3 "Figure 3 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that, the auxiliary LoRA module continues broad exploration, while the primary module converges rapidly to a nearly fixed point; once the primary stabilizes, the auxiliary then begins to converge quickly. This behavior suggests that at each parameter point, there exists a (though inefficient) convergent multi-step scheme for the auxiliary module, which is consistent with the intuition of Proposition[3](https://arxiv.org/html/2508.19564#Thmproposition3 "Proposition 3 (Bi-LoRA is a Regularized LoRA). ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +## Appendix G Reasons of Using Global Clipping + +We clip using the _total_ Frobenius norm over all auxiliary LoRA modules in Bi-LoRA for the following reasons. + +* • +Faithfulness to SAM. The SAM-style inner maximization is defined with a _single_\rho-ball over the full parameter vector, i.e., \max_{\lVert\epsilon\rVert\leq\rho}\mathcal{L}(w+\epsilon). In Bi-LoRA, the perturbation is implemented entirely through the auxiliary LoRA modules, so constraining the _global_ Frobenius norm of all auxiliary modules (Eqn.([8](https://arxiv.org/html/2508.19564#S3.E8 "In 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) is the faithful analogue of SAM’s original formulation. + +* • +Avoiding scale explosion and instability from per-layer normalization. If each layer is normalized independently, every layer receives a perturbation with norm \rho, so the overall perturbation magnitude roughly scales with the number of adapters. This not only leads to training instability, but also makes the same \rho incomparable across architectures with different depths or adapter placements. SAM’s original global constraint is designed precisely to avoid such architecture-dependent scaling effects. + +* • +Empirical evidence. In Table[5](https://arxiv.org/html/2508.19564#S4.T5 "Table 5 ‣ 4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), global clipping consistently outperforms both per-layer clipping and no clipping in terms of final performance and training stability under the same \rho. + +In summary, global clipping over all auxiliary LoRA modules yields a scale-aware neighborhood that is consistent with SAM’s objective, comparable across architectures, and empirically more stable, which motivates our choice of this clipping scheme. + +## Appendix H More baselines with Llama 3.1-8B on instruction following tasks + +Table A1: Results of fine-tuning Llama 3.1-8B on instruction-following tasks with more LoRA-variants and SAM-based methods. + +Method MMLU DROP HEval BBH Avg +LoRA 63.38±0.39 49.82±0.54 43.15±0.93 42.82±0.27 49.79 +PiSSA 63.59±0.14 50.20±0.20 43.86±0.12 43.25±0.51 50.22 +DoRA 63.58±0.22 50.53±0.10 44.10±0.73 42.98±0.79 50.30 +DeLoRA 63.49±0.32 51.40±0.28 45.53±0.73 42.72±0.29 50.78 +HiRA 63.62±0.33 51.17±0.26 45.12±0.61 42.89±0.20 50.70 +WSAM 63.42±0.12 50.63±0.14 42.88±0.81 43.16±0.52 50.02 +S^{2}-SAM 63.73±0.38 50.80±0.23 43.09±0.20 42.44±0.19 50.02 +Bi-LoRA 63.67±0.15 51.53±0.33 46.12±0.89 43.45±0.15 51.19 + +Table[A1](https://arxiv.org/html/2508.19564#A8.T1 "Table A1 ‣ Appendix H More baselines with Llama 3.1-8B on instruction following tasks ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") further compares Bi-LoRA with four LoRA variants (PiSSA, DoRA, HiRA, DeLoRA) and two SAM-based methods (WSAM, S^{2}-SAM) on Llama 3.1-8B instruction following tasks. Bi-LoRA attains the highest average performance of 51.19, outperforming the strongest LoRA variant by about 0.4% and vanilla LoRA by roughly 1.4% points. + +## Appendix I Results on Natural Language Understanding + +Setting. We fine-tune the T5-base model(Raffel et al., [2020](https://arxiv.org/html/2508.19564#bib.bib49 "Exploring the limits of transfer learning with a unified text-to-text transformer")) on multiple datasets from GLUE(Wang et al., [2019b](https://arxiv.org/html/2508.19564#bib.bib51 "GLUE: a multi-task benchmark and analysis platform for natural language understanding")) and SuperGLUE(Wang et al., [2019a](https://arxiv.org/html/2508.19564#bib.bib52 "Superglue: a stickier benchmark for general-purpose language understanding systems")) benchmarks, including MNLI, SST2, CoLA, QNLI, MRPC, BoolQ, CB, COPA, RTE, and WIC, following Wang et al. ([2024](https://arxiv.org/html/2508.19564#bib.bib45 "LoRA-ga: low-rank adaptation with gradient approximation")); Li et al. ([2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape")). Performance is evaluated on the development set using accuracy as the primary metric, except for CoLA, where the Matthews correlation coefficient is used. + +Results. We first focus on the GLUE datasets. From the results in Table[A2](https://arxiv.org/html/2508.19564#A9.T2 "Table A2 ‣ Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), we observe that Bi-LoRA outperforms both LoRA and LoRA-SAM, achieving an average improvement of 0.47% and 0.32%, respectively. It is worth noting that Bi-LoRA achieves these gains with the same training speed as LoRA, whereas LoRA-SAM doubles the training time. Moreover, the gains are more pronounced on smaller datasets, with Bi-LoRA outperforming LoRA by 1.36% on CoLA and 0.82% on MRPC. In contrast, LoRA-SAM shows no clear advantage over vanilla LoRA, indicating its limited ability to enhance generalization due to constrained sharpness optimization. + +Next, we evaluate on the SuperGLUE datasets, which feature more challenging language understanding tasks. To reduce variance from the limited evaluation sizes of CB (56), COPA (100), and RTE (277) samples, we increase the number of runs from 3 to 20 for CB and COPA, and to 5 for RTE. Table[A3](https://arxiv.org/html/2508.19564#A9.T3 "Table A3 ‣ Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows Bi-LoRA demonstrates more significant advantages over LoRA and LoRA-SAM by 0.69% and 0.60% on average, while LoRA-SAM yields minimal gains and even hurts some datasets (e.g., CB). These results confirm the effectiveness of Bi-LoRA in enhancing generalization. + +Table A2: Results on fine-tuning the T5-base model on a subset of GLUE datasets. “Cost” indicates the gradient steps per training iteration, e.g., one step (Cost \times 1) for Full FT, LoRA, and Bi-LoRA. + +Dataset Size Cost MNLI (393k)SST2 (67k)CoLA (8.6k)QNLI (105k)MRPC (3.7k)Avg. +Full FT\times 1 85.57±0.09 94.27±0.24 56.60±0.62 93.18±0.09 87.30±0.79 83.38 +LoRA\times 1 86.25±0.16 94.23±0.30 59.41±0.52 93.25±0.06 88.56±0.26 84.34 +LoRA-SAM\times 2 86.25±0.09 94.46±0.17 59.80±0.85 93.21±0.16 88.73±0.52 84.49 +Bi-LoRA\times 1 86.33±0.08 94.34±0.05 60.77±0.39 93.25±0.06 89.38±0.26 84.81 + +Table A3: Results on fine-tuning T5-base with a subset of SuperGLUE datasets. “Cost” indicates the gradient steps per training iteration, e.g., one step (Cost \times 1) for Full FT, LoRA, and Bi-LoRA. + +Dataset Size Cost BoolQ (9.4k)CB (0.25k)COPA (0.4k)RTE (2.5k)WIC (5.4k)Avg. +Full FT\times 1 72.19±0.14 92.26±0.24 64.67±0.26 84.48±0.12 68.34±0.19 76.39 +LoRA\times 1 72.20±0.21 92.86±0.27 63.80±0.23 83.10±0.29 68.60±0.32 76.11 +LoRA-SAM\times 2 72.47±0.42 92.32±0.28 64.20±0.42 83.47±0.47 68.55±0.76 76.20 +Bi-LoRA\times 1 72.25±0.30 92.86±0.32 64.60±0.25 83.61±0.37 70.69±0.40 76.80 + +## Appendix J Hyperparameter sensitivity + +### J.1 Effect of the auxiliary rank on CoLA + +![Image 11: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/effect_rank2_cola_t5base.png) + +Figure A1: Performance on CoLA under varying ranks of primary (r_{1}) and auxiliary (r_{2}) LoRA modules. + +To further examine the sensitivity to the auxiliary rank r_{2}, we fine-tune T5-base on the CoLA dataset with r_{1}\in\{4,8,16\} and r_{2}\in\{2,4,8,16,32,64\}. + +Figure[A1](https://arxiv.org/html/2508.19564#A10.F1 "Figure A1 ‣ J.1 Effect of the auxiliary rank on CoLA ‣ Appendix J Hyperparameter sensitivity ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows the performance of Bi-LoRA with varying the primary (r_{1}) and auxiliary (r_{2}) LoRA ranks. We observe that integrating the auxiliary module consistently enhances performance. Notably, r_{2}=8 generally delivers good results, which replicates the pattern we reported in Section[4.6](https://arxiv.org/html/2508.19564#S4.SS6 "4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), confirming that the sweet‑spot generalizes from NLU tasks to instruction‑tuning. Therefore, we suggest using r_{2}=8 as a task-agnostic default. + +### J.2 Effect of the neighborhood radius \rho + +From Table[A4](https://arxiv.org/html/2508.19564#A10.T4 "Table A4 ‣ J.2 Effect of the neighborhood radius 𝜌 ‣ Appendix J Hyperparameter sensitivity ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), Bi-LoRA performs consistently well across \rho=\{0.01,0.05,0.1\}, showing its effectiveness to moderate perturbations during training. The ability to achieve strong performance across a range of \rho indicates that Bi-LoRA can effectively adapt to varying levels of perturbation, demonstrating its practicality for diverse applications. + +Table A4: Results on CoLA, SST2 and GSM8K under different values of the neighborhood radius \rho. + +\rho CoLA SST2 GSM8K +0.01 60.15±0.36 60.15±0.36 59.36 +0.05 60.77±0.39 60.15±0.36 59.74 +0.1 60.18±0.59 60.15±0.36 60.65 +0.2 58.95±0.81 60.15±0.36 58.83 + +## Appendix K Qualitative Results on SDXL + +As discussed in Section[4.4](https://arxiv.org/html/2508.19564#S4.SS4 "4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), Bi-LoRA improves both image–text (I2T) and text–text (T2T) CLIP similarity on the 3D Icons dataset. To visualize these gains, Figure[A2](https://arxiv.org/html/2508.19564#A11.F2 "Figure A2 ‣ Appendix K Qualitative Results on SDXL ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") presents generated images from SDXL models fine-tuned by LoRA and Bi-LoRA under identical promptss. + +Observations. Across instances, Bi-LoRA tends to preserve the intended icon identity and style more faithfully. These qualitative differences align with the measured CLIP improvements, i.e., larger average CLIP I2T/T2T scores shown in Table[3](https://arxiv.org/html/2508.19564#S4.T3 "Table 3 ‣ 4.4 Results on Diffusion Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") also exhibit clearer textual correspondence in Figure[A2](https://arxiv.org/html/2508.19564#A11.F2 "Figure A2 ‣ Appendix K Qualitative Results on SDXL ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +![Image 12: Refer to caption](https://arxiv.org/html/2508.19564v2/x2.png) + +(a) LoRA - Fluffy Rabbit + +![Image 13: Refer to caption](https://arxiv.org/html/2508.19564v2/x3.png) + +(b) Bi-LoRA - Fluffy Rabbit + +![Image 14: Refer to caption](https://arxiv.org/html/2508.19564v2/x4.png) + +(c) LoRA - Curly Cloud + +![Image 15: Refer to caption](https://arxiv.org/html/2508.19564v2/x5.png) + +(d) Bi-LoRA - Curly Cloud + +![Image 16: Refer to caption](https://arxiv.org/html/2508.19564v2/x6.png) + +(e) LoRA - Roaring Lion + +![Image 17: Refer to caption](https://arxiv.org/html/2508.19564v2/x7.png) + +(f) Bi-LoRA - Roaring Lion + +![Image 18: Refer to caption](https://arxiv.org/html/2508.19564v2/x8.png) + +(g) LoRA - Dripping Water Drop + +![Image 19: Refer to caption](https://arxiv.org/html/2508.19564v2/x9.png) + +(h) Bi-LoRA - Dripping Water Drop + +![Image 20: Refer to caption](https://arxiv.org/html/2508.19564v2/x10.png) + +(i) LoRA - Hopping Frog + +![Image 21: Refer to caption](https://arxiv.org/html/2508.19564v2/x11.png) + +(j) Bi-LoRA - Hopping Frog + +![Image 22: Refer to caption](https://arxiv.org/html/2508.19564v2/x12.png) + +(k) LoRA - Smiling Sun + +![Image 23: Refer to caption](https://arxiv.org/html/2508.19564v2/x13.png) + +(l) Bi-LoRA - Smiling Sun + +Figure A2: Images generated with SDXL fine-tuned using LoRA (left) and Bi-LoRA (right) on the 3D icon dataset. Each row corresponds to a different prompt in the form “a ToK icon of a , in the style of ToK”. Images in the same row are generated with the same random seed for fair comparison. + +## Appendix L Convergence Analysis of Main and Auxiliary LoRA Modules + +We measure the cosine similarity between the weights of the main and auxiliary LoRA modules and their inal weights on CoLA and SST2 with T5-base. The results indicate that the auxiliary module converges substantially slower than the main one, preserving flexibility for sharpness-aware updates. + +![Image 24: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/cola_wd0.0_lr5e-4_bilora.png) + +(a) CoLA: \eta_{2}=\eta_{1}=5e-4 + +![Image 25: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/cola_wd0.0_lr5e-3_bilora.png) + +(b) CoLA: \eta_{2}=10\eta_{1}=5e-3 + +![Image 26: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/sst2_wd0.0_lr5e-4_bilora.png) + +(c) SST2: \eta_{2}=\eta_{1}=5e-4 + +![Image 27: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/sst2_wd0.0_lr5e-3_bilora.png) + +(d) SST2: \eta_{2}=10\eta_{1}=5e-3 + +Figure A3: Cosine similarity between the main and auxiliary LoRA weights and each final weights during training on CoLA and SST2 using T5-base, under different auxiliary learning rates (\eta_{2}). The auxiliary module (perturbation, blue lines) converges in the final 20% of steps, significantly slower than the main (optimization, red lines) module. A larger \eta_{2}=10\eta_{1} further slows down convergence. Moreover, the auxiliary module’s trajectory remains independent of the main module throughout training (green lines). + +## Appendix M Bi-LoRA in Quantized Settings + +We apply Bi-LoRA in a quantized setting, where the transition to lower precision can cause the model to jump from a flat region into a sharper region with higher local loss, reducing generalization on unseen data. Following QLoRA(Dettmers et al., [2023](https://arxiv.org/html/2508.19564#bib.bib33 "QLoRA: efficient finetuning of quantized llms")), we quantize the pretrained model to NF4 and adopt the setup described in Section 5.3. + +As shown in Table[A5](https://arxiv.org/html/2508.19564#A13.T5 "Table A5 ‣ Appendix M Bi-LoRA in Quantized Settings ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), Bi-QLoRA outperforms QLoRA and QLoRA-SAM by 2.58% and 2.36% on GSM8K, respectively. This gain is attributed to Bi-LoRA’s ability to maintain flatness under quantization by decoupling optimization and perturbation during training. + +Table A5: Performance on GSM8K under quantized NF4 settings. Results are averaged over 3 runs. + +Method GSM8K (%) +QLoRA 57.77±0.39 +QLoRA-SAM 57.89±0.29 +Bi-QLoRA 59.26±0.17 + +## Appendix N Effect on Adversarial Robustness + +Figure[4](https://arxiv.org/html/2508.19564#S3.F4 "Figure 4 ‣ 3 Bi-LoRA: Bi-directional Low-Rank Adaptation ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") already shows that Bi‑LoRA is more robust than vanilla LoRA against random perturbations applied both in the LoRA subspace and in the full parameter space. We now investigate robustness to adversarial perturbations. Concretely, we fine‑tune T5‑base on CoLA and report Matthew’s Correlation Coefficient. Following (Kim et al., [2022](https://arxiv.org/html/2508.19564#bib.bib4 "Fisher sam: information geometry and sharpness aware minimisation")), we apply the worst‑case parameter perturbation + +\theta\rightarrow\theta+\alpha\frac{\nabla_{\theta}\mathcal{L}}{\left\|\nabla_{\theta}\mathcal{L}\right\|},(A22) + +evaluating attacks restricted to (i) the LoRA subspace (\theta\in\{B,A\}) and (ii) the full parameter space (\theta=W), with perturbation strength \alpha from 0.5 to 5.0. Results are shown in Table[A6](https://arxiv.org/html/2508.19564#A14.T6 "Table A6 ‣ Appendix N Effect on Adversarial Robustness ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +In the LoRA subspace, Bi-LoRA consistently suffers smaller performance degradation than vanilla LoRA. LoRA-SAM shows the strongest resilience, which aligns with its explicit optimization of sharpness in this subspace. In the full parameter space, degradation is more substantial across all methods, yet Bi-LoRA still mitigates the impact more effectively than LoRA-SAM. + +These findings are consistent with our random perturbation results, demonstrating that Bi-LoRA provides robustness to adversarial parameter perturbations, particularly when the attack targets the full parameter space. + +Table A6: Results on CoLA with adversarial parameter perturbation. Perturbation strength \alpha from 0 to 5.0. + +Perturb in LoRA space +Method\alpha=0\alpha=0.5\alpha=1.0\alpha=2.5\alpha=4.0\alpha=5.0 +LoRA 60.62 59.42 57.47 42.07 24.22 19.19 +LoRA-SAM 60.95 59.14 58.44 52.12 41.33 30.61 +Bi-LoRA 61.12 59.36 59.14 48.47 33.24 24.20 +Perturb in full space +Method\alpha=0\alpha=0.5\alpha=1.0\alpha=2.5\alpha=4.0\alpha=5.0 +LoRA 60.62 52.69 44.42 19.08 11.39 6.55 +LoRA-SAM 60.95 56.16 45.67 24.77 18.02 13.34 +Bi-LoRA 61.12 56.29 49.00 32.91 24.83 18.06 + +## Appendix O Training and test metric curves + +We fine-tune T5-base on MRPC for 10 epochs, tracking both training and eval losses and metrics. The results in [A4](https://arxiv.org/html/2508.19564#A15.F4 "Figure A4 ‣ Appendix O Training and test metric curves ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows nearly identical training losses across methods, yet Bi‑LoRA attains a lower validation loss and higher accuracy. Importantly, the smaller train–val generalization gap demonstrates Bi-LoRA’s improved generalization. + +![Image 28: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/train_loss_plot.png) + +(a) Training Loss + +![Image 29: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/eval_loss_plot.png) + +(b) Evaluation Loss + +![Image 30: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/train_accuracy_plot.png) + +(c) Training Accuracy + +![Image 31: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/eval_accuracy_plot.png) + +(d) Evaluation Accuracy + +![Image 32: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/generalization_gap.png) + +(e) Generalization Gap + +Figure A4: Training and evaluation loss, accuracy, and generalization gap of T5-base fine-tuned on the MRPC dataset. + +## Appendix P Compute Resources + +We utilize two types of GPUs: the NVIDIA RTX 4090 24GB and the A100 80GB. For NLU and Text-to-Image generation tasks, computations are performed on a single RTX 4090. All other experiments are conducted on a single A100. + +## Appendix Q Learning Rates for Tuning the Two Modules in Bi-LoRA + +In our experiments, both modules in Bi-LoRA use the same learning rate of 5\times 10^{-4}. Initially, we explored an unequal learning rate strategy. However, results in Table[A7](https://arxiv.org/html/2508.19564#A17.T7 "Table A7 ‣ Appendix Q Learning Rates for Tuning the Two Modules in Bi-LoRA ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") show that tuning \eta_{2} provides only a limited improvement in performance. Considering the trade-off between gains and the additional effort for hyperparameter tuning, we recommend using the same learning rate for both modules. + +Table A7: CoLA accuracy under different learning rate of the auxiliary module \eta_{2} settings (with \eta_{1}=5{\times}10^{-4} for the main LoRA module). + +\eta_{2}1\times 10^{-4}5\times 10^{-4} (main)5\times 10^{-3} +CoLA 60.23±0.35 60.77±0.55 60.44±0.87 + +## Appendix R Empirical Details + +This section provides detailed empirical setups used for the previous experiments. + +We tune the neighborhood radius \rho over \{0.005, 0.01, 0.05, 0.1, 0.2, 0.5\} for LoRA-SAM and Bi-LoRA and search learning rates among \{2e-5, 5e-5, 1e-4, 2e-4\} for Full FT and \{5e-5, 1e-4, 2e-4, 5e-4, 1e-3\} for LoRA and its variants. + +### R.1 Experiments on the GLUE and SuperGLUE datasets + +The hyperparameter configurations for all GLUE and SuperGLUE experiments (Sections[I](https://arxiv.org/html/2508.19564#A9 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [4.5](https://arxiv.org/html/2508.19564#S4.SS5 "4.5 Integration with Other LoRA Variants ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [4.6](https://arxiv.org/html/2508.19564#S4.SS6 "4.6 Ablations and Hyperparameter Sensitivity ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) are detailed in Tables[A8](https://arxiv.org/html/2508.19564#A18.T8 "Table A8 ‣ R.1 Experiments on the GLUE and SuperGLUE datasets ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [A9](https://arxiv.org/html/2508.19564#A18.T9 "Table A9 ‣ R.1 Experiments on the GLUE and SuperGLUE datasets ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). As discussed in Section[I](https://arxiv.org/html/2508.19564#A9 "Appendix I Results on Natural Language Understanding ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), we set the neighborhood radius \rho to 0.01 and 0.05 for LoRA-SAM and Bi-LoRA, respectively, in the GLUE experiments. For SuperGLUE, we use \rho values of 0.05 and 0.1 for LoRA-SAM and Bi-LoRA, respectively. For full fine-tuning, we set the learning rate to 1e-4 while keeping all other hyperparameters unchanged. We use the same learning rate of 5e-4 for both the primary and auxiliary LoRA modules in Bi-LoRA. + +Table A8: Hyperparameter configurations for fine-tuning T5-base using LoRA-based methods on the GLUE datasets. + +Hyperparameter MNLI SST-2 COLA QNLI MRPC +Learning Rate 5e-4 +Batch Size 32 +Epochs 3 3 10 3 10 +Max Sequence Length 256 +LoRA Rank 8 +LoRA Alpha 16 +LR Scheduler Cosine +Target Modules All +Warmup Ratio 0.03 +Evaluation Metric Accuracy Accuracy Matthews Corr.Accuracy Accuracy + +Table A9: Hyperparameter settings for fine-tuning T5-base using LoRA-based methods on the SuperGLUE datasets. + +Hyperparameter BoolQ CB COPA RTE WIC +Learning Rate 5e-4 +Batch Size 32 +Epochs 10 50 50 10 10 +Max Sequence Length 256 +LoRA Rank 8 +LoRA Alpha 16 +LR Scheduler Cosine +Target Modules All +Warmup Ratio 0.03 +Evaluation Metric Accuracy + +### R.2 Experiments on Llama Models + +The hyperparameter settings for Llama 2-7B and Llama 3.1-8B are listed in Tables[A10](https://arxiv.org/html/2508.19564#A18.T10 "Table A10 ‣ R.2 Experiments on Llama Models ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [A11](https://arxiv.org/html/2508.19564#A18.T11 "Table A11 ‣ R.2 Experiments on Llama Models ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +The neighborhood radius \rho is set to 0.01 and 0.05 for LoRA-SAM and Bi-LoRA. For full fine-tuning, we set the micro-batch size to 4, 2, 2, and 8 for the respective four tasks, and the learning rate is set to 5e-5. We use the same learning rate of 5e-4 for both the primary and auxiliary LoRA modules in Bi-LoRA. + +Table A10: Hyperparameter configurations for fine-tuning Llama 2-7B on Math, Code, and Chat tasks using LoRA-based methods. “Math”, “Code” and “Chat” correspond to mathematical reasoning, code generation, and dialogue generation tasks, respectively, as described in Section[4.2](https://arxiv.org/html/2508.19564#S4.SS2 "4.2 Results on Llama Models ‣ 4 Experiments ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +Hyperparameter Math Code Chat +Learning rate 5e-4 +Batch Size 32 +Micro-Batch size 4 +Epochs 2 +Max Sequence Length 1024 +LoRA Rank 8 +LoRA Dropout 0.0 +LoRA Alpha 16 +Target Modules All +LR Scheduler Cosine +Warmup Ratio 0.03 + +Table A11: Hyperparameter configurations for fine-tuning Llama 3.1-8B on Cleaned Alpaca dataset. + +Hyperparameter Cleaned Alpaca +Learning rate 3e-4 +Batch Size 128 +Micro-Batch size 4 +Epochs 3 +Max Sequence Length 256 +LoRA Rank 8 +LoRA Dropout 0.0 +LoRA Alpha 16 +Target Modules All +LR Scheduler Cosine +Warmup Ratio 0.03 + +### R.3 Experiments on Qwen model + +We first grid-search the learning rate \in\{1\mathrm{e}-5,5\mathrm{e}-5,1\mathrm{e}-4,3\mathrm{e}-4,5\mathrm{e}-4\} for LoRA, Bi-LoRA and LoRA-SAM, and perturbation radius \rho\in\{0.1,0.05,0.01\}. For all random noise baselines we tune the variance \rho^{2}\in\{0.1,0.01,0.005,0.001\}. + +We set the perturbation radius to \rho=0.1 for Bi-LoRA and \rho=0.05 for LoRA-SAM and other SAM variants, use variance \rho^{2}=0.001 for all random-noise baselines. Other hyperparameters are summarized in Table[A12](https://arxiv.org/html/2508.19564#A18.T12 "Table A12 ‣ R.3 Experiments on Qwen model ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). + +Table A12: Hyperparameter configurations for fine-tuning Qwen 2.5-14B on Cleaned Alpaca dataset. + +Hyperparameter Cleaned Alpaca +Learning rate 5e-5 +Batch Size 128 +Micro-Batch size 4 +Epochs 3 +Max Sequence Length 256 +LoRA Rank 8 +LoRA Dropout 0.0 +LoRA Alpha 16 +Target Modules All +LR Scheduler Cosine +Warmup Ratio 0.03 + +### R.4 Experiments on Diffusion Models + +The hyperparameter settings for Diffusion Models are listed in Table[A13](https://arxiv.org/html/2508.19564#A18.T13 "Table A13 ‣ R.4 Experiments on Diffusion Models ‣ Appendix R Empirical Details ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"). We finetune the model for 500 steps with a constant learning rate of 2e-4, and set the \rho for Bi-LoRA to 0.05. The LoRA rank is set to 4, with an auxiliary rank of 4 for Bi-LoRA. We use the same learning rate of 2e-4 for both the primary and auxiliary LoRA modules in Bi-LoRA. The instance prompt used for DreamBooth training is “a TOK icon, in the style of TOK”, and the validation prompt used for generating the specific icons is formatted as "a ToK icon of a , in the style of ToK". We follow the scripts implemented by Hugging Face 2 2 2[https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_sdxl.md](https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/README_sdxl.md). + +Table A13: Hyperparameter settings for fine-tuning SDXL using Dreambooth on the 3D Icons dataset + +Hyperparameter Value +Learning Rate 2e-4 +Batch Size 4 +Micro-Batch Size 1 +Train Steps 500 +Resolution 1024 +Validation Epochs 25 +LoRA Rank r 4 +LoRA Dropout 0.0 +LoRA Alpha \alpha 8 +Target Modules W_{K},W_{Q},W_{V},W_{O} +LR Scheduler Constant +Warmup Steps 0 + +### R.5 Efficient SAM Variants for LoRA Fine-Tuning + +To ensure a fair comparison, we adopt the results of Flat-LoRA from Li et al. ([2025](https://arxiv.org/html/2508.19564#bib.bib64 "Flat-lora: low-rank adaptation over a flat loss landscape")), and follow the default setup of Li et al. ([2024a](https://arxiv.org/html/2508.19564#bib.bib53 "Implicit regularization of sharpness-aware minimization for scale-invariant problems")) for LoRA-oBAR/nBAR with \alpha=0.25. And following Du et al. ([2022a](https://arxiv.org/html/2508.19564#bib.bib59 "Efficient sharpness-aware minimization for improved training of neural networks")), (Liu et al., [2022b](https://arxiv.org/html/2508.19564#bib.bib60 "Towards efficient and scalable sharpness-aware minimization")), LoRA-ESAM perturbs 50% of the parameters on the top-50% sharpness-sensitive data, while LoRA-LookSAM applies SAM perturbation every five steps. All other hyperparameters are aligned with those used in our prior experiments. + +## Appendix S Time Comparisons Across Tasks + +In this section, we compare the per-step optimization time for LoRA, LoRA-SAM, and Bi-LoRA across a range of benchmarks. The results are organized based on different datasets and tasks, including T5 experiments on the GLUE and SuperGLUE datasets, as well as Llama experiments on domain-specific datasets. All experiments were conducted on a single NVIDIA RTX 4090 GPU 24GB, using the same hyperparameters as in the previous experiments. + +Table A14: Optimization time (s) per step for T5-base on GLUE benchmarks. The average time across tasks is also reported. + +Dataset MNLI SST2 CoLA QNLI MRPC Avg. +Time +Full FT 0.758 0.538 0.549 0.756 0.681 0.654 +LoRA 0.458 0.310 0.273 0.508 0.442 0.398 +LoRA-SAM 1.036 0.573 0.546 0.950 1.039 0.829 +Bi-LoRA 0.476 0.325 0.305 0.513 0.508 0.425 + +Table A15: Optimization time (s) per step for T5-base on SuperGLUE benchmarks. The average time across tasks is also reported. + +Dataset BoolQ CB COPA RTE WIC Avg. +Time +Full FT 0.622 0.415 0.451 0.689 0.544 0.544 +LoRA 0.276 0.276 0.253 0.263 0.286 0.271 +LoRA-SAM 0.519 0.534 0.568 0.688 0.484 0.559 +Bi-LoRA 0.304 0.285 0.278 0.276 0.262 0.281 + +Table A16: Optimization time (s) per step for training Llama2 on math (MetaMathQA), code (Code-Feedback) and chat (WizardLM) tasks. + +Method MetaMathQA Code-Feedback WizardLM +LoRA 2.81 5.10 5.13 +LoRA-SAM 5.52 10.10 10.53 +ROP 2.88 5.41 5.24 +RWP_full 3.21 6.02 5.88 +RWP_LoRA 2.97 5.56 5.38 +LoRA-oBAR 2.90 5.41 5.21 +LoRA-nBAR 2.89 5.41 5.24 +Flat-LoRA 3.42 5.92 5.78 +LoRA-ESAM 4.20 8.13 7.35 +LoRA-LookSAM 3.47 6.49 6.29 +Bi-LoRA 3.12 5.68 5.43 + +Table A17: Optimization time (s) per step for training Llama 3-8B on the instruction-following dataset Cleaned Alpaca. + +Method Cleaned Alpaca +LoRA 8.93 +LoRA-SAM 19.23 +ROP 9.43 +RWP_full 12.35 +RWP_LoRA 10.10 +LoRA-oBAR 9.35 +LoRA-nBAR 9.35 +Flat-LoRA 9.90 +LoRA-ESAM 13.44 +LoRA-LookSAM 11.49 +Bi-LoRA 9.71 + +Table A18: Optimization time (s) per step for training Qwen 2.5-14B on the instruction-following dataset Cleaned Alpaca. + +Method Cleaned Alpaca +LoRA 15.47 +LoRA-SAM 32.37 +ROP 15.87 +RWP_full 21.28 +RWP_LoRA 16.67 +LoRA-oBAR 15.87 +LoRA-nBAR 15.87 +Flat-LoRA 17.06 +LoRA-ESAM 21.79 +LoRA-LookSAM 19.07 +Bi-LoRA 16.63 + +## Appendix T Norm and Ratio Analysis Across Layers + +In addition to Figures[2(a)](https://arxiv.org/html/2508.19564#S2.F2.sf1 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [2(b)](https://arxiv.org/html/2508.19564#S2.F2.sf2 "In Figure 2 ‣ 2.1 Preliminaries ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), we extend the analysis of LoRA-SAM’s training dynamics to additional self-attention layers across different encoder blocks to verify the consistency of our observations. + +![Image 33: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_layer_25.png) + +(a) Norm of different components in the key of the first self-attention layer in the fifth encoder block + +![Image 34: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_ratio_layer_25.png) + +(b) Norm ratio of the key of the first self-attention layer in the fifth encoder block + +![Image 35: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_layer_45.png) + +(c) Norm of different components in the output projection the second self-attention layer in the eighth encoder block + +![Image 36: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_ratio_layer_45.png) + +(d) Norm ratio of the output projection in the second self-attention layer in the eighth encoder block + +![Image 37: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_layer_50.png) + +(e) Norm of different components in the value of the first self-attention layer in the ninth encoder block + +![Image 38: Refer to caption](https://arxiv.org/html/2508.19564v2/_figures/t5_cola_norm_ratio_layer_50.png) + +(f) Norm ratio of the value of the first self-attention layer in the ninth encoder block + +Figure A5: Analyses of training statistics for LoRA-SAM. (a) and (c) and (e): Frobenius norms of different terms. And (b), (d) and (f): ratio of the Frobenius norms of the first two terms (B\epsilon_{A}+\epsilon_{B}A) to that of the third term (\epsilon_{B}\epsilon_{A}) in Eqn.([5](https://arxiv.org/html/2508.19564#S2.E5 "In 2.2 LoRA-SAM ‣ 2 Issues of LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models")) during fine-tuning. Models are fine-tuned on CoLA with T5 for 10 epochs. + +Figures[5(a)](https://arxiv.org/html/2508.19564#A20.F5.sf1 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5(c)](https://arxiv.org/html/2508.19564#A20.F5.sf3 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [5(e)](https://arxiv.org/html/2508.19564#A20.F5.sf5 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") show the Frobenius norms of different terms in the key, output projection and value of self-attention layers in the 5th, 8th, and 9th encoder blocks, respectively. And Figures[5(b)](https://arxiv.org/html/2508.19564#A20.F5.sf2 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models"), [5(d)](https://arxiv.org/html/2508.19564#A20.F5.sf4 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") and [5(f)](https://arxiv.org/html/2508.19564#A20.F5.sf6 "In Figure A5 ‣ Appendix T Norm and Ratio Analysis Across Layers ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") present corresponding norm ratios. As observed, the norm of the third higher-order term remains several orders of magnitude smaller than the first two terms, confirming its negligible impact. + +## Appendix U Restricting perturbation in the optimization subspace limits LoRA-SAM + +We evaluate the performance of LoRA-SAM on Llama 3.1-8B for the instruction following benchmark, sweeping the rank over r\in\{4,8,16,32,64\}. Table[A19](https://arxiv.org/html/2508.19564#A21.T19 "Table A19 ‣ Appendix U Restricting perturbation in the optimization subspace limits LoRA-SAM ‣ Bi-LoRA: Efficient Sharpness-Aware Minimization for Fine-Tuning Large-Scale Models") shows that while the average score rises monotonically with rank, the gain per doubling quickly reduces. The upward trend confirms that restricting perturbation in the optimization subspace does limit LoRA-SAM, and that increasing the rank partially mitigates it. Ideally, if r=\max\{m,n\} for each weight matrix W\in\mathbb{R}^{m\times n}, LoRA-SAM becomes full parameter SAM, and the limitation would disappear. Our results show that even at r=64, LoRA-SAM still lags behind our Bi-LoRA with primary rank =8, auxiliary rank =8, whose wider exploratory subspace is achieved with decoupling the (auxiliary) perturbation and (primary) optimization subspace. Moreover, LoRASAM’s SAM-style perturbation doubles computation, whereas Bi-LoRA keeps the cost close to vanilla LoRA. + +Table A19: Instruction-tuning results on fine-tuning Llama 3.1-8B using LoRA-SAM with varying rank. + +Rank MMLU DROP HEval BBH Avg. +4 64.16 48.66 40.24 43.37 49.11 +8 65.50 50.93 40.85 43.14 50.11 +16 66.05 51.92 40.85 42.91 50.43 +32 64.45 50.93 44.51 42.43 50.58 +64 65.81 50.72 43.90 42.37 50.70 + +## Appendix V Theoretical Comparison with Other Efficient SAM Variants Designed for LoRA Fine-tuning + +We provide a brief theoretical comparison showing the difference between Bi-LoRA and two efficient SAM variants, LoRA-o/nBAR and Flat-LoRA. + +Starting from the common SAM objective \min_{W}\max_{\|\epsilon\|\leq\rho}\mathcal{L}(W+\epsilon), all three methods approximate this minimax objective under a low-rank constraint, yet they differ in how the approximation is carried out. LoRAo/nBAR utitlizes SAM’s implicit regularization, it replace the inner maximization with an explicit balanceness penalty, leading to \min_{B_{1},A_{1}}\mathcal{L}\left(W_{0}+BA\right)+\rho\left\|B^{\top}B-AA^{\top}\right\| (take oBAR and \|B\|\geq\|A\| for example) with a looser O\left(\rho^{2}\right) approximation. Flat-LoRA and Bi-LoRA both target LoRA-SAM’s restricted-subspace issue, i.e., the perturbation \epsilon\approx BB^{\top}\left(\nabla_{W}\mathcal{L}\right)+\left(\nabla_{W}\mathcal{L}\right)A^{\top}A is confined to \operatorname{Col}(\mathrm{B}) and \operatorname{Row}(\mathrm{A}). Yet they pursue flatness differently. Flat-LoRA replaces the inner maximisation with a random weight perturbation (RWP) expectation, optimising \min_{B_{1},A_{1}}\mathbb{E}_{\epsilon\sim\mathcal{N}\left(0,\sigma^{2}\right)}\mathcal{L}\left(W_{0}+B_{1}A_{1}+\epsilon\right), thus encouraging expected flatness rather than worst-case sharpness. Bi-LoRA introduces a gradient-driven auxiliary adapter ( B_{2},A_{2} ) so that \min_{B_{1},A_{1}}\max_{\left\|B_{2}A_{2}\right\|\leq\rho}\mathcal{L}\left(W_{0}+B_{1}A_{1}+B_{2}A_{2}\right), allowing a more targeted exploration of the sharp regions through learnable perturbations. From the perspective of gradient-norm regularization, Bi-LoRA essentially replaces SAM’s \ell_{2}-norm constraint with a rank r nuclear norm constraint, thereby inheriting similar theoretical guarantees on generalization and convergence. diff --git a/docs/disam_domain_shift_sharpness.md b/docs/disam_domain_shift_sharpness.md new file mode 100644 index 0000000..3dbe35e --- /dev/null +++ b/docs/disam_domain_shift_sharpness.md @@ -0,0 +1,6918 @@ +Title: Domain-Inspired Sharpness-Aware Minimization Under Domain Shifts + +URL Source: https://arxiv.org/html/2405.18861 + +Markdown Content: +Back to arXiv + +This is experimental HTML to improve accessibility. We invite you to report rendering errors. +Use Alt+Y to toggle on accessible reporting links and Alt+Shift+Y to toggle off. +Learn more about this project and help improve conversions. + +Why HTML? +Report Issue +Back to Abstract +Download PDF + Abstract +1Introduction +2Preliminaries +3Method +4Experiments +5Conclusion +                                               Appendix + References + +HTML conversions sometimes display errors due to content that did not convert correctly from the source. This paper uses the following packages that are not yet supported by the HTML conversion tool. Feedback on these issues are not necessary; they are known and are being worked on. + +failed: titletoc +failed: minitoc + +Authors: achieve the best HTML results from your LaTeX submissions by following these best practices. + +License: CC BY-SA 4.0 +arXiv:2405.18861v1 [cs.CV] 29 May 2024 +\doparttoc\faketableofcontents +Domain-Inspired Sharpness-Aware Minimization Under Domain Shifts +Ruipeng Zhang†,‡, Ziqing Fan†,‡, Jiangchao Yao†,‡,🖂, Ya Zhang†,‡, Yanfeng Wang†,‡,🖂 +† Cooperative Medianet Innovation Center, Shanghai Jiao Tong University +‡ Shanghai Artificial Intelligence Laboratory +{zhangrp, zqfan_knight, Sunarker, ya_zhang, wangyanfeng}@sjtu.edu.cn +Abstract + +This paper presents a Domain-Inspired Sharpness-Aware Minimization (DISAM) algorithm for optimization under domain shifts. It is motivated by the inconsistent convergence degree of SAM across different domains, which induces optimization bias towards certain domains and thus impairs the overall convergence. To address this issue, we consider the domain-level convergence consistency in the sharpness estimation to prevent the overwhelming (deficient) perturbations for less (well) optimized domains. Specifically, DISAM introduces the constraint of minimizing variance in the domain loss, which allows the elastic gradient calibration in perturbation generation: when one domain is optimized above the averaging level w.r.t. loss, the gradient perturbation towards that domain will be weakened automatically, and vice versa. Under this mechanism, we theoretically show that DISAM can achieve faster overall convergence and improved generalization in principle when inconsistent convergence emerges. Extensive experiments on various domain generalization benchmarks show the superiority of DISAM over a range of state-of-the-art methods. Furthermore, we show the superior efficiency of DISAM in parameter-efficient fine-tuning combined with the pretraining models. The source code is released at https://github.com/MediaBrain-SJTU/DISAM. + +1Introduction + +Although deep learning has achieved remarkable advances in various areas (He et al., 2016; Dosovitskiy et al., 2020), it remains a challenge for optimization in pursuit of strong generalization. Especially, a lower training loss does not necessarily guarantee a better generalization, as there exist numerous local minima in the complex and non-convex hypothesis space. Recent empirical and theoretical investigations (Dziugaite & Roy, 2017; Chaudhari et al., 2019; Jiang et al., 2020; 2023; Dinh et al., 2017b; Keskar et al., 2017b) have identified a significant correlation between generalization and the sharpness of the loss landscape. This correlation suggests that generalizability can be interpreted as flatness in the loss surface, leading to a wide range of explorations that have contributed to the rapid development of Sharpness-Aware Minimization (SAM) (Foret et al., 2021). + +Existing SAM-based methods predominantly focus on the narrowly defined generalizability between training and test data under the Independent and Identically Distributed (i.i.d) assumption, which can be summarized as two categories. The first strives to improve the performance by creating a more effective estimation of sharpness like the enhanced minimization in GSAM (Zhuang et al., 2022), PGN (Zhao et al., 2022), SAGM (Wang et al., 2023b) and VaSSO (Li & Giannakis, 2023), as vanilla perturbation in SAM fails to accurately capture the geometric flatness of the loss landscape. The other category targets to improve computational efficiency by reducing perturbation directions (Liu et al., 2022) or using a more efficient perturbation surrogate (Du et al., 2022a; b), as the original SAM incurs double the computational overhead compared to Empirical Risk Minimization (ERM). Nonetheless, these methods cannot solve generalizability scenarios that involve training data of multiple domains with domain shifts like Domain Generalization (DG) (Ben-David et al., 2010; Li et al., 2017). + +In this study, we observed that sometimes SAM even has a detrimental impact in situations where there exist domain shifts across multiple domains as shown in Figure 1(1(a)). While a few studies have incorporated SAM-based methods in domain generalization tasks (Wang et al., 2023b; Foret et al., 2021), they cannot ensure consistent improvements in generalizability during domain shifts due to their reliance on the i.i.d assumption. Upon a thorough analysis of the behavior of SAM under domain shifts, we discovered that the degradation of the training process caused by SAM from the disparity in convergence degree among different domains as shown in Figure 1(1(a)). Given the inconsistency in the degree and direction of convergence among different domains during training (Arjovsky et al., 2019; Krueger et al., 2021), the straightforward application of SAM for perturbations may not only disrupt convergence but also generate perturbation directions that are not adequately coherent to the geometric characteristics of the entire loss landscape. + +(a)Performance under domain shifts. +(b)Convergence curves under domain shifts. +Figure 1:Illustration of SAM’s degradation of the training process under domain shifts. (a) Performance comparison between ERM and SAM, where SAM consistently performs worse than ERM across all hyperparameters +𝜌 +. (b) Convergence curves of SAM and ERM for each domain during training, with the convergence degree normalized to [0,1]. SAM exacerbates the disparity in convergence degree among different domains in domain shift scenarios, resulting in inferior generalization performance. The dataset used here is TerraInc from the DomainBed benchmark, and the backbone is ResNet50. Further experimental details are provided in Section 4.1 and Appendix C.5. + +To solve the aforementioned problem, we propose a Domain-Inspired Sharpness-Aware Minimization (DISAM) algorithm. As the degradation origins from the inconsistency in convergences of these domains, DISAM incorporates domain-level convergence information to intervene in the perturbation of the vanilla SAM: the perturbation direction should focus more on domains with higher convergence degree while being mild to domains with lower convergence degree. Under such balancing in perturbation, the gradient update actually implements the domain-level personalization, thus mitigates the impact of domain shifts and enhances the generalization performance. Technically, we ingeniously accomplish the adaptive adjustment of the perturbation direction in accordance with the degree of convergence through the domain loss variance minimization constraint. The perturbation of DISAM directs towards a location with a more consistent convergence degree, enabling a better global view of the loss landscape for gradient update. We summarize our contributions as follows: + +• + +We identify that the use of SAM has a detrimental impact on training under domain shifts, thereby compromising generalizability, and further analyze that the reason is the inconsistent convergence of training domains that deviates from the underlying i.i.d assumption of SAM. + +• + +We introduce a novel approach called Domain-Inspired Sharpness-Aware Minimization to mitigate the problem above. DISAM incorporates domain-level convergence consistency by imposing a variance minimization constraint on domain loss during the sharpness estimation process, thereby enabling a more representative perturbation location and enhancing generalization. + +• + +Extensive experiments show the superiority of DISAM in improving the current state-of-the-art methods on several benchmarks. We also provide a comprehensive analysis of its merit of faster convergence compared to SAM, and show its persistent generalization capabilities under parameter-efficient fine-tuning with large models like CLIP. + +2Preliminaries +2.1Basic Notations +• + +𝒮 += +{ +𝐷 +1 +, +𝐷 +2 +, +⋯ +, +𝐷 +𝑀 +} +: Overall training set of a +𝑀 +-source domain generalization task. We denote each domain by +𝐷 +𝑖 + and the number of samples in +𝐷 +𝑖 + by +𝑛 +𝑖 += +| +𝐷 +𝑖 +| +. + +• + +𝜉 +, +𝜉 +𝑗 +𝑖 +: A specific sample and the j-th sample in i-th domain +𝐷 +𝑖 +, respectively. + +• + +ℒ +, +ℒ +⁢ +( +𝑤 +) +, +ℒ +⁢ +( +𝑤 +; +𝜉 +) +: A loss function, expected loss under +𝑤 + and the specific loss of +𝜉 +, respectively. + +• + +ℒ +𝑖 +⁢ +( +𝑤 +) += +𝔼 +𝜉 +∈ +𝐷 +𝑖 +⁢ +ℒ +⁢ +( +𝑤 +; +𝜉 +) +: Expected loss under +𝑤 + for each domain +𝐷 +𝑖 +. + +• + +Var +⁢ +{ +⋅ +} +𝑖 += +1 +𝑀 +: The variance among +𝑀 + training domains, which holds: +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +) +} +𝑖 += +1 +𝑀 += +1 +2 +⁢ +𝑀 +2 +⁢ +∑ +𝑖 += +1 +𝑀 +∑ +𝑗 += +1 +𝑀 +( +ℒ +𝑖 +⁢ +( +𝑤 +) +− +ℒ +𝑗 +⁢ +( +𝑤 +) +) +2 +. + +• + +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +) += +ℒ +⁢ +( +𝑤 +) +− +𝜆 +⁢ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +) +} +𝑖 += +1 +𝑀 +: A loss function with domain-inspired regularizer +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +) +} +𝑖 += +1 +𝑀 + on +ℒ +. +𝜆 + is a constant value that controls the strength of the constraint. + +• + +ℒ +𝑝 +⁢ +( +𝑤 +) += +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +ℒ +⁢ +( +𝑤 ++ +𝜖 +) +: The perturbed loss and the objective of SAM. + +• + +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 += +𝑤 +𝑡 ++ +𝜌 +⁢ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +𝑡 +) +‖ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +𝑡 +) +‖ +: The sharpness estimation of DISAM with gradient ascend at step +𝑡 +. + +• + +𝜂 +𝑡 +: Learning rate at step +𝑡 +. + +• + +𝑤 +: Parameters of a neural network +∈ +ℝ +𝑘 +, where +𝑘 + is the dimension. + +• + +𝜖 +∈ +ℝ +𝑘 +: A perturbation on the parameters +𝑤 + with scale +𝜌 +∈ +ℝ +. + +2.2Sharpness-Aware Minimization + +In general, simply minimizing ERM tends to overfit training data and extensive studies show the correlation between generalizability and the sharpness of minima (Dinh et al., 2017b; Hochreiter & Schmidhuber, 1994b; McAllester, 1999; Chaudhari et al., 2019). We clarify the concepts as below. + +Sharpness. + +The sharpness on parameter +𝑤 + with a dataset +𝐷 + and loss function +ℒ + is: + + +𝑠 +⁢ +( +𝑤 +, +𝐷 +) +≜ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +𝔼 +𝜉 +∈ +𝐷 +⁢ +[ +ℒ +⁢ +( +𝑤 ++ +𝜖 +; +𝜉 +) +− +ℒ +⁢ +( +𝑤 +; +𝜉 +) +] +. + +(1) +Sharpness-Aware Minimization (SAM). + +Foret et al. (2021) proposed SAM to improve the generalization by simultaneously minimizing the loss and the sharpness of the overall loss surface. The objective is defined as: + + +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +𝔼 +𝜉 +∈ +𝐷 +⁢ +[ +ℒ +⁢ +( +𝑤 ++ +𝜖 +; +𝜉 +) +] +. + +(2) + +From the above equation, we can see SAM minimizes a perturbed loss “ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +𝔼 +𝜉 +∈ +𝐷 +⁢ +[ +ℒ +⁢ +( +𝑤 ++ +𝜖 +; +𝜉 +) +] +”, which aims to maximize the loss +ℒ + within radius +𝜌 + centered at the parameter +𝑤 +. + +3Method +3.1Motivation + +Although existing SAM-based methods that minimize the sharpness have achieved good generalization, in the case of multiple domains with shifts, the inherent heterogeneity in quantity and task difficulty among domains can considerably distort their sharpness estimation, yielding a degradation in the performance. Concretely, with a collection +𝒮 + of +𝑀 + domains, each of which contains a set of +𝑛 +𝑖 + samples, i.e., +{ +𝜉 +𝑗 +𝑖 += +( +𝑥 +𝑗 +𝑖 +, +𝑦 +𝑗 +𝑖 +) +} +𝑗 += +1 +𝑛 +𝑖 +, the training objective can be then formulated as follows: + + +min +𝑤 +⁡ +𝔼 +𝜉 +∈ +𝒮 +⁢ +[ +ℒ +⁢ +( +𝑤 +; +𝜉 +) +] += +1 +𝑁 +⁢ +∑ +𝑖 += +1 +𝑀 +∑ +𝑗 += +1 +𝑛 +𝑖 +ℒ +⁢ +( +𝑤 +; +𝜉 +𝑗 +𝑖 +) += +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 +) +, + +(3) + +where +𝑁 += +∑ +𝑖 += +1 +𝑀 +𝑛 +𝑖 +, +𝛼 +𝑖 += +𝑛 +𝑖 +𝑁 + and +ℒ +𝑖 +⁢ +( +𝑤 +) += +1 +𝑛 +𝑖 +⁢ +∑ +𝑗 += +1 +𝑛 +𝑖 +ℒ +⁢ +( +𝑤 +; +𝜉 +𝑗 +𝑖 +) +. Note that, we clarify here that we will ignore the notations of data properly in some subsequent equations to avoid clutter. Then, on the basis of Eq. (3), the corresponding objective of SAM under domain shifts is defined as: + + +min +𝑤 +⁡ +𝔼 +𝜉 +∈ +𝒮 +⁢ +[ +ℒ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +⁢ +( +𝑤 +; +𝜉 +) +] += +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +𝔼 +𝜉 +∈ +𝒮 +⁢ +[ +ℒ +⁢ +( +𝑤 ++ +𝜖 +; +𝜉 +) +] +⁢ +≈ +? +⁢ +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +. + +(4) + +The core that we should point out is whether the approximation from +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +𝔼 +𝜉 +∈ +𝒮 +⁢ +[ +ℒ +⁢ +( +𝑤 ++ +𝜖 +; +𝜉 +) +] + to +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) + in Eq. (4) is reasonable. There is no harm when samples in +𝒮 + are intrinsically independently and identically distributed. However, this is actually ill-posed under domain shifts. Differences in the amount of data or inconsistency in task difficulty can result in biased sharpness estimation towards specific domains, hindering the overall convergence. As shown in Figure 2, neglecting the domain shifts and the consequent convergence inconsistency issue, using SAM directly at this point, significantly misdirects the perturbation direction towards the domain with the largest gradient vectors (implying a lower degree of convergence). Consequently, it does not help find a better convergence path and conversely leads to a suboptimal sharp minima. + +Figure 2:Toy example illustrating the problem of SAM under domain shifts. Left: Domain shifts on the loss surface of training domains, which causes the inconsistency of convergence degree. Middle: Differences between SAM and DISAM in the perturbation generation and convergence. Specifically, SAM is affected by the inconsistent degree of convergence. Right: Visualization of loss landscape for ERM, SAM, and DISAM on unseen test domain. DISAM is flatter than SAM and ERM. +3.2Domain-Inspired SAM + +To address the problem described in Eq. (4) and Figure 2, we need to design an adjustment mechanism that takes into account the convergence degree of each domain during the perturbation generation. Specifically, we should make the perturbation direction +∇ +ℒ +𝑝 +⁢ +( +𝑤 +) + to efficiently pull domains that are close to convergence out of sharp minima while minimizing the negative impact on domains that have not yet converged. Here, we define the convergence degree of domain +𝑖 + with model parameter +𝑤 + as +𝐶 +𝑖 +⁢ +( +𝑤 +) += +ℒ +𝑖 +∗ +− +ℒ +𝑖 +⁢ +( +𝑤 +) +, where +ℒ +𝑖 +∗ + represents the optimal minimum of domain +𝑖 + ( +ℒ +𝑖 +∗ +≥ +0 +). The design principle is to prioritize the contribution of domains with larger +𝐶 +𝑖 +⁢ +( +𝑤 +) + to the overall perturbation direction. To achieve this, a simple approach involves directly adding +𝐶 +𝑖 +⁢ +( +𝑤 +) + to the weight +𝛼 +𝑖 + of SAM with a controlling coefficient term +𝛽 +. + + +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) +→ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 ++ +𝛽 +⁢ +𝐶 +𝑖 +⁢ +( +𝑤 +) +∑ +𝑗 += +1 +𝑀 +( +𝛼 +𝑗 ++ +𝛽 +⁢ +𝐶 +𝑗 +⁢ +( +𝑤 +) +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) += +∑ +𝑖 += +1 +𝑀 +𝛽 +⁢ +( +𝐶 +𝑖 +⁢ +( +𝑤 +) +− +𝛼 +𝑖 +⁢ +∑ +𝑗 += +1 +𝑀 +𝐶 +𝑗 +⁢ +( +𝑤 +) +) +1 ++ +𝛽 +⁢ +∑ +𝑗 += +1 +𝑀 +𝐶 +𝑗 +⁢ +( +𝑤 +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) + +(5) + +However, we observe that the weight adjustment in Eq. (5) that is affected by the convergence degree, is constrained by the magnitude of +𝛼 +𝑖 +. That is, domains with higher +𝛼 +𝑖 + values can tolerate lower convergence degrees, which may not accurately satisfy our goals. To refine this, we propose to use an adaptive way to ensure fairness by calculating the average convergence at the domain level, namely, +𝐶 +𝑖 +⁢ +( +𝑤 +) +− +𝛼 +𝑖 +⁢ +∑ +𝑖 += +1 +𝑀 +𝐶 +𝑖 +⁢ +( +𝑤 +) +→ +ℒ +𝑖 +⁢ +( +𝑤 +) +− +1 +𝑀 +⁢ +∑ +𝑖 += +1 +𝑀 +ℒ +𝑖 +⁢ +( +𝑤 +) +. With this intuition, we introduce a method called Domain-Inspired Sharpness-Aware Minimization (DISAM) that incorporates a variance constraint between domain losses to estimate sharpness. It enables the adaptive adjustment of the perturbation direction similar to our spirit, which we will provide a detailed analysis in the following Eq. (8). First of all, we give the definition of the variance between different domain losses as: + + +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +} +𝑖 += +1 +𝑀 += +1 +2 +⁢ +𝑀 +2 +⁢ +∑ +𝑖 += +1 +𝑀 +∑ +𝑗 += +1 +𝑀 +( +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +− +ℒ +𝑗 +⁢ +( +𝑤 ++ +𝜖 +) +) +2 +. + +(6) + +Then, putting the above variance term into the loss, the new training objective can be defined as: + + +min +𝑤 +⁡ +𝔼 +𝜉 +∈ +𝒮 +⁢ +[ +ℒ +𝐷 +⁢ +𝐼 +⁢ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +⁢ +( +𝑤 +; +𝜉 +) +] +≜ +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +[ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +− +𝜆 +⁢ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +^ ++ +𝜖 +) +} +𝑖 += +1 +𝑀 +] + +(7) + +Here +𝑤 +^ + is +𝑤 + without derivative taken during backpropagation, and it only makes effect in the +max +‖ +𝜖 +‖ +2 +≤ +𝜌 + loop without affecting the optimization of the first term w.r.t. +𝑤 +. Following the computing way of perturbation +𝜖 + in SAM, namely, using first-order Taylor expansion (Foret et al., 2021), we will have +𝜖 +≈ +𝜌 +⁢ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +‖ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +‖ +, where +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +𝑆 +⁢ +𝐴 +⁢ +𝑀 + w.r.t. +𝑤 + has the form: + + +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +𝑆 +⁢ +𝐴 +⁢ +𝑀 + += +∑ +𝑖 += +1 +𝑀 +( +𝛼 +𝑖 +− +2 +⁢ +𝜆 +𝑀 +⁢ +( +ℒ +𝑖 +⁢ +( +𝑤 +) +− +1 +𝑀 +⁢ +∑ +𝑗 += +1 +𝑀 +ℒ +𝑗 +⁢ +( +𝑤 +) +) +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) + +(8) + + += +∇ +ℒ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +− +∑ +𝑖 += +1 +𝑀 +2 +⁢ +𝜆 +𝑀 +⁢ +( +ℒ +𝑖 +⁢ +( +𝑤 +) +− +1 +𝑀 +⁢ +∑ +𝑗 += +1 +𝑀 +ℒ +𝑗 +⁢ +( +𝑤 +) +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) +⏟ +Adaptive adjustment: +  +increase + weights for smaller losses,  +reduce + for larger ones. + + +The first term in the RHS of Eq. (8) recovers the gradient term for perturbation generation in SAM, and the second term characterizes the working mechanism for the adaptive adjustment. As can be seen, when the loss +ℒ +𝑖 +⁢ +( +𝑤 +) + of one certain domain is above the averaging level, the second term will generate a residual gradient for this domain to cancel out the gradient contribution in +∇ +ℒ +𝑆 +⁢ +𝐴 +⁢ +𝑀 +, and vice versa. It means to have a mild perturbation for the domain that is not well optimized, and have an aggressive perturbation for the domain that is well optimized. In total, the variance constraint ensures that the perturbation location is at a more consistent point, enabling a better global view of the loss landscape for gradient update. The complete algorithm is described in Appendix B. Regarding +𝜆 +, a default value of +0.1 + is relatively stable, and we provide more discussion about +𝜆 + in Appendix B.3. + +Difference and Compatibility. Similarly, the current SAM variants will meet the same challenge, if they are directly applied to this scenario. Different from existing state-of-the-art methods like GSAM (Zhuang et al., 2022) and SAGM (Wang et al., 2023b) that modify the optimization objective based on the second derivative of SAM, DISAM rectifies the domain shift issue by the domain-level adjustment in the perturbation generation, which actually alleviates the negative impacts on the training objective. In Appendix A.2.5, we present a table to comprehensively characterize the difference between DISAM and other domain-invariant robust optimization methods. Besides, DISAM can be easily extended into other SAM-based methods to improve the generalization performance. We have provided a comparison of the similarities and differences between DISAM and general convergence consistency methods (such as V-REx(Krueger et al., 2021) and Fishr(Rame et al., 2022)) in Appendix B.1. + +3.3Understanding Domain-Inspired SAM +Complexity. + +Compared to SAM-based methods, our algorithm only additionally computes the loss variance between different domains as Eq. (6) and requires no extra storing space. Therefore it has the same space complexity and the time complexity can be represented as +𝑂 +DISAM += +𝑂 +SAM ++ +𝑂 +Var +. Since it only needs to additionally count the domain loss and the corresponding variance according to the domain label when calculating the empirical loss, the overall cost on +𝑂 +Var + is negligible. + +Convergence. + +In the following, we provide the convergence analysis of SAM and DISAM. Similar to (Zhuang et al., 2022; Jiang et al., 2023), our theorem is established on assumptions that a non-convex function +ℒ +⁢ +( +𝑤 +) + is +𝐿 + Lipschitz-smooth, the lower bound of the empirical loss is bounded by +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 +, and the norm of noisy stochastic gradients is bounded ( +‖ +∇ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 +) +‖ +2 +≤ +𝐺 +) at the t-step. + +Theorem 1. + +Consider a non-convex function +ℒ +⁢ +( +𝑤 +) + with Lipschitz-smooth constant +𝐿 + and lower bound +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 +. With the bounded norm assumption of noisy stochastic gradients ( +‖ +∇ +ℒ +𝑝 +⁢ +( +𝑤 +) +‖ +2 +≤ +𝐺 +) at the t-step, the learning rate +𝜂 +𝑡 += +𝜂 +0 +/ +𝑡 + and a fixed perturbation amplitude +𝜌 +, we have: + + +1 +𝑇 +⁢ +∑ +𝑡 += +1 +𝑇 +𝔼 +⁢ +‖ +∇ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 +) +‖ +2 +2 +≤ +ℒ +𝑝 +⁢ +( +𝑤 +0 +) +− +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 +𝜂 +0 +⁢ +1 +𝑇 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +⁢ +log +⁡ +( +𝑇 +) +𝑇 +, + + +where in SAM, +Γ += +𝐿 + and in our DISAM, +Γ +≤ +𝐿 +. + +The complete proof is presented in Appendix B. As can be seen in Theorem 1, the critical convergence difference between SAM and DISAM is on +Γ +, and especially the +Γ + in our DISAM is smaller than that in SAM due to the canonically correlated perturbations during training (see the proof for the details), which thus leads to a faster convergence rate. Note that, the overall +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 + in Theorem 1 indicates that a larger perturbation amplitude +𝜌 + will result in larger upper bound of convergence. However, as analyzed in SAM (Foret et al., 2021), a larger perturbation amplitude +𝜌 + has the merit of reaching a smaller upper bound on generalization error. This means that +𝜌 + actually has a trade-off between accelerating the convergence and improving the generalization. Fortunately, when in the same overall value of +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +, as DISAM enjoys a smaller +Γ + than SAM, DISAM can permit potential larger +𝜌 + than that in SAM, thus yielding a better generalization (Please refer to Appendix B.4 for more discussion). + +(a) +(b) +(c) +(d) +Figure 3:Convergence curves and Max +𝜌 + search for SAM and DISAM. (a) & (b) show the trend of +ℒ +⁢ +( +𝑤 +) + during the training process on PACS dataset, while (c) & (d) search for the maximum perturbation amplitude +𝜌 + of SAM and DISAM on PACS and OfficeHome datasets. + +To verify the theoretical analysis, we present our empirical results of DISAM and SAM in Figure 3. As shown in Figures 3(3(a)) and 3(3(b)), the training curves on the PACS dataset show that DISAM achieves faster and steeper convergence than SAM. In addition, as DISAM has a smaller +Γ +, it is able to utilize a larger perturbation amplitude +𝜌 +. In Figures 3(3(c)) and 3(3(d)), we show the experimental support that DISAM allows larger +𝜌 + values (0.2 and 0.13) than SAM (0.09 and 0.06) on PACS and OfficeHome datasets, while achieving a better performance. In total, these experiments confirm the advantage of DISAM that allows larger +𝜌 + for better generalization. + +4Experiments +4.1Experiment Setups +Datasets. + +We evaluate DISAM on five datasets PACS (Li et al., 2017), VLCS (Fang et al., 2013) OfficeHome (Venkateswara et al., 2017), TerraIncognita (Beery et al., 2018) (abbreviated as TerraInc), and DomainNet (Peng et al., 2019), following the DomainBed benchmark (Gulrajani & Lopez-Paz, 2021). For fair comparison, we adhere to the training and evaluation protocol outlined in DomainBed. + +Evaluation. + +The standard leave-one-domain-out strategy is used in evaluation. Specially, the unseen domain is used to evaluate the out-of-domain generalization, and the validation sets of source domains are used to measure the in-domain generalization, while the others are used for training. Final accuracy is averaged across all settings, and the performance is the averaging over three trials with distinct random seeds. Detailed statistics for each case of all datasets are provided in Appendix C. + +Implementation details. + +Our backbones are ResNet50 pretrained on ImageNet (He et al., 2016) and a pretrained CLIP (Radford et al., 2021) with ViT-B/16 structure (Dosovitskiy et al., 2020). For model hyperparameters, we adopt settings in (Wang et al., 2023b) for experiments using ResNet50 and in (Shu et al., 2023) for experiments using CLIP. As the default, we set the perturbation hyperparameter +𝜌 + to 0.05 (Wang et al., 2023b) (Fixed value during training), and the weight of the variance constraint +𝜆 + to 0.1. For a detailed description of the hyparameter settings, please see Appendix C. + +Table 1:Comparison with state-of-the-art domain generalization methods based on ResNet50. In-domain and Out-of-domain accuracies on five datasets from DomainBed. +Algorithm PACS VLCS OfficeHome TerraInc DomainNet Avg. +In-domain results +ERM +96.6 +± +0.2 + +84.6 +± +0.4 + +84.2 +± +0.3 + +93.6 +± +0.3 + +67.1 +± +1.6 + +85.2 + +SAM +97.3 +± +0.1 + +84.8 +± +0.3 + +85.8 +± +0.2 + +88.9 +± +0.2 + +68.5 +± +0.1 + +85.1 + +Domain-Inspired +97.8 +± +0.1 + +84.4 +± +0.3 + +86.3 +± +0.2 + +94.8 +± +0.2 + +70.2 +± +0.1 + +86.7 + +GSAM +97.8 +± +0.2 + +83.9 +± +0.2 + +85.9 +± +0.2 + +92.1 +± +0.2 + +69.1 +± +0.1 + +85.8 + +Domain-Inspired +97.9 +± +0.1 + +85.1 +± +0.4 + +86.2 +± +0.2 + +94.8 +± +0.3 + +70.0 +± +0.1 + +86.8 + +SAGM +97.6 +± +0.1 + +84.6 +± +0.3 + +86.1 +± +0.2 + +92.0 +± +0.2 + +69.2 +± +0.1 + +85.9 + +Domain-Inspired +97.9 +± +0.1 + +85.0 +± +0.2 + +86.5 +± +0.3 + +94.9 +± +0.2 + +70.5 +± +0.1 + 87.0 +Out-of-domain results +ERM +85.5 +± +0.2 + +77.3 +± +0.4 + +66.5 +± +0.3 + +46.1 +± +1.8 + +43.8 +± +0.1 + +63.9 + +CORAL (SOTA) +86.2 +± +0.3 + +78.8 +± +0.3 + +68.7 +± +0.3 + +47.6 +± +1.0 + +41.5 +± +0.1 + +64.5 + +SAM +85.8 +± +0.2 + +79.4 +± +0.1 + +69.6 +± +0.1 + +43.3 +± +0.7 + +44.3 +± +0.0 + +64.5 + +Domain-Inspired +87.3 +± +0.2 + +80.1 +± +0.5 + +70.7 +± +0.2 + +47.9 +± +0.8 + +45.8 +± +0.2 + +66.4 + +GSAM +85.9 +± +0.1 + +79.1 +± +0.2 + +69.3 +± +0.0 + +47.0 +± +0.8 + +44.6 +± +0.2 + +65.1 + +Domain-Inspired +87.2 +± +0.3 + +80.0 +± +0.3 + +70.8 +± +0.3 + +50.6 +± +1.2 + +45.6 +± +0.1 + +66.8 + +SAGM +86.6 +± +0.2 + +80.0 +± +0.3 + +70.1 +± +0.2 + +48.8 +± +0.9 + +45.0 +± +0.2 + +66.1 + +Domain-Inspired +87.5 +± +0.3 + +80.7 +± +0.2 + +71.0 +± +0.2 + +50.0 +± +1.2 + +46.0 +± +0.1 + 67.0 +   + CORAL +88.4 +± +0.3 + +81.2 +± +0.4 + +71.7 +± +0.2 + +51.7 +± +0.3 + +46.3 +± +0.2 + +67.9 +Table 2:Comparison with state-of-the-art domain generalization methods based on CLIP with ViT-B/16. Out-of-domain accuracies on five datasets from DomainBed. +Algorithm PACS VLCS OfficeHome TerraInc DomainNet Avg. +Zero-shot +96.2 + +81.7 + +82.0 + +33.4 + +57.5 + +70.2 + +CoOp +96.8 + +81.2 + +84.2 + +44.9 + +59.9 + +73.4 + ++ SAM +97.1 +± +0.1 + +81.3 +± +0.8 + +84.6 +± +0.2 + +47.7 +± +1.3 + +60.3 +± +0.2 + +74.2 + ++ DISAM +97.2 +± +0.1 + +81.8 +± +0.4 + +84.8 +± +0.2 + +49.5 +± +1.2 + +60.6 +± +0.2 + +74.8 + +ERM +96.1 +± +0.5 + +83.0 +± +0.2 + +83.3 +± +0.3 + +60.9 +± +0.2 + +59.9 +± +0.1 + +76.7 + +CLIPOOD1 +97.3 +± +0.1 + +85.0 +± +0.4 + +87.0 +± +0.2 + +60.4 +± +0.7 + +63.5 +± +0.1 + +78.6 + + +CLIPOOD +∗ +2 +96.6 +± +0.4 + +84.1 +± +0.3 + +86.1 +± +0.2 + +59.7 +± +0.8 + +63.1 +± +0.1 + +77.9 + ++ SAM +96.9 +± +0.2 + +84.3 +± +0.6 + +84.4 +± +0.4 + +60.0 +± +1.4 + +58.6 +± +0.2 + +76.9 + ++ DISAM +97.1 +± +0.1 + +85.6 +± +0.2 + +86.6 +± +0.0 + +61.1 +± +0.7 + +63.6 +± +0.1 + 78.8 +4.2Performance under ResNet50 backbone + +We propose incorporating our domain-inspired adaptive adjustment into three SAM-based methods: SAM (Foret et al., 2021), GSAM (Zhuang et al., 2022), and SAGM (Wang et al., 2023b) on five datasets of DomainBed with ResNet50 backbone. Table 1 shows that our Domain-Inspired SAM can mitigate issues arising from SAM’s training under domain shifts, by comparing averaged in-domain and out-of-domain performance of leading SAM methods, with and without DISAM. In-domain results show domain-inspired perturbations enhance convergence, especially on the TerraInc dataset with substantial domain gaps. In Out-of-domain results, DISAM consistently improves generalization, with average improvements of 1.9% for SAM, 1.7% for GSAM, and 1.9% for SAGM. Notably, SAM performs well when the performance gap between in-domain and out-of-domain is small but worse than ERM on datasets like TerraInc with large gaps, which proves our analysis of SAM’s shortcomings under domain shifts. This shows SAM’s inconsistent convergence for large domain shifts, which DISAM addresses by incorporating domain-inspired adaptive adjustments based on domain-level convergence degree. Incorporating CORAL constraints, a recognized effective traditional DG method on DomainBed improves SAGM with DISAM and sets new state-of-the-art results on all settings. + +4.3Performance under CLIP-based pretrained large model + +The CLIP-based large pretrained models (Radford et al., 2021) show great zero-shot performance but struggle with domain shifts in downstream tasks. We assess DISAM’s out-of-domain results on CLIP using the experimental setup of CLIPOOD (Shu et al., 2023). We test two downstream adaptation methods: CoOp (Zhou et al., 2022a), an effective prompt learning approach, and CLIPOOD, an image encoder finetuning approach for DG problems. For CoOp, we use a 16-length learnable generic prompt and 5000 training steps, and For CLIPOOD settings, we follow Shu et al. (2023). As shown in Table 2, DISAM effectively mitigates the impact of domain shifts on model generalization during downstream task adaptation. In addition, as CoOp and CLIPOOD∗ primarily focus on rapid adaptation with limited parameters, the overfitting risk can be alleviated through early stopping, resulting in the relatively marginal improvements for DISAM in Table 2. Despite this, when handling challenging tasks like TerraInc and DomainNet, our approach still exhibits substantial enhancements. + +Table 3:Accuracy on OfficeHome and DomainNet with both domain shifts and open classes. +        Split Algorithm         OfficeHome         DomainNet Avg. +A C P R C I P Q R S +        Base         Zero-shot +86.7 + +75.9 + +89.6 + +92.2 + +72.6 + +51.8 + +65.4 + +13.6 + +83.5 + +67.2 + +72.6 + +         +CoOp +∗ + +87.3 + +76.7 + +92.2 + +92.5 + +74.6 + +58.2 + +67.9 + +15.0 + +83.7 + +69.9 + +74.4 + +        +SAM +89.2 + +79.6 + 93.0 +93.7 + +73.5 + +58.4 + +67.8 + +14.8 + +83.6 + +69.5 + +75.1 + +        +DISAM +88.0 + 80.5 +92.7 + +92.4 + +75.0 + +59.9 + +68.7 + +14.9 + +84.4 + +70.5 + +75.3 + +         +CLIPOOD +∗ + +88.9 + +79.5 + +92.2 + +94.0 + +76.3 + +58.7 + +69.9 + +17.5 + +85.6 + +72.4 + +76.0 + +        +SAM +89.1 + +78.9 + +92.3 + +94.1 + 78.7 62.1 72.0 +19.9 + 86.5 73.5 77.0 +        +DISAM 89.7 +79.4 + +92.7 + 94.2 +77.1 + +61.8 + +71.5 + 20.0 +86.0 + +73.1 + 77.0 +        New         Zero-shot +76.8 + +59.7 + +88.7 + +86.4 + +69.7 + +45.0 + +67.0 + +14.3 + +83.9 + +60.8 + +67.4 + +         +CoOp +∗ + +73.7 + +56.4 + +86.6 + +85.0 + +69.7 + +47.4 + +67.0 + +15.2 + +82.5 + +61.5 + +66.3 + +        +SAM +75.2 + +59.1 + +89.6 + +86.0 + +71.1 + 49.2 +69.3 + +15.4 + +82.2 + +62.9 + +67.9 + +        +DISAM +79.3 + +61.5 + 90.9 +88.4 + 72.1 +49.0 + 69.6 +15.5 + 85.5 +62.9 + +69.6 + +         +CLIPOOD +∗ + +75.2 + +58.6 + +87.5 + +85.8 + +69.3 + +46.4 + +67.2 + +15.2 + +83.2 + +60.6 + +66.9 + +        +SAM +77.2 + +60.0 + +89.8 + +87.6 + +66.8 + +45.4 + +64.9 + +14.8 + +82.0 + +57.1 + +66.9 + +        +DISAM 79.7 62.0 +90.5 + 89.0 +71.8 + +48.7 + +68.7 + 17.5 +84.7 + 63.0 69.7 +        Total         Zero-shot +82.6 + +67.3 + +88.8 + +89.5 + +71.4 + +47.1 + +66.2 + +13.8 + +83.4 + +63.4 + +69.8 + +         +CoOp +∗ + +81.4 + +65.7 + +88.9 + +88.8 + +71.9 + +51.3 + +67.4 + +15.1 + +83.1 + +65.1 + +70.1 + +        +SAM +83.5 + +69.1 + +91.3 + +90.1 + +72.3 + +52.8 + +68.6 + +15.1 + +82.9 + +65.7 + +71.5 + +        +DISAM +84.2 + 70.1 91.7 +90.4 + +73.4 + 53.2 +69.1 + +15.2 + 85.0 +65.8 + +72.2 + +         +CLIPOOD +∗ + +83.3 + +68.8 + +89.9 + +90.1 + +72.5 + +51.2 + +68.5 + +16.4 + +84.4 + +65.6 + +71.4 + +        +SAM +84.2 + +69.2 + +91.0 + +91.0 + +72.3 + +52.0 + +68.4 + +17.3 + +84.3 + +64.0 + +71.8 + +        + DISAM 84.6 +69.5 + +91.3 + 91.2 73.5 53.2 69.4 18.3 +84.9 + 66.3 72.6 +4.4Performance under open-class generalization +Figure 4:Comparison of CoOp based ERM, SAM and DISAM on accuracy curves for base/new classes. (Top: In-Domain; Bottom: Out-of-Domain) + +In this part, we evaluate the performance of our DISAM in a more realistic in-the-wild setting, where both domain shifts and open-class scenarios may arise in the test domain. This setting was first proposed by Shu et al. (2023). OfficeHome and DomainNet are selected to conduct related experiments because they offer an ample number of classes suitable for evaluating open-class situations. To delineate, we segregate the classes within each dataset into two categories, based on the class ID. The initial half denotes the base classes, and the latter half signifies the new classes. Based on Section 4.1, we eliminate the data corresponding to new classes in the training domains. Due to CLIP’s open vocabulary property, we can evaluate the new classes on the unseen test domain. + +As presented in Table 3, we evaluated the classification accuracy of "Base" classes, "New" classes, and "Total" classes in the test domain, where total classes represent the overall test domain. It revealed that the existing adaptation approach while performing well on base classes, lacks generalization on new classes during the fitting process. Our DISAM mitigates open-class overfitting using domain-level convergence constraints, improving performance by 3.3% over CoOp and 3.1% over CLIPOOD. Figure 4 provides a detailed analysis of open classes and domain shifts dimensions. ERM tends to overfit to both in-domain and base class. While SAM outperforms ERM, it struggles with sharp minima perturbations, failing to effectively escape from them. This difficulty hampers its generalization capabilities in larger models. Please refer to Appendix C.6 for more discussion about DISAM and other methods for open-class generalization. + +4.5Ablation Studies +(a) +𝜌 + on PACS. +(b) +𝜌 + on OfficeHome. +(c) +𝜆 + on PACS. +(d) +𝜆 + on OfficeHome. +Figure 5:Ablation study investigating the sensitivity of hyperparameters, namely perturbation amplitude +𝜌 + and variance constraint weight +𝜆 + in DISAM. +(a) +(b) +(c) +Figure 6: (a) & (b): Sharpness curves for SAM and DISAM trained on PACS dataset, which show the trend of the estimated sharpness of the model on the test domain. (c): Computation cost with and without Domain-Inspired SAM used on ResNet50 and ViT-B/16 backbone. +Hyperparameter sensitivity. + +We performed a sensitivity analysis of the perturbation amplitude +𝜌 +, and the variance constraint weight +𝜆 +, on the PACS and OfficeHome datasets. The default +𝜌 + and +𝜆 + is set to 0.05 (Zhuang et al., 2022) and 0.1, respectively. As illustrated in Figure 5(5(a)) and 5(5(b)) within a wide range of +𝜌 +, DISAM consistently achieves stable and superior results compared to SAM. However, when +𝜌 + is too large or small, experimental results worsen. Large +𝜌 + hiders convergence, while small +𝜌 + weakens sharpness constraint, both affecting generalization. As for +𝜆 +, Figure 5(5(c)) and 5(5(d)) show stable results when +𝜆 +∈ +[ +0.1 +, +0.7 +] +. However, larger +𝜆 + values increase the variance due to excessive over-conditioning weight, which can also influence the convergence. + +Estimated sharpness on unseen test domain. + +Estimating sharpness has a high computational cost. Early methods (Dinh et al., 2017b; Hochreiter & Schmidhuber, 1994b) relied on Monte Carlo sampling, but recent advancements (Jiang et al., 2023; 2020) use gradient-based approximations for efficiency. We assess model sharpness on unseen test domains at each epoch’s end, based on the work of Jiang et al. (2023). As depicted in Figure 6(6(a)) and 6(6(b)), our DISAM achieves much smaller gradient variance +Var +⁢ +{ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +; +𝐵 +𝑡 +) +} + than SAM during the whole training, indicating the incorporation of domain-inspired information can further reduce the sharpness of the loss surface. + +Computation cost of DISAM. + +In Figure 6(6(c)), we show the extra computational cost from adding domain-inspired perturbation direction generation versus the original algorithm (time cost/step, batch size 32, RTX 3090 GPU). Empirical findings show DISAM integration incurs minimal overhead ( 0.01s) across algorithms/backbones, linked solely to domain number and batch size, not model size, via strategic domain loss variance constraints for domain-level convergence consistency. + +5Conclusion + +This paper presents Domain-Inspired Sharpness-Aware Minimization (DISAM), an algorithm that incorporates domain-level convergence consistency into the generation of SAM’s perturbations, to address the dilemma under multiple domains. DISAM mitigates SAM’s bias in domain shifts that can detrimentally impact the convergence during training, yielding perturbations towards highly converged domains and limiting those in less optimized ones. This is achieved by minimizing the variance of domain loss during perturbation generation, enabling an adaptive weight adjustment for each domain based on its convergence degree, thereby enhancing the convergence across training domains and generalization on unseen domains. Extensive experiments on the domain generalization benchmarks prove DISAM’s superiority over existing methods. In addition, DISAM persistents generalization capabilities under parameter-efficient fine-tuning with large models like CLIP. + +Ethics statement + +This paper does not raise any ethics concerns. This study does not involve any human subjects, practices to data set releases, potentially harmful insights, methodologies and applications, potential conflicts of interest and sponsorship, discrimination/bias/fairness concerns, privacy and security issues, legal compliance, and research integrity issues. + +Reproducibility Statement + +All experiments were conducted using NVIDIA GeForce RTX 3090 GPU, Python 3.9.15, Pytorch 1.12.1, and clip 1.0. Further details regarding experimental setups and implementation can be found in Section 4.1 and Appendix C, while theoretical proofs are provided in Appendix B. The principal code for implementing Domain-Inspired SAM is presented in Appendix D. + +Acknowledgments + +This work is supported by the National Key R&D Program of China (No. 2022ZD0160702), STCSM (No. 22511106101, No. 18DZ2270700, No. 21DZ1100-100), 111 plan (No. BP0719010), and State Key Laboratory of UHD Video and Audio Production and Presentation. Ruipeng Zhang and Ziqing Fan are partially supported by Wu Wen Jun Honorary Doctoral Scholarship, AI Institute, Shanghai Jiao Tong University. + +References +Arjovsky et al. (2019) +↑ + Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz.Invariant risk minimization.arXiv preprint arXiv:1907.02893, 2019. +Balaji et al. (2018) +↑ + Yogesh Balaji et al.Metareg: Towards domain generalization using meta-regularization.In NeurIPS, pp.  998–1008, 2018. +Beery et al. (2018) +↑ + Sara Beery, Grant Van Horn, and Pietro Perona.Recognition in terra incognita.In Proceedings of the European conference on computer vision (ECCV), pp.  456–473, 2018. +Ben-David et al. (2010) +↑ + Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan.A theory of learning from different domains.Machine learning, 79:151–175, 2010. +Carlucci et al. (2019) +↑ + Fabio Maria Carlucci, Paolo Russo, Tatiana Tommasi, and Barbara Caputo.Hallucinating agnostic images to generalize across domains.In 2019 IEEE/CVF International Conference on Computer Vision Workshop (ICCVW), pp.  3227–3234. IEEE, 2019. +Cha et al. (2021) +↑ + Junbum Cha, Sanghyuk Chun, Kyungjae Lee, Han-Cheol Cho, Seunghyun Park, Yunsung Lee, and Sungrae Park.Swad: Domain generalization by seeking flat minima.Advances in Neural Information Processing Systems, 34:22405–22418, 2021. +Chang et al. (2019) +↑ + Woong-Gi Chang, Tackgeun You, Seonguk Seo, Suha Kwak, and Bohyung Han.Domain-specific batch normalization for unsupervised domain adaptation.In Proceedings of the IEEE/CVF conference on Computer Vision and Pattern Recognition, pp.  7354–7362, 2019. +Chaudhari et al. (2019) +↑ + Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina.Entropy-sgd: Biasing gradient descent into wide valleys.Journal of Statistical Mechanics: Theory and Experiment, 2019(12):124018, 2019. +Dinh et al. (2017a) +↑ + Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.Sharp minima can generalize for deep nets.In International Conference on Machine Learning, pp.  1019–1028. PMLR, 2017a. +Dinh et al. (2017b) +↑ + Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio.Sharp minima can generalize for deep nets.In International Conference on Machine Learning, pp.  1019–1028. PMLR, 2017b. +Dosovitskiy et al. (2020) +↑ + Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.An image is worth 16x16 words: Transformers for image recognition at scale.arXiv preprint arXiv:2010.11929, 2020. +Dou et al. (2019) +↑ + Qi Dou, Daniel Coelho de Castro, Konstantinos Kamnitsas, and Ben Glocker.Domain generalization via model-agnostic learning of semantic features.In NeurIPS, pp.  6450–6461, 2019. +Du et al. (2022a) +↑ + Jiawei Du, Hanshu Yan, Jiashi Feng, Joey Tianyi Zhou, Liangli Zhen, Rick Siow Mong Goh, and Vincent Tan.Efficient sharpness-aware minimization for improved training of neural networks.In International Conference on Learning Representations, 2022a. +Du et al. (2022b) +↑ + Jiawei Du, Daquan Zhou, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou.Sharpness-aware training for free.Advances in Neural Information Processing Systems, 35:23439–23451, 2022b. +Dziugaite & Roy (2017) +↑ + Gintare Karolina Dziugaite and Daniel M. Roy.Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.In Proceedings of the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAI), 2017. +Fang et al. (2013) +↑ + Chen Fang, Ye Xu, and Daniel N Rockmore.Unbiased metric learning: On the utilization of multiple datasets and web images for softening bias.In Proceedings of the IEEE International Conference on Computer Vision, pp.  1657–1664, 2013. +Finn et al. (2017) +↑ + Chelsea Finn, Pieter Abbeel, and Sergey Levine.Model-agnostic meta-learning for fast adaptation of deep networks.In International conference on machine learning, pp.  1126–1135. PMLR, 2017. +Foret et al. (2021) +↑ + Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.Sharpness-aware minimization for efficiently improving generalization.In International Conference on Learning Representations, 2021. +Gulrajani & Lopez-Paz (2021) +↑ + Ishaan Gulrajani and David Lopez-Paz.In search of lost domain generalization.In International Conference on Learning Representations, 2021.URL https://openreview.net/forum?id=lQdXeXDoWtI. +He et al. (2016) +↑ + Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.Deep residual learning for image recognition.In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.  770–778, 2016. +Hochreiter & Schmidhuber (1994a) +↑ + Sepp Hochreiter and Jürgen Schmidhuber.Simplifying neural nets by discovering flat minima.Advances in neural information processing systems, 7, 1994a. +Hochreiter & Schmidhuber (1994b) +↑ + Sepp Hochreiter and Jürgen Schmidhuber.Simplifying neural nets by discovering flat minima.Advances in neural information processing systems, 7, 1994b. +Izmailov et al. (2018) +↑ + Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson.Averaging weights leads to wider optima and better generalization.arXiv preprint arXiv:1803.05407, 2018. +Jiang et al. (2023) +↑ + Weisen Jiang, Hansi Yang, Yu Zhang, and James Kwok.An adaptive policy to employ sharpness-aware minimization.In The Eleventh International Conference on Learning Representations, 2023.URL https://openreview.net/forum?id=6Wl7-M2BC-. +Jiang et al. (2020) +↑ + Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio.Fantastic generalization measures and where to find them.In International Conference on Learning Representations, 2020.URL https://openreview.net/forum?id=SJgIPJBFvH. +Jin et al. (2022) +↑ + Xin Jin, Cuiling Lan, Wenjun Zeng, and Zhibo Chen.Style normalization and restitution for domain generalization and adaptation.IEEE Transactions on Multimedia, 24:3636–3651, 2022.doi: 10.1109/TMM.2021.3104379. +Keskar et al. (2017a) +↑ + Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.On large-batch training for deep learning: Generalization gap and sharp minima.In International Conference on Learning Representations, 2017a. +Keskar et al. (2017b) +↑ + Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang.On large-batch training for deep learning: Generalization gap and sharp minima.In International Conference on Learning Representations, 2017b.URL https://openreview.net/forum?id=H1oyRlYgg. +Kim et al. (2021) +↑ + Daehee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jaekoo Lee.Selfreg: Self-supervised contrastive regularization for domain generalization.In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp.  9619–9628, 2021. +Kim et al. (2022) +↑ + Minyoung Kim, Da Li, Shell X Hu, and Timothy Hospedales.Fisher sam: Information geometry and sharpness aware minimisation.In International Conference on Machine Learning, pp.  11148–11161. PMLR, 2022. +Krueger et al. (2021) +↑ + David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville.Out-of-distribution generalization via risk extrapolation (rex).In International Conference on Machine Learning, pp.  5815–5826. PMLR, 2021. +Kwon et al. (2021) +↑ + Jungmin Kwon, Jeongseop Kim, Hyunseo Park, and In Kwon Choi.Asam: Adaptive sharpness-aware minimization for scale-invariant learning of deep neural networks.In International Conference on Machine Learning, pp.  5905–5914. PMLR, 2021. +Li & Giannakis (2023) +↑ + Bingcong Li and Georgios B Giannakis.Enhancing sharpness-aware optimization through variance suppression.arXiv preprint arXiv:2309.15639, 2023. +Li et al. (2017) +↑ + Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales.Deeper, broader and artier domain generalization.In Proceedings of the IEEE international conference on computer vision, pp.  5542–5550, 2017. +Li et al. (2018a) +↑ + Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M. Hospedales.Learning to generalize: Meta-learning for domain generalization.In AAAI, 2018a. +Li et al. (2018b) +↑ + Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein.Visualizing the loss landscape of neural nets.Advances in neural information processing systems, 31, 2018b. +Li et al. (2018c) +↑ + Haoliang Li, Sinno Jialin Pan, Shiqi Wang, and Alex C Kot.Domain generalization with adversarial feature learning.In CVPR, pp.  5400–5409, 2018c. +Li et al. (2018d) +↑ + Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao.Deep domain generalization via conditional invariant adversarial networks.In Proceedings of the European conference on computer vision (ECCV), pp.  624–639, 2018d. +Li et al. (2019) +↑ + Yiying Li, Yongxin Yang, Wei Zhou, and Timothy Hospedales.Feature-critic networks for heterogeneous domain generalization.In International Conference on Machine Learning, pp.  3915–3924. PMLR, 2019. +Liu et al. (2023) +↑ + Yajing Liu, Zhiwei Xiong, Ya Li, Xinmei Tian, and Zheng-Jun Zha.Domain generalization via encoding and resampling in a unified latent space.IEEE Transactions on Multimedia, 25:126–139, 2023.doi: 10.1109/TMM.2021.3121564. +Liu et al. (2022) +↑ + Yong Liu, Siqi Mai, Xiangning Chen, Cho-Jui Hsieh, and Yang You.Towards efficient and scalable sharpness-aware minimization.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  12360–12370, 2022. +McAllester (1999) +↑ + David A McAllester.Pac-bayesian model averaging.In Proceedings of the twelfth annual conference on Computational learning theory, pp.  164–170, 1999. +Mi et al. (2022) +↑ + Peng Mi, Li Shen, Tianhe Ren, Yiyi Zhou, Xiaoshuai Sun, Rongrong Ji, and Dacheng Tao.Make sharpness-aware minimization stronger: A sparsified perturbation approach.In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. +Motiian et al. (2017) +↑ + Saeid Motiian, Marco Piccirilli, Donald A Adjeroh, and Gianfranco Doretto.Unified deep supervised domain adaptation and generalization.In Proceedings of the IEEE international conference on computer vision, pp.  5715–5725, 2017. +Niu et al. (2023) +↑ + Ziwei Niu, Junkun Yuan, Xu Ma, Yingying Xu, Jing Liu, Yen-Wei Chen, Ruofeng Tong, and Lanfen Lin.Knowledge distillation-based domain-invariant representation learning for domain generalization.IEEE Transactions on Multimedia, pp.  1–11, 2023.doi: 10.1109/TMM.2023.3263549. +Norton & Royset (2021) +↑ + Matthew D Norton and Johannes O Royset.Diametrical risk minimization: Theory and computations.Machine Learning, pp.  1–19, 2021. +Peng et al. (2019) +↑ + Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang.Moment matching for multi-source domain adaptation.In Proceedings of the IEEE/CVF international conference on computer vision, pp.  1406–1415, 2019. +Radford et al. (2021) +↑ + Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.Learning transferable visual models from natural language supervision.In International conference on machine learning, pp.  8748–8763. PMLR, 2021. +Rame et al. (2022) +↑ + Alexandre Rame, Corentin Dancette, and Matthieu Cord.Fishr: Invariant gradient variances for out-of-distribution generalization.In International Conference on Machine Learning, pp.  18347–18377. PMLR, 2022. +Seo et al. (2020) +↑ + Seonguk Seo, Yumin Suh, Dongwan Kim, Geeho Kim, Jongwoo Han, and Bohyung Han.Learning to optimize domain specific normalization for domain generalization.In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XXII 16, pp.  68–83. Springer, 2020. +Shao et al. (2019) +↑ + Rui Shao, Xiangyuan Lan, Jiawei Li, and Pong C Yuen.Multi-adversarial discriminative deep domain generalization for face presentation attack detection.In CVPR, pp.  10023–10031, 2019. +Shi et al. (2021) +↑ + Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve.Gradient matching for domain generalization.arXiv preprint arXiv:2104.09937, 2021. +Shu et al. (2023) +↑ + Yang Shu, Xingzhuo Guo, Jialong Wu, Ximei Wang, Jianmin Wang, and Mingsheng Long.Clipood: Generalizing clip to out-of-distributions.In International Conference on Machine Learning, 2023. +Sun & Saenko (2016) +↑ + Baochen Sun and Kate Saenko.Deep coral: Correlation alignment for deep domain adaptation.In ECCV, pp.  443–450. Springer, 2016. +Venkateswara et al. (2017) +↑ + Hemanth Venkateswara, Jose Eusebio, Shayok Chakraborty, and Sethuraman Panchanathan.Deep hashing network for unsupervised domain adaptation.In Proceedings of the IEEE conference on computer vision and pattern recognition, pp.  5018–5027, 2017. +Wang et al. (2023a) +↑ + Jindong Wang, Cuiling Lan, Chang Liu, Yidong Ouyang, Tao Qin, Wang Lu, Yiqiang Chen, Wenjun Zeng, and Philip S. Yu.Generalizing to unseen domains: A survey on domain generalization.IEEE Transactions on Knowledge and Data Engineering, 35(8):8052–8072, 2023a.doi: 10.1109/TKDE.2022.3178128. +Wang et al. (2023b) +↑ + Pengfei Wang, Zhaoxiang Zhang, Zhen Lei, and Lei Zhang.Sharpness-aware gradient matching for domain generalization.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  3769–3778, 2023b. +Wang et al. (2020) +↑ + Shujun Wang, Lequan Yu, Caizi Li, Chi-Wing Fu, and Pheng-Ann Heng.Learning from extrinsic and intrinsic supervisions for domain generalization.In ECCV, 2020. +Xu et al. (2021) +↑ + Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yanfeng Wang, and Qi Tian.A fourier-based framework for domain generalization.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  14383–14392, 2021. +Xu et al. (2023a) +↑ + Qinwei Xu, Ruipeng Zhang, Ziqing Fan, Yanfeng Wang, Yi-Yan Wu, and Ya Zhang.Fourier-based augmentation with applications to domain generalization.Pattern Recognition, 139:109474, 2023a. +Xu et al. (2023b) +↑ + Qinwei Xu, Ruipeng Zhang, Yi-Yan Wu, Ya Zhang, Ning Liu, and Yanfeng Wang.Simde: A simple domain expansion approach for single-source domain generalization.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  4797–4807, 2023b. +Xu et al. (2023c) +↑ + Qinwei Xu, Ruipeng Zhang, Ya Zhang, Yi-Yan Wu, and Yanfeng Wang.Federated adversarial domain hallucination for privacy-preserving domain generalization.IEEE Transactions on Multimedia, pp.  1–13, 2023c.doi: 10.1109/TMM.2023.3257566. +Yan et al. (2020) +↑ + Shen Yan, Huan Song, Nanxiang Li, Lincan Zou, and Liu Ren.Improve unsupervised domain adaptation with mixup training.arXiv preprint arXiv:2001.00677, 2020. +Zhang et al. (2023a) +↑ + Lei Zhang, Yingjun Du, Jiayi Shen, and Xiantong Zhen.Learning to learn with variational inference for cross-domain image classification.IEEE Transactions on Multimedia, 25:3319–3328, 2023a.doi: 10.1109/TMM.2022.3158072. +Zhang et al. (2022) +↑ + Ruipeng Zhang, Qinwei Xu, Chaoqin Huang, Ya Zhang, and Yanfeng Wang.Semi-supervised domain generalization for medical image analysis.In 2022 IEEE 19th International Symposium on Biomedical Imaging (ISBI), pp.  1–5. IEEE, 2022. +Zhang et al. (2023b) +↑ + Ruipeng Zhang, Qinwei Xu, Jiangchao Yao, Ya Zhang, Qi Tian, and Yanfeng Wang.Federated domain generalization with generalization adjustment.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  3954–3963, 2023b. +Zhang et al. (2023c) +↑ + Xingxuan Zhang, Renzhe Xu, Han Yu, Hao Zou, and Peng Cui.Gradient norm aware minimization seeks first-order flatness and improves generalization.In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.  20247–20257, 2023c. +Zhao et al. (2022) +↑ + Yang Zhao, Hao Zhang, and Xiuyuan Hu.Penalizing gradient norm for efficiently improving generalization in deep learning.In International Conference on Machine Learning, pp.  26982–26992. PMLR, 2022. +Zhou et al. (2020a) +↑ + Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang.Deep domain-adversarial image generation for domain generalisation.In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp.  13025–13032, 2020a. +Zhou et al. (2020b) +↑ + Kaiyang Zhou, Yongxin Yang, Timothy Hospedales, and Tao Xiang.Learning to generate novel domains for domain generalization.In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XVI 16, pp.  561–578. Springer, 2020b. +Zhou et al. (2021) +↑ + Kaiyang Zhou, Yongxin Yang, Yu Qiao, and Tao Xiang.Domain generalization with mixstyle.arXiv preprint arXiv:2104.02008, 2021. +Zhou et al. (2022a) +↑ + Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu.Learning to prompt for vision-language models.International Journal of Computer Vision, 130(9):2337–2348, 2022a. +Zhou et al. (2022b) +↑ + Zhihan Zhou, Jiangchao Yao, Yan-Feng Wang, Bo Han, and Ya Zhang.Contrastive learning with boosted memorization.In International Conference on Machine Learning, pp.  27367–27377. PMLR, 2022b. +Zhou et al. (2023) +↑ + Zhihan Zhou, Jiangchao Yao, Feng Hong, Ya Zhang, Bo Han, and Yanfeng Wang.Combating representation learning disparity with geometric harmonization.In Thirty-seventh Conference on Neural Information Processing Systems, 2023. +Zhuang et al. (2022) +↑ + Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C Dvornek, sekhar tatikonda, James s Duncan, and Ting Liu.Surrogate gap minimization improves sharpness-aware training.In International Conference on Learning Representations, 2022. +                                               Appendix +\parttoc +Appendix ARelated Work +A.1Sharpness-Aware Minimization (SAM) + +Numerous studies (Hochreiter & Schmidhuber, 1994a; Li et al., 2018b; Dinh et al., 2017b) have been conducted to enhance our understanding of the generalization capabilities of deep learning models through an exploration of the geometric properties of the loss landscape. These investigations have consistently demonstrated that deep neural networks exhibiting a flat minimum tend to exhibit superior generalization performance. In order to obtain a flat minimum, the Sharpness-Aware Minimization (SAM) approach (Foret et al., 2021) was proposed, which utilizes a base optimizer to simultaneously minimize both the vanilla training loss and the sharpness metric. The sharpness metric, as defined by (Keskar et al., 2017a), quantifies the flatness of a minimum through the eigenvalues of the Hessian matrix. In practice, SAM involves obtaining a fixed-length perturbation through gradient ascent on the initial parameter, followed by updating the gradient based on this perturbed parameter with respect to the initial parameter. Although SAM can result in a flat minimum and substantially enhance the generalization capability, it incurs a twofold increase in computational overhead. The variants of SAM have been extensively investigated from two perspectives: the first pertains to the enhancement of SAM’s generalizability (Kwon et al., 2021; Zhuang et al., 2022; Zhang et al., 2023c; Zhao et al., 2022; Wang et al., 2023b), while the second focuses on improving its efficiency (Liu et al., 2022; Du et al., 2022a; b; Mi et al., 2022). + +A.1.1Generalizability improvement of SAM + +One key problem of SAM is that the perturbation obtained by gradient ascent might disagree with sharpness since gradient ascent is only a first-order approximation of the sharpness calculation. Zhuang et al. (2022) introduced a surrogate gap to enhance the evaluation of sharpness, while (Wang et al., 2023b) integrated the perturbed loss and the surrogate gap from (Zhuang et al., 2022) into a unified objective. Additionally, (Zhao et al., 2022) revealed that SAM inherently optimizes both the empirical risk loss and the corresponding gradient norm. Besides, FisherSAM (Kim et al., 2022) and ASAM (Kwon et al., 2021) achieved improved perturbations by adjusting the scales of the perturbation magnitudes. (Zhang et al., 2023c) further proposed Gradient norm Aware Minimization (GAM), which regularized the Hessian of the gradient norm. VaSSO (Li & Giannakis, 2023) focuses on addressing the issue of SAM’s subpar performance in perturbation direction generation due to the noise introduced by mini-batch sampling. + +A.1.2Efficiency improvement of SAM + +Due to the doubled overhead of SAM in comparison to a base optimizer like SGD (Stochastic Gradient Descent), considerable efforts have been devoted to mitigating this overhead. (Liu et al., 2022) introduced LookSAM as a means to reduce the number of perturbations. Meanwhile, (Mi et al., 2022) achieved sparse perturbations through the use of a binary mask. Furthermore, Du et al.explored various proxy methods (ESAM (Du et al., 2022a), SAF (Du et al., 2022b)) for computing perturbations, thereby replacing the gradient ascent derivation process employed in SAM. + +A.2Domain Generalization + +Domain generalization is a vital research direction that focuses on training models capable of generalizing well to unseen domains by leveraging knowledge from multiple source domains (Wang et al., 2023a). Over the past decade, several methods have been proposed to address the challenges of domain generalization. These methods can be broadly categorized into five main approaches: domain alignment, meta-learning, domain hallucination, domain disentanglement, and robustness training. In this section, we provide a brief overview of each of these categories. + +A.2.1Domain alignment-based method + +The goal of domain alignment is to mitigate discrepancies among distinct source domains by aligning the marginal feature distributions to extract domain-invariant representations. This objective can be accomplished using various strategies, including adversarial training (Li et al., 2018d; Shao et al., 2019), maximum mean discrepancy (Li et al., 2018c), moment matching (Sun & Saenko, 2016), self-supervised learning (Wang et al., 2020), or contrastive learning (Kim et al., 2021; Motiian et al., 2017; Zhou et al., 2022b; 2023). All of these methods improve generalization across unseen domains by either directly or indirectly reducing the discrepancy between different feature distributions and imposing domain-invariant constraints on these discriminative features. + +A.2.2Meta Learning-based Methods + +These approaches aim to address unforeseen domain shifts and enhance the generalizability of models to such shifts through meta-optimization, achieved by partitioning the training domains into distinct meta-train and meta-test domains. (Li et al., 2018a) first introduced meta learning into DG, following the concept of Modal-Agnostic Meta-Learning (MAML) (Finn et al., 2017). Subsequently, (Balaji et al., 2018) designed a weight regularizer based on the meta-learning framework, while (Li et al., 2019) chose to meta-learn a feature critic loss. (Dou et al., 2019) constrained the invariance of learned semantic relations between the meta-train and meta-test domains. Additionally, (Zhang et al., 2023a) integrated meta learning into a Bayesian framework and enforced the model to learn a meta-variational distribution to enhance knowledge transfer. + +A.2.3Domain hallucination-based methods + +Domain hallucination, also known as data augmentation in the presence of domain shifts, aims to encompass a wider range of domain variations by generating additional training samples from fictional domains while preserving their semantic integrity. Early approaches such as (Xu et al., 2021; 2023a; Zhang et al., 2022; Xu et al., 2023b; Zhou et al., 2020a; Yan et al., 2020; Zhou et al., 2020b; Carlucci et al., 2019; Xu et al., 2023c) involve cross-domain data augmentation in the input space and can be categorized into non-parametric and adversarial sample-based approaches. Non-parametric methods (Xu et al., 2021; 2023a; Yan et al., 2020; Zhang et al., 2022) employ traditional image transformations to achieve enhancement, while adversarial sample-based methods (Xu et al., 2023b; Zhou et al., 2020a; b; Carlucci et al., 2019; Xu et al., 2023c) generate samples from a new domain through adversarial training. Adversarial training ensures the quality of generation by enforcing consistency in terms of category among the samples from the generated fictional domain. Some recent work focuses on augmentation in the latent space (Liu et al., 2023; Zhou et al., 2021), which achieves more efficient augmentation perturbations by applying perturbations to the latent features to improve the generalization of the model. + +A.2.4Domain disentanglement-based methods + +In contrast to the majority of domain generalization approaches that aim for domain invariance, disentanglement-based approaches focus on separating the domain-invariant and domain-specific components. To achieve this, Seo et al. (2020) introduced domain-specific batch normalization (Chang et al., 2019) for each training domain, effectively balancing feature discrimination and invariance. In a similar vein,Jin et al. (2022) designed a style restitution module that encourages the separation of task-relevant and task-irrelevant features. Furthermore, Niu et al. (2023) proposed a two-stage distillation approach, aimed at learning a domain-invariant representation while preserving domain-specific features. + +A.2.5Robustness training-based methods +Table 4:Comparison of SAM-based methods and other robustness training-based methods on the optimization objective. +Method Total Optimization Function Optimization on +𝑤 + Optimization on +𝜖 + +ERM +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 +) + Same to left +× + +V-REx +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 +) ++ +𝛽 +⁢ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +) +} +𝑖 += +1 +𝑀 + Same to left +× + +Fish +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 +) +− +𝛾 +⁢ +2 +𝑀 +⁢ +( +𝑀 +− +1 +) +⁢ +∑ +𝑖 +, +𝑗 +∈ +[ +1 +, +𝑀 +] +𝑖 +≠ +𝑗 +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) +⋅ +∇ +ℒ +𝑗 +⁢ +( +𝑤 +) + Same to left +× + +Fishr +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 +) +− +𝜆 +⁢ +1 +𝑀 +⁢ +∑ +𝑖 += +1 +𝑀 +‖ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +) +− +∇ +ℒ +⁢ +( +𝑤 +) +‖ +2 + Same to left +× + +SAM +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) + +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) + +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) + +DISAM +min +𝑤 +⁡ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +[ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +− +𝜆 +⁢ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +^ ++ +𝜖 +) +} +𝑖 += +1 +𝑀 +] + +min +𝑤 +⁢ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) + +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +[ +∑ +𝑖 += +1 +𝑀 +𝛼 +𝑖 +⁢ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +− +𝜆 +⁢ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 ++ +𝜖 +) +} +𝑖 += +1 +𝑀 +] + +The objective of robustness training-based methods is to incorporate constraints that enhance the model’s robustness or flatness during the training process. Robustness-related methods aim to learn domain-invariant representations by employing a technique known as Invariant Risk Minimization (IRM) (Arjovsky et al., 2019). By minimizing the risk across different domains, these methods (Arjovsky et al., 2019; Krueger et al., 2021; Norton & Royset, 2021; Shi et al., 2021; Rame et al., 2022; Li & Giannakis, 2023) seek to learn features that are insensitive to domain variations, thereby improving the model’s ability to generalize. On the other hand, a separate class of flatness-related methods (Izmailov et al., 2018; Cha et al., 2021; Zhang et al., 2023b; Foret et al., 2021; Wang et al., 2023b) aims to address the effects of domain shifts by identifying flat minima. These methods strive to find regions in the loss landscape where small perturbations in the input have minimal impact on the model’s predictions. By leveraging flat minima, these methods enhance the model’s robustness to domain variations. + +In Table 4, we provide the comparison of optimization objectives between representative algorithms in the two categories. Domain-invariant methods solely concentrate on optimizing the parameter +𝑤 +. For instance, V-REx (Krueger et al., 2021) directly minimizes the variance of the domain loss, which can have a detrimental effect on convergence. Similarly, Fish (Shi et al., 2021) and Fishr (Rame et al., 2022) impose constraints on gradient updates. SAM-based methods require the estimation of sharpness, so in addition to optimizing the parameters +𝑤 +, they also need to optimize the perturbation directions +𝜖 +. + +This paper primarily concentrates on the flatness-based method, which encompasses two main approaches for enhancing the flatness of the model. The first approach involves leveraging the self-ensemble of multiple minima attained during the training process to passively acquire a result that favors flatness minima. Notable examples of this approach include Stochastic Weight Averaging (SWA) (Izmailov et al., 2018) and Stochastic Weight Averaging Densely (SWAD) (Cha et al., 2021). The second approach involves directly optimizing for flatness and is referred to as Sharpness-Aware Minimization (SAM) (Foret et al., 2021). In the subsequent section, we will delve into a comprehensive review of the relevant literature pertaining to these approaches. + +Appendix BDetails of DISAM +B.1Comparative Analysis of DISAM Versus General Convergence Consistency + +Here, we present a thorough examination of the distinctions between our proposed DISAM framework and the broader, conventional convergence consistency issue like V-REx(Krueger et al., 2021) and Fishr(Rame et al., 2022). Specifically, we address the following aspects: + +• + +Distinct Focus: DISAM focuses on the issue where SAM-based methods are unable to accurately estimate sharpness in domain shift scenarios, leading to the ineffective sharpness minimization and reduction in generalization performance. + +• + +Enhancing on Top of General Methods: While traditional solutions(Krueger et al., 2021; Rame et al., 2022; Shi et al., 2021) aim at convergence consistency in parameter optimization, DISAM’s methodology is distinct and orthogonal. It builds upon methods like V-REx(Krueger et al., 2021) and Fishr(Rame et al., 2022), but goes further in enhancing out-of-domain generalization through better sharpness minimization. This is evident in our experiments, where combining DISAM with Fishr results in significant performance gains (shown in Table 5). + +We also provide extensive experimental results to validate DISAM’s effectiveness and its practical implications in various domain-shift scenarios. + +Table 5:Comparison with other general convergence consistency methods. +Algorithm PACS VLCS OfficeHome TerraInc DomainNet Avg. +V-REx 84.9 78.3 66.4 46.4 33.6 61.9 +V-REx + DISAM 85.8 78.4 70.5 45.9 42.3 64.6 +Fishr 86.9 78.2 68.2 53.6 41.8 65.7 +Fishr + DISAM 87.5 79.2 70.7 54.8 43.9 67.2 + +It is imperative to reiterate the contributions of our DISAM. We provide a detailed exposition of how simplistic applications of SAM compromise training robustness, especially when dealing with domain shifts. DISAM strategically mitigates these issues by finely tuning the perturbation vectors and their location points, thus significantly enhancing model generalization. Furthermore, we underscore the notable enhancements achieved with DISAM, as corroborated by comprehensive experimental analyses and the ensuing performance metrics. + +B.2Algorithm of DISAM + +We give specific algorithmic details for DISAM in Algorithm 1, and the python code implementation is in Appendix D. + +0:  Source Domains +𝒮 += +{ +𝐷 +1 +, +⋯ +, +𝐷 +𝑀 +} +, initial model +𝑤 +1 +, perturbation ratio +𝜌 +, variance constraint weight +𝜆 +, learning rate +𝜂 +𝑡 +, training iterations +𝑇 +. +0:  Generalization model +𝑤 +𝑇 +. +1:  for  +𝑡 + in +1 +⁢ +⋯ +⁢ +𝑇 + do +2:     Sample mini-batch +𝐵 += +{ +𝐵 +1 +, +⋯ +, +𝐵 +𝑀 +} +⊆ +𝒮 +, where +𝐵 +𝑖 +⊆ +𝐷 +𝑖 + and +| +𝐵 +𝑖 +| +≥ +0 +. +3:     Compute the domain-inspired loss gradient: +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +𝑡 +; +𝐵 +) += +∇ +ℒ +⁢ +( +𝑤 +𝑡 +; +𝐵 +) +− +𝜆 +⁢ +∇ +Var +⁢ +{ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +; +𝐵 +𝑖 +} +𝑖 += +1 +𝑀 +. +4:     Get the perturbation weight: +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 += +𝑤 +𝑡 ++ +𝜌 +⁢ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +𝑡 +; +𝐵 +) +‖ +∇ +ℒ +𝐷 +⁢ +𝐼 +⁢ +( +𝑤 +𝑡 +; +𝐵 +) +‖ +. +5:     Update weights: +𝑤 +𝑡 ++ +1 += +𝑤 +𝑡 +− +𝜂 +𝑡 +∇ +ℒ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +; +𝐵 +) +} +. +6:  end for +Algorithm 1 Domain-Inspired Sharpness-Aware Minimization (DISAM). +B.3Proof of DISAM’s Convergence +Theorem 1. + +(Convergence During Training). Consider a non-convex function +ℒ +⁢ +( +w +) + with Lipschitz-smooth constant +L + and lower bound +ℒ +m +⁢ +i +⁢ +n +. With the bounded norm assumption of noisy stochastic gradients ( +‖ +∇ +ℒ +p +⁢ +( +w +) +‖ +2 +≤ +G +) at the t-step, the learning rate +η +t += +η +0 +/ +t + and a fixed perturbation amplitude +ρ +, we have: + + +1 +𝑇 +⁢ +∑ +𝑡 += +1 +𝑇 +𝔼 +⁢ +‖ +∇ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 +) +‖ +2 +2 +≤ +ℒ +𝑝 +⁢ +( +𝑤 +0 +) +− +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 +𝜂 +0 +⁢ +1 +𝑇 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +⁢ +log +⁡ +( +𝑇 +) +𝑇 + +(9) + +where in SAM, +Γ += +𝐿 + and when use DISAM +Γ +≤ +𝐿 +. + +Proof. + +For simplicity of notation, we denote the update at step +𝑡 + as +𝑑 +𝑡 += +− +𝜂 +𝑡 +⁢ +𝑔 +𝑝 +( +𝑡 +) +, where +𝜂 +𝑡 + is the decayed learning rate and +𝑔 +𝑝 +𝑡 + is the expected gradient of perturbation loss +ℒ +𝑝 +. By +𝐿 +-smoothness of the loss function +ℒ + and the definition of +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 +) += +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, where +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 + represents the parameters after the perturbation of gradient ascent, we have: + + +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 ++ +1 +) += +ℒ +⁢ +( +𝑤 +𝑡 ++ +1 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) ++ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝑤 +𝑡 ++ +1 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +− +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +⟩ ++ +𝐿 +2 +⁢ +‖ +𝑤 +𝑡 ++ +1 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +− +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +‖ +2 + +(10) + +where +𝐿 + is the Lipschitz constant of loss +ℒ + and with the definition of +𝑑 +𝑡 += +𝑤 +𝑡 ++ +1 +− +𝑤 +𝑡 + and +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 += +𝑤 +𝑡 ++ +𝜖 +𝑡 +, we have: + + +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 ++ +1 +) + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) ++ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝑤 +𝑡 ++ +1 ++ +𝜖 +𝑡 ++ +1 +− +𝑤 +𝑡 +− +𝜖 +𝑡 +⟩ ++ +𝐿 +2 +⁢ +‖ +𝑤 +𝑡 ++ +1 ++ +𝜖 +𝑡 ++ +1 +− +𝑤 +𝑡 +− +𝜖 +𝑡 +‖ +2 + +(11) + + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) ++ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝑑 +𝑡 +⟩ ++ +𝐿 +⁢ +‖ +𝑑 +𝑡 +‖ +2 ++ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +⟩ ++ +𝐿 +⁢ +‖ +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +‖ +2 + + +Let us take the expectation conditioned on observations up to step +𝑡 +. For the sake of simplicity, we use the symbol +𝔼 + to denote the expectation over all possible data points on the training data distribution. Moreover, given the observations up to step +𝑡 +, we can use the definition of +𝑑 +𝑡 + to obtain: + + +𝔼 +⁢ +[ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 ++ +1 +) +] + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +− +𝜂 +𝑡 +⁢ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝔼 +⁢ +[ +𝑔 +𝑝 +( +𝑡 +) +] +⟩ ++ +𝜂 +𝑡 +2 +⁢ +𝐿 +⁢ +𝔼 +⁢ +‖ +𝑔 +𝑝 +( +𝑡 +) +‖ +2 + +(12) + + ++ +𝔼 +⁢ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +⟩ ++ +𝐿 +⁢ +𝔼 +⁢ +‖ +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +‖ +2 + + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +− +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +𝜂 +𝑡 +2 +⁢ +𝐿 +⁢ +𝐺 +2 + + ++ +𝔼 +⁢ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +⟩ ++ +𝐿 +⁢ +𝔼 +⁢ +‖ +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +‖ +2 + + +By the definition of +𝜖 +𝑡 +, we have: + + +𝜖 +𝑡 += +𝜌 +⁢ +𝑔 +( +𝑡 +) +‖ +𝑔 +( +𝑡 +) +‖ +, +𝜖 +𝑡 ++ +1 += +𝜌 +⁢ +𝑔 +( +𝑡 ++ +1 +) +‖ +𝑔 +( +𝑡 ++ +1 +) +‖ + +(13) + +where +𝑔 +( +𝑡 +) + is the gradient of +ℒ + at +𝑤 +𝑡 + with the domain-inspired gradient in Eq.( 8). We denote +𝜖 +𝑡 += +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) += +∑ +𝑖 += +1 +𝑀 +𝛽 +𝑡 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +, where +𝛽 +𝑡 +𝑖 += +𝛼 +𝑖 +− +2 +⁢ +𝜆 +𝑀 +⁢ +( +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +− +1 +𝑀 +⁢ +∑ +𝑗 += +1 +𝑀 +ℒ +𝑗 +⁢ +( +𝑤 +𝑡 +) +) +. Since both +𝜖 +𝑡 + and +𝜖 +𝑡 ++ +1 + are unit length vectors, +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 + can be bounded by the arc length +𝜙 +𝑡 + between them. Here the difference vector between +𝜖 +𝑡 ++ +1 + and +𝜖 +𝑡 + can be regarded as a random noise in the gradient direction and in SAM +𝜌 +≫ +𝜂 +𝑡 +, so the expectation of the inner product with the gradient direction +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) + can be approximated as 0 ( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 + is updated from +𝑤 +𝑡 + with a larger step size +𝜌 +, and its gradient direction can be considered approximately independent of the gradient direction in the neighborhood of +𝑤 +𝑡 +, so its difference with the inner product between +𝜖 +𝑡 ++ +1 + and +𝜖 +𝑡 + is negligible). Therefore, we have: + + +𝔼 +⁢ +[ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 ++ +1 +) +] + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +− +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +𝜂 +𝑡 +2 +⁢ +𝐿 +⁢ +𝐺 +2 + +(14) + + ++ +𝔼 +⁢ +[ +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝜖 +𝑡 ++ +1 +⟩ +− +⟨ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +, +𝜖 +𝑡 +⟩ +] ++ +𝐿 +⁢ +𝜌 +2 +⁢ +𝔼 +⁢ +‖ +𝑔 +( +𝑡 ++ +1 +) +‖ +𝑔 +( +𝑡 ++ +1 +) +‖ +− +𝑔 +( +𝑡 +) +‖ +𝑔 +( +𝑡 +) +‖ +‖ +2 + + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +− +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +𝜂 +𝑡 +2 +⁢ +𝐿 +⁢ +𝐺 +2 ++ +𝐿 +⁢ +𝜌 +2 +⁢ +𝜙 +𝑡 +2 + + +Because of the continuity of the optimization, the angle between the gradient perturbations before and after is small. Therefore, we can get: + + +𝜙 +𝑡 + +≈ +tan +⁡ +𝜙 +𝑡 += +‖ +𝜖 +𝑡 ++ +1 +− +𝜖 +𝑡 +‖ +‖ +𝜖 +𝑡 +‖ ++ +𝑂 +⁢ +( +𝜙 +𝑡 +2 +) += +‖ +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 ++ +1 +) +− +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) +‖ +‖ +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) +‖ ++ +𝑂 +⁢ +( +𝜙 +𝑡 +2 +) + +(15) + + += +‖ +∑ +𝑖 += +1 +𝑀 +( +𝛽 +𝑡 ++ +1 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) +− +𝛽 +𝑡 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +) +‖ +‖ +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) +‖ ++ +𝑂 +⁢ +( +𝜙 +𝑡 +2 +) + + += +‖ +∑ +𝑖 += +1 +𝑀 +( +( +𝛽 +𝑡 ++ +1 +𝑖 +− +𝛽 +𝑡 +𝑖 +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) ++ +𝛽 +𝑡 +𝑖 +⁢ +( +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) +− +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +) +) +‖ +‖ +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) +‖ ++ +𝑂 +⁢ +( +𝜙 +𝑡 +2 +) + + +Here we consider the effect of the weight coefficients generated by DISAM in the perturbation of +∇ +ℒ +𝑑 +, for the part of +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) + that is large, +𝛽 +𝑡 +𝑖 + is smaller, we assume that the larger +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) + is, the larger the corresponding gradient +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) + is also, and after one optimization process, the variability between the domains will be reduced, so +𝛽 +𝑡 ++ +1 +𝑖 + is a little bit smaller than the weight of +𝛽 +𝑡 +𝑖 +, in the place where the gradient is large, and by the rearranging inequality, we can obtained: + + +∑ +𝑖 += +1 +𝑀 +𝛽 +𝑡 ++ +1 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) +≤ +∑ +𝑖 += +1 +𝑀 +𝛽 +𝑡 ++ +1 +𝑖 +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) + +(16) + +So bring Eq.( 16) to Eq.( 15), and with +∇ +ℒ +⁢ +( +𝑤 +𝑡 ++ +1 +) += +∇ +ℒ +⁢ +( +𝑤 +𝑡 ++ +𝑑 +𝑡 +) += +∇ +ℒ +⁢ +( +𝑤 +𝑡 +) ++ +𝐻 +⁢ +𝑑 +𝑡 ++ +𝑂 +⁢ +( +‖ +𝑑 +𝑡 +‖ +2 +) + we can get: + + +𝜙 +𝑡 +≤ +‖ +∑ +𝑖 += +1 +𝑀 +( +𝛽 +𝑡 ++ +1 +𝑖 +− +𝛽 +𝑡 +𝑖 +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) ++ +𝐻 +⁢ +𝑑 +𝑡 ++ +𝑂 +⁢ +( +‖ +𝑑 +𝑡 +‖ +2 +) +‖ +‖ +∇ +ℒ +𝑑 +⁢ +( +𝑤 +𝑡 +) +‖ ++ +𝑂 +⁢ +( +𝜙 +𝑡 +2 +) +≤ +𝜂 +𝑡 +⁢ +Γ + +(17) + +Here since +∑ +𝑖 += +1 +𝑀 +( +𝛽 +𝑡 ++ +1 +𝑖 +− +𝛽 +𝑡 +𝑖 +) +⁢ +∇ +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 ++ +1 +) +≤ +0 +, we use +Γ + to denote an upper bound that is smaller than +𝐿 +. + +Plug Eq.( 17) into Eq.( 14), we have: + + +𝔼 +⁢ +[ +ℒ +𝑝 +⁢ +( +𝑤 +𝑡 ++ +1 +) +] + +≤ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +− +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +𝜂 +𝑡 +2 +⁢ +𝐿 +⁢ +𝐺 +2 ++ +𝐿 +⁢ +𝜌 +2 +⁢ +𝜂 +𝑡 +2 +⁢ +Γ +2 + +(18) + +Perform telescope sum and note that +𝜂 +𝑇 += +𝜂 +0 +𝑇 +, we have: + + +𝔼 +⁢ +ℒ +𝑝 +⁢ +( +𝑤 +𝑇 +) +− +ℒ +𝑝 +⁢ +( +𝑤 +0 +) + +≤ +− +∑ +𝑡 += +1 +𝑇 +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +2 +⁢ +∑ +𝑡 += +1 +𝑇 +1 +𝑡 + +(19) + + +≤ +− +∑ +𝑡 += +1 +𝑇 +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +2 +⁢ +log +⁡ +( +𝑇 +) + + +Hence, + + +𝜂 +𝑇 +⁢ +∑ +𝑡 += +1 +𝑇 +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 +≤ +∑ +𝑡 += +1 +𝑇 +𝜂 +𝑡 +⁢ +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 +≤ +ℒ +𝑝 +⁢ +( +𝑤 +0 +) +− +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +2 +⁢ +log +⁡ +( +𝑇 +) + +(20) + +Note that +𝜂 +𝑇 += +𝜂 +0 +𝑇 +, we have: + + +1 +𝑇 +⁢ +∑ +𝑡 += +1 +𝑇 +𝔼 +⁢ +‖ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +𝑎 +⁢ +𝑠 +⁢ +𝑐 +) +‖ +2 +2 +≤ +ℒ +𝑝 +⁢ +( +𝑤 +0 +) +− +ℒ +𝑚 +⁢ +𝑖 +⁢ +𝑛 +𝜂 +0 +⁢ +1 +𝑇 ++ +( +𝐿 +⁢ +𝐺 +2 ++ +𝜌 +2 +⁢ +𝐿 +⁢ +Γ +2 +) +⁢ +𝜂 +0 +⁢ +log +⁡ +( +𝑇 +) +𝑇 + +(21) + +∎ + +The influence of +𝜆 +: + +In the proof of Theorem 1, specifically in Eq. (15), +𝜆 + is integrated into +𝛽 +, serving as a hyperparameter that regulates the weight adjustment in DISAM. It functions by modulating the degree of correction for domain shifts: + + +𝛽 +𝑡 +𝑖 += +𝛼 +𝑖 +− +2 +⁢ +𝜆 +𝑀 +⁢ +( +ℒ +𝑖 +⁢ +( +𝑤 +𝑡 +) +− +1 +𝑀 +⁢ +∑ +𝑗 += +1 +𝑀 +ℒ +𝑗 +⁢ +( +𝑤 +𝑡 +) +) + + +The choice of +𝜆 + influences how aggressively DISAM responds to variance or domain shifts, with a higher +𝜆 + leading to more pronounced adjustments in +𝛽 +. Our experimental analysis in Figure 5(5(c)) and 5(5(d)), reveals that DISAM’s performance remains relatively stable across a wide range of +𝜆 + values. However, choosing too large +𝜆 + can result in overly aggressive early training adjustments, yielding the increased variance in repeated experiments. Consequently, we adopted a default +𝜆 + value of 0.1 in all experiments. + +B.4Discussion of the role of +𝜌 + in DISAM + +Here, we provide a detailed discussion on how +𝜌 + affects both generalization and convergence. First, we introduce the generalization theorem of the upper bound on generalization error, which is only related to the magnitude of +𝜌 +, and DISAM follows the same upper bound on generalization error as SAM. In the SAM framework, the parameter +𝜌 + plays a crucial role in determining generalizability. As established in SAM (Foret et al., 2021), there exists an upper bound on the generalization error for SAM, suggesting that a larger +𝜌 + could potentially enhance generalization, provided that convergence is not impeded. Here is the relevant generalization bound from SAM (Foret et al., 2021): + +Theorem 2. + +(Generalization Bound of SAM). For any +ρ +> +0 + and any distribution +𝒟 +, with probability +1 +− +δ + over the choice of the training set +S +∼ +𝒟 +, + + +ℒ +𝒟 +⁢ +( +𝑤 +) +≤ +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +ℒ +𝑆 +⁢ +( +𝑤 ++ +𝜖 +) ++ +𝑘 +⁢ +log +⁡ +( +1 ++ +‖ +𝑤 +‖ +2 +2 +𝜌 +2 +⁢ +( +1 ++ +log +⁡ +( +𝑛 +) +𝑘 +) +2 +) ++ +4 +⁢ +log +⁡ +𝑛 +𝛿 ++ +𝑂 +~ +⁢ +( +1 +) +𝑛 +− +1 + +(22) + +where +𝑛 += +| +𝑆 +| +, +𝑘 + is the number of parameters and we assumed +ℒ +𝒟 +⁢ +( +𝑤 +) +≤ +𝔼 +𝜖 +𝑖 +≈ +𝒩 +⁢ +( +0 +, +𝜌 +) +⁢ +[ +ℒ +𝒟 +⁢ +( +𝑤 ++ +𝜖 +) +] +. This theorem’s proof focuses solely on the magnitude of +𝜌 +, thus affirming the applicability of this theoretical framework to DISAM. + +When considering domain shift, the upper bound on generalization error for the test domain is: + +Theorem 3. + +(PAC-Bayesian Generalization Bound). For any +ρ +> +0 + and the unseen domain +T +, suppose we have multi-source domains +S += +{ +S +1 +, +S +2 +, +⋯ +} + with a total of +N + samples. Let +ℋ + be the hypothesis space and +Ω + be the corresponding parameter space, where the VC dimension of +ℋ + is +d +. We denote the domain divergence between two domains +D +i + and +D +j + on the hypothesis space +ℋ + as +d +ℋ +⁢ +Δ +⁢ +ℋ +⁢ +( +D +i +, +D +j +) +. Then, for any +δ +∈ +( +0 +, +1 +) +, with probability at least +1 +− +δ +, for all +w +∈ +Ω +, we have: + + +ℒ +𝑇 +⁢ +( +𝑤 +) +≤ + +max +‖ +𝜖 +‖ +2 +≤ +𝜌 +⁡ +ℒ +𝑆 +⁢ +( +𝑤 ++ +𝜖 +) ++ +1 +2 +⁢ +𝑑 +ℋ +⁢ +Δ +⁢ +ℋ +⁢ +( +𝑆 +, +𝑇 +) ++ +log +⁡ +𝑑 ++ +log +⁡ +1 +𝛿 +2 +⁢ +𝑁 ++ +𝜆 + +(23) + + ++ +𝑘 +⁢ +log +⁡ +( +1 ++ +‖ +𝑤 +‖ +2 +2 +𝜌 +2 +⁢ +( +1 ++ +log +⁡ +( +𝑁 +) +𝑘 +) +2 ++ +4 +⁢ +log +⁡ +𝑁 +𝛿 ++ +𝑂 +~ +⁢ +( +1 +) +) +𝑁 +− +1 + + +where +𝜆 + is the optimal combined risk on +𝑇 + and +𝑆 + that can be achieved by the parameters in +Ω +. + +Combining this with the convergence theorem (Theorem 1), there is a trade-off with respect to +𝜌 +. A larger +𝜌 + might theoretically enhance generalization but poses greater challenges for convergence. This reflects the intuitive notion that searching for flatter minima across a broader range is inherently more complex, which can potentially affect training efficiency. However, if +ℒ +𝑆 +⁢ +( +𝑤 ++ +𝜖 +) + can be converged with a sufficiently small value, a larger +𝜌 + corresponds to better generalization. DISAM, compared to SAM, converges faster, which means that under the same convergence speed, a larger +𝜌 + can be used to achieve better generalization. This advantage is empirically showcased in Figure 3(3(c)) and (3(d)), where we demonstrate that DISAM effectively employs a larger +𝜌 + compared to traditional SAM. This ensures both convergence and enhanced generalization. Such a capability to balance between convergence efficiency and generalization is a distinguishing feature of DISAM over conventional SAM methods. + +Appendix CDetailed Experiments +C.1Detailed Experiment Setups + +We present the detailed results obtained from five datasets, namely PACS (Li et al., 2017) (9,991 images, 7 classes, 4 domains), VLCS (Fang et al., 2013) (10,729 images, 5 classes, 4 domains), OfficeHome (Venkateswara et al., 2017) (15,588 images, 65 classes, 4 domains), TerraIncognita (Beery et al., 2018) (abbreviated as TerraInc, 24,788 images, 10 classes, 4 domains), and DomainNet (Peng et al., 2019) (586,575 images, 345 classes, 6 domains), following the DomainBed benchmark (Gulrajani & Lopez-Paz, 2021) with the ResNet50 backbone architecture. We set the hyperparameters for the Domain-Inspired + SAM method as follows: +𝜌 += +0.5 + and +𝜆 += +0.1 + for PACS, VLCS, OfficeHome, and DomainNet; for TerraInc, we use +𝜌 += +0.01 + and +𝜆 += +0.2 +. Both Domain-Inspired + GSAM and Domain-Inspired + SAGM employ the strategy described in the supplementary material of SAGM (Wang et al., 2023b). As for the CoOp with CLIP, we set the batch size as 32 and the default learning rate as 2e-3. Given the detailed experimental hyperparameter settings provided in the SAGM supplement (Wang et al., 2023b) and the official open-source CLIPOOD code (Shu et al., 2023), we directly applied these official settings. The results, replicated using the official open-source CLIPOOD code, are presented in Table 2 of the main text. + +As for the experiments on open class, we found that CLIPOOD requires a lower learning rate and correspondingly lower +𝜌 +, and therefore used learning rate 1e-07 and +𝜌 + 1e-05 as default settings. + +C.2Detailed Experimental Results + +We present the specific out-of-domain experimental results for each dataset in Table 1, corresponding to each leave-one-domain-out setting. + +Table 6:Comparison with state-of-the-art domain generalization methods. Out-of-domain accuracies on the PACS dataset with ResNet50 backbone. +Algorithm Art Cartoon Photo Sketch Avg. +ERM +84.7 +± +0.4 + +80.0 +± +0.6 + +97.2 +± +0.3 + +79.3 +± +1.0 + +85.5 + +SAM +85.6 +± +2.1 + +80.9 +± +1.2 + +97.0 +± +0.4 + +79.6 +± +1.6 + +85.8 + +Domain-Inspired +87.1 +± +0.4 + +81.9 +± +0.5 + +96.2 +± +0.3 + +83.1 +± +0.7 + +87.1 + +GSAM +86.9 +± +0.1 + +80.4 +± +0.2 + +97.5 +± +0.0 + +78.7 +± +0.8 + +85.9 + +Domain-Inspired +88.4 +± +0.2 + +81.1 +± +0.3 + +97.0 +± +0.0 + +82.3 +± +0.6 + +87.2 + +SAGM +87.4 +± +0.2 + +80.2 +± +0.3 + +98.0 +± +0.2 + +80.8 +± +0.6 + +86.6 + +Domain-Inspired +89.7 +± +0.6 + +81.5 +± +0.0 + +97.0 +± +0.1 + +81.8 +± +0.6 + 87.5 +   +CORAL +89.8 +± +0.5 + +82.9 +± +0.2 + +97.4 +± +0.2 + +83.4 +± +0.2 + 88.4 +Table 7:Comparison with state-of-the-art domain generalization methods. Out-of-domain accuracies on the VLCS dataset with ResNet50 backbone. +Algorithm Caltech LabelMe Pascal Sun Avg. +ERM +98.0 +± +0.3 + +64.7 +± +1.2 + +71.4 +± +1.2 + +75.2 +± +1.6 + +77.3 + +SAM +99.1 +± +0.2 + +65.0 +± +1.0 + +73.7 +± +1.0 + +79.8 +± +0.1 + +79.4 + +Domain-Inspired +99.3 +± +0.0 + +66.3 +± +0.5 + +81.0 +± +0.1 + +73.2 +± +0.1 + +79.9 + +GSAM +98.7 +± +0.3 + +64.9 +± +0.2 + +74.3 +± +0.0 + +78.5 +± +0.8 + +79.1 + +Domain-Inspired +99.8 +± +0.0 + +66.6 +± +0.1 + +74.2 +± +0.9 + +79.3 +± +0.1 + +80.0 + +SAGM +99.0 +± +0.2 + +65.2 +± +0.4 + +75.1 +± +0.3 + +80.7 +± +0.8 + +80.0 + +Domain-Inspired +99.9 +± +0.1 + +66.1 +± +0.6 + +75.1 +± +0.3 + +81.8 +± +0.0 + 80.7 +   +CORAL +99.7 +± +0.1 + +67.8 +± +0.7 + +75.5 +± +0.8 + +81.6 +± +0.2 + 81.2 +Table 8:Comparison with state-of-the-art domain generalization methods. Out-of-domain accuracies on the OfficeHome dataset with ResNet50 backbone. +Algorithm Art Clipart Product Real World Avg. +ERM +61.3 +± +0.7 + +52.4 +± +0.3 + +75.8 +± +0.1 + +76.6 +± +0.3 + +66.5 + +SAM +64.5 +± +0.3 + +56.5 +± +0.2 + +77.4 +± +0.1 + +79.8 +± +0.4 + +69.6 + +Domain-Inspired +65.8 +± +0.2 + +55.6 +± +0.2 + +79.2 +± +0.2 + +80.6 +± +0.1 + +70.3 + +GSAM +64.9 +± +0.1 + +55.2 +± +0.2 + +77.8 +± +0.0 + +79.2 +± +0.2 + +69.3 + +Domain-Inspired +65.7 +± +0.3 + +57.4 +± +0.3 + +79.4 +± +0.1 + +80.7 +± +0.3 + +70.8 + +SAGM +65.4 +± +0.4 + +57.0 +± +0.3 + +78.0 +± +0.3 + +80.0 +± +0.2 + +70.1 + +Domain-Inspired +67.2 +± +0.0 + +56.3 +± +0.3 + +79.6 +± +0.2 + +81.0 +± +0.3 + 71.0 +   +CORAL +68.5 +± +0.1 + +57.6 +± +0.1 + +79.3 +± +0.4 + +81.3 +± +0.2 + 71.7 +Table 9:Comparison with state-of-the-art domain generalization methods. Out-of-domain accuracies on the TerraInc dataset with ResNet50 backbone. +Algorithm L100 L38 L43 L46 Avg. +ERM +49.8 +± +4.4 + +42.1 +± +1.4 + +56.9 +± +1.8 + +35.7 +± +3.9 + +46.1 + +SAM +46.3 +± +1.0 + +38.4 +± +2.4 + +54.0 +± +1.0 + +34.5 +± +0.8 + +43.3 + +Domain-Inspired +46.2 +± +2.9 + +41.6 +± +0.1 + +58.0 +± +0.5 + +40.5 +± +2.2 + +46.6 + +GSAM +50.8 +± +0.1 + +39.3 +± +0.2 + +59.6 +± +0.0 + +38.2 +± +0.8 + +47.0 + +Domain-Inspired +56.7 +± +1.5 + +46.7 +± +1.0 + +59.2 +± +0.7 + +39.9 +± +1.5 + 50.6 +SAGM +54.8 +± +1.3 + +41.4 +± +0.8 + +57.7 +± +0.6 + +41.3 +± +0.4 + +48.8 + +Domain-Inspired +57.6 +± +1.6 + +44.8 +± +1.5 + +58.6 +± +1.2 + +38.9 +± +0.6 + +50.0 + +   + CORAL +57.9 +± +0.3 + +46.6 +± +0.6 + +59.9 +± +0.3 + +42.5 +± +0.1 + +51.7 +Table 10:Comparison with state-of-the-art domain generalization methods. Out-of-domain accuracies on the DomainNet dataset with ResNet50 backbone. +Algorithm Clipart Infograph Painting Quickdraw Real Sketch Avg. +ERM +62.8 +± +0.4 + +20.2 +± +0.3 + +50.3 +± +0.3 + +13.7 +± +0.5 + +63.7 +± +0.2 + +52.1 +± +0.5 + +43.8 + +SAM +64.5 +± +0.3 + +20.7 +± +0.2 + +50.2 +± +0.1 + +15.1 +± +0.3 + +62.6 +± +0.2 + +52.7 +± +0.3 + +44.3 + +Domain-Inspired +65.9 +± +0.2 + +20.7 +± +0.2 + +51.7 +± +0.3 + +16.6 +± +0.3 + +62.8 +± +0.5 + +54.8 +± +0.4 + +45.4 + +GSAM +64.2 +± +0.3 + +20.8 +± +0.2 + +50.9 +± +0.0 + +14.4 +± +0.8 + +63.5 +± +0.2 + +53.9 +± +0.2 + +44.6 + +Domain-Inspired +65.7 +± +0.1 + +21.3 +± +0.1 + +52.2 +± +0.1 + +15.6 +± +0.0 + +64.5 +± +0.2 + +54.1 +± +0.2 + +45.6 + +SAGM +64.9 +± +0.2 + +21.1 +± +0.3 + +51.5 +± +0.2 + +14.8 +± +0.2 + +64.1 +± +0.2 + +53.6 +± +0.2 + +45.0 + +Domain-Inspired +65.9 +± +0.2 + +21.4 +± +0.0 + +52.6 +± +0.1 + +15.8 +± +0.0 + +65.3 +± +0.0 + +54.8 +± +0.2 + 46.0 +   +CORAL +66.4 +± +0.3 + +21.9 +± +0.2 + +53.1 +± +0.1 + +16.1 +± +0.0 + +65.3 +± +0.0 + +55.0 +± +0.0 + +46.3 +C.3Details about Estimated Sharpness on Unseen Test Domain + +Estimating sharpness involves a significant computational overhead. In the earliest methods, Monte Carlo random sampling was the only viable approach (Dinh et al., 2017b; Hochreiter & Schmidhuber, 1994b). However, recent advancements have introduced efficient approximation techniques based on gradients to estimate sharpness (Jiang et al., 2023; 2020). Based on the work of Jiang et al. (2023), we assess the sharpness of the training model on the unseen test domain at the end of each epoch. Sharpness is commonly characterized by the eigenvalues of the Hessian matrix (Keskar et al., 2017b; Dinh et al., 2017a), but direct computation incurs substantial overhead. To address this, a computationally efficient measurement of sharpness is proposed by Jiang et al. (2020), which utilizes the gradient variance +Var +⁢ +{ +∇ +ℒ +⁢ +( +𝑤 +𝑡 +; +𝐵 +𝑡 +) +} + as an estimate ( +𝐵 +𝑡 + represent the batch data sampled at step +𝑡 +). + +C.4Details about Comparison of Computation Cost + +We selected the PACS dataset for experimentation, using a platform with a 16-core CPU, a single RTX3090 GPU, and 64GB RAM. The time overhead for one training step was calculated and averaged over 500 iterations. Due to the lack of optimization for parallel acceleration in the variance calculation code, which employs a simple ’for’ loop approach, the actual overhead is larger than theoretically expected. Nonetheless, DISAM’s advantage lies in its overhead being unrelated to gradient size, but only to batch size and domain number. This drawback can be addressed through parallel code optimization, and no additional memory overhead is present. + +C.5Details about Convergence Curves of SAM and ERM + +In this section, we provide a detailed analysis of the convergence curves depicted in Figure 1(1(b)). Figure 7(7(a)) presents the same as Figure 1(1(b)), with a normalized representation of the loss curves, ranging from 0 to 1, achieved by subtracting the minimum loss value and dividing by the maximum loss value. Our intention is to emphasize the inconsistency in convergence trends across different SAM domains, as illustrated by the optimization overshoot observed in Figure 7(7(a)). Figure 7(7(b)) showcases the actual loss change curve. It is apparent that due to the consistency issue encountered during the early phase of convergence, the in-domain convergence is compromised, resulting in poor generalization performance in the out-of-domain scenario. + +(a)Convergence curves under domain shifts +(b)Loss curves under domain shifts +Figure 7:Illustration of SAM’s degradation of the training process under domain shifts. (a) Convergence curves of SAM and ERM for each domain during training, with the convergence degree normalized to [0,1]. (b) Loss curves of SAM and ERM for each domain during training. +C.6Detailed Analysis about Open-Class Generalization + +In the experiments of open-class generalization, as presented in Table 3 and Figure 4 of section 4.4, we specifically explore the effectiveness of DISAM for parameter-efficient fine-tuning (PEFT). Our quantitative analysis compares the performance of ERM, SAM, and DISAM in fine-tuning scenarios. As shown in Table 3, although CoOp and CLIPOOD perform better on base classes than zero-shot, their performance on new classes is worse than zero-shot. This suggests that the fine-tuned parameters overfit to the existing training data distribution from both the domain and class perspectives. This overfitting is particularly detrimental to the generalization of large VLM models, which often have feature representations too rich for the downstream task, especially when only a small number of parameters are fine-tuned. Figure 4 visualizes the change in performance trends during the training process, and we observe a trend where ERM initially performs well on base classes but then exhibits a decline on new classes, suggesting a collapse of the feature space onto the training data classes. Although SAM offers some relief from overfitting, its performance on new classes does not match zero-shot levels. In contrast, DISAM, by minimizing sharpness more effectively, shows improved performance on new classes, especially in domain shift scenarios. + +C.7Detailed Analysis of Convergence Speed Comparison + +We presented a comparison of the convergence speed with the inclusion of ERM in Figure 8. It can be observed that although DISAM converges much faster compared to SAM, the overall convergence speed is still slower than ERM due to the introduction of +𝜌 +. + +(a) +(b) +Figure 8:Convergence curves for ERM, SAM and DISAM. (a) & (b) show the trend of +ℒ +⁢ +( +𝑤 +) + during the training process on PACS dataset. +Appendix DPseudo code of DISAM + +We present pseudo-code for DISAM using Python syntax. PyTorch is utilized as the deep learning experimental framework. The code for the optimizer in the SAM-based method can be referenced from the provided open source links in the relevant papers. + +Listing 1: Training Code for DISAM +def train_epoch_disam(dataloader, model, optimizer): +""" +Train the DISAM model for one epoch. +Args: +dataloader (DataLoader): The training dataloader. +model (nn.Module): The training model. +optimizer (Optimizer): The SAM-based optimizer, such as SAM, GSAM, and SAGM. +""" +model.train() +for i, data_list in tqdm(enumerate(dataloader)): +imgs, labels = data_list +imgs, labels = imgs.cuda(), labels.cuda() +preds = model(imgs) +# Calculate domain losses and total loss using the cross-entropy loss function +domain_loss_list, total_loss = get_domain_loss(preds, labels, domain_labels, loss_func) +loss_variance = compute_variance(domain_loss_list) +loss = total_loss - lamda * loss_variance +optimizer.zero_grad() +loss.backward() +# Perform the first step of SAM: gradient ascent with a fixed length rho +optimizer.first_step(zero_grad=True) +output = model(imgs) +loss = loss_func(output, labels) +loss.backward() +# Obtain the actual gradient from the perturbation location of DISAM +optimizer.second_step(zero_grad=True) +def get_domain_loss(preds, labels, domain_labels, loss_func): +""" +The function to compute the loss for each domain. +Args: +preds (Tensor): The predictions of the training model in one batch. +labels (Tensor): The labels of batch data. +domain_labels (Tensor): The domain labels of batch data. +loss_func: (Function): The loss function. +""" +# Get a list of all domains +domain_list = list(set(domain_labels)) +domain_loss_list = [] +total_loss = 0. +for domain_name in domain_list: +# Get the mask for the current domain +domain_mask = domain_labels == domain_name +labels_per_domain = labels[domain_mask] +preds_pre_domain = preds[domain_mask] +# Compute the loss for the current domain +single_domain_loss = loss_func(preds_pre_domain, labels_per_domain) +domain_loss_list.append(single_domain_loss) +# Add the loss for the current domain to the total loss, taking into account the number of samples in the domain +total_loss += len(labels_per_domain) * single_domain_loss +total_loss /= len(labels) +return domain_loss_list, total_loss +def compute_variance(domain_loss_list): +""" +The function to compute the variance of the list of domain losses +Args: +domain_loss_list (List): the list of each domain’s loss. +""" +loss_variance = 0. +for domain_i_loss in domain_loss_list: +for domain_j_loss in domain_loss_list: +# Compute the square of the difference in loss between each pair of elements and add it to the loss variance +loss_variance += (domain_i_loss - domain_j_loss)**2 +loss_variance /= (2*len(domain_loss_list)**2) +return loss_variance +Report Issue +Report Issue for Selection +Generated by L A T E xml +Instructions for reporting errors + +We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below: + +Click the "Report Issue" button. +Open a report feedback form via keyboard, use "Ctrl + ?". +Make a text selection and click the "Report Issue for Selection" button near your cursor. +You can use Alt+Y to toggle on and Alt+Shift+Y to toggle off accessible reporting links at each section. + +Our team has already identified the following issues. We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all. + +Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion, and welcome developer contributions. diff --git a/docs/fmlora_flat_minima_lora.md b/docs/fmlora_flat_minima_lora.md new file mode 100644 index 0000000..8126166 --- /dev/null +++ b/docs/fmlora_flat_minima_lora.md @@ -0,0 +1,598 @@ +Title: Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond + +URL Source: https://arxiv.org/html/2508.00522 + +Published Time: Tue, 16 Dec 2025 02:15:18 GMT + +Markdown Content: +Jiaxin Deng 1\equalcontrib, Qingcheng Zhu 2\equalcontrib, Junbiao Pang 1, Linlin Yang 3, Zhongqian Fu 4, Baochang Zhang 5 + +###### Abstract + +Little research explores the correlation between the expressive ability and generalization ability of the low-rank adaptation (LoRA). Sharpness-Aware Minimization (SAM) improves model generalization for both Convolutional Neural Networks (CNNs) and Transformers by encouraging convergence to locally flat minima. However, the connection between sharpness and generalization has not been fully explored for LoRA due to the lack of tools to either empirically seek flat minima or develop theoretical methods. In this work, we propose Flat Minima LoRA (FMLoRA) and its efficient version i.e., EFMLoRA, to seek flat minima for LoRA. Concretely, we theoretically demonstrate that perturbations in the full parameter space can be transferred to the low-rank subspace. This approach eliminates the potential interference introduced by perturbations across multiple matrices in the low-rank subspace. Our extensive experiments on large language models and vision-language models demonstrate that EFMLoRA achieves optimize efficiency comparable to that of LoRA while simultaneously attaining comparable or even better performance. For example, on the GLUE dataset with RoBERTa-large, EFMLoRA outperforms LoRA and full fine-tuning by 1.0% and 0.5% on average, respectively. On vision-language models e.g., Qwen-VL-Chat, there are performance improvements of 1.5% and 1.0% on the SQA and VizWiz datasets, respectively. These empirical results also verify that the generalization of LoRA is closely related to sharpness, which is omitted by previous methods. + +## Introduction + +Parameter-Efficient Fine-Tuning (PEFT) methods only update a small subset of parameters, e.g., adapters (hu2022lora) or prompt weights (li2021prefix) for Large language models (LLMs) with substantially lower memory and computational costs. Specifically, Low-Rank Adaptation (LoRA) (hu2022lora) stands out for achieving performance comparable to full fine-tuning (FT) while being considerably more efficient. + +![Image 1: Refer to caption](https://arxiv.org/html/2508.00522v3/x1.png) + +Figure 1: Comparison of Methods: LoRA, FMLoRA, and EFMLoRA. + +Many works have been proposed to enhance the performance of LoRA by introducing more dedicated budgets for rank allocation (zhang2023adaptive), decomposing optimization for direction and magnitude updates (liu2024dora), or designing better initialization strategies for LoRA parameters (meng2024pissa), etc. These studies demonstrate the significant potential to improve LoRA performance. However, most existing approaches fail to effectively address bias inheritance, where LLMs may propagate and amplify their inherent biases, significantly impacting model performance and robustness on downstream tasks (li2025understanding). Therefore, a natural question is: how to model and understand the generalization of LoRA for various LLMs and beyond, e.g., vision-language models? + +It is widely believed that a flatter loss landscape can lead to better generalization performance(hochreiter1994simplifying)(hochreiter1997flat). For instance, Foret et al. proposed Sharpness-Aware Minimization (SAM)(foret-2020-SAM-ICLR), which seeks parameter regions where the training loss remains uniformly flat. SAM and its variants have demonstrated State-Of-The-Art (SOTA) performances across various applications, such as classification(kwon-2021-asam-ICML), transfer learning(zhuang-2022-GSAM-ICLR), domain generalization(dong2024implicit) and federated learning(FedGAMMA). + +To the best of our knowledge, compared to theoretical analysis, e.g.,(neyshabur2017exploring), empirically connecting sharpness and generalization ability of LoRA is a practical approach, e.g.,(andriushchenko2023modern). For the second line of research, a naive approach is to combine SAM with LoRA. However, if perturbations in SAM are applied simultaneously to two low-rank subspaces of LoRA, they may change the maximum loss within the neighborhood of LoRA’s full parameter space(dinh2017sharp); besides, SAM incurs a computational cost twice that of Stochastic Gradient Descent (SGD)(deng2024effective). The key question in the second line of research is how to efficiently find flat minima in LoRA, aiming to better understand the connection between sharpness and generalization. + +In this paper, we propose a novel PEFT method, FMLoRA, that promotes convergence toward flatter minima. Specifically, we theoretically uncover that perturbations in the full parameter space can be equivalently re-parameterized as perturbations within the low-rank space. In addition, we propose EFMLoRA to accelerate FMLoRA by an Exponential Moving Average (EMA) strategy. We validate that EFMLoRA improves generalization performance on downstream tasks while maintaining computational efficiency comparable to that of LoRA. Fig.[1](https://arxiv.org/html/2508.00522v3#Sx1.F1 "Figure 1 ‣ Introduction ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") compares three methods: LoRA, FMLoRA, and EFMLoRA. We conducted comprehensive experiments on diverse tasks (fine-tuning, few-shot learning) and various model types (RoBERTa (liu2019roberta), GPT-2 (radford2019language), CLIP (zanella2024low), Qwen-VL-Chat (Bai2023QwenVLAV)) and scales. We find that EFMLoRA achieves model accuracy very close to, or even surpass both full fine-tuning and LoRA across many tasks. Our main contribution can be summarized as follows: + +* •We propose FMLoRA, a novel PEFT training method that integrates SAM into the LoRA framework. Furthermore, EFMLoRA provides an efficient tool for empirically understanding the connection between sharpness and generalization in LLMs and beyond. We empirically show that reducing sharpness is highly correlated with improved generalization in PEFT tasks, which has been rarely explored in PEFT studies before. +* •We conduct comprehensive experiments on LLMs (e.g., RoBERTa, GPT-2) and vision-language models (e.g., CLIP, Qwen-VL-Chat) across various tasks including fine-tuning and few-shot learning. Results show that EFMLoRA achieves optimize efficiency comparable to that of LoRA while simultaneously attaining comparable or even better performance. + +## Related Works + +### Low-rank Adaption + +Hu et al. proposed LoRA (hu2022lora) as a PEFT method that introduced low-rank adapters into each layer of a pre-trained model. Recent advancements in LoRA can be broadly categorized into two directions: 1) advanced architectures and 2) optimization methods. In the first research line, for example, LoraHub (huang2023lorahub) trained multiple adapters and strategically combined them based on the domain during inference. LoRA-FA (zhang2023lora) chose to freeze the projection-down weight of \mathbf{A} and update the projection-up weight of \mathbf{B} in each LoRA layer. DoRA (liu2024dora) improved LoRA by incorporating a learnable magnitude vector to re-scale the normalized product of low-rank matrices. HydraLoRA (tian2024hydralora) extended the LoRA framework with an asymmetric architecture that shared a common \mathbf{A} matrix for efficiency while dynamically assigning samples to multiple \mathbf{B} matrices via a MoE mechanism. In the second line, for example, LoRA+ (hayou2024lora+) applied different learning rates to the two low-rank matrices. Additionally, Galore (zhao2024galore) employed SVD to compress the gradients and its first and second momentum of full training into a low-rank space, thereby reducing the memory footprint during pre-training and fine-tuning. Recently, Li et al. (li2024flat) proposed combining SAM with LoRA for better generalization, but they used random perturbation. Our method belongs to the second research line. Different from (li2024flat), our method transfers the perturbation from the full parameter space to a single low-rank parameter space without changing the maximum perturbed loss, avoiding misalignment with SAM’s training behavior. + +### Sharpness and Generalization Ability + +Research on the relationship between sharpness and generalization could be traced back to (hochreiter1997flat). Following the observation by (keskar-2016-large_batch-ICLR) that larger batch sizes tended to increase sharpness and generalization error. (jastrzkebski2017three) extended this by finding a correlation between the sharpness and the ratio of learning rate to batch size. (dinh-2017-sharp_minima-ICML) showed that one can easily construct networks with good generalization but with arbitrary large sharpness by reparameterization. (jiang-2019-fantastic-ICLR) performed a large-scale empirical study on various generalization measures and showed that sharpness-based measures have the highest correlation with generalization. Theoretical understandings on the generalization error using sharpness-related measures were provided in (neyshabur2017exploring), (wanggeneralization). Collectively, these studies justified the goal of seeking flatter minima to improve generalization. However, to the best of our knowledge, the correlation between sharpness and generalization for LoRA has barely been discussed due to the lack of theoretical understanding or efficient tools for empirical analysis. Our method provides an efficient tool for empirical analysis in this domain. + +### Recap of SAM + +Foret et al.(foret-2020-SAM-ICLR) proposed the SAM to enhance model generalization as follows: + +\displaystyle\mathop{\min}\limits_{\mathbf{w}}[(\mathop{\max}\limits_{||\bm{\varepsilon}||\leq\rho}L(\mathbf{w}+\bm{\varepsilon})-L(\mathbf{w}))+L(\mathbf{w})+\lambda||\mathbf{w}||_{2}^{2}],(1) + +where \mathbf{w} represents the weights of the network, \bm{\varepsilon} represents the perturbation of weights \mathbf{w} in a Euclidean ball with the radius \rho(\rho>0), L(\cdot) is the loss function, and \lambda||\mathbf{w}||_{2}^{2} is a standard L2 regularization term. + +SAM utilizes Taylor expansion to search for the maximum perturbed loss (\mathop{\max}\limits_{||\bm{\varepsilon}||\leq\rho}L(\mathbf{w}+\bm{\varepsilon})) in local parameter space as follows: + +\displaystyle\mathop{\arg\max}\limits_{||\bm{\varepsilon}||\leq\rho}\;L(\mathbf{w}+\bm{\varepsilon})\approx\mathop{\arg\max}\limits_{||\bm{\varepsilon}||\leq\rho}\;{\bm{\varepsilon}^{\top}}{\nabla_{\mathbf{w}}}L(\mathbf{w}).(2) + +By solving Eq.([2](https://arxiv.org/html/2508.00522v3#Sx2.E2 "Equation 2 ‣ Recap of SAM ‣ Related Works ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), SAM obtains the perturbation as follows: + +\displaystyle\hat{\bm{\varepsilon}}=\rho{\nabla_{\mathbf{w}}}L(\mathbf{w})/||{\nabla_{\mathbf{w}}}L(\mathbf{w})||.(3) + +Substituting the perturbation \hat{\bm{\varepsilon}} back into Eq.([1](https://arxiv.org/html/2508.00522v3#Sx2.E1 "Equation 1 ‣ Recap of SAM ‣ Related Works ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), we then have: + +\displaystyle{\nabla_{\mathbf{w}}}\mathop{\max}\limits_{||\bm{\varepsilon}||\leq\rho}L(\mathbf{w}+\bm{\varepsilon})\approx{\nabla_{\mathbf{w}}}L({\mathbf{w}}+\hat{\bm{\varepsilon}}({\mathbf{w}}))(4) +\displaystyle={\nabla_{\mathbf{w}}}L({\mathbf{w}}){|_{{\mathbf{w}}+\hat{\bm{\varepsilon}}({\mathbf{w}})}}+\frac{{d\hat{\bm{\varepsilon}}({\mathbf{w}})}}{{d{\mathbf{w}}}}{\nabla_{\mathbf{w}}}L({\mathbf{w}}){|_{{\mathbf{w}}+\hat{\bm{\varepsilon}}({\mathbf{w}})}}. + +By dropping the second-order terms in Eq.([4](https://arxiv.org/html/2508.00522v3#Sx2.E4 "Equation 4 ‣ Recap of SAM ‣ Related Works ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), SAM calculates the gradient at \mathbf{w}+\bm{\hat{\varepsilon}} as follows: + +\displaystyle{\nabla_{\mathbf{w}}}\mathop{\max}\limits_{||\bm{\varepsilon}||\leq\rho}L(\mathbf{w}+\bm{\varepsilon})\approx{\nabla_{\mathbf{w}}}L(\mathbf{w}){|_{\mathbf{w}+\bm{\hat{\varepsilon}}}}.(5) + +Finally, SAM uses the gradients from Eq.([5](https://arxiv.org/html/2508.00522v3#Sx2.E5 "Equation 5 ‣ Recap of SAM ‣ Related Works ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) for optimization. + +### SAM Variants + +Recently, SAM variants could be broadly categorized into three groups: 1) studies on the perturbation radius \rho in SAM, 2) studies that speed up the optimization process of SAM, and 3) redefinitions of sharpness in SAM. For the first direction, Kwon et al.(kwon-2021-asam-ICML) proposed Adaptive SAM (ASAM), which adapted the perturbation radius in a scale-aware manner, allowing SAM to be effectively applied to scale-invariant neural networks. For the second group, Kim et al.(kim2023exploring) introduced a multi-step ascent approach to improve SAM. Li et al.(li2024friendly) introduced Friendly SAM (F-SAM), which improved generalization by removing the detrimental influence of the full gradient component and instead utilizing batch-specific gradients to guide optimization more effectively. For the third group, Zhuang et al.(zhuang-2022-GSAM-ICLR) pointed out that SAM did not always favor flat minima. Consequently, they proposed GSAM, which minimized the surrogate gap and the perturbed loss to better encourage flatness. Zhang et al. introduced the first-order flatness(zhang-2023-gradient-CVPR), which assessed the maximal gradient norm within a perturbation radius. Consequently, they proposed GAM which explicitly seeks minima characterized by uniformly small curvature. + +## Method + +### SAM on LoRA + +LoRA achieves parameter efficiency by modeling the low-rank decomposed weight(li2022low). Specifically, the weight change for each layer \mathbf{W}_{0}\in\mathbb{R}^{n\times m} is represented as \Delta\mathbf{W}=s\mathbf{B}\mathbf{A}, where s is a scaling factor, \mathbf{B}\in\mathbb{R}^{n\times r}, \mathbf{A}\in\mathbb{R}^{r\times m}, with rank r\ll\min(n,m). Given an input \mathbf{x}, the forward is as follows: + +\displaystyle\mathbf{y}=\mathbf{W}_{0}\mathbf{x}+\Delta\mathbf{W}\mathbf{x}=(\mathbf{W}_{0}+s\mathbf{B}\mathbf{A})\mathbf{x},(6) + +where matrix \mathbf{A} is typically initialized by the Kaiming’s method(he2015delving), \mathbf{B} is set to zeros. \mathbf{W}_{0} remains unchanged during fine-tuning, while \mathbf{B} and \mathbf{A} are trained. During inference, \Delta\mathbf{W} is merged into \mathbf{W_{0}}. + +If SAM is naively combined with LoRA, the optimization loss can be rewritten as follows: + +\displaystyle\min_{\mathbf{A},\mathbf{B}}~~\mathop{\max}\limits_{\scriptstyle||{{\bf{E}}^{\bf{A}}}|{|_{F}}\leq\rho,\hfill\atop\scriptstyle||{{\bf{E}}^{\bf{B}}}|{|_{F}}\leq\rho\hfill}L({\mathbf{W_{0}}}+{s}(\mathbf{B}+{\mathbf{E}^{\mathbf{B}}})(\mathbf{A}+{\mathbf{E}^{\mathbf{A}}})),(7) + +where \mathbf{E}^{\mathbf{B}}\in\mathbb{R}^{n\times r} and \mathbf{E}^{\mathbf{A}}\in\mathbb{R}^{r\times m} represent the perturbations applied to the parameters \mathbf{B} and \mathbf{A}, respectively, and \rho is the radius of perturbations. There are two key challenges: + +* •Two separate perturbations in two low-rank subspaces interfere with each other, leading to an inconsistency between the maximum loss obtained when perturbing in the low-rank subspaces and the maximum loss obtained when perturbing in the full parameter space. +* •SAM requires computing gradients twice per iteration, resulting in approximately twice the computational cost compared to LoRA. + +### FMLoRA + +To deal with the first challenge, we propose to re-parameterize the perturbation from the full parameter space to a single low-rank parameter space. Concretely, the loss in the full parameter space can be formulated as follows: + +\displaystyle\min_{\mathbf{A},\mathbf{B}}~~\max_{\|\mathbf{E}^{\mathbf{W}}\|_{F}\leq\rho}~~L(\mathbf{W_{0}}+s\mathbf{B}\mathbf{A}+\mathbf{E}^{\mathbf{W}}).(8) + +To solve the minimax problem in Eq.([8](https://arxiv.org/html/2508.00522v3#Sx3.E8 "Equation 8 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), it is necessary to first find optimal \hat{\mathbf{E}}^{\mathbf{W}}\in\mathbb{R}^{n\times m}. Analogous to SAM, we approximate the optimal perturbation \hat{\mathbf{E}}^{\mathbf{W}} to maximize L(\mathbf{W}+\mathbf{E}^{\mathbf{W}}) where \mathbf{W}=\mathbf{W_{0}}+s\mathbf{B}\mathbf{A} as follows: + +\displaystyle\hat{\bm{\varepsilon}}^{\mathbf{w}}=\rho\text{sign}(\mathbf{g}^{\mathbf{w}})\frac{\mathbf{g}^{\mathbf{w}}}{||\mathbf{g}^{\mathbf{w}}||},(9) + +where \mathbf{g}^{\mathbf{w}}=\text{Vector}(\nabla L_{\mathbf{W}}(\mathbf{W})) and \hat{\bm{\varepsilon}}^{\mathbf{w}}=\text{Vector}(\hat{\mathbf{E}}^{\mathbf{W}}), in which the \text{Vector}(\cdot) function represents a vectorized operation. However, the solution for \hat{\mathbf{E}}^{\mathbf{W}} explicitly depends on the gradient of the matrix \mathbf{W}. That is, the form of solution in Eq.([9](https://arxiv.org/html/2508.00522v3#Sx3.E9 "Equation 9 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) is undesirable since \nabla L_{\mathbf{W}}(\mathbf{W}) is unknown during LoRA optimization. + +In this paper, we propose to approximate the unknown gradient \nabla L_{\mathbf{W}}(\mathbf{W}) using standard LoRA gradients, which can be computed in two ways: + +\displaystyle(1)\quad\nabla L_{\mathbf{W}}(\mathbf{W})=\frac{1}{s}\nabla L_{\mathbf{B}}(\mathbf{W_{0}}+s\mathbf{BA})(\mathbf{A}^{\top})^{+},(10) +\displaystyle(2)\quad\nabla L_{\mathbf{W}}(\mathbf{W})=\frac{1}{s}(\mathbf{B}^{\top})^{+}\nabla L_{\mathbf{A}}(\mathbf{W_{0}}+s\mathbf{BA}),(11) + +where (\mathbf{A}^{\top})^{+} and (\mathbf{B}^{\top})^{+} represent the pseudo-inverse of \mathbf{A}^{\top} and \mathbf{B}^{\top}, respectively. The accuracy of the pseudo-inverse depends on the condition number of matrix. A smaller condition number leads to a more accurate pseudo-inverse. Matrices with lower condition numbers are better suited for stable representation. In LoRA, we found that the condition number is typically low, around 3. + +To obtain a more accurate estimate of the gradient of the full weights, we combine the above two approaches to compute \nabla L_{\mathbf{W}}(\mathbf{W}) as follows: + +\displaystyle\overline{\nabla{L}_{\mathbf{W}}(\mathbf{W})}\displaystyle=0.5*(\frac{1}{s}\nabla L_{\mathbf{B}}(\mathbf{W_{0}}+s\mathbf{BA})(\mathbf{A}^{\top})^{+} +\displaystyle+\frac{1}{s}(\mathbf{B}^{\top})^{+}\nabla L_{\mathbf{A}}(\mathbf{W_{0}}+s\mathbf{BA})).(12) + +Let {\bar{\mathbf{g}}^{\mathbf{W}}}=\text{Vector}(\overline{\nabla{L}_{\mathbf{W}}(\mathbf{W})}). Then the perturbation in Eq.([9](https://arxiv.org/html/2508.00522v3#Sx3.E9 "Equation 9 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) could be rewritten as follows: + +\displaystyle\bar{\mathbf{E}}^{\mathbf{W}}=\text{Matrix}(\rho\text{sign}({\bar{\mathbf{g}}^{\mathbf{W}}})\frac{{\bar{\mathbf{g}}^{\mathbf{W}}}}{||{\bar{\mathbf{g}}^{\mathbf{W}}}||}),(13) + +where \text{Matrix}(\cdot) denotes the operation that converts a vector into a matrix. We transfer the perturbation from the full parameter space to a single low-rank parameter space without changing the maximum loss in the local region of the parameters. We apply no perturbation to matrix \mathbf{A}, i.e., {\mathbf{E}^{\mathbf{A}}}=\mathbf{0}, and ensure that the loss under perturbations in the low-rank subspace in Eq.([7](https://arxiv.org/html/2508.00522v3#Sx3.E7 "Equation 7 ‣ SAM on LoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) matches the inner maximum loss in Eq.([8](https://arxiv.org/html/2508.00522v3#Sx3.E8 "Equation 8 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), as follows: + +\displaystyle L({\mathbf{W_{0}}}\displaystyle+{s}(\mathbf{B}+{\mathbf{E}^{\mathbf{B}}})\mathbf{A})(14) +\displaystyle=\max_{\|\mathbf{E}^{\mathbf{W}}\|_{F}\leq\rho}L(\mathbf{W_{0}}+s\mathbf{B}\mathbf{A}+\mathbf{E}^{\mathbf{W}}). + +Substituting \bar{\mathbf{E}}^{\mathbf{W}} into Eq.([14](https://arxiv.org/html/2508.00522v3#Sx3.E14 "Equation 14 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), we obtain: + +\displaystyle{\mathbf{E}}^{\mathbf{B}}\approx\frac{1}{s}\bar{\mathbf{E}}^{\mathbf{W}}\mathbf{A}^{+},(15) + +where \mathbf{A}^{+} is the pseudo-inverse of \mathbf{A}. An alternative approach is to transfer the perturbation to matrix \mathbf{A}. Following the observations from HydraLoRA(tian2024hydralora), matrix \mathbf{A} shows high parameter similarity across heads, likely due to initialization, making it capture domain-common features, while matrix \mathbf{B} remains distinct and domain-specific. Since different tasks require different perturbations, we adopt the approach of transferring the perturbation to the matrix \mathbf{B}, as expressed in Eq.([14](https://arxiv.org/html/2508.00522v3#Sx3.E14 "Equation 14 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")). The detailed derivation of Eq.([10](https://arxiv.org/html/2508.00522v3#Sx3.E10 "Equation 10 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) and the pseudo-algorithm for FMLoRA are provided in the supplementary file. + +#### Balancedness of FMLoRA. + +Balancedness is well-appreciated in domains such as matrix factorization/sensing (ge2017no)(du2018algorithmic). It is also observed that balanced neural networks are easier to optimize relative to unbalanced ones (neyshabur2015path). Recently, Balancedness B_{t}:=\frac{1}{2}(||\mathbf{x}_{t}||^{2}-||\mathbf{y}_{t}||^{2}) (where \mathbf{x}_{t} and \mathbf{y}_{t} are variables) turns out to be an intriguing alternative to sharpness on the scale-invariant problem (li2024implicit). + +To investigate the balancedness of our proposed method, we express the update process of FMLoRA analogously to Eq.(4) in (li2024implicit) as follows: + +\displaystyle{\tilde{\mathbf{x}}_{t}}={{\mathbf{x}}_{t}}+\rho\frac{1}{s}\frac{{{\mathbf{G}_{t}}}}{{\left\|{{\mathbf{G}_{t}}}\right\|}}{\mathbf{y}_{t}}^{+}\displaystyle,\quad{\tilde{\mathbf{y}}_{t}}={{\mathbf{y}}_{t}},(16) +\displaystyle{\mathbf{g}_{{\tilde{\mathbf{x}}_{t}}}}={{\tilde{\mathbf{G}}}_{t}}\tilde{\mathbf{y}}_{t}\displaystyle,\quad{\mathbf{g}_{{\tilde{\mathbf{y}}_{t}}}}={{\tilde{\mathbf{G}}}_{t}}^{\top}\tilde{\mathbf{x}}_{t}, +\displaystyle{{\mathbf{x}}_{t+1}}={{\mathbf{x}}_{t}}-\eta{\mathbf{g}_{{\tilde{\mathbf{x}}_{t}}}}\displaystyle,\quad{{\mathbf{y}}_{t+1}}={{\mathbf{y}}_{t}}-\eta{\mathbf{g}_{{\tilde{\mathbf{y}}_{t}}}}, + +where {\mathbf{x}}_{t}=\text{Vector}(\mathbf{B}_{t}), {\mathbf{y}}_{t}=\text{Vector}(\mathbf{A}_{t}), {\mathbf{G}_{t}}=\nabla L({\mathbf{x}}_{t}{\mathbf{y}}_{t}^{\top}) is the gradient of the full parameter space at the original parameter point, {\tilde{\mathbf{G}}_{t}}=\nabla L(\tilde{\mathbf{x}}_{t}\tilde{\mathbf{y}}_{t}^{\top}) is the gradient of the full parameter space at the perturbed parameter point, and \mathbf{y}_{t}^{+} is the pseudo inverse of \mathbf{y}_{t}. + +###### Theorem 1. + +Let B_{t}:=\frac{1}{2}(||\mathbf{x}_{t}||^{2}-||\mathbf{y}_{t}||^{2}). For the learning rate \eta\Rightarrow 0, the limiting flow of FMLoRA guarantees that: + +\displaystyle\left|{\frac{1}{2}\frac{{d({{\left\|{{\mathbf{x}_{t}}}\right\|}^{2}}-{{\left\|{{\mathbf{y}_{t}}}\right\|}^{2}})}}{{dt}}}\right|\leq\left|{\rho\frac{1}{s}\frac{1}{{\left\|{\mathbf{y}_{t}}\right\|}}\left\|{{\mathbf{g}_{{{{\rm{\tilde{\mathbf{x}}}}}_{t}}}}}\right\|}\right|.(17) + +Theorem 1 indicates that the balancedness of FMLoRA is influenced by the perturbation range \rho, the norm of the gradient at the perturbed point, the \ell_{2}-norm of \mathbf{y}_{t}, and the scale constraint of LoRA. To ensure that the balancedness of FMLoRA gradually decreases during training, we reduce \rho progressively. In addition, the norm of the gradient with respect to \mathbf{y}_{t} at the perturbed point also decreases due to the weight decay. The \ell_{2}-norm of \mathbf{y}_{t} is bounded within a certain range, these factors collectively contribute to the reduction in the balancedness of FMLoRA. + +![Image 2: Refer to caption](https://arxiv.org/html/2508.00522v3/x2.png) + +Figure 2: Parameter update process for EFMLoRA. + +### Efficient FMLoRA + +The optimization processes of FMLoRA also require two gradient computations per iteration. To enhance optimization efficiency, we propose Efficient FMLoRA (EFMLoRA), which estimates the subsequent perturbation {\mathbf{E}}^{\mathbf{B}} in Eq.([15](https://arxiv.org/html/2508.00522v3#Sx3.E15 "Equation 15 ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) by maintaining an Exponential Moving Average (EMA) of previous perturbations as follows: + +\displaystyle{{\hat{\mathbf{E}}}^{\mathbf{B}}_{t}}=(1-\beta){{\hat{\mathbf{E}}}^{\mathbf{B}}_{t-1}}+\beta{\mathbf{E}^{\mathbf{B}}_{t}},(18) + +where \beta\in(0,1) is the momentum coefficient that determines the update rate of the exponential moving average. {\mathbf{E}^{\mathbf{B}}_{t}} is the perturbation on matrix \mathbf{B}_{t} at t-th iteration, {{\hat{\mathbf{E}}}^{\mathbf{B}}_{t}} is the EMA perturbation at t-th iteration. Fig.[2](https://arxiv.org/html/2508.00522v3#Sx3.F2 "Figure 2 ‣ Balancedness of FMLoRA. ‣ FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") illustrates the parameter update process of EFMLoRA: (1) Calculate the gradient at the perturbed point (\mathbf{W}_{0}, \mathbf{B}_{t-1}+\hat{\mathbf{E}}^{\mathbf{B}}_{t-1}, \mathbf{A}_{t-1}). (2) Calculate the perturbation {\mathbf{E}}^{\mathbf{B}}_{t}=\frac{1}{s}\bar{\mathbf{E}}^{\mathbf{W}}\mathbf{A}^{+}_{t-1}. (3) Return to the original parameter point (\mathbf{W}_{0},\mathbf{B}_{t-1},\mathbf{A}_{t-1}). (4) Update the parameters to (\mathbf{W}_{0},\mathbf{B}_{t},\mathbf{A}_{t}). (5) Calculate the EMA perturbation by Eq.([18](https://arxiv.org/html/2508.00522v3#Sx3.E18 "Equation 18 ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")) and update the parameters to the next perturbed point (\mathbf{W}_{0},\mathbf{B}_{t}+\hat{\mathbf{E}}^{\mathbf{B}}_{t},\mathbf{A}_{t}). During this optimization process, each optimization step requires only a single forward and backward. The algorithmic pseudocode is provided in the supplementary file. + +Table 1: Experiments on few-shot RoBERTa (355M). Results marked with * are taken from (li2024implicit). + +To theoretically analyze the error of EFMLoRA, some necessary assumptions are listed below, all of which are common and standard when analyzing SAM optimization(du-2022-ESAM-ICLR)(zhuang-2022-GSAM-ICLR). + +###### Assumption 1. + +(Smooth) L(\mathbf{w}) is \tau-Lipschitz smooth in \mathbf{w}, i.e., \left\|{\nabla L(\mathbf{w})-\nabla L(\mathbf{v})}\right\|\leq\tau\left\|{\mathbf{w}-\mathbf{v}}\right\|. + +###### Assumption 2. + +(Bounded gradients). By the assumption that an upper bound exists on the gradient of each mini-batch. There exists G>0 for each mini-batch such that \mathbb{E}\left[{\left\|{\nabla L(\mathbf{w})}\right\|}\right]\leq G. + +###### Assumption 3. + +(Bounded variance of stochastic gradients). Given the training set \mathbf{D} and a mini-batch \mathbf{B}\in\mathbf{D}. There exists \sigma\geq 0, the variance of stochastic gradient L_{\mathbf{B}}(\mathbf{w}) is bounded by \mathbb{E}\left[{{{\left\|{\nabla{L_{\mathbf{B}}}(\mathbf{w})-\nabla{L_{\mathbf{D}}}(\mathbf{w})}\right\|}^{2}}}\right]\leq\sigma^{2}. + +###### Assumption 4. + +(Convex) We assume that the loss function f:\mathbb{R}^{n}\rightarrow\mathbb{R} is convex and twice differentiable over an open domain. That is, for all x,y\in\text{dom}(f), it satisfies: f(y)\geq f(x)+\nabla f(x)^{\top}(y-x). + +This convexity assumption is reasonable in the fine-tuning stage, as the model is typically close to a local minimum and the loss landscape is approximately convex in a local neighborhood (jang2024lora). + +###### Theorem 2. + +[EMA perturbation approximate perturbation of SAM due to the convex of the loss landscape] Assume that during fine-tuning, the solution is already close to a local minimum and the local loss function is convex. Let the model weights at i-th iteration be \mathbf{w}_{t}. Under Assumptions [1](https://arxiv.org/html/2508.00522v3#Thmassumption1 "Assumption 1. ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"), [2](https://arxiv.org/html/2508.00522v3#Thmassumption2 "Assumption 2. ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"), and [3](https://arxiv.org/html/2508.00522v3#Thmassumption3 "Assumption 3. ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"), let {\rho_{t}}=\frac{{{\rho_{0}}}}{{\sqrt{t}}}, the error between the sharpness calculated using the EMA perturbation (S^{\text{EMA}}) and that calculated using the original SAM perturbation (S^{\text{SAM}}) is bounded as follows: + +\displaystyle|\underbrace{\left[{L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right]}_{S^{\text{EMA}}}-\underbrace{\left[{L({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right]}_{S^{\text{SAM}}}|(19) +\displaystyle\leq\left({\left({1+{{(1-\beta)}^{t-1}}}\right)\tau{\rho_{0}}+G+{\sigma^{2}}}\right) +\displaystyle\quad\quad\cdot\left({\left({1+{{(1-\beta)}^{t-1}}}\right){\rho_{0}}+\frac{{{\rho_{0}}}}{{\sqrt{t}}}}\right). + +Theorem [2](https://arxiv.org/html/2508.00522v3#Thmtheorem2 "Theorem 2. ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") demonstrates that as t increases, the difference between S^{\text{EMA}} and S^{\text{SAM}} gradually decreases. The perturbation estimated by the EMA can effectively approximate the original SAM perturbation. + +Table 2: Experiments on finetuning RoBERTa (355M). Results marked with \dagger are taken from (hu2022lora), and those with * are taken from (li2024implicit). + +Table 3: GPT-2 medium (M) and large (L) with different adaptation methods on the E2E NLG Challenge. Results marked with \dagger are taken from (hu2022lora). + +### Memory and Time Complexity + +LoRA reduces the number of trainable parameters by decomposing weight updates as \Delta\mathbf{W}\approx\mathbf{B}\mathbf{A}, where \mathbf{B}\in\mathbb{R}^{n\times r} and \mathbf{A}\in\mathbb{R}^{r\times m} with r\ll\min(n,m). Both FMLoRA and EFMLoRA retain this parameter efficiency: + +\displaystyle\text{P}_{\text{LoRA}}\displaystyle=\text{P}_{\text{FMLoRA}}=\text{P}_{\text{EFMLoRA}}(20) +\displaystyle=O(nr+rm)\ll O(nm). + +However, FMLoRA and EFMLoRA introduce additional memory overhead. Specifically, FMLoRA temporarily stores the original values of \mathbf{B} and \mathbf{A}, as well as the gradients of \mathbf{A}. The memory usage of FMLoRA is calibrated as follows: + +\displaystyle\text{M}_{\text{FMLoRA}}=\text{M}_{\text{LoRA}}+O(5\times(nr+rm)),(21) + +where \text{M}_{\text{LoRA}} indicates the memory required by LoRA. The memory of EFMLoRA needs to maintain the EMA perturbation on \mathbf{B} as follows: + +\displaystyle\text{M}_{\text{EFMLoRA}}=\text{M}_{\text{LoRA}}+O(2\times(nr+rm)).(22) + +Notably, modern optimizers like AdamW already require O(2\times(nr+rm)) memory for momentum and second-moment statistics when applied to LoRA. + +For time complexity, suppose that the time complexity of optimizing the model with LoRA is O(T), which mainly includes the time for forward and backward. Theoretically, the time complexity of FMLoRA is approximately as follows: + +\displaystyle\text{T}_{\text{FMLoRA}}\approx O(2T)=2\times\text{T}_{\text{LoRA}}.(23) + +In contrast, the time complexity of EFMLoRA can be approximated as follows: + +\displaystyle\text{T}_{\text{EFMLoRA}}\approx O(T)=\text{T}_{\text{LoRA}}.(24) + +We implement QR decomposition by Householder transformations, with time complexity of O(r^{2}n) for an r\times n matrix, e.g., r is rank, n is the input dimension in LORA. + +## Experiments and Discussions + +The best and second-best results are highlighted in bold and underline, respectively. Additional experimental details are provided in the supplementary file. + +### Experiments on Large Language Models + +Few-shot with RoBERTa-large. We first consider few-shot learning with EFMLoRA. Following the setup of (li2024implicit), we adopt RoBERTa-large—a 355M-parameter language model—as the backbone. The results in Table [1](https://arxiv.org/html/2508.00522v3#Sx3.T1 "Table 1 ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") show that FMLoRA outperforms all other methods with the highest average score (83.1), particularly excelling on SST-2, SNLI, and MNLI. EFMLoRA follows closely with an average score of 82.3. It consistently surpasses baseline LoRA (+2.3), LoRA-SAM (+1.0), and both BAR variants. These results highlight its superior generalization ability under distribution shift and limited supervision. We conjecture that the performance gap between SAM and EFMLoRA comes from EFMLoRA eliminating the mutual interference between perturbations in the two low-rank subspaces. + +Fine-tuning with RoBERTa-large. We apply EFMLoRA to finetune RoBERTa-large. Our implementation follows (hu2022lora), using the same hyperparameters as those in its GitHub repository. The results can be found in Table [2](https://arxiv.org/html/2508.00522v3#Sx3.T2 "Table 2 ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). we observe that EFMLoRA achieves the highest scores on all datasets, and achieves the highest accuracy on average over these datasets. Specifically, on average over these datasets, EFMLoRA surpasses standard LoRA with a margin of 1.0. Additionally, EFMLoRA even achieve better performance than full fine-tuning on some datasets. This superior performance may be attributed to overfitting in full fine-tuning, where optimizing all model parameters can lead to overfitting on the training data, thus reducing the model’s generalization to the test set. This effect is particularly pronounced on small datasets, such as MRPC, which contains only 3.7k training data. + +Fine-tuning with GPT-2. Having shown that FMLoRA is effective for NLU tasks, we now explore whether EFMLoRA can improve LoRA in NLG models like GPT-2 Medium and Large (radford2019language). To enable a direct comparison, we adopt the experimental setup of (li2021prefix) with minimal deviation. Table[3](https://arxiv.org/html/2508.00522v3#Sx3.T3 "Table 3 ‣ Efficient FMLoRA ‣ Method ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") demonstrates the effectiveness of EFMLoRA on the E2E NLG Challenge (novikova2017e2e) with GPT-2 Medium and Large models. Compared with existing PEFT methods such as Adapter and LoRA, EFMLoRA consistently achieves superior performance across all metrics. Notably, it achieves this improvement without increasing the number of trainable parameters, maintaining the same efficiency as standard LoRA. + +### Experiments on Vision Language Models + +Few-shot with CLIP. Recent advances in few-shot adaptation of Vision-Language Models (VLMs) have significantly enhanced their generalization. CLIP-LoRA (zanella2024low) explores the application of LoRA in this few-shot VLM setting. In our work, we also apply FMLoRA and EFMLoRA to VLMs to evaluate their effectiveness. For a fair comparison, our experimental setup follows that of CLIP-LoRA. We consider five datasets for fine-grained classification of satellite imagery (EuroSAT (helber2019eurosat), Ox-fordPets (parkhi2012cats), Flower102 (nilsback2008automated), Caltech101 (fei2004learning), DTD (cimpoi2014describing)). These datasets offer a thorough benchmarking framework for evaluating few-shot visual classification tasks. Table[4](https://arxiv.org/html/2508.00522v3#Sx4.T4 "Table 4 ‣ Experiments on Vision Language Models ‣ Experiments and Discussions ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") demonstrates that FMLoRA and EFMLoRA outperformed Adapter and LoRA in most settings. In the low-data regimes (1-shot and 4-shot), EFMLoRA shows clear advantages. These results highlight the effectiveness of EFMLoRA in improving generalization in few-shot adaptation of vision-language models. + +Table 4: Detailed results for five datasets with CLIP-Adapter, CLIP-LoRA and EFMLoRA. + +Fine-tuning with Qwen-VL-Chat. Qwen-VL-Chat (Bai2023QwenVLAV) is a multimodal conversational large language model capable of understanding both images and text. We apply EFMLoRA to fine-tune Qwen-VL-Chat, following the same experimental setup as in (zhou2024empirical). Table [5](https://arxiv.org/html/2508.00522v3#Sx4.T5 "Table 5 ‣ Experiments on Vision Language Models ‣ Experiments and Discussions ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") presents the results on the ScienceQA (lu2022learn) and VizWiz (gurari2018vizwiz) datasets. The results in Table [5](https://arxiv.org/html/2508.00522v3#Sx4.T5 "Table 5 ‣ Experiments on Vision Language Models ‣ Experiments and Discussions ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") demonstrate that the perturbation size \rho significantly influences the performance of EFMLoRA when fine-tuning Qwen-VL-Chat. By tuning \rho, EFMLoRA adapts to different tasks, enabling improved generalization—achieving higher accuracy than LoRA. Specifically, a larger \rho (e.g., \rho=0.2) yields the best accuracy on ScienceQA, while a smaller \rho (e.g., \rho=0.05) performs better on VizWiz. This suggests that different tasks benefit from different levels of perturbation. Therefore, selecting an appropriate \rho based on the task characteristics is crucial for achieving optimal fine-tuning performance on multimodal large language models. + +Table 5: EFMLoRA Fine-Tuning Results on Qwen-VL-Chat with different \rho. + +Table 6: Runtime (Hour) and memory (GB) of LoRA, FMLoRA and EFMLoRA on fine-tuning GPT-2 Medium/Large. + +### Runtime and Memory Consumption + +The results in Table[6](https://arxiv.org/html/2508.00522v3#Sx4.T6 "Table 6 ‣ Experiments on Vision Language Models ‣ Experiments and Discussions ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") confirm the theoretical time complexity analysis. As expected, FMLoRA has approximately double the runtime of LoRA (2.1× on both GPT-2 Medium and Large), consistent with its theoretical complexity of O(2T) due to two forward and backward passes for sharpness optimization. In contrast, EFMLoRA operates with near-LoRA efficiency, requiring only 1.1× and 1.2× more time on GPT-2 Medium and Large, respectively. This supports the theoretical claim that EFMLoRA maintains a time complexity close to O(T) while benefiting from sharpness-aware optimization. In addition, EFMLoRA maintains a memory usage almost identical to that of LoRA, with only negligible increases (less than 0.4 GB across both model scales). These results demonstrate that EFMLoRA achieves near-LoRA efficiency in both memory and runtime. + +### Conclusion + +In this work, we propose FMLoRA, a novel PEFT method that integrates sharpness-aware optimization into the LoRA framework to promote convergence toward flatter minima. We theoretically demonstrate that perturbations in the full parameter space can be equivalently represented within the low-rank subspace. To improve computational efficiency, we introduce EFMLoRA, which leverages an exponential moving average to approximate perturbations, significantly reducing runtime overhead while maintaining effectiveness. Extensive experiments across various large language and vision-language models demonstrate that EFMLoRA achieves comparable or even superior generalization performance to full fine-tuning and LoRA. Our results emphasize the importance of reducing sharpness to improve generalization in PEFT methods, offering valuable insights and practical tools for future research on the link between sharpness and generalization in LLMs and beyond. + +## A. Proofs + +### A.1 Proof of Eq.(10) and Eq.(11) + +###### Proof. + +we propose to approximate the unknown gradient \nabla L_{\mathbf{W}}(\mathbf{W}) using standard LoRA gradients, which can be computed in two ways: + +\displaystyle(1)\nabla L_{\mathbf{B}}\displaystyle(\mathbf{W_{0}}+s\mathbf{BA})=s\nabla L_{\mathbf{W}}(\mathbf{W})\mathbf{A}^{\top} +\displaystyle\Rightarrow\quad\nabla L_{\mathbf{W}}(\mathbf{W})=\frac{1}{s}\nabla L_{\mathbf{B}}(\mathbf{W_{0}}+s\mathbf{BA})(\mathbf{A}^{\top})^{+},(25) +\displaystyle(2)\nabla L_{\mathbf{A}}\displaystyle(\mathbf{W_{0}}+s\mathbf{BA})=s\mathbf{B}^{\top}\nabla L_{\mathbf{W}}(\mathbf{W}) +\displaystyle\Rightarrow\quad\nabla L_{\mathbf{W}}(\mathbf{W})=\frac{1}{s}(\mathbf{B}^{\top})^{+}\nabla L_{\mathbf{A}}(\mathbf{W_{0}}+s\mathbf{BA}),(26) + +∎ + +### A.2 Proof of Theorem 1 + +###### Proof. + +The update process of the FMLoRA is as follows: + +\displaystyle{\tilde{\mathbf{x}}_{t}}={{\mathbf{x}}_{t}}+\rho\frac{1}{s}\frac{{{\mathbf{G}_{t}}}}{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}{\mathbf{y}_{t}}^{+}\displaystyle,\quad{\tilde{\mathbf{y}}_{t}}={{\mathbf{y}}_{t}}(27) +\displaystyle{\mathbf{g}_{{\tilde{\mathbf{x}}_{t}}}}={{\tilde{\mathbf{G}}}_{t}}\tilde{\mathbf{y}}_{t}\displaystyle,\quad{\mathbf{g}_{{\tilde{\mathbf{y}}_{t}}}}={{\tilde{\mathbf{G}}}_{t}}^{\top}\tilde{\mathbf{x}}_{t} +\displaystyle{{\mathbf{x}}_{t+1}}={{\mathbf{x}}_{t}}-\eta{\mathbf{g}_{{\tilde{\mathbf{x}}_{t}}}}\displaystyle,\quad{{\mathbf{y}}_{t+1}}={{\mathbf{y}}_{t}}-\eta{\mathbf{g}_{{\tilde{\mathbf{y}}_{t}}}} + +where {\mathbf{x}}_{t}=\text{Vector}(\mathbf{B}_{t}) is the vectorized form of matrix \mathbf{B}_{t}, {\mathbf{y}}_{t} is the vectorized form of matrix \mathbf{A}_{t}, {\mathbf{G}_{t}}=\nabla L({\mathbf{x}}_{t}{\mathbf{y}}_{t}^{\top}) is the gradient of the full parameter space at the original point during gradient descent, {\tilde{\mathbf{G}}_{t}}=\nabla L(\tilde{\mathbf{x}}_{t}\tilde{\mathbf{y}}_{t}^{\top}) is the gradient of the full parameter space at the perturbed point, and \mathbf{y}^{+} is the pseudo inverse of \mathbf{y}. Let balancedness B_{t}:=\frac{1}{2}(||\mathbf{x}_{t}||^{2}-||\mathbf{y}_{t}||^{2}). Then, we have that: + +\displaystyle\frac{1}{2}\frac{{d({{\left\|{{\mathbf{x}_{t}}}\right\|}^{2}}-{{\left\|{{\mathbf{y}_{t}}}\right\|}^{2}})}}{{dt}}(28) +\displaystyle=\frac{1}{2}\frac{{d({{\left\|{{\mathbf{x}_{t}}}\right\|}^{2}})}}{{dt}}-\frac{1}{2}\frac{{d({{\left\|{{\mathbf{y}_{t}}}\right\|}^{2}})}}{{dt}} +\displaystyle={\mathbf{x}_{t}}^{\top}\frac{{d{\mathbf{x}_{t}}}}{{dt}}-{{\mathbf{y}}_{t}}^{\top}\frac{{d{\mathbf{y}_{t}}}}{{dt}} +\displaystyle=-{\mathbf{x}_{t}}^{\top}({\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}})+({\mathbf{y}_{t}}^{\top}({\mathbf{\tilde{G}}_{t}}^{\top}({{\mathbf{x}}_{t}}+\rho\frac{1}{s}\frac{{{\mathbf{G}_{{t}}}}}{{{{\left\|{\mathbf{G}_{t}}\right\|}_{F}}}}\mathbf{y}_{t}^{+}))) +\displaystyle=-{\mathbf{x}_{t}}^{\top}({\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}})+({\mathbf{y}_{t}}^{\top}({\mathbf{\tilde{G}}_{t}}^{\top}{\mathbf{{x}}_{t}}+\rho\frac{1}{s}{\mathbf{\tilde{G}}_{t}}^{\top}\frac{{{\mathbf{G}_{{t}}}}}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\mathbf{y}_{t}^{+})) +\displaystyle=-{\mathbf{x}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}}+({\mathbf{x}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}}){{}^{\top}}+\rho\frac{1}{s}{\mathbf{y}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}^{\top}\frac{{{\mathbf{G}_{{t}}}}}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\mathbf{y}_{t}^{+} +\displaystyle=\rho\frac{1}{s}{\mathbf{y}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}^{\top}\frac{{{\mathbf{G}_{{t}}}}}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\mathbf{y}_{t}^{+} +\displaystyle=\rho\frac{1}{s}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{\mathbf{y}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}^{\top}{\mathbf{G}_{{t}}}\mathbf{y}_{t}^{+}}\right] + +Because \frac{1}{{{s}}}\mathbf{g_{x}}={\mathbf{G}_{{t}}}{\mathbf{y}_{t}} and {\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}={\mathbf{\tilde{G}}_{t}}{\mathbf{\tilde{y}}_{t}}, we have: + +\displaystyle\frac{1}{2}\frac{{d({{\left\|{{\mathbf{x}_{t}}}\right\|}^{2}}-{{\left\|{{\mathbf{y}_{t}}}\right\|}^{2}})}}{{dt}}(29) +\displaystyle=\rho\frac{1}{s}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{\mathbf{y}_{t}}^{\top}{\mathbf{\tilde{G}}_{t}}^{\top}{\mathbf{G}_{{t}}}\mathbf{y}_{t}^{+}}\right] +\displaystyle=\rho\frac{1}{{{s^{2}}}}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{{({\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}})}^{\top}}\mathbf{{g}_{x}}{{(\mathbf{y}_{t}^{\top})}^{+}}\mathbf{y}_{t}^{+}}\right] +\displaystyle=\rho\frac{1}{{{s^{2}}}}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{{({\mathbf{\tilde{G}}_{t}}{{}\mathbf{y}_{t}})}^{\top}}\mathbf{g_{x}}{{(\mathbf{y}_{t}^{+})}^{\top}}\mathbf{y}_{t}^{+}}\right] +\displaystyle=\rho\frac{1}{{{s^{2}}}}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{{({\mathbf{\tilde{G}}_{t}}{\mathbf{y}_{t}})}^{\top}}\mathbf{g_{x}}{{\left\|{\mathbf{y}_{t}^{+}}\right\|}^{2}}}\right] +\displaystyle=\rho\frac{1}{{{s^{2}}}}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{{({\mathbf{g}_{{{\mathbf{{\tilde{x}}}}_{t}}}}({\mathbf{\tilde{y}}_{t}}^{+})^{\top}{\mathbf{y}_{t}})}^{\top}}\mathbf{g_{x}}{{\left\|{\mathbf{y}_{t}^{+}}\right\|}^{2}}}\right] +\displaystyle=\rho\frac{1}{{{s^{2}}}}\frac{1}{{{{\left\|{{\mathbf{G}_{t}}}\right\|}_{F}}}}\left[{{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}^{\top}\mathbf{g_{x}}{{\left\|{\mathbf{y}_{t}^{+}}\right\|}^{2}}}\right] + +Taking the absolute value of balancedness B_{t} gives: + +\displaystyle\left|{\frac{1}{2}\frac{{d({{\left\|{{\mathbf{x}_{t}}}\right\|}^{2}}-{{\left\|{{\mathbf{y}_{t}}}\right\|}^{2}})}}{{dt}}}\right|(30) +\displaystyle=\left|{\rho\frac{1}{s}\frac{{{{\left\|{\mathbf{y}_{t}^{+}}\right\|}^{2}}}}{{{{\left\|{{\mathbf{{g}}_{\mathbf{x}}}{{(\mathbf{y}_{t}^{+})}^{\top}}}\right\|}_{F}}}}({{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}^{\top}\mathbf{g_{x}}})}\right| +\displaystyle=\left|{\rho\frac{1}{s}\frac{{{{\left\|{\mathbf{y}_{t}^{+}}\right\|}^{2}}}}{{\left\|{\mathbf{{g}_{x}}}\right\|\left\|{{{(\mathbf{y}_{t}^{+})}^{\top}}}\right\|}}({{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}^{\top}\mathbf{g_{x}}})}\right| +\displaystyle=\left|{\rho\frac{1}{s}\frac{{\left\|{\mathbf{y}_{t}^{+}}\right\|}}{{\left\|{\mathbf{{g}_{x}}}\right\|}}({{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}^{\top}\mathbf{g_{x}}})}\right| +\displaystyle\leq\left|{\rho\frac{1}{s}\frac{{\left\|{\mathbf{y}_{t}^{+}}\right\|}}{{\left\|{\mathbf{g_{x}}}\right\|}}\left\|{{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}}\right\|\left\|{\mathbf{{{g}}_{x}}}\right\|}\right| +\displaystyle=\left|{\rho\frac{1}{s}\left\|{\mathbf{y}_{t}^{+}}\right\|\left\|{{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}}\right\|}\right| +\displaystyle=\left|{\rho\frac{1}{s}\left\|{\frac{{\mathbf{y}_{t}^{\top}}}{{{{\left\|{\mathbf{y}_{t}}\right\|}^{2}}}}}\right\|\left\|{{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}}\right\|}\right| +\displaystyle=\left|{\rho\frac{1}{s}\frac{1}{{\left\|{\mathbf{y}_{t}}\right\|}}\left\|{{\mathbf{g}_{{{{\mathbf{\tilde{x}}}}_{t}}}}}\right\|}\right| + +The proof is thus completed. ∎ + +###### Lemma 1. + +Let A_{t+1}=\alpha A_{t}+\beta with some \alpha\in(0,1), then we have + +A_{t+1}\leq\alpha^{t+1}A_{0}+\frac{\beta}{1-\alpha}. + +###### Proof. + +The proof can be completed by simply unrolling A_{t+1} and using the fact 1+\alpha+\alpha^{2}+\dots+\alpha^{t}\leq\frac{1}{1-\alpha}. ∎ + +### A.3 Proof of Theorem 2 + +###### Proof. + +Assume that \bm{\varepsilon}_{t} is the perturbation at time step t, and \bm{\hat{\varepsilon}}_{t-1} is the EMA perturbation from the previous step. Let \nabla L(\mathbf{w}_{t}+\bm{\hat{\varepsilon}}_{t-1}) denote the gradient used for updating at time t. The standard SAM perturbation at step t is defined as \bm{\tilde{\varepsilon}}_{t}=\rho_{t}\frac{\nabla L(\mathbf{w}_{t})}{\|\nabla L(\mathbf{w}_{t})\|}, and the EMA perturbation at step t is computed as \bm{\hat{\varepsilon}}_{t}=(1-\beta)\bm{\hat{\varepsilon}}_{t-1}+\beta\bm{\varepsilon}_{t}. Based on Assumption 4, we have that: + +\displaystyle\left[{L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})-L({\mathbf{w}_{t}})}\right]-\left[{L({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right](31) +\displaystyle=L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})-L({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}}) +\displaystyle\leq-\nabla L{({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})^{\top}}({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}}-{\mathbf{w}_{t}}-{\bm{\hat{\varepsilon}}_{t-1}}) +\displaystyle=\nabla L{({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})^{\top}}({\bm{\hat{\varepsilon}}_{t-1}}-{\bm{\tilde{\varepsilon}}_{t}}) +\displaystyle\leq\left|{\nabla L{{({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})}^{\top}}({\bm{\hat{\varepsilon}}_{t-1}}-{\bm{\tilde{\varepsilon}}_{t}})}\right| +\displaystyle\leq\left\|{\nabla L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})}\right\|\left\|{{\bm{\hat{\varepsilon}}_{t-1}}-{\bm{\tilde{\varepsilon}}_{t}}}\right\|(32) + +For the first term \left\|{\nabla L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})}\right\| in Eq.([32](https://arxiv.org/html/2508.00522v3#Sx5.E32 "Equation 32 ‣ A.3 Proof of Theorem 2 ‣ A. Proofs ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), Based on Assumption 1, Assumption 2 and Lemma 1, we have: + +\displaystyle\left\|{\nabla L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})}\right\|(33) +\displaystyle=\left\|{\nabla L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})-\nabla L({\mathbf{w}_{t}})+\nabla L({\mathbf{w}_{t}})}\right\| +\displaystyle\leq\left\|{\nabla L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})-\nabla L({\mathbf{w}_{t}})}\right\|+\left\|{\nabla L({\mathbf{w}_{t}})}\right\| +\displaystyle\leq\tau\left\|{{\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}}-{\mathbf{w}_{t}}}\right\|+\left\|{\nabla L({\mathbf{w}_{t}})}\right\| +\displaystyle=\tau\left\|{{\bm{\hat{\varepsilon}}_{t-1}}}\right\|+\left\|{\nabla L({\mathbf{w}_{t}})-\nabla{L_{\rm{D}}}({\mathbf{w}_{t}})+\nabla{L_{\rm{D}}}({\mathbf{w}_{t}})}\right\| +\displaystyle=\tau\left\|{{\bm{\hat{\varepsilon}}_{t-1}}}\right\|+\left\|{\nabla{L_{\rm{D}}}({\mathbf{w}_{t}})}\right\|+{\sigma^{2}} +\displaystyle=\tau\left\|{(1-\beta){\bm{\hat{\varepsilon}}_{t-2}}+\beta\bm{\varepsilon}_{t-1}}\right\|+G+{\sigma^{2}} +\displaystyle\leq\tau((1-\beta)\left\|{{\bm{\hat{\varepsilon}}_{t-2}}}\right\|+\beta{\rho_{0}})+G+{\sigma^{2}} +\displaystyle\leq\tau{(1-\beta)^{t-1}}\left\|{{\bm{\hat{\varepsilon}}_{0}}}\right\|+\tau{\rho_{0}}+G+{\sigma^{2}} + +For the second term \left\|{{\bm{\hat{\varepsilon}}_{t-1}}-{\bm{\tilde{\varepsilon}}_{t}}}\right\| in Eq.([32](https://arxiv.org/html/2508.00522v3#Sx5.E32 "Equation 32 ‣ A.3 Proof of Theorem 2 ‣ A. Proofs ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond")), we have: + +\displaystyle\left\|{{\bm{\hat{\varepsilon}}_{t-1}}-{\bm{\tilde{\varepsilon}}_{t}}}\right\|(34) +\displaystyle\leq\left\|{{\bm{\hat{\varepsilon}}_{t-1}}}\right\|+\left\|{{\bm{\tilde{\varepsilon}}_{t}}}\right\| +\displaystyle=\left\|{{\bm{\hat{\varepsilon}}_{t-1}}}\right\|+{\rho_{\rm{t}}} +\displaystyle\leq{(1-\beta)^{t-1}}\left\|{{\bm{\hat{\varepsilon}}_{0}}}\right\|+{\rho_{0}}+{\rho_{t}} + +Let {\bm{\hat{\varepsilon}}_{0}}={\bm{\tilde{\varepsilon}}_{0}}=\rho_{0}\frac{\nabla L(\mathbf{w}_{0})}{\|\nabla L(\mathbf{w}_{0})\|}, {\rho_{t}}=\frac{{{\rho_{0}}}}{{\sqrt{t}}}, we have: + +\displaystyle\left[{L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t-1}})-L({\mathbf{w}_{t}})}\right]-\left[{L({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right](35) +\displaystyle\leq\left(\tau{(1-\beta)^{t-1}}\left\|{{\bm{\hat{\varepsilon}}_{0}}}\right\|+\tau{\rho_{0}}+G+{\sigma^{2}}\right) +\displaystyle\quad\cdot\left({(1-\beta)^{t-1}}\left\|{{\bm{\hat{\varepsilon}}_{0}}}\right\|+{\rho_{0}}+{\rho_{t}}\right) +\displaystyle=\left({\left({1+{{(1-\beta)}^{t-1}}}\right)\tau{\rho_{0}}+G+{\sigma^{2}}}\right) +\displaystyle\quad\quad\cdot\left({\left({1+{{(1-\beta)}^{t-1}}}\right){\rho_{0}}+\frac{{{\rho_{0}}}}{{\sqrt{t}}}}\right) + +The proof is thus completed. ∎ + +## B. Experimental Details + +### B.1 Details on datasets + +Our evaluations are carried out on commonly-used datasets in the literature. + +Datasets for few-shot learning of RoBERTa-large. We consider classification datasets: SST-2 (socher2013recursive), SST-5 (socher2013recursive), TREC (voorhees2000building), MNLI (williams2018broad), SNLI (bowman2015large), and RTE (dagan2005pascal). We follow Malladi et al. (malladi2023kernel) in limiting the test set to 1, 000 examples for fast iteration. For training and validation, we set k = 512, which mean that we have 512 examples per class for both training and validation. + +Table 7: The hyperparameters used for RoBERTa large with LoRA on the GLUE benchmark. + +Table 8: Hyperparameters used for few-shot learning with RoBERTa-large. + +Table 9: Hyperparameters used for GPT2. + +GLUE benchmark. GLUE is designed to provide a general-purpose evaluation of language understanding (wangglue). Those adopted in our work include MNLI (inference, (williams2018broad)), SST-2 (sentiment analysis, (socher2013recursive)), MRPC (paraphrase detection, (dolan2005automatically)), CoLA (linguistic acceptability (warstadt2019neural)), QNLI (inference (rajpurkar2018know)), QQP 1 1 1 https://quoradata.quora.com/First-Quora-Dataset-Release-Question-Pairs (question-answering), RTE 2 2 2 https://paperswithcode.com/dataset/rte (inference), and STS-B (textual similarity (cer2017semeval)). These datasets are released under different permissive licenses. + +E2E NLG Challenge. The E2E NLG Challenge dataset (novikova2017e2e) is a standard benchmark for end-to-end data-to-text natural language generation. It consists of around 42,000 training instances, along with 4,600 each for validation and testing, all within the restaurant domain. Inputs are structured as sequences of slot-value pairs and paired with one or more reference texts. The dataset is released under the Creative Commons BY-NC-SA 4.0 license. + +Datasets for few-shot learning of CLIP. We consider five datasets for fine-grained classification of satellite imagery (EuroSAT (helber2019eurosat)), pet breeds (Ox-fordPets (parkhi2012cats)), flowers (Flower102 (nilsback2008automated)), general objects (Caltech101 (fei2004learning)), textures (DTD (cimpoi2014describing)). These datasets offer a thorough benchmarking framework for evaluating few-shot visual classification tasks. + +Datasets for fine-tuning with Qwen-VL-Chat. We use two representative datasets: ScienceQA (lu2022learn) and VizWiz (gurari2018vizwiz). ScienceQA is a multimodal multiple-choice QA dataset covering elementary science, with questions accompanied by text and images. VizWiz is a real-world visual QA dataset collected from blind users, featuring diverse and often low-quality images, posing challenges for robust multimodal understanding. + +### B.2 Details on models + +We summarize the adopted language models in our evaluation. All model checkpoints are obtained from HuggingFace. + +RoBERTa-large. This is a 355 M parameter model. The model checkpoint 3 3 3 https://huggingface.co/FacebookAI/roberta-large is released under the MIT license. + +GPT2-medium. This is a 345 M parameter model. Its checkpoint 4 4 4 https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-pytorch˙model.bin is under MIT License. + +GPT2-large. This is a 774 M parameter model. Its checkpoint 5 5 5 https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-pytorch˙model.bin is under MIT License. + +CLIP. This is a model that learns to connect images and text by mapping them into a shared semantic space using contrastive learning. + +Qwen-VL-Chat. Qwen-VL-Chat (Bai2023QwenVLAV) is a multimodal conversational large language model capable of understanding both images and text. + +Table 10: Hyperparameters used for few-shot learning with CLIP. + +Algorithm 1 Pseudocode of the FMLoRA + +Require: The training dataset, the learning rate \eta, the batch size b, parameters \rho and \beta. + +1:for + +t=1,2,\cdot\cdot\cdot +do + +2: Randomly sample a mini-batch; + +3: Evaluate the gradient at the current point; + +4: Apply Equation (12) to compute the gradient in the full parameter space + +{\bar{\mathbf{g}}^{\mathbf{W}}} +; + +5: Use Equation (13) to calculate the perturbation + +\bar{\mathbf{E}}^{\mathbf{W}} +; + +6: Compute the perturbation + +\bar{\mathbf{E}}^{\mathbf{B}}=\frac{1}{s}\bar{\mathbf{E}}^{\mathbf{W}}\mathbf{A}^{+} +on matrix + +\mathbf{B} +according to Equation (14); + +7: Evaluate the gradient at the perturbed point ( + +\mathbf{W}_{0} +, + +\mathbf{B}+\bar{\mathbf{E}}^{\mathbf{B}} +, + +\mathbf{A} +); + +8: Return to the original (unperturbed) parameter point ( + +\mathbf{W}_{0} +, + +\mathbf{B} +, + +\mathbf{A} +); + +9: Update the weights using the gradient obtained in Step 6; + +10:end for + +### B.3 Details on hyperparameters + +Few-shot Learning with RoBERTa. We adopt the k-shot learning setup from (malladi2023fine), focusing on classification tasks with k=512 training samples per class and 1000 samples for testing. Prompt-based finetuning is used, following the same prompt templates as in (malladi2023fine, Table 13). We use AdamW as the optimizer and tune hyperparameters based on Table [8](https://arxiv.org/html/2508.00522v3#Sx6.T8 "Table 8 ‣ B.1 Details on datasets ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). All results are averaged over three random seeds. + +Fine-tuning with RoBERTa-large. AdamW is adopted as the base optimizer, and hyperparameters are in Table [7](https://arxiv.org/html/2508.00522v3#Sx6.T7 "Table 7 ‣ B.1 Details on datasets ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). However, we employ single GPU rather than multiple ones and use gradient accumulation rather than parallelism due to memory constraint. We consider the GLUE benchmark and report the mismatched accuracy for MNLI, Matthew’s correlation for CoLA, Pearson correlation for STS-B, and accuracy for other datasets. Larger values indicate better results for all datasets. Experiments are conducted over three random trials for all datasets. + +GPT2 medium/large on E2E NLG Challenge. We use the batch size, learning rate, and beam search beam size described in (hu2022lora). AdamW is adopted as base optimizer. The hyperparameters can be found in Table [9](https://arxiv.org/html/2508.00522v3#Sx6.T9 "Table 9 ‣ B.1 Details on datasets ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). The result for each run is taken from the last epoch. + +Few-shot Learning with CLIP. We follow the setting of previous work (zanella2024low). The hyperparameters are tuned from those in Table [10](https://arxiv.org/html/2508.00522v3#Sx6.T10 "Table 10 ‣ B.2 Details on models ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). We only apply low-rank matrices on the query, key and value matrices with r=2. We regularize the input of the LoRA module by a dropout layer with p=0.25. The number of iterations is set equal to 500 times N/K (the number of labeled samples per class). + +Fine-tuning with Qwen-VL-Chat. We conduct experiments follow the setting of previous work (zhou2024empirical). The hyperparameters can be found in Table [11](https://arxiv.org/html/2508.00522v3#Sx6.T11 "Table 11 ‣ B.3 Details on hyperparameters ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). + +Table 11: Hyperparameters used for fine-tuning with Qwen-VL-Chat. + +![Image 3: Refer to caption](https://arxiv.org/html/2508.00522v3/x3.png) + +Figure 3: Approximation ability of EMA perturbations across datasets + +![Image 4: Refer to caption](https://arxiv.org/html/2508.00522v3/x4.png) + +Figure 4: Evolution of balancedness across layers during training with Adam and FMLoRA. + +## C. Algorithm + +The two algorithms presented in [1](https://arxiv.org/html/2508.00522v3#alg1 "Algorithm 1 ‣ B.2 Details on models ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") and [2](https://arxiv.org/html/2508.00522v3#alg2 "Algorithm 2 ‣ C. Algorithm ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") describe the training procedures of the proposed FMLoRA and its accelerated variant EFMLoRA. + +Algorithm 2 Pseudocode of the EFMLoRA + +Require: The training dataset, the learning rate \eta, the batch size b, parameters \rho and \beta. + +1:for + +t=1,2,\cdot\cdot\cdot +do + +2: Randomly sample a mini-batch; + +3:if + +t=1 +then + +4: Evaluate the gradient at the current point; + +5: EMA perturbation + +\hat{\mathbf{E}}^{\mathbf{B}}_{1}=\bar{\mathbf{E}}^{{\mathbf{B}}}_{1} +; + +6: Update the weights using the gradient obtained in Step 4; + +7: Update the parameters to the next perturbation point + +(\mathbf{W}_{0} +, + +\mathbf{B}_{1}+\hat{\mathbf{E}}^{\mathbf{B}}_{1} +, + +\mathbf{A}_{1}) +. + +8:else + +9: Calculate the gradient at the perturbation point + +(\mathbf{W}_{0} +, + +\mathbf{B}_{t-1}+\hat{\mathbf{E}}^{\mathbf{B}}_{t-1} +, + +\mathbf{A}_{t-1}) +. + +10: Compute the perturbation + +\bar{\mathbf{E}}^{\mathbf{B}}_{t}=\frac{1}{s}\bar{\mathbf{E}}^{\mathbf{W}}\mathbf{A}^{+}_{t-1} +on matrix + +\mathbf{B} +according to Equation (14); + +11: Return to the original parameter point + +(\mathbf{W}_{0} +, + +\mathbf{B}_{t-1} +, + +\mathbf{A}_{t-1}) +. + +12: Calculate the EMA perturbation + +{{\hat{\mathbf{E}}}^{\mathbf{B}}_{t}}=(1-\beta){{\hat{\mathbf{E}}}^{\mathbf{B}}_{t-1}}+\beta{\bar{\mathbf{E}}^{\mathbf{B}}_{t}} +. + +13: Update the weights to + +(\mathbf{W}_{0} +, + +\mathbf{B}_{t} +, + +\mathbf{A}_{t}) +using the gradient obtained in Step 9; + +14: Update the parameters to the next perturbation point + +(\mathbf{W}_{0} +, + +\mathbf{B}_{t}+\hat{\mathbf{E}}^{\mathbf{B}}_{t} +, + +\mathbf{A}_{t}) +. + +15:end if + +16:end for + +## D. More experiments + +### D.1 The approximate ability of EMA perturbation + +We consider few shot learning with LoRA on RoBERTa-large. Fig.[3](https://arxiv.org/html/2508.00522v3#Sx6.F3 "Figure 3 ‣ B.3 Details on hyperparameters ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond") illustrates the evolution of the difference in sharpness, \left[{L({\mathbf{w}_{t}}+{\bm{\hat{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right]-\left[{L({\mathbf{w}_{t}}+{\bm{\tilde{\varepsilon}}_{t}})-L({\mathbf{w}_{t}})}\right], as described in Theorem 2, during training on six datasets (SNLI, SST-2, SST-5, MNLI, RTE, and TREC). S^{\text{EMA}} denotes the sharpness computed using EMA perturbations, while S^{\text{SAM}} refers to the original SAM sharpness. As training progresses, the absolute difference consistently decreases across all datasets, demonstrating that the EMA perturbation becomes increasingly effective at approximating the SAM perturbations. This validates the use of EMA perturbations as a computationally efficient surrogate for SAM perturbations. This result empirically supports Theorem 2. + +### D.2 The change in balancedness during FMLoRA training + +We consider few shot learning with LoRA on RoBERTa-large. For dataset MNLI, 1st, 12th and 24th query layers’ 2|B_{t,l}| are plotted, where t denotes the iteration and l denotes the layer index. The layers are chosen to represent early, middle, and final stages of RoBERTa. Balancedness of FMLoRA and Adam on different layers are plotted in Fig.[4](https://arxiv.org/html/2508.00522v3#Sx6.F4 "Figure 4 ‣ B.3 Details on hyperparameters ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"). Balancedness may increase or decrease across different layers. As shown in Fig.[4](https://arxiv.org/html/2508.00522v3#Sx6.F4 "Figure 4 ‣ B.3 Details on hyperparameters ‣ B. Experimental Details ‣ Efficiently Seeking Flat Minima for Better Generalization in Fine-Tuning Large Language Models and Beyond"), the balancedness of FMLoRA in the first query layer of RoBERTa-large gradually decreases during training, while in the 12th layer, it first decreases and then increases. In contrast, the balancedness in the 24th layer continuously increases. An increase typically occurs when parameter magnitudes in both low-rank subspaces grow simultaneously. This behavior can be influenced by factors such as the learning rate, optimization algorithm, weight decay, and other regularization strategies. Despite these occasional increases, FMLoRA generally maintains lower balancedness than Adam in most layers, suggesting its capacity to induce implicit regularization during training. diff --git a/docs/lora_mgpo_momentum_perturbation.md b/docs/lora_mgpo_momentum_perturbation.md new file mode 100644 index 0000000..0b6885c --- /dev/null +++ b/docs/lora_mgpo_momentum_perturbation.md @@ -0,0 +1,465 @@ +Title: 2025.findings-emnlp.34.pdf + +URL Source: https://aclanthology.org/2025.findings-emnlp.34.pdf + +Published Time: Fri, 31 Oct 2025 19:44:11 GMT + +Number of Pages: 12 + +Markdown Content: +> Findings of the Association for Computational Linguistics: EMNLP 2025 , pages 648–659 November 4-9, 2025 ©2025 Association for Computational Linguistics + +# LoRA-MGPO: Mitigating Double Descent in Low-Rank Adaptation via Momentum-Guided Perturbation Optimization + +Yupeng Chang 1 Chenlu Guo 1 Yi Chang 1,2,3 Yuan Wu 1* + +> 1 + +School of Artificial Intelligence, Jilin University + +> 2 + +Engineering Research Center of Knowledge-Driven Human-Machine Intelligence, MOE, China + +> 3 + +International Center of Future Science, Jilin University {changyp23, guocl23}@mails.jlu.edu.cn, {yichang, yuanwu}@jlu.edu.cn + +Abstract + +Parameter-efficient fine-tuning (PEFT), partic-ularly Low-Rank Adaptation (LoRA), adapts large language models (LLMs) by training only a small fraction of parameters. However, as the rank of the low-rank matrices used for adap-tation increases, LoRA often exhibits an un-stable "double descent" phenomenon, charac-terized by transient divergence in the training loss, which delays convergence and impairs generalization by causing instability due to the attraction to sharp local minima. To address this, we introduce LoRA-MGPO , a framework that incorporates Momentum-Guided Pertur-bation Optimization (MGPO). MGPO stabi-lizes training dynamics by mitigating the dou-ble descent phenomenon and guiding weight perturbations using momentum vectors from the optimizer’s state, thus avoiding dual gra-dient computations. Additionally, an adaptive normalization scheme scales the magnitude of perturbations based on an exponential mov-ing average (EMA) of gradient norms, further enhancing stability. While EMA controls the magnitude of the perturbations, MGPO guides their direction, ensuring a more stable opti-mization trajectory. Experiments on a suite of natural language understanding and genera-tion benchmarks show that LoRA-MGPO con-sistently achieves superior performance over LoRA and other PEFT methods. The analysis indicates that LoRA-MGPO leads to smoother loss curves, faster convergence, and improved generalization by stabilizing the training pro-cess and mitigating the attraction to sharp min-ima. The code is publicly available at https: //github.com/llm172/LoRA-MGPO . + +1 Introduction + +Large language models (LLMs) have driven signif-icant advancements in natural language process-ing, establishing new performance benchmarks + +> *Corresponding authors + +on tasks ranging from text generation to seman-tic understanding (Chang et al., 2024b; Wei et al., 2022). However, the conventional method of full-parameter fine-tuning (Full FT) requires updating billions of parameters, incurring prohibitive mem-ory and computational costs. To overcome this limitation, parameter-efficient fine-tuning (PEFT) methods have emerged as an effective alternative, enabling efficient adaptation by optimizing only a small subset of model parameters (Lester et al., 2021; Fu et al., 2023). Among these methods, Low-Rank Adaptation (LoRA) (Hu et al., 2021) is distinguished by its computational efficiency and architectural simplic-ity. LoRA approximates the weight update ma-trix ∆W as a low-rank decomposition, where the original pre-trained weights W0 remain frozen. The trainable matrices B and A, with rank r ≪ + +min( m, n ), drastically reduce the number of train-able parameters, improving efficiency without al-tering the model architecture. Despite its efficiency, LoRA’s training dynam-ics can be unstable. As shown in Figure 1, fine-tuning LLaMA-2-7B (Touvron et al., 2023) on MetaMathQA (Yu et al., 2024) often exhibits a "double descent" trajectory with initial conver-gence, transient divergence, and eventual stabi-lization. This phenomenon worsens with higher ranks and is not unique to LoRA; Full FT can ex-hibit even more severe double descent, highlighting the general challenge of stabilizing fine-tuning in high-capacity models (Nakkiran et al., 2019). Such non-monotonic behavior delays convergence and impairs generalization due to unstable gradients and the attraction to sharp local minima (Li et al., 2024a). Addressing these stability issues is crucial. Sharpness-Aware Minimization (SAM) (Foret et al., 2020) improves generalization by seeking flatter minima. However, its application is hin-dered by the dual gradient computation require- + +648 0 1000 2000 3000 4000 5000 6000 + +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> Training Loss +> LoRA-rank32 +> 01000 2000 3000 4000 5000 6000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> LoRA-rank64 +> 01000 2000 3000 4000 5000 6000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> LoRA-rank128 +> 01000 2000 3000 4000 5000 6000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> 1.4 +> 1.6 +> Full FT + +Figure 1: Training loss curves of Full FT and LoRA (Hu et al., 2021) methods with LLaMA-2-7B (Touvron et al., 2023) on the MetaMathQA dataset (Yu et al., 2024). For LoRA, rank ( r) and alpha ( α) are set to the same values (r = α ∈ { 32 , 64 , 128 }), with a fixed learning rate of 5e − 4. + +ment, which doubles the training cost (Becker et al., 2024; Li et al., 2024b). More efficient variants like momentum-guided SAM reuse optimizer states to avoid this overhead but may not guarantee stable convergence. To further enhance stability, comple-mentary techniques such as applying an exponen-tial moving average (EMA) to smooth optimization dynamics have been shown to suppress parameter oscillations and improve convergence in certain scenarios (Wang et al., 2021). Building on these insights, we propose LoRA-MGPO , a novel framework that integrates Momentum-Guided Perturbation Optimization (MGPO) into LoRA to mitigate the detrimental effects of double descent. Our contributions are twofold: 1. Mitigating Double Descent: MGPO stabi-lizes training by addressing double descent, typically observed at higher ranks in LoRA. By reusing momentum vectors, it guides weight perturbations towards flatter minima, preventing transient divergences in loss. 2. Adaptive Perturbation Normalization: + +MGPO introduces an adaptive scheme that scales perturbation magnitude based on an exponential moving average (EMA) of gradient norms, decoupling perturbation intensity from optimization dynamics and further enhancing stability. We evaluate LoRA-MGPO on a suite of natural lan-guage understanding (NLU) and generation (NLG) benchmarks. Our results show that it consistently achieves superior performance over standard LoRA and other state-of-the-art PEFT methods. Crucially, we demonstrate that LoRA-MGPO effectively mit-igates the double descent phenomenon, leading to more stable training dynamics, smoother loss curves, and faster convergence, all of which con-tribute to better generalization and the avoidance of sharp minima. + +2 Method + +In this section, we first provide a concise overview of the Low-Rank Adaptation (LoRA) framework. We then introduce LoRA-MGPO , an extension of LoRA that integrates Momentum-Guided Perturba-tion Optimization (MGPO) to enhance its stability and efficiency. We describe how MGPO reuses op-timizer momentum for guided perturbations of the trainable parameters and incorporates an adaptive normalization scheme to stabilize training. + +2.1 Review of LoRA + +While full fine-tuning directly updates the entire pre-trained weight matrix W0 ∈ Rm×n, its pro-hibitive computational cost makes it impractical for large-scale models. Low-Rank Adaptation (LoRA) (Hu et al., 2021) offers a parameter-efficient al-ternative. LoRA freezes W0 and injects a train-able low-rank decomposition, ∆W = BA , where + +B ∈ Rm×r and A ∈ Rr×n are trainable matrices with rank r ≪ min( m, n ). The weight update is incorporated into the forward pass as: + +Y = X(W0 + αr BA ), (1) where X is the input, α is a scaling hyperparameter, and r is the rank of the decomposition. Typically, + +A is initialized with a Kaiming normal distribution, and B with zeros. While effective, LoRA can suf-fer from training instability, particularly the double descent phenomenon, when r increases without appropriate optimization strategies to maintain sta-bility (Li et al., 2024a). 649 2.2 LoRA with Momentum-Guided Perturbation Optimization + +To address the training instabilities in LoRA, we propose LoRA-MGPO , which integrates Momentum-Guided Perturbation Optimization (MGPO). Inspired by Sharpness-Aware Minimiza-tion (SAM), MGPO is redesigned for computa-tional efficiency and parameter efficiency. It di-rectly perturbs the trainable LoRA parameters by reusing the optimizer’s first-moment estimate, guid-ing the perturbations toward stable directions. Ad-ditionally, MGPO incorporates adaptive normaliza-tion to dynamically scale the perturbation, enhanc-ing training stability. + +2.2.1 Motivation: SAM for LoRA and Its Limitations + +The goal of SAM (Foret et al., 2020) is to find parameters in flat loss regions to improve gener-alization. A direct application to LoRA would involve perturbing the full weight matrix, solv-ing min A,B max ∥ϵ∥F ≤ρ L (W0 + BA + ϵ). This approach is ill-suited for PEFT due to two critical flaws: (1) its dual gradient computation require-ment doubles the training cost, and (2) creating and storing the full-space perturbation ϵ counteracts the memory savings of LoRA. MGPO is explicitly designed to resolve these inefficiencies. + +2.2.2 Momentum-Guided Perturbation of LoRA Parameters + +MGPO achieves the stability benefits of SAM by perturbing the trainable parameters θ = ( A, B ) + +directly, using information readily available in the optimizer’s state. At each training step t, instead of computing a new gradient for the perturbation direction, it reuses the optimizer’s first-moment vector (momentum) from the previous step, mt−1.The optimization objective is: + +min + +> θ + +L(θt + ϵθt ), (2) where the perturbation ϵθt applied to the LoRA parameters θt = ( At, B t) is constructed using the state from step t − 1: + +ϵθt = ρ · mt−1 + +∥mt−1∥2 + +· 1¯g(t−1) . (3) Here, ρ is the perturbation radius. Using the histori-cal momentum vector is a deliberate design choice, as it represents a smoothed average of past gradi-ents, filtering out the noise from any single mini-batch and providing a more stable direction for assessing landscape sharpness. This vector is main-tained by the optimizer itself. After computing the gradient on the perturbed parameters, the momen-tum for the current step is updated as: + +mt = μmt−1 + ∇˜θt L. (4) The decay factor μ (e.g., ‘beta1‘ in AdamW) is reused from the optimizer’s standard settings. The scalar ¯g(t−1) is a global normalization factor, de-tailed next. This formulation entirely avoids the second gradient computation and any operations in the full weight space. + +Two-Stage Update Mechanism MGPO is imple-mented efficiently within each training step t. First, using the state from step t − 1, we compute the per-turbation ϵθt and apply it to the current parameters + +θt to get a perturbed version, ˜θt: + +˜θt = θt + ϵθt . (5) Second, the loss and its gradient are computed with respect to these perturbed parameters: ∇˜θt L. This single gradient is then used by the optimizer to up-date both the original parameters from θt to θt+1 + +and the momentum from mt−1 to mt. For infer-ence, the final, unperturbed parameters θT are used. + +2.2.3 Adaptive Perturbation Normalization + +To ensure robustness across training stages, we in-troduce an Adaptive Perturbation Normalization (APN) scheme. The normalization factor ¯g(t) used in Equation 3 is a scalar computed via an exponen-tial moving average (EMA) of the global L2-norm of the LoRA parameter gradients. Following the principle of using the actually computed gradient, the update rule is: + +¯g(t) = β¯g(t−1) + (1 − β)∥∇ ˜θt L∥ 2, (6) where β is the EMA decay rate. This mechanism makes the perturbation scale-invariant relative to the gradient dynamics. For instance, during early training with large gradients, the normalization fac-tor increases, reducing the effective perturbation size to prevent destabilization. Conversely, in later stages, it ensures the perturbation remains suffi-ciently large to be effective. This adaptive scaling enhances training stability. + +3 Experiments + +3.1 Experimental Setup Baselines To provide a comprehensive evalua-tion, we compare LoRA-MGPO against a carefully 650 Table 1: Performance of T5-Base on five GLUE tasks, comparing LoRA-MGPO with full fine-tuning and other LoRA variants (rank r = 8 ). Scores are reported for the primary metric of each task, averaged over 3 runs, with standard deviations shown in subscripts. Bold indicates the best score, while underlining denotes the second best. + +Method MNLI SST2 CoLA QNLI MRPC Avg + +Train Size 393k 67k 8.5k 105k 3.7k + +Full FT 86.33 ±0.00 94.75 ±0.21 80.70 ±0.24 93.19 ±0.22 84.56 ±0.73 87.91 LoRA 85.30 ±0.04 94.04 ±0.11 69.35 ±0.05 92.96 ±0.09 68.38 ±0.01 82.08 + +LoRA Variants with Modified Structure + +DoRA 85.67 ±0.09 94.04 ±0.53 72.04 ±0.94 93.04 ±0.06 68.08 ±0.51 82.57 AdaLoRA 85.45 ±0.11 93.69 ±0.20 69.16 ±0.24 91.66 ±0.05 68.14 ±0.28 81.62 + +LoRA Variants with Original Structure + +PiSSA 85.75 ±0.07 94.07 ±0.06 74.27 ±0.39 93.15 ±0.14 76.31 ±0.51 84.71 rsLoRA 85.73 ±0.10 94.19 ±0.23 72.32 ±1.12 93.12 ±0.09 52.86 ±2.27 79.64 LoRA+ 85.81 ±0.09 93.85 ±0.24 77.53 ±0.20 93.14 ±0.03 74.43 ±1.39 84.95 LoRA-GA 85.70 ±0.09 94.11 ±0.18 80.57 ±0.20 93.18 ±0.06 85.29 ±0.24 87.77 LoRA-MGPO 86.58 ±0.11 94.72 ±0.46 82.32 ±0.18 93.79 ±0.46 86.62 ±0.68 88.81 + +selected set of baselines. These include Full Fine-Tuning (Full FT), serving as a strong performance benchmark, and vanilla LoRA (Hu et al., 2021), our primary point of comparison. We further in-clude two categories of state-of-the-art LoRA vari-ants. The first category, variants with architec-tural modifications , comprises methods that alter the LoRA structure itself, such as DoRA (Liu et al., 2024), which introduces learnable magnitude vec-tors, and AdaLoRA (Zhang et al., 2023), which dynamically allocates rank budgets. The second category, variants improving the training process or initialization , includes rsLoRA (Kalajdzievski, 2023), which stabilizes update magnitudes; LoRA+ (Hayou et al., 2024), which employs different learn-ing rates for the LoRA matrices; and PiSSA (Meng et al., 2024), which refines initialization using SVD. Finally, we compare against methods focused on + +gradient alignment , such as LoRA-GA (Wang et al., 2024a) and LoRA-Pro (Wang et al., 2024b), which aim to align LoRA’s gradient updates more closely with those of full fine-tuning. + +Datasets Our experiments span a range of tasks in natural language understanding and generation. For NLU, we evaluate on five tasks from the widely-used General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2018): MNLI, SST-2, CoLA, QNLI, and MRPC. These tasks cover natural language inference, sentiment analy-sis, grammatical acceptability, and paraphrase iden-tification. For NLG, we fine-tune the LLaMA-2-7B (Tou-vron et al., 2023) model on a 52k randomly sam-pled subset of the WizardLM dataset (Xu et al., 2024). We evaluate the model on the MT-Bench dataset (Zheng et al., 2024a), which consists of 80 multi-turn questions designed to assess conversa-tional abilities across various aspects. The quality of the responses is evaluated by GPT-4, and we report the first-turn score as the primary evaluation metric. For mathematical reasoning, we use a 100k ran-dom sample from MetaMathQA (Yu et al., 2024), with evaluation on the GSM8K test set (Cobbe et al., 2021). For code generation, fine-tuning is performed on a 100k randomly sampled subset of the CodeFeedback dataset (Zheng et al., 2024b), with evaluation on HumanEval (Chen et al., 2021). + +Implementation Details For fair comparison, our experimental setup closely follows that of LoRA-GA (Wang et al., 2024a). Across all experi-ments, we use the AdamW optimizer (Loshchilov and Hutter, 2019) with weight decay set to 0 and a cosine learning rate schedule with a warm-up ratio of 0.03. LoRA adapters are applied to all linear layers within the transformer blocks, with the rank + +r set to 8 and scaling factor α to 16 by default. For our two task families, the settings are as follows. For Natural Language Understanding (NLU) on GLUE, we fine-tune T5-base (Raffel et al., 2020) with a learning rate of 1 × 10 −4, a sequence length of 128, and a batch size of 32. The MGPO hy-651 Table 2: Fine-tuning results of LLaMA-2-7B on MT-Bench, GSM8K, and HumanEval. Performance is evaluated using primary task metrics: MT-Bench score, GSM8K accuracy, and HumanEval Pass@1. PEFT methods are tested with rank r = 8 , and additional tests at ranks 32 and 128 are included to evaluate performance scaling. Results are averaged over three random seeds, with standard deviations provided. Bold and underlining denote the best and second-best scores, respectively. + +Method MT-Bench GSM8K HumanEval Avg + +Full FT 5.30 ±0.11 59.36 ±0.85 35.31 ±2.13 33.32 + +LoRA 5.61 ±0.10 42.08 ±0.04 14.76 ±0.17 20.82 DoRA 5.97 ±0.02 53.07 ±0.75 19.75 ±0.41 26.26 AdaLoRA 5.57 ±0.05 50.72 ±1.39 17.80 ±0.44 24.70 PiSSA 5.30 ±0.02 44.54 ±0.27 16.02 ±0.78 21.95 rsLoRA 5.25 ±0.03 45.62 ±0.10 16.01 ±0.79 22.29 LoRA+ 5.71 ±0.08 52.11 ±0.62 18.17 ±0.52 25.33 LoRA-GA 5.95 ±0.16 53.60 ±0.30 19.81 ±1.46 26.45 LoRA-GA (rank=32) 5.79 ±0.09 55.12 ±0.30 20.18 ±0.19 27.03 LoRA-GA (rank=128) 6.13 ±0.07 55.07 ±0.18 23.05 ±0.37 28.08 LoRA-MGPO 6.27 ±0.12 54.56 ±0.44 21.02 ±0.39 27.28 LoRA-MGPO (rank=32) 6.21 ±0.15 55.74 ±0.21 21.34 ±0.47 27.76 LoRA-MGPO (rank=128) 6.48 ±0.23 56.96 ±0.35 24.87 ±0.54 29.44 perparameters are ρ = 0 .05 , μ = 0 .9 (AdamW’s ‘beta1‘), and β = 0 .9. For Natural Language Gener-ation (NLG), we fine-tune LLaMA-2-7B (Touvron et al., 2023) with a learning rate of 2 × 10 −5 and a sequence length of 1024. We use a per-device batch size of 4 with 8 gradient accumulation steps for an effective batch size of 32. The MGPO hyperparam-eters are ρ = 0 .01 , μ = 0 .8 (AdamW’s ‘beta1‘), and β = 0 .8. All experiments were conducted on NVIDIA H20 96GB GPUs, repeated three times with different random seeds, and we report the av-erage and standard deviation of the results. Further details on optimizer settings, specific LoRA target modules, and the software environment are pro-vided in the Appendix. + +3.2 Main Results Performance on Natural Language Understand-ing (NLU) We first evaluated LoRA-MGPO on a standard suite of NLU tasks from the GLUE bench-mark (Wang et al., 2018), using the T5-base model. As detailed in Table 1, our method demonstrates strong and consistent performance. The improve-ments are particularly notable on challenging, low-resource benchmarks such as CoLA and MRPC, where LoRA-MGPO surpasses not only all other PEFT methods but also full fine-tuning. Success on these tasks often hinges on capturing subtle lin-guistic nuances. The stability afforded by LoRA-MGPO likely prevents the fine-tuning process from corrupting the rich knowledge encoded in the base model; by preventing erratic weight updates, our method may better preserve the pre-trained model’s nuanced understanding of syntax and semantics. Quantitatively, LoRA-MGPO achieves the highest scores among all PEFT methods on five out of five tasks, obtains the best average score, and outper-forms the next-best PEFT method, LoRA-GA, by a margin of 1.04 points. + +Performance on Natural Language Genera-tion (NLG) We further assessed our method on three challenging NLG tasks using the LLaMA-2-7B model, with results summarized in Ta-ble 2. LoRA-MGPO consistently secures top performance among all PEFT baselines. On the conversational MT-Bench, its top score suggests that stable training helps maintain the model’s coherence and instruction-following capabilities. For structured reasoning tasks like mathemati-cal problem-solving (GSM8K) and code genera-tion (HumanEval), where logical consistency is paramount, LoRA-MGPO again emerges as the strongest PEFT method. A stable optimization tra-jectory may reduce the risk of the model deviating from a correct reasoning path during fine-tuning, as each update step is more measured, preventing 652 0 500 1000 1500 2000 2500 3000 + +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> 1.4 +> 1.6 +> Training Loss +> rank 16 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> 1.4 +> 1.6 +> rank 32 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> 1.4 +> 1.6 +> rank 64 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> 1.4 +> 1.6 +> rank 128 +> Full FT +> LoRA-MGPO +> LoRA + +Figure 2: Training loss dynamics across different rank configurations: A comparative analysis of LoRA, LoRA-MGPO, and full fine-tuning on LLaMA-2-7B with MetaMathQA. Rank ( r) and alpha ( α) follow r = α ∈{16 , 32 , 64 , 128 } with a fixed learning rate of 5e − 4.0 500 1000 1500 2000 2500 3000 + +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> Training Loss +> lr = 2e-4 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> lr = 3e-4 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.0 +> 0.2 +> 0.4 +> 0.6 +> 0.8 +> 1.0 +> 1.2 +> lr = 4e-4 +> Full FT +> LoRA-MGPO +> LoRA +> 0500 1000 1500 2000 2500 3000 +> Steps +> 0.00 +> 0.25 +> 0.50 +> 0.75 +> 1.00 +> 1.25 +> 1.50 +> 1.75 +> 2.00 +> lr = 6e-4 +> Full FT +> LoRA-MGPO +> LoRA + +Figure 3: Learning rate sensitivity analysis: A comparison of training loss for LoRA, LoRA-MGPO, and full fine-tuning on LLaMA-2-7B with MetaMathQA. The analysis spans learning rates {2e − 4, 3e − 4, 4e − 4, 6e − 4},with rank ( r) and alpha ( α) fixed at 128. Full FT + +> LoRA (Baseline) +> LoRA + MGPO +> LoRA-MGPO (Full) +> 0 +> 10 +> 20 +> 30 +> 40 +> 50 +> 60 + +(a) Performance on NLG Tasks + +> MT-Bench +> GSM8K +> HumanEval +> 5.30 5.61 5.69 6.27 +> 59.36 +> 42.08 +> 54.12 54.56 +> 35.31 +> 14.76 +> 20.43 21.02 +> Full FT +> LoRA (Baseline) +> LoRA + MGPO +> LoRA-MGPO (Full) +> 82 +> 84 +> 86 +> 88 +> 90 + +(b) Average Performance on GLUE + +> 87.91 +> 82.08 +> 86.76 +> 88.81 +> GLUE Average + +Figure 4: Ablation study of LoRA-MGPO on NLG and NLU tasks. (a) LLaMA-2-7B performance across three NLG tasks. (b) T5-Base performance on the GLUE benchmark. "LoRA (Baseline)" refers to standard LoRA, "LoRA + MGPO" refers to an ablation with only momentum-guided perturbation, and "LoRA-MGPO (Full)" includes both momentum-guided perturbation and adaptive normalization. + +catastrophic error accumulation common in multi-step generation. While full fine-tuning still holds an edge on the reasoning tasks, our method narrows the gap and outperforms it on MT-Bench. Notably, as the LoRA rank increases from 8 to 128, the performance of LoRA-MGPO scales gracefully, validating its ability to effectively leverage a higher parameter budget while maintaining the training stability that standard LoRA often lacks at higher ranks. + +3.3 Analysis and Ablation Studies Effectiveness in Mitigating Double Descent To empirically validate LoRA-MGPO’s core claim of mitigating double descent, we conducted a con-trolled analysis of its training dynamics, focusing on the impacts of rank and learning rate. The re-sults, presented in Figure 2 and Figure 3, offer compelling visual evidence of our method’s sta-bility. Figure 2 illustrates that as the LoRA rank + +r increases, the double descent phenomenon in standard LoRA becomes progressively more se-653 Table 3: Comparison of computational efficiency and performance across LoRA, LoRA-MGPO, and Full FT methods, trained for one epoch on the WizardLM dataset using LLaMA-2-7B. + +Method #Params Memory Cost Training Time MT-Bench GSM8K HumanEval Full FT 6738M >96 GB - 5.30 ±0.11 59 .36 ±0.85 35 .31 ±2.13 + +LoRA 320M 81.73 GB 5h 48min 5.61 ±0.10 42 .08 ±0.04 14 .76 ±0.17 + +LoRA-MGPO 320M 90.56 GB 6h 52min 6.27 ±0.12 54 .56 ±0.44 21 .02 ±0.39 + +Table 4: Ablation study of LoRA-MGPO vs. random noise perturbation on three NLG benchmarks. Exper-iments use LLaMA-3.1-8B-Base (Dubey et al., 2024) with rank r = 8 . Scores are averaged over three random seeds, with standard deviations in subscripts. Bold indi-cates the best method. + +> Method MTBench GSM8k HumanEval +> Full FT 5.88 ±0.23 73.69 ±0.28 51.63 ±1.27 +> LoRA 6.15 ±0.02 67.78 ±1.25 43.09 ±0.35 +> LoRA + Random Noise 6.43 ±0.26 68.05 ±1.12 42.92 ±0.41 +> LoRA-MGPO 7.51 ±0.07 70.23 ±1.08 45.13 ±0.63 + +vere, exhibiting a sharp rebound at r = 128 . In stark contrast, LoRA-MGPO’s loss curve remains smooth and monotonically decreasing across all ranks. Similarly, Figure 3 shows that while higher learning rates induce significant oscillations in stan-dard LoRA, LoRA-MGPO maintains a stable con-vergence path. These findings provide strong empir-ical evidence that our method effectively stabilizes fine-tuning and potentially broadens the effective learning rate window. + +Ablation Study To rigorously dissect the indi-vidual and combined contributions of our method’s two key components—Momentum-Guided Pertur-bation (MGPO) and Adaptive Perturbation Normal-ization (APN)—we conducted a detailed ablation study, with results shown in Figure 4. The findings clearly validate our design choices. The first abla-tion step, labeled ‘LoRA + MGPO‘ , applies only the MGPO component and yields a substantial per-formance lift over the vanilla ‘LoRA (Baseline)‘ .On the NLU task suite, for instance, this single component boosts the average score from 82.08 to 86.76, demonstrating that the core strategy of using momentum to guide perturbations towards flatter loss regions is fundamentally effective. However, the full potential is unlocked when introducing APN. Our complete model, labeled + +‘LoRA-MGPO (Full)‘ , combines both compo-nents and achieves the final NLU score of 88.81. The significant improvement from 86.76 to 88.81 underscores the critical role of adaptive normal-ization. It suggests that while MGPO provides a stable perturbation direction , its effectiveness is maximized only when the perturbation magnitude + +is dynamically scaled in response to the gradient landscape. The consistent superiority of the full model across all NLU and NLG tasks confirms that these two components are not merely additive but work in synergy, fulfilling the design goals of our framework. + +Comparison with Random Noise Perturbation + +To further validate that our performance gains stem from a principled optimization strategy rather than simple regularization, we compared LoRA-MGPO to LoRA augmented with undirected, isotropic ran-dom noise. The results in Table 4 are revealing: adding random noise provides only inconsistent and marginal benefits, and can even be detrimen-tal in some cases (e.g., HumanEval). In contrast, LoRA-MGPO yields consistent and significant im-provements across all tasks. This disparity highlights a fundamental differ-ence in mechanism. Random noise acts as a general regularizer by pushing parameters out of their im-mediate trajectory, which can occasionally help escape sharp minima by chance. However, the di-rection is arbitrary and uncorrelated with the loss landscape’s structure. Our momentum-guided per-turbation, conversely, is informed . It leverages the recent history of the optimization path—a strong indicator of relevant high-curvature directions—to perform a targeted exploration. This principled approach makes the search for flat minima non-stochastic and significantly more effective and reli-able than undirected noise injection. + +Computational Cost Analysis Finally, we an-alyzed the practical overhead of our method (Ta-ble 3). As expected, LoRA-MGPO operates with the same minimal number of trainable parameters as standard LoRA, making it vastly more memory-efficient than Full FT. In terms of training time, LoRA-MGPO introduces a modest and acceptable 654 overhead compared to vanilla LoRA (6h 52m vs. 5h 48m in our NLG setup). Given the significant performance improvements it delivers, this analy-sis confirms that LoRA-MGPO presents a highly favorable trade-off between computational cost and model performance, underscoring its practical via-bility. + +4 Related Work + +Parameter-Efficient Fine-Tuning (PEFT) The prohibitive computational and storage costs of full-parameter fine-tuning (Howard and Ruder, 2018; Devlin, 2018) have spurred the development of PEFT techniques for adapting large language mod-els (Houlsby et al., 2019; Ding et al., 2023). By selectively updating a small subset of parameters, PEFT methods can achieve performance competi-tive with full fine-tuning while being significantly more efficient (Han et al., 2024). Among the diverse PEFT strategies, Low-Rank Adaptation (LoRA) (Hu et al., 2021) has gained prominence for its simplicity and effectiveness. Recent works have enhanced LoRA along several directions. One line of work introduces architectural modifications ; for instance, DoRA (Liu et al., 2024) integrates learn-able magnitude vectors, while AdaLoRA (Zhang et al., 2023) dynamically allocates rank budgets. Another direction focuses on improving the train-ing process and initialization , such as adjusting scaling factors in rsLoRA (Kalajdzievski, 2023), us-ing separate learning rates in LoRA+ (Hayou et al., 2024), or refining initialization with PiSSA (Meng et al., 2024) and NLoRA (Guo et al., 2025). A third direction aims to improve the quality of the param-eter updates, for instance by alleviating training biases with BA-LoRA (Chang et al., 2024a) or by more closely aligning LoRA’s gradients with those of full fine-tuning, as seen in LoRA-GA (Wang et al., 2024a) and LoRA-Pro (Wang et al., 2024b). Additional work has further explored LoRA’s ap-plication in multi-task learning, such as (Liu et al., 2025b,a). Distinct from these approaches, our work focuses directly on the underlying optimization dy-namics. Rather than altering LoRA’s architecture or mimicking full fine-tuning gradients, we intro-duce a novel training framework to stabilize the optimization process itself. + +Optimization Stability in PEFT The training stability of PEFT methods, particularly LoRA, is a critical concern. Empirical studies have revealed that as LoRA’s rank increases, performance can de-grade after an initial improvement, a behavior anal-ogous to the double descent phenomenon (Belkin et al., 2019; Nakkiran et al., 2019). This insta-bility highlights the challenge of navigating high-dimensional and non-convex loss landscapes dur-ing fine-tuning. To promote smoother optimization and find flatter minima, Sharpness-Aware Mini-mization (SAM) (Foret et al., 2020) has been influ-ential. However, its requirement for dual gradient computations imposes a significant computational burden (Becker et al., 2024; Li et al., 2024b). More recent work has explored more efficient directional perturbation strategies. Momentum-guided meth-ods, for example, reuse optimizer momentum to avoid the extra gradient step, reducing computa-tional cost without sacrificing the directional guid-ance (Becker et al., 2024). Other techniques, such as applying an exponential moving average (EMA) to model weights, also contribute to stability by smoothing the trajectory of parameter updates (Wang et al., 2021). While these components— efficient perturbation and smoothing—are individu-ally effective, they are typically studied in isolation. This leaves a clear gap for a unified framework that synergistically combines these strategies to enhance both the efficiency and stability of PEFT. Our work, LoRA-MGPO, is designed to fill this gap. + +5 Conclusion + +In this work, we addressed the double descent phe-nomenon in Low-Rank Adaptation (LoRA), an in-stability that can affect the fine-tuning of large lan-guage models. We proposed LoRA-MGPO , an optimization framework that integrates Momentum-Guided Perturbation Optimization (MGPO). This method aims to find flatter minima by reusing optimizer momentum to guide weight perturba-tions, combined with an adaptive normalization scheme to improve robustness. Our experimental results across a range of natural language under-standing (NLU) and natural language generation (NLG) tasks show that LoRA-MGPO provides im-proved performance over standard LoRA and other common PEFT baselines. This improvement is re-flected in more stable convergence trajectories and reduced training instability. LoRA-MGPO offers a practical approach to overcoming some of the op-timization challenges in LoRA while maintaining its parameter efficiency. Future research may ex-plore extending this framework to other parameter-655 efficient methods or adapting it for different do-mains, such as vision and speech. + +Limitations + +First, LoRA-MGPO’s use of momentum vectors for perturbation directions assumes relatively stable optimizer dynamics, which might limit its effective-ness during early training stages or in the presence of highly non-stationary gradient conditions. Sec-ond, while the adaptive perturbation normalization via EMA-smoothed gradients improves robustness, its performance may be sensitive to sudden changes in gradient magnitude distributions, potentially re-quiring adjustments to the smoothing hyperparam-eters depending on the specific task. + +Ethics Statement + +Our research focuses on LoRA-MGPO, a general-purpose optimization algorithm designed to im-prove the stability of parameter-efficient fine-tuning (PEFT). The experiments use publicly avail-able, pre-trained models (LLaMA-2-7B, T5-base) and standard academic benchmarks. We acknowl-edge that these foundational models may inherit and potentially amplify societal biases present in their training data. The primary goal of this work is to provide a more reliable and resource-efficient tool for adapting and studying such models within the research community. By enhancing PEFT tech-niques, our work contributes to broader efforts aimed at reducing the computational costs involved in large-scale model adaptation. + +Acknowledgments + +This work is supported by the National Key Research and Development Program of China (No.2023YFF0905400), the National Natural Sci-ence Foundation of China (No.U2341229) and the Reform Commission Foundation of Jilin Province (No.2024C003). + +References + +Marlon Becker, Frederick Altrock, and Benjamin Risse. 2024. Momentum-sam: Sharpness aware minimiza-tion without computational overhead. arXiv preprint arXiv:2401.12033 .Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. 2019. Reconciling modern machine-learning practice and the classical bias–variance trade-off. Proceedings of the National Academy of Sciences , 116(32):15849–15854. Yupeng Chang, Yi Chang, and Yuan Wu. 2024a. Ba-lora: Bias-alleviating low-rank adaptation to mitigate catastrophic inheritance in large language models. + +arXiv preprint arXiv:2408.04556 .Yupeng Chang, Xu Wang, Jindong Wang, Yuan Wu, Linyi Yang, Kaijie Zhu, Hao Chen, Xiaoyuan Yi, Cunxiang Wang, Yidong Wang, et al. 2024b. A sur-vey on evaluation of large language models. ACM Transactions on Intelligent Systems and Technology ,15(3):1–45. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Ka-plan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. 2021. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 .Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 .Jacob Devlin. 2018. Bert: Pre-training of deep bidi-rectional transformers for language understanding. + +arXiv preprint arXiv:1810.04805 .Ning Ding, Yujia Qin, Guang Yang, Fuchao Wei, Zonghan Yang, Yusheng Su, Shengding Hu, Yulin Chen, Chi-Min Chan, Weize Chen, et al. 2023. Parameter-efficient fine-tuning of large-scale pre-trained language models. Nature Machine Intelli-gence , 5(3):220–235. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 .Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. 2020. Sharpness-aware min-imization for efficiently improving generalization. + +arXiv preprint arXiv:2010.01412 .Zihao Fu, Haoran Yang, Anthony Man-Cho So, Wai Lam, Lidong Bing, and Nigel Collier. 2023. On the effectiveness of parameter-efficient fine-tuning. In Proceedings of the AAAI conference on artificial intelligence , volume 37, pages 12799–12807. Chenlu Guo, Yuan Wu, and Yi Chang. 2025. Nlora: Nystr \" om-initiated low-rank adaptation for large language models. arXiv preprint arXiv:2502.14482 .Zeyu Han, Chao Gao, Jinyang Liu, Jeff Zhang, and Sai Qian Zhang. 2024. Parameter-efficient fine-tuning for large models: A comprehensive survey. + +arXiv preprint arXiv:2403.14608 .Soufiane Hayou, Nikhil Ghosh, and Bin Yu. 2024. Lora+: Efficient low rank adaptation of large models. + +Preprint , arXiv:2402.12354. 656 Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. 2019. Parameter-efficient transfer learning for nlp. In In-ternational conference on machine learning , pages 2790–2799. PMLR. Jeremy Howard and Sebastian Ruder. 2018. Universal language model fine-tuning for text classification. + +arXiv preprint arXiv:1801.06146 .Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021. Lora: Low-rank adap-tation of large language models. arXiv preprint arXiv:2106.09685 .Damjan Kalajdzievski. 2023. A rank stabilization scaling factor for fine-tuning with lora. Preprint ,arXiv:2312.03732. Brian Lester, Rami Al-Rfou, and Noah Constant. 2021. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 .Tao Li, Zhengbao He, Yujun Li, Yasheng Wang, Lifeng Shang, and Xiaolin Huang. 2024a. Flat-lora: Low-rank adaption over a flat loss landscape. arXiv preprint arXiv:2409.14396 .Tao Li, Qinghua Tao, Weihao Yan, Zehao Lei, Yingwen Wu, Kun Fang, Mingzhen He, and Xiaolin Huang. 2024b. Revisiting random weight perturbation for efficiently improving generalization. arXiv preprint arXiv:2404.00357 .Jinda Liu, Yi Chang, and Yuan Wu. 2025a. R-lora: Random initialization of multi-head lora for multi-task learning. arXiv preprint arXiv:2502.15455 .Jinda Liu, Bo Cheng, Yi Chang, and Yuan Wu. 2025b. Align, don’t divide: Revisiting the lora architecture in multi-task learning. arXiv preprint arXiv:2508.05078 .Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, and Min-Hung Chen. 2024. Dora: Weight-decomposed low-rank adaptation. Preprint ,arXiv:2402.09353. Ilya Loshchilov and Frank Hutter. 2019. Decoupled weight decay regularization. In ICLR .Fanxu Meng, Zhaohui Wang, and Muhan Zhang. 2024. Pissa: Principal singular values and singular vec-tors adaptation of large language models. Preprint ,arXiv:2404.02948. Preetum Nakkiran, Gal Kaplun, Yamini Bansal, Tristan Yang, Boaz Barak, and Ilya Sutskever. 2019. Deep double descent: Where bigger models and more data hurt. Preprint , arXiv:1912.02292. Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. 2020. Exploring the lim-its of transfer learning with a unified text-to-text transformer. Journal of machine learning research ,21(140):1–67. Hugo Touvron, Louis Martin, Kevin Stone, Peter Al-bert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. 2023. Llama 2: Open founda-tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 .Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. 2018. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations .Shaowen Wang, Linxi Yu, and Jian Li. 2024a. Lora-ga: Low-rank adaptation with gradient approximation. + +Preprint , arXiv:2407.05000. Yizhou Wang, Yue Kang, Can Qin, Huan Wang, Yi Xu, Yulun Zhang, and Yun Fu. 2021. Rethinking adam: A twofold exponential moving average approach. arXiv preprint arXiv:2106.11514 .Zhengbo Wang, Jian Liang, Ran He, Zilei Wang, and Tieniu Tan. 2024b. Lora-pro: Are low-rank adapters properly optimized? Preprint , arXiv:2407.18242. Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. 2022. Emergent abilities of large language models. + +arXiv preprint arXiv:2206.07682 .Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, Qingwei Lin, and Daxin Jiang. 2024. Wizardlm: Empowering large pre-trained language models to follow complex instructions. In ICLR .Longhui Yu, Weisen Jiang, Han Shi, YU Jincheng, Zhengying Liu, Yu Zhang, James Kwok, Zhenguo Li, Adrian Weller, and Weiyang Liu. 2024. Metamath: Bootstrap your own mathematical questions for large language models. In ICLR .Qingru Zhang, Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. 2023. Adalora: Adap-tive budget allocation for parameter-efficient fine-tuning. Preprint , arXiv:2303.10512. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. 2024a. Judging llm-as-a-judge with mt-bench and chatbot arena. In NeurIPS .Tianyu Zheng, Ge Zhang, Tianhao Shen, Xueling Liu, Bill Yuchen Lin, Jie Fu, Wenhu Chen, and Xiang Yue. 657 2024b. OpenCodeInterpreter: Integrating code gen-eration with execution and refinement. In Findings of ACL . + +# Appendix + +Contents + +A Models and Datasets 11 + +A.1 Details of Models . . . . . . . . . 11 A.2 Details of Datasets . . . . . . . . 11 + +B Baselines and Implementation 11 + +B.1 Baseline Methods . . . . . . . . . 11 B.2 Implementation Details . . . . . . 12 B.3 Hyperparameter Settings for Base-lines . . . . . . . . . . . . . . . . 12 + +A Models and Datasets + +A.1 Details of Models + +In this work, we primarily utilize two pre-trained language models: LLaMA-2-7B and T5-base. • LLaMA-2-7B : A 7-billion parameter, decoder-only transformer model from the LLaMA-2 series, primarily used for genera-tion tasks. More details are available at its Hugging Face repository *.• T5-base : A 220-million parameter encoder-decoder transformer model, widely used for a variety of natural language understanding tasks. More details are available at its Hug-ging Face repository †.Our experiments were conducted using the im-plementations of these models provided by the Hug-ging Face Transformers library. + +A.2 Details of Datasets + +Table 5 summarizes the GLUE benchmark datasets (Wang et al., 2018). For our Natural Language Gen-eration (NLG) experiments, we used the following evaluation metrics: Accuracy for GSM8K; Pass@1 for HumanEval; and a score based on GPT-4 evalu-ation for MT-Bench. + +B Baselines and Implementation + +B.1 Baseline Methods + +Our study includes several baseline methods for a comprehensive comparison. Full Fine-Tuning + +serves as a strong performance benchmark. Vanilla + +> *https://huggingface.co/meta-llama/LLaMA-2-7B +> †https://huggingface.co/t5-base + +658 Table 5: GLUE Benchmark Datasets and Evaluation Metrics + +> Dataset Task Type Classes Train Examples Metric Description +> CoLA Acceptability 28.5k Matthews Corr. Grammatical acceptability SST-2 Sentiment 267k Accuracy Sentiment analysis MRPC Paraphrase 23.7k Accuracy/F1 Paraphrase detection MNLI NLI 3393k Accuracy Multi-genre NLI QNLI NLI/QA 2108k Accuracy QA/NLI converted from SQuAD + +LoRA (Hu et al., 2021) is our primary point of com-parison from the PEFT literature. We also compare against LoRA variants that introduce structural modifications (DoRA (Liu et al., 2024), AdaLoRA (Zhang et al., 2023)) and those that refine the training process or initialization (rsLoRA (Kala-jdzievski, 2023), LoRA+ (Hayou et al., 2024), PiSSA (Meng et al., 2024)). Finally, we include methods focused on gradient alignment (LoRA-GA (Wang et al., 2024a), LoRA-Pro (Wang et al., 2024b)). + +B.2 Implementation Details LoRA Configuration. As stated in the main text, LoRA adapters were applied to all linear layers within the transformer blocks for both LLaMA-2-7B and T5-base models. + +Initialization of MGPO. The implementation of our method requires an initial state for the momen-tum vector and the adaptive normalization factor. Following standard optimizer practice, the momen-tum ‘ m‘ is initialized to zeros. The adaptive nor-malization factor ‘ ¯g‘ is initialized using the L2-norm of the gradient computed in the first training step. + +Hyperparameters. Our method introduces two primary hyperparameters: the perturbation radius ‘ρ‘ and the EMA decay rate ‘ β‘. ‘ ρ‘ controls the magnitude of the weight perturbation, influencing the search for flatter minima. ‘ β‘ controls the tem-poral smoothing window for the adaptive normal-ization. The values used in our main experiments were effective across the evaluated tasks, as evi-denced by the strong performance reported in Sec-tion 3. + +B.3 Hyperparameter Settings for Baselines + +To ensure a fair and robust comparison, we adhered to the hyperparameter settings recommended in the original papers or official codebases of our baseline methods wherever possible. General settings, such as the learning rate schedule and batch size, were kept consistent across all methods as described in Section 3. Key method-specific hyperparameters are detailed below. • DoRA (Liu et al., 2024): We utilized the offi-cial implementation provided by the authors, maintaining its default configuration for the magnitude and directional components. • AdaLoRA (Zhang et al., 2023): We followed the setup from the original paper, with the rank budget dynamically allocated starting from a higher initial rank and pruned during training. • LoRA+ (Hayou et al., 2024): Following the authors’ recommendation, the learning rate for the LoRA matrix A was set to our default value ( 1 × 10 −4 for NLU, 2 × 10 −5 for NLG), while the learning rate for matrix B was set 16 times higher. • LoRA-GA and LoRA-Pro (Wang et al., 2024a,b): For these methods focused on gra-dient alignment, we used the hyperparameter settings as specified in their respective papers and official implementations to ensure a faith-ful comparison. For all other baselines, we used their standard, publicly available implementations without modifi-cation to their core components. 659