Provable Benefits of Task-Specific Prompts for In-context Learning

2025-03-03

AI Summary

This work addresses the interpretability of task prior modeling in in-context learning (ICL) for language models. We propose a collaborative “task prompt + dedicated prediction head” mechanism: task prompts, optimized via prompt tuning, explicitly model the mean of the conditional task distribution, while the prediction head—jointly optimized with attention weights—implicitly captures its variance. This constitutes the first theoretical explanation of ICL from a covariance-mean decoupling perspective. Leveraging analysis of a one-dimensional linear attention model and projected gradient descent, we rigorously prove that our design achieves exact separation of mean and variance estimation. Empirical results demonstrate that the method significantly improves ICL’s generalization performance and convergence speed, outperforming standard fine-tuning paradigms.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: (Large) Language Models
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

The in-context learning capabilities of modern language models have motivated a deeper mathematical understanding of sequence models. A line of recent work has shown that linear attention models can emulate projected gradient descent iterations to implicitly learn the task vector from the data provided in the context window. In this work, we consider a novel setting where the global task distribution can be partitioned into a union of conditional task distributions. We then examine the use of task-specific prompts and prediction heads for learning the prior information associated with the conditional task distribution using a one-layer attention model. Our results on loss landscape show that task-specific prompts facilitate a covariance-mean decoupling where prompt-tuning explains the conditional mean of the distribution whereas the variance is learned/explained through in-context learning. Incorporating task-specific head further aids this process by entirely decoupling estimation of mean and variance components. This covariance-mean perspective similarly explains how jointly training prompt and attention weights can provably help over fine-tuning after pretraining.

Problem

• Explores task-specific prompts for in-context learning in language models.
• Investigates covariance-mean decoupling using task-specific prompts and heads.
• Demonstrates benefits of joint training of prompts and attention weights.

Innovation

• Task-specific prompts enhance in-context learning.
• One-layer attention model learns conditional task distributions.
• Covariance-mean decoupling improves task-specific predictions.