Summary
This paper addresses the in-context uncertainty quantification for pre-trained transformers, focusing on predicting both the mean and variance of the conditional distribution of a target variable given input features. They derive a Bayes-optimal solution for the task and demonstrate that transformers are capable of performing in-context uncertainty quantification.
1. The authors rigorously analyze the uncertainty quantification problem when Transformers are constrained to a context window of size \( S \). They derive a generalization bound of \( \tilde{O}(\sqrt{\min\{S, T\}/(nT)}) \) for pre-training over \( n \) tasks with sequence length \( T \), extending to other cases under the assumption of bounded and Lipschitz loss functions. They laim that this bound is tighter than existing analyses, particularly when \( S < T \), and is the first of its kind to utilize the context-window structure to model a Markov chain over the prompt sequence and construct an upper bound on mixing time. They also discuss the additional approximation error due to a finite context window and quantify the trained Transformer's risk convergence to the Bayes-optimal risk.
2. The paper provides an empirical evaluation of the in-context learning abilities of transformers for uncertainty quantification across three distribution shift scenarios: task shift, covariate shift, and prompt length shift. The results suggest that transformers can learn both mean and uncertainty predictions in-context, even under moderate task distribution shifts, when trained on data with substantial task diversity. Moreover, the authors show that increasing task diversity through meta-learning enhances robustness to covariate shifts. Finally, they observe that removing positional encoding from the embedding vector significantly improves generalization performance.
Strengths
1. This paper is fairly written even though I do not understand some parts. Their problem setup is easy to follow and their organization is good (except the section 2.1). However, I could not understand why they write transformers with capital T.
2. The authors develop a good framework for studying the uncertainty quantification. Their theoretical analysis of the generalization properties of the loss function with the uncertainty quantification is interesting. They empirically observe the performance of transformers under task shift, distribution shift, and length shift.
3. The observation that the positional embedding worsens the performance of length generalization is valuable.
4. The way they define the loss function with respect to $\hat{\sigma}$ and $\hat{y}$ is valuable and they provide the Bayesian optimal of these parameters.
Weaknesses
1. I do not understand some parts of section 2.1 ( I think this section is written poorly). I think there is a gap between their definitions of IWL and ICL and the loss function the authors use. It would be better if the authors explain the relationship between their loss function and their definitions of ICL and IWL.
1.1. I am not able to follow the discussion of the pattern and sequence in section 2.1. I will ask some questions in the questions section.
1.2. The authors claim that the understanding of uncertainty quantification objective is important for understanding the contribution of in-context learning compared to the contribution of in-weight learning. This is fair. However, towards end of the subsection, they define IWL(M) and ICL(W) based on the worst and best case test distribution, which are not directly related to variance of the model. Therefore, the way of defining the loss function with the variance does not perfectly align with the definition of IWL(M) and ICL(M). Furthermore, they do not utilize (at least I did not notice) the definitions of IWL(M) and ICL(M) in their paper. When two functions are defined at the beginning of the paper (specifically in problem setup), I expect that the functions are frequently used, which is not the case in the paper.
2. Theorem 3.2 is suboptimal, and their comparison to previous work, specifically Li et al. (2023b), is not fair.
2.1. I think that Theorem 3.2 is suboptimal because I do not understand why the generalization performance of the transformers decreases with the context length. Consider the case when $S = T/4$ and $S = T/2$. From Theorem 3.2, I observe that the generalization performance is worse when $S = T/2$, but I think the generalization performance should improve when the context length increases. How can increasing the context length decreases the generalization performance. Please correct me if my reasoning fails.
2.2. The authors compare their bound with Li et al. (2023b) and Zhang et al. (2023b) in an unfair way. The bounds obtained in Li et al. (2023b) is better than the current paper when $S \geq T$. Li et al. (2023b) does not cover the case where $S << T$. In order to compare the results in the case of $S<T$, the authors change the algorithmic stability assumption (in line 297-298) and compare their results. However, when the algorithmic stability term is $\mathcal{O}(1/T)$, then I think the bounds in Li et al. (2023b) still better than the current paper.
Questions
1. Is the loss function devised by the authors and is there any previous works that study the same loss function? Could the authors provide a brief discussion for this loss function whether a similar loss function has been utilized in the literature?
2. What is pattern and sequence? Could the authors define these terms rigorously? Is the pattern a distribution? The authors use $\mathcal{P}$ for distribution first, then use the same notation for a pattern.
3. How is the loss function defined based on a pattern? It is used in the definition of IWL and ICL, but I could not find the definition.
4. Could you provide a more rigorous argument about the definitions of IWL and ICL?