He and Hofmann [1]detailed a skipless transformer without the V and P (post-attention projection) linear layers, which reduces the total number of weights. However, this scheme is only applicable to MHA (multi-head attention) [ 2], but not for MQA (multi-query attention) [ 3] and GQA (grouped- query attention) [ 4]. The latter schemes are used by many popular LLMs such as Llama 2, Mistral, Mixtral, PaLM, and Gemma [ 5,6,7,8,9]. Therefore, this micro-paper [ 10] proposes mathematically equivalent versions that are suitable for MQA and GQA. For example, removing Q and P from a skipless version of Mistral-7B would remove 15% of its weights (and thus reduce its compute and memory complexity). See [11, 12] for code and more transformer tricks. 1 Vanilla transformer without skip connections He et al. [13] have shown how transformers without skip connections and normalization (see Figure 1(a)) can be trained successfully. Qi Ki Vi Attention (a)transformer block i e dde Oi-1 Pi Mi Oi ddffblock i-1 FFN Attention (b)e dde Oi-1* Mi* dff Vi* Ki* Activation ActivationOi* Attention (c)d ddd Oi-1* Mi* dff Vi* Qi* ActivationOi* Attention (d)d ddd Oi-1* Mi* dff Ki* ActivationOi* Qi* merge Figure 1: (a) Skipless vanilla transformer; equivalent versions with (b) Q and P merged into the FFN (feedforward network); (c) K and P merged into FFN; (d) V and P merged into FFN. M∗ i,Q∗ i,K∗ i,V∗ i,O∗ i−1are defined in table 1. Removing skip connections and normalization allows us to merge linear layers in a mathematically identical way as shown in Figures 1(b) to (d). This reduces the number of weights without changing the functionality as follows: •Figure 1(b) is mathematically identical to Figure 1(a) and eliminates 2d2weights per transformer block by merging PiintoM∗ iandQiintoO∗ i−1. •For MHA where e=d, Figures 1(c) and (d) are mathematically identical to Figure 1(a) and eliminate 2d2 weights per transformer block by merging PiintoM∗ iandKiorViintoO∗ i−1. •This requires that Qi,Ki,Viare invertible (i.e. nonsingular). It is extremely rare that a square matrix with random values is not invertible [14] (which requires its determinant to be exactly 0). Figure 1 uses the following dimensions and weight matrices, based on the type of attention: •d: embedding dimension •e:e=dfor MHA. For MQA, e=d/n heads . And for GQA, e=d·nkv_heads/nheads . •f: hidden dimension of the FFN. f= 4din the vanilla transformer; Shazeer [3]usesf > 4d. For models that use a GLU variant [ 15] (such as Llama and Mistral), the effective f′for the first linear layer M is f′= 2f, because the GLU variant uses two linear layers that are combined (via pointwise multiplication) with a non-linear activation function. •Qi,Ki,Vi,Pi: The weight matrices of the linear layers for query, keys, values, and the post-attention projection of transformer block i. •Mi,Oi: The weight matrices of the FFN input and output linear layers. PiMiddf Pi Midf Oi-1Qi Ki Vidd e or d e or dux y zOi-1QiQi-1 Ki Qi-1 Viux y zdd e e Oi-1KiKi-1 Qi Ki-1 Viux y zdd d d Oi-1ViVi-1 Qi Vi-1 Kiux y zdd d d(a) (b) (c) (d) Figure 2: (a) Merging P and M; (b) eliminating Q; (c) eliminating K; (d) eliminating V . Figure 2 details how the linear layers are merged: •Figure 2(a) shows how the two linear layers with weight matrices PiandMiare collapsed and replaced by a single linear layer with weight matrix M∗ i=PiMi, which eliminates d2weights. •Figure 2(b) illustrates how to merge Qiinto the preceding Oi−1-matrix, which eliminates d2weights and requires Qito be invertible. Note that ⃗ y=⃗ uOi−1(QiQ−1 i)Ki=⃗ uOi−1Kiand⃗ z=⃗ uOi−1(QiQ−1 i)Vi= ⃗ uOi−1Vi. 2 of 6 •For MHA where e=d,Kican be removed as shown in Figure 2(c), which eliminates d2weights. Note that ⃗ x=⃗ uOi−1(KiK−1 i)Qi=⃗ uOi−1Qiand⃗ z=⃗ uOi−1(KiK−1 i)Vi=⃗ uOi−1Vi. This requires that Kiis invertible. •For MHA where e=d,Vican be removed as shown in Figure 2(d), which eliminates d2weights. Note that ⃗ x=⃗ uOi−1(ViV−1 i)Qi=⃗ uOi−1Qiand⃗ y=⃗ uOi−1(ViV−1 i)Ki=⃗ uOi−1Ki. This requires that Viis invertible. Table 1 specifies how the new weight matrices ( M∗ i,Q∗ i,K∗ i,V∗ i,O∗ i−1) of Figure 1 are calculated from the original ones. For the first transformer block ( i= 1), we use the input embedding instead of Oi−1(because there is no Oi−1 fori= 1). Figure 1(b) Figure 1(c) Figure 1(d) O∗ i−1Oi−1Qi Oi−1Ki Oi−1Vi Q∗ i 1 (eliminated) K−1 iQi V−1 iQi K∗ i Q−1 iKi 1 (eliminated) V−1 iKi V∗ i Q−1 iVi K−1 iVi 1 (eliminated) M∗ i PiMi Table 1: How to calculate the new weight matrices from the original ones for Figure 1. 2 Parallel transformer without skip connections Similar to the parallel transformer [ 16], Figure 3 shows parallel versions of Figures 1(b) to (d). Here, “parallel” refers to having the attention (including its linear layers) in parallel to the FFN. Attention (a) e dde Mi* d ff Vi* Ki* ActivationOi* + d Attention ddd Mi* d ff Vi* Qi* ActivationOi* + d Attention ddd Mi* d ff Ki* ActivationOi* + d (b) d (c) dQi* Figure 3: Parallel skipless transformers (a) without Q and P; (b) without K and P; (c) without V and P. Figures 3(b) and (c) require that e=d, so they are only suitable for MHA, but not for MQA and GQA. Figure 3(a) is suitable for MHA, MQA, and GQA. Figure 3(c) is identical to the simplified transformer proposed in [1]. 3 Examples The table below lists the configurations and weight counts for Pythia-6.9B and Mistral-7B. For a skipless version of Mistral-7B we would save 15% of weights after merging the Q and P linear layers into the FFN layers. For a batch 1 system that is limited by memory bandwidth, these 15% weight savings can speed up inference by 1.17x during the autoregressive next-token-generation phase, see the table below. 3 of 6 Parameter Pythia-6.9B Mistral-7B Notes Parallel attention/FFN? parallel serial [16] MHA, MQA, or GQA? MHA GQA [2, 3, 4] dim(akad) 4,096 embedding dimension n_layers 32 number of layers n_heads 32 number of heads n_kv_heads 32 8 number of KV-heads e(output dim. of K, V) 4,096 1,024 e = d * n_kv_heads / n_heads FFN type MLP MLP with SwiGLU [15] FFN hidden_dim 16,384 14,336 FFN hidden dimension vocab_size 50,400 32,000 vocabulary size Number of weights (calculated from above parameters): Q+P weights per layer 33,554,432 2 * dim * dim K+V weights per layer 33,554,432 8,388,608 2 * dim * dim / n_heads * n_kv_heads FFN weights per layer 134,217,728 176,160,768 (2 or 3) * dim * hidden_dim Input+output embed. 412,876,800 262,144,000 2 * dim * vocab_size Total weights: 6.9B 7.2B Weight savings and speedup after removing Q and P: Total w/o Q+P weights: 5.8B 6.2B total after removing Q and P Weight savings: 16% 15% Possible speedup: 1.19x 1.17x assumes batch size 1 4 Experiments Refer to [ 11] for Python code that demonstrates the numerical equivalency of the weight reduction illustrated in Figures 1(b) and 2(b). The code also confirms that all square matrices of Mistral-7B are invertible. 5 Future work Because skipless transformers are not very popular right now, future work should investigate whether removing P and Q (or K or V) is also beneficial for transformers with normalization and skip connections as illustrated in Figure 4. Adding normalization and skip connections again could simplify and speed up training relative to skipless transformers. Acknowledgements We would like to thank Bobby He (ETH Zürich) and James Martens (DeepMind) for helpful discussions on this work.