📋 Paper Information
| Item | Content | |------|------| | Title | Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning | | Authors | Jihoon Kwon, Jiwon Choi, Jy-yong Sohn | | Affiliations | KAIST / Korea University | | arXiv | 2605.20730 | | Date | 2026-05-20 | | Categories | cs.CL / cs.AI | | Core claim | In-context learning (ICL) lets LLMs adapt quickly to new tasks, but inference cost grows linearly with the number of demonstrations. Task vectors can compress demonstrations into a compact hidden-state representation, yet prior work only evaluated their quality with indirect criteria like downstream accuracy. This paper takes "the task vector's predictive distribution should align with that of full ICL" as a direct criterion, proposes the d_NTP metric and Linear Task Vectors (LTV), improving average accuracy by 9.2%. It also finds that task vectors from large models transfer to smaller models (+6.4%). |
"Show me a few examples and I can learn this new task."
That sentence captures the entire appeal of In-Context Learning (ICL). Give a large language model a handful of (input, output) demonstrations—sometimes five, sometimes ten—and it rapidly adjusts its behavior and starts making correct judgments on new inputs. No fine-tuning, no gradients. Just the statistical structure of the context.
But there is a very unpleasant cost hidden here.
Every example is a string of tokens. Ten examples, in multi-task classification, can consume thousands of tokens. If you have ten thousand queries to run—each carrying the same set of examples—your inference cost explodes a hundredfold. The same examples are repeatedly encoded, repeatedly computed, repeatedly written into the KV cache, with no reuse.
A natural idea emerges: can the examples be "compressed" into something persistent? Distill the inductive bias of ten examples into one vector, and simply load that vector at inference time—eliminating the redundancy of repeated encoding.
This thing is called a task vector in the literature.
---
🧭 Existing Approaches to Task Vectors and Their Problems
Existing task-vector extraction methods broadly follow two paths.
The first path is finding directions in weight space. For certain tasks, researchers have found that a subspace of model parameters carries task-specific knowledge. Adjusting weights along this direction shifts the model's output toward the corresponding task. The drawbacks: modifying weights for every new task is inflexible at inference time, and vectors for different tasks may conflict.
The second path is extracting representations from hidden states. For each set of demonstrations, an encoder (or the model's own forward pass) produces a fixed-length vector representing the task's "semantics." At inference time, this vector is injected into the model's hidden layers to steer its predictions.
The second path is more flexible. But its core problem is as old as time: how do you know whether the extracted task vector is actually good?
Until now, the only answer has been: look at downstream accuracy.
That is the problem. Downstream accuracy is a distant, indirect signal. Along which dimensions does the vector diverge from the original ICL distribution? What is the mechanism linking that divergence to final performance? Nobody has answered.
---
📏 d_NTP: A Direct Alignment Metric
This paper does something refreshingly direct.
They define a metric called d_NTP (Next-Token Probability discrepancy). It precisely measures the difference between two things: the next-token probability distribution the model outputs when reasoning with the task vector, versus the distribution when reasoning with the full ICL examples.
The closer the two, the more faithfully the task vector represents the examples' inductive bias.
The definition is direct because it asks the most fundamental question: does the task vector actually make the model make the same choices it would with the full examples? Unpolluted by indirect, noisy signals like downstream accuracy, benchmark scores, or human evaluation.
The paper makes a key empirical finding: d_NTP is strongly negatively correlated with downstream accuracy. This means d_NTP can serve as a "performance proxy"—you don't need to run the entire test set to judge a task vector's quality. Just compare its probability distribution at sampling points to the full ICL distribution. The smaller the gap, the higher the final accuracy.
This finding turns task-vector evaluation from "guessing" into "measuring."
---
📐 Linear Task Vectors (LTV): A Closed-Form Solution
With d_NTP as a clear optimization target, extracting a task vector becomes a precisely solvable problem.
Existing methods rely on black-box distillation; you can't tell what information from the examples actually landed in the vector. LTV's approach: use regression to estimate each example's contribution to the model's predictive distribution.
Mathematically, LTV learns a closed-form linear mapping for each example. Given a set of demonstrations (input-output pairs), LTV directly computes their corresponding influence directions in the model's hidden-state space. The computation has an analytical solution—no iteration, no hyperparameter tuning, no reinforcement learning. One pass and you're done.
The core assumption is that each example's contribution to the model's final prediction is approximately linearly additive in hidden-state space. When the assumption holds, LTV's analytical solution naturally minimizes d_NTP; when it doesn't fully hold, the linear approximation is still good enough for many tasks.
The paper validates LTV on 8 classification benchmarks and 5 LLMs: a 9.2% average accuracy improvement over existing task-vector baselines, along with reduced inference latency. More importantly, LTV also outperforms baselines on regression tasks (continuous outputs)—something prior task-vector methods could not do at all.
---
🔗 Task Vectors from Large Models, Fed to Small Models
This is the paper's most tantalizing finding.
The researchers ran cross-model transfer experiments. LTV task vectors extracted from large models were injected into smaller models—without the small model ever seeing the examples, without any fine-tuning on the target model.
Result: the small model's performance improved by 6.4%.
The significance of that number far exceeds "6.4%." It shows that a task vector is not just a compression of examples—it captures, to some degree, abstract knowledge of "how to execute this task," and that knowledge transfers across model scales.
It's like a math professor summarizing a problem's solution on a note card. Not only can the professor use the card to solve the problem faster later—an undergraduate holding that same card can solve a problem they originally couldn't. The card carries not just the problem itself, but the professor's "way" of understanding it.
This cross-model transferability points to an extremely attractive future: use one very strong model to precompute LTV vectors for all common tasks, store them in a vector library, and any model of any size can call them at inference time—consuming no example tokens, occupying no KV cache, requiring no fine-tuning, directly boosting performance.
---
🤔 To Be Honest: Things We're Not Sure About
The limits of the linearity assumption. LTV's core assumption is that each example's contribution to the predictive distribution is approximately linearly additive. This holds reasonably well with few examples (≤10); but when the number of examples grows to 50 or 100, interaction effects between examples (some combinations synergize, some cancel out) will necessarily break the linearity assumption. At what example count does d_NTP begin to deviate from the linearly predictable range? The paper does not validate at very long context lengths.
Alignment strength of d_NTP for specific task types. d_NTP correlates strongly negatively with downstream accuracy on the classification and regression tasks tested. But for some tasks—those demanding creativity rather than accuracy, diversity rather than precision—d_NTP's logic may break down. For open-ended generation, what does "aligning with the full ICL predictive distribution" even mean? ICL itself has inherently high variance in open-ended generation—align with which version of ICL?
Consequences of LTV in safety-sensitive settings. If a task vector extracted from a large model is injected into a small model, does the small model also inherit the large model's safety training? Or, conversely—could some "undesirable preferences" of the large model leak to the downstream small model through the task vector? The paper does not discuss the possibility of task vectors carrying values or safety preferences.
Temporal stability. An LTV vector faithfully represents the examples' inductive bias at extraction time. But a month later, when the underlying model is updated (API upgrades, weight fine-tuning), is the old vector still valid? The paper does not test task vectors' robustness to model version drift.
d_NTP's counterintuitive trap. The paper shows lower d_NTP means higher accuracy. But could there be an "overfitting" regime—where pushing d_NTP extremely low makes the model over-imitate certain statistical noise of the full ICL, actually harming robustness when generalizing to unseen examples? The paper does not explicitly test this minimum paradox.
---
📚 References
1. Kwon, J., Choi, J., & Sohn, J. (2026). *Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning.* arXiv:2605.20730.
2. Hendel, R. et al. (2023). *In-Context Learning Creates Task Vectors.* EMNLP 2023.
3. Todd, E. et al. (2024). *Function Vectors in Large Language Models.* ICLR 2024.
4. Ilharco, G. et al. (2023). *Editing Models with Task Arithmetic.* ICLR 2023.
5. Brown, T. et al. (2020). *Language Models are Few-Shot Learners.* NeurIPS 2020.
6. Dong, Q. et al. (2023). *A Survey on In-Context Learning.* arXiv:2301.00234.