On the Convergence of Gradient Descent on Learning Transformers with Residual Connections

TL;DR

This paper provides a theoretical analysis of the convergence behavior of gradient descent on single-layer and multi-layer Transformers with residual connections, highlighting their role in improving optimization stability.

Contribution

It offers the first convergence analysis of complete Transformer architectures with residuals, demonstrating their impact on training dynamics and stability.

Findings

Gradient descent converges linearly under proper initialization.

Residual connections improve conditioning of the attention output matrix.

Empirical results support the theoretical benefits of residuals in training stability.

Abstract

Transformer models have emerged as fundamental tools across various scientific and engineering disciplines, owing to their outstanding performance in diverse applications. Despite this empirical success, the theoretical foundations of Transformers remain relatively underdeveloped, particularly in understanding their training dynamics. Existing research predominantly examines isolated components--such as self-attention mechanisms and feedforward networks--without thoroughly investigating the interdependencies between these components, especially when residual connections are present. In this paper, we aim to bridge this gap by analyzing the convergence behavior of a structurally complete yet single-layer Transformer, comprising self-attention, a feedforward network, and residual connections. We demonstrate that, under appropriate initialization, gradient descent exhibits a linear convergence rate, where the convergence speed is determined by the minimum and maximum singular values of the output matrix from the attention layer. Moreover, our analysis reveals that residual connections serve to ameliorate the ill-conditioning of this output matrix, an issue stemming from the low-rank structure imposed by the softmax operation, thereby promoting enhanced optimization stability. We also extend our theoretical findings to a multi-layer Transformer architecture, confirming the linear convergence rate of gradient descent under suitable initialization. Empirical results corroborate our theoretical insights, illustrating the beneficial role of residual connections in promoting convergence stability.