Timezone: »

Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse
Sotiris Anagnostidis · Luca Biggio · Lorenzo Noci · Antonio Orvieto · Sidak Pal Singh · Aurelien Lucchi

Tue Nov 29 02:00 PM -- 04:00 PM (PST) @ Hall J #121

Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers — the distinctive architectural component of Transformers — can result in rank collapse of the tokens’ representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens’ representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide a thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth-dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries, keys and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.

Author Information

Sotiris Anagnostidis (ETH Zurich)
Luca Biggio (ETH Zürich)
Lorenzo Noci (ETH Zürich)
Antonio Orvieto (ETH Zurich)

PhD Student at ETH Zurich. I’m interested in the design and analysis of optimization algorithms for deep learning. Interned at DeepMind, MILA, and Meta. All publications at http://orvi.altervista.org/ Looking for postdoc positions! :) antonio.orvieto@inf.ethz.ch

Sidak Pal Singh (ETH Zürich)
Aurelien Lucchi (Swiss Federal Institute of Technology)

More from the Same Authors