Official Review · ICLR 2024

This paper offers empirical evidence highlighting parallels between shallow linearized transformer models and the dynamics observed during practical Transformer training. These properties span various well-known aspects: the superiority of adaptive methods over SGD, the presence of heavy-tailed stochastic gradient noise, a robust condition number characterizing the optimization landscape, and the directional smoothness gap. The findings advocate for the potential of simple linear transformers as an insightful and accessible model to elucidate the mystery of Transformer optimization trajectories. The presented work is overall clear and of notable quality. Its innovative aspect is the unexpected empirical evidence indicating that, under certain conditions, shallow linear transformers can closely resemble conventional attention models, revealing a significant similarity between the simple model and its real-world analogs. Given the historical challenges associated with analyzing attention mechanisms involving softmax or other nonlinearities, it is commendable to identify such marked resemblances under broad observations using a simplified framework. This study illuminates the theoretical analysis for optimization on linear transformers, which makes a valuable contribution to the literature. The paper is somewhat limited in its data distribution and optimization objective setting to support its conclusion. The data distribution discussed in this setting focuses on in-context linear regression tasks—a largely statistical setting, rather than one that closely resembles language. This might be attributed to the limited expressivity of linear transformers. I fully understand that the authors are trying to *find the simplest abstraction* as a representative of the transformer’s landscape. However, since the landscape is likely heavily dependent on the data distribution, it would be more convincing to see a variety of data models that exhibit the same phenomenon in simple models. Furthermore, it's difficult to assert that the linearized model is sufficient to understand Transformer optimization, especially since the properties verified in this paper represent only a portion of the optimization narrative. - Are there other possible tasks that can be expressed by linear transformers? Do those tasks on simple linear transformers also exhibit those behaviors? - If we add other linearities (such as softmax) to the shallow neural networks, can you also show those properties along the training trajectory?