Spent some time reading this over the weekends and honestly I wish it existed a few years ago.

every AI tutorial we watch says : “just import torch and don’t worry about it, it’s magic!” this paper: “here is the exact linear algebra of what a transformer is actually doing.”

if you’re sick of reading vague tech blogs and just want to see the actual math behind attention mechanisms, tokenization, and KV caching, someone finally wrote a proper guide.

it basically strips away all the software engineering fluff and translates LLMs into pure math. you can finally stop treating these models like black boxes and actually understand the math making them tick.

genuinely a must-read if you want to know what’s really going on under the hood:

https://arxiv.org/pdf/2604.00965

Understanding Transformers and Attention Mechanisms: An Introduction for Applied Mathematicians

Source: https://arxiv.org/html/2604.00965 Michel Fabrice Serret

Abstract

This document provides a brief introduction to the attention mechanism used in modern language models based on the Transformer architecture. We first illustrate how text is encoded as vectors and how the attention mechanism processes these vectors to encode semantic information. We then describe Multi-Headed Attention, examine how the Transformer architecture is built and look at some of its variants. Finally, we provide a glimpse at modern methods to reduce the computational and memory cost of attention, namely KV caching, Grouped Query attention and Latent Attention. This material is aimed at the applied mathematics community and was written as introductory presentation in the context of the IPAM Research Collaboration Workshop entitled“Randomized Numerical Linear Algebra” (RNLA), for the project:Randomization in Transformer models.

1Introduction

Figure 1:Example of character tokenization of at most two characters.Modern Natural Language Processing (NLP) methods based on attention mechanisms process textual information not in the form of strings of characters, but as sequences of vectors. To create these sequences, the text is divided into successive sub-strings called tokens, an example of which is given in Figure1. The choice of the complete set of tokens used by the model, known as the vocabulary, is important. Indeed, it needs to allow the encoding of enough of the semantic information contained in a text to achieve, given a model, the natural language processing tasks this model is designed to address. For example, the task consisting of predicting the next token in the text might require a finer tokenization than one classifying the complete sentence into two categories. However, it should also be as small as possible to avoid diluting the information contained in the sentence and to avoid unnecessarily increasing memory and computational costs. The tokenization procedures vary from language to language and are often based on the structure of the language itself. Ideally, it would encode the smallest set of strings with inherent meaning, i.e. what would be defined as the set of all morphemes in linguistics. Second, once the text has been decomposed into a sequence of tokens, a vectorization procedure is required. In what follows, tokens can be taken without loss of generality to be words.

Now, given an ordered setT=(Ti)1≤i≤NTT=(T_{i})_{1\leq i\leq N_{T}}of theNTN_{T}possible distinct tokens, which we call the vocabulary, we can represent our text as a sequence of elements ofTTthrough their index. This sequence of indices associated to a text is known as the tokenization.

Given a vocabulary, an embedding step allows to represent it in a lower dimensional space, through a trained linear application. This is also referred to as vectorization, since each token is thus represented by a vector. Generally this embedding is given as a matrixE∈ℝNT×dE\in\mathbb{R}^{N_{T}\times d}such that, given an embedding dimensiond∈ℕ∗d\in\mathbb{N}^{*}, thedd-dimensional vectorization of theii-th tokenTiT_{i}in the vocabulary is given by theii-th row vectorEiE_{i}ofEE, as is illustrated in Figure2. These embedding matrices can either be pretrained or obtained through special methods and fine-tuned for the use cases of the language model or, as in most modern Large Language Models (LLMs), the embedding matrix can be trained from scratch as model weights. To get an idea of the size of matrices at play, the vocabulary size of Llama 3 70b is∼128\sim 128k with an embedding dimension81928192while for Gemma 3 27b it is∼262\sim 262k of embedding dimension53765376. Note that these are relatively small open source models.

Figure 2:Example of sentence embedding for a word-based tokenizer for the phrase “the lazy dog”.Additionally, positional or feature embeddings (such as an embedding of the sentence to which the token belongs) can be applied to the sequence of token embedding vectors in order to enrich the semantic information contained therein. In some models, such as DeBERTa[He+21]or DeepSeek V2[Dee+24], the positional encoding and token embedding are treated separately, either through concatenation or by design.

2Attention mechanism

Figure 3:Example of an attention layer for the embeddings obtained from the tokenization example in Figure2. Note that in this example we take the kernel to be linear for simplicity, i.e. for anyx∈ℝx\in\mathbb{R}, we setfκ​(x)=xf_{\kappa}(x)=x.Attention mechanisms, the building blocks of the Transformer architecture[Vas+17], allow encoding of semantic information between tokens through a database-like structure[Zha+24]. Indeed, a database can be seen as a set of (key, value) tuples,(k,v)(k,v), such that when we submit a queryqqto the database, the database returns the valuevvassociated to the tuple whose key is equal to the query,q=kq=k. Analogously, letX∈ℝNK​V×dinX\in\mathbb{R}^{N_{KV}\times d_{\text{in}}}be the set ofNK​VN_{KV}embedded tokens of dimensiondind_{\text{in}}such as those obtained in Figure2. As shown in Figure3, given an input vectorx∈Xx\in X, associated to one of the tokens, the output is given as a linear combination of the “value” embeddings in the database, associated with each tokenx′∈Xx^{\prime}\in X, weighted by a similarity metric relating the query embeddingqqofxxto the “key” embeddingk′k^{\prime}ofx′x^{\prime}. As we shall see in the next section, in modern language models, attention mechanisms are used as modules, known in the literature as layers, and sequentially interweaved with other types of layers in order to construct the language model. Let us now describe more formally the basic mechanism of a single “layer” of attention.

2.1General attention mechanism

Consider the given matrixXQ∈ℝNQ×dinQX^{Q}\in\mathbb{R}^{N_{Q}\times d_{\text{in}}^{Q}}ofNQ∈ℕ∗N_{Q}\in\mathbb{N}^{*}input query vectors, and the matricesXK∈ℝNK​V×dinKX^{K}\in\mathbb{R}^{N_{KV}\times d_{\text{in}}^{K}}andXV∈ℝNK​V×dinVX^{V}\in\mathbb{R}^{N_{KV}\times d_{\text{in}}^{V}}ofNK​V∈ℕ∗N_{KV}\in\mathbb{N}^{*}input key and value vectors respectively, withdinQ,dinK,dinV∈ℕ∗d_{\text{in}}^{Q},d_{\text{in}}^{K},d_{\text{in}}^{V}\in\mathbb{N}^{*}the respective dimensions of the vectors. In the following, we take for simplicity, as the query, key and value inputs come from the same embeddings in most cases,dinQ=dinK=dinV=dind_{\text{in}}^{Q}=d_{\text{in}}^{K}=d_{\text{in}}^{V}=d_{\text{in}}.

Example:In the context of a conversation with chatGPT, theNQN_{Q}query vectors, given by the rows ofQQ, would correspond to theNQN_{Q}tokens in the last question asked to the model while theNK​VN_{KV}key, encoded in the rows ofKK, and value vectors, encoded in the rows ofVV, would correspond to every token of the complete conversation. While in general the key and value vectors can be embedded from different sources for the same token, in the context of a chat, they would correspond to the same input vector, the embeddingXXof the associated token as in the Figure2.

Furthermore, note that in general, the key and value inputs will be shared as they generally correspond to two embeddings of the same “object”, while the query key can correspond to another object. As we shall see in the cross-attention module of the decoder of the original Transformer, the key and value vectors correspond to the tokens of a sentence in one language, while the query vectors correspond to the predicted tokens of the translation of the sentence in another language. For GPT models, such as chatGPT, the attention mechanism’s inputs all come from the same source. The only difference in the query and key/value vectors arises from the fact that when a new question is asked to the model, only the new query vectors from the question need to be calculated if the key/value vectors from the discussion have been stored; this is known as streaming attention[Han+25].

The first step of the attention mechanism is to construct embeddings by applying a linear transformation to the query, key, and value input vectors. We denote the associated model weights asWQ∈ℝdinQ×dQ​KW^{Q}\in\mathbb{R}^{d_{\text{in}}^{Q}\times d_{QK}},WK∈ℝdinK×dQ​KW^{K}\in\mathbb{R}^{d_{\text{in}}^{K}\times d_{QK}}andWV∈ℝdinV×doutW^{V}\in\mathbb{R}^{d_{\text{in}}^{V}\times d_{\text{out}}}wheredQ​Kd_{QK}anddVd_{V}are the shared dimensions for the query and key embeddings and the value embedding dimension respectively, while these dimensions can be different in theory, in practice they are often equal to one another. We denote the associated query, key and value matrices as

Q=XQ​WQ∈ℝNQ×dQ​K,K=XK​WK∈ℝNK​V×dQ​K,V=XV​WV∈ℝNK​V×dV.Q=X^{Q}W^{Q}\in\mathbb{R}^{N_{Q}\times d_{QK}},\quad K=X^{K}W^{K}\in\mathbb{R}^{N_{KV}\times d_{QK}},\quad V=X^{V}W^{V}\in\mathbb{R}^{N_{KV}\times d_{V}}.The similarity metric is usually given in the form of a kernelκ:ℝdQ​K×ℝdQ​K→ℝ\kappa:\mathbb{R}^{d_{QK}}\times\mathbb{R}^{d_{QK}}\to\mathbb{R}which allows to obtain theattention scoresfor each tuple(i,j)∈⟦1,NQ⟧×⟦1,NK​V⟧(i,j)\in\llbracket 1,N_{Q}\rrbracket\times\llbracket 1,N_{KV}\rrbracketof theNQN_{Q}query andNK​VN_{KV}key/value indices,

A=(κ​(Q​[i,:],K​[j,:]))1≤i≤NQ1≤j≤NK​V∈ℝNQ×NK​V,A=\Big(\kappa(Q[i,:],K[j,:])\Big)_{\begin{subarray}{c}1\leq i\leq N_{Q}\\ 1\leq j\leq N_{KV}\end{subarray}}\in\mathbb{R}^{N_{Q}\times N_{KV}},where we denote byB​[i,:]B[i,:]theii-th row vector of the matrixBB. Note that the kernel will often be of the form

κ​(v,w)=fκ​(⟨v,w⟩),\kappa(v,w)=f_{\kappa}\left(\langle v,w\rangle\right),for a given functionfκ:ℝ→ℝf_{\kappa}:\mathbb{R}\to\mathbb{R}and we will use this notation throughout the document to simplify some equations. Furthermore, the element-wise application offκf_{\kappa}to a matrixMMwill be denoted asfκ​(M)f_{\kappa}(M). As we wish to emulate a probability distribution on the set of key/value inputs for a given query, a normalization step is required; the normalization coefficients are obtained as the inverse of the diagonal matrixZ∈ℝNQ×NQZ\in\mathbb{R}^{N_{Q}\times N_{Q}}, whose non-zero elements are the sum of the rows ofAA, i.e.

Z​[i,i]=∑1≤j≤NK​Vκ​(Q​[i,:],K​[j,:]),∀i∈{1,…,NQ}.Z[i,i]=\sum\limits_{1\leq j\leq N_{KV}}\kappa\left(Q[i,:],K[j,:]\right),\quad\forall i\in\{1,\dots,N_{Q}\}.Normalized attention scores, i.e., the productZ−1​AZ^{-1}A, are known in the literature asattention weights. The final output of the attention mechanism is then given by

Y=Z−1​A​V∈ℝNQ×dout.Y=Z^{-1}AV\in\mathbb{R}^{N_{Q}\times d_{\text{out}}}.Note that the output dimension of the attention mechanism is given by the value dimension, i.e.dout=dVd_{\text{out}}=d_{V}.

The most common kernel is the scaled exponential kernel used in the original Transformer[Vas+17]given, for anyv,w∈ℝdQ​Kv,w\in\mathbb{R}^{d_{QK}}, byκ​(v,w)=exp⁡(⟨v,w⟩dQ​K).\kappa(v,w)=\exp\left(\frac{\langle v,w\rangle}{\sqrt{d_{QK}}}\right).It originates from the fact that if the entries ofvvandwwcome from an i.i.d. distribution with mean zero and variance11, then⟨v,w⟩\langle v,w\ranglehas variancedQ​Kd_{QK}. Hence, the denominator rescales the scalar product so that its variance is 1. Given the matricesQQandKK, the attention scores can be obtained asA=exp⁡(Q​K𝖳dQ​K),A=\exp\left(\frac{QK^{\mathsf{T}}}{\sqrt{d_{QK}}}\right),where the exponential is applied element-wise. The attention weights in this context are equal to the softmax function,σ:ℝN→ℝN\sigma:\mathbb{R}^{N}\to\mathbb{R}^{N}, applied to the query-key scalar product, the softmax function is defined asσ:x∈ℝN↦1∑ℓ=1Nexℓ​(ex1…exN)𝖳,\sigma:x\in\mathbb{R}^{N}\mapsto\frac{1}{\sum\limits_{\ell=1}^{N}e^{x_{\ell}}}\begin{pmatrix}e^{x_{1}}&\dots&e^{x_{N}}\end{pmatrix}^{\mathsf{T}},and hence, for anyi∈{1,…,NQ}i\in\{1,\dots,N_{Q}\},(Z−1​A)​[i,:]=σ​(Q​[i,:]​K𝖳dQ​K).(Z^{-1}A)[i,:]=\sigma\left(\frac{Q[i,:]K^{\mathsf{T}}}{\sqrt{d_{QK}}}\right).

2.2Multi-headed attention

In order to allow the attention mechanism to focus on more types of semantic information, multiple parallel attention heads can be used on the same input vectors. The simplest way to do so is to consider stacked single attention heads with the constraint that they must all have the same output dimensiondheadd_{\mathrm{head}}, i.e. for anyh∈⟦1,Nheads⟧h\in\llbracket 1,N_{\mathrm{heads}}\rrbracket,WhV∈ℝdinV×dheadW^{V}_{h}\in\mathbb{R}^{d_{\text{in}}^{V}\times d_{\mathrm{head}}}. LetNheads∈ℕ∗N_{\mathrm{heads}}\in\mathbb{N}^{*}be the number of attention heads, for anyh∈⟦1,Nheads⟧h\in\llbracket 1,N_{\mathrm{heads}}\rrbracket, we denote byWhQ,WhKW^{Q}_{h},W_{h}^{K}andWhVW_{h}^{V}the query, key and value embedding matrices associated to thehh-th head. For each headh∈⟦1,Nheads⟧h\in\llbracket 1,N_{\mathrm{heads}}\rrbracket, the attention score can then be obtained as

Ah=(κ​(Qh​[i,:],Kh​[j,:]))1≤i≤NQ1≤j≤NK​V=fκ​(Qh​Kh𝖳),A_{h}=\Big(\kappa\left({Q_{h}[i,:],K_{h}[j,:]}\right)\Big)_{\begin{subarray}{c}1\leq i\leq N_{Q}\\ 1\leq j\leq N_{KV}\end{subarray}}=f_{\kappa}(Q_{h}K_{h}^{\mathsf{T}}),whereQh=XQ​WhQ∈ℝNQ×dQ​KQ_{h}=X^{Q}W^{Q}_{h}\in\mathbb{R}^{N_{Q}\times d_{QK}},Kh=XK​WhK∈ℝNK​V×dQ​KK_{h}={X^{K}}{W_{h}^{K}}\in\mathbb{R}^{N_{KV}\times d_{QK}}, andkkis applied element-wise. We note that in generaldQ​K=dheadd_{QK}=d_{\mathrm{head}}for each headhh. Similarly, we set

Vh=XVWhV,Zh=diag(∑1≤j≤NK​Vκ(Qh[i,:],Kh[j,:]))1≤i≤NQ.V_{h}=X^{V}W^{V}_{h},\quad Z_{h}=\operatorname{diag}\left(\sum\limits_{1\leq j\leq N_{KV}}\kappa\left({Q_{h}[i,:],K_{h}[j,:]}\right)\right)_{1\leq i\leq N_{Q}}.The output of each attention head can then be obtained as

Yh=Zh−1​Ah​Vh∈ℝNQ×dhead.Y_{h}=Z_{h}^{-1}A_{h}V_{h}\in\mathbb{R}^{N_{Q}\times d_{\mathrm{head}}}.Given the outputs of each head, the output of the multi-headed attention mechanism is combined into a single outputY∈ℝNQ×doutY\in\mathbb{R}^{N_{Q}\times d_{\text{out}}}through a linear applicationWO∈ℝNheads​dhead×doutW^{O}\in\mathbb{R}^{N_{\mathrm{heads}}d_{\mathrm{head}}\times d_{\text{out}}}, which is trained, as

Y=concatcol(Yh)1≤h≤NheadsWO,Y=\operatorname{concat}_{\mathrm{col}}\left(Y_{h}\right)_{1\leq h\leq N_{\mathrm{heads}}}W^{O},whereconcatcol(Yh)1≤h≤Nheads=(Y1⋯YNheads)∈ℝNQ×Nheads​dhead\operatorname{concat}_{\mathrm{col}}\left(Y_{h}\right)_{1\leq h\leq N_{\mathrm{heads}}}=\begin{pmatrix}Y_{1}&\cdots&Y_{N_{\mathrm{heads}}}\end{pmatrix}\in\mathbb{R}^{N_{Q}\times N_{\mathrm{heads}}d_{\mathrm{head}}}.

3Transformer Architecture

3.1Encoders, Decoders and the Transformer architecture

In natural language processing, a common method used for machine translation of one language to another is the use of an encoder-decoder architecture to first encode an input sentence in the first language into an intermediary state and then use a decoder to transform this intermediary state into the equivalent sentence in the second language.

Figure 4:Encoder-Decoder modelThe originalTransformermodel[Vas+17]was based on such a system. Let us now look at the individual components of the different layers which compose the Transformer architecture.

3.1.1Encoder layers

Figure 5:Encoder layer of the original Transformer architecture. Note that the⨁\bigoplussign in the diagram corresponds to the skip connections.The encoder is composed of a series ofNLN_{L}layers of encoder layers . Each encoder layer consists of the following components, assembled as shown in Figure5,

1. 1.Multi-headedself-attention sublayer: This sublayer consists in a multi-headed attention mechanism exactly as described in subsection2.2. Theselfpart in the name refers to the fact that both its query and key/value input vectors are from the same input, i.e. we haveXQ=XK=XVX^{Q}=X^{K}=X^{V}. In Figure5this corresponds to the fact that the three input arrows originate from the previous layer.
2. 2.Layer normalization sublayer: Given the inputx∈ℝdx\in\mathbb{R}^{d}to a layer, layer normalization[BKH16]constructs a vectorx~=x−μσ2+ϵ\tilde{x}=\frac{x-\mu}{\sqrt{\sigma^{2}+\epsilon}}of mean0and variance11, whereμ=1d​∑i=1dxi\mu=\frac{1}{d}\sum\limits_{i=1}^{d}x_{i}is the input vector’s mean andσ=1d​∑i=1d(xi−μ)2\sigma=\sqrt{\frac{1}{d}\sum_{i=1}^{d}(x_{i}-\mu)^{2}}is its standard deviation. It outputsy=γ⊙x~+βy=\gamma\odot\tilde{x}+\beta, whereγ,β∈ℝd\gamma,\beta\in\mathbb{R}^{d}are known as the gain and shift and correspond to trainable parameters that rescale the variance and mean of the output to a fixed value and we denote by⊙\odotthe Hadamard or term-wise product.
3. 3.Feed forward sublayer: It is the basic building block of neural networks, it outputsy=f​(W​x+b)y=f(Wx+b)whereffis a nonlinear function,W∈ℝd×d′W\in\mathbb{R}^{d\times d^{\prime}}is the weight matrix of the sublayer andb∈ℝd′b\in\mathbb{R}^{d^{\prime}}is the bias vector, whered′d^{\prime}is the dimension of the output of the layer. In general the nonlinear functionffis a variant of the Rectified Linear Unit (ReLU) family of functions; recently, parametric Gated Linear Unit (GLU) variants[Sha20]have been favored in state-of-the-art models[Tea+24,The25,Yan+25,Gra+24].
4. 4.Skip connections: Represented by the arrows that go under the components in Figure5, skip connections consist in adding to the output of a sublayer its input. They allow the preservation of the flow of information through the neural network and reduces the problem of vanishing gradients which arises in deep neural networks.

3.2Decoder layers, Causal/Masked attention and Cross-attention

Figure 6:Decoder layer of the original Transformer architecture. Notice that the main differences with the encoder are the cross-attention sublayer where the key/value vectors are constructed not from the decoder state but from the final encoder state and the masked causal self-attention restricting the query vectors from seeing the key vectors of future tokens.Decoder layers in the Transformer architecture have two differences with the encoder layers, as can be seen in Figure6.

Remark:The architecture shown in Figure5is the original Transformer architecture. An alternative[Xio+20], reducing the problem of the vanishing gradients and allowing easier training but at the cost of performance, consists in applying the layer normalization before the attention sublayer and is shown in Figure3.2. Note how the skip connections are not normalized in this architecture and thus allow the gradients to skip through unscathed. Recently, state-of-the-art large language models such as Gemma 3\citesteamGemma3Technical2025, have made use of both Pre-Layer and Post-Layer Normalization[Kim+25]. Furthermore, layer-normalization has often been replaced with the less computationally intensive RMS-Normalization procedure which consists in only rescaling the input with respect to its Root-mean-squared norm without recentering or adding a bias[ZS19]. Figure 7:Encoder layer with Pre-Layer Normalization

Masked and causal attention

The first difference is the use of masked attention. As the decoders construct the new sentence token by token, the self-attention mechanism that produces the token at stepttis not able to take into account the key-value vectors of tokens that will be created afterwards, hence, for consistency, a mask is applied to the attention to limit its vision of the tokens that come afterwards. More formally, masked attention consists in restricting the subset of key-value tokens seen by each query token, i.e. defining, for each1≤i≤NQ1\leq i\leq N_{Q}, a subset of key-value tokensSi⊂⟦1,NK​V⟧S_{i}\subset\llbracket 1,N_{KV}\rrbracketsuch that the attention matrix is given by

A=(𝟏j∈Si​κ​(Q​[i,:],K​[j,:]))1≤i≤NQ1≤j≤NK​V,A=\Big(\mathbf{1}_{j\in S_{i}}\kappa\left(Q[i,:],K[j,:]\right)\Big)_{\begin{subarray}{c}1\leq i\leq N_{Q}\\ 1\leq j\leq N_{KV}\end{subarray}},where, for any setAA,𝟏A\mathbf{1}_{A}is the associated indicator function associated toAA. In causal self-attention, the principle is that as the sentences are formed token by token in a sequential manner and the previous tokens are fixed, the addition of new tokens should not affect the previous parts of the sentence. Hence, the query should only see itself and the tokens that have been defined before itself, i.e.Si={j∈ℕ,j≤i}S_{i}=\{j\in\mathbb{N}\,,\,j\leq i\}and thus

A=(𝟏j≤i​κ​(Q​[i,:],K​[j,:]))1≤i≤NQ1≤j≤NK​V.A=\left(\mathbf{1}_{j\leq i}\kappa\left({Q[i,:],K[j,:]}\right)\right)_{\begin{subarray}{c}1\leq i\leq N_{Q}\\ 1\leq j\leq N_{KV}\end{subarray}}. Equivalently, given the contextual maskMMsuch thatMi​j={0ifj∈Si,∅κotherwise,M_{ij}=\begin{cases}0&\text{if j\\in S\_\{i\},}\\ \varnothing_{\kappa}&\text{otherwise,}\end{cases}(where∅κ=−∞\varnothing_{\kappa}=-\inftyfor the exponential kernel). In the case of softmax attention, masked terms are given the value−∞-\inftyso as to be attributed a zero probability after the exponential kernel function. Masked attention can then be seen as addingMMto theQ​K𝖳QK^{\mathsf{T}}matrix before applying the kernel function givingA=κ​(Q​K𝖳+M).A=\kappa(QK^{\mathsf{T}}+M).Equivalently, the mask can be applied after the kernel by setting the masked elements to zero. By settingM~i​j=𝟏Si​(j),\widetilde{M}_{ij}=\mathbf{1}_{S_{i}}(j),, with𝟏Si\mathbf{1}_{S_{i}}, the indicating function for the setSiS_{i}, the masked attention is obtained asA=κ​(Q​K𝖳)⊙M~.A=\kappa(QK^{\mathsf{T}})\odot\widetilde{M}.

Cross-attention

The second architectural difference in the decoder layers is that they use the encoder’s output,ENLE_{N_{L}}in Figure6, as the input of a second attention sublayer, which allows the decoder to relate its output sentence to the information contained in the original input. To do so, the intermediate state is used as the key and value inputs,XKX^{K}andXVX^{V}, of a multi-headed attention sublayer as shown in Figure6.

Remark: Cross-attention[Ala+22,Rom+22], or more generally, shared embedding spaces[KSK21,Rad+21]are at the root of modern multimodal LLMs in which embedding subspaces can be shared through attention, or directly through a scalar product in the embedding space, between texts and images.

3.3The Transformer and its variants

Figure 8:Transformer architecture.The complete Transformer architecture, as can be seen in Figure8, is modelled for, and trained on, the task of machine translation and is at the origin of most LLMs today, albeit with several modifications. Indeed, while the Transformer network is well suited for sequence to sequence translation, a change in architecture is required for less structured problems. Traditionally, these changes, which vary according to the task at hand, can be classified into two main groups, Encoder-only architectures such as BERT (Bidirectional Encoder Representation from Transformers)[Dev+19,He+21,War+24], more suited for tasks requiring the extraction of information from text such as token or sentence classification and Decoder-only architectures such as GPT(Generative Pre-Trained Transformers)[Rad+18,Bro+20,Tou+23], suited to generative tasks such as next-token prediction for text generation. The current trend, put forth in[Rad+18], lies in pretraining such models on tasks that require an understanding of the inherent structure of the text and with a huge amounts of available data, such as next-token prediction for the GPT architecture, using the architecture shown in Figure9, and masked text completion for Bert, and using the pretrained weights as an initialization for a fine-tuning a model with the same weights but a different output head for other tasks using smaller datasets.

Figure 9:GPT architecture.

4KV Caching, compression and attention optimization

One of the main memory bottlenecks of modern LLMs arises due to what is known asKV caching. In the attention mechanism, the computation of the next output for the last token depends on the key and value vectors of all tokens. Indeed, this property is evident if we look at the last output of single-headed attention layer

Y​[NQ,j]=(Z−1​A​V)​[NQ,j]=∑i=1NK​Vκ​(Q​[NQ,:],K​[i,:])∑ℓ=1NK​Vκ​(Q​[NQ,:],K​[ℓ,:])​V​[i,j],Y[N_{Q},j]=\left(Z^{-1}AV\right)[N_{Q},j]=\sum\limits_{i=1}^{N_{KV}}\frac{\kappa\left(Q[N_{Q},:],K[i,:]\right)}{\sum\limits_{\ell=1}^{N_{KV}}\kappa\left(Q[N_{Q},:],K[\ell,:]\right)}V[i,j],where theNQN_{Q}index corresponds to the last query token. Note that the computational cost of the attention mechanism is of the order ofO​(NQ​NK​V​(dQ​K+dV))O(N_{Q}N_{KV}(d_{QK}+d_{V}))and thatdQ​Kd_{QK}anddVd_{V}are generally small with respect toNQN_{Q}andNK​VN_{KV}. In the context of a chat for example, the token embeddings for the query, key and value terms are calculated in a stream, first you have the first sentence sent, then the answer, second sentence, etc… As the conversation goes on, the query, key and value vectors can either be recalculated at each new sentence, but with a computational cost ofO​(Ntokens3​d)O(N_{\text{tokens}}^{3}d), whereNK​V=NQ=NtokensN_{KV}=N_{Q}=N_{\text{tokens}}as all the vectors need to be recomputed each time a new token appears. Alternatively, the key and value vectors can be stored after each call to be used again for the next step, as we have seen above, once the key and value vectors are calculated, the query terms can be discarded as their only impact on subsequent predictions is through the KV vectors. KV caching consists in keeping these key and value vectors for each token in memory. This allows to obtain a linear cost in the number of tokens for each new token added( thus overall anO​(Ntokens2​d)O(N_{\text{tokens}}^{2}d)complexity in total but anO​(NQ​Ntokens​d)O(N_{Q}N_{\text{tokens}}d)cost for each new sentence of sizeNQN_{Q}). This does not come for free as the vectors need to be stored. As each we have to cache the KV vectors of each attention head for each layer the memory cost of such caching is2​NL​Nh​NK​V​d2N_{L}N_{h}N_{KV}dtimes the number of bits used per float. This can quickly become prohibitive, especially in the long-context case.

We callstreaming attentionthe calculation of attention where the query input tokens are a subset of the key-value input tokens, as opposed tofull attention. Note that the key-value vectors need to be cached for this to work as the previous query tokens are required for their construction.

4.1Grouped Query Attention

The size of the cache in memory is a bottleneck both in terms of storage and of communication overheads, methods have thus been developed to palliate to this. One of the simplest, and most ubiquitous, such method consists in sharing the same key-value heads for multiple query heads, hence reducing the size of the vector that needs to be kept in memory. This is known as Grouped Query Attention (GQA), and in the limit of a single key-value head for all query heads it is sometimes referred to as Multi-Query Attention (MQA).

One recent method which generalizes the use of lower-dimensional embedding spaces is the Multi-Headed Latent Attention mechanism developed by DeepSeek[Dee+24]. Let us now look at how this fits into our framework.

4.2Latent Attention

The main idea behind Latent Attention is to construct a shared low-rank latent embedding from which the key and value vectors are formed for each head through a linear application. In the following, we setXK=XV=X∈ℝNK​V×dinX^{K}=X^{V}=X\in\mathbb{R}^{N_{KV}\times d_{\text{in}}}anddQ​K=dV=dheadd_{QK}=d_{V}=d_{\mathrm{head}}. Before diving into the formalism of Multi-headed Latent Attention(MLA), let us look back at Multi-Headed Attention and notice that the per-head embedding calculations can be seen as a single matrix multiplication,

𝐐=(Q1⋯QNheads)=XQ​(W1Q⋯WNheadsQ)≔XQ​W𝐐∈ℝNQ×Nheads​dhead𝐊=(K1⋯KNheads)=X​(W1K⋯WNheadsK)≔X​W𝐊∈ℝNK​V×Nheads​dhead𝐕=(V1⋯VNheads)=X​(W1V⋯WNheadsV)≔X​W𝐕∈ℝNK​V×Nheads​dhead,\begin{aligned} \mathbf{Q}&=\begin{pmatrix}Q_{1}&\cdots&Q_{N_{\mathrm{heads}}}\end{pmatrix}&=X_{Q}\begin{pmatrix}W^{Q}_{1}&\cdots&W^{Q}_{N_{\mathrm{heads}}}\end{pmatrix}&\coloneq X_{Q}W^{\mathbf{Q}}\in\mathbb{R}^{N_{Q}\times N_{\mathrm{heads}}d_{\mathrm{head}}}\\ \mathbf{K}&=\begin{pmatrix}K_{1}&\cdots&K_{N_{\mathrm{heads}}}\end{pmatrix}&=X\begin{pmatrix}W^{K}_{1}&\cdots&W^{K}_{N_{\mathrm{heads}}}\end{pmatrix}&\coloneq XW^{\mathbf{K}}\in\mathbb{R}^{N_{KV}\times N_{\mathrm{heads}}d_{\mathrm{head}}}\\ \mathbf{V}&=\begin{pmatrix}V_{1}&\cdots&V_{N_{\mathrm{heads}}}\end{pmatrix}&=X\begin{pmatrix}W^{V}_{1}&\cdots&W^{V}_{N_{\mathrm{heads}}}\end{pmatrix}&\coloneq XW^{\mathbf{V}}\in\mathbb{R}^{N_{KV}\times N_{\mathrm{heads}}d_{\mathrm{head}}}\end{aligned},where𝐐\mathbf{Q},𝐊\mathbf{K}, and𝐕\mathbf{V}denote the column-wise concatenations of the head-specific query, key, and value matrices. Similarly, we set

W𝐐=(W1Q​⋯WNheadsQ)∈ℝdi​n×Nheads​dhead,W𝐊=(W1K⋯WNheadsK)∈ℝdi​n×Nheads​dheadW^{\mathbf{Q}}=\begin{pmatrix}W^{Q}_{1}\cdots&W^{Q}_{N_{\mathrm{heads}}}\end{pmatrix}\in\mathbb{R}^{d_{in}\times N_{\mathrm{heads}}d_{\mathrm{head}}},\quad W^{\mathbf{K}}=\begin{pmatrix}W^{K}_{1}&\cdots&W^{K}_{N_{\mathrm{heads}}}\end{pmatrix}\in\mathbb{R}^{d_{in}\times N_{\mathrm{heads}}d_{\mathrm{head}}}and​W𝐕=(W1V⋯WNheadsV)∈ℝdi​n×Nheads​dhead.\text{ and }W^{\mathbf{V}}=\begin{pmatrix}W^{V}_{1}&\cdots&W^{V}_{N_{\mathrm{heads}}}\end{pmatrix}\in\mathbb{R}^{d_{in}\times N_{\mathrm{heads}}d_{\mathrm{head}}}.to denote the column-wise concatenations of the corresponding weight matrices.

Using this formulation, the latent attention method uses two latent subspaces, of dimensionsdLQd_{L_{Q}}anddL≤di​nd_{L}\leq d_{in}, to construct a low-rank factorization of the query embeddings and a shared low-rank factorization of the key and value embeddings:

W𝐐\displaystyle W^{\mathbf{Q}}=WLQ​W𝐋𝐐​𝐐=WLQ​(W1LQ​Q⋯WNheadsLQ​Q),\displaystyle=W^{L_{Q}}W^{\mathbf{L_{Q}Q}}=W^{L_{Q}}\begin{pmatrix}W^{{L_{Q}}Q}_{1}&\cdots&W^{{L_{Q}Q}}_{N_{\mathrm{heads}}}\end{pmatrix},W𝐊\displaystyle W^{\mathbf{K}}=WL​W𝐋𝐊=WL​(W1L​K⋯WNheadsL​K),\displaystyle=W^{L}W^{\mathbf{LK}}=W^{L}\begin{pmatrix}W^{{LK}}_{1}&\cdots&W^{{LK}}_{N_{\mathrm{heads}}}\end{pmatrix},W𝐕\displaystyle W^{\mathbf{V}}=WL​W𝐋𝐕=WL​(W1L​V⋯WNheadsL​V),\displaystyle=W^{L}W^{\mathbf{LV}}=W^{L}\begin{pmatrix}W^{{LV}}_{1}&\cdots&W^{{LV}}_{N_{\mathrm{heads}}}\end{pmatrix},whereWL∈ℝdin×dLW^{L}\in\mathbb{R}^{d_{\text{in}}\times d_{L}}is the shared latent subspace weight matrix for the key and value embeddings, and we define for any1≤h≤Nheads1\leq h\leq N_{\mathrm{heads}}, the matricesWhL​K∈ℝdL×dheadW^{LK}_{h}\in\mathbb{R}^{d_{L}\times d_{\mathrm{head}}}, andWhL​V∈ℝdL×dheadW_{h}^{LV}\in\mathbb{R}^{d_{L}\times d_{\mathrm{head}}}to be the matrices that send the latent vectors onto the key and value embeddings of thehh-th head. The matricesW𝐋𝐊∈ℝdL×Nheads​dheadW^{\mathbf{LK}}\in\mathbb{R}^{d_{L}\times N_{\mathrm{heads}}d_{\mathrm{head}}}andW𝐋𝐕∈ℝdL×Nheads​dheadW^{\mathbf{LV}}\in\mathbb{R}^{d_{L}\times N_{\mathrm{heads}}d_{\mathrm{head}}}are the the column-wise concatenation of(WhL​K)1≤h≤Nheads\Big(W^{LK}_{h}\Big)_{1\leq h\leq N_{\mathrm{heads}}}, and respectively(WhL​V)1≤h≤Nheads\Big(W^{LV}_{h}\Big)_{1\leq h\leq N_{\mathrm{heads}}}. Analogously,WLQ∈ℝdin×dLQW^{L_{Q}}\in\mathbb{R}^{d_{\text{in}}\times d_{L_{Q}}}is latent subspace weight matrix for the query embeddings andW𝐋𝐐​𝐐W^{\mathbf{L_{Q}Q}}is the column-wise concatenation of(WhLQ​Q)1≤h≤Nheads\Big(W^{L_{Q}Q}_{h}\Big)_{1\leq h\leq N_{\mathrm{heads}}}, the matrices that send the latent query embeddings back to the shared query-key subspace.

In practice, in the DeepSeek V2 model[Dee+24], the latent model is trained directly and is not created from the key and value matrices using the above factorizations.

This idea allows to keep a single cache vector per token,shared between all heads; the latent embeddingL=X​WLL=XW^{L}, instead ofKKandVV. Furthermore, the latent formulation allows to merge weight matrices by rewriting the attention mechanism. First, for each head1≤h≤Nheads1\leq h\leq N_{\mathrm{heads}}, the Latent-to-Query weightsWhLQ​QW_{h}^{L_{Q}Q}and the Latent-to-Key weightsWhL​KW_{h}^{LK}can be merged, reducing the computational cost asdLd_{L}anddLQ<di​nd_{L_{Q}}<d_{in}. Indeed, we can rewrite, using (4.2),

Qh​Kh𝖳=LQ​WhLQ​Q​WhL​K𝖳​WL𝖳​X𝖳=LQ​WhLQ​Q​WhL​K𝖳⏟WhL​Q​K∈ℝdLQ×dL​L𝖳,Q_{h}K_{h}^{\mathsf{T}}=L_{Q}W^{L_{Q}Q}_{h}{W^{LK}_{h}}^{\mathsf{T}}{W^{L}}^{\mathsf{T}}X^{\mathsf{T}}=L_{Q}\underbrace{W^{L_{Q}Q}_{h}{W^{LK}_{h}}^{\mathsf{T}}}_{{W^{LQK}_{h}}\in\mathbb{R}^{d_{L_{Q}}\times d_{L}}}L^{\mathsf{T}},which gives us, by indexing on each headhh, the merged query tensor(WhL​Q​K)1≤h≤Nheads∈ℝdLQ×dL×Nheads\left({W^{LQK}_{h}}\right)_{1\leq h\leq N_{\mathrm{heads}}}\in\mathbb{R}^{d_{L_{Q}}\times d_{L}\times N_{\mathrm{heads}}}.

Second, the Latent-to-Value weightsWL​VW^{LV}can be merged with the head combining output weightsWOW^{O}. Indeed, we can rewrite the output of the attention mechanism as follows

Y=concatcol(Zh−1AhVh)1≤h≤NheadsWO=concatcol(Zh−1AhL)1≤h≤Nheads⏟Y′(W1L​V0⋱0WNheadsL​V)​WO⏟WL​O,Y=\operatorname{concat}_{\mathrm{col}}\left(Z_{h}^{-1}A_{h}V_{h}\right)_{1\leq h\leq N_{\mathrm{heads}}}W^{O}=\underbrace{\operatorname{concat}_{\mathrm{col}}\left(Z_{h}^{-1}A_{h}L\right)_{1\leq h\leq N_{\mathrm{heads}}}}_{Y^{\prime}}\underbrace{\begin{pmatrix}W^{LV}_{1}&&0\\ &\ddots&\\ 0&&W^{LV}_{N_{\mathrm{heads}}}\end{pmatrix}W^{O}}_{{W^{LO}}},where we highlight again the fact thatLLdoes not depend onhhas it is shared for all heads. This gives us the new latent output matrixWL​O∈ℝNheads​dL×do​u​tW^{LO}\in\mathbb{R}^{N_{\mathrm{heads}}d_{L}\times d_{out}}.

The latent attention mechanism can thus be simplified to

Y=concatcol(Zh−1AhL)1≤h≤NheadsWL​O,Y=\operatorname{concat}_{\mathrm{col}}\left(Z_{h}^{-1}A_{h}L\right)_{1\leq h\leq N_{\mathrm{heads}}}{W^{LO}},withAh=exp⁡(LQ​WhL​Q​K​L𝖳)A_{h}=\exp(L_{Q}W_{h}^{LQK}L^{\mathsf{T}})andZh=diag(∑j=1NK​Vexp((LQWhL​Q​KL𝖳)[i,j]))1≤i≤NQZ_{h}=\operatorname{diag}\left(\sum\limits_{j=1}^{N_{KV}}\exp\Big((L_{Q}W^{LQK}_{h}L^{\mathsf{T}})[i,j]\Big)\right)_{1\leq i\leq N_{Q}}. With this formulation, note that the only model weights that need to be kept areWL​OW^{LO},WLW^{L},WLQW^{L_{Q}}andWL​Q​K{W^{LQK}}and only the latent vectors need to be kept in memory.

In the absence of positional embeddings, latent attention models can be represented exactly in the form of an equivalent GQA/MHA model thanks to the low-rank factorizations shown above. However, in practice, when positional embeddings are applied, this equivalence is lost. This loss of equivalence is due to the fact that positional embeddings, such as the Rotary Position Embeddings[Su+23](RoPE), in classical attention are applied after constructingKK. For RoPE, in latent attention, this would amount to applying a position-dependent matrixRm∈ℝdhead×dheadR_{m}\in\mathbb{R}^{d_{\mathrm{head}}\times d_{\mathrm{head}}}, as followsQh​[i,:]​Ri​Rj𝖳​Kh​[j,:]𝖳=(LQ​[i,:])​WhLQ​Q​Ri​Rj𝖳​WhL​K𝖳​(L​[j,:])𝖳.Q_{h}[i,:]R_{i}R_{j}^{\mathsf{T}}K_{h}[j,:]^{\mathsf{T}}=(L_{Q}[i,:])W_{h}^{L_{Q}Q}R_{i}R_{j}^{\mathsf{T}}{W_{h}^{LK}}^{\mathsf{T}}(L[j,:])^{\mathsf{T}}.This would not allow the merging of theWhL​KW^{LK}_{h}andWhLQ​QW^{L_{Q}Q}_{h}matrices and would require recomputing the application of the positional embedding at each evaluation, not only preventing a speedup but incurring an additional computational overhead. The latent attention RoPE variant instead consist in appending to the key and query vectors “non-latent” part on which the positional embedding is applied, breaking the equivalence but keeping the computational advantage of the latent model. Recently, methods to approximately convert pretrained GQA/MHA models of attention, from previously learnt models, to the Latent attention formalism have been developed[MYZ25]and allow for the use of the DeepSeek’s optimizations on other models.

The tables2and2summarize the matrices and caches that need to be stored in memory for the MHA and MLA cases respectively. Example of model sizes are given in table3.

Table 1:Dimensions of tensors and matrices stored in multi-headed attention. The total number of floats to be stored isNheads​NK​V⋅(dQ​K+dhead)+2​Nheads​din​dQ​K+Nheads​dhead​(din+dout)N_{\mathrm{heads}}\ N_{KV}\cdot\left(d_{QK}+d_{\mathrm{head}}\right)+2N_{\mathrm{heads}}d_{\text{in}}d_{QK}+N_{\mathrm{heads}}d_{\mathrm{head}}\left(d_{\text{in}}+d_{\text{out}}\right). For GQA models, theNheadsN_{\mathrm{heads}}in the dimensions needs to be replaced with the number of KV heads forWKW^{K},WVW^{V}and the cache terms. Table 2:Dimensions of tensors and matrices stored in multi-headed latent attention. The total number of floats to be stored is(NK​V⋅dL)+(din⋅dL)+(Nheads⋅dL⋅(din+dout))(N_{KV}\cdot d_{L})+(d_{\text{in}}\cdot d_{L})+(N_{\mathrm{heads}}\cdot d_{L}\cdot(d_{\text{in}}+d_{\text{out}})).

ParameterLlama 3 70BGemma 3 27BDeepseek V2Number of Layers806260Number of HeadsNheadsN_{\mathrm{heads}}6432128Number of Key/Value Heads816-Hidden Dimensionsdin=doutd_{\text{in}}=d_{\text{out}}8,1925,3765120KV Dimensiondhead=dQ​Kd_{\mathrm{head}}=d_{QK}128128512(dLd_{L})Table 3:Comparison of Llama 3 70B[Gra+24]and Gemma 3 27B[The25]and DeepSeek V2 Model[Dee+24]Specifications. Note that DeepSeekV2 is a mixture-of-experts model with latent attention while Gemma 3 and Llama 3 are Grouped Query attention models. These models use a Grouped Query attention model wherein multiple attention mechanisms are used in parallel, these are known as attention heads. To reduce the computational overhead, the key and value matrices are shared between heads, the number of key/value matrices is denoted as the number of KV heads above.

5Acknowledgments

The author would like to thank Alice Cortinovis and Laura Grigori for their valuable feedback on this document.

References

• [Ala+22]Jean-Baptiste Alayrac et al.“Flamingo: A Visual Language Model for Few-Shot Learning”InarXiv preprint, 2022DOI:10.48550/arXiv.2204.14198
• [BKH16]Jimmy Ba, J. Kiros and Geoffrey E. Hinton“Layer Normalization”InarXiv preprintabs/1607.06450, 2016
• [Bro+20]Tom B. Brown et al.“Language Models Are Few-Shot Learners”InarXiv preprint, 2020DOI:10.48550/arXiv.2005.14165
• [Dee+24]undef DeepSeek-AI et al.“DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model”InarXiv preprint, 2024DOI:10.48550/arXiv.2405.04434
• [Dev+19]Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova“BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”InarXiv preprint, 2019arXiv:1810.04805
• [Gra+24]Aaron Grattafiori et al.“The Llama 3 Herd of Models”InarXiv preprint, 2024DOI:10.48550/arXiv.2407.21783
• [Han+25]Insu Han et al.“Streaming Attention Approximation via Discrepancy Theory”, 2025DOI:10.48550/arXiv.2502.07861
• [He+21]Pengcheng He, Xiaodong Liu, Jianfeng Gao and Weizhu Chen“DeBERTa: Decoding-enhanced BERT with Disentangled Attention”InarXiv preprint, 2021DOI:10.48550/arXiv.2006.03654
• [Kim+25]Jeonghoon Kim et al.“Peri-LN: Revisiting Normalization Layer in the Transformer Architecture”InForty-second International Conference on Machine Learning, 2025
• [KSK21]Wonjae Kim, Bokyung Son and Ildoo Kim“ViLT: Vision-and-Language Transformer Without Convolution or Region Supervision”InarXiv preprint, 2021DOI:10.48550/arXiv.2102.03334
• [MYZ25]Fanxu Meng, Zengwei Yao and Muhan Zhang“TransMLA: Multi-Head Latent Attention Is All You Need”InarXiv preprint, 2025DOI:10.48550/arXiv.2502.07864
• [Rad+18]Alec Radford, Karthik Narasimhan, Tim Salimans and Ilya Sutskever“Improving Language Understanding by Generative Pre-Training”InOpenAI, 2018
• [Rad+21]Alec Radford et al.“Learning Transferable Visual Models From Natural Language Supervision”InarXiv preprint, 2021DOI:10.48550/arXiv.2103.00020
• [Rom+22]Robin Rombach et al.“High-Resolution Image Synthesis with Latent Diffusion Models”InarXiv preprint, 2022DOI:10.48550/arXiv.2112.10752
• [Sha20]Noam Shazeer“GLU Variants Improve Transformer”InarXiv preprint, 2020DOI:10.48550/arXiv.2002.05202
• [Su+23]Jianlin Su et al.“RoFormer: Enhanced Transformer with Rotary Position Embedding”arXiv, 2023DOI:10.48550/arXiv.2104.09864
• [Tea+24]Gemma Team et al.“Gemma 2: Improving Open Language Models at a Practical Size”InarXiv preprint, 2024DOI:10.48550/arXiv.2408.00118
• [The25]undef The Gemma 3 Team“Gemma 3 Technical Report”InarXiv preprintabs/2503.19786, 2025DOI:10.48550/arXiv.2503.19786
• [Tou+23]Hugo Touvron et al.“LLaMA: Open and Efficient Foundation Language Models”InarXiv preprint, 2023DOI:10.48550/arXiv.2302.13971
• [Vas+17]Ashish Vaswani et al.“Attention is All you Need”InAdvances in Neural Information Processing Systems, 2017
• [War+24]Benjamin Warner et al.“Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference”InarXiv preprint, 2024DOI:10.48550/arXiv.2412.13663
• [Xio+20]Ruibin Xiong et al.“On Layer Normalization in the Transformer Architecture”InProceedings of the 37th International Conference on Machine Learning, 2020
• [Yan+25]An Yang et al.“Qwen3 Technical Report”InarXiv preprint, 2025DOI:10.48550/arXiv.2505.09388
• [Zha+24]Aston Zhang, Zachary Lipton, Mu Li and Alexander J. Smola“Dive into Deep Learning”Cambridge University Press, 2024
• [ZS19]Biao Zhang and Rico Sennrich“Root Mean Square Layer Normalization”InAdvances in Neural Information Processing Systems, 2019