In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

We study the in-context learning (ICL) ability of a Linear Transformer Block (LTB) that combines a linear attention component and a linear multi-layer perceptron (MLP) component. For ICL of linear regression with a Gaussian prior and a non-zero mean, we show that LTB can achieve nearly Bayes optimal ICL risk. In contrast, using only linear attention must incur an irreducible additive approximation error. Furthermore, we establish a correspondence between LTB and one-step gradient descent estimators with learnable initialization (𝖦𝖣-β), in the sense that every 𝖦𝖣-β estimator can be implemented by an LTB estimator and every optimal LTB estimator that minimizes the in-class ICL risk is effectively a 𝖦𝖣-β estimator. Finally, we show that 𝖦𝖣-β estimators can be efficiently optimized with gradient flow, despite a non-convex training objective. Our results reveal that LTB achieves ICL by implementing 𝖦𝖣-β, and they highlight the role of MLP layers in reducing approximation error.