Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth

Transformer feature maps undergo rank collapse when going deeper.

“Dong et al. (2021) mathematically reveals the rank collapse phenomenon when Transformer goes deeper.”

Dong et al. (2021) is the sole work in the literature that systematically and rigorously analyzes the depth scalability issue with deep ViT.

“However, to our best knowledge, Dong et al. (2021) is the sole work in the literature that systematically and rigorously analyzes this issue with deep ViT.”

ViT have been witnessed performance saturation when going deeper.

“Likewise, ViT have been witnessed performance saturation when going deeper.”

In Dong et al. (2021), the proof of Lemma A.3 implicitly assumes E_ij - E_ij' is bounded around zero; the authors confirmed a missed assumption is max_{i,j}(E_ij - E_ij') ≤ 1.

“In the proof of Lemma A.3, Taylor expansion was used to approximate and upper bound an exponentiation. However, to let the right-hand side upper bound satisfied, we conjectured that the authors implicitly assumed E ij -E ij ′ is bounded around zero. After directly communicating with the authors, they confirmed that a missed assumption here is max i,j ( E ij -E ij ′ ) ≤ 1 .”

In Dong et al. (2021), the proof of Lemma A.1 may assume R, W_V ≥ 0 entry-wise to let Eqn. (8)-(9) hold.

“In the proof of Lemma A.1, to let Eqn. (8)-(9) hold, the authors may have assumed R , W V ≥ 0 , where ≥ denotes entry-wise inequality.”

In Dong et al. (2021), Eqn. (12) may not be satisfied in general; a counterexample gives ||D1 softmax(r)^T R W_V||_1 = 4 while ||D1||_∞ ||R||_1 ||W_V||_1 = 3, disproving the claim.

“In the proof of Lemma A.1, we find Eqn. (12) may not be satisfied in general. We can raise the following counterexample: Since E , r , R , W V can be any matrices, we simply let D = diag(2 , 3) , softmax( r ) = [0 . 8 0 . 2] T , R = W V = I . Then ‖ D 1 softmax( r ) T RW V ‖ 1 = 4 while ‖ D 1 ‖ ∞ ‖ R ‖ 1 ‖ W V ‖ 1 = 3 , which disproves the claim.”