On the Role of Depth and Looping for In-Context Learning with Task Diversity

TL;DR

It is demonstrated that Looped Transformers -- a special class of multilayer Transformers with weight-sharing -- not only exhibit similar expressive power but are also provably robust under mild assumptions, and it is shown that Looped Transformers are the only models that exhibit a monotonic behavior of loss with respect to depth.

Abstract

The intriguing in-context learning (ICL) abilities of deep Transformer models have lately garnered significant attention. By studying in-context linear regression on unimodal Gaussian data, recent empirical and theoretical works have argued that ICL emerges from Transformers' abilities to simulate learning algorithms like gradient descent. However, these works fail to capture the remarkable ability of Transformers to learn multiple tasks in context. To this end, we study in-context learning for linear regression with diverse tasks, characterized by data covariance matrices with condition numbers ranging from $[1, κ]$, and highlight the importance of depth in this setting. More specifically, (a) we show theoretical lower bounds of $\log(κ)$ (or $\sqrtκ$) linear attention layers in the unrestricted (or restricted) attention setting and, (b) we show that multilayer Transformers can indeed solve such tasks with a number of layers that matches the lower bounds. However, we show that this expressivity of multilayer Transformer comes at the price of robustness. In particular, multilayer Transformers are not robust to even distributional shifts as small as $O(e^{-L})$ in Wasserstein distance, where $L$ is the depth of the network. We then demonstrate that Looped Transformers -- a special class of multilayer Transformers with weight-sharing -- not only exhibit similar expressive power but are also provably robust under mild assumptions. Besides out-of-distribution generalization, we also show that Looped Transformers are the only models that exhibit a monotonic behavior of loss with respect to depth.

Figures (2)

• Figure 1: We evaluate the test loss of looped models and multilayer model as the function of loops and depth respectively. In (a), we plot the test loss as a function of number of layers for three different covariance ranges. As predicted by theory, the larger the range ($\kappa$), the more layers are needed to get a small loss. In (b) we observe that the number of loops required to solve the problem is very close to the number of layers required for multilayer model.
• Figure 2: We evaluate the robustness of looped models and multilayer model as the level of deviation of the test distribution from the train distribution increases. The three plots differ in the number of distinct covariances used in the train distribution. The test distribution is supported on a window of size $w$ around the train covariances, and $w$ is varied on the x-axis. We see that in all settings the looped model is more robust than the multilayer model, especially on the left plot where there is least diversity in the training distribution, as predicted by the theory.

Theorems & Definitions (25)

• Definition 1
• Lemma 1: Effect of scaling on accuracy - restricted attention
• Theorem 1: Lower bound on the representation power of Transformers
• Remark 3.1
• Theorem 2
• Theorem 3: Termination guarantee
• Theorem 4: Restatement of Theorem 5.1 in fu2023transformers
• Theorem 5: Follows from Proposition 1 in von2023transformers
• Definition 2
• Theorem 6: Multilayer Transformers blow up out of distribution