In-Context Learning State Vector with Inner and Momentum Optimization

Summary

This work introduces the novel concept of the state vector. It encapsulates the information of ICL examples from separate tokens as anchors. Drawing inspiration from the duality between transformer attention and gradient descent, the authors implement inner and momentum optimization to generate task-specific state vectors. To address the context length limitation of LLMs, a divide-and-conquer strategy is employed to aggregate information from multiple examples into a single state vector. Experimental results demonstrate that patching the state vector during the inference forward pass outperforms previously proposed task vector and function vectors.

Strengths

1. Although the idea of using vectors to capture in-context learning is not new, the proposed approach of optimizing the state vector using inner/momentum optimization is well-motivated. 2. The paper is well-organized and sound with experiments and ablation studies. 3. The experimental results and the analyses provide some useful insights.

Weaknesses

1. It is not clear whether larger models or more complex datasets can still benefit from this work. Is this work only applicable to relatively small models? 2. It is unclear why a dummy query has to be included when extracting the state vector, whose task is to encapsulate the examples only. What if not? 3. Can we use this vector to align LLMs further, considering encapsulating them into memory/vector bases for reuse and continuous learning of new knowledge?

Questions

## Summary of review In general, the paper presents an interesting concept, and the overall submission is technically sound. Therefore I am positive to this paper at this point. However, I hope the authors can address the several concerns raised in the weakness section by providing extra explanations and discussions regarding the unclear aspects. Also, please response to the additional questions shown as follows. If the authors can well address my concerns/questions, I will keep the positive accept rating. ## Additional questions 1. I would like to know whether the ICL state vector can be used in actual alignment scenarios. If we switch to a long-context version of the model, how will the ICL state vector be different? 2. The previous study demonstrates that the label words are information anchors [1], while this work aggregates the information of separate tokens in the state vector, is there any evidence that the separate tokens can also gather in-context information in a progressive manner? [1]Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. 2023. Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 9840–9855, Singapore. Association for Computational Linguistics.

Rating

7

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

n/a

Summary

The paper presents an in-depth analysis and optimization of In-Context Learning (ICL) within large language models (LLMs). In particular, it proposes inner and momentum-based optimization techniques to increase the performance of in-context learning.

Strengths

1. In-depth analysis of ICL vectors across 12 different datasets (linguistic, translation, and knowledge). The analysis provides insights into ICL vectors and motivates the optimization. 2. The optimization method can be divided into inner and momentum-based stages, each of which has shown promising improvement. 3. The ablation study is abundant. 4. The presentation of the research question is clear to read and follow.

Weaknesses

1. The method itself is somewhat simple. Specifically, the inner optimization is merely a uniform averaging. 2. The definition of state vectors lacks a rigorous theoretical basis, making it difficult to infer the approach's generalizability and reliability across different NLP tasks. 3. The experiments are conducted on Llama-2-7B and GPT-J-6B, which are limited in terms of scale and performance.

Questions

In principle, ICL plays a more significant role in larger models (also with better long-sequence performance). Does this method benefit these models?

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

NA

Summary

The authors focus on the compressed vectors in In-Context Learning (ICL). They highlights the similarities between the compressed vectors and parameters trained via gradient descent, proposing the formulation of the state vector. Then they applies two optimization algorithm to progressively refine state vector. The proposed method achieves SotA performance on diverse tasks.

Strengths

1.The authors introduce the formulation of state vector that encapsulates the processing state of ICL stored in the attention activations. 2.The proposed two optimization methods for the state vector seem to be effective. 3. The research topic is significant and intriguing

Weaknesses

1.The proposed state vector seems to be benefit from examples (as shown in Eqn. 5), could you explain how it works without demonstration in the zero-shot setting (as shown in Table 1)? 2.As ICL baseline makes predictions conditioned on the demonstrations, could you explain why its performance remains unchanged with different number of examples in Figure 2? 3.The proposed method is mainly draw inspiration from trainable parameters and gradient optimization algorithm. However, the performance of applying other classical gradient optimization algorithms has deteriorated significantly in Table 2. And the authors argue that “The outcome indicates a discrepancy between the state vector and updated parameters with gradient descent.” This confuses me a lot. The authors should state their similarities and differences more clearly and precisely.

Questions

See weakness.

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

See weakness.

Summary

This paper aims to revealing the working mechanism of the compressed state vector in context learning. They first prove the state vectors are similar with parameters trained via gradient descent. Then, they propose three methods to optimize such kind of parameters: (1) inner optimization averaging each vector of the separate token; (2) momentum optimization using momentum-based optimization algorithm; (3) divide-and-conquer strategy aggregating several vector groups. They test these methods on some medium-sized LLMs. Experiments show some improvements against other compressed vector baselines.

Strengths

1. The paper is well written and easy to follow. 2. The contribution of this paper is significant. It attempts to provide understanding of the mechanism of ICL via a gradient descent perspective. The proposed state vector can be viewed as trained parameters so that they can be optimized with future methods. 3. The proposed optimization methods are reasonable. Experiments with these methods support their claims. The performance can be improved in most scenarios.

Weaknesses

A substantive assessment of the weaknesses of the paper. Focus on constructive and actionable insights on how the work could improve towards its stated goals. Be specific, and avoid generic remarks. For example, if you believe the contribution lacks novelty, provide references and an explanation as evidence; if you believe experiments are insufficient, explain why and exactly what is missing, etc. Please keep in mind that the rebuttal period is not necessarily enough to run new experiments. As a result, asking for new results is rather unrealistic.

Questions

1. Does the state vector without optimization actually function the same as the task vector? If not, I think there should be another baseline called state vector (w/o optimization). The authors should claim whether this is a new compressed vector method or just a prove of similarity to gradient descent. 2. Can the authors provide more rationality of optimizing the state vector? I understand the state vector is a compression of previous context, but what it would present after the optimization? 3. Experiments show that the inner optimization and momentum optimization prefer different tasks. Is there a conclusion when to use which? Or is there a combination to adaptively cover all the situations? 4. Can the authors elaborate more on the differences between averaging and D&C? Why D&C is worse when the number of examples equals 10 (Figure 2)? In my understanding, they should be the same since the authors use 10-shot ICL (there should be only 1 group). 5. The authors should check the correctness of the experiment results. For example, in Table 5, they report Llama-2 Zero-shot Regular with six zeros but with an average of 0.2.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See questions

Summary

The paper considers the problem of using compressed vectors to replace explicit demonstrations, with potential benefits of providing new perspectives to better understand the mechanism of In-Context Learning (ICL), and addressing the issue of overly long demonstrations by summarizing them into a vector. The main challenge is the drop in accuracy. The paper proposes new methods to enhance compressed vectors, using ideas of ensemble averaging, momentum (though details are insufficient), and Divide-and-Conquer aggregation. Many ablation experiments are conducted to prove the effectiveness and answer different research questions.

Strengths

1. The paper conducts a large number of ablation experiments. 2. The methods proposed are quite novel, including the heuristic approaches and Divide-and-Conquer aggregation. 3. In practice, the methods perform very well on certain tasks. 4. The work provides new perspectives to better understand the mechanism of ICL.

Weaknesses

1. The development is still in its early stages, mostly heuristic, lacking any theoretical guarantees. However, I think this is not a critical issue, as the final performance is good in practice for certain tasks, and a rigorous analysis is very hard due to the complexity of transformers. 2. The presentation needs improvement, e.g., Figure 4 should have a color legend for different positions. The specific parameter $L$ used should be clarified in Section 5.1 3. The strict definition of the momentum gradient optimization algorithm used needs to be clearly stated.

Questions

1. What is the strict definition of the momentum gradient optimization algorithm used as opt? Is it Polyak Momentum or Nesterov momentum? 2. In the code, what do FixOptimizer, OneStepOptimizer, and FixOneStepOptimizer represent respectively? Which one represents momentum? 3. For Divide-and-Conquer aggregation, how does the performance compare to directly averaging the different group-specific state vectors?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

1. The presentation needs improvement. 2. It is worth considering larger models, e.g., Llama2-70B. But existing experiments on Llama-2-7B, Llama-2-13B and GPT-J-6B is strong enough. 3. It is worth discussing the performance on more complex ICL tasks, such as math and text summarization.