Research map / Optimization / Optimizers
In-Context Learning State Vector with Inner and Momentum Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes compressed ICL attention-activation vectors as gradient-descent-like weights and proposes state vectors refined by inner (averaging), momentum optimization, and divide-and-conquer aggregation for test-time intervention.
Problem
Compressed task/function vectors extracted from ICL demonstrations are poorly understood and not well optimized, and ICL cannot directly handle many examples due to context length limits.
Approach
Views attention activations of the last separate token as a dual form of gradient-descent-updated weights, defines a state vector from concatenated early-layer activations, then refines it with uniform averaging over per-example state vectors (inner optimization), momentum applied to differences between adjacent state vectors (momentum optimization), and aggregates many examples via a divide-and-c
Key results
- In zero-shot on Llama-2-7B across six tasks, state vector (mom.) averages 65.2% versus 53.7% for task vector and 0.3% for regular prompting.
- In few-shot on Llama-2-7B, state vector (mom.) averages 73.8% versus 72.5% for the ICL baseline.
- Zero-shot state vector (mom.) reaches about 90% of ICL performance on Llama-2 and 78% on GPT-J, per the authors.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Average state vectors from separate tokens (inner optimization) (ablation) | State vector · Six-task average, zero-shot; Llama-2-7B | first-token accuracy: 53.7 -> 63.9; average improvement of 10.2% | Helped |
| Average state vectors from separate tokens (inner optimization) (ablation) | State vector · Six-task average, zero-shot; GPT-J-6B | first-token accuracy: 35.6 -> 41.5; average improvement of 5.9% | Helped |
| Add momentum optimization to the inner-optimized state vector (ablation) | State vector · Six-task average, zero-shot; Llama-2-7B | first-token accuracy: 63.9 -> 65.2; average increase of 1.3% | Helped |
| Add momentum optimization to the inner-optimized state vector (ablation) | State vector · Six-task average, zero-shot; GPT-J-6B | first-token accuracy: 41.5 -> 43.9; average increase of 2.4% | Helped |
| Add momentum optimization to the inner-optimized state vector (ablation) | State vector · Six-task average, few-shot; Llama-2-7B | first-token accuracy: 73.0 -> 73.8; average increase of 0.8% | Helped |
| Add momentum optimization to the inner-optimized state vector (ablation) | State vector · Six-task average, few-shot; GPT-J-6B | first-token accuracy: 56.4 -> 57.4; average increase of 1.0% | Helped |
| Use Adagrad, RMSprop, or Adam instead of momentum optimization (ablation) | Inner-optimized state vector · State vector optimization tasks | first-token accuracy | Hurt |
| Increase the number of layers used to extract the state vector, beyond the 14th layer (ablation) | State vector · Zero-shot, average across datasets | first-token accuracy | Hurt |
| Use alternative separate tokens instead of the arrow token (ablation) | State vector method · Experiments across tasks | first-token accuracy | No clear effect |
| Apply momentum optimization directly to task vectors without inner optimization (ablation) | Task vector · Preliminary experiment across tasks | first-token accuracy | No clear effect |
| Aggregate many examples using divide-and-conquer rather than average aggregation (ablation) | State vector aggregation · Multiple-example setting | accuracy | Helped |
| Increase the number of examples used for state vector aggregation (ablation) | Divide-and-conquer and average aggregation · Multiple-example setting | accuracy | Helped |
Surprising findings
- Applying momentum optimization directly to task vectors did not give average improvements, which the authors attribute to poor task vector robustness.
- First-order optimizers (Adagrad, RMSprop, Adam) substantially degraded state vector performance, unlike momentum.
Limitations
- Definition of state vectors rests on strong assumptions (linear attention approximation) and lacks rigorous theoretical grounding.
- Experiments limited to moderate-sized models (7B/6B, 13B) and datasets.
- Few-shot inner optimization can underperform ICL baseline due to noise from dummy queries.
Experimental setup
Datasets: Antonym Capitalize Present-Past Singular-Plural English-French Country-Capital AG News Person-Sport Person-Instrument Person-Occupation Product-Company Landmark-Country
Benchmarks: Antonym English-French Person-Instrument Person-Occupation Product-Company Landmark-Country Capitalize Country-Capital Present-Past Singular-Plural Person-Sport AG News
Metrics: first output token accuracy
Baselines: Regular zero-shot ICL baseline Function vector Task vector Average aggregation
Builds on: task vector function vector model soup momentum gradient descent
Closest papers by meaning
- Vector-ICL: In-context Learning with Continuous Vector Representations (ICLR 2025 · similarity 0.89)
- In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering (ICML 2024 · similarity 0.89)
- Position: Do pretrained Transformers Learn In-Context by Gradient Descent? (ICML 2024 · similarity 0.89)
- In-Context Language Learning: Architectures and Algorithms (ICML 2024 · similarity 0.89)
- Enhancing In-Context Learning Performance with just SVD-Based Weight Pruning: A Theoretical Perspective (NeurIPS 2024 · similarity 0.88)
- Implicit In-context Learning (ICLR 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.