Research map / Large language models / Language, safety and interpretability

Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning

NeurIPS 2026 · Large language models › Language, safety and interpretability · approach: In-context learning analysis · 0 citations (OpenAlex)

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.

Proposes d_NTP, a next-token-probability discrepancy metric between task-vector and ICL inference that correlates with accuracy, and Linear Task Vector (LTV), a closed-form ridge-regression linear mapping from zero-shot hidden states to ICL-induced hidden-state shifts.

Problem

Task vectors for in-context learning are evaluated only by downstream accuracy, which gives little insight into why one extraction method works better or how to design better ones, while ICL incurs growing inference cost with demonstration length.

Approach

Measure task-vector quality as KL divergence between TV-based and ICL next-token distributions (d_NTP), show it correlates negatively with accuracy, and derive a training-free closed-form linear mapping (LTV) that minimizes a hidden-state MSE proxy bounding d_NTP.

Key results

• LTV achieves average accuracy 79.46% on LLaMA-2-13B vs 70.22% for the best baseline (State Vector), a 9.2-point gain, and ranks best in 31 of 40 model-benchmark cases.
• On LLaMA-3.1-8B, LTV reaches 75.46% average accuracy with 0.0226 s inference latency, close to zero-shot latency (0.0220 s) and far below ICL (0.2238 s).
• Transferring LTV task vectors from Qwen-2.5-72B to Qwen-2.5-7B improves the small model's average accuracy by 6.43 points (73.57% to 80.00%).

Surprising findings

• Transferred task vectors from a 72B model outperform the 7B model's own LTV on most benchmarks, and the 80.00% average even exceeds the 72B model's own LTV (79.96%).
• Task-vector methods can exceed their own ICL accuracy on some individual tasks while overall ICL remains higher on average (e.g., LTV on Qwen-2.5-72B reaches 79.96% vs ICL 84.25%).

Limitations

• Linear mapping is the only evaluated form; the paper acknowledges more expressive mappings may help and when the linear approximation holds is unexplored.
• Evaluated mainly on short outputs (single-token classification and scalar regression), not longer generation or reasoning.
• Cross-model transfer requires the source and target models to share a tokenizer.

Experimental setup

Datasets: AGNews DBPedia HateSpeech18 MR SST-2 SST-5 SUBJ TREC

Benchmarks: AGNews DBPedia HateSpeech18 MR SST-2 SST-5 SUBJ TREC linear regression (synthetic) ReLU regression (synthetic)

Metrics: accuracy d_NTP (KL divergence) mean squared error inference latency extraction time

Baselines: Zero-shot ICL Function Vector Task Vector State Vector I2CL

Builds on: Task Vector Function Vector State Vector I2CL

Closest papers by meaning

• In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering (ICML 2024 · similarity 0.90)
• Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.90)
• Capturing In-Context Learning Dynamics with Task Operators (NeurIPS 2026 · similarity 0.90)
• Vector-ICL: In-context Learning with Continuous Vector Representations (ICLR 2025 · similarity 0.89)
• Variational Task Vector Composition (NeurIPS 2025 · similarity 0.89)
• Provable In-Context Vector Arithmetic via Retrieving Task Concepts (ICML 2025 · similarity 0.89)