TRAM: Bridging Trust Regions and Sharpness Aware Minimization
Authors: Tom Sherborne
Naomi Saphra
Pradeep Dasigi
Hao Peng
Venue: ICLR 2024 spotlight (2024)
Keywords: sharpness-aware minimization
sam
trust region
optimization
cross-lingual transfer
language modeling
Primary area: optimization
TL;DR: We propose a trust region motivated variant of SAM which scales learning with a trust region constraint to jointly optimize for low sharpness parameters and low curvature representations.

Abstract: Sharpness-aware minimization (SAM) reports improving domain generalization by
reducing the loss surface curvature in the parameter space. However,
generalization during _fine-tuning_ is often more dependent on the
transferability of _representations_ in the function space. Trust-region
methods (TR) target this goal by regularizing representation curvature to reduce
catastrophic forgetting of pre-trained task-agnostic information while adopting
task-specific skills. We consider unifying these strategies for low curvature in
both parameter space and function space to improve out-of-domain (OOD)
generalization. We propose **Trust Region Aware Minimization** (TRAM), a
SAM algorithm fine-tuning for low parameter sharpness and smooth, informative
representations preserving pre-trained structure. TRAM uses a trust region bound
to inform the SAM adversarial neighborhood, introducing an awareness of function
curvature within optimization for flatter minima. We empirically validate TRAM
in vision (cross-dataset adaptation) and text (OOD language modeling, zero-shot
cross-lingual transfer) tasks where robust domain transfer and representation
generality are critical. TRAM outperforms SAM- and TR-based optimization across
all tasks, notably surpassing competing methods for hard transfer between
_anticorrelated_ domains. TRAM establishes a novel standard in
fine-tuning for domain-generalizable models with minimal additional computation
over previous sharpness-aware methods.

Published as a conference paper at ICLR 2024
TRAM: BRIDGING TRUST REGIONS AND SHARPNESS AWARE MINIMIZATION
Tom Sherborne Naomi Saphra² Pradeep Dasigi"
Hao Peng
1 University of Edinburgh *eempner Institute, Harvard University
3 Allen Institute for Al
Universiit Illinois Urbana- Champpaan
tom. www.sshrrrreeeeeee ac uk, nsaphra@fas harvard.
pradeepdeal wew.leenni... org, haoperg@illinois..
ABSTRACT
Sharpness-aware minimization (SAM) reports improving domain generalization by reducing the loss surface curvature in the parameter space. However, gen- eralization during fine-tuning is often more dependent on the transferability of representations in the function space. Trust-region methods (TR) target this goal by regularizing representation curvature reduce catastrophic forgetting of pre- trained task-agnostic information while adopting task-specific skills. We consider unifying these strategies for low curvature in both parameter space and function space to improve out-of-domain (OOD) generalization. We propose Trust Region Aware Minimization (TRAM), SAM algorithm fine-tuning for low parameter sharpness and smooth, informative representations preserving pre-trained structure. TRAM uses a trust region bound to inform the SAM adversarial neighborhood, introducing an awareness of function curvature within optimization for flatter min- ima. We empirically validate TRAM in vision (cross- -atosed adaptation) and text (OOD language modeling, zero-shot cross- lons-linguual transfer) tasks where robust domain transfer and representation generality are critical. TRAM outperforms SAM- and TR-based optimization across all tasks, notably surpassing competing methods for hard transfer between anticorrelated domains. TRAM establishes novel standard in fine-tuning for domain domain-generalinn generalizzble models with minimal additional computation over previous sharpness-awar methods.
INTRODUCTION
Neural model training requires navigating over complex, non-convex loss surface (Frankle 2020) towards good local minimum. Studying loss surfaces and training dynamics has led to many algorithmic advances (Izmailov et al. 2018 Foret et al. 2022 Chen et al. (023) and regularization schemes (Srivastava et al.] 2014; loffe & Szegedy 2015) to improve optimization. One such strategy is exploit association between generalization and flat minima, defined by Hochreiter & Schmidhuber (994) as "region[s] in weight space with the property that each weight vector from that region has [a] similar small error". Intuitively, a flatter, or less sharp (Keskar et al.) (017), minimum will generalize better, as the loss function will be non-increasing under distribution shift. Recent work has developed family of sharpness aware minimization (SAM) algorithms targeting flat minima by jointly minimizing a worst-case generalization bound and local parameter sharpness (Foret et al.
2021 K Kwon al. 202T Kim et al. 2022 Möllenhoff & Khan 2023).
While flat minima methods report widespread improvement over conventional optimizers (Kaddour ct al. 2022), we argue that they are not fully connected the modern fine-tuning paradigm, wherein task-specific model inherits parameters from a pre-trained model instead of being trained from scratch (Wang et al. 2019 Liang al., 2020). In these settings, focusing on local properties of the loss landscape (e.g., sharpness) may fail by suboptimally exploiting useful generic task agnossgc structures within pre-trained representations.
this work, we propose to combine sharpness-awan aware minimization with the robust transfer of pre-trained information (in representation space) for fine-tuning scenarios requiring -atoo-ddistriuutin knowledge for successful adaptation.
This work done while Tom Sherbome and Hao Peng were the Allen Institute for AT.

---

Published as a conference paper at ICLR 2024
Trust region
82
SAM Ascent
Existing methods to improve leveraging pre- trained structure during fine-tuning include trust
(f(z)
far))
<E
ASAM
region regularization (Schulman et al., [2015]
f(x)
TRAM scoled by
Jiang et al. 20200 Aghajanyan et al., 2021 ) or
adversarial perturbation (Zhu et al. 2020; He
crst region size
et al. 2021) D. These methods focus on the cur-
Function space
21
Parameter space
vature of the function itself e.g., by encouraging smooth local changes in representations, The intuition is that lower representation curvature during fine-tuning limits a function from catas- trophically forgetting (French) 1999] inter alia) useful information from pre-training. This rep- resentation smoothing approach contrasts with SAM-style optimization for parameter smooth- ness. Both perspectives show empirical improve- ment in downstream tasks (Aghajanyan al.) 2021 Bahri et al.) 2022), but a fusion of these strategies is presently under-explored.
Figure I: TRAM introduces an awareness of func- tion curvature (i.e., the trust region) into sharpness- aware minimization. (left) TRAM estimates the size of the trust region, d, around (ax) in green. (right) the loss contour in parameter space follow- ing Kwon al. (02I) where blue is the typi- cal loss; red is the maximized worst-case loss for ASAM; and green is the maximized loss within the subdomain constrained for function smoothness.
To this end, we propose TRAM: Trust Region Aware Minimization, a fine-tuning algorithm for out-of-distribution generalization combining the success of both sharpness-aware and trust region opti- mization. TRAM uses a trust region bound to inform the SAM adversarial neighborhood. introducing an awareness of function curvature within optimization for flatter minima. The resulting algorithm yields low-sharpness parameters and improved adaptation of pre-traines models to downstream tasks. To illustrate TRAM's s advantage over strong baselines in retaining generic representations, we focus on distribution transfer challenges within Transformer-base models. Our contributions are:
We propose a new optimization algorithm: Trust Region Aware Minimization integrates representation smoothing regularization into sharpness-awan minimization. We propose and contrast multiple variants of TRAM based on differing perspectives on trust region estimation and efficiency trade-offs (Section 3
We highlight that TRAM empirically improves generalization for multiple nut-of-distribution adaptation tasks across vision and natural language: cross- crass-dataset adaptation for image clas- sification, cross-domain language modeling and zero-shot cross-lingual transfer (Section 4
We analyze how TRAM limits catastrophic forgetting and optimizes flatter minima to improve fine-tuning. By characterizing major and minor distribution shifts, we identify how TRAM outperforms the trend in anticorrelated generalization scenarios. Our analysis verifies that TRAM optimizes a smoother loss surface for both in-domain and out-of-domain distributions. TRAM also improves representation similarity between seen and unseen distributions to improve cross-domain classification (Section
2 BACKGROUND
We describe SAM and trust region optimization, highlighting how these approaches have similar goals. Our motivation for TRAM is the unifying features of each approach outlined in Table
Notation: We consider function f X + Y parameterized by weights 0 and evaluated by loss func-
tion E: YXY - R4. The expected loss on true distribution D L-D (0) B( [e (v. f (x; 0)) and the empirical estimate is Lg = 1 Es E (yi, ff ()) sampling n training samples, S = {(ris yi ) from D. Functional distance on model outputs is measured by the Kullback-Leibler divergence DxL. (pllz) between target p and estimate q. We describe successful domain transfer to distribution D' non-inereasing loss for sample S" D'.
Sharpness-Aware Minimization: Foret et al. (021) define local sharpness as
regularizes parameter magnitude to minimize this sharpness metric jointly with loss within local parameter neighborhood p.
max 1p Ls ee) Ls (0). The SAM objective (Equation [
Coode atgii ott.attt.c...ccooheeerrerrrrrrrr
2

---

Published as a conference paper at ICLR 2024
A S SAM = min max Ls (0 + 2 (1) 1ee
CAAAM = LL
(2)
L
This min-max optimization problem is solved in alternating stages. Initial ascent perturbs parameters 0 to 0 + G, where is a perturbation maximizing loss (to minimize local sharpness). The feasible region for perturbation is a Euclidean spherical neighborhood with radius p > 0. Successive descent evaluates gradients at 0 + e for gradient descent at 0 using the local worst-case loss.
The optimal c, the perturbation for worst-case loss within the p-ball, the source of ongoing debate. Forct et al. (2021) express closed-form solution setting as the radius scaled by the normalized gradient. Kwon et al. (2021) propose Adaptive SAM (ASAM) to improve SAM with invariance to the loss scaling. For ASAM, each parameter within 0 is perturbed by p scaled by parameter gradient and the parameter norm (Equation? 2) TRAM follows SAM in setting e with scale invariance and also augments such that the update & respects maximum divergence in the function space.
Trust Region Regularization:
Trust region regularization encourages low curvature during opti- mization by regularizing the function output distribution with respect to previous step's distribution. fine-tuned model with high curvature (i.e., distance) to pre-trained representations may struggle to connect task-specific knowledge with novel domains. This approach proves successful in penalizing large policy updates in reinforcement learning (Schulman et al. 2013)) encouraging local smoothness to adversarial perturbation (Jiang et al. 2020) and minimizing catastrophic forgetting for domain transfer (Aghajanyan et al. 2021).
Equation B defines the objective under Trust Region Policy Optimization (TRPO; Schulman et al.
2015) constraining loss, T.s. with a regularization term dg. TRPO idealizes smoothness in / (*r by regularizing local function similarity to the previous iterate. The update is constrained such that changes in probability density, ps (-I 0) are no larger than some measured by divergence
: y x y - R4. There are several ways of defining d-we consider options in Equations to
min Ls (0) subject dg S
Equation Destimates the trust region as the KL divergence between predictive distributions the previous and current step. Intuitively, penalizing divergence from prior steps encourages the function to stay "close" to the previous distribution i.e., within the trust region of equivalent output. Across training, de encourages small updates with low curvature between fine-tuned and pre-trained models.
de (0,-1, 0,) = Ez=D [DEL (py (-1 ær, (,,, Ily (-1 x, 0,))
Equation 5 provides the penalty from R3F (Aghajanyan et al. 2021) where d2 estimates the trust region by sampling from inputs under parametric noise. This penalizes the divergence between Py (-l x, 8,) and Py (-1 x + z, (,) for some gero-mean noise z ~ A' (0, d²). R3F proposes that sampling estimates the trust region by simulating distribution shift in Py corresponding to perturbed x + 2. This encourages similarity to a neighborhood around f (xx f) with equivalent output. Either approach estimates the permissible distance for update in 0 without increasing local representation curvature. We focus on trust region methods to improve generalization across distributions via improved leveraging of pre-trained structure (Tiang et al. 2020, inter alia).
d (x + 2,)) Ege [DLL (py (+) 2, 0) Ipyy (-1 x. 0))
(5)
Comparison: SAM, TRPO, and R3F have similar goals in searching for generalizable solutions while appearing superficially distinct. We compare the broad motivations and qualities of methods in Table highlighting both perspectives optimize for smoothness in different spaces.
SAM minimizes sharpness within neighborhood 0 set by sealar parameter p. Trust region regular- ization penalizes loss by scalar distance do or dz. We hypothesize that this regularization can inform the size of the SAM neighborhood. Can we jointly minimize sharpness and penalize high curvature in representations? Considering the sharpness objective (maxjell:$p Ls (0 + Ls (0) we consider this can also satisfy the Equation 3 constraint of de or d, < Our intuition here is to minimize
parameter sharpness (i.e., SAM) only within update promoting low representation curvature. Combining the features of these solutions could improve generalization to unseen distributions during fine-tuning.
3

---

Published as a conference paper at ICLR 2024
Table 1:
Comparison between SAM-style, trust region and TRAM learning. SAM optimizes param- eters for low sharpness, trust region methods optimize for low-curvature representations. TRAM combines these strategies to bound SAM-sttyle learning within a trust region neighborhood.
Goal
Distance
Domain
Grudient
Forward//Bakwant
SAM-style
Low-sharpness
Equation
p-hall
Vig at +
2-,24
Trust region
Low-curvature f (g/12,0)
dg or d2 do or dl
DxL over Distance do of da-ball
Ls + or
2-h1e
TRAM
Both
Equation
7
Vhg 0
3 34,26
TRAM: TRUST REGION AWARE MINIMIZATION
We consider methods improving generalization by encouraging low-sharpness parameters and task transfer by encouraging low curvature representation space. We introduce TRAM: Trust Region Awan Aware Minimization unifying sharpness-awar and trust region optimization. Kim et al. (0022) raise that the hyperparameter defining the ascent neighborhood in SAM is an "ad hoc" scaling with little relationship to the loss landscape or parameter geometry. We propose to instead define the ascent region by a trust region in representation space.
TRAM substitutes in Equation 2 with the trust region metric, d Y y - Rt. as defined in Section 2 We estimate trust regions using the divergence from prior model distribution (dg. Equa- tion 4 or divergence from the current distribution under parametric noise (dz, Equation 5). TRAM constrains the maximization domain for ascent (i.c., 0 + 0 + E) to the parameter corollary for the trust region i.e., max d substituted within Equation I TRAM perturbs 0 with loss perturbation only within the parameter neighborhood constrained for low representation curvature. This introduces function curvature awareness within TRAM in addition to the sharpness- awareness objective for flatter minima. In contrast, the maximization region, P in SAM/ASAM has no sensitivity to function curvature. We build TRAM on ASAM, and not SAM, after observing strictly better performance in our preliminary experiments.
V LTRAM (0) = 0Ls
d L (0.)
00
(6)
TRAM LOVLS (0,)l2
10=0+CTAMM
The gradient descent update in TRAM Equation where CRAM is solved as Equation by direct substitution of in ASAM. Algorithm Appendix B.6 details the full training algorithm for TRAM based on the SAM-style min-max optimization routine. TRAM does not require tuning P hyperparameter for stable training. TRAM using de introduces no new hyperparameters, and using dy requires only tuning a for additive noise z. We hypothesize that TRAM jointly minimizes parameter sharpness and representation curvature to minimize catastrophic forgetting of pre-trained structure. Our results in Section Tempiricllly validate this hypothesis.
The geometric interpretation of TRAM frames the maximization domain defined by as a subdomain of the p-radius Euclidean ball defined in ASAM. Whereas ASAM defines fixed radius by P each step, TRAM instead uses the nonzero d radius constraining the maximization domain to additionally satisfy the trust region constraint outlined in Equations 3 to Foret et al. (2021] Theorem 2) defines a PAC-Bayesial generalization bound for SAM on LD assuming p >> 0. Kwon et al. [202 Theorem 3) identify a similarly valid bound when considering the We assume d S p for similar asymptotic behavior for ETRAM to ASAM- We infer that TRAM inherits the existing generalization bound of ASAM for any p > directly substituted for d i.. TRAM is a subsolution of ASAM. We can constrain d such that maxp_, do (0-e; 0t) S p max; d, (x + 2, S p to enforce this bound d E (0, p). However, we empirically observe this constraint is satisfied for the optimal setting of p in ASAM.
Connection to ASAM:
norm-adaptive scaling on ÄSAM as in Equation 2
Improving Efficiency with TRAM-Fisher: Kim et al.] (2022) propose an alternative to SAM removing the Euclidean assumption for parameter geometry. Fisher SAM (FSAM) instead exploits the statistical manifold induced by the Fisher Information metric of predictive distribution of the
function, py (yl xx 9) (Amari 1998) to set c. This measures statistical divergence between @ and 0 + resulting in CEAAM in Equation & defining an ellipsoid around 0 scaled by the Fisher Information matrix, F (0). F (8) is prohibitively expensive at scale and is approximated with Equation 9 the diagonal of the squared gradient for each batch B.
4

---

Published as a conference paper at ICLR 2024
Table 2: We propose four variants TRAM based on different trust region estimations. TRAM-8,11 uses divergence against the previous step; TRAM-60 is a simplifying heuristic of this divergence against the pre-trained model only; and TRAM-Z uses noised imput divergence, dz. TRAM-Fisher extends FSAM by measuring the Fisher Information metric around the trust region.
Variant
Trust region measurement
Domain
Forward/Backwant
TRAM-0:-1
dø (fz-1.0r) 0i)
Equation 7
dg-ball
+, +
TRAM-00
dg (fo, 0z)
Equation
dg-ball
3 -, 2 <-
TRAM-2
dg (x + 2,z) = N (0,
Equation
dg-ball
3 2 <-
TRAM-Fisher
F (z 2;0),
(0,02)
Equation 8
F-ellipse
2-,,2
CPSAM F(0) VLL VLS
2
(8)
(0) Diag
1 B
(log Pr 0))
(9)
VLSFF()) LL
We propose TRAM- as an efficient variant of TRAM inspired by Fisher SAM. Where FSAM measures the Fisher Information geometry of 9 under input x, we instead sample the geometry of 9 under the trust region estimation from x + z. Our proposal minimal: replace p (y, 0) with p (yi 22 + 2i, 0) to estimate the Fisher Information Matrix of the trust region neighborhood
as ES- F (ar (x)) We sample parametric noise {= identically to TRAM and now scale
learning with the information geometry of the low curvature neighborhood, f (x + +2). TRAM-Fisher uses the same number of forward/backward passes as FSAM and only requires additional processing to sample and compute x + 2. TRAMM-isher matches FSAM runtime efficiency (with marginal additional operations) and performs competitively across our experiments. The full TRAM-Fisher algorithm is shown in Appendix B.6
Summary: We propose three variants of TRAM, and TRAM-Fisher, summarized in Table 2 TRAM- 0 follows TRPO (Schulman et al.) 2015) in using previous step parameters, 0,-1, to measure the trust region. We also propose simplification of TRAM-0,- estimating the trust region using dø between current 0r and pre-trained model 6. TRAM-0, improves training efficiency by removing an updating 0 state. TRAM-x follows R3F (Aghajanyan et al., (021) using noise-based trust region measurement with additional hyperparameter for sampling parametric noise. Practically, TRAM requires one additional forward pass adding marginal overhead the extant complexity of SAM-style training. Despite this additional cost, Section identifies empirical benefits to TRAM and targeted improvement to out-of-domain loss surface sharpness and cross-domain representation similarity.
We outline our datasets in Appendix A and experiment design in Appendix B for both vision and language modalities. We compare to gradient descent methods (SGD, Adam), sharpness aware methods (SAM, ASAM, FSAM), and trust region methods (TRPO, R3F, MESA) further detailed in Appendix B.2 Broadly, we investigate the hypothesis that distrr-iittii generalization improves by jointly minimizing parameter sharpness and representation curvature in the function.
RESULTS
4.1
CROSS-DATASET IMAGE CLASSIFICATION
First, we validate the performance of TRAM in a standardized setting for comparison to other SAM-style optimizers. We evaluate adapting ViT-base (Dosovitskiy al.] [02T from ImageNet pre-training to image classification fine-tuning. We follow the setup of Kim et al. 2022 Section 5.1) evaluating adaptation to CIFAR- 100 (Krizhevsky) 2009), Stanford Cars (Krause al. (013). and
Oxford Flowers (Nilsback & Zisserman 2008). Appendix B.3 details our experiment design.
Table 3 details the Top-1 accuracy results for this experiment with direct comparison to Kim et al. 2022 Table 3). The best-performing variant of TRAM (TRAM-0,- or TRAM-z) is significantly superior to the closest FSAM competitor (p < 0.01). Other variants of TRAM, TRAM-9, or TRAM- Fisher, are largely competitive with prior methods. Our observations validate the hypothesis that TRAM improves adaptation across datasets during fine-tuning for image classification. This comparison acts as sanity check and demonstrates the utility of method compared to other SAM-
5

---

Published as a conference paper at ICLR 2024
Table Cross-dataset adaptation from ImageNet CIFAR-I00, Stanford Cars and Oxford Flowers. We report Top-1 classification accuracy averaged over five runs, + the 95% confidence interval, for direct comparison Kim al. 2022).
CIFAR-100 (1)
Cars (1)
Flowers (1)
SGD
87.9710.12
92.85=0.31
94.5310.20
SAM
87.99=0.009
93.29+00.01
95.05+00.66
ASAM
87.97=0.08
93.28+0002
95.08+0.
FSAM
88.39:00133
93.42++0.01
95.26+0..3
TRAM-0,-1
88.47-10.16
93.4910.04
97.07-10.10
TRAM-0b
88.31
93.1610.07
95.53.10. 10
TRAM-2
88.78=0.01
93.32+0.11
96.34+0.05
88.02=00.18
93.12:t0.13
94.90+0.11
Table M2D2 perplexity (lower better) on Wikipedia (upper) S2ORC (lower) splits. TRAM- 0,-1 significantly improves over prior work (p < 0.01 Kolmogorov-Smirnov test). Results are grouped as: (i) optimizers; (ii) trust region methods; and (iii) TRAM variants. The leftmost column the training domain and we evaluate zero-shot perplexity ten domains unseen during fine-tuning (full details in Appendix A). ZS AvG. is the macro-average all zero-shot domains.
Wiki
Soc.
CULT.
GEN.
HEALTH.
HIST.
HUMAN.
MATIL
NAT.
PHIL.
REL
TECH.
ZS AvG.
GPT-2
27.2
27.7
27.8
24.5
29.2
28.8
28.6
29.4
27.8
27.7
28.7
28.0
Adam
24.8
26.3
26.4
23.6
27.2
27.0
27.4
27.6
26.3
25.8
27.4
26.5
SAM
24.5
25.9
26.0
26.9
26.6
26.6
27.2
25.8
25.5
27.0
26.1
ASAM
21.8
25.4
25.6
22.5
27.1
26.4
26.3
26.7
25.5
25.5
28.1
25.9
TRPO
21.8
23.0
23.3
20.7
24.0
23.7
23.8
24.0
23.1
22.8
24.1
23.3
R3F
21.8
23.0
23.3
20.7
24.0
23.7
23.8
24.0
23.1
22.8
24.1
23.3
23.1
24.0
25.4
24.9
25.1
TRAM-2
21.9
24.0
23.3
23.9
23.2
23.9
23.2
TRAM-0.-1
20.9
22.4
22.7
20.1
23.1
22.9
23.2
23.3
22.4
22.0
23.44
22.5
TRAM-00
23.9
23.3
23.9
TRAM-Fisher
22.5
23.7
24.0
21.3
24.6
24.0
24.7
24.6
23.8
23.3
24.6
23.9
S2ORC
MATH
ART
ASTRO
CONDM.
CS
NLIN.
PHIL.
PHYS.
QB:0
Ave.
GPT-2
27.6
32.4
30.9
27.9
27.6
33.7
33.5
30.9
30.6
Adam
11.4
44.2
33.9
20.1
21.2
14.7
41.9
29.5
30.8
16.9
SAM
10.5
45.3
33.2
18.7
20.3
20.0
13.7
42.4
28.3
30.2
16.1
26.8
ASAM
10.3
13.5
42.6
28.2
FSAM
10.4
45.6
33.3
18.5
20.2
19.9
13.5
42.7
28.3
30.2
15.9
26.8
TRPO
10.4
46.0
33.4
18.6
20.3
20.0
13.6
42.9
28.4
30.4
16.0
26.9
10.4
46.0
33.4
18.6
20.2
20.0
13.6
42.9
28.4
30.4
16.0
26.9
MESA
11.9
44.1
34.1
20.8
21.7
21.6
15.3
41.7
30.0
31.0
17.4
27.8
TRAM-r
10.4
44.9
33.0
18.6
20.1
19.9
13.6
42.0
28.1
30.0
15.9
26.6
TRAM-B--1
9.6
46.8
32.5
17.2
19.2
18.9
12.6
43.33
27.0
29.6
15.0
26.2
TRAM-6
10.4
33.0
20.1
13.6
42.0
28.2
30.0
TRAM Fisher
10.5
46.1
32.4
18.7
20.3
20.0
13.6
43.0
28.2
30,3
16.0
26.9
style optimizers. TRAM yields improved fine-tuned image classification models by encouraging smoothness in parameter and function space.
4.2
CROSS-DOMAN LANGUAGE MODELING
We now consider zero-shot cross-domain language modeling using the M2D2 Corpus (Reid et al. 2022) outlined in Appendix A We hypothesize that TRAM can improve domain transfer in language modeling by retaining domain-agnostic information from pre- pre-training when fine-tuning domain. We train a GPT-2 Base model (Radford et al. 2019) on the largest domain each split of M2D2 (Soc. domain 379M tokens for Wikipedia and MATH 1.4B tokens for S2ORC) and evaluate perplexity across ten domains unseen during fine-tuning. Appendix Bariiixx..4eeeee detiiis our complete experiment design.
specific
Our results in Table validate our hypothesis for TRAM in the cross tross-domain setting to improve out-of-domain language modeling fine-tuning on single domain. All TRAM variants (excluding

---

Published as a conference paper at ICLR 2024
Figure 2: Perplexity on S2ORC training domain (MATH) and zero-shot domains. We report perplexity
across: a domains correlated with MATH STEM domains (see Appendix C.T. b ART domain, and C the Philosophy (PHIL..) domain. Each figure includes linear regression trends: the blue dotted trend is for prior work and green dashed line includes all TRAM variants. Positive slope (p 0) represents correlated domains, negative slope (p < 0) represents anticorrelated domains. We report Pearson p correlation for the blue trend noting p 0.01 significance.
(a): STEM Avg. p=0.99
(b): ART p=-0.
(c): PIIL. p.-0.86
Adam
24
SAM
ASAM
43
OOD Parplexity
46
FSAM
28
TRPO
45
MESA
42
TRAM-6
22
TRAM-8--1
44
TRAM-r
TRAM-Fisher
21
Tread
11
10 12
Tread w/o TRAM
MATH Perplexity
TRAM-Fisher) perform comparably or above competitors in zero-shot transfer across both splits of M2D2. TRAM improves domain transfer in fine-tuned models by better leveraging pre-trained information from unseen domains within the smoother minima idealized by SAM-style training. Generally, the naive Adam baseline the MESA trust region comparison perform poorest at cross- domain language modeling for Wikipedia or S2ORC splits respectively. As with image classification, FSAM is the strongest competitor to TRAM. The best variant in both splits is TRAM-9,-1 improving in-domain and average zero-shot perplexity. TRAM-6,-1 uses the TRPO method of estimating the trust region using the parameters of the previous step. This variant always yields the lowest perplexity in the training domain and the majority of similar and distant zero-shot domains. We additionally verify that TRAM performs competitively at a larger model scale using GPT2-XL (1.5B parameters) in Table in Appendix C.2 We also compare against a naive combination of methods (e.g., ASAM+TRPO) in Appendix
TRAM improves perplexity for all domains in the Wikipedia split, where all zero-shot domains are positively correlated with the training domain perplexity. However, we observe that perplexity degrades for domains distant from the fine-tuning domain in S2ORC (MATH) which benefit less from shared features. Given that neither SAM-style nor trust region methods inverted this anticorrelation trend, it is unsurprising that TRAM follows suit. This confounder results in the overall best model, TRAM-0;-1- reAMoing.. MESA, reports the best performance. We suggest that optimization alone may be insufficient to improve zero-shot domain adaptation for larger distribution shifts. We discuss further the correlation between domain-specine sppeiffce perplexity in Appendix C.T
reporting the worst performance for the distant domains where the overall poorest model,
4.2. EASY AND HARD GENERALIZATION
When evaluating performance variation between different distributional shifts we find that TRAM improves on all prior work for minor shifts (e.g., MATI to Physics/Puys..) and generally matches or improves on negative trend for major shifts (e.g., MATII ART). Discussion of out-of-domain generalization often overlooks differences between major and minor shifts. In practice, in-domain performance has very different relationship to performance when generalizing to a major domain shift rather than minor shift. Considering minor distribution shifts, accuracy is strongly correlated on in-domain and out-of-domain datasets (Miler et al. (20TT). However, major distribution shifts may lead to scenarios where performance is instead anticorrelated with in-domain accuracy (Teney et al.) 2022). Considering these scenarios in the S2ORC task, we observe that models trained using TRAM often perform better
domains than their in-domain performance would predict. Furthermore, TRAM improves perplexity across both minor and major distribution shifts.
Figure 2a shows the close positive correlation between performance the training domain (MATH) and the average across all other STEM disciplines, considering all optimization approaches. As detailed in Appendix C.1 performance correlates with p >> 0.8 between MATH and each individual
7

---

Published as a conference paper at ICLR 2024
Table 5: XNLI accuracy (higher is better) for training language (EN) and 14 zero-shot target languages summarised by ZS AVG. (key in Appendix A. All TRAM variants significantly outperform other methods (p < 0.01 Wilcoxon test). Results are grouped as: (i) optimizers: (ii) trust region methods; and (ii) TRAM variants. We report the mean across 20 seeds with standard deviation in Table 13
EN
AR
BG
DE
EL
Es
FR
HI
Ru
Sw
TB
TR
UR
ZH
ZS AvG.
Adam
83.9
71.2
77.1
75.7
76.2
78.3
77.6
69.6
74.9
64.6
71.2
72.2
65.8
74.1
73.1
72.9
SAM
84,8
72.1
78.1
76.7
75.7
79.0
77.9
69.8
75.7
65.2
71.8
73.1
66.8
75.1
74.2
73.7
ASAM
85.0
72.0
78.4
76.9
76.1
79.5
78.5
70.4
76.1
65.2
72.5
73.4
66.9
75.5
74.2
74.0
FSAM
84.7
72.2
78.1
76.9
76.0
79.3
78.4
70.0
76.1
65.1
72.2
73.0
66.8
75.3
74.2
73.8
TRPO
84.9
71.3
77.7
76.2
75.3
78.6
77.3
69.2
75.2
64.4
71.6
72.4
05.3
73.8
73.3
73.0
85.5
72.7
78.9
77.5
79.9
79.2
70.7
76.8
66.2
72,9
73.9
66.6
74.6
74.5
MESA
84,9
71.9
77.9
76.7
75.7
78.8
77.8
69.6
75.8
64.1
72.1
72.4
65.7
74.4
73.9
73.3
TRAM-1
86.2
73.5
79.8
78.3
77.5
80.9
79.6
71.4
77.5
66.0
73.8
74.3
67.6
76.7
75.9
75.2
TRAM-0-11
86.2
73.1
79.5
78.2
77.0
80.2
79.7
71.5
77.5
66.4
73.3
74.2
67.5
76.7
75.8
75.0
TRAM-9,
85.6
72.9
79.3
77.8
77.4 4
80.2
79.6
71.2
77.1
65.9
73.3
74.2
67.5
76.7
75.8
74.9
Fisher
84.3
73.1
78.7
77.1
76.2
79.5
78.4
71.4
76.6
65.7
73.2
73.6
67.5
75.5
75.5
74.4
STEM category. Considering the blue dotted trend for previous optimization methods (excluding TRAM), we see that all TRAM optimizers fall on or marginally below the line. This result suggests that TRAM not only supports in-domain performance but specifically improves generalization to similar domains.
By contrast, we find there is generally a trade-off between performance MATH and the hardest anticorrelated domains: ART (Figure 26) and Philosophy (PHIL, Figure (c). Both TRAM-2 and TRAM-0, fall far below the trend for previous algorithms where in-domain improvement worsens out-of-domain perplexity. TRAM not only matches or outperforms existing methods on easier generalization cases, but exhibits lesser trade-off between easy and hard generalization compared to all previous approaches.
4.3 ZERO-SHOT CROSS-I.INGUA, TRANSFFR
Finally, now consider if TRAM improves cross-lingual adaptation during monolingual fine-tuning. We adapt a multilingual pre-trained model to an English entailment classification task (NLI) and then evaluate the zero-shot cross-lingust lingull capability for the model to classify entailment from inputs in 14 unseen languages. We hypothesize that TRAM benefits gross-lingul transfer via improved application of multilingual pre-trained information to task with only English training data. In general, languages closer to English (e.g.. French, German) are "easier" for transfer than distant or low-resource languages (e.g., Urdu, Swahili) (Ahmad et al., 2019). An ideal system will produce equivalent cross-lingual transfer for all zero-shot languages. Our complete experiment design is outlined in Appendix [.33 We train an FL.M-Robert--basee model (Conneau et al. 2020a) on English MultiNLI (Williams et al. 2018) and report accuracy results for the XNLI cross- lrngullingual entailment benchmark (discussed in Appendix A).
Table 5 highlights that TRAM improves over all competing methods for the cross-lingual transfer objective, similar to our findings for gross-dataset image classification and cross-domain language modeling. Similar to the above tasks, TRAM-r and TRAM-9;-1 are the best-performing algorithms reporting both the strongest in-domain and average out-of-domain accuracy. Either TRAM variant is the best method across all individual languages. We identify that all methods worsen for languages distant from English in a similar trend to language modeling for anticorrelated domains. However, here TRAM is strictly superior to any other method for both near and distant languages to English. Notably, TRAM-Fisher significantly improves upon FSAM (p < 0.01) despite the close similarity in methods. Given the additional forward pass required for TRAM-x, TRAM-Fisher represents a better performance-comples trade-off which is competitive in some tasks. We analyze the loss surface and representation transfer in Table [ to verify that TRAM extends a low-curvature loss surface and representation smoothness to all zero-shot languages. In Appendix C.5] we train a model using TRAM with alternative distances for trust region measurement to analyze the criticality of using KL divergence. We observe that TRAM is robust to multiple distances with marginal degradation. These results empirically verify our hypothesis that training with complementary SAM-style and trust region methods improves the language transferability of fine- mune model.
8

---

Published as a conference paper at ICLR 2024
Table Analysis of (a) e-sharpness and (b) CKA representation similarity for TRAM. We measure each metric using the XNLI validation set and report for the training language (EN) and the zero-shot languages (Z.S). We report mean of 20 runs standard deviation across languages and the Pearson correlation between EN and ZS AVG. e-sharpness across runs.
(a) e-sharpness $
EN
ZS AVG.
Pearson p
(b) CKA
EN
ZS AVG.
Adam
2.16
1.98 0.79
0.29::0.20
Adam
0.69
0.44 0.10
SAM
1.43
3.32 0.96
0.26:10.3
SAM
0.69
0.42+ 0.10
ASAM
2.57
2.22 0.79
0.38:1.12
ASAM
0.69
0.42+ 0.10
FSAM
2.34
2.62 0.29
0.27
FSAM
0.73
0.48:00111
TRPO
6.17
2.36+ 1.02
0.52++0.25
TRPO
0.70
0,45
R3F
6.22
2.56+ 1.21
0.50++0112
R3F
0.66
0.40
MESA
2.76
5.48 0.75
0.21+:0255
MESA
0.67
0.42
TRAM-0 --
0.50
1.1910.38
0.60:10.15
TRAM-0--1
0.77
0.57.10.10
TRAM-00
0.75
1.92.10:24
0.58.10.27
TRAM-00
0.69
0.45 0.11
TRAM-2
0,61
1,49
0.75=0.18
TRAM-E
0.75
0.541 0.11
1.67
2.02 0.40
0.42
TRAM-Fisher
0.72
0.49 0.10
Investigating the loss surface, we test the hypothesis that TRAM leads to flatter minima on both in-domain and out-of-domain data. We evaluate validation set e-sharpness (Keskar al. 2017), defined in Appendix B.7, across 20 trained models. We report in-domain (for English) and oft-oi-domain (zero-shot languages) e-sharpness in Table across TRAM and baselines (omitting models which under-performedd Most methods unsurprisingy demonstrate lower in- domain sharpness but poorer out-of -ff--an----inn sharpness. TRAM yields a smoother solution for both the in-domain and out-of-domain regions of the loss surface. We also observe an improved average Pearson correlation (and lower variance) between in in-distribution and out-of of-iiiiiiiuttiionn sharpness using TRAM. This infers that the relationship between loss surfaces of different distributions is more desirably predictable with TRAM. Notably, other SAM-style methods are worse than Adam for out-of-domain sharpness suggesting that current SAM algorithms (excluding TRAM) are possibly "sharpness sewripeeess-aare"""" only within the training distribution.
Loss surface dynamics:
sharnness-Suggeetingg
We analyze the similarity of pre-trained and fine-tuned representations for the same setup of XNLI. We hypothesize that if TRAM optimizes within the trust region, pre- and post-fine-tuned representations will more similar to allow better usage of pre-trained structure. We measure this relationship using CKA similarity (Kornblith et al. 2019) defined in Appendix [B.8] Similar to the previous analysis, we observe that TRAM produces representations that are more similar to pre- treitrained XLM-Robert representations than any competitor. This applies both the En case and the ZS AvG. case, with other models performing similarly to the Adam baseline. Counterintuitively, trust region methods perform no better than SAM-style methods which do not explicitly target representational similarity. This observation could be related to recent insight into the smoothness side effects of training with SAM (Wen et al., 2023). We additionally raise that neither metric in Table 6 shows similar trend to our empirical findings-compariionn comparisons here do not strictly reflect similar performance variation on specific tasks. Despite empirical improvement. recent work questions if sharpness meaningfully correlates with generalization (Juneja et al.) 2023] Andriushchenko al., 2003). Extending TRAM should further evaluate this relationship and investigate how trust region measurement could inform better predictors of generalization capability.
Representation transfer:
CONCLUSION
We present TRAM:
Trust Region Aware Minimization. TRAM optimizes for smoothness in both parameter and function spaces to improve domain generalization during fine-tuning. TRAM inherits the capability of SAM to optimize towards flatter minima and integrates trust region awareness to ensure low local curvature between output representations. We evaluate TRAM on out-of-distrribttio scenarios, where the model must generalize to new distributions unseen during training. In this setup, TRAM proves more effective than SAM-style optimization trust region methods. Our analysis identifies how TRAM bucks the anticorrelated trend for major distribution shifts, learns a flatter out-of-domain loss surface, and improves representation similarity for data unseen during fine-tuning.

---

Published as a conference paper at ICLR 2024
6 ACKNOWLEDGMENTS
TS gratefully acknowledges the support of the UK Engincering and Physical Sciences Research Council (grant EP/W002876/1). This work has been made possible in part by gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence.
REFERENCES
Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zettlemoyer, and Sonal Gupta. Better fine-tuning by reducing representational collapse. In International Confer- ence on Learning Representations, 2021. https:://pppnttt ew net. / forumm d=
0Q085N70MIV
Wasi Ahmad, Zhisong Zhang, Xuezhe Ma, Eduard Hovy, Kai-Wei Chang, and Nanyun Peng. On difficulties of cross-lingual transfer with order differences: A case study on dependency parsing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 1440-2455 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1011665//1//119111 / http:///i...aoooooooooooo///a/// org/N19-1253
Shun-ichi Amari. Natural Gradient Works Efficiently in Learning. Neural Computation, 10(2):
251-276, 02 1998. ISSN 0899-7667. 7/doi.org /ttp///...1160099996600000000000
doi: 10. N0.16688877769880007774444 URL https
Maksym Andriushchenko, Francesco Croce, Maximilian Müller, Matthias Hein, and Nicolas Flam- marion. A modern look at the relationship between sharpness and generalization. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.). Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 840-902. PMLR, 23-29 Jul 2023. URL htttps:/pproceedinn:
m.r prwss//s////e..... (eress/v2000 v22/2/222/////////2 /2t////2222////// htd0002/nn2ssshhhhkkkk2222.2 html
Dara Bahri, Hossein Mobahi, and Yi Tay. Sharpness-aware minimization improves language model generalization. Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7360-7371, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 101186653///022.2.oolooogg URL org/ 20t.....lloo.m......e 508
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoc V. Le. Symbolic discovery of optimization algorithms, 2023. URL
/tte:///rxiivi http://awwx...orrra////32222.00
Alexandra Chronopoulou, Matthew Peters, and Jesse Dodge. Efficient hierarchical domain adaptation for pretrained language models. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 1336-1355 1351, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10. 1/203./// waw.:22nnaalmmaaaa 9................. URL https::/aclaatthlogggr www.orgg/22222 naac naa.l-mainnmm
96
Alexandra Chronopoulou, Matthew Peters, Alexander Fraser, and Jesse Dodge. AdapterSoup: Weight averaging to improve generalization of pretrained language models. In Findings of the Association for Computational Linguistics: EACL 2023, pp. 2054-2060 Dubrovnik, Croatia, http::::///c..annoooo...oo
May 2023. Association for Computational Linguistics. 2023. findings eacl. 153
Alexis Conneau, Ruty Rinott, Guillaume Lample, Adina Williams, Samuel Bowman, Holger Schwenk, and Veselin Stoyanov. XNLI: Evaluating cross-lingual sentence representations. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 2485,, Brussels, Belgium, October- Novembe 2018. Association for Computational Linguistics. doi: 10. 10.18653/v//11--11 1269. URL wtt.ttttss/./ htttss://acll antho wwtholoog..org/iiiee.
10

---

Published as a conference paper at ICLR 2024
Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, Edouard Grave, Myle Ott, Luke Zettlemoyer, and Veselin Stoyanov. Un- supervised cross posss-lnguall representation learning at scale. In Proceedings of the 58th An- nual Meeting of the Association for Computational Linguisties, pp. 8440-8451, Online, July 2020a. Association for Computational Linguistics. doi: 10. 30.18633/1/20222acl-mm 747. URL
attipss//aclannholooyooooo org/
main. 747
Alexis Conneau, Shijie Wu, Haoran Li, Luke Zettlemoyer, and Veselin Stoyanov. Emerging cross- lingual structure in pretrained language models. In Proceedings of the 58th Anmual Meeting of the Association for Computatioonal Linguistics, pp. 6022-6034, Online, July 2020b. Asso- ciation for Computational Linguistics. doi: 1/2200///00 //aclanchology..ooo 2twp22200..cc1mmmm 536
URL https
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEFE Conference on Computer Vision and Pattern Recognition, pp. 248-2522 2009. doi: 10. 109/CVPR22009 52068488.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL hattps: net/ / Eorum? YichF FicTT
Jiawei Du, Zhou Daquan, Jiashi Feng, Vincent Tan, and Joey Tianyi Zhou. Sharpness-aware training for free. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL htppsss:///eeeeeeeee..... http:/ewwwwwi...eeffoeeeee dd=
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness- aware minimization for efficiently improving generalization. In International Conference on Learning Representations, 2021. URL
6= rM.
Jonathan Frankle. Revisiting "qualitatively characceiinn neural network optimization problems". ArXiv, abs/2012.06899 2020. URL htttpst/:///pp wwwp:pi..eeeetnnniniiieeeec
Org/CorpusID3
229152287
Robert M. French. Catastrophic forgetting in connectionist networks.
Trends in Cog-
nitive Sciences, 3(4):128-135,
1999. ISSN 1364-6613. doi: http:///oi.org/10
1016/813644-611309000144444 URL http://wwwwwcccecedeeeeeeeeee cew.e//eeeerrrcc...ooo//ccieeeee com.s/ienc/cceec/
articlee/11/11366611900022
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola.
A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012. URL
httt::///mlr..rg//eee papers/ httppp/pwesr//sptttoreomo
html
Pengcheng He, Xiaodong Liu. Jianfeng Gao, and Weizhu Chen. (DEBERTA): (DECODING)-
(enhanced) (bert) [with) (disentangled) (attention). In International Conference on Learning
Representations, 2021.
Sepp Hochreiter and Jürgen Schmidhuber. Simplifying neural nets by discovering flat minima. In G. Tesauro, D. Touretzky, and T. Leen (eds.), Advances in Neural Information Processing Syy- tems, volume 7. MIT Press, 1994. URL http://htss//prrceeeinggggg neuri pss. cc/paper............eeeo
Files/paper httpa///....................2.33 www.::12212---pp...--- papar.prfffi....p..
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 448 456, Lille, France, 07-09 Jul 2015. PMLR. http::::: htttps::/proceccingg mlr. w wwt.pppess/3////m http:///33//offfllll/ htmm
Pavel Izmailov, Dmitrii Podoprikhín, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. In Conference on Uncertainty in Artificial Intelligence, 2018. URI. /ttpp::://axiv..../.e wwwpwwwwwwww./////t/ea//oaaa/e/e 05407
11

---

Published as a conference paper at ICLR 2024
Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. SMART: Robust and efficient fine-tuning for pre-trained natural language models through principled regular- ized optimization. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguisties, pp. 2177-2190, Online, July 2020. Association for Computational Linguistics. doi: (0.18653/v1/2022.aac aal-aaan.m19999 main.197. URL Ut..ittees//aa...accloooo..ogog wwwp:llooy.org//2222/ ac -ma
197
Jeevesh Juneja, Rachit Bansal, Kyunghyun Cho, João Sedoc, and Naomi Saphra. Linear connec- tivity reveals generalization strategies. In The Eleventh International Conference on Learning Representations, 2023, URL httttt:t/ppenreeeteesse www.:/ewwn.e.e.o.iiiiiiiiiti
Jean Kaddour, Linqing Liu, Ricardo Silva, and Matt Kusner. When do flat minima optimiz- ers work? In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 16577-16595 Curran Asso- neur www.eeeeriseee.eeeeeeeeeee cctpeippscccpppppppa.ceeccc /pppe files/paper pdf
ciates, Inc., 2022, URL
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Gencralization gap and sharp minima. In International Conference on Learning Representations, 2017. URL /tttps://openreeee
nett/oorum?Id=Hlgyye
Minyoung Kim, Da Li, Shell Hu, and Timothy Hospedales. Fisher SAM: Information geom- etry and sharpness aware minimisation. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.). Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learming Research, pp. 11148-11161. PMLR, 17-23 Jul 2022. URL httpt::p:///
kim22f. htmll
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3519-3529. PMLR, 09-15 Jun 2019. URL https://
proceedings wwm......... wwa.s.press/vvee /ttesssvvv// /oo
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained categorization. In 2013 IEEE International Conference on Computer Vision Workshops, pp. 2013. doi: 10.1 00.10///CCW::2013....
Alex Krizhevsky, Learning multiple layers of features from tiny images. pp. 32-33, 2009. URL http:://wwwwcc. htt::::/www..c..oorottt..
www.ttoooooooooooo httpo//tw.....//////////aaaaaaat 2009
Jungmin Kwon, Jeongseop Kim, Hyunsco Park, and In Kwon Choi. Asam: Adaptive charpness-aware minimization for icale-invariant learning of deep neural networks. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5905 5914. PMLR, 18-24 24 Jul 2021. URL https: https://proceedingg
mw...presse. https////vv///ww22ll// wtt.wom22mbmmllllll.1mmmmmmmmmm.
39
Yaobo Liang, Nan Duan, Yeyun Gong, Ning Wu, Fenfei Guo, Weizhen Qi, Ming Gong, Linjun Shou, Daxin Jiang, Guihong Cao, Xiaodong Fan, Ruofei Zhang, Rahul Agrawal, Edward Cui, Sining Wei, Taroon Bharti, Ying Qiao, Jiun-Hung Chen, Winnie Wu, Shuguang Liu, Fan Yang, Daniel Campos, Rangan Majumder, and Ming Zhou. XGLUE: A new benchmark dataset for cross-lingual pre-training, understanding and generation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 2020. Association for Computational Linguistics. doi: 18.18//3/v wtw.1/.e22mmnn///.. main.484. URL https: a/cll antho ogy .rgg 2020 emn p-ma i n 484
1818, Online, November
Kyle Lo, Lucy Lu Wang, Mark Neumann, Rodney Kinney, and Daniel Weld. S2ORC: The semantic scholar open research corpus. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 4969 4983, Online, July 2020. Association for Computational
12

---

Published as a conference paper at ICLR 2024
Linguistics. doi: (006633/1//222.allmamiiiia URL https::/aclaatthooooyyyoo
John Miller, Rohan Taori, Aditi Raghunathan, Shiori Sagawa, Pang Wei Koh, Vaishaal Shankar, Percy Liang. Yair Carmon, and Ludwig Schmidt. Accuracy on the line: On the strong correlation between distr---ittiittiionn and in in-distribution generalization, 2021.
Thomas Möllenhoff and Mohammad Emtiyaz Khan. SAM as optimal relaxation of bayes. In The Eleventh International Conference on Learning Representations, 2023. URL https?
//openreview.oo lttee//fooommmdaaefeeeeeeeee
Maria-Elene Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing. pp. 722-729, 2008.
doi: 10. 10.1100000 N0.110//CVGIPP22000..1
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1667, 2020. URL
http::/ http:://jm...orr//ppprse/// http:/paaprr///11/20/ httpeerrs////20000..tttt
Machel Reid, Victor Zhong, Suchin Gururangan, and Luke Zettlemoyer. M2D2: A massively multi- domain language modeling dataset. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 964-975, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. URL 2022. www.eemain.nee.eeeeeemmmmmmm 63
http///oaaaattooo.....ooooo org/
John Schulman, Sergey Levine, Picter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889-1897, Lille, France, 07-09 Jul 2015. PMLR. URL press/v37/stt hmam3///cchllmmmmmmmmmmmmm 5. h
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov: Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):129-1198, 2014. URL http: wttp:/mlr....rrreeeee ortpppr.pp//ppp/////////// srivastaval4a.htm.
Damien Teney, Seong Joon Oh, and Ehsan Abbasnejad. Id and ood performance are sometimes inversely correlated real-world datasets. ArXiv, abs/2209.0061. 2022. URL http::://eatt
semanticscholaz wtg/Corppsspiippoooooooopp : 251999599 979643
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 International Conference Learning Representations, 2019. URL
nett/oorum?1d==xJ4m222222
Kaiyue Wen, Zhiyuan Li, and Tengyu Ma. Sharpness minimization algorithms do not only minimize sharpness to achieve better generalization, 2023.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume (Long Pupers), pp. 1112-1122 New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10. 00.18653/1/N188-1000 URL https
Chen Zhu, Yu Cheng, Zhe Gan, Siqi Sun, Tom Goldstein, and Jingjing Liu. Freelb: Enhanced adversarial training for natural language understanding. In International Conference on Learning Representations, 2020. URL /ttps::/openreviee net/
wtwpreewwwweeeeeeeooomeeeeeeeeee forum ddw..nindn-yyabairyy.
13

---

Published as a conference paper at ICLR 2024
Table 7: Data splits for M2D2 (Riid al. 2022) across Wikipedia and S2ORC (Lo al. 2020). For simplicity, we do not consider the fine- -naaind subdomains in each domain. All data sourced from Huggingface (huggingfacee co hts./datssee..... https//reeeeeeeeeeeeeeeee
Split
Domain
Abbrev.
(Tokens)
Truining
Validation
Calture and the arts
CLLT:
289M
34,33M
2.38M
Health and fitacss
HEALTH.
116M
6.83M
History and events
HIST.
226M
1166M
Human activities
HUMAN.
343M
12.41M
Wiki
Mathematics and logic
MATH.
52M
165M
Natural physicul sciences
1345M
Philosophy and thinking
PHIL
165M
2.32M
Religion belief systems Society and social sciences
64M
5.4MM
Soc.
397M
380M
118M
11,74M
Technology and applied sciences
TECH.
297M
11.7M
Art Astrophysics Condensed matter
ART
98M
LO6M
ASTRO
728M
L14M
CONDM.
Computer science
CS
1,1B
1.7MM
Economics
ECON.
116M
S2ORC
Mathematics
MATH
1.4B
1.IB
1.46M
1,4MM
Nonlineur sciences
NLIN.
134M
1.2MM
Philosophy
PHIL.
156M
1,6M
Physics
PHYS.
737M
1.12M
Quantitative biology
336M
1,8M
Statistics
STAT
450M
1.19M
A
DATA SPLITS
VISION DATASETS
For vision modality experiments, we evaluate cross-dataset transfer from ImageNet (Deng et al. 2009)
to CIFAR-100 (Krizhessky) 2009), Stanford Cars (Krause et al. 2013), and Oxford Flowers Nilssack
Zisserman 2008). partitions.
We source all datasets from Hugging Face using the default training/testing
LANGUAGE DATASETS
We evaluate the M2D2 dataset (Reid et al. 2022) for cross-domain language modeling. M2D2 contains two groups: 11 domains from the S2ORC corpus of ArXiv listings (Lo al. (020) and an archive of Wikipedia articles. We train a language model each split's 's largest domain and evaluate zero-shot generalization to ten domains unseen during fine-tuning. Evaluation uses token-level perplexity across each domain. Table details the partition sizes (in tokens) for each domain in M2D2.
Zero- shot gross-lingual transfer is evaluated using MultiNLI and XNLI for entailment classification. In this task, model predicts an entailment label (neutral, entailment, contradiction) between sentence pairs. We use only English language MultiNLI (Williams et al. 2018) for training data and evaluate trained model on the unseen natural languages in XNLI (Conneau et al. (2018) during test time. These datasets are balanced in label classes and report accuracy per language in results. A complete breakdown partition sizes shown Table 3
B
ADDITIONAL EXPERIMENTAL DETAILS
B.1 MODEL TRAINING
We fine-tune each pre-trained model without any freezing or additional task-specific parameters where possible. We also do not explore fine-tuning with low-rank approximations or adapters i.e., 'full fine-tuning'. This setup isolates the contribution of the optimization algorithm over additional capacity in the model. For image classification and cross-lingual entailment classification, we follow
"huggingfee co co/datasstt//ifa110 huggi http:///cceeeee///iiotttaaaaaaaa ctt...c//aaae htt.fff..../accaaaat/////ttttttr
14

---

Published as a conference paper at ICLR 2024
Table 8: Data splits for XNLI (Conneau et al. 2018). The Training data in English is sourced from the MultiNLI dataset (Williams et al., 2018) with translations provided for XNLI. Model selection during training uses only the English validation data. Validation data for other languages is used to measure e-sharness in our analysis. We omit data splits not used in this work. All data sourced from HuggingFace (hugging ngaace co pooddaaaeesss xnli).
XNLI
Abbrev.
Train Sentences
Validation Sentences
Test Sentences
English
EN
393K
2.5K
5K
Arabic
AR
2.5K
5K
Bulgarian
BG
2.5K
5K
German
DE
2.5K
5K
Greek
E1.
2.5K
5K
Spanish
Es
2.5K
5K
French
FR
2.5K
5K
Hindi
HI
2.5K
5K
Russian
RU
2.5K
5K
Swahili
Sw
2.5K
5K
Thai
TH
2.5K
5K
Turkish
TR
2.5K
5K
Urdu
UR
2.5K
5K
Vietnamese
VI
2.5K
5K
Chinese (Simplified)
ZH
5K
fine-tuning norms and only introduce a new datasct-specific 'head' to predict dataset-specific logits. For language tasks, we fine-tune each pre-trained model for 50,000 steps using an initial learning rate 2 x 10 polynomial decay schedule, and 10,000 step learning rate warmup. We use Adam (Kingma & Ba 2017). with a decay factor setting (31, .22 = (0.9, 0.99), as the base optimizer for each SAM-style and TR method unless mentioned otherwise. When using validation loss for model selection, we use only the validation partition of the training domain to reflect a stricter evaluation setup without access additional domains during training. All models are trained Ax1000 80GB GPU for under 72 hours except for GPT2-XL experiments in Appendix
B.2 BASELINES
We compare naive SGD baseline for vision experiments following Kim al. (2022). Our naive baseline for language experiments is Adam (Kingma & Ba 2017) without any augmentation setting decay factors (B1, 32) = (0.9, 0.99). All algorithms listed below use Adam as the inner optimizer for the final update (e.g., Algorithm| Step 6).
For sharpness-aware methods: we compare to SAM (p = 0.05, Foret al. 202 Adaptive SAM (ASAM, p = 0.5, Kwon et al. 2021) and Fisher SAM (FSAM, 7 = 0.1, 0.1, Kim et al. (022).
For trust region methods: we compare Trust Region Policy Optimization (TRPO, Schulman et al. 2015). R3F (a = 0.1. Aghajanyan et al. 2021). and MESA (Du et al. [0022). MESA is variant of TRPO regularizing output representation divergence between current 6, and the exponential moving
average of previous 0cr with decay factor 0.999. For trust-region methods, we add the regularizer directly to the task-specific loss function with a weighting coefficient of = 0.1 (in Equation 3).
B.3 CROSS-DATASET TRANSFER FOR IMAGE CLASSIFICATION
We implement the same cross- datsset adaptation setup as Kim et al. (022) as a 'sanity check' directly comparing TRAM to prior methods in the same setting. This setup is not strictly similar to the "out-of-distributiin" scenario we report for language tasks-this experiment verifies that
TRAM is performant on standard benchmarks and valuably evaluates TRAM the vision modality. The objective is to adapt ViT-base (Dosovitskiy et al. 2021) from ImageNet pre-training (Deng) et al. 2009) to additional image classification tasks. We evaluate dataset adaptation to CIFAR-100 (Krizhevsky| 2009). Oxford Flowers (Nilsback & Zisserman, 2008) and Stanford Cars (Krause et al.)
2013) datasets. Our hypothesis is that TRAM can improve applying information from ImageNet to
additional datasets with different labels and input data.
15

---

Published as a conference paper at ICLR 2024
Algorithm Trust Region Aware Minimization
Input:
Training set S = [(Ir, 1.)}. loss function e, learning rate model parameters 0, noise standard deviation (if moise-estimated trust region). for 1, 2,... do
(1) Sample batch of {(s, w)) 2 data from S.
(2) Compute the predictive distribution, py (-1a,, 0t), and gradient of the batch loss (8).
(3) Compute trust region distance d as:
da using P/ (-(ZB, 0c-1) ] (Equation or
de using Py (-|2BB + z, 0z)
N [0, σ²) (Equation 5
(4) Compute ETRAM=
ÉTRAM 0" (0.)
(5) Ascent step perturbing to + ÉTRAM-
(6) Compute gradient at 0 CTRAM as Equation
V LTRAM (0)
2L2 00
0-0+
(7) Gradient descent update: 8 + o (TRAM(()).
end for
We match the experimental setting Kim et al. 2022): fine-tuning ViT-base-16 for 200 epochs with a base optimizer of SGD, an initial learning rate of 5 x 10-4, and cosine learning rate decay with no warmup or restarts. We do not use early stopping to match prior work and use the final model regardless of validation loss. We report the average Top-1 accuracy over 5 runs, = the 95% confidence interval, in Table for direct comparison to Kim al. (2022 Table 3).
B.4 CROSS-DOMAIN LANGUAGE MODELING
We consider zero-shot cross-domain language modeling using the M2D2 Corpus (Reid et al. 2022). Our hypothesis is that TRAM can better apply language modeling information from large text corpora improve out-of-domain perplexity when fine-tuning on the "Math" domain ATH, 1.4B tokens) and for Wikipedia, we train on the "Society and social sciences" domain (Soc., 379M tokens). We use the 112M parameter GPT-2 base model (Radford cet al. 2019) with a batch size of 16 blocks of 1024 tokens following the setup of prior work (Reid al. 2022 Chronopoulou al. 2022 2023). We evaluate generalization via perplexity for each test domain. We also evaluate zero-shot baseline (i.e., GPT-2 before fine-tuning) to contrast with the same model before specifc adaptation. To reduce computation, we train one model with one random seed per algorithm.
specific domain. For S2ORC, we train
B.5 ZERO-SHOT CROSS-LINGUAL TRANSFER
We test zero-shot cross-lingual transfer by fine-tuning a multilingual model on an English task and then evaluating the model in other languages. We hypothesize that TRAM can improve task transfer across languages by improving the usage of information from multilingual pre-training during monolingual fine-tuning. A poorer model may 'forget' other languages during the adaptation process. We evaluate transfer from English to additional languages by predicting labels for the XNLI test set after training the model for NLI only in English. We use the 250M X-M-Roberta Base multilingual pre-trained model (Conneau et al., (2020a) with classification head trained from scratch. This model uses batch size of 32 examples using only English validation loss for model selection. Each reported result is averaged across 20 runs of varying random seeds to control for variation in loss surface.
B.6 TRAINING ALGORITTMS
The training algorithm for TRAM is outlined in Algorithm Jusing different metrics for trust region estimation, d, outlined in Section Algorithm@ modifies Algorithm|] in in removing one forward pass to estimate the trust region distance and instead approximate the Fisher Information Matrix of the trust region neighborhood in representation space.
the TRAM-Fisher algorithm. Practically, this
16

---

Published as a conference paper at ICLR 2024
Algorithm 2 Trust Region Aware Minimization with Fisher Information Matrix (TRAM-Fisher)
Input: Training set S [(rr, W:)3, loss function e, learning rate model parameters 0, noise standard deviation for 1......
1) Sample batch of ((r,, 10 data from S.
2) Compute the predictive distribution, Py (-Ix,, 0t), and gradient of the batch loss V (0))
3) Sample input noise z ~ (0, o² [m).
4) Approximate the Fisher Information Matrix at x + 2:
F (I + 2; 0) Diag
[3] EEED (logps (yilz, 21, 01)
5) Compute GTRAM-P:
P(ata: ILL
ÉTRAM-F VDLzFle+s; 0)-I
6) Ascent step perturbing to 0 ETRAM-F-
7) Compute gradient
6TRAM-F as Equation 8Ls 09
LTTRM-- (0)
@-0--TRAM
8) Gradient descent update: 0 < 0 a LTRAM-P(O))
end for
B.7 MEASURING SHARPNESS
We follow Keskar et (017) in evaluating model e-sharpness Equation I0 where é the loss function, 2 E R are n model parameters, A E R"x is a matrix restricting the e-sharpness to subspace of p parameters (A* is the pseudo-inverse of A) and C2 defined as Equation denoting "box" region around the solution over which loss is maximized.
&r (e, maXyEc, ((x + Ay) 8 (xx)
100
(10)
(+((x)
C= E
Rd -c((A-c((TT))))) 1) s 2, (I(A 1) bp)
(11)
For our measurement of e-sharpness, we set A to the identity matrix Iuxn to measure over the complete model. We measure e-sharpness over the validation set of XNLI all languages comparing between original loss é (x) and maximized loss max,ec, {(x + Ay). We follow the e-sharpness setup of Juneja et al. (2023) using an SGD optimizer, learning rate of 8 x 10-5, a 32 example batch size, accumulation over steps and of 1 x 5
B.8 MEASURING REPRESENTATION SIMILARITY
We follow Kornblith et al. (2019) and Conneau et al. (2020b) in evaluating cross-lingual similarity using Centered Kernel Alignment (CKA). At a language level, CKA computes a similarity score between matrix X and Y where X, Y R d are dense matrices of outputs of d-dimensional representations from each model. We compute linear CKA similarity as Equation 12 using the Frobenius norm. For cross- cross-lingual transfer experiments, example (i.c., the representation before the classification head) to evaluate similarity.
the base model output for each
XI
CKA (X,Y)
(12)
XX X ITY
C ADDITIONAL RESULTS
C. DOMAIN CORRELATIONS FOR S2ORC
Table details the correlation between zero-shot and training domain perplexity across methods. We omit the combination approaches (e.g., ASAM+R3F) due to poor performance. For Wikipedia, all domains correlated with the training domain indicating that the domain-specific fine tuning on
17

---

Published as a conference paper at ICLR 2024
Table Pearson correlation between training domains and zero-shot domains for M2D2. We report how the change in training domain correlates with changes in zero-shot perplexity to analyze how different domains improve or worsen during fine-tuning. All domains are correlated with Soc. for the Wikipedia split. ART and PIIIL. domains are anti-correlated with MATII training domain for S2ORC indicating a major distribution shift.
Wiki Domain
p to Soc.
< 0.01?
S2ORC Domain
p MATII p < 0.01?
STEM?
CULT. 0.982
ART
-0,861
GEN. 0.983
/
ASTRO
0.812
HEALTH.
0.970
V
CONDM.
0.999
HIST.
0.998
CS
0.996
HUMAN.
0.980
V
ECON.
0.997
V
MATH.
0.976
NLIN.
1.000
V
NAT.
0.982
PIIIL
-0.825
PIIL.
0.985
PIYYS.
0.991
4
REL
0.994
QBIO
0,932
TECH. 0.983
STAT
0.998
7S AVG. 0.990
ZS AVG.
0.968
STEM AvG
0.998
Table 10: M2D2 perplexity across training algorithms for GPT2-XI. We fine-tune on the MATH do- main M2D2 S2ORC split and evaluate in-domain and out-of-domain perplexity. We evaluate TRAM, competitive comparisons and GPT2-XL zero-shot baseline. We omit algorithms demonstrating poorer results smaller scale experiments to limit computation demands. As in Table 4 TRAM performs strongly compared to all comparisons. We report the average zero-shot perplexity (ZS AVG.) as the summary metric to judge domain transfer capability (lower is better). Worst perplexity (excluding zero-shot) is red, best is green.
S2ORC
MATH
ART
PuIL.
CONDM.
CS
ECON.
NLIN.
PHYS.
QB10
STAT
ZS AvG.
GPT2-XL
16.9
22.8
21.2
19.8
19.0
17.3
18.5
17.8
20.7
19.8
14.9
19.2
Adam
8.7
30.4
28.2
24.0
14.9
15.4
15.4
11.4
21.4
22.1
12.6
19.6
SAM
8.7
29.3
28.0
22.6
14.6
15.1
15.1
11.2
20.5
21.4
12.3
19.0
ASAM
7.9
28.0
26.1
21.8
13.4
14.1
10.4
19.4
11.4
17.9
FSAM
7.8
26.7
25.0
21.1
13.1
13.7
13.7
10.2
15.8
19.6
11.2
17.3
TRPO
8.9
27.9
26.4
23.0
14.9
15.3
15.3
11.5
20.8
21.3
12.5
18.9
R3F
8.9
27.9
26.4
23.0
14.9
15.3
15.3
11.5
20.8
21.3
12.5
18.9
MESA
9.1
28.7
26.7
23.7
15.0
13.1
20.7
23.2
12.8
19.5
TRAM-0
25.2
20.1
13.7
14.0
14.2
10.7
17.2
TRAM-2
8.3
25.3
23.8
20.2
13.8
14.1
14.2
10.8
19.0
19.5
11.6
17.2
Soc. domain has net positive improvement all zero-shot domains. This trend is not consistent for S2ORC where we observe that ART and PHIL. domains are anti-correlated with the MATH training domain. Improvement to MATH perplexity worsens the performance these domains across all methods. As discussed in Section 4.2.1 TRAM reports perplexity below this trend to perform better than expected for negatively correlated trend. For comparison, contrast the correlations between positively correlated domains (grouped as an average entitled STEM) and anticorrelated domains Figure]
C.2 TRAINING GPT2-XL WITH TRAM
Bahri et al. (022) report that training with SAM is effective over all sizes of T5 (Raffel et al.) 2020). We verify this improvement trend extends to TRAM by training a GPT2-XL model (1.5B parameters) the same language modeling task for 100,000 steps. The setup is the same as described in Appendix B but we use 4 A100 GPUs for training each with a batch size per device of 4 blocks x 1024 tokens. Perplexity for S2ORC domains is shown in Table where we observe similar trends to the 112M parameter GPT2 model. We choose not to run these larger experiments on methods with poor performance in Table (e.g., combined approaches, TRAM-Fisher) to limit computation
18

---

Published as a conference paper at ICLR 2024
demands. Zero-shot GPT2-XL is stronger baseline here which some methods struggle to improve upon despite improvement in the training domain. TRAM-8--1 and TRAM-r perform similarly reporting the lowest perplexity in four domains. The most competitive adjacent algorithm is FSAM reporting the lowest perplexity in seven domains. The difference between FSAM and either TRAM algorithm is not significant here, observed for smaller models in Table
C.3 RESULTS FROM COMBINING OPTIMIZATION ALGORITHMS
Table 11: M2D2 perplexity (lower is better) on Wikipedia (upper) & S2ORC (lower) splits. TRAM- 8,-1 significantly improves over prior work (p < 0.01 Kolmogorov-Smirma
test). Results are grouped as: (i) optimizers; (ii) trust region methods; (iii) combined SAM optimizers and trust region methods; and (iv) TRAM variants, The leftmost column is the training domain and we evaluate zero-shot perplexity on domains unseen during fine-tuning (full details in Appendix A ZS AVG. is the macro- average of all zero-shot domains.
WIki
Soc.
CULT.
GEN.
HEALTH.
HIST.
HUMAN.
MATH.
NAT.
PHIL
REL
TECH.
ZS Ave.
GPT-2
27.2
27.7
27.8
24.5
29.2
28.8
28.6
29.4
27.8
27.7
28.7
28.0
Adam
24.8
26.3
26.4
23.6
27.2
27.0
27.4
27.6
26.3
25.8
27.4
26.5
SAM
24.5
25.9
26.6
27.2
25.8
25.5
27.0
26.1
ASAM
24.8
25.4
25.6
22.5
27.1
26.4
26.3
26.7
25.5
25.5
28.1
25.9
23.0
TRPO
21.8
23.0
23.3
20.7
24.0
23.7
23.8
24.0
23.1
22.8
24.1
23.3
R3F
21.8
23.0
23.3
20.7
24.0
23.7
23.8
24.0
23.1
22.8
24.1
23.3
MESA
24.0
25.4
24.9
24.8
25.2
24.1
24.0
25.1
ASAM-TRPO
26.8
27.6
28.2
26.5
27.9
ASAM+R3F
25.0
26.0
26.2
23.2
27.4
26.9
26,8
27.4
26.1
25.9
27.1
26.3
ASAM+MESA
26.3
27.1
27.7
26.3
27.4
TRAM-6;-1
22.1
22.9
23.3
22.0
23.4
TRAM-1
21.9
23.1
23.4
20.7
23.9
23.3
23.9
23.8
23.1
22.7
23.9
23.2
TRAM-F
23.1
20.7
23.9
23.9
23.2
22.7
23.9
TRAM-Fisher
22.5
23.7
24.0
21.3
24.6
24.0
24.7
24.6
23.8
23.3
24.6
23.9
S2ORC
MATH
ART
ASTRO
CONDM.
CS
ECON.
NLIN.
PHIL.
PHYS.
QBi0
STET
ZS AvG.
GPT-2
27.6
35.8
32.4
30.9
27.9
29.5
27.6
33.7
33.5
30.9
23.4
30.6
Adam
11.4
44.2
33.9
20.1
21.2
21.0
14.7
41.9
29.5
30.8
16.9
27.4
SAM
10.5
45.3
33.2
18.7
20.3
20.0
13.7
42.4
28.3
30.2
16.1
26.8
ASAM
10.3
45.6
33.2
18.5
20.1
19.8
13.5
42.6
28.2
30.2
15.9
26.8
TRPO
33.4
20.3
28.4
R3F
10.4
46.0
33.4
18.6
20.2
20.0
13.6
42.9
28.4
30.4
16.0
26.9
MESA
11.9
44.1
34.1
20.8
21.7
21.6
15.3
41.7
30.0
31.0
17.4
27.8
ASAM-TRPO
36.9
23.6
23.8
23.8
17.4
33.1
33.5
ASAM-R3F
13.5
46.2
36.5
23.3
23.6
23.5
17.2
43.4
32.7
33.22
19.0
29.9
ASAM+MESA
13.4
45.9
36.3
23.1
23.4
23.3
17.0
43.2
32.5
33.0
18.9
29.7
TRAM-0-1
9.6
16.5
32.5
17.2
19.2
18.9
12.6
43.3
27.0
29.6
15.0
26.2
TRAM-00
44,8
33.0
18.6
20.1
19.9
13.6
42.0
28.2
30.0
15.9
26.6
TRAM-2
10.4
44.9
33.0
18.6
20.1
19.9
13.6
42.0
28.1
30.0
15.9
26.6
TRAM-Fisher
20.3
30.3
Given that TRAM builds on integrating SAM-style optimization with trust-region regularization, additionally compare to naive combination of each of these methods. We replace the standard loss function in ASAM with the loss function adding trust region regularization.
Our full results featuring these systems are shown in Table I for language modeling and Table 12 for zero-shot cross-lingual transfer. Across both tasks, naive combination approaches are some the weakest approaches. When we directly combine ASAM with each trust region regularizer (TRPO, R3F, MESA), we find that the naive combination approaches perform worse than Adam alone, even with extensive hyperparameter tuning. We conjecture that the constituent methods fail to compound constructively because the trust region regularizer does not interact with (or respect) the p-ball neighborhood of ASAM. Therefore, each component may contribute cross-feature interference, with a disadvantageous net effect on training. TRAM instead offers to combine strategies with complementary features without interference.
19

---

Published as a conference paper at ICLR 2024
Table 12: XNLI accuracy (higher is better) for training language (EN) and 14 zero-shot target languages summarised by ZS AVG. (key in Appendix(A) All TRAM variants significantly outperform other methods (p 0.01 Wilcoxon test). Results are grouped as: (i) optimizers; (ii) trust region methods; (iii) combined SAM optimizers and trust region regularization; and (iv) TRAM variants. We report the mean across 20 seeds with standard deviation in Table 13
En
AR
BG
DE
EL
Es
FR
HI
RU
Sw
Th
TR
UR
Z#
ZS AvG 1
Adam
7T.1
75.7
77.6
69.6
71.2
72.2
74.1
T3.1
SAM
84,8
72.1
78.1
76.7
75.7
79.0
77.9
69.8
75.7
65.2
71.8
73.1
66.8
75.1
74.2
73.7
ASAM
72.0
78.4
79.5
78.5
70.4
72.5
73.4
74.2
74.0
FSAM
84.7
72.2
78.1
76.9
76.0
79.3
78.4
70.0
76.1
65.1
72.2
73.0
66.8
75.3
74.2
73.8
TRPO
84.9
71.3
77.7
76.2
75.3
78.6
77.3
69.2
75.2
64.4
71.6
72.4
65.3
73.8
73.3
73.0
R3F
85.5
72.7
78.9
77.5
76.8
79.9
79.2
70.7
76.8
66.2
72.9
73.9
66.6
75.8
74.6
74.5
MESA
71.9
77.9
76.7
75.7
78.8
77.8
69.6
64.
72.1
72.4
65.7
74.4
73.9
73.3
ASAM+TRPO
85.0
72.4
78.5
77.2
76.4
79.7
78.9
70.4
76.4
65.3
72.4
73.2
66.8
75.7
74.6
74.1
ASAM+R3F
85.1
72.1
78.3
76.9
75.9
79.3
78.4
70.3
76.0
65.1
72.4
73.3
66.3
75.1
74.3
73.8
ASAM+MESA
84.7
71.7
77.8
76.3
75.7
78.8
77.9
69.5
75.4
61.1
71.6
72.7
65.6
74.3
73.4
73.2
TRAM-0--1
86.2
73.1
79.5
78.2
77.0
80.2
79.7
71.5
77.5
66.1
73.3
74.2
67.5
76.7
75.8
75.0
TRAM-0,
79.3
77.8
80.2
79.6
71.2
65.9
73.3
74,2
75.8
TRAM-F
56.2
73.5
79.8
78.3
77.5
80.9
79.6
71.4
77.5
66.0
73.8
74.3
67.6
76.7
75.9
75.2
TRAM-Fisher
73.1
78.7
79.5
78.4
71.4
65.7
73.2
73.6
75.5
Table 13: Standard deviation of accuracy for the XNLI dataset across 20 training runs with varying random seed. Results are split into groups for: (i) optimizers, (ii) trust region methods, (iii) combined methods, (iv) TRAM variants and (v) TRAM using d, with varying metrics for computing divergence. This accompanies Table and Table 14 which report average values across seeds.
BG
DE
AR
Es
H1
RU
T1
TR
V1
Zii
Adam
0,34
0.42
0.65
0.47
0.50
0.39
0.51
0.51
0.51
0.65
0.40
0.41
0.43
0.51
0.55
SAM
0.24
0.31
0.35
0.34
0.31
0.32
0.49
0.33
0.50
0.36
0.40
0.35
0.44
0.39
0.39
ASAM
0.33
0.33
0.45
0.39
0.47
0.36
0.51
0.44
0.44
0.51
0.56
0.42
0.47
0.45
0.50
PSAM
0.35
0.31
0.51
0.56
0.35
0.37
0.47
0.41
0.45
0.53
0.37
0.39
0.44
0.35
0,38
TRPO
0,24
0.35
0.37
0.34
0.30
0.39
0,34
0.46
0.40
0,38
0.36
0.34
0,53
0.38
0.34
R3F
0.40
0.44
0.35
0.35
0.42
0.46
0.41
0.35
0.39
0.47
MESA
0.34
0.34
0.44
0.24
0.40
0.52
0.37
0.67
0.43
0.26
0.40
0.45
0.59
0.43
0.34
ASAM+TRPO
0,30
0.28
0.36
0.29
0.26
0.28
0.35
0.34
0.44
0.36
0.39
0.38
0.34
0.32
0.32
ASAM+R3F
0.32
0.45
0.45
0.40
0.46
0.36
0.49
0.53
0.50
0.52
0.40
0.49
0.60
0.45
0.49
ASAM+MESA
0,34
0.31
0.30
0.44
0.42
0.39
0.38
0.51
0.58
0.46
0.44
0.31
0,51
0.50
0.46
TRAM-0,-1
0.40
0.31
0.40
0.30
0.36
0.31
0.43
0.50
0.53
0.48
0.43
0.34
0,36
0.49
0.42
TRAM-6,
0,34
0.38
0.41
0.44
0.48
0.40
0.43
0.53
0.63
0.47
0.66
0.39
0.57
0.50
0.54
TRAM-1
0.31
0.29
0.45
0.44
0.37
0.33
0.38
0.37
0.48
0.39
0.44
0.32
0.43
0.39
0.42
TRAM-Fisher
0.29
0.65
0.67
0.55
0.58
0.60
0.49
0.72
0.69
0.64
0.86
0.49
0.68
0.55
0.73
TRAM-r M- (MMD)
0.42
0.38
0.46
0.43
0.47
0.35
0.42
0.43
0.48
0.44
0.59
0.49
0.36
0.37
0.59
TRAM-r (L2)
0.30
0.27
0.27
0.26
0.24
0.28
0.27
0.21
0.29
0.26
0.24
0.28
0,21
0.22
0.22
C.4 RUN VARIATION IN CROSS-LINGUAL TRANSFER
For XNLI experiments, we report the mean over 20 runs varying random seed in Table We report the respective standard deviation values for each reported mean in Table [3]
5 and Table
C.5 CHOOSING DISTANCE METRIC
TRAM relies on KL divergence to estimate the trust region around the pre-trained function (i.e.. pf (+|æ + 2,0) 0) or Py (-|z, 0.-1)). We propose TRAM with forward KL on the intuition that the perturbed distribution (i.e., estimated point in the trust region) is the target (i.e., true) output which the current outputs (i.e., estimate) should match. We empirically verify this setup as the optimal arrangement (i.e., forward KL). While reverse KL or symmetric KL report only marginally poorer results, we report only forward KL for simplicity. We also consider alternative distance metrics in Table 14 We evaluate modifying the best pestt-erformin model for XNLI with different distances to examine if the divergence for trust region estimation influential in performance. We evaluate maximum mean discrepancy using an inverse multiquadratic kernel (MMD; Gretton et al. 2012). or La distance within dz. Even using the worst-pe metric, T,2 distance, TRAM is still
20

---

Published as a conference paper at ICLR 2024
Table 14: XNLI accuracy across varying the divergence metric estimating the trust region distance in TRAM. We compare to using maximum mean discrepancy (MMD) and L2 distance. TRAM is generally robust different estimates for the trust region between py (-|z, 0) and ps (+12r + 2, 2)
Avo.
TRAM-r (KL)
86.2
79.8
78.3
77.5
73.5
80.9
79.0
71.4
77.5
06.0
73.8
74.3
67.6
76.7
75.9
75.2
TRAM-r (MM))
TRAM-r (La)
85.1
78.7
76.8
76.2
72.
79.4
78.8
70.4
76.2
65.
12.6
73.5
67.1
75.8
74.6
74.1
competitive to methods in Table Characterizing the best trust region estimate for TRAM is outside the scope of this work. Future work should explore the suitability of different distances (e.g., Renyi divergence) to improve the estimation of the trust region space.
21


Reviews:
[{"review_id": "MBkRWxvRGw", "reviewer": "Reviewer_AFpQ", "summary": "This paper develops a new SAM-style optimizer for better representation learning (especially for language modeling).\nThe key idea is simple-to-state (although the proposed algorithms are a bit more complicated than the original SAM).\nThe authors attempt to combine the best of both worlds between SAM and Trust-region:\n1. SAM encourages the solution to be \"flat\" in the parameter space (assuming that the flat minima is desired)\n2. Trust-region methods encourages the representation to be \"smooth\" or to stay close to the good \"pre-trained\" model initialization.\nIn particular, with the Trust-region feature, the authors claim that the proposed method has the benefit of not forgetting task-agnostic representations from pre-trained model and also learning \"smooth\" representation which is good for the transferability of representations.\n\nAlso, building on more advanced SAM algorithms like ASAM and FSAM, the authors develop other variants of TRAM which are used for the experiments and show better performance than the previous approaches.", "questions": "- As far as I understood, the original SAM paper became popular because of its extensive experiments over standard benchmark datasets. In particular, I remember SAM achieving state-of-the art for various vision tasks (CIFAR, ImageNet etc...). Given that the effectiveness of SAM was first demonstrated on these benchmark tasks, **for future research I think it is required for follow-up works to sanity check the performance of the proposed methods on the same setting as the original SAM paper**. In particular, if the new methods end up giving a worse performance than SAM for the settings considered in the SAM paper, that would be an important information for practitioners.\n(As I mentioned in the strength part, the authors' quite comprehensive experiments on the language modeling tasks of choice look great; however, since this work follows up on the original SAM paper, **some experiments that benchmark the new method against the original SAM on the task that SAM did great seems required**.)\n\n- Also, it's great that the authors built their default algorithm on ASAM. But, given that SAM has been quite popular, as a reader, I'm quite curious how the SAM version of TRAM (instead of ASAM) performs. In particular, **is ASAM type of updates really necessary?**\n\n- Given that the main motivation of this work is to develop an optimizer for learning good-representation, I think at least one experiment is needed for pre-training from scratch. In particular, do you think having a reasonable pre-trained model is necessary for TRAM to work well? As far as I know, it's still debated in the literature **whether SAM-type of updates are required in the beginning of the training or at the end of the training**. Some theoretical works have claimed that it's only effective at the end of the training (**as did in this work**), but empirical works also have claimed that it's required from the beginning of the training.\n\n- In Figure 1, could you clarify what the negative and positive slopes are supposed to be interpreted as? Also I can't really understand how to interpret this plot.\n\n- In Table 5, why didn't you present the statistics for the other two variants of TRAM?\n\nI acknowledge the novelty of this work. However, given the extensive experiments in the previous works (e.g. original SAM work), in order to make a case about the effectiveness of the proposed methods, **I think some more \"sanity-check\" experiments are needed. Especially, because this work is empirical in nature.** I'm voting for weak accept at the moment, but I'll make the final decision based on how the authors address my questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper did a good job summarizing the existing approaches and how the proposed method builds on top of them.\n- The experimental settings are detailed, and reasonable.\n- Experiments seem quite comprehensive at least for the settings considered in this work.\n- It's quite remarkable that the proposed methods achieve best performance across different fine-tuning tasks.", "weaknesses": "- See the question section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aB7WnXOkh3", "reviewer": "Reviewer_MZmv", "summary": "This paper proposes a new optimization algorithm called TRAM that combines sharpness-aware minimization (SAM) with trust region regularization. SAM methods like ASAM optimize for low sharpness (flat minima) in parameter space. Trust region methods constrain optimization to a local neighborhood in representation space. TRAM unifies these approaches by bounding the SAM perturbation region using the trust region distance. This encourages flat minima while retaining representation smoothness.", "questions": "1. How well does TRAM transfer to other modalities like images?\n2. There are several hyper-parameters of the proposed method. How to select them for out-of-distribution generalization?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is intuitive and well-motivated. The combination of SAM and Trust region methods is reasonable and interesting.\n2. Extensive experiments on multiple NLP tasks demonstrate the effectiveness of the proposed method.", "weaknesses": "1. Theoretical motivation for unifying SAM and Trust region methods is not provided. \n2. Some results have high variance across runs. More runs may better characterize the performance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0SlKGUf7xx", "reviewer": "Reviewer_2Ruf", "summary": "In this paper, the authors aim to propose a SAM variant to contribute to the training in the area of model fine-tuning, where they propose the Trust Region Aware Minimization. In the method, specific distance measures in TRR are employed as the neighbourhood radius in SAM rather than the manually-set pre-defined radius. The authors claim that the proposed method can optimizer for informative representations without forgetting pre-trained structure. And the authors investigate the perplex on M2D2 Corpus with GPT2 to show the effectiveness of the proposed method.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "**Strengths**\n\n1. The paper is clearly written and easy to follow.\n2. I think the paper aims to contribute to SAM from a very interesting perspective, i.e. fine-tuning techniques. Considering that fine-tuning has become a nearly necessary procedure in NLP tasks, the paper may provide some promising instructions further.\n3. Combine the proposed method with Fisher-SAM can reduce extra forward-propagation count when implementing to the same count as in vanilla SAM.", "weaknesses": "**Weakness**\n\n1. The core of this proposed method is to adaptively change the neighbourhood radius in SAM (or ASAM) based on certain distance measure. This somehow does not follow the idea of Trust Region Regularization which adds additional constraint on top of the loss according to the measure. More accurately, they are two different things. And, I could not find a clear meaning why using such a distance as the neighbourhood radius could give the \"Trust\". Several questions arise: what does the \"trust\" indicate in the proposed method? Why we should not trust the region that is not in the proposed method but in SAM (and ASAM, Fisher-SAM)? Why the given region would not harm the pre-trained models while SAM could? Clear answers are missing in the current paper. Also, it is highly recommended that the authors use figures to illustrate this and the core of the presented method.\n\n2. The Stochastic Weight Averaging (SWA) could also lead to a similar effect as TTR methods. The authors may need to also consider or compare SWA with the TRR and the proposed method. The following papers may be helpful.\n\n    [1] Kaddour, Jean, et al. \"When do flat minima optimizers work?.\" Advances in Neural Information Processing Systems 35 (2022): 16577-16595.\n\n    [2] Wortsman, Mitchell, et al. \"Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.\" International Conference on Machine Learning. PMLR, 2022.\n\n3. From the paper, I see nearly no discussions regarding why and how the proposed method could contribute the fine-tuning in the related section, given that the authors claim \"their method could not forget the pre-trained structure\". BTW, I think their abstract may be somewhat over-claimed. It is interesting to see that the authors are aiming to study the effect of SAM specifically in fine-tuning. But unfortunately, the current version could not present sufficient helpful insights.\n\n4. It would be more impactful that the the authors present results using the proposed method on some recent popular scalable pre-trained llm such as llama. \n\n\n5. It is highly recommended that the authors release their code.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CBLkHEBNUv", "reviewer": "Reviewer_cqYa", "summary": "This paper considered adopted trust regions and sharpness aware minimization. Specifically, the gradient update is based on two key steps. Equation (4) illustrates the update rule by KL divergence fine-tuning. Equation (5) illustrates the update of the input by adding a gaussian random variable. Then the model is further validated in GPT2 on different dataset benchmarks.\n\n****Post-Rebuttal****\n\nI would appreciate the hard work and additional experiments by authors. I increased the soundness score. \nHowever, I still feel very confused about several core arguments and positioning within the paper. \n\n- If this paper is a pure empirical paper, this paper should be revised in a major form to highlight the practical contribution rather than new algorithmic contribution in **optimization**. I do not think the contribution is significant in the context of optimization theory. \n- If current paper aims to tackle an optimization problem, I do not think the entire analysis is rigorous from a theoretical perspective. I read the related works by authors and I still could not understand why this kind of optimization could have better transferability. Please take note the mentioned theoretical papers did not discuss this point, but rather on the generalization property. (They are different concepts in theory. One is IID and another is OOD). Based on the unclear fundamental points, I could not increase my score, despite numerous additional experiments.", "questions": "I noticed the primary domain is about optimization.  While there are so many missing points in terms of rigorous analysis in the optimization. If this paper is an applied NLP paper, the paper should be revised in a major form.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- This paper considered an improved method in TRPO. Through smoothly updating the gradient, this method seems to transfer the information.\n- Empirical validation in the foundation model and large scale NLP dataset are done.", "weaknesses": "- I could not understand *Why* such as transfer could encourage a better transfer. When it does not work? I would like to see a *rigor* mathematical analysis. The sharpness aware minimization could achieve better generalization. How is this mathematically to ensure a better transfer?\nI noticed the experiments illustrated \n> domain transfer in fine-tuned models by better leveraging the pre-trained structure from unseen domains within the smoother minima idealized by SAM-style training. \n\nHowever, without clear analysis. This reviewer feels quite difficult to understand why SAM could achieve this objective. Does this approach only work for the selected dataset? **When it fails?** \n\n- I would like to see a computational/memory complexity analysis. How it compares with other methods.\n\n- This paper proposes a general machine learning method while it is only validated in the NLP dataset. Unless the author clearly revised the title and contributions, I would like to see the results in other modalities such as image. \n\n- Equation (3) is not clearly defined. What does it mean by d_{\\theta, x}? It is not a rigorous expression.\n- Eq(4), Eq(5) why forward KL divergence is considered? Why not reverse KL divergence? Or Other general forms such as Renyi divergence?\n- In eq(5), how important is the noise variable? Does the variance of the noise matter?\n- Eq(12) may not be effective to correctly estimate the similarity in high dimensional regime. I think there is a complexity issue here (this is not a sample efficient estimator). I could think this value does not make sufficient sense to me in a high-dimensional case.\n- Table 4 Why only accuracy is considered a metric? Is this dataset balanced?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)