TRAM: Bridging Trust Regions and Sharpness Aware Minimization

We thank the reviewer for their time and attention in reviewing our work. We hope our additional results and insight will consider the reviewer in revising their score.

Q1. On formal analysis

Please see our response below adapting ViT-base-16 to CIFAR-100 / Cars / Flowers dataset for validation that TRAM also is applicable for computer vision tasks similar to [1,2,3]. We also address that TRAM requires some initial pre-training for applicability in fine-tuning. We can therefore broadly state that TRAM does not work training from scratch at this stage.

We did not discuss the generalization bound of TRAM as we generally consider TRAM as a subsolution to ASAM [2]. The generalization bound of ASAM is described as Theorem 3 from [2] and also derived formally in [4]. On the assumption that $d$ is normally distributed, we expect the generalization bound of TRAM to converge to the same bound as ASAM. ASAM already describes a bound valid for any $\rho>0$ where we ensure that the TR distance is similarly nonzero.

[1] Sharpness-Aware Minimization for Efficiently Improving Generalization https://arxiv.org/pdf/2010.01412.pdf [2] ASAM: Adaptive Sharpness-Aware Minimization for Scale-Invariant Learning of Deep Neural Networks https://arxiv.org/pdf/2102.11600.pdf [3] Fisher SAM: Information Geometry and Sharpness Aware Minimisation https://proceedings.mlr.press/v162/kim22f/kim22f.pdf [4] The intriguing role of module criticality in the generalization of deep networks https://arxiv.org/abs/1912.00528

Q2. Complexity

We generally address this by discussing the number of forward and backward passes required for each algorithm (Table 2). The complexity of SAM / ASAM is not discussed in detail by the original authors, similarly referring to the number of forward and backward passes [1,2,3].

At a high level, SAM-style training also needs to permute each parameter in the network. Permuting each parameter requires some operations to permute $\theta$ to $\theta + \epsilon$. This scales with the number of parameters $N$ with some $k$ operations per parameter (e.g., to compute Equation (2) / Equation (7)). Computing $\epsilon^{\ast}$ (e.g., Eq(2) or (7)) requires one computation of the model parameter global norm, $g$, for all algorithms. ASAM / TRAM requires one matrix-matrix product to compute $\theta^2$ per parameter —we denote $M^{3}$ for worst case complexity of this operation. TRAM does not require additional computation here as the distance replaces the scalar $\rho$.

We can approximate this as:

SGD: 1 forward pass, 1 backward pass
SAM: 2 forward pass, 2 backward pass. $Nk + g$ operations for permuting $\theta$
ASAM: 2 forward pass, 2 backward pass. $N(k + M^{3}) + g$ operations for permuting $\theta$
TRAM: 3 forward pass, 2 backward pass. $N(k + M^{3}) + g$ operations for permuting $\theta$ + constant operation to compute the TR distance. 

We empirically observe that the 2 backward passes dominate the learning process as the complexity bottleneck. TRAM has minimal additional overhead beyond ASAM. The permutation operations for $\theta$ to $\theta + \epsilon$ can be parallelized to reduce wall clock runtimes. The memory overhead of TRAM is minimal compared to SAM or ASAM. There is no gradient passed through the permuted parameters for estimating the TR distance, and the final TR distance is a scalar replacing an existing hyper-parameter. Our empirical conclusion is that TRAM has minimal (<10%) increase in wall clock runtime and no additional memory requirements beyond what is required for ASAM.

[1] https://arxiv.org/abs/2010.01412 [2] https://arxiv.org/pdf/2102.11600.pdf [3] https://arxiv.org/abs/2206.04920

Q2: commentary on SWA We thank the reviewer for highlighting this recent work. We did examine SWA in our early experiments but did not find the results to be very competitive with Trust region and SAM-style training in the cross-lingual transfer context. These results were omitted from the paper purely as we considered this additional discussion to be a distractor from the contrast between TRAM and the parent SAM and TR methods. Third party feedback has also requested this comparison and we will include this in future revisions.

Here we show our results for SWA on XNLI compared to SAM/ASAM/FSAM and TRAM.

We identified several trends within our SWA results which warrant further study and we are actively exploring. The TL;DR of these findings is that weight averaging for optimal in-domain and out-of-domain performance is extremely variable across optimizers. We observe the same pattern as prior work in “turning on” SWA after some training, to average later checkpoints, works best for ID performance. However, we also identify that the best model for OOD uses only early checkpoints and training can stop after 30% of training. We also observe negligible correlation between performances for ID, OOD and where averaging starts/stops. We find SWA+SAM to perform below SAM alone but the inverse trend for Adam or SGD. We consider in-depth study of these phenomena as a different contribution which we plan to address in future work.

Q3: on the abstract claims

We thank the reviewer for their feedback on the abstract and claims. We will revise these for the scope of our claims to be clearer. We ask the reviewer if they can clarify where they consider more insight to be necessary.

Q4: Larger model scales

We agree with the reviewer that results on larger Llama scale models would be insightful. At present, we don’t have the resources to run full fine-tuning of larger models with TRAM. In Appendix C.2, we highlight that training GPT2-XL (1.5B params) with TRAM yields a more generalizable model than other baseline algorithms including SAM, ASAM, FSAM and Trust Region methods. This is intended to provide some insight into if larger scales diminish the benefit of optimization contributions (where we find that TRAM is still beneficial at the 1.5B model scale). This echoes other recent work which finds that Sharpness-Aware optimization is still beneficial at many model scales [1]. We conjecture that TRAM would benefit a larger model as it has for 100M – 1.5B scales.

[1] Sharpness-Aware Minimization Improves Language Model Generalization https://aclanthology.org/2022.acl-long.508/

Q5: Code release

We include an anonymised code repository for critical code elements of TRAM https://anonymous.4open.science/r/tram_optimizer-0448/tram.py.

We thank the reviewer for their time and attention in considering TRAM. We provide the first part of our response to address your questions. We are still working on experiments for the vision domain which we hope will provide further insight into TRAM for the reviewer. We will update our response when these are finalized.

Q1: the notion of "Trust" in the model.

We thank the reviewer for highlighting where the explanation of trust regions can be improved in the paper. We will expand on this here and include this in future revisions to the paper. We will also improve the visual/graphical description of the method in future revisions.

The notion of the trust region is the local neighborhood around the current function where we can assume an adequately similar representation of the objective function [1,2]. Prior neural TR methods will include this constraint as a direct regularizer on the loss [2,3] without access to the representational probability density [2]. The brief interpretation of this constraint is to iteratively encourage smooth updates to the objective which retain local function similarity to the previous iterate. This regularizer translates the TR objective constraining large changes in the function space into an adjustment to the change in the parameters for a given learning step (i..e, a constraint in $\delta f$ translated into $\delta \theta$). The prior work inspiring TRAM proposes two efficient strategies to approximate the region around the current objective function [2,3] and translating this to a penalty on $\theta$. In either case (i.e., TRAM using Eq (4) or (5)), we introduce a measurement on change in the function space to inform a change in $\theta$.

Our connection to SAM [4] is based on the similarity in geometry between defining the aforementioned change in $\theta$, permitting safe change in the function space, and the $\rho$-ball domain containing sharpness maximization. As noted in [5], the domain of the SAM $\rho$-ball is somewhat arbitrary in that the setting of $\rho$ has little relationship to the parameter geometry or training dynamics. FSAM [4] reformulates SAM to respect the parameter geometry. However, in TRAM we instead propose to define this maximization domain by the space in $\theta$ which permits only “trusted” change in the function. The connection to “trust” in TRAM is to reinterpret what is originally a regularizer in [3,4] to define a search space where SAM-style optimization is permitted. SAM-type algorithms do not contain this notion of “trust”, as the maximization domain has no relationship to the function space and cannot ensure similar smooth changes in function density (e.g., there is no control for staying close to the original space within SAM/ASAM). Therefore, training with SAM does not ensure, or optimize for, smooth curvature between pre-trained and fine-tuned functions. We note that SAM-style algorithms could yield this assurance empirically with a large $\rho$ to always contain the trust region (i.e., Eq (2) > Eq (7) for all $\theta$ and $t$). However, we empirically observe that this is not the case — a larger setting of $\rho$ for SAM produces poorer results than TRAM. TRAM could be considered a subsolution of ASAM—our contribution highlights that sensitivity to the function space trust region produces better empirical outcomes with less catastrophic forgetting.

[1] Better Fine-Tuning by Reducing Representational Collapse https://arxiv.org/abs/2008.03156 [2] Natural Gradient Revisited https://openreview.net/forum?id=jbLdjjxPd-b2l [3] Trust Region Policy Optimization https://arxiv.org/abs/1502.05477 [4] Sharpness-Aware Minimization for Efficiently Improving Generalization https://arxiv.org/pdf/2010.01412.pdf [5] Fisher SAM: Information Geometry and Sharpness Aware Minimisation https://proceedings.mlr.press/v162/kim22f/kim22f.pdf

We thank the reviewer for their time and attention in reviewing our work. We now address each question in turn.

1. Theoretical Motivation

Please see our response to 2RUf Q1 for further insight on the notion of Trust and cqYa Q1 for discussion on the generalisation bound.

2. Run variance

We appreciate the concern for the run variation. At present, we do not have sufficient compute bandwidth for repeated runs of the LM objective but will endeavour for this in future revisions. All results for XNLI are averaged over 20 runs, which we consider sufficient for our claims and many more runs than is often claimed for experiments on cross-lingual transfer. The variance across different languages is more indicative of the task challenges than the run variance. We note that our vision experiments provided below do not report notably higher confidence intervals than other methods.

3. Vision modality

For comparison, we implement the same experiments as [1] for FisherSAM training ViT-base-16 [2] for CIFAR-100 [3], Flowers [4] and Cars [5]. We fine-tune ViT-base-16 for 200 epochs training with TRAM and report the Top-1 accuracy(+ 95% confidence interval) averaged over 5 runs to compare directly to Table 3 in [1]. We match the hyperparameter setup of [1] — the base optimizer is SGD with an initial learning rate of 5e-4 and a cosine decay LR schedule. Due to computational constraints, we only have the bandwidth to report the two main variants of TRAM on these tasks during the rebuttal. We plan to update with the other variants of TRAM (TRAM-$d_{\theta_0}$ and TRAM-Fisher) in future revisions.

We observe that TRAM performs competitively across all datasets, with one or both variants of TRAM performing above all other methods. The largest improvement for TRAM is Flowers where we perform +1.81% above FSAM. The smallest improvement is for Cars where we perform 0.07% above FSAM, but we not that the confidence intervals for these results do not overlap.

[1] Fisher SAM: Information Geometry and Sharpness Aware Minimisation https://proceedings.mlr.press/v162/kim22f/kim22f.pdf [2] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale https://arxiv.org/abs/2010.11929 [3] Learning Multiple Layers of Features from Tiny Images https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf [4] C3D Object Representations for Fine-Grained Categorization http://vision.stanford.edu/pdf/3drr13.pdf [5] Automated Flower Classification over a Large Number of Classes https://ieeexplore.ieee.org/document/4756141

We thank the reviewer for their time and attention in considering TRAM. We provide the first part of our response to address your questions. We are still working on experiments for the vision domain (Q1) and will update our response when these are finalized.

Q2. "is ASAM type of updates really necessary?"

We thank the reviewer for correctly pointing out that TRAM can build on either ASAM or SAM. We pursue building on ASAM due to empirical benefits of ASAM vs SAM in our early experiments. We experimented with both techniques and found that just as ASAM generally outperforms SAM, TRAM building on ASAM outperforms TRAM built on SAM. We show XNLI accuracy for TRAM without and with adaptive scaling. This corresponds to TRAM using SAM or ASAM respectively. TRAM using SAM scaling still performs above most prior methods (SAM, ASAM, FSAM, TRPO, MESA) but below TRAM using ASAM.

Q3: "whether SAM-type of updates are required in the beginning of the training or at the end of the training."

We thank the reviewer for highlighting the ongoing debate of when to apply SAM-style optimization. Given that “pretrain then fine-tune” is a standard paradigm in NLP, and more commonly in CV, we focus our efforts on applying TRAM here to learn a task-specific model. Our understanding of pre-training dynamics (in NLP) is that an algorithm requiring multiple forward and backward passes would be too expensive for large scale pre-training from scratch. Furthermore, the benefits of these steps may be weaker than simple additional pre-training on more data (i.e., the bitter lesson). We consider the value of TRAM greatest when adapting a task-agnostic model (or mostly task agnostic) to a task-specific model. A principle of TRAM is to maintain generality from pre-training, therefore requiring some task-agnostic initial training to obtain a “trustable” original model.

Q4: Fig 1.

The slope of these curves represents the correlation between the training domain (ArXiv Math) and different groups of zero-shot evaluation domains. These curves represent the trend of domain correlation across all optimization algorithms (SAM-style, trust region and TRAM). For a positive correlation, a point above the curve indicates performing better than the trend. For a negative correlation, a point below the curve indicates a better performance than the trend. We do this to highlight that TRAM beats expectations of how negatively-correlated domain transfer impacts performance for hard OOD cases.

Q5: Table 5

Table 5 includes both “main” variants of TRAM which directly compute a TR distance. We omitted the other variants due to last minute space issues. We provide the numbers below and will include these updates in revisions. In general, TRAM-$d_{\theta_{0}}$ and TRAM-Fisher align with our trends comparing TRAM and other algorithms. All variants of TRAM have among the lowest sharpness and the highest CKA similarity across all variants.

Q1. Sanity check experiments

For comparison, we implement the same experiments as [1] for FisherSAM training ViT-base-16 [2] for CIFAR-100 [3], Flowers [4] and Cars [5]. We fine-tune ViT-base-16 for 200 epochs training with TRAM and report the Top-1 accuracy(+ 95% confidence interval) averaged over 5 runs to compare directly to Table 3 in [1]. We match the hyperparameter setup of [1] — the base optimizer is SGD with an initial learning rate of 5e-4 and a cosine decay LR schedule. Due to computational constraints, we only have the bandwidth to report the two main variants of TRAM on these tasks during the rebuttal. We plan to update with the other variants of TRAM (TRAM-$d_{\theta_0}$ and TRAM-Fisher) in future revisions.

We observe that TRAM performs competitively across all datasets, with one or both variants of TRAM performing above all other methods. The largest improvement for TRAM is Flowers where we perform +1.81% above FSAM. The smallest improvement is for Cars where we perform 0.07% above FSAM, but we not that the confidence intervals for these results do not overlap.

• [1] Fisher SAM: Information Geometry and Sharpness Aware Minimisation https://proceedings.mlr.press/v162/kim22f/kim22f.pdf
• [2] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale https://arxiv.org/abs/2010.11929
• [3] Learning Multiple Layers of Features from Tiny Images https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf
• [4] C3D Object Representations for Fine-Grained Categorization http://vision.stanford.edu/pdf/3drr13.pdf
• [5] Automated Flower Classification over a Large Number of Classes https://ieeexplore.ieee.org/document/4756141

We hope these experiments provide additional utility to the benefits of TRAM. We ask the reviewer to consider revising their score if these additional responses have sufficiently addressed your questions.