Learning Task Representations from In-Context Learning (2025)

Abstract: Large language models (LLMs) have demonstrated remarkable proficiency in in-context learning (ICL), where models adapt to new tasks through example-based prompts without requiring parameter updates. However, understanding how tasks are internally encoded and generalized remains a challenge. To address some of the empirical and technical gaps in the literature, we introduce an automated formulation for encoding task information in ICL prompts as a function of attention heads within the transformer architecture. This approach computes a single task vector as a weighted sum of attention heads, with the weights optimized causally via gradient descent. Our findings show that existing methods fail to generalize effectively to modalities beyond text. In response, we also design a benchmark to evaluate whether a task vector can preserve task fidelity in functional regression tasks. The proposed method successfully extracts task-specific information from in-context demonstrations and excels in both text and regression tasks, demonstrating its generalizability across modalities.

## introduction large language models (llms) based on the transformer architecture @xcite have seen dramatic improvements in recent years. a notable feature of these models, such as @xcite , is their capability for in-context learning (icl). this process involves the model receiving a prompt that includes demonstrations of a task in the form of input-output pairs. when presented with a new query input, the model can generate the appropriate output by extrapolating from the provided examples. for instance, after being prompted with a few examples, these models are capable of producing the antonyms of given input words. a * corresponding author: baturay.saglam@yale.edu concrete example is @xmath0, @xmath1, examples @xmath2, (1) where the blue text is the prompt and the red text is the completion provided by the model. not limited to linguistic tasks, transformers can also in-context learn a general class of functions f @xcite . specifically, for any function @xmath3, the model is capable of approximating f (x query ) for a new query input x query . this class may include linear or nonlinear relationships, which can be represented by various machine learning models such as linear regression, multi-layer relu networks, and decision trees. the icl capability is particularly intriguing as it allows models to adapt to a wide range of downstream tasks on-the-fly-i.e., without requiring any parameter updates post-training @xcite . this indicates that models can extract abstract task information (the relationship between input-output pairs) from the prompt and use it to generate accurate responses for query inputs. however, due to the complex nature of the transformer architecture, the computational mechanisms that facilitate how models internally extract and encode task information are not well understood. task vectors have been identified as key mechanisms for encoding task information during icl @xcite @xcite . these vectors capture how models process prompts and embed the abstract incontext task as a numerical representation. while the effectiveness of these vectors is empirically well-supported, the precise computation of a task vector within the transformer architecture remains debated. early findings suggest that layer or attention activations significantly influence icl performance @xcite @xcite . however, there is no consensus on the optimal way to conceptualize or compute these vectors, resulting in varied methodologies that often excel in linguistics but fail to generalize to other tasks, such as functional regression. since these methods rely on the internal components of transformers and architectural differences due to output modality (e.g., language head or numeric outputs) do not change their operation, any meaningful task representation should, in principle, be agnostic to output modality. poor non-language performance therefore suggests that they may be capturing modality-specific patterns rather than a generalizable task representation. additionally, these methods require adaptation to individual transformer architectures. for example, the layers over which the task vector is computed must be explicitly identified for each model @xcite @xcite . these limitations raise questions about the generalizability and practical utility of existing task vector formulations. to pursue a more standardized and automated pipeline for task encodings, we build on the recent finding that attention heads are the primary driver for vectorizing icl tasks and hypothesize that: in-context learning tasks can be effectively represented as a weighted sum of all attention heads in a transformer, with the weights learned causally through gradient descent. ## technical preliminaries we focus on decoder-only autoregressive transformer architectures. ## conclusion we investigated the internal representation of incontext learning (icl) tasks in large language models (llms). we propose the learnable task vector (ltv) framework, which encodes icl tasks as a weighted sum of attention head activations optimized through causal gradient descent. ltv generalizes across both text and regression tasks, addressing the limitations of existing representations. through empirical validation using a novel benchmark designed to analyze task encodings, we demonstrate that ltv preserves task fidelity when models fail to extract task information from demonstrations. our findings further suggest that this preservation is achieved by maintaining the last hidden state distribution similar to that of an optimally in-context-learned model, bringing us closer to understanding how task steering works in general.