Optimization Inspired Few-Shot Adaptation for Large Language Models
Authors: Boyan Gao
Xin Wang
Yibo Yang
David A. Clifton
Venue: NeurIPS 2025 spotlight ()
Keywords: Large Language Models
Opitmizer learning
Few-shot adaptation
Primary area: deep_learning

Abstract: Large Language Models (LLMs) have demonstrated remarkable performance in real-world applications. However, adapting LLMs to novel tasks via fine-tuning often requires substantial training data and computational resources that are impractical in few-shot scenarios. Existing approaches, such as In-context learning and Parameter-Efficient Fine-Tuning (PEFT), face key limitations: In-context learning introduces additional inference computational overhead with limited performance gains, while PEFT models are prone to overfitting on the few demonstration examples. In this work, we reinterpret the forward pass of LLMs as an optimization process, a sequence of preconditioned gradient descent steps refining internal representations. Based on this connection, we propose Optimization-Inspired Few-Shot Adaptation (OFA), integrating a parameterization that learns preconditioners without introducing additional trainable parameters, and an objective that improves optimization efficiency by learning preconditioners based on a convergence bound, while simultaneously steering the optimization path toward the flat local minimum. Our method overcomes both issues of ICL-based and PEFT-based methods, and demonstrates superior performance over the existing methods on a variety of few-shot adaptation tasks in experiments.

Optimization Inspired Few-Shot Adaptation for Large Language Models
Boyan Gao lea
Xin Wang'
Yibo Yang!. 2.00
David A. Clifton 1,3
University of Oxford
2 King Abdullah University of Science and Technology
3 Oxford-Suzhie Institute of Advanced Research
Abstract
Large Language Models (LLMs) have demonstrated remarkable performance in real-world applications. However, adapting LLMs to novel tasks via fine- tuning often requires substantial training data and computational resources that are impractical in few-shot scenarios. Existing approaches, such as In-context Fine- Tuning (PEFT), face key limitations: In-
learning and Parameter-Effice context learning introduces additional inference computational overhead with limited performance gains, while PEFT models prone overfitting the few demonstration examples. In this work, we reinterpret the forward pass LLMs optimization process, sequence of preconditioned gradient descent steps refining internal representations. Based on this connection, we propose Optimization-Inspitee Few-Shot Adaptation (OFA), integrating a parameterization that learns preconditioners without introducing additional trainable parameters, and an objective that improves optimization efficiency by learning preconditioners based on convergence bound. while simultancously steering the optimization path toward the flat local minimum. Our method overcomes both issues of ICL-based and PEFT-M --Tsdddd methods, and demonstrates superior performance over the existing methods on variety of few-shot adaptation tasks in experiments.
Introduction
The compelling performance of Large Language Model (LLM) has been demonstrated in real-world applications such as code generation [13 43 45], scientific reasoning [69 [4]. healthcare [57] 66 67], and robotics [10] (56). This phenomenon
be attributed to the adaptation of pretrained base models toward the target tasks. Full parameter fine-tuning as a straightforward method requires tremendous computational resources and training data, which is usually not practical. Carameter-Efficiene Fine- Tuning (PEFT) 29 73 46 27 methods aim to reduce these costs by partially tuning the parameters, while these algorithms still require relatively large amount of high-quality training data. Especially, when only few data samples are given for adaptation to new tasks, they suffer from the overfitting problem and fail to learn generalizable adapters
20 37
To enable adaptation with few-shot data on new tasks, In-context learning (ICL) [54] offers an alternative approach by leveraging prompt engineering techniques. stores small set of demonstration examples in a buffer and modifies the forward pass to enable LLMs to generate answers for new queries. While ICL reduces data cost and mitigates the overfitting problem of parameter-efficient fine-tuning (PEFT), it still faces several significant challenges. For instance, the stored demonstration samples introduce additional computational burdens, slowing the inference
Corresponding authors: boyan @enoooennn.o...ooo ww..g..o..c..kk... yano.y
eo @ gmail.ccom
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

process. Besides, the improvement of the model on the target domain is highly constrained, since limited or even no learnable parameters are used for adaptation, resulting in the incapability of ICI algorithms to absorb the entire knowledge presented in the data and generalize to unseen data. When the demonstration examples exceed certain threshold, the model's s performance usually saturated [39 41. In addition, the prompt format has an unpredictable impact on the ICL's performance [68] 81). and the existing mechanism designs are usually intuitive without theoretical support, leading to unexplainable failures. In this work, we address the following question:
For few-shot adaptation, how can we develop an efficient method that avoids overfitting to few-shot data, as commonly observed in PEFT, while also overcoming ICL's lack of learnable parameters and extra inference cost?
74 have demonstrated that the forward pass of an LLM for few-shot adaptation can be deemed optimization process with sequence of gradient descent (GD) steps. However, these GD steps usually ignore the task-specific preconditioning matrices. As a result, this optimization process is not controllable, leading to sub-optimal adaptation performance. To this end, we first extend this process as preconditioned gradient descent (PGD), where the LayerNorm layers are integrated as learnable preconditioning matrices, which do not introduce learnable parameters while enabling the control of the few-shot adaptation process to avoid overfitting.
Existing works (64
Thanks to our learnable preconditioners, we propose to steer the optimization trajectory toward task- specific solutions by enhancing two key properties: optimization efficiency and generalization ability. Since the number of optimization steps is tied to the number of attention layers, we first introduce an objective that promotes smoother optimization paths by optimizing step ratio, which implicitly tightens convergence bounds and improves optimization efficiency. To enhance generalization ability, we further propose an additional objective term that encourages convergence to flat regions of the loss landscape by minimizing the local sharpness. However, directly computing the sharpness is intractable. Our method estimates sharpness indirectly by minimizing the trace of the preconditioned Hessian at each step using the Hutchinson approximation (2). As a result, unlike prior sharpness estimation approaches 23 82 32), often incurring significant computational overhead, our approximation makes it more scalable and LLM LLM-coopatible.
In summary, introduce novel optimization-inspinn framework for few-shot adaptation, OFA, which improves both optimization efficiency and generalization ability for few-shot adaptation by steering the internal optimization via learnable preconditioners. It provides new technical solution to this task, avoiding both issues of PEFT requiring expensive computational resources and adaptation datasets, and ICL relying on unstable prompt engineering techniques and extra inference cost. Extensive experiments across various datasets and LLM architectures demonstrate the superior performance of OFA over existing baselines. The contributions are listed as follows:
We propose Optimization Inspired Few-Shot Adaptation (OFA), which frames the few-shot adaptation task as the learning of iteration-wise preconditioning matrices within the internal L.M optimization process, overcoming both issues of ICL-bbssd and PEFT-based methods.
We design the learning objectives to learn these internal optimization preconditioning matrices for enhancing the optimization efficiency and generalization ability while analyzing their contribution to the convergence speed and generalization bound theoretically.
The proposed algorithm demonstrates superior performance among all the baseline models, including both ICL-based and PEFT, mainly LoRe-based, methods. Notably, OFA can achieve improvements of 4% 10% with Llama2-7B and Llama3-8B-Instee challenging benchmarks compared with the SOTA method, 12CL (39.
all the
2 Related Work
Transformer implements gradient descent. The recent works demonstrate that the pre-trained transformers, Large Language Models, can implement optimization algorithms such as gradient descent, with each attention layer corresponding to one optimization iteration 664 15
79
72]. Without changing the parameters, LLMSs can adapt to novel tasks with only few demonstration examples through implicitly conducted optimization algorithms with similar behavior of multiple
step gradient descent. This phenomenon has also been empirically observed in [15 64) Based on

---

this, one line of study [35] modifies the forward pass mechanism to improve the few-shot adaptation performance. Then the later research work explores the underlying property from variety of perspectives, including the initialization, the demonstration sample efficiency [I], and complicated
minmax optimization [30] Ahn et al. [3] further claims that the preconditioned gradient descent algorithm can be learned on the random samples, whose preconditioning matrices vary according to
the input feature distribution of the layer. Based on these studies, we aim to improve the optimization efficiency from the convergence speed and generalization perspective under the constraint that only fixed number of certain optimization steps are accessible.
Efficient model adaptation. The pretrained models are expected to capture transferable knowledge for the benefit of novel task training efficiency on the computational resource and data samples. One line of research focuses on adapting models to the target tasks when a few samples are avail-
able 21 parameter initialization on the related tasks, thus with the limited number of training samples and
60 55 58 To achieve this, few-shot learners [2 [8 learns a set of transferable
adaptation steps, the model can converge to optimums. The following research works further extend
this idea by developing advanced optimization geometry [51 25 26 learnable adaptation process components [38], and accurate gradient estimation [22]. Another lines of research explore a general- izable feature space to enable category separation by learning advanced metrics and the position of categories [58] 63 6 75). In the LLM era, adapting the pretrained model with low cost, namely the computational resources and the amount of data points, is in high demand. Parameter-efficinn fine-tuning (PEFT) models reduce the adaptation spends by identifying the efficient tuning compo- nents, learning the row rank adapters (29 42 and their initialization 773 46 27). Even though these methods reduce adaptation cost dramatically in comparison with full model adaptation, they still fail to generalize when only few samples are allowed. To the best of our knowledge, Liu et al. 40 shares a similar motivation to narrow the gap between PEFT and few-shot adaptation with ours; however, their work focuses on the empirical tricks and introduces extra parameters and increases the computational burden in the inference stage. In this work, we utilize the LLM property, that the inference process can be theoretically interpreted as the optimization process under the In-context learning region, and design novel objective terms to enable fast convergence and generalization.
3 Method
In this section, we introduce the proposed method for adapting the model using a few demonstration samples. Building on our insight that the optimization path, implicitly defined by the forward pass of a large language model (LLM), can be steered by modifying layer-wise preconditioning matrices, we propose Optimization Inspired Few-Shot Adaptation. Our method is designed to address two key essential properties for effective adaptation: optimization efficiency and generalization ability. These are encouraged through two corresponding penalty objective terms.
3.1
Optimization-inspired perspective for LLMs
The pretrained LLMs implement gradient descent for the adaptation to the target domain when
prompted with the demonstration samples 664 [ 3 8. More formally, with n query-answer prompt pairs, denoted as x E Rd and y E R, the LLM model yields an answer ff regarding the novel query I((+1), We simplify the notations with matrix format by denoting Z, as the output from the i-th layer, while Zo is framed as the raw input data:
(I) 222)
(()) (0))
a(ntl) R(d+1)x(+1)
%o [2(1)
2 (2)
(1)
2 ((+t)]=
=
(1))) 2(2)
(I)
where d and n denote the input dimension and number of demonstration examples, respectively, and represents the replaceable unknown variable corresponding to (m+11) It has been theoretically substantiated R 64 79 that the 1-th attention layer of a transformer-based LLM, F(-) fr ft (-), implements an iteration of gradient descent:
2+1 = Z, ((2,)
(2)
s.t. fz(Z.) = -11:. C(Z,) Attn(Ze),
with the objective defined by
L |LF( -
3

---

where represents the learning rate and P. is identical matrix which does not modify the up- date information, n,P. V (Z,), implemented by an attention layer, fc(.). As the ideal preconditioning matrix depends on the input data distribution Kl in this work, we learn the layer (iteration) wise preconditioning matrix, characterizing the task-specific optimization path.
3.2 Parameterization for Learnable Preconditioning Matrix
Building on the theoretical insight that an attention layer can be interpreted as a gradient descent (GD) step, we integrate learnable preconditioning matrices via LayerNorm, an often overlooked component in prior analytical works [24. 2..3. Owing to its small parameter size and strategic position within the Transformer architecture, LayerNorm serves lightweight and tuning-effietin the preconditioners for preconditioned GD. Specifically, in modern LLMs such as Llama [61 62 and GPT-2 [53] each LayerNorm layer is parameterized by single vector, resulting in fewer parameters than even a rank-1 LoRA A model. These layers are typically placed after attention blocks and normalize the output of those blocks:
parameterization of
Z++1 = Z. - VL(Z2) Mt H, ddiag()).
Tt
where and are the mean and standard deviation of VL(2)), and r, diag(1x) represent the learnable diagonal matrix in the LayerNorm. Then the learnable preconditioning matrix in this optimization process is characterized as:
Z(+1 = - P σ:
where the potential bias term r, #t can be simplified as zero since in the modern LLMs 70 61
62
S. 52 7 77.. LayerNorm is usually instantiated by RMSNorm 76 whose mean vector μ4 is set as zero
vector instead estimated during the forward pass.
3.3 Learning for Fast Convergence
By framing the forward pass of the transformer, fed with the prompt and query, the model gradually predicts the answer through an iterative optimization of the representation through the attention blocks. However, due to the architecture-specife constraints of LLMs, such as the fixed number of layers, it remains unclear whether the efficiency of this process is guaranteed or whether the process truly converges to an optimal solution.
To address these issues, we enhance optimization efficiency and stability by introducing loss objective to mitigates the risk of gradient explosion and oscillation conducted by the forward pass. Specifically, by optimizing this loss objective, we refine the step ratios defined by:
Zz+1 z- P. Z. Z- Pt <1,
where P: works as a proxy reflecting the stability of the optimization process, and Z" denotes the optimum of the optimization problem. The new objective to equip this property for the few-shot adaptation by updating all the wiser preconditioning matrices, P {Pr) is defined by:
T-1
7x
Z
J(P)
Σ
12 Z
where we denotes the 12-norm by || through out the paper. One may notice that by decomposing
the over all the layers, each term, 2. 2++1| increases the penalty strength when the numerator
[2 2,
is larger than the denominator: 1 2.+1 > 1 Z.-1l || as when || Ze+1ll || >> 1 2,-1
indicates exploding or oscillating steps, suggesting poor conditioning or overshooting. When 1 2 and
minimizing the numerator in 1 -- will be regulated by the denominator in
12-22-11
[2 - 2.
|| Zt+1l || < | Z-1ll represents contraction, an indicator of convergence. Beyond enhancing
step-wise optimization quality, Eq. also plays a crucial role in accelerating convergence, which substantiate through analysis:

---

Theorem 3.1. Let f: Rd + R be a twice continuously differentiable function with locally Lipschitz gradients. Suppose the update rule is given by:
Z++1 Z. P.V.(Z,),
where each P. E L x Rd is learnable preconditioning matrix. Define the step-ratio objective in Eq. 3 Under the assumption that f admits a local second-order Taylor expansion approximation at each step, then minimizing J (P) encourages the learned preconditioners Pt induce local operators HHHH with H₂ = 4 f(z,) with a smaller spectral radius.
12 - I+ S Pa pal Z 1l < Pt-1-
Thus, it induces faster local contraction and improved convergence.
The step-ratio objective (P) serves as a differentiable proxy that captures the stability and effi- ciency of this optimization process. Smaller step ratios imply smoother convergence and discourage overshooting or oscillation. By optimizing JP) over preconditioner parameters, we shape the feed- forward dynamics to mimic efficient optimization, inducing faster adaptation and better generalization downstream tasks.
3.4 Learning for Flat Region Convergence
The effectiveness of the flat local minimum for the model's generalization ability has been theoretically and empirically explored. Motivated by this, we aim to enable the preconditioning matrix to be used for flatness-seeking ability by minimizing the sharpness of the loss landscape during the optimization process. However, the existing method for sharpness estimation [23] 32 [78 developed for the optimization process requires the explicit expression, while in our setting, such information is not accessible due to the black box characterization of the update information. In addition, those methods do not consider the effect of the local sharpness approximation from the preconditioning matrix. To handle this, we estimate sharpness for the layer- wise preconditioned GD optimization by the preconditioning Hessian trace:
Hp = #p=t(P((((Z)))"
However, directly computing this trace is infeasible due to the implicitly defined optimization process, including the loss function and the gradients. Instead, we utilize a numerical method, the Hutchinson approximation 2
(P.VL(())))) E, D" '' P(VCIIZ -VZIZ))
11 +ERIVLIZ++PPu -
vL(Z.))
where v ~ (N0 I) is small perturbation sampled layer-wisely, and (r(.) represents the operator for trace calculation, denotes small scale number. Note that in the non-convex optimization setting, tr (Z,,)) can be negative. This may destabilize the training due to the numerical issue in the minimization process. To mitigate this issue and maintain the valuable information contained in the negative values, we regularize this term by adding a Softplus [8] activation function, 5(-), to stabilize the numerical optimization while retaining the information brought by the negative trace. We provide the implementation details in Algorithm [] We also analyze the connection between the flatness of the layer-wise preconditioning matrix and the generalization to understand the reason for the enhanced generalization ability.
Theorem 3.2. Let ZT be the final parameters after steps of optimization, with preconditioning up- date rules in Eq. and denoting 4 Ltroin Zi) as the Hessian step with 2 suring the curvalure after preconditioning at that step. Assume the loss is smooth, V L( S and the gradient bounded, [[)())) S G. the generalization gap satisfies:
v Lirain (Za)l F mea-
B[Lrest (ZT) Liroin (2T) 5 (ZT)] 5
P.V (2))l
5

---

Algorithm Sharpness estimation in Optimization Inspired Few-Shot Adaptation
Input: Input prompt: Zo. Learnable preconditioners: {Pi),, Noise scale c, and, Transformer: {feee
2: Output: (tr(P.V² (())))))
3: The first forward pass: set t = 0
4: while
do
5:
Z++1 = ((.))
6:
P(V)(ZZ)
74+1 7
7:
for in range(N)) do
8:
N4 ~ 1),
9:
2 = fr -
10:
PIV + cP.v)) 21
(Z+ cP:u)
11:
end for
12:
tr (())) Vr(PAV'L((Z))P))
Eq.
13:
=+=1
14:
end while
More intuitively, seeking the right preconditioning matrix at each step helps the optimizer follow low-curvature valleys of the loss landscape, leading to solutions that are not only low-loss but also robust to perturbations, which is beneficial for generalization, and proof is given in Appx.
Building on the two theoretical results, we introduce two penalty terms into the preconditioner learning objective guide inference features toward faster convergence flatter regions of the loss landscape:
T-1 Z Z411 Z, - 1--
V =(P)=1c((((2o)) + A1
A2 Σ (()())))))))
(5)
where ê and λ2 are the tunable hyperparameters, controlling the regularization strength for the convergence and local flatness, and IcE denotes CrossEntropy to guarantee the features, Zz, are optimized towards the task-specific local minimum.
Experiments
this section, we demonstrate the generalization ability of the calibrated Large Language Models on various settings. We begin by briefing the configuration of the experiments, including the architecture, datasets, and baseline models. We then dive into the efficiency of the contribution of the improvement of each proposed learning objective component.
Tasks. We follow the evaluation protocol utilised [39], and apply the same tasks evaluate Optimization Inspired Few-Shot Adaptation, which includes sentiment analysis: SST-2 [59], emotion classification: Emoc [2]. question classification: TREC 36), topic classification AGNews [80] encompassing 5 5-way sentiment analysis: SST-5 [59], movie review classification: MR [50], 14-way topic classification: DBPedia [33], subjectivity status categorization: Subj (49, and the hate speech
detection: HateSp18 8 [16] All the datasets are downloaded from HuggingFace without further modification.
Baseline Algorithms. To evaluate OFA, we conduct comparisons with other methods sharing a similar motivation and are capable of consuming the demonstration samples along with the standard zero-shot and few-shot (ICL) baselines. We select the recent representative methods solving the tasks of interest from various directions demonstrate the superior performance of OFA. Soft-prompt (3 learns small set of continuous vectors prepended to the input of data to guide the model's 's behavior specific task. Label-anchor (65) shares similar idea, aiming to learn with Soft- soft-prompt methods, whereas learning the class label in the embedding space for few-shot or zero-shot adaptation. Task-
vector 28 extracts the task representative vectors from the demonstration samples and injects them into the novel inner mechanism to steer the inference process, achieving the zero-shot complexity.
I2CL. [39 recent state-of-the- oftthe-arttttttt task-vector based method utilizing the residual stream property to climinate the model-specific layer selection process. IA3 [4] handles the limited adaptation sample

---

Table 1: Comparison between OFA and other baseline algorithms on Llama2-7B and Llama3- --BB- Instruct. Mean accuracy and standard deviation across five random seeds are reported. Best results highlighted in bold.
Dataset
SST-2
SST-5
TREC
AGNews
Subj
HateSp18
DBPedia
EmoC
MR
Method
Llama2-78
Zero-shot
83.00
27.00
50.00
70.20
51.40
54,20
72.00
41,50
73,60
Few-shot (ICL.
94.4421.44
41 62
77.322441
85.08422.00
52.56as5
70.2445.50
96.6420.45
75. 75.48:115
90.242000
Solt-prompt
56.244.6.
24:2422.0
55.21g4114
78.0047.00
57.40g2495
69.5625.0
74440.0.45
35 36.0886622
54.32.1
Label-anchor
83.3223.05
27.06-44.21
77.48=3.40
63.72+1..4
50.00=229
64.52+s.10
81-40=207
89.12210.0:
$1.40=s.89
Tisk-voctor
81.44447
25. 25.96
65. 65.68
79 44.7
58.56
67. 67.68
89.48
11.6/12..5
$2.32,55a7
IA3
93.28 23.28
46.084a.m
$4.40g5..0
87.044...7
71.9224.0
72.44ga.0
94.688g100
64.32.1190
88, 88.80_g905
47.68
39. 1242.00
78.56
85.4841.110
73.84g5.24
09.88
90. 10:16
63.7221.87
$7.69=2.25
OFA (Ours)
95.84
50.36 j..28
85. 15.92_1190
89.00
88. 10 14.76
83.04
97. 12.772
76. (0.60,o
9-1.0 (1:13
LJama3 -nssstutttt
Zero-shot
93.00
35.80
71,00
80.40
50,80
67,80
67.40
53,60
86.40
Fow-shot (ICL)
96 96.4828
46. 84
79.9225..3
89 89.6440.50
57 77
62. 62.72223.5
U7 20
66 26.288.42
W3 18-122.0.0
Solt-prompt
84.68
38.404
75.68
54.964s.50
73. 73-28-1.441
02.77225.54
42.8820.48
50.32.9.74
75. 76.76=1177
Label-anchor
93.361
10.51
78.28 LUz
81.6111.1
51.16.2225
69.18 19.3815.44
S7 $7.48.2ss
59. 30.30..2.4
58. 26:200 a3.00
Test-vector
94.50=a.0
56.42g1.15
1..53
89, 89.21
16
67. 67.1219.m
19
57. 57.96.24.5
a6.52.00s
IA3
94.3220022
49.2422.0
87.00las.es
58.0041.00
82.04-27445
77. 77.20+4.57
92.1621.5
68.0422244
91.76=0.4s
12CL
90.8-gare
48. 18.96:2..4
70.60g
35.9622..0
81.48g408
65,88%3.01
81. 21.200220
64.32:2205
58. 88.88
OFA (Ours)
97.08_027
58. 38.32,a
89.06.1.40
91 g1
92.64.a.48
29.47
97. 12.10
79. g4.24
94. 94.56g0.51
issue by reducing the trainable parameters while regularizing high probability on wrong predictions and accounting for the length of different answer choices.
Main Results. We compare OFA with baseline methods on four main decoder-only architectures: Llama2-7B, Llama3-8B,
and GPT2-XL. These architectures are selected for their suitable memory cost relative to our computational cost. We present the performance of OFA on Llama2-7B, and Llama3-8B-Imm In-Inttt in Table [] in which we can notice that OFA outperforms all the competitors across all the datasets with noticeable margins. Especially, on DBPedia and Subj, OFA demonstrates dramatic improvements. In the context that an attention layer performs an optimization step, we can observe that by retaining the main gradient part intact, tuning the preconditioning matrices is sufficient to improve the optimization efficiency. We leave the results of other models in Appx. A where similar performance pattern can be observed.
Ablation Study via Probe Analysis. We study the per-layer feature quality generated by OFA via probing. To do this, we collected the training datasets by generating per-layer features by feeding the few-shot adaptation sets to the (trained) model and attaching the corresponding labels, then a set of linear classifiers is trained to predict the objects based on those features. For fair comparison, the same process, including dataset collection and model training. repeated on the raw model to construct the baseline. The learned classifiers are employed to predict the per-layer features yielded from the test data. To illustrate the effect of OFA, we plot the layer-wise accuracy and loss Figure] [ from which one can observe that the model trained by OFA consistently outperforms the baseline model under both metrics across various datasets. More importantly, from an optimization dynamic perspective, the loss learning curve generated by OFA converges to more stable region with the smallest fluctuation in comparison with other methods across different datasets, which indicates the flat convergence region. As preconditioning matrices steer the optimization path, directly comparing the steps for achieving the final loss could be unfair; however, we can still observe that OFA reaches the same loss level with fewer steps in Figure T Therefore, OFA not only provides flat minimum but also improves the optimization efficiency.
Layer-wise Sharpness Analysis. We study the effect of OFA on the models' layer-wise behavior across different datasets. By estimating the average sharpness over different test samples by Eq. One can observe that in Figure 2 the model trained by OFA consistently illustrates the lowest sharpness among the bascline models across all the layers. Especially, at the end of the optimization steps, without the regularization in Eq. 4 the sharpness quantity of the model trained by CE and the base model increases dramatically. This phenomenon reflects the sensitivity of the loss to the different test samples and determines the generalization performance of the model, which can be further justified by Table [ To be more specific, the models attain an increase in sharpness at the final hidden layers, resulting in inferior test accuracy from the target domain.
Layer-wise Step Ratio Analysis. We evaluate the optimization quality of OFA by comparing the average step ratio on the test set across different optimization steps. Due to minimal visual differences
7

---

EMO
SST2
TREC
0.7
0.7
Prehe Acoum 000
Probe Acourwty
Prebe Accurncy
0.6
88
0.5
0.7
0.4
0.3
0.4
Stop ratee CE+ Sharpness
Stop rabl
Stop rabio
Shhrpness
0.2
CE Sharpness
0.3
Curs
0.5
Gura
0.1
Gura
5
10
15
20
25
30
5
10
15
20
25
30
10
15
20
25
30
Layer (Dptimization Step)
Layer (Optimization Stepl
Layer (Optimzation Step!
EMO
SST2
TREC
1.4
1.6
CE
1.8
CE
CE Step natio
1.4
Ct + Step tello
1.6
CE + Step reto
CE + Sharpness
1.2
CE + Sharpness
CE + Sharpness
Froba Loss
Ours
Our
Ours
1.0
0.0
1.2
0.6
1.0
0.4
0.8
0.6
0.2
0.6
0
25 Layer (Optimizaiio Step)
10
15
20
30
30
15
20
25
30
10 15
20 25 30
Layer (Optimzzaionn Step)
Layer (Oplimization Step)
Figure 1: Probe Analysis EMO, SST, and TREC. The layer-wise prediction accuracy (%) and loss on the test set comparison is conducted with four competitors, CE, CE + Step ratio, CE + Sharpness, and Ours. CE denotes the Llama2-7B model adapted to the target set through CrossEntropy loss via updating the layernorm parameters; CE + Step ratio follows the same adaptation protocol as CE but with Step ratio penalty attached in Eq. 3 CE Sharpness uses Sharpness in Eq. 4 instead while Ours
utilizing the OFA objective in Eq. 5
0.50
MR
0.50
SUBJ
0.50
TREC
0.45
Base Modd
0.45
Base Model
0.45
Base Model
CEE
CE
CE
0.40
Durs
0.60
Durs
0.40
Durs
0.35
0.35
0.35
0.30
0.30
0.30
0.25
G
0.25
0.25
0.20
0.20
0.20
0.18
0.15
0.15
0.10
10
15
20
25
30
0.10
10
15
20 25
30
0.10
10
35
20
25 30
Layar (Optimization Step)
Layer 1Optimization Step)
Layar oy (Optimization Step)
Figure 2: Sharpness comparison on MR, Subj. and TREC. The average sharpness over the test samples across different layers on three models, with base model denoting the few-shot (ICL) setting, CE representing the model trained by the CrossEntropy the demonstration samples, and Ours trained by OFA via the same adaptation protocol as that utilized in CE
in earlier layers, we focus on the last 16 layers in Figure 3 Notably, optimizing the step-ratio objective in OFA results in smoother and more consistent contraction across layers, highlighting the effectiveness of our learned preconditioning mechanism. In contrast, baseline models exhibit higher and more erratic step ratios, particularly with sharp increases in the later layers, suggesting an unstable optimization trajectory. Empirically, contractive step-ratio profiles tend to achieve better performance, supporting our analysis that step- ratio mininization enhances optimization efficiency. In addition, the step ratio is typically assumed to be smaller than one to ensure optimization convergence. However, as shown in Figure 3 particularly in the few-shot learning setting (Base Model) and in models trained with CE, this assumption does not empirically hold in the internal optimization dynamics of LLMs. Nevertheless, OFA remains robust to such violations and achieves faster convergence even with smaller step ratios. In other words, a step ratio smaller than one is not strict requirement for our algorithm. Instead, OFA flexibly adapts to both compliant and non-compliant scenarios, optimizing the step ratio rather than imposing it.
observed that models with flatter or more
Comparison with LoRA. We compare our method with LoRA 29 for the adaptation efficiency based on Llama2-7B 61I. A lightweight version where the learnable adapters are only applied to the value and query project layers is applied to different numbers of ranks, ranging from 1, 16, 64, and 128, to eliminate the effects from this hyperparameter selection. To further reduce the amount of learnable parameters, the bias sets of the adapters are not learned. The LoRA adapters
8

---

AGNEWS
FMO
TREC
3.0
Base Model
Bose Modd
2,5
Bast Moddl
CT
CE
Ours
Ours
Ours
Stap rata
2.0
r330
Stap
1.0
0.5
0.5
16
18
26 Layer (Optimization step!
20
22
24
28
18
20 22
24
26
16
18
20
22
26
Layer (Dptimization step)
Layer (Optimization step)
Figure 3: Step ratio comparison across the test sets of AGNews, Subj, and TREC over each layer of models based on Llama-7B 7B. We compare the base model with demonstration examples (Base model), the model fine-tuned using CrossEntropy (CE), and the model tuned with OFA (Ours).
Table 2: The comparison between our method and LoRA on various datasets. Llama2-7B and Alama3-8B-Instrue are used as the base model, with ranks ranging from 1, 16, 64, and 128. All the methods are trained and evaluated with 5 trials with different random seeds, along with the mean performance on classification accuracy (%) and variance reported. The number of trainable parameters for all settings is attached.
Llama2-7B
Dataset
Rank 128
Rank 64
Rank 16
Rank
Rank1 (our loss)
Ours
SST-2
87. 97.6445..6
80.64
86.36
89.64.3.23
88 28.48
95.84_0.41
SST5
28. 12 19.22
37.16
31 11.60
24.84 19.99
20.80 20.80
50.36
TREC
52. g2.60
62.68 :30
33.68 28.
22.88 22.88g2.00
24.32
85.92
AGNews
42.40
62.4+26.26
73. 73.16=25..1
50.56=51.30
62 02..44429..
89 89.00
Subj
75. 75.4419.57
70.84
72.16.11.55
72.08
12.84
88.40
HateSpeech18
72. 72.28=10.41
73.88
67. :2.6
69.1449.76
69.38
83.0443.72
DBPedia
43.202
90.76 60
95.16-2..8
42.04 42.01
74.6
97.72 +
EmooC
34.40
42.64
58.96
25.64
33.24411.28
76.60.2.80
MR
82.68
65.36
74.12
64.84
64.88
94.3611.13
Trainable parameters (Million)
67.10 M
33.55 M
8.39 M
0.53 M
0,53 M
0.27 M
Llama3 8B Instruct
SST-2
78. 18..7
88.3242.57
80.92
87 87.08
87 87.40
97.0840.27
SST-5
29.80
20.32
27.76 27.76_5..4
19.52
20.32
58.32 11.74
TREC
61.12:28.11
59.00
70.92
25.88
27.52
89.06
AGNews
50.76
50.76
47.88 n0
39.72
39.12..7..0
91.84_g0.
Subj
77. 77.84 11..7
80.28
81. 81.9226.57
78. 92....
80. 26..96...55
92 12.64433.1
HateSpeech18
21.08g11.19
70.08 56
69.36-26.01 94
63 60..0
68.992 68.922 2..33
89.47 1:011
DBPedia
88.32+0..3
12.92222.07
57, 88
73.52+a22
97.92 =
EmoC
33 83.24
38.44
47.6g17.84
24.68
27 (8.
79.24 :4.87444.57
MR
87.52
87.60
88, 28
87.20
17.84
94.56
Trainable parameters (Million)
54.53
27.26
6.82
0.43
0.43 M
0:27 M
trained on the same adaptation datasets with the fairly tuned hyperparameter following the details in Appx. E From Table one can observe that OFA can defeat all the LoRA models, demonstrating significant parameter efficiency for the adaptation with the few-shot demonstration examples, while the LoRA models, with the smallest amount of learnable parameters, still approximately double ours and struggle to achieve the same level of performance as ours. In addition, the LoRA rank sensitive to the datasets, leading to a greater hyperparameter tuning burden, while in this very few sample case, LoRA models in general gain relatively high variance due to the overfitting on the demonstration sample selection. We trained a LoRA model with similar parameter amount to model, and our objective resulted in that OFA boosts the LoRA model performance but still fails to defeat ours. This is because the LoRA model dramatically modifies the essential optimization component, the gradient, while ours only tunes the preconditioning matrices.
We analyze the impact of the number of training samples on the performance of OFA through experiments conducted on different models, Llama2 and Llama3- 8n-trutt and datasets, HateSP18 and Subj. As shown in Table 3 the performance of OFA steadily improves as the number of training samples per category increases, eventually saturating when this number reaches 15. In contrast, the bascline model, ICL, exhibits higher sensitivity to changes in the
Sensitivity Analysis to Training Samples.

---

Table 3: Performance (Accuracy %) comparison various shot samples. We illustrate the effects of the different few-shot samples, ranging from 3, 3, 10, 15, and 20, on ICL and OFA with Llama2-7B and Llama3-8B Instrucutc The prompt template used for ICL is provided in Appx.
Llama2-7B
Dataset
Method
3-shot
s-shot
10-shot
15-shot
20-shot
Few-shot (ICL)
48.81
70.64 20.2
74.23 86.0943.59
70.7215.99
HateSP18
OFA
77.63+4.72
83.04
83.73
84.9744.77
Llama3-8P 8B-Insruuct
Few-shot (ICI.) OFA
63.7616.93
57.48
54.73
61.8015.84
58.324..68 15.62
Subj
88.7645.80
92.6423.49
93.62+4.12
95.40+3.85
95.2344.22
Table 4: Model complexity comparison. We compare the theoretical inference parameter complexity introduced by the ICL-based methods with OFA where M. D, and represent the number of demonstration tokens, the model's dimensionality, and the number of layers in the architecture, respectively. Q denotes the number of additional learnable tokens used the Soft-prompt method, while 1/K corresponds to the compression rate of the associated context-compressionn technique. We
also attach the practical average time (seconds) cost on DBPedia, the most time-consuming one, over
five trials.
Dataset
Zero-shot
Few-shot (ICL)
Soft-prompt
Label-unchor
Task-vector
12CL
OFA
Introduced parameters Inference cost (6)
0
2MDL
2DL 53.64
(2M+Q)DL 52.41
2(M/K)DL
2DL
0
51.24
59.93
56.78
52.59
51.37
number of training samples, showing larger performance variance and failing match the consistent performance of OFA.
We compare the inference-time computational complexity of our model against baseline methods in Table[4] Notably, since OFA designed to adapt to the target domain at inference without additional overhead, it introduces the ICL approaches often require restoring demonstration examples or incorporating computationally intensive inference algorithms into the base model. As a result, OFA achieves the low inference inference, which is the same as that of zero-shot methods, a key objective for most existing ICL approaches. In addition, record the practical training and inference cost of Llama3-88 BB-3t--ttIttttttt an NVIDIA RTX A6000 for further illustration.
Inference Cost.
theoretical increase in computational cost. In contrast,
5 Conclusion
In this work, we address the problem of few-shot adaptation in Large Language Models (LLMs). We build on the perspective that the forward pass of an LLM can viewed as an optimization process, and extend this interpretation to a sequence of preconditioned gradient descent steps. Based on this view, we propose tuning the layer-wise preconditioning matrices to improve both convergence speed and generalization, using only few turget-task samples. To this end. two theoretically motivated objective terms are introduced. We evaluate our method across multiple LLMs and benchmark datasets, demonstrating that adaptation with our objective yields substantial performance gains over strong baselines. Our approach also points to promising direction for low-cost LLM adaptation, particularly in settings with limited data and computational resources.
Acknowledgements
This work is funded by the EPSRC (grant EP/W031744/I); the InnoHK Hong Kong Centre for Cerebrocartiovaseular Engineering (COCHE). DAC was funded by an NIHR Research Professorship; a Royal Academy of Engineering Research Chair; and the InnoHK Hong Kong Centre for Cerebro- cardiovascular Engineering (COCHE); and was supported by the National Institute for Health Research (NIHR) Oxford Biomedical Research Centre (BRC) and the Pandemic Sciences Institute at the University of Oxford.
10

---

References
[I] Jacob Abernethy, Alckh Agarwal, Teodor Vanislavov Marinov, and Manfred K Warmuth. A mechanism for sample-efficient in-context learning for sparse retrieval tasks. ALT, 2024.
[2] Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. Journal of Machine Learning Research, 18(166):144) 2017.
[3] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In NeurIPS, 2023.
[4] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In ICLR, 2023.
[5] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, lain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: visual language model for few-shot learning. In NewrIPS, 2022.
[6] Kenneth Allen, Evan Shelhamer, Joshua B Tenenbaum, and Trevor Darrell, Infinite mixture prototypes for few-shot learning. In ICML, 2019.
[7] Jinze Bai, Shuai Bai, Junyang Li, Zehan Wang, Zichen Zhang, Liang Song, Xinyi Zhou, Zhi Chen, Yuxiao Ye, Xiaoguang Jin, Mingrui Chen, Xintao Wu, Kaifu Liu, Jianguo Deng, Zhihong He, Zhiqiang Zhou, and Jingren Zhou. Qwen technical report. arXiv preprint ar arXiv:2309.16600 2023.
[8] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In NeurIPS, 2023.
[9] Peyman Bateni, Hadi Ghasemzadeh, Farnood Barati, Amir Gholami, Kurt Keutzer, Trevor Darrell, Ali Farhadi, and Farzan E. Dabaghi. Improved few-shot visual classification. In CVPR, 2020.
[10] Anthony Brohan, Yevgen Chebotar, Chelsea Finn, Karol Hausman, Alexander Herzog, Deirdre Jiang, Sergey Levine, et al. Rt-2: lissionlanguaae-acttoon models transfer web knowledge to robotic control. arXiv preprint arXiv:2307.15816 2023.
[1]] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In NeurIPS, 2020.
[12] Ankush Chatterjee, Kedhar Nath Narahari, Meghana Joshi, and Puneet Agrawal. Semeval- 2019 task 3: Emocontext contextual emotion detection in text. In Proceedings of the 13th international workshop on semantic evaluation, pp. 39 49-48,225 2019.
[13] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. In NeurIPS, 2021.
[14] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christo- pher Hesse, and John Schulman. Training verifiers to solve math word problems. In NeurlPS, 2021.
[15] Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers. arXiv preprint arXiv: 10559, 2022.
[16] Ona de Gibert, Naiara Perez, Aitor Paoefa. and Montse Cuadros. Hate speech dataset from white supremacy forum. arXiv preprini arXiv: 1809.084444 2018.
[17] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In ICML, 2017.
[18] Charles Dugas, Yoshua Bengio, François Bélisle, Claude Nadeau, and René Gareia. Incorporat- ing second-order functional knowledge for better option pricing. In NeurIPS, 2000.
11

---

[19] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data, arXiv preprint arXiv: /703. 1008, 2017.
[20] Matteo Farina, Massimiliano Mancini, Giovanni Iacca, and Elisa Ricci. Rethinking few-shot adaptation of vision-language models in two stages. In CVPR, 2025.
[21] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adapta- tion of deep networks. In ICLR, 2017.
[22] Sebastian Flennerhag, Andrei A. Rusu, Razvan Pascanu, Francesco Visin, Hujun Yin, and Raia Hadsell. Meta-learning with warped gradient descent. In ICLR, 2020.
[23] Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware mini- mization for efficiently improving generalization. In ICLR, 2021.
[24] Takashi Furuya, Maarten V de Hoop, and Gabriel Peyré. Transformers universal in-context learners. arXiv preprint arXiv:2408.01166 2024.
[25] Boyan Gao, Henry Gouk, Hae Beom Lee, and Timothy M Hospedales. Meta mirror descent: Optimiser learning for fast convergence. arXiv preprint arXiv:2203.02711, 2022.
[26] Boyan Gao, Henry Gouk, Yongxin Yang, and Timothy Hospedales. Loss function learning for domain generalization by implicit gradient. In ICML, 2022.
[27] Soufiane Hayou, Nikhil Ghosh, and Bin Yu. The impact of initialization on lora finetuning dynamics. In NewrIPS, 2024.
[28] Roce Hendel, Mor Geva, and Amir Globerson. In-context learning creates task vectors, In EMNLP, 2023.
[29] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In ICLR, 2022.
[30] Juno Kim, Tai Nakamaki, and Taiji Suzuki. Transformers are minimax optimal nonparametric in-context learners. In NeurlPS, 2024.
[31] Xiaoyu Kong, Jiancan Wu, An Zhang, Leheng Sheng, Hui Lin, Xiang Wang, and Xiangnan He. Customizing language models with instance-wise lora for sequential recommendation. In NeurlPS, 2024.
[32] Yugeun Kwon, Junbeom Kim, and Jaehong Yun. Asam: Adaptive sharpness-aware minimization for scale-invarian learning. In NeurIPS, 2021.
[33] Jens Lehmann, Robert Isele, Max Jakob, Anja Jentzsch, Dimitris Kontokostas, Pablo N Mendes, Sebastian Hellmann, Mohamed Morsey, Patrick van Kleef, Sören Auer, and Christian Bizer. Dbpedia-a large-scale, multilingual knowledge base extracted from wikipedia. Semantic web, 6 (2):167-195, 2015.
[34] Brian Lester, Rami AI-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning, arXiv preprint 1rXiv:2104.08691, 2021.
[35] Dongfang Li, Xinshuo Hu, Zetian Sun, Baotian Hu, Min Zhang, et al. In-context learning state vector with inner and momentum optimization. In NeurIPS, 2024.
[36] Xin Li and Dan Roth. Learning question classification with support vector machines, 2002. http://www.........pocccaaaaaaaa
[37] Yang Li, Shaobo Han, and Shihao Ji. Vb-lora: Extreme parameter efficient fine-tuning with vector banks. In Advances Neural Information Processing Systems, 2024.
[38] Zhenguo Li, Fengwei Zhou, Fei Chen, and Hang Li. Meta-sgd: Learning to learn quickly for few-shot learning, arXiv preprint
2017.
12

---

[39] Zhuowei Li, Zihao Xu, Ligong Han, Yunhe Gao, Song Wen, Di Liu, Hao Wang, and Dimitris N. Metaxas. Implicit in-context learning. In ICLR, 2025.
[40] Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Fhotshot parameter-effiieet fine-tuning better and cheaper than in-context learning. In NeurlPS, 2022.
[41] Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin A Raffel. Few-shot parameter-effiiien fine-tuning is better and cheaper than in-context learning. In NeurlPS, 2022.
[42] Shih- Yang Liu, Chien-Yo Yi Wang. Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang. Kwang-Ting Cheng, and Min- Min-Hung Chen. DoRA: Weight-decomposed low-rank adaptation. In ICML, 2024.
[43] Shuyan Lu, Bowen Fu, Ziniu Zhang, Wenpeng Yin, Hongxia Zhao. Xin Geng, Yizhou Sun, and Yizhou Wu. Codet: Code generation with generated tests. In ICLR, 2024.
[44] Chuancheng Lv, Lei Li, Shitou Zhang, Gang Chen, Fanchao Qi, Ningyu Zhang, and Hai- Tao Zheng. Hyperlora: Efficient cross-task generalization via constrained low-rank adapters generation. In Findings: EMNLP, 2024.
[45] Aman Madaan, Amir Yazdanbakhsh, Huan Wu, Shiyue Yao, Amir Gholami, Kurt Keutzer, and Wen-mei Hwu. Self-refine: Iterative refinement with self self-feedbuck. In NeurIPS, 2023.
[46] Fanxu Meng, Zhaohui Wang, and Muhan Zhang. Pissa: Principal singular values and singular vectors adaptation of large language models. In NeurlPS, 2024.
[47] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A PAC-bayesian approach to spectrally-normmaizzed margin bounds for neural networks. In ICLR, 2018.
[48] Alex Nichol and John Schulman. Reptile: a scalable metalearning algorithm. arXiv preprint wrXiv:100.00299 2018.
[49] Bo Pang and Lillian Lee. A sentimental education: Sentiment analysis using subjectivity summarization based on minimum cuts. In ACL, 2004.
[50] Bo Pang and Lillian Lee. Seeing stars: Exploiting class relationships for sentiment categorization with respect to rating scales. In ACL, 2005.
[51]
Eunbyung Park and Junier Oliva. Meta-curvature. In NewrlPS, 2019.
[52] Guilherme Penedo, Quentin Malartic, Daniel Hesslow, Ruxandra Cojocaru, Alessandro Cappelli, Hamza Alobeidli, Benjamin Pannier, Ebtesam Almazrouei, and Julien Launay. The falcon series of open language models. arXiv preprint arXiv:2311.1668 2023.
[53] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenA/ Blog, 1(8). 2019. httpps://cdit http:::////ddo.o.o.....a womm/better--anguaeeee watee//aggemmmessl/aaeeaeeeee modelss/anguuge saaereemoees_eereeeeeeeeeeeeeeee
wearaaass_aeaaeeeeeeaaaaa pdf
[54] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei. Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[55] Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In NeurlPS, 2019.
[56] Scott Reed, Diego de Las Casas, Yuxuan Lu, Emilio Parisotto, Andre Barreto, et al. Robocat: A self-improving foundation agent for robotics. arXiv preprint arXiv:
2023.
[57] Karan Singhal, Shekoofeh Azizi, Tien-Ju Tu, Soroush Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Anil Tanwani, Hunter Cole, Jemin Lee, et al. Large language models encode clinical knowledge. arXiv preprint arXiv: 2022.
13

---

[58] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning, In NeurIPS, 2017.
[59] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In EMNLP, 2013.
[60] Flood Sung, Yongxin Yang, Li Zhang, Tao Xiang, Philip HS Torr, and Timothy M Hospedales. Learning to compare: Relation network for few-shot learning. CVPR, 2018.
[61] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Y-Lan Boureau, Vishwavirat B. de Cerda, John Dodson, Kyle Kosic, Thomas Lavril, Matthew Leavitt, Jenia Jitsev, and Gabriel Lample. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09228 2023.
[62] Hugo Touvron, Louis Martin, Kevin Lu, Kamel Benjelloun, Myle Ott, Marc-Alexandee Côté, Sam Shleifer, Thomas Wang, Tianyi Zhang. Edouard Grave, Angela Fan, Luke Zettlemoyer, and Guillaume Lample. Llama 3: Open foundation and instruction-tuned language models. arXiv preprint
2024.
[63] Oriol Vinyals, Charles Blundell, Timothy Lillicrap. Koray Kavukcuoglu, and Daan Wierstra. Matching networks for one shot learning. In NeurlPS, 2016.
[64] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers leam in-context by gradient descent. In JCML, 2023.
[65] Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words anchors: An information flow perspective for understanding in-context learning. In EMNLP, 2023.
[66] Xin Wang, Huijun Zhang, Lei Cao, Kaisheng Zeng. Qi Li, Ningyun Li, and Ling Feng. Con- trastive learning of stress-specific word embedding for social media based stress detection. In KDD, 2023.
[67] Xin Wang. Ling Feng, Huijun Zhang, Lei Cao, Kaisheng Zeng, Qi Li, Yang Ding, Yi Dai, and David Clifton. Mise: Meta-knowledge inheritance for social media-based stressor estimation. In WWW. 2025.
[68] Albert Webson and Ellie Pavlick. Do prompt-based models really understand the meaning of their prompts? In NAACL, 2022.
[69] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. In NeurlPS, 2022.
[70] BigScience Workshop. Bloom: 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 2023.
[71] Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? In ICLR, 2024.
[72] Yibo Yang, Hong Wang. Haobo Yuan, and Zhouchen Lin. Towards theoretically inspired neural initialization optimization. In NewrIPS, 2022.
[73] Yibo Yang, Xiaojie Li, Zhongzhu Zhou, Shuaiwen Leon Song. Jianlong Wu, Liqiang Nie, and Bernard Ghanem. Corda: Context-oriented decomposition adaptation of large language models for task-aware parameter-efficien fine-tuning. In NeurlPS, 2024.
[74] Yibo Yang, Sihao Liu, Chuan Rao, Bang An, Tiancheng Shen, Philip HS Torr, Ming- Yang, and Bernard Ghanem. Dynamic crrtextd decomposition for task-aware low-rank adaptation with less forgetting and faster convergence. arXiv preprint arXiv:2506.13188 2025.
14

---

[75] Yingxiang Yang, Zhi Zhang, Timothy Hospedales, and Tao Xiang. Free lunch for few-shot learning: Distribution calibration. In ICLR, 2021.
[76] Biao Zhang and Rico Sennrich. Root mean square layer normalization. NeurIPS, 2019.
[77] Bo-Wen Zhang, Liangdong Wang, Jijie Li, Shuhao Gu, Xinya Wu, Zhengduo Zhang, Boyan Gao, Yulong Ao, and Guang Liu. Aquila2 technical report. arXiv preprint 10014:2408.07410, 2024.
[78] Kaidi Zhang, Ruijia Zhang, Behnam Neyshabur, and Olivier Bousquet. Improving generalization by controlling label-noise information neural network weights. In ICLR, 2022.
[79] Ruiqi Zhang, Jingfeng Wu, and Peter Bartlett. In-context learning of linear transformer block: benefits of the mlp component and one-step gd initialization. In NeurlPS, 2024.
[80] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. NeurIPS, 2015.
[81] Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In ICML, 2021.
[82] Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha C. Dvorekk, Sekhar Tatikonda, James S. Duncan, and Ting Liu. Surrogate gap minimization improves sharpness-awe aware training. In ICLR, 2022.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claimed contribution in the abstract and introduction sections is described in the method part, and with empirical justification in the experiment section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have a limited section in Appendix G
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

We propose two analysis results for our model with two theorems in the method section, the corresponding proofs are given in the Appendix along with the assumptions required.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
We provide the source of the dataset and the hyperparameter pool used for model tuning. Besides, the implementation details are carefully described in the algorithm form in the main paper. All the base models are publicly available with clear source. We will release our code for all the implementations the paper.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
17

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
All the datasets we used are public, and we give the detailed resource.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. ec/
Sublic/guidess htaiai//guides/oeeeubbmmmss
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips ccpps.cccppp pttppppllllllppieeee https///oooooooassssooooooleee.e
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We follow the standard benchmark for the experiments, and the tuning details are given in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We follow the standard training and evaluation protocol with multiple trials to generate error bars and justify that the improvements introduced by our model are significant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We give the computational resource for the submission in terms of the GPU index in the paper.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification:
The author information is not released in the paper.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We have Broader impacts section in Appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
we cite all the datasets and models used in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper doess not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy pettps://neurr CcCC/oofeeeeecccs///2////// for what should or should be described.
22

---

A Few-shot performance
We report the entire few-shot performance of all the models, Llama2-7B, Carmma3--B-Insttute Llama3-8B, and GPT2-XL, in Table 5 to comprehensively evaluate the effectiveness of OFA.
Table 5: Comparison between OFA and other baseline algorithms I.Lama-7BB, ILLam3-8B: Instruct, LLama3-8B, and GPT2-XI.. Mean accuracy and standard deviation across five random seeds reported. AGnews and DBPedia are not evaluated for GPT22XI. due to its limitation of context window size. Best results are highlighted in bold.
SST.5
AGNews
Subj
88teSp18
Mothod
LJama2-7B
Zero-shoe
83.00
27.00
50,00
70.20
51,40
$4.20
72.00
41,80
73.60
Few-shot (ICL)
84.4921.4
11.78=a as
77.32g441
85 25.08
32..66ga..o
70 70.2445.00
06.642..44 is
75.48g1.0
03:244...00
Soft-prompt
56.2426m0
94.24-229
55.20g4414
78.00 78.00=7:00
57. 27.400
59. 5620.6
74. 74.40.2..s
35.08852.2
54.322+.100
Label-anchor
83.322t.m
77.18
83 0..72.1.04
53 2..00....0
641 64
81 81.40,aa7
59. 59.12.:1..0
81.10;3.0
81.444475
25.90-00m
65.68
07
58.56=4.91
07 88.88..
48..2..
44.64=3.03
82.02
IA3
93.28 12.9
16.08 1:
11.10
87.04_1.9
L1.92
72. 12.01
91.08,
64.32
88, 1..ss
12CL
87.68
30. 30.12.22.
78.56g5.0
65.48_2114
23.84.a.54
69. 45
00. 00.16.g1
60.72g1.57
87.68
OFA (Ours)
95 05.84.0.1
36
85.92
89 00
04
97. 07.7220.02
94.36..111
Llamz3 Instrrrt
Zero-shot
93,00
35.80
71.00
80:40
50.80
67.80
67.40
53.60
86,40
Few-shot (ICL)
96.48.n.4s
16.7 16.72=2
79.92 as
69.61-20.s
52 72-2.2
97.00.:0.2
65.28=4.29
12
Soft-prompt
81.68 a7.6
75.68 15.17
81.96..a.e
73.28
62. 02.72.5..4
80.88
55.32
75.76
Label-ancher
93.36-2.m
40.54:0044
78.28
64 64.64::141
54.16
69. 89.48:2543
87 07.4842204
58.36
$8.20=3.m
Tesk-vector
56.42_1155
19.63
89
20.08
67. 10.12.0.3
79 11
11.96_1.0
86.52
94.3220a2
40.24-2205
87.60 a
38 10
62.04g7244
77. 47
20.56/21 88
68.04
91.70
I2CL
00.84.00
48 96
70.6026 se
88 96.
81. 31.488...0
66 (6-88899.)
91 20
64.32g2
88. 88.88..00
OFA (Ours)
07.08
58 32.2.74
89.06.1.co
91 84
92.64.3.43
29.47.0.47
97. 07.92.0..0
79.2 79.24
94.56.0aa:
Llama3-88
Zero-shot
56.00
33,20
66.40
85,80
50.60
$0.80
55.80
40.60
53.80
Few-shot (ICL)
95.32
11.36
14.18
87.20)
63.8 63.84
70.00.5.42
85 13.5
52.30
91.88,
Soft. prompt
50.4421255
28.44_6
10.52_11049
65.68
69.12g9.66
63. 20
85.36.2a.s
54.20
60.28
Label-anchor
84.14.000
35.44_o.s
77.6819.1
16.20.1141
64. c.
68.09 11
74.241am
50.7223.61
84.28 178...00
Tisk-vector
94 94.28
37 37.20-22s
75.802..00
85 55.00=a04
68. 18.440ac.o
55 45.802844
73 73.28.1.29
$4.0420.0
75. 75.2884.00
IA3
92.72=1.5
80 24.0
44
60.244e.11
62 64asas
83 20.205
64 64.36223.0
89 291.522.3.44
12CL
87.30
39.3 39.32.4.0
77.72s6.00
65.20022.1
20.03
58. 68.0829.79
86, 86.4422241
02. 02.644...0
80.84
OFA (Ours)
96. jo.3
54 966
87.52
90 36
91. 44
86. 15.7
97. Lo.45
7H. 8fi
94.04-_
Method
GPT2-XL
Zero-shoe
24.76
80:44
35.40
64 88
70.84
37 88
71.36
Few-shot (BCI))
73.65
35. 35.95_9
60 60.64.5.00
60 60.82
$1.8862922
39. 14.02.5..
75. 75.79
Sofl-prompt
61.04a3.44
23.966-2.0
40.60)-10.5
44.44.12
63 1227000
37. 37.001g3.68
Label-anchor
63.40=882
22. 22.30=3.. an
66.30 16.36-11029
35.56=4.30
51.88244.55
20..68221 no
20..20
Tisk-voctor
81.0824.77
28 a7
41. 41.40,3.30
71.80
62 02.48822.53
S7. (0.60=2..
78. 2..000
IA3
86. 86.64=280
20
70.96ea.51
71.52es.4
70.54s843
02.24e.5.0
88.24e1.0
12CL
30.16.3..
25.01_2a00
$1.4815.6
65.96_444
68. 14.722.77
47.92_.144
63.200.a.m
OFA (Ours)
88. 88.68.2266
42 48
70.60
86. 26.11
71 71.444a.56
65.30g4.18
84.80
B Proof of Theorem 3.1
Theorem 3.1. Let f RRd R be a twice continuously differentiable funcciooo with locally Lipschitz
gradients. Suppose the update rule is given by:
= Z
where each P. E Rd x Rd is learnable preconditioning matrix. Define the step-ratio objective in Eq. 3 Under the assumption that f admits local second-order Taylor expansion approximation at each step, then minimizing J (PP) encourages the learned preconditioners Pt to induce local operators T H, with H; ((z) with smaller spectral radius.
1z++1 z S p. 1l < Pt-1-
Thus, it induces faster local contraction and improved convergence.
Proof. By Taylor's theorem, for smooth function f, near point Tt, have;
f(z) f(x)) + vf(x:) - xx) + 2 (z 21) H2(x I²).
23

---

Given the preconditioned gradient descent:
It+1 34 = -
with the quadratic approximation, approximate the gradient:
Vf(z+) = H1( - x x)),
then
T++1 1t - - xx)),
and
T++1 -' (
Then the step-ratio objective becomes:
J(0) = T-1
2t T141
nPPHe( x²)"
[t Iz-1l
1 Tz-1
and operator 7 P. H2 governs convergence. Assuming:
P: spectral radius(/ -1); < 1.
Then minimizing J(P) ensures Pr decreases over time:
I++1 - = (1 -1- (z) -x)- x²),
which leads to
| - 2 = - xx'1l s
pillzee 1 x²11 P: < PC-1-
C Proof of Theorem
Theorem 3.2. Let Zr be the final parameters after T steps of optimization, with preconditioning up- date rules in Eq. and denoting V Ltrain (Zi) as the Hessian at step t with I'. Lirein (Ze)ll mea- suring the curvature after preconditioning at that step. Assume the loss is smooth, | L(Z)) S and the gradient bounded, [Z(Z)) S G, the generalization gap satisfies:
BiLrest (TT) Leroin (Z+) S 50
(1
T
P.V Lirain (z)]
Proof. The proof follows from stability-basee based generalization bounds and Taylor expansion. Let A, 0,+1 0r -nP. nPP.V. (O,)
By second-order Taylor approximation, for perturbation E:
x V(Z)) _ 2
2 L(Z).
Consider the increase in loss under Gaussian perturbation e ~~ ~O, E), used PAC-Bayes analysis. The expected curvature-basee increase is:
- ~ TT<EV><(Zr))
Since E shaped by optimization history through (A.) encountered is influenced by the preconditioned curvature norm:
Then, the effective curvature
|| 2 Emain (0,),1) = CV..) P.V Liiin (L. < P (0,
24

---

Summing over t to T, we obtain:
T Σ (mnn S G22
Σ
IP.. Irain (
Applying a Rademacher complexity or PAC-Bayes-basee argument 19] this leads to:
[[[ (Zr) Curaia (Zr-) 0
P.V L (2))
train
m
D Prompting Templates
Table 6: Illustration of prompting templates and label spaces in our setting. The input prompt template decomposed into multiple [Sentence) and [Label] pairs, which are placeholders for the input sentence and its corresponding label. The template containing a single example for each dataset is generated for the illustration, while in the multiple demonstration example setting, the sentence-label pairs are stacked and separated by newline character: "In".
Dataset
Template
Label Space
SST-2
Review: (Sentence)
negative positive
Sentiment: (Label}
SST-5
Sentence: (Sentence)
terrible / negative / neutral positive / great
Sentiment: (Label)
MR
Review: (Sentence)
negative / positive
Sentiment: [Label}
Subj
Sentence: (Sentence)
Mbjective/suhjecctee
Label: (Label)
DBPedia
Input: (Sentence)
company school / artist athlete politics/ transportation / building nature / village animal plant album film book
Label: (Label)
AGNews
News: (Sentence)
World / Sports Business Technology
Type: [Label]
TREC
Question: (Sentence)
Abhreviation Entity / Person Location / Number
Answer Type: (Label)
HateSpeech18
Text: [Sentence]
neutral hate
Label: (Label)
EmoC
Dialogue: (Sentence)
others happy sad angry
Emotion: (Label)
Extra Details We follow the dataset preprocessing protocol from 339 for our experiments setting. Regarding HateSpeech 18,, only the first two categories- (0: neutral) and [1: hate) are used, since the very few number of samples in the other two may impede a comprehensive evaluation of the model in the test stage.
E LoRA experiment settings
describe the details the LoRA implementation in our experiments. For fair comparison, the LoRA model trained for each individual dataset is tuned by grid search according to the hy- perparameter pool, including LoRA alpha, LoRA dropout, optimizer, and learning rate in Table
25

---

Table 7: Hyperparameter Pool for the LoRA model tuning.
Hyperparameter
Values
LoRA alpha
8, 16, 32, 64
LoRA A dropout
0.0. 0.05, 0.1 AdamW
Optimizer
Learning rate
0.001, 0.0001, 0.00001
F Hyperparameter Pool
We conduct the grid search for fair comparison over all the models, including all the baseline models and ours. The hyperparameter pool for the model tuning is give in Table g
Table 8: Hyperparameter Pool for the LoRA model tuning.
Hyperparameter
Values
0.1, 0.001, 0.0001, 0.00001, 0.000001 0.1, 0.001, 0.0001, 0.00001, 0.000001 AdamW 0.001, 0,0001, 0.00001. 0.001. 0.0001, 0,00001, 0,000001
X2
Optimizer
Learning rate
Weight decay
Training epoch
20, 50, 60, 80, 100
G Limitation and Future Work
In this work, we address the problem of few-shot adaptation within the LLM framework by en- hancing both the internal optimization efficiency and the generalization capability of pretrained models. Specifically, we introduce two distinct objective terms, each targeting one of these properties. While this design improves performance, it also increases the burden of hyperparameter tuning and computational overhead. We leave the unification of these objectives into a single term, enabling joint optimization of both properties, community by developing a rigorous theoretical foundation for this adaptation problem and further improving our method based these insights.
future work. Moving forward, aim to contribute the
H Broader impacts
This paper aims to contribute to the advancement of Machine Learning, especially to the few-shot adaptation of LLMs. While our work may have various societal implications, none require specific emphasis in this context.
26


Reviews:
[{"review_id": "V9zwzSfV6c", "reviewer": "Reviewer_6QXn", "summary": "The paper proposes Optimization-Inspired Few-shot Adaptation (OFA), a lightweight way to dapat a language model with only a small portion of labeled examples. The paper views the transformer block as one step of gradient descent and uses the gradient descent to adapt the model. OFA uniquely adapts an LLM by learning the LayerNorm scales as diagonal pre-conditioners that rescale each layer’s gradient, which is different from the previous optimization-inspired method. Only the scale parameters are fine-tuned. After training, the adapted model runs with zero extra tokens. The experimental results showed that OFA outperforms ICL-based and PEFT-based methods while requiring only fewer parameters and matching zero-shot latency.", "questions": "1. Theorem 3.1 requires $\\rho_t < 1$ for local contraction. However the paper does not provide and proof or intuition for why we can assume this. Can you provide a theoretical argument or a lemma with conditions on $\\eta$ and $P_t$ that help to guarantee this? Or is this possible to supply some empirical evidence on demonstrating $\\rho_t < 1$ in most of the cases.\n\n2. Please give a precise definition of $Z^*$ and state where the symbol first appears.\n\n3. Did the paper look at tasks other than single-sentence classification? Does the method work for other tasks, such as generation, or name entity recognition? \n\n4. When do the authors plan to release the implementation and pre trained checkpoints for reproducibility?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "For quality, the paper conducted thorough experiments comparing OFA with ICL-based and PEFT-based methods and other optimization-inspired methods. The tasks however are limited to single sentence classification, no generation, or name entity recognition. For clarity, the paper explains clearly how OFA is compared to ICL and PEFT, and the algorithm and figures are easy to follow. The work showed that tweaking only LayerNorm scales can beat ICL and PEFT. Offer practical insights for the community. It is also the first to turn the LayerNorm scales into a preconditioner and bridges the optimization-inspired methods with PEFT. Other weaknesses are discussed in the Questions section.", "was_revised": true, "final_justification": "The authors explained that their methods works for the case when $\\rho >1$.  This makes the work more general. It would be helpful to state explicitly that the proposed method applies in both $\\rho < 1$ and $\\rho > 1$ in the paper, while noting that the current proof covers only the $\\rho <1$ case. As the work meets the next acceptance threshold, I have accordingly raised my score.", "rebuttal": "> Question 1: Local contraction with $\\rho_t < 1$. Can the author justify this assumption? \n\nOur work is motivated by an optimization perspective, where the assumption $\\rho < 1$ implies that the optimization process converges [1,2]. In the in-context learning framework introduced in Section 3.1, the forward pass of the LLM can be interpreted as an optimization process, with each attention layer corresponding to one step of gradient descent. Our analysis relies on this assumption, as a rigorous proof that optimizing our loss leads to fast convergence is only feasible under the convergence condition $\\rho < 1$. Specifically, the preconditioning matrix must satisfy $\\rho = || I - \\eta P_t H || < 1$.\n\nConversely, if this assumption does not hold, the internal optimization implemented by the forward LLM pass cannot be strictly interpreted as an optimization process. Nevertheless, our algorithm is designed to find the optimal task-specific $\\rho$ by searching for the corresponding preconditioning matrices $P_t$​, aiming to accelerate convergence regardless of whether $\\rho$ is initially greater or less than 1. In other words, $\\rho < 1$ is not a strict constraint for our algorithm. Instead, our method applies to both $\\rho < 1$ and $\\rho > 1$ scenarios, with the goal of optimizing $\\rho$ rather than requiring it.\n\nIn Figure 3 of our submission, one can observe that our algorithm produces the lowest $\\rho$ among all the baseline models, which implies the fastest convergence speed. \n\n>  Question 2: Definition of $Z^{*}$\n\nWe clarify that $Z^{\\*}$ denotes the optimum of the objective function, defined as $Z^{\\*}= \\arg\\min_{Z} \\mathcal{L}$, where $\\mathcal{L}$ represents the objective of the internal optimization performed during the forward pass of the LLM under the in-context learning setting. This optimization occurs between lines 131 and 132 of our submission. $Z^{\\*}$ is first introduced between lines 156 and 157. We sincerely thank the reviewer for the helpful feedback and will ensure this clarification is explicitly highlighted in the revised version.\n\n\n> Question 3: Can the proposed algorithm work on the few-shot generalization task and name entity recognition? \n\nWe sincerely thank the reviewer for the insightful suggestion to evaluate our algorithm on few-shot generation and named entity recognition (NER) tasks. Our work follows in the mainstream in-context learning setting [3, 4, 5, 6, 7], where sentence classification serves as the main benchmark for evaluation. Furthermore, our current experiments already encompass a diverse set of models, such as Llama2-7B, Llama3-8B, Llama3-8B-Instruct, and GPT2-XL, and datasets, SST-2, SST-5, TREC, AGNews, Subj, HateSp18, DBPedia, EmoC, and MR. While it is theoretically possible to extend our method to few-shot generation and NER, doing so would entail significant methodological modifications and implementation effort for both our approach and the corresponding baselines, which would be difficult to complete in the rebuttal phase given the the limited time available. We agree that these directions are valuable extension work, and we intend to explore them more thoroughly in future work.\n\n> Question 4: Releasing code and model\n\nCertainly, we will release our code and trained models once our paper is accepted, to support implementation transparency and ensure reproducibility.\n\n[1] Nesterov, Yurii. Lectures on convex optimization. Vol. 137. Berlin: Springer International Publishing, 2018.\n\n[2] Boyd, Stephen P., and Lieven Vandenberghe. Convex optimization. Cambridge University Press, 2004.\n\n[3] Huang, Yu, Yuan Cheng, and Yingbin Liang. \"In-context Convergence of Transformers. In ICML. 2024.\n\n[4] Bai, Yu, et al. \"Transformers as statisticians: Provable in-context learning with in-context algorithm selection.\" In NeurIPS. 2023.\n\n[5] Ahn, Kwangjun, et al. \"Transformers learn to implement preconditioned gradient descent for in-context learning.\" In NeurIPS. 2023.\n\n[6] Akyürek, Ekin, et al. \"​​ What learning algorithm is in-context learning? Investigations with linear models.\" In ICLR. 2023.\n\n[7] Li, Zhuowei, et al. \"Implicit In-context Learning.\" In ICLR. 2025"}, {"review_id": "Rv5PIUlpBH", "reviewer": "Reviewer_HYpd", "summary": "This paper proposes Optimization-Inspired Few-Shot Adaptation (OFA), an approach for adapting LLMs to new tasks using only a few demonstration examples, addressing key limitations in existing few-shot adaptation methods. The paper builds upon previous finding which view the forward pass of LLMs as gradient descent optimization. The key insight is that LayerNorm parameters can be finetuned per task to learn preconditioning matrices for each gradient descent step, which happens during the forward pass. The method introduces two objective terms: one for optimization efficiency (step-ratio minimization) and another for generalization (sharpness minimization via Hutchinson approximation). The authors evaluate OFA across multiple classification tasks and LLM architectures, demonstrating consistent improvements over strong baselines.", "questions": "- Could the authors comment on the sensitivity of the OFA hyperparameters?\n- What is the overall cost of finetuning an LLM with OFA? How does it compare to the baselines?\n- In the transition from equation L144 to the preconditioning matrix formulation (equation L147), a bias term $\\frac{\\mu_t}{\\sigma_t}$ appears to be dropped. Could the authors elaborate on this step?", "limitations": "Yes, but only in Appendix.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 2, "originality": 4, "strengths_and_weaknesses": "### Strengths\n\nOverall, the contribution of this work is solid.\n\n(Novelty): To the best of my knowledge, this line of research is novel. The authors rely on optimization and in-context learning theory to guide their design choices. Importantly, the paper also includes compelling empirical evidence that is consistent with what theory predicts.\n\n(Strong empirical performance): OFA demonstrates clear improvements over baselines across multiple classification datasets and LLM architectures.\n\n### Weaknesses\n\nThe practical usability of the approach is still unclear to me since the paper lacks some details about the experimental validation.\n\n(Lack of hyperparameter ablation): The method introduces two additional hyperparameters (λ₁, λ₂) requiring grid search across a substantial range (6 different values each), potentially making the method less practical and adding tuning burden. An ablation study over the two hyperparameters seems crucial to me. It would be informative to clarify whether the values of these hyperparameters are specific for each experiment or shared across all.\n\n(Lack of details about the training setup / overhead): Another missing detail is the runtime during training/finetuning. It would be interesting to get insights on the difficulty of the optimization when considering the two loss terms.", "was_revised": true, "final_justification": "My recommendation for acceptance is based on the novelty and the strong empirical performance of the work. The issues I raised regarding runtimes and hyperparameter sensitivity were addressed during the rebuttal.", "rebuttal": "> Question 1: hyperparameter sensitivity of our algorithm, OFA.\n\nCertainly. We provide details of the grid search over the hyperparameters $\\lambda_1$​ and $\\lambda_2$ on the Emoc and Hatesp18 datasets. The results demonstrate that our proposed loss function is robust to variations in these hyperparameters within the moderate range of $[10^{-5}, 10^{-3}]$, with only limited performance fluctuations. Accordingly, we set $\\lambda_1$ and $\\lambda_2$​ within this range for all datasets in our experiments. Therefore, this confirms that the range $[10^{-5}, 10^{-3}]$ is effective and generalizes well across all experimental settings.\n\n### EmoC \n\n| $\\lambda_1$ \\  $\\lambda_2$ | 0.00001 | 0.0001 | 0.001 | 0.01  | 0.1   |\n|----------------------------|---------|--------|--------|--------|--------|\n| 0.00001                   | 77.15   | 77.58  | 76.19  | 77.47  | 77.30  |\n| 0.0001                    | 78.05   | 78.49  | 78.71  | 77.52  | 74.11  |\n| 0.001                     | 78.09   | **79.24** | 78.31  | 76.27  | 76.09  |\n| 0.01                      | 77.66   | 78.89  | 79.38  | 78.29  | 75.68  |\n| 0.1                       | 75.77   | 74.31  | 74.91  | 74.88  | 74.22  |\n\n### HateSp18 \n\n| $\\lambda_1$ \\ $\\lambda_2$ | 0.00001 | 0.0001 | 0.001 | 0.01  | 0.1   |\n|---------------------------|---------|--------|--------|--------|--------|\n| 0.00001                   | 87.58   | 88.06  | 88.97  | 87.43  | 85.45  |\n| 0.0001                    | 88.25   | **89.47** | 88.95  | 88.11  | 85.09  |\n| 0.001                     | 87.89   | 88.71  | 89.30  | 87.88  | 85.10  |\n| 0.01                      | 87.53   | 88.17  | 88.61  | 87.01  | 85.49  |\n| 0.1                       | 87.15   | 86.75  | 87.88  | 86.03  | 86.20  |\n\n\n> Question 2:  Training cost of our algorithm? \n\nWe compare the training costs of various algorithms on DBPedia, the dataset with the largest number of categories in our experimental suite, using LLaMA3-8B-Instruct.\nCompared to LoRA-based models, our method achieves faster training times and better performance, while maintaining the same inference cost. When compared with embedding-based approaches such as Label-Anchor, Task-Vector, and I2CL, our method incurs slightly higher training time but achieves significantly better performance during inference, along with reduced inference time. Finally, in comparison to IA3 and Soft-Prompt, our method demonstrates substantial improvements in both inference efficiency and accuracy.\n\n| Method           | Soft-prompt | Label-anchor | Task-vector | IA3     | I2CL    | LoRA (Rank16) | LoRA (Rank1) | Ours    |\n|------------------|-------------|---------------|--------------|---------|---------|---------------|----------------|---------|\n| Training cost (s) | 196.08      | 50.41         | 213.25       | 324.12  | 312.50  | 483.52       |  449.23       | 367.05  |\n| Inference cost (s) | 53.64       | 52.41         | 56.78        | 52.40   | 52.59   | 51.38         | 51.35          | 51.37   |\n|Performance (Accuracy %)| 82.88 | 87.48| 79.52 | 92.56 | 91.20 | 92.62 | 57.88  | 97.92 | \n\n> Question 3: The omitted bias term.\n\nWe appreciate this issue being identified by the reviewer. In the mainstream open-source LLMs, such as Llama, BLOOM, Falcon, Qwen, and MPT, a special type of LayerNorm, RMSNorm, is usually used, with the definition as: $\\text{RMSNorm}(x) =  \\gamma \\cdot \\frac{x}{\\sqrt{\\frac{1}{d} \\sum_{i=1}^d x_i^2 + \\epsilon}} $ where $\\gamma$ is learnable parameter. In this way, the computation of RMSNorm can be interpreted by setting the mean  $\\mu$ in LayerNorm to 0. As a result, the bias term is removed by multiplying by 0. In the normal LayerNorm case, such as GPT2-XL, the bias term can be bounded by a constant; then omitting this will simplify our analysis without affecting the conclusion that optimizing our proposed loss leads to faster convergence. To further test whether our algorithm is robust to the normal layer norm, the GPT2-XL (with LayerNorm) trained with our algorithm on the few-shot learning tasks also demonstrates superior performance over all the baseline algorithms in Table 4 of our submission. Therefore, we omit this bias to highlight the main focus, $P_t \\nabla\\mathcal{L}(Z_t)$, of our work. We will explain these reasons clearly in our revised version. \n\n### Few-shot performance of GPT2-XL. We quote the experiment results below from Table 4 in the appendix of our original submission. \n\n| Dataset / Method   | SST-2            | SST-5            | TREC             | Subj             | HateSp18         | EmoC             | MR               |\n|--------------------|------------------|------------------|------------------|------------------|------------------|------------------|------------------|\n| Zero-shot          | 74.76            | 30.44            | 35.40            | 64.88            | 70.84            | 37.88            | 71.36            |\n| Few-shot (ICL)     | 73.65 $\\pm$ 8.89 | 35.95 $\\pm$ 2.39 | 60.64 $\\pm$ 5.00 | 63.82 $\\pm$ 10.55| 51.86 $\\pm$ 3.22 | 38.62 $\\pm$ 6.87 | 75.79 $\\pm$ 9.25 |\n| Soft-prompt        | 61.04 $\\pm$ 3.45 | 23.96 $\\pm$ 2.09 | 40.60 $\\pm$ 10.15| 55.44 $\\pm$ 4.12 | 63.92 $\\pm$ 7.06 | 36.68 $\\pm$ 2.70 | 57.60 $\\pm$ 3.53 |\n| Label-anchor       | 63.40 $\\pm$ 8.82 | 22.36 $\\pm$ 3.37 | 66.36 $\\pm$ 10.69| 55.56 $\\pm$ 4.26 | 54.88 $\\pm$ 4.53 | 36.68 $\\pm$ 2.70 | 60.20 $\\pm$ 3.32 |\n| Task-vector        | 81.08 $\\pm$ 4.87 | 28.52 $\\pm$ 1.37 | 41.40 $\\pm$ 5.35 | 71.81 $\\pm$ 1.86 | 62.48 $\\pm$ 2.83 | 37.60 $\\pm$ 2.48 | 78.40 $\\pm$ 2.26 |\n| IA3                | 86.64 $\\pm$ 2.89 | 40.52 $\\pm$ 2.25 | 70.96 $\\pm$ 8.61 | 71.52 $\\pm$ 8.46 | 70.84 $\\pm$ 3.63 | 62.24 $\\pm$ 3.50 | 83.24 $\\pm$ 1.09 |\n| I2CL               | 80.16 $\\pm$ 3.98 | 35.04 $\\pm$ 2.60 | 51.48 $\\pm$ 5.26 | 65.96 $\\pm$ 4.83 | 68.32 $\\pm$ 4.76 | 47.92 $\\pm$ 1.84 | 83.20 $\\pm$ 3.29 |\n| **Ours**           | **88.68 $\\pm$ 2.66** | **42.48 $\\pm$ 2.51** | **70.60 $\\pm$ 6.44** | **86.11 $\\pm$ 4.29** | **71.44 $\\pm$ 8.65** | **65.30 $\\pm$ 4.18** | **84.80 $\\pm$ 6.21** |"}, {"review_id": "Q8NnDBp3Z7", "reviewer": "Reviewer_CPuh", "summary": "The paper proposes a new few-shot adaptation algorithm that adjusts the parameters of layer-norm on few-shot data. This approach is motivated by the gradient descent properties of transformer layers to encourage smooth gradient behaviors across layers (step ratio regularization). The proposed approach also aims to learn flat minima using the Hutchinson approximation method. Empirical results show the proposed approach outperforms other baselines: ICL, soft prompt, label anchor, task vector, IA3, I2CL, and LoRA. In particular, it outperforms LoRA by a large margin while using less parameters. Ablation study through probe classifier analysis shows the importance of step ratio and sharpness regularization.", "questions": "1. Are the proposed terms applicable to any SFT or PEFT training algorithm?\n\n2. How many exemplars are used in the test sets? How the proposed approach would work if we increase the number of exemplars, and at what point it would become limited by the representation of the layer-norm parameters.\n\n3. How to set the hyperparameters of different loss terms and is the model sensitive to these parameters?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength\n\n1. The paper is well-motivated. The proposed regularization terms on step ratio and loss sharpness are interesting and technically sound.\n2. The proposed approach demonstrates strong empirical results, outperforming LoRA by a large margin while using less parameters.\n3. The paper carries out comprehensive ablation studies and visualization to demonstrate the value of step ratio and sharpness regularization.\n\nWeakness\n\n1. Loose connection between regularization and parameter efficient fine-tuning. The proposed approach is comprised of cross entropy loss on the layer norm parameters, and two regularization terms. It seems these regularization terms could be applied just as easily to other PEFT (or any training) methods, such as soft prompt or LoRA.\n2. Scaling with exemplar size. It is unclear how the proposed approach would work if we increase the number of exemplars, and at what point it would become limited by the representation of the layer-norm parameters.", "was_revised": true, "final_justification": "My final rating is borderline accept. The paper presents clear motivation and shows strong empirical results. The authors' response addressed my concerns regarding the distinctions between soft prompting and LoRA. The additional experiments exploring sample size variations provide valuable insights for practitioners weighing different approaches based on the level of available training data.", "rebuttal": "> Weakness 1: Loose connection between the introduced loss terms and the efficient fine-tuning. Application of our loss to Soft Prompt and LoRA.\n\nAdmittedly, soft prompt, LoRA, and our proposed regularization terms are performed in different aspects. However, LoRA and soft prompts have drawbacks in few-shot adaptations. LoRA adapter introduces a substantial number of additional parameters, making it less suitable for few-shot learning scenarios, a limitation that our loss function does not inherently resolve. The in-context learning based algorithms usually have limited flexibility, such as suitable parameters, for adaptation to new tasks. \n\nIn our work, we primarily focus on the few-shot learning setting, where the training data is too limited to support effective supervised fine-tuning (SFT) or parameter-efficient fine-tuning (PEFT) methods such as LoRA. Our experimental results show that LoRA often fails to outperform even the zero-shot and few-shot baselines. This is because the LoRA adapter introduces a substantial number of additional parameters, making it less suitable for few-shot learning scenarios, a limitation that our loss function does not inherently resolve. Our experimental results support this assessment: although LoRA models trained with our loss consistently outperform those trained with the LoRA only models, they still fall short of outperforming ours. This pattern is observed consistently across both the LLaMA-2 and LLaMA-3 model families.\n\n### LoRA comparsion on Llama2-7B. We quote the results from Table 2 in our original submission. \n\n| Method              | SST-2             | SST-5             | TREC              | AGNews            | Subj              | HateSpeech18      | DBPedia           | EmoC              | MR                |\n|---------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|\n| Rank 1              | $89.64_{\\pm 3.23}$ | $24.84_{\\pm 9.92}$ | $22.88_{\\pm 9.09}$ | $50.56_{\\pm 31.36}$ | $72.08_{\\pm 8.74}$ | $69.14_{\\pm 9.76}$ | $59.44_{\\pm 42.01}$ | $25.64_{\\pm 2.95}$ | $64.84_{\\pm 13.86}$ |\n| Rank1 (our loss)    | $88.48_{\\pm 3.34}$ | $20.80_{\\pm 0.89}$ | $24.32_{\\pm 6.02}$ | $62.04_{\\pm 29.5}$  | $72.84_{\\pm 11.33}$ | $69.38_{\\pm 10.77}$ | $74.6_{\\pm 33.23}$  | $33.24_{\\pm 18.28}$ | $64.88_{\\pm 18.48}$ |\n| Ours                | $\\mathbf{95.84}_{\\pm0.41}$ | $\\mathbf{50.36}_{\\pm3.28}$ | $\\mathbf{85.92}_{\\pm1.90}$ | $\\mathbf{89.00}_{\\pm1.26}$ | $\\mathbf{88.40}_{\\pm4.76}$ | $\\mathbf{83.04}_{\\pm3.72}$ | $\\mathbf{97.72}_{\\pm0.52}$ | $\\mathbf{76.60}_{\\pm2.39}$ | $\\mathbf{94.36}_{\\pm1.13}$ |\n\n\n### LoRA comparison on Llama3-8B-Instruct. We quote the results from Table 2 in our original submission. \n\n| Method              | SST-2             | SST-5             | TREC              | AGNews            | Subj              | HateSpeech18      | DBPedia           | EmoC              | MR                |\n|---------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|-------------------|\n| Rank 1              | $87.08_{\\pm 4.81}$ | $19.52_{\\pm 0.45}$ | $25.88_{\\pm 9.17}$ | $39.72_{\\pm 25.1}$ | $78.92_{\\pm 8.43}$ | $63.60_{\\pm 9.07}$ | $57.88_{\\pm 39.37}$ | $24.68_{\\pm 1.17}$ | $87.20_{\\pm 3.38}$ |\n| Rank1 (our loss)    | $87.40_{\\pm 8.05}$ | $20.32_{\\pm 1.72}$ | $27.52_{\\pm 6.09}$ | $39.12_{\\pm 24.40}$ | $80.96_{\\pm 5.55}$ | $68.92_{\\pm 8.33}$ | $73.52_{\\pm 32.69}$ | $27.68_{\\pm 3.65}$ | $87.84_{\\pm 13.59}$ |\n| Ours                | $\\mathbf{97.08}_{\\pm0.27}$ | $\\mathbf{58.32}_{\\pm2.74}$ | $\\mathbf{89.06}_{\\pm1.49}$ | $\\mathbf{91.84}_{\\pm0.61}$ | $\\mathbf{92.64}_{\\pm3.43}$ | $\\mathbf{89.47}_{\\pm0.47}$ | $\\mathbf{97.92}_{\\pm1.06}$ | $\\mathbf{79.24}_{\\pm4.87}$ | $\\mathbf{94.56}_{\\pm0.51}$ |\n\nFrom an optimization perspective, Soft Prompt modifies the initialization state, whereas our loss function influences the optimization trajectory itself by adapting the preconditioning matrices. To this end, the gradient generated by our loss would not be informative for the target learnable parameters. However, we still apply our loss to the soft prompt. Therefore, the soft prompt problem cannot be solved by our algorithm, and in the empirical results, it harms the performance of our model. \n\n| Method                     | SST‑2  | SST‑5 | TREC  | AGNews | Subj  | HateSp18 | DBPedia | EmoC  | MR    |\n|---------------------------|--------|-------|-------|--------|-------|-----------|---------|-------|-------|\n| Soft prompt with our loss | 88.2   | 43.4  | 66.6  | 89.6   | 69.2  | 54.0      | 89.6    | 44.2  | 85.2  |\n| Soft‑prompt               | 84.68  | 38.40 | 75.68 | 84.96  | 73.28 | 62.72     | 82.88   | 55.32 | 75.76 |\n| **Ours**                  | **97.08** | **58.32** | **89.06** | **91.84** | **92.64** | **89.47** | **97.92** | **79.24** | **94.56** |\n\n> Question 1: Application to SFT and PEFT algorithm. \n\nAs noted in our response to Weakness 1, while the proposed regularization terms are compatible with any SFT or PEFT method, these methods are generally tailored for fine-tuning in data-rich scenarios. In contrast, our approach is specifically designed for few-shot adaptation, where methods like SFT and LoRA are prone to overfitting.\n\n> Weakness 2 and Question 2: The effects of different numbers of exemplars. \n\nWe conduct few-shot learning with 5 samples per category. To explore the effects of different numbers of samples, we conduct experiments on various models, Llama2 and Llama3-Instruct, and datasets, HateSP18 and Subj, with the experiment results below. In general, our method benefits from the increasing number of samples, but the model’s performance saturates after 15 shots. \n\n### Llama2-7B\n\n| HateSP18 | 3-shot | 5-shot | 10-shot | 15-shot | 20-shot |\n|----------|--------|--------|---------|---------|---------|\n| ICL      | 68.81  | 70.24  | 70.64   | 74.23   | 70.72   |\n| Ours     | 77.63  | 83.04  | 83.73   | 86.09   | 84.97   |\n\n\n### Llama3-Instruct\n\n| Subj | 3-shot | 5-shot | 10-shot | 15-shot | 20-shot |\n|-------|--------|--------|---------|---------|---------|\n| ICL   | 63.76  | 57.48  | 54.73   | 61.80   | 58.32   |\n| Ours  | 88.76  | 92.64  | 93.62   | 95.40   | 95.23   |\n\n> Question 3: Hyperparameter sensitivity analysis\n\nWe provide details of the grid search over the hyperparameters $\\lambda_1$​ and $\\lambda_2$ on the Emoc and Hatesp18 datasets. The results demonstrate that our proposed loss function is robust to variations in these hyperparameters within the moderate range of $[10^{-5}, 10^{-3}]$, with only limited performance fluctuations. Accordingly, we set $\\lambda_1$ and $\\lambda_2$​ within this range for all datasets in our experiments. Therefore, this confirms that the range $[10^{-5}, 10^{-3}]$ is effective and generalizes well across all experimental settings.\n\n### EmoC\n| $\\lambda_1$ \\  $\\lambda_2$ | 0.00001 | 0.0001 | 0.001 | 0.01  | 0.1   |\n|----------------------------|---------|--------|--------|--------|--------|\n| 0.00001                   | 77.15   | 77.58  | 76.19  | 77.47  | 77.30  |\n| 0.0001                    | 78.05   | 78.49  | 78.71  | 77.52  | 74.11  |\n| 0.001                     | 78.09   | **79.24** | 78.31  | 76.27  | 76.09  |\n| 0.01                      | 77.66   | 78.89  | 79.38  | 78.29  | 75.68  |\n| 0.1                       | 75.77   | 74.31  | 74.91  | 74.88  | 74.22  |\n\n### HateSp18\n\n| $\\lambda_1$ \\ $\\lambda_2$ | 0.00001 | 0.0001 | 0.001 | 0.01  | 0.1   |\n|---------------------------|---------|--------|--------|--------|--------|\n| 0.00001                   | 87.58   | 88.06  | 88.97  | 87.43  | 85.45  |\n| 0.0001                    | 88.25   | **89.47** | 88.95  | 88.11  | 85.09  |\n| 0.001                     | 87.89   | 88.71  | 89.30  | 87.88  | 85.10  |\n| 0.01                      | 87.53   | 88.17  | 88.61  | 87.01  | 85.49  |\n| 0.1                       | 87.15   | 86.75  | 87.88  | 86.03  | 86.20  |"}, {"review_id": "za1r8g0iOa", "reviewer": "Reviewer_7e2o", "summary": "This paper introduces Optimization-Inspired Few-Shot Adaptation (OFA), an approach for adapting LLMs to new tasks with limited data. The paper reinterpret the forward pass of LLMs as a sequence of preconditioned gradient descent steps and leverage this insight to propose a parameter-efficient adaptation method that modifies only the LayerNorm parameters. OFA combines three objectives: cross-entropy loss, a step-ratio objective that improves optimization efficiency by promoting smoother convergence paths, and a sharpness minimization objective that enhances generalization by steering optimization toward flat minima. The method is theoretically grounded with convergence and generalization bounds, while being extremely parameter-efficient (using only 0.27M parameters). Extensive experiments across multiple datasets and model architectures demonstrate that OFA consistently outperforms existing ICL and PEFT methods, achieving 4-10% improvements over baseline approaches without introducing additional inference overhead.", "questions": "- In your experimental evaluation, you use IA3 as one of the baseline methods for comparison. However, the original IA3 paper (citation [35]) introduces additional training components beyond the parameter-efficient architecture, specifically Unlikelihood Training and Length Normalization, which are designed to address challenges in few-shot learning scenarios. Did you implement these additional loss terms when training your IA3 baseline? If not, do you believe including these components might narrow the performance gap between IA3 and your proposed OFA method?\n- The paper introduces novel loss terms (step-ratio objective and sharpness minimization) that are theoretically motivated to improve optimization efficiency and generalization. Did you consider applying these loss terms to other parameter-efficient fine-tuning methods such as IA3 or LoRA to isolate their effect?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n- The paper provides an interesting theoretical interpretation of transformer layers as implementing gradient descent steps, creating a good foundation for their method.\n- OFA achieves impressive results while updating only LayerNorm parameters (0.27M parameters), making it significantly more efficient than competing approaches.\n- The experiments cover multiple datasets and model architectures (Llama2-7B, Llama3-8B, GPT2-XL), demonstrating consistent performance improvements.\n- The paper provides theoretical analysis connecting their objectives to convergence speed and generalization bounds.\n\n### Weaknesses\n- citation [35] on line 27 doesn't adequately support the claim that PEFT methods suffer from overfitting in few-shot scenarios. This weakens one of their motivating arguments.\n- The Hutchinson approximation for sharpness estimation adds implementation complexity that isn't fully addressed in terms of computational overhead during training.\n- Some of the theoretical results rely on assumptions like locally Lipschitz gradients and quadratic approximations that may not always hold in practice for complex LLMs.", "was_revised": true, "final_justification": "The authors have addressed my concerns regarding computational cost and have provided additional experiments on adding those loss terms during LoRA training, showing that their method consistently outperforms the baseline methods considered. Accordingly, I have decided to increase my score by one point.", "rebuttal": "> Weakness 1: Not sufficient support for the claim that the PEFT models, such as LoRA, suffer from overfitting in few-shot scenarios. \n\nWe appreciate the reviewer for pointing out this issue. Actually, there are more studies that have pointed out the overfitting problem of LoRA in few-shot settings. \n\nSpecifically, [1] notes that “the matrix decomposition approximation might impact the adaptability” of LoRA, as its low-rank constraint can hinder generalization when data is extremely scarce. To address this, [2] proposes HyperLoRA, a meta-adaptation method designed to overcome the challenges of “few-shot adaptation” with static LoRA modules. Similarly, [3] introduces BiLoRA, a bi-level optimization framework motivated by the observation that standard LoRA “often overfit training data, resulting in sub-optimal generalization on test data.” Furthermore, several studies [4–9] conclude that LoRA alone is insufficient for generalization when the adaptation data is scarce. This highlights LoRA’s limited adaptability when learning from very small datasets. We also provide practical evidence in our response to Question 2 to support this argument. \n\n> Weakness 2: Computational cost from our sharpness approximation.\n\nWe acknowledge that estimating layer-wise sharpness incurs additional computational overhead. However, this cost is confined to the training phase and remains acceptable when compared to baseline models. At inference time, our method maintains the same efficiency as standard zero-shot settings, offering significantly lower complexity than many baselines.\nWe compare the training costs of various algorithms on DBPedia, the dataset with the largest number of categories in our experimental suite, using LLaMA3-8B-Instruct.\n\nCompared to LoRA-based models, our method achieves faster training times and better performance, while maintaining the same inference cost. When compared with embedding-based approaches such as Label-Anchor, Task-Vector, and I2CL, our method incurs slightly higher training time but achieves significantly better performance during inference, along with reduced inference time. Finally, in comparison to IA3 and Soft-Prompt, our method demonstrates substantial improvements in both inference efficiency and accuracy.\n\n| Method           | Soft-prompt | Label-anchor | Task-vector | IA3     | I2CL    | LoRA (Rank16) | LoRA (Rank1) | Ours    |\n|------------------|-------------|---------------|--------------|---------|---------|---------------|----------------|---------|\n| Training cost (s) | 196.08      | 50.41         | 213.25       | 324.12  | 312.50  | 483.52       |  449.23       | 367.05  |\n| Inference cost (s) | 53.64       | 52.41         | 56.78        | 52.40   | 52.59   | 51.38         | 51.35          | 51.37   |\n|Performance (Accuracy %)| 82.88 | 87.48| 79.52 | 92.56 | 91.20 | 92.62 | 57.88  | 97.92 | \n\n\n> Weakness 3: The assumption about locally Lipschitz gradients and quadratic approximations for analysis. \n\nOur local Lipschitz gradient assumption is supported by recent theoretical work [10], which shows that the attention mechanism is smooth under standard conditions. In our framework, the gradient is implicitly computed as each attention layer corresponds to a single gradient step in the internal optimization process. \n\nSimilarly, our assumption of second-order smoothness in in-context learning is theoretically supported by [11,12]. In practice, computing gradients by backpropagation with respect to the input of each transformer layer equals the second-order differentiability of the loss with respect to internal representations, thereby justifying our use of a second-order Taylor expansion for approximation. Moreover, since the internal optimization is implicitly performed by the attention layers, the backpropagation used to update LLM parameters effectively captures second-order information. Besides the in-context learning setting, the second-order assumption for conventional neural network settings is also broadly studied and applied [13, 14], which can be seen as third order in our case. \n\nWe will include more discussion about the validity of the assumptions in the revised version.\n\n> Question 1: Do we follow the implementation of IA3?\n\nWe implement IA3 following its default architectural design and training objective, including unlikelihood training and length normalization terms, to ensure a fair comparison. Nevertheless, our model consistently demonstrates a clear performance advantage across all experiments.\n\n> Question 2: Can our loss be applied to other algorithms, such as LoRA?\n\nWe appreciate this question. Regarding the performance of combining LoRA with our proposed loss in the few-shot learning setting, please refer to the tables below as well as Table 2 in our submission, which includes additional results across different LoRA ranks. From these results, we observe that LoRA generally underperforms in the few-shot setting due to limited data availability for effectively fitting the model. Compared to LoRA trained with standard cross-entropy loss, our loss helps mitigate this overfitting to some extent, though it does not fully resolve the problem.\n\n The results from the table below show that the LoRA models suffer in the few-shot settings.\n\n### LoRA comparsion on Llama2-7B\n\n| Method  | SST-2| SST-5| TREC| AGNews  | Subj | HateSpeech18| DBPedia| EmoC| MR |\n|-------|------|-----|-----|-------|-------|-----|-----|----|----|\n| Rank 1  | $89.64_{\\pm 3.23}$ | $24.84_{\\pm 9.92}$ | $22.88_{\\pm 9.09}$ | $50.56_{\\pm 31.36}$ | $72.08_{\\pm 8.74}$ | $69.14_{\\pm 9.76}$ | $59.44_{\\pm 42.01}$ | $25.64_{\\pm 2.95}$ | $64.84_{\\pm 13.86}$ |\n| Rank1 (our loss)    | $88.48_{\\pm 3.34}$ | $20.80_{\\pm 0.89}$ | $24.32_{\\pm 6.02}$ | $62.04_{\\pm 29.5}$  | $72.84_{\\pm 11.33}$ | $69.38_{\\pm 10.77}$ | $74.6_{\\pm 33.23}$  | $33.24_{\\pm 18.28}$ | $64.88_{\\pm 18.48}$ |\n| Ours | $\\mathbf{95.84}_{\\pm0.41}$ | $\\mathbf{50.36}_{\\pm3.28}$ | $\\mathbf{85.92}_{\\pm1.90}$ | $\\mathbf{89.00}_{\\pm1.26}$ | $\\mathbf{88.40}_{\\pm4.76}$ | $\\mathbf{83.04}_{\\pm3.72}$ | $\\mathbf{97.72}_{\\pm0.52}$ | $\\mathbf{76.60}_{\\pm2.39}$ | $\\mathbf{94.36}_{\\pm1.13}$ |\n\n### LoRA comparison on Llama3-8B-Instruct\n\n| Method  | SST-2 | SST-5| TREC | AGNews | Subj | HateSpeech18 | DBPedia | EmoC | MR |\n|-------|------|-----|--------|-----|-----|------|-----|-----|------|\n| Rank 1 | $87.08_{\\pm 4.81}$ | $19.52_{\\pm 0.45}$ | $25.88_{\\pm 9.17}$ | $39.72_{\\pm 25.1}$ | $78.92_{\\pm 8.43}$ | $63.60_{\\pm 9.07}$ | $57.88_{\\pm 39.37}$ | $24.68_{\\pm 1.17}$ | $87.20_{\\pm 3.38}$ |\n| Rank1 (our loss)    | $87.40_{\\pm 8.05}$ | $20.32_{\\pm 1.72}$ | $27.52_{\\pm 6.09}$ | $39.12_{\\pm 24.40}$ | $80.96_{\\pm 5.55}$ | $68.92_{\\pm 8.33}$ | $73.52_{\\pm 32.69}$ | $27.68_{\\pm 3.65}$ | $87.84_{\\pm 13.59}$ |\n| Ours  | $\\mathbf{97.08}_{\\pm0.27}$ | $\\mathbf{58.32}_{\\pm2.74}$ | $\\mathbf{89.06}_{\\pm1.49}$ | $\\mathbf{91.84}_{\\pm0.61}$ | $\\mathbf{92.64}_{\\pm3.43}$ | $\\mathbf{89.47}_{\\pm0.47}$ | $\\mathbf{97.92}_{\\pm1.06}$ | $\\mathbf{79.24}_{\\pm4.87}$ | $\\mathbf{94.56}_{\\pm0.51}$ |\n\nWe hypothesize three main reasons for this behavior: 1) Parameter count: Even in the rank-1 setting, LoRA introduces twice as many parameters as our method, thus requiring more data to generalize effectively. 2) Initialization: Our parameters are pretrained on a large-scale dataset, whereas LoRA parameters are initialized from scratch, making our method more data-efficient for adaptation. 3) Optimization behavior: LoRA significantly alters the gradient flow, the core component of the optimization process, whereas our method preserves the original gradients and instead adapts the preconditioning matrices to each task.\n\n[1] Srinivasan, Krishna Prasad Varadarajan, et al. \"Comparative analysis of different efficient fine tuning methods of large language models (llms) in low-resource setting.\" arXiv preprint arXiv:2405.13181.2024..\n\n[2]Lv, Chuancheng, et al. \"Hyperlora: Efficient cross-task generalization via constrained low-rank adapters generation.\" Findings: EMNLP 2024.\n\n[3] Qiang, Rushi, Ruiyi Zhang, and Pengtao Xie. \"Bilora: A bi-level optimization framework for overfitting-resilient low-rank adaptation of large pre-trained models.\" arXiv preprint arXiv:2403.13037. 2024.\n\n[4]Kong, Xiaoyu, et al. \"Customizing language models with instance-wise lora for sequential recommendation.\" NeurIPS. 2024.\n\n[5] Asadi, Nader, et al. \"Does combining parameter-efficient modules improve few-shot transfer accuracy?.\" arXiv preprint arXiv:2402.15414. 2024.\n\n[6] Baklouti, Ghassen, et al. \"Regularized Low-Rank Adaptation for Few-Shot Organ Segmentation.\" arXiv preprint arXiv:2507.15793. 2025.\n\n[7] Asadi, Nader, et al. \"Combining Pre-trained LoRA Modules Improves Few-shot Adaptation of Foundation Models to New Tasks.\" ICML 2024 Workshop on Foundation Models in the Wild.\n\n[8]Farina, Matteo, et al. \"Rethinking Few-Shot Adaptation of Vision-Language Models in Two Stages.\" CVPR. 2025.\n\n[9] Li, Yang, Shaobo Han, and Shihao Ji. \"Vb-lora: Extreme parameter efficient fine-tuning with vector banks.\" NeurIPS. 2024. \n\n[10] Castin, Valérie, Pierre Ablin, and Gabriel Peyré. \"How Smooth Is Attention?\" ICML. 2024. \n\n[11] Fu, Deqing, et al. \"Transformers learn to achieve second-order convergence rates for in-context linear regression.\" NeurIPS. 2024.\n\n[12] Giannou, Angeliki, et al. \"How Well Can Transformers Emulate In-context Newton's Method?.\" arXiv preprint arXiv:2403.03183. 2024.\n\n[13] Tatzel, Lukas, et al. \"Debiasing Mini-Batch Quadratics for Applications in Deep Learning.\" ICLR. 2025.\n\n[14] Petersen, Felix, et al. \"Newton losses: Using curvature information for learning with differentiable algorithms.\" NeurIPS. 2024.\n\n[15] Shin, Dahun, et al. \"Sassha: Sharpness-aware Adaptive Second-order Optimization with Stable Hessian Approximation.\" ICML. 2025."}]

Author rebuttal:

Meta-review: Based on my reading the reviews, author responses and discussion, I see the following strengths of this manuscript:

- **S1.** Based on an insightful connection between optimization and the transformer forward pass, the submission proposes a novel and unique low-parameter LLM adaptation for few-shot problems, and provides a strong and elaborate empirical evaluation (multiple datasets, architectures).

- **S2.** Along with the parameterization, the submission also proposes novel loss terms such as step-size ratio and sharpness regularization, which significantly improve the performance on few-shot problems. Furthermore, these loss can also be applied to other PEFT and LoRA based adaptation schemes though these variations of the baselines still underperform the proposed scheme.

I also saw various limitations such as (i) unknown sensitivity to new hyperparameters for the new loss terms, (ii) additional computation time due to the layer-wise sharpness estimation, and (iii) the additional smoothness assumptions in the analysis. However, it seems to me that all the limitations are somewhat minor, and adequately addressed in the discussion with additional empirical evaluations (for the hyperparameter dependencies and runtime comparisons).

Based on this, I am recommending an accept for this paper.
Decision: Accept (spotlight)