Learning Task Representations from In-Context Learning

2025-02-08

AI Summary

This work addresses the fundamental challenges of opaque implicit task encoding in large language models (LLMs) under in-context learning (ICL) and the poor cross-modal generalizability of existing methods. We propose a learnable task vector framework that constructs task-semantic representations via attention-head-weighted aggregation—enabling, for the first time, unified task extraction across textual and functional regression modalities. To enhance interpretability and fidelity, we introduce causal gradient-based optimization of attention weights and design the first benchmark for task fidelity evaluation. Our method integrates Transformer attention analysis, differentiable modeling, causal gradient descent, and latent-state distribution alignment. Experiments demonstrate consistent and significant improvements over strong baselines on both text and regression tasks; learned task vectors faithfully reproduce ICL behavior; and ablation studies confirm that aligning the final-layer latent-state distributions is the primary driver of performance gains.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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. Moreover, ablation studies show that our method's effectiveness stems from aligning the distribution of the last hidden state with that of an optimally performing in-context-learned model.

Problem

• Understanding task encoding in large language models.
• Generalizing task vectors across different modalities.
• Evaluating task fidelity in functional regression tasks.

Innovation

• Automated task encoding via attention heads
• Causal optimization of task vector weights
• Cross-modal task vector generalizability benchmark