Task Descriptors Help Transformers Learn Linear Models In-Context

Authors: Ruomin Huang, Rong Ge
Primary area: learning theory
Keywords: Transformer, in-context learning, linear regression, task descriptor

Abstract:
Large language models (LLMs) exhibit strong in-context learning (ICL) ability, which allows the model to make predictions on new examples based on the given prompt. Recently, a line of research (Von Oswald et al., 2023; Aky¨urek et al., 2023; Ahn et al., 2023; Mahankali et al., 2023; Zhang et al., 2024) considered ICL for a simple linear regression setting and showed that the forward pass of Transformers is simulating some variants of gradient descent (GD) algorithms on the in-context examples. In practice, the input prompt usually contains a task descriptor in addition to in-context examples. We investigate how the task description helps ICL in the linear regression setting. Consider a simple setting where the task descriptor describes the mean of input in linear regression. Our results show that gradient flow converges to a global minimum for a linear Transformer. At the global minimum, the Transformer learns to use the task descriptor effectively to improve its performance. Empirically, we verify our results by showing that the weights converge to the predicted global minimum and Transformers indeed perform better with task descriptors.

Decision: accept (poster)
Reviewer scores: 8, 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730023394778

thread_type:
official_review

<thread_object>
creator: reviewer_usgc

<content>
summary:
This paper studies in-context learning with task descriptions. It uses a synthetic setup of linear regression, where the means from which the inputs are samples serve as the task description. Theory shows that a single-layer linear attention model can reach the optimal solution and the optimal solutions are characterized. Experiments show that a single-layer can indeed reach the optimal solution, but deeper models still perform better.

strengths:
- Important research question on the effect of task descriptions in ICL
- Simple synthetic setup to study 
- Theory showing optimality of a specific model class

weaknesses:
- The biggest issue I have is whether the synthetic setup is a good proxy to study ICL, and specifically task descriptors. I understand similar synthetic setups were used to study ICL. But why is giving the mean a good task description? How does it simulate the natural language case, for example the one cited from Brown et al.? 
- It would be helpful to explicitly highlight the insights drawn from the various lemmas and theorems throughout section 4. Unfortunately I had a hard time following this part and so I indicate this is a low-confidence review. 
- The experiments are a nice edition to the theory, but I'm confused about the single-layer being inferior to the deeper models. 
- See other comments and questions below.

questions:
1. It's quite puzzling to understanding the embedding matrix formulation of the problem. I'm used to seeing in-context examples as running examples in text, which are made of words that are embedded. But here, the embedding matrix itself is the object to learn. More specifically, given this embedding matrix, the objective is to learn the bottom-right entry, which contains the prediction y_query. What are the trainable parameters? Is E itself updated during training? How does this formulation relate to the usual ICL setup of running text? 

2. What are multiple tasks mentioned in equation 6 and does it make sense to take the expectation over all of them?  What is a task specifically?

3. What are the different sequences near equation 7? Are there multiple sequences per task? One sequence per task? 

4. Can you motivate the initialization in page 4? any clearer motivation besides having two matrices have the same norm?

5. I am possibly missing some background to understand this, but what do you mean by "We run gradient flow on the population loss"? How is the optimization done exactly? 

6. What is the significance of characterizing the optimal solutions in section 4? How should they be interpreted and what does it tell us about ICL with descriptions more broadly? 

7. If a single linear layer transformer can reach the optimal solution, then why do the experimental results show deeper models to perform better? Is it because of training difficulties with the single layer case? Is it just a sample complexity issue, with the single-layer model not having enough samples?

8. Since there's no pre-training and fine-tuning going on here, it seems like all instances of "pretraining" could just be changed to "training".
</content>
</thread_object>

</thread>
<thread>
creation date:
1730091784239

thread_type:
official_review

<thread_object>
creator: reviewer_yqse

<content>
summary:
This paper presents a theory analysis of the task descriptor usage in in-context learning (ICL)  for Transformer models.
The analysis is conducted based on a simple linear regression problem (i.e., given $x$ drawn from Gaussian $(\mu, \mathcal{I})$ predict $y$  = $w^Tx$ where $w$ is the underlying generative weights), and the task descriptors here are the $\mu$ for the input distribution.

Theoretically, the analysis results show that the provided task descriptors could help the Transformer directly use the $\mu$ provided in specialized input embedding given infinite samples $n$. Furthermore, the authors show that the model guaranteed to converge to the global optima from a a reasonable initialization. Empirical results show that the task descriptor could help a 1-layer Transformer achieve significantly lower loss, and the findings generalize to multi-layer Transformer model as well.

strengths:
- Overall solid analysis (theory and empirical ) for the task descriptor usage in ICL of Transformer models.

weaknesses:
- The analysis and results are overly narrow, focusing solely on linear regression. This limited scope makes it difficult to generalize the findings to fundamental tasks like sentiment classification with in-context learning (ICL).

- The paper's restricted scope can be partially attributed to defining the task descriptor as the mean $\mu$. This is a non-standard approach, as task descriptions in practical ICL applications typically convey semantic instructions (e.g., `Classify the review as positive or negative.`) rather than characterizing the input distribution $x$ (e.g., the lexical and syntactic patterns of review text).

- (Minor) Figure 2 reveals that the performance gap between models with and without task descriptors narrows as model complexity increases (from 1-layer to 2-layer architectures). This raises questions about the utility of task descriptors in modern large-scale models. The diminishing effect of task descriptors with increasing model size may limit the long-term relevance of this analysis, particularly given the trend toward ever-larger models.

questions:
- Would the loss curve comparison change with larger models and complex data samples (higher dimension than 5)? 12-Layer Transformer such as GPT-2?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730357960519

thread_type:
official_review

<thread_object>
creator: reviewer_wnxg

<content>
summary:
This paper proposes a new mathematical setting for theoretical in-context learning analysis. It propose a setting which models the task description in in-context learning. The author discussed the optimal soluations for a 1-layer linear self-attention network in two situations, where samples are sufficient or limited, and prove that such optimal soluations can be achieved under gradient flow training with a reasonable initialization. Also, empirical verification for 1-layer transformers are also provided.

strengths:
1.	The author models the task description part in in-context learning and provide comprehensive analysis.
2.	Both limited-sample and infinite-sample situations are considered. 
3.	Experiments further strengthened their analysis.

weaknesses:
1.	The conclusions based on a 1-layer linear self-attention model and gradient-flow might be a little weak.
2.	Though introducing task description is very interesting, the techniques used and the theorems got might be a little incremental.

questions:
1.	Might the author provide some insights/comparisons about what real difference the task description brings for ICL? It would be very helpful.
2.	Proofs of more complex situation than gradient flow / 1-layer LSA would be very helpful, especially for cases where global optimal can not achieved (which I guess is more often in reality?)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730647416980

thread_type:
official_review

<thread_object>
creator: reviewer_13fd

<content>
summary:
The paper explores the role of task descriptors in enhancing the in-context learning (ICL) capabilities of Transformers. The authors show that incorporating task descriptors into the input prompt leads to improved ICL performance for linear models. They present theoretical findings indicating that gradient flow converges to a global minimum for simple linear models when task descriptors are used, and they empirically verify these results by demonstrating weight convergence to the predicted global minimum. The paper includes a theoretical analysis of how task descriptors facilitate in-context learning in linear regression settings, illustrates how Transformers can leverage these descriptors to enhance performance, and provides empirical validation of the theoretical findings through experiments with various Transformer configurations.

strengths:
- The paper offers a novel perspective on how task descriptors can improve the ICL performance of Transformers.
- The paper presents a theoretical analysis of how task descriptors enhance in-context learning in linear regression settings, showing how Transformers can utilize these descriptors to boost performance. 
- It also offers empirical validation of the theoretical findings through experiments conducted with different Transformer configurations.

weaknesses:
- While the paper provides a theoretical analysis, it may lack a deeper exploration into the underlying mechanisms of how Transformers integrate task descriptors during the learning process.
- The experiments, although comprehensive, are limited to a specific setup. More diverse datasets and problem settings could provide a more robust evaluation.

questions:
- What are the specific mechanisms through which task descriptors improve ICL in Transformers? Further analysis or experiments that shed light on this process would be valuable.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732158194276

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
We sincerely thank all reviewers for constructive comments. We revised our paper according to the issues and concerns that reviewers proposed. Here is a summary of the updates.

1.  We add descriptions of how transformers use task descriptors in our setting. 
* For the setting with infinite in-context examples, the task descriptors enable transformers to do data standardization (subtract the mean from each $x_i$). Specifically,  the optimal solution given in Theorem 3.1 subtracts the mean $\mu_\tau$ from the $x_i$ in key values and do one step of preconditioned GD to output $$\begin{aligned}
    \hat y_{query}
&=x_{query}^{\top} \Lambda^{-1} \frac{1}{n} \sum_{i=1}^n  (x_i-\mu_\tau) y_i\\\\
&= x_{query}^{\top} \Lambda^{-1} \left(\frac{1}{n} \sum_{i=1}^n (x_i-\mu_\tau)x_i^\top\right) w \\\\
&\to x_{query}^{\top} w \text{ ~~as } n\to\infty, 
\end{aligned}$$ which is correct. While without descriptors, one can only get $$\begin{aligned}
    \hat y_{query}
&=x_{query}^{\top} A \frac{1}{n} \sum_{i=1}^n  x_i y_i\\\\
&= x_{query}^{\top}A \left(\frac{1}{n} \sum_{i=1}^n x_i x_i^\top\right) w \\\\
&\to x_{query}^{\top} A (\Lambda + \mu_\tau\mu_\tau^\top)w \text{ ~~as } n\to\infty,
\end{aligned}$$ which would have a non-zero loss for any preconditioner $A$ in the presence of an unfixed $\mu_\tau$. 
* For the setting with finite in-context examples, transformers learned a more delicate way to use the descriptors in not only keys but also queries to minimize both the bias and variance. This implies that the transformers trained on fewer in-context examples could leverage descriptors more.

2. We also emphasize in the intro that in the more general setting the transformer may have many other potential ways of leveraging the task descriptor, this paper only gives one example.

3. We add a new experiment varying the data dimension $d$. The experiment shows that as the ICL linear regression has higher dimension $d$, the performance gap widens between transformers trained with and without descriptors. Please see table 1 in Section B.2 at appendix for details.

4. We extend our experiments by introducing one-hot vectors as descriptors to indicate different distributions of $x$. We consider another mean-varying linear regression setting where there are $k$ spherical Gaussian distributions with different means $\mu_1,\dots,\mu_k$. Each ICL linear regression task $\tau$ independently samples a $\mathcal{N}(\mu_i, I_d)$ from them as the input distribution of $x$. Instead of directly using the mean $\mu_i$ as the task descriptor, we assign an one-hot ``task-identity'' vector $e_i$ to each of the Gaussian distributions $\mathcal{N}(\mu_i, I_d)$ and use this task identity vector as the task descriptor. 
As a summary, there are two key modifications:
- The mean $\mu_\tau$ for each task $\tau$ is uniformly drawn from a fixed set of $k$ vectors ${\mu_1, \mu_2, \cdots, \mu_k}$. 
- When $\mu_\tau = \mu_i$ for some $1 \leq i\leq k$, we use the one-hot vector $e_i$ (1 at position $i$, 0 elsewhere) as the task descriptor.
 The corresponding input embedding matrix becomes $$E_\tau=\left(\begin{array}{ccccc} e_i & e_i & \dots & e_i& e_i\\\\x_{\tau, 1} & x_{\tau, 2} & \dots & x_{\tau, n} & x_{\tau, \text { query }} \\\\y_{\tau,1} & y_{\tau,2} & \dots & y_{\tau,n} & 0\end{array}\right).$$
In this experiment, the descriptors also improve the performance of ICL. Please see Section B.5 at the appendix for details. 

5. We move Lemma 4.2 to the appendix as it is a bit orthogonal to the main discussion in Section 4.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors study the linear attention, in-context learning dynamics problem, and ask whether task descriptions affect the gradient-descent like learning dynamics in these papers. The paper presents nice theoretical analysis of these linear ICL dynamics and also presents some simple empirical experiments to back this up. The reviewers are pretty consistent on both the strengths and the weaknesses, with the strength being the theory and toy experiments, and the weakness being the very limited setting of linear attention ICL.

The weakness is generally minor - there's been fairly extensive (both empirical and theory) interest in understanding the linear ICL setting, and the paper seems to be a solid contribution to this area.