# In-Context Learning State Vector with Inner and Momentum Optimization

arXiv: https://arxiv.org/abs/2404.11225 · HTML: https://mlsanity.com/paper/2404.11225/

## Key facts

- Cited by 15 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 44 references that are in the map.
- First posted on arXiv 17 Apr 2024.

## Builds on (most cited first, 30 of 32)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Language Models are Super Mario: Absorbing Abilities from Homologous Models as a Free Lunch](https://mlsanity.com/paper/2311.03099/) (2023)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [Word Translation Without Parallel Data](https://mlsanity.com/paper/1710.04087/) (2017)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [In-Context Learning Creates Task Vectors](https://mlsanity.com/paper/2310.15916/) (2023)
- [Data Distributional Properties Drive Emergent In-Context Learning in Transformers](https://mlsanity.com/paper/2205.05055/) (2022)
- [Function Vectors in Large Language Models](https://mlsanity.com/paper/2310.15213/) (2023)
- [In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering](https://mlsanity.com/paper/2311.06668/) (2023)
- [Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning](https://mlsanity.com/paper/2305.14160/) (2023)
- [Linearity of Relation Decoding in Transformer Language Models](https://mlsanity.com/paper/2308.09124/) (2023)
- [The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives](https://mlsanity.com/paper/1909.01380/) (2019)
- [Task-Specific Skill Localization in Fine-tuned Language Models](https://mlsanity.com/paper/2302.06600/) (2023)
- [Language Models Implement Simple Word2Vec-style Vector Arithmetic](https://mlsanity.com/paper/2305.16130/) (2023)
- [AdapterSoup: Weight Averaging to Improve Generalization of Pretrained Language Models](https://mlsanity.com/paper/2302.07027/) (2023)
- [Separate the Wheat from the Chaff: Model Deficiency Unlearning via Parameter-Efficient Module Operation](https://mlsanity.com/paper/2308.08090/) (2023)
- [In-context Learning and Gradient Descent Revisited](https://mlsanity.com/paper/2311.07772/) (2023)

## Built on by (most cited first, 15 of 15)

- [CODI: Compressing Chain-of-Thought into Continuous Space via Self-Distillation](https://mlsanity.com/paper/2502.21074/) (2025)
- [Learning Task Representations from In-Context Learning](https://mlsanity.com/paper/2502.05390/) (2025)
- [Understanding Task Vectors in In-Context Learning: Emergence, Functionality, and Limitations](https://mlsanity.com/paper/2506.09048/) (2025)
- [ELICIT: LLM Augmentation via External In-Context Capability](https://mlsanity.com/paper/2410.09343/) (2024)
- [From Compression to Expression: A Layerwise Analysis of In-Context Learning](https://mlsanity.com/paper/2505.17322/) (2025)
- [Label Words as Local Task Vectors in In-Context Learning](https://mlsanity.com/paper/2406.16007/) (2024)
- [Optimization-Inspired Few-Shot Adaptation for Large Language Models](https://mlsanity.com/paper/2505.19107/) (2025)
- [Context Learning for Multi-Agent Discussion](https://mlsanity.com/paper/2602.02350/) (2026)
- [Take Off the Training Wheels Progressive In-Context Learning for Effective Alignment](https://mlsanity.com/paper/2503.09958/) (2025)
- [Task-driven Layerwise Additive Activation Intervention](https://mlsanity.com/paper/2502.06115/) (2025)
- [Multi-Layer Attention is the Amplifier of Demonstration Effectiveness](https://mlsanity.com/paper/2508.00385/) (2025)
- [Unleashing Hierarchical Reasoning: An LLM-Driven Framework for Training-Free Referring Video Object Segmentation](https://mlsanity.com/paper/2509.05751/) (2025)
- [Hyperbolic Optimization](https://mlsanity.com/paper/2509.25206/) (2025)
- [Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning](https://mlsanity.com/paper/2605.20730/) (2026)
- [Capturing In-Context Learning Dynamics with Task Operators](https://mlsanity.com/paper/2610.01054/) (2026)

## Abstract (arXiv)

Large Language Models (LLMs) have exhibited an impressive ability to perform In-Context Learning (ICL) from only a few examples. Recent works have indicated that the functions learned by ICL can be represented through compressed vectors derived from the transformer. However, the working mechanisms and optimization of these vectors are yet to be thoroughly explored. In this paper, we address this gap by presenting a comprehensive analysis of these compressed vectors, drawing parallels to the parameters trained with gradient descent, and introduce the concept of state vector. Inspired by the works on model soup and momentum-based gradient descent, we propose inner and momentum optimization methods that are applied to refine the state vector progressively as test-time adaptation. Moreover, we simulate state vector aggregation in the multiple example setting, where demonstrations comprising numerous examples are usually too lengthy for regular ICL, and further propose a divide-and-conquer aggregation method to address this challenge. We conduct extensive experiments using Llama-2 and GPT-J in both zero-shot setting and few-shot setting. The experimental results show that our optimization method effectively enhances the state vector and achieves the state-of-the-art performance on diverse tasks. Code is available at https://github.com/HITsz-TMG/ICL-State-Vector