[Paper Review] On the Role of Attention in Prompt-tuning

Samet Oymak, Ankit Singh Rawat+2 more|arXiv (Cornell University)|Jun 6, 2023

Topic Modeling4 citations

TL;DR

This paper provides the first theoretical analysis of prompt-tuning in attention-based models, showing that softmax-prompt-attention is provably more expressive than self-attention or linear-prompt-attention under a contextual mixture model. It demonstrates that gradient descent on prompt-attention achieves near-optimal sample complexity and enables precise attention to sparse, context-relevant tokens, with exponentially decreasing test error as context strength increases.

ABSTRACT

Prompt-tuning is an emerging strategy to adapt large language models (LLM) to downstream tasks by learning a (soft-)prompt parameter from data. Despite its success in LLMs, there is limited theoretical understanding of the power of prompt-tuning and the role of the attention mechanism in prompting. In this work, we explore prompt-tuning for one-layer attention architectures and study contextual mixture-models where each input token belongs to a context-relevant or -irrelevant set. We isolate the role of prompt-tuning through a self-contained prompt-attention model. Our contributions are as follows: (1) We show that softmax-prompt-attention is provably more expressive than softmax-self-attention and linear-prompt-attention under our contextual data model. (2) We analyze the initial trajectory of gradient descent and show that it learns the prompt and prediction head with near-optimal sample complexity and demonstrate how prompt can provably attend to sparse context-relevant tokens. (3) Assuming a known prompt but an unknown prediction head, we characterize the exact finite sample performance of prompt-attention which reveals the fundamental performance limits and the precise benefit of the context information. We also provide experiments that verify our theoretical insights on real datasets and demonstrate how prompt-tuning enables the model to attend to context-relevant information.

Motivation & Objective

• To theoretically understand the expressive power and optimization dynamics of prompt-tuning in attention-based models.
• To isolate the role of the softmax attention mechanism in enabling effective, sparse attention to context-relevant tokens.
• To characterize the finite-sample performance of prompt-attention when the prompt is known but the prediction head is unknown.
• To compare prompt-tuning's generalization and optimization behavior to self-attention and linear-prompt-attention under controlled data models.
• To validate theoretical insights with experiments on synthetic and real-world image classification datasets.

Proposed method

• Proposes a self-contained prompt-attention model based on a one-layer attention architecture with a soft prompt and linear classifier head.
• Introduces a contextual mixture-model where input tokens are partitioned into context-relevant (R) and irrelevant (R^c) sets, enabling controlled analysis of attention behavior.
• Analyzes the initial trajectory of gradient descent using a population-level model, showing that early steps learn the prompt and head with near-optimal sample complexity.
• Derives a population-level theorem (Theorem 2) proving that with sufficient context strength Q and step-size γ, test error decays exponentially as e^{-cQ²/σ²}.
• Provides a finite-sample analysis (Theorem 3) with de-biasing to correct for context vector influence, showing high-probability attention to relevant tokens and low test error.
• Uses a de-biasing step in training to isolate the role of the prompt in token selection from the final prediction, simplifying generalization analysis.

Experimental results

Research questions

• RQ1 How does prompt-tuning with softmax-attention compare in expressive power to self-attention and linear-prompt-attention in a controlled data model?
• RQ2 Can gradient descent on prompt-attention provably attend to sparse, context-relevant tokens while ignoring irrelevant ones during early training?
• RQ3 What is the finite-sample performance of prompt-attention when the prompt is known but the prediction head is unknown, and how does it depend on context strength and sparsity?
• RQ4 How does the softmax nonlinearity in attention facilitate optimization and generalization in prompt-tuning, particularly in high-dimensional, sparse settings?
• RQ5 What are the fundamental trade-offs between sparsity (fraction of relevant tokens), context strength, and correlation between prompt and classifier vectors?

Key findings

• Softmax-prompt-attention is provably more expressive than both self-attention and linear-prompt-attention under the contextual mixture model.
• The initial trajectory of gradient descent learns the prompt and prediction head with near-optimal sample complexity, achieving high accuracy in few steps.
• With sufficient context strength Q and step-size γ, the test error decays exponentially as e^{-cα²Q²/σ²}, achieving near-perfect attention to relevant tokens.
• In finite-sample settings, with n ≳ d(Q/ζW)^4 log⁵(nd), the model achieves high-probability attention to relevant tokens: aₜ ≥ (1−ε)/(ζT) for t ∈ R and aₜ ≤ ε/((1−ζ)T) for t ∉ R.
• The test error of the model is bounded by 2Te^{-cα²Q²}, demonstrating exponential improvement over linear models under sparse settings (ζ = O(1/√T)).
• The analysis reveals that effective attention to relevant tokens requires the prompt vector to dominate in energy relative to the classifier, especially when ρ ≠ 0.