Problem $4$: Why output projection on MHA? Consider the standard multi$-$head self$-$attention

$($MHA$)$ layer defined by

$)1)H$

where

$W^{O}in{R}^{H{d}_{\text{h}\text{e}\text{n}\text{d}x{d}_{\text{o}\text{u}\text{t}}}}$

$W_h^{Q}in{R}^{dq{d}_{\text{a}\text{t}\text{t}\text{n}}},W_h^{K}in{R}^{d_K{d}_{stts}},W_h^{V}in{R}^{dV{d}_{\text{h}\text{e}\text{n}\text{d}}}$

$Qin{R}^{L{d}_Q},Kin{R}^{L{d}_K},Vin{R}^{L{d}_V}.$

$($Of course, it is often the case that $Q=K=V=xin{R}^{Ld}.)$ Let us call this model MHA$1.$

Next, consider a variant that we call MHA$2.$

$)1)H)h$

where

$W_h^{Q}in{R}^{d_Q{d}_{sttn}},W_h^{K}in{R}^{d_K{d}_{\text{a}\text{t}\text{t}\text{n}}},W_h^{V}in{R}^{d_V{d}_{\text{o}\text{u}\text{t}}}$

$Qin{R}^{L{d}_Q},Kin{R}^{L{d}_K},Vin{R}^{L{d}_V}.$

$($a$)$ Given an MHA$1$ model, decompose the rows of $W^O$ as

$W^O=\left[\begin{array}{c}{W}_1^O\\ W_2^O\\ \text{v}\text{d}\text{o}\text{t}\text{s}\\ W_H^O\end{array}\right]in{R}^{H{d}_{\text{h}\text{e}\text{s}\text{d}}{d}_{\text{o}\text{u}\text{t}}}$

such that $W_1^O,W_2^O,$dots,$W_H^{O}in{R}^{d_{\text{b}\text{e}\text{e}\text{d}}{d}_{\text{o}\text{a}\text{t}}}.$ Show that if we set the parameters of an MHA$2$

model as $W_h^{V}larr{W}_h^V{W}_h^O$ for $h=1,$dots,$H$ and keep all other parameters the same, then

the MHA$1$ and MHA$2$ models are equivalent, i$.$e$.,($MHA$1(Q,K,V)=$MHA$2(Q,K,V)$

for all inputs $Q,K,V.$

$($b$)$ How many trainable parameters do MHA$1$ and MHA$2$ have?

$($c$)$ If $d_V=d_{\text{o}\text{u}\text{t}}=512$ and $d_{\text{b}\text{e}\text{a}\text{d}}=64,$ what is the difference in the number of trainable

parameters?