Optimization Inspired Few-Shot Adaptation for Large Language Models

Weakness 1: Not sufficient support for the claim that the PEFT models, such as LoRA, suffer from overfitting in few-shot scenarios.

We 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.

Specifically, [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.

Weakness 2: Computational cost from our sharpness approximation.

We 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. We 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.

Compared 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.

Weakness 3: The assumption about locally Lipschitz gradients and quadratic approximations for analysis.

Our 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.

Similarly, 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.

We will include more discussion about the validity of the assumptions in the revised version.

Question 1: Do we follow the implementation of IA3?

We 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.

Question 2: Can our loss be applied to other algorithms, such as LoRA?

We 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.

The results from the table below show that the LoRA models suffer in the few-shot settings.

LoRA comparsion on Llama2-7B

LoRA comparison on Llama3-8B-Instruct

We 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.

[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..

[2]Lv, Chuancheng, et al. "Hyperlora: Efficient cross-task generalization via constrained low-rank adapters generation." Findings: EMNLP 2024.

[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.

[4]Kong, Xiaoyu, et al. "Customizing language models with instance-wise lora for sequential recommendation." NeurIPS. 2024.

[5] Asadi, Nader, et al. "Does combining parameter-efficient modules improve few-shot transfer accuracy?." arXiv preprint arXiv:2402.15414. 2024.

[6] Baklouti, Ghassen, et al. "Regularized Low-Rank Adaptation for Few-Shot Organ Segmentation." arXiv preprint arXiv:2507.15793. 2025.

[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.

[8]Farina, Matteo, et al. "Rethinking Few-Shot Adaptation of Vision-Language Models in Two Stages." CVPR. 2025.

[9] Li, Yang, Shaobo Han, and Shihao Ji. "Vb-lora: Extreme parameter efficient fine-tuning with vector banks." NeurIPS. 2024.

[10] Castin, Valérie, Pierre Ablin, and Gabriel Peyré. "How Smooth Is Attention?" ICML. 2024.

[11] Fu, Deqing, et al. "Transformers learn to achieve second-order convergence rates for in-context linear regression." NeurIPS. 2024.

[12] Giannou, Angeliki, et al. "How Well Can Transformers Emulate In-context Newton's Method?." arXiv preprint arXiv:2403.03183. 2024.

[13] Tatzel, Lukas, et al. "Debiasing Mini-Batch Quadratics for Applications in Deep Learning." ICLR. 2025.

[14] Petersen, Felix, et al. "Newton losses: Using curvature information for learning with differentiable algorithms." NeurIPS. 2024.

[15] Shin, Dahun, et al. "Sassha: Sharpness-aware Adaptive Second-order Optimization with Stable Hessian Approximation." ICML. 2025.

Weakness 1: Loose connection between the introduced loss terms and the efficient fine-tuning. Application of our loss to Soft Prompt and LoRA.

Admittedly, 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.

In 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.

LoRA comparsion on Llama2-7B. We quote the results from Table 2 in our original submission.

LoRA comparison on Llama3-8B-Instruct. We quote the results from Table 2 in our original submission.

From 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.

Question 1: Application to SFT and PEFT algorithm.

As 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.

Weakness 2 and Question 2: The effects of different numbers of exemplars.

We 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.

Llama2-7B

Llama3-Instruct

Question 3: Hyperparameter sensitivity analysis

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.

EmoC

HateSp18

Question 1: Local contraction with $\rho_t < 1$. Can the author justify this assumption?

Our 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$.

Conversely, 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.

In 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.

Question 2: Definition of $Z^{*}$

We 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.

Question 3: Can the proposed algorithm work on the few-shot generalization task and name entity recognition?

We 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.

Question 4: Releasing code and model

Certainly, we will release our code and trained models once our paper is accepted, to support implementation transparency and ensure reproducibility.

[1] Nesterov, Yurii. Lectures on convex optimization. Vol. 137. Berlin: Springer International Publishing, 2018.

[2] Boyd, Stephen P., and Lieven Vandenberghe. Convex optimization. Cambridge University Press, 2004.

[3] Huang, Yu, Yuan Cheng, and Yingbin Liang. "In-context Convergence of Transformers. In ICML. 2024.

[4] Bai, Yu, et al. "Transformers as statisticians: Provable in-context learning with in-context algorithm selection." In NeurIPS. 2023.

[5] Ahn, Kwangjun, et al. "Transformers learn to implement preconditioned gradient descent for in-context learning." In NeurIPS. 2023.

[6] Akyürek, Ekin, et al. "​​ What learning algorithm is in-context learning? Investigations with linear models." In ICLR. 2023.

[7] Li, Zhuowei, et al. "Implicit In-context Learning." In ICLR. 2025

Dear Reviewer,

We appreciate the time and effort you have taken to review our submission and provide valuable feedback. We have carefully addressed the comments raised during the first round and submitted our responses accordingly. As the discussion period is ending soon, we would greatly appreciate any further feedback you might have. Please let us know if there are any remaining concerns or questions, so we can address them in time.

Thank you again for reviewing our paper.

Best regards,

Authors

Question 1: hyperparameter sensitivity of our algorithm, OFA.

Certainly. 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.

EmoC

HateSp18

Question 2: Training cost of our algorithm?

We 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. Compared 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.

Question 3: The omitted bias term.

We 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.

Few-shot performance of GPT2-XL. We quote the experiment results below from Table 4 in the appendix of our original submission.