![Revisit consent button](../images/body-df29beacd7496220cc94.svg) We value your privacy We use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking "Accept All", you consent to our use of cookies. Customise Reject All Accept All Customise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg) We use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below. The cookies that are categorised as "Necessary" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more NecessaryAlways Active Necessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data. * Cookie BCTempID * Duration 10 minutes * Description No description available. * Cookie __cf_bm * Duration 1 hour * Description This cookie, set by Cloudflare, is used to support Cloudflare Bot Management. * Cookie AWSALBCORS * Duration 7 days * Description Amazon Web Services set this cookie for load balancing. * Cookie _cfuvid * Duration session * Description Cloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services * Cookie li_gc * Duration 6 months * Description Linkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes. * Cookie __hssrc * Duration session * Description This cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session. * Cookie __hssc * Duration 1 hour * Description HubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie. * Cookie wpEmojiSettingsSupports * Duration session * Description WordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly. * Cookie BCSessionID * Duration 1 year 1 month 4 days * Description Blueconic sets this cookie as a unique identifier for the BlueConic profile. * Cookie _octo * Duration 1 year * Description No description available. * Cookie logged_in * Duration 1 year * Description No description available. * Cookie __Secure-YEC * Duration past * Description YouTube sets this cookie to stores the user's video player preferences using embedded YouTube video * Cookie __eoi * Duration 6 months * Description Description is currently not available. * Cookie AWSALBTGCORS * Duration 7 days * Description No description available. * Cookie login-status-p * Duration past * Description Description is currently not available. * Cookie AWSALBTG * Duration 7 days * Description No description available. * Cookie csrf_token * Duration session * Description No description available. * Cookie token_v2 * Duration 1 day * Description Description is currently not available. * Cookie D * Duration 1 year * Description Description is currently not available. * Cookie PHPSESSID * Duration session * Description This cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed. * Cookie VISITOR_PRIVACY_METADATA * Duration 6 months * Description YouTube sets this cookie to store the user's cookie consent state for the current domain. * Cookie cookietest * Duration session * Description The cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience. * Cookie __Host-airtable-session * Duration 1 year * Description This cookie is used to enable us to integrate the services of Airtable. * Cookie __Host-airtable-session.sig * Duration 1 year * Description This cookie is used to enable us to integrate the services of Airtable. * Cookie m * Duration 1 year 1 month 4 days * Description Stripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly. * Cookie BIGipServer* * Duration session * Description Marketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. * Cookie __cfruid * Duration session * Description Cloudflare sets this cookie to identify trusted web traffic. * Cookie _GRECAPTCHA * Duration 6 months * Description Google Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks. * Cookie __Secure-YNID * Duration 6 months * Description Google cookie used to protect user security and prevent fraud, especially during the login process. * Cookie cookieyes-consent * Duration 1 year * Description CookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors. Functional Functional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features. * Cookie lidc * Duration 1 day * Description LinkedIn sets the lidc cookie to facilitate data center selection. * Cookie brw * Duration 1 year * Description No description available. * Cookie brwConsent * Duration 5 minutes * Description Description is currently not available. * Cookie WMF-Uniq * Duration 1 year * Description Description is currently not available. * Cookie loom_anon_comment * Duration 1 year * Description No description available. * Cookie loom_referral_video * Duration session * Description Description is currently not available. * Cookie VISITOR_INFO1_LIVE * Duration 6 months * Description A cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface. * Cookie yt-remote-connected-devices * Duration Never Expires * Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos. * Cookie ytidb::LAST_RESULT_ENTRY_KEY * Duration Never Expires * Description The cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future. * Cookie yt-remote-device-id * Duration Never Expires * Description YouTube sets this cookie to store the user's video preferences using embedded YouTube videos. * Cookie yt-remote-session-name * Duration session * Description The yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video. * Cookie yt-remote-fast-check-period * Duration session * Description The yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos. * Cookie yt-remote-session-app * Duration session * Description The yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player. * Cookie yt-remote-cast-available * Duration session * Description The yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player. * Cookie yt-remote-cast-installed * Duration session * Description The yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video. * Cookie cp_session * Duration 3 months * Description Codepen sets this cookie for Help systems found in the website. * Cookie loid * Duration 1 year 1 month 4 days * Description This cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform. Analytics Analytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc. * Cookie __hstc * Duration 6 months * Description Hubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session). * Cookie hubspotutk * Duration 6 months * Description HubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts. * Cookie _ga * Duration 1 year 1 month 4 days * Description Google Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors. * Cookie _ga_* * Duration 1 year 1 month 4 days * Description Google Analytics sets this cookie to store and count page views. * Cookie __Host-psifi.analyticsTrace * Duration 6 hours * Description Description is currently not available. * Cookie __Host-psifi.analyticsTraceV2 * Duration 6 hours * Description Description is currently not available. * Cookie _gh_sess * Duration session * Description GitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form. * Cookie YSC * Duration session * Description YSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages. * Cookie ajs_anonymous_id * Duration 1 year * Description This cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before. * Cookie vuid * Duration 1 year 1 month 4 days * Description Vimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. Performance Performance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors. * Cookie AWSALB * Duration 7 days * Description AWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target. * Cookie acq * Duration past * Description Description is currently not available. * Cookie acq.sig * Duration past * Description Description is currently not available. * Cookie ptc * Duration 2 years * Description No description available. Advertisement Advertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns. * Cookie muc_ads * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to collect user behaviour and interaction data to optimize the website. * Cookie guest_id_marketing * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. * Cookie guest_id_ads * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. * Cookie personalization_id * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting. * Cookie guest_id * Duration 1 year 1 month 4 days * Description Twitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences. * Cookie bcookie * Duration 1 year * Description LinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs. * Cookie __Secure-ROLLOUT_TOKEN * Duration 6 months * Description YouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment. * Cookie yt.innertube::nextId * Duration Never Expires * Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen. * Cookie yt.innertube::requests * Duration Never Expires * Description YouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen. * Cookie session_tracker * Duration session * Description This cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website. * Cookie edgebucket * Duration session * Description Reddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content. * Cookie did * Duration 1 year * Description Arbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year. Uncategorised Other uncategorised cookies are those that are being analysed and have not been classified into a category as yet. No cookies to display. Reject All Save My Preferences Accept All Toggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation [](https://towardsdatascience.com/) Publish AI, ML & data-science insights to a global community of data professionals. Sign in[Submit an Article](https://write.insightmediagroup.io) [Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/) * * * [Write For TDS](https://towardsdatascience.com/submissions/) [LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) # The Math Behind Multi-Head Attention in Transformers Deep Dive into Multi-Head Attention, the secret element in Transformers and LLMs. Let's explore its math, and build it from scratch. [Cristian Leo](https://towardsdatascience.com/author/cristianleo120/) July 16, 202416 min read ### Deep Dive into Multi-Head Attention, the secret element in Transformers and LLMs. Let's explore its math, and build it from scratch in Python ![Image generated by DALL-E](../images/000.jpg) Image generated by DALL-E Learn this step by step with the interactive [AI Engineer](https://roadmap.sh/ai-engineer), [Computer Science](https://roadmap.sh/computer-science) and [Machine Learning](https://roadmap.sh/machine-learning) roadmaps. ## 1: Introduction ### 1.1: Transformers Overview The Transformer architecture, introduced by Vaswani et al. in their paper "[Attention is All You Need](https://arxiv.org/abs/1706.03762)" has transformed deep learning, particularly in natural language processing (NLP). Transformers use a self-attention mechanism, enabling them to handle input sequences all at once. This parallel processing allows for faster computation and better management of long-range dependencies within the data. This doesn't sound familiar? Don't worry as it will be at the end of this article. Let's first take a brief look at what a Transformer looks like. ![Transformer Architecture \(Architecture in "Attention is all you need"\) - Image by Author](../images/001.jpeg) Transformer Architecture (Architecture in "[Attention is all you need](https://arxiv.org/abs/1706.03762)") - Image by Author A Transformer consists of two main parts: an **encoder** and a **decoder**. The encoder processes the input sequence to create a continuous representation, while the decoder generates the output sequence from this representation. Both the encoder and the decoder have multiple layers, each containing two essential components: a multi-head self-attention mechanism and a position-wise feed-forward network. In this article, we'll focus on the multi-head attention mechanism, but we'll explore the entire Transformer architecture in future articles. ### 1.2: Multi-Head Attention Overview Multi-head attention enables the model to focus on different parts of the input sequence at the same time, capturing various aspects of the data. Think of it like having multiple spotlights shining on different parts of a stage. Each spotlight (or "head") can illuminate a different performer (or data feature), allowing the audience (or model) to see the whole scene more clearly. By splitting the input into multiple subspaces, each with its attention mechanism, multi-head attention provides the model with several views of the input data. This setup helps the model understand complex relationships within the data more effectively. This mechanism allows the _Transformer_ to capture different relationships in the data by attending to various parts of the sequence. This improves the learning process by offering multiple perspectives of the input, enhancing the model's ability to generalize. It also increases the model's expressiveness by enabling it to learn different aspects of the input data simultaneously. These capabilities make multi-head attention a crucial component in the success of Transformer models across a range of applications, from language translation to image processing. ## 2: The Mathematical Foundations ![Multi-Head Attention Architecture - Image by Author](../images/002.jpeg) Multi-Head Attention Architecture - Image by Author ### **2.1: Attention Mechanism** The attention mechanism in neural networks is designed to mimic the human ability to focus on specific parts of information while processing data. Imagine reading a book: your eyes don't pay equal attention to every word on the page. Instead, they focus more on the important words that help you understand the story. Similarly, in neural networks, attention allows the model to dynamically weigh the importance of different input elements. This means the model can prioritize parts of the input sequence that are more relevant for generating the output, improving its performance in tasks like language translation, text summarization, and more. Mathematically, the attention mechanism can be described using a set of queries, keys, and values. Let's denote the input as a set of queries **Q** , keys **K** , and values **V**. These are typically linear transformations of the input data. The attention scores are calculated by taking the dot product of the query with all keys, which gives a measure of similarity. For a query **q** and a set of keys **k1, k2, ..., kn** ​, the attention scores are given by: ![Attention Score - Image by Author](../images/003.png) Attention Score - Image by Author Think of this as comparing how similar each word (key) in a sentence is to the word (query) you're focusing on. Higher scores indicate greater similarity. To prevent the dot products from becoming too large, especially when dealing with high-dimensional vectors, we scale the scores by the square root of the dimension of the keys, **d_k** : ![Scaled Score Formula - Image by Author](../images/004.png) Scaled Score Formula - Image by Author This is like adjusting the intensity of the spotlight based on how large the stage is. It ensures that the scores remain manageable and helps maintain stable gradients during training. This scaling ensures that the values passed to the softmax function have a standard deviation close to 1, which helps maintain stable gradients during training. To see why this is necessary, consider the properties of dot products and high-dimensional vectors. When we compute the dot product of two vectors **q** and **k_i** ​ of dimension **d_k** ​, the expected value of their dot product is proportional to **d_k​**. Without scaling, as **d_k** ​ increases, the variance of the dot product grows, leading to very large values which can cause the softmax function to produce near-binary outputs (i.e., probabilities close to 0 or 1). This sharpness reduces the model's ability to learn effectively because it makes the gradients very small. By dividing the dot product by **d_k** ​​, we normalize the input to the softmax function, ensuring that the scores remain within a reasonable range. This normalization helps the model to maintain a balance, enabling it to learn more effectively and stably. These scaled scores are then passed through a softmax function to obtain the attention weights. The softmax function converts the scores into probabilities, which represent the importance of each key relative to the query: ![Attention Weight using Softmax Function - Image by Author](../images/005.png) Attention Weight using Softmax Function - Image by Author This step is like converting the adjusted spotlight intensities into a clear ranking, highlighting the most relevant parts of the scene more brightly. The final attention output is obtained by taking a weighted sum of the values, using the attention weights: ![Attention Output Formula - Image by Author](../images/006.png) Attention Output Formula - Image by Author Here, **v_i** ​ represents the value corresponding to key **k_i**. This weighted sum combines the most relevant information from the values, much like focusing your attention on the most important parts of the book to understand the story better. ### **2.2: Multi-Head Attention** Multi-head attention is an advanced form of the attention mechanism that allows a model to focus on different parts of the input sequence simultaneously, capturing various relationships within the data. Instead of having a single attention mechanism, multi-head attention processes the input into multiple "heads," each with its own set of queries, keys, and values. Each head performs the attention operation independently, and their outputs are then combined. This enhances the model's ability to understand complex patterns and dependencies in the data. Imagine you're trying to understand a complex scene with many elements. If you had multiple pairs of eyes, each looking at different parts of the scene, you'd get a more comprehensive understanding. Similarly, multi-head attention allows the model to focus on different parts of the input data at once, providing a richer and more detailed representation. Given an input sequence **X** , we project it into queries **Q** , keys **K** , and values **V** using learned linear transformations. For each head **i** , we have separate weight matrices **W_Q** ​, **W_K** ​, and **W_V** ​: ![Queries, Keys, and Values Linear Transformations - Image by Author](../images/007.png) Queries, Keys, and Values Linear Transformations - Image by Author These projections allow each head to focus on different aspects of the input data. For each head **i** , we compute the attention scores using the scaled dot-product attention mechanism. The attention output for head **i** is: ![Attention Formula - Image by Author](../images/008.png) Attention Formula - Image by Author Here, **d_k** ​ is the dimension of the key vectors, ensuring the scores are appropriately scaled. After computing the attention outputs for all heads, we concatenate them along the feature dimension. If we have **h** heads, each producing an output of dimension **d_v** ​, the concatenated output will have a dimension of **h** ×**d_v** ​: ![Multi-Head Concatenation - Image by Author](../images/009.png) Multi-Head Concatenation - Image by Author The concatenated output is then projected back to the original input dimension **d** ​ using a learned weight matrix **W_O** : ![Final Layer of Multi-Head Attention - Image by Author](../images/010.png) Final Layer of Multi-Head Attention - Image by Author This final linear transformation combines the outputs from all heads into a single representation. The core idea behind combining multiple attention heads is to allow the model to capture different types of information from the input sequence simultaneously. By having multiple heads, each head can learn to attend to different parts of the input or different features. This diversity in attention leads to a richer and more nuanced representation of the data. ### **2.3: Position-wise Feed-Forward Networks** In the Transformer architecture, each layer consists of a multi-head attention mechanism followed by a position-wise feed-forward network. These feed-forward layers are applied independently to each position in the sequence, hence the term "position-wise." Essentially, they are simple fully connected neural networks applied separately and identically to each position of the input sequence. Imagine a factory where every product on the conveyor belt goes through the same set of machines. Each machine processes the product in a specific way, adding something new or refining it. Similarly, each position in the sequence is processed independently by the feed-forward layers, transforming and enhancing the representation. The purpose of these feed-forward layers is to introduce non-linearity and additional learning capacity to the model. After the attention mechanism has aggregated information from different parts of the sequence, the feed-forward network processes this information to further transform and refine the representation. Mathematically, a position-wise feed-forward network consists of two linear transformations with a ReLU activation function in between. Given an input **x** at a particular position, the feed-forward network can be represented as: ![Feed-Forward Network Formula - Image by Author](../images/011.png) Feed-Forward Network Formula - Image by Author Here: * **W1** ​ and **W2** ​ are learned weight matrices. * **b1** ​ and **b2** ​ are learned bias vectors. * **max(0, xW1​+b1​)** represents the ReLU activation function applied element-wise. The input **x** is first linearly transformed using the weight matrix **W1** ​ and bias **b1** ​: ![Linear Transformation of the input - Image by Author](../images/012.png) Linear Transformation of the input - Image by Author Think of this step as passing the input through the first machine in the factory, which adds initial modifications based on learned weights and biases. The linear transformation is followed by a ReLU activation function, which introduces non-linearity: ![ReLU Formula - Image by Author](../images/013.png) ReLU Formula - Image by Author ReLU (Rectified Linear Unit) sets all negative values to zero, allowing the model to capture non-linear relationships in the data. This step is like ensuring only positive contributions from the first machine are passed on. The activated output is then passed through a second linear transformation using weight matrix **W2** ​ and bias **b2** ​: ![Second FFN - Image by Author](../images/body-52896ff70587b76501e3.png) Second FFN - Image by Author This final step further refines the output, much like the second machine in the factory making additional modifications to produce a finished product. The position-wise feed-forward network in the Transformer architecture further processes the information captured by the multi-head attention mechanism. While the attention mechanism allows the model to focus on different parts of the sequence and aggregate context-specific information, the feed-forward network refines and transforms this information at each position. This enhances the model's ability to capture complex patterns and dependencies. ## 3: Building Multi-Head Attention from Scratch In this section, we will break down and explain the implementation of a multi-head attention mechanism from scratch using Python and `numpy`. The goal is to understand how the input is modified during the process. Before proceed with your reading, take a look at the code we will cover in this section. You should be able to get a general understanding, but don't worry if now, as we will go over each line in this section. > [**models-from-scratch-python/Multi-Head Attention/demo.py at main ·...**](https://github.com/cristianleoo/models-from-scratch-python/blob/main/Multi-Head%20Attention/demo.py) To begin, we define the `MultiHeadAttention` class, which is responsible for managing the parameters needed for the multi-head attention mechanism. Let's go through this step-by-step to understand how we set it up. python ``` import numpy class MultiHeadAttention: def __init__(self, num_hiddens, num_heads, dropout=0.0, bias=False): self.num_heads = num_heads self.num_hiddens = num_hiddens self.d_k = self.d_v = num_hiddens // num_heads ``` In the initialization method, we first set the number of attention heads and the total number of hidden units in the model. These values are provided as arguments when the class is instantiated. * `num_hiddens`: This represents the total number of hidden units in the model. It's a crucial parameter because it determines the size of the linear transformations applied to the input data. * `num_heads`: This indicates the number of attention heads. Each head will independently learn to focus on different parts of the input, enabling the model to capture various aspects of the data. * `dropout`: This is the dropout rate, which is not used in this particular implementation but is included for completeness. * `bias`: This is a boolean flag indicating whether to include biased terms in the linear transformations. We then calculate the dimensions of the queries and values for each head. Since the total number of hidden units (`num_hiddens`) is split across all heads (`num_heads`), each head will have a query and value dimension of `num_hiddens // num_heads`. python ``` self.W_q = np.random.rand(num_hiddens, num_hiddens) self.W_k = np.random.rand(num_hiddens, num_hiddens) self.W_v = np.random.rand(num_hiddens, num_hiddens) self.W_o = np.random.rand(num_hiddens, num_hiddens) ``` Next, we initialize the weight matrices for the queries, keys, values, and output transformations. These weight matrices are randomly initialized: * `W_q` is used to transform the input data into queries. It has dimensions `num_hiddens x num_hiddens`, meaning it maps the input features to the query space. * `W_k` is used to transform the input data into keys. It also has dimensions `num_hiddens x num_hiddens`, mapping the input features to the key space. * `W_v` is used to transform the input data into values, with the same dimensions as the previous matrices. * `W_o` is used to transform the concatenated output of all heads back to the original input dimensions. python ``` if bias: self.b_q = np.random.rand(num_hiddens) self.b_k = np.random.rand(num_hiddens) self.b_v = np.random.rand(num_hiddens) self.b_o = np.random.rand(num_hiddens) else: self.b_q = self.b_k = self.b_v = self.b_o = np.zeros(num_hiddens) ``` Finally, we initialize the bias vectors for the queries, keys, values, and output transformations. If the `bias` parameter is set to `True`, these biases are randomly initialized. Otherwise, they are set to zero: * `b_q`: Bias for the query transformation. * `b_k`: Bias for the key transformation. * `b_v`: Bias for the value transformation. * `b_o`: Bias for the output transformation. The biases have dimensions equal to the number of hidden units, `num_hiddens`. By setting up these weights and biases, we ensure that each attention head can independently learn to focus on different parts of the input data. Next, we define methods to prepare and transform the data for multi-head attention. First, let's look at the `transpose_qkv` method: python ``` def transpose_qkv(self, X): X = X.reshape(X.shape[0], X.shape[1], self.num_heads, -1) X = X.transpose(0, 2, 1, 3) return X.reshape(-1, X.shape[2], X.shape[3]) ``` This method is responsible for reshaping and transposing the input data to prepare it for multi-head attention. In particular: python ``` X = X.reshape(X.shape[0], X.shape[1], self.num_heads, -1) ``` This line reshapes the input tensor `X` to have four dimensions: `(batch_size, sequence_length, num_heads, depth_per_head)`. * `X.shape[0]` is the batch size. * `X.shape[1]` is the sequence length (number of positions in the input sequence). * `self.num_heads` is the number of attention heads. * `-1` automatically infers the size of the last dimension (depth per head) so that the total number of elements remains the same. python ``` X = X.transpose(0, 2, 1, 3) ``` This line transposes the tensor so that the dimensions are reordered to `(batch_size, num_heads, sequence_length, depth_per_head)`. This rearrangement ensures that each attention head processes its part of the input sequence independently. python ``` return X.reshape(-1, X.shape[2], X.shape[3]) ``` This final reshape flattens the batch and head dimensions into a single dimension, resulting in a tensor of shape `(batch_size * num_heads, sequence_length, depth_per_head)`. By doing this, `transpose_qkv` ensures that the input data is split correctly among the multiple heads, with each head having the appropriate dimensions to process its segment of the data. Next, we have the `transpose_output` method: python ``` def transpose_output(self, X): X = X.reshape(-1, self.num_heads, X.shape[1], X.shape[2]) X = X.transpose(0, 2, 1, 3) return X.reshape(X.shape[0], X.shape[1], -1) ``` This method reverses the transformation done by `transpose_qkv` to combine the outputs from all heads back into the original shape. After transposing our matrices, we can process with the scaled dot-product attention mechanism, which allows the model to focus on different parts of the input sequence with varying degrees of importance. python ``` def scaled_dot_product_attention(self, Q, K, V, valid_lens): d_k = Q.shape[-1] scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(d_k) if valid_lens is not None: mask = np.arange(scores.shape[-1]) < valid_lens[:, None] scores = np.where(mask[:, None, :], scores, -np.inf) attention_weights = np.exp(scores - np.max(scores, axis=-1, keepdims=True)) attention_weights /= attention_weights.sum(axis=-1, keepdims=True) return np.matmul(attention_weights, V) ``` The inputs to this method are the query (**Q**), key (**K**), and value (**V**) matrices. These matrices are derived from the input data through linear transformations. python ``` d_k = Q.shape[-1] ``` Here, we extract the dimension of the key vectors, **d_k** ​, from the last dimension of the query matrix **Q**. This value is used for scaling the attention scores. python ``` scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(d_k) ``` We calculate the attention scores by performing a matrix multiplication of **Q** and the transpose of **K**. The scores are then scaled by the square root of **d_k** ​. This scaling helps prevent the scores from growing too large, which can lead to issues during the softmax calculation. Next, we define the forward pass method to process the input data through the multi-head attention mechanism. This method is crucial as it orchestrates the entire multi-head attention process, from transforming the input data to combining the outputs from multiple heads. python ``` def forward(self, queries, keys, values, valid_lens): queries = self.transpose_qkv(np.dot(queries, self.W_q) + self.b_q) keys = self.transpose_qkv(np.dot(keys, self.W_k) + self.b_k) values = self.transpose_qkv(np.dot(values, self.W_v) + self.b_v) if valid_lens is not None: valid_lens = np.repeat(valid_lens, self.num_heads, axis=0) output = self.scaled_dot_product_attention(queries, keys, values, valid_lens) output_concat = self.transpose_output(output) return np.dot(output_concat, self.W_o) + self.b_o ``` Let's break down the forward method: python ``` queries = self.transpose_qkv(np.dot(queries, self.W_q) + self.b_q) keys = self.transpose_qkv(np.dot(keys, self.W_k) + self.b_k) values = self.transpose_qkv(np.dot(values, self.W_v) + self.b_v) ``` First, the input queries, keys, and values are projected into their respective subspaces using the learned weight matrices (`W_q`, `W_k`, `W_v`) and biases (`b_q`, `b_k`, `b_v`). This is done by performing matrix multiplication with the weight matrices and adding the biases. The results are then transformed for multi-head attention using the `transpose_qkv` method, which reshapes and transposes the data to ensure each head processes the input independently. **Queries, keys, and values** are the transformed inputs, now prepared for multi-head attention. python ``` if valid_lens is not None: valid_lens = np.repeat(valid_lens, self.num_heads, axis=0) ``` If `valid_lens` (valid lengths) are provided, they are repeated for each head. This ensures that the appropriate mask is created for each attention head, allowing the model to focus only on valid positions within the sequences. python ``` output = self.scaled_dot_product_attention(queries, keys, values, valid_lens) ``` The method then calls `scaled_dot_product_attention` with the transformed queries, keys, values, and repeated valid lengths. This function calculates the attention scores, applies the softmax function to obtain attention weights, and computes the weighted sum of the values to produce the attention output for each head. python ``` output_concat = self.transpose_output(output) return np.dot(output_concat, self.W_o) + self.b_o ``` After obtaining the attention outputs from all heads, the method concatenates these outputs along the feature dimension using `transpose_output`. This method reverses the initial transformation, combining the outputs from all heads back into a single representation. The concatenated output is then transformed back to the original input dimension using a final linear transformation with weight matrix `W_o` and bias `b_o`. Lastly, we test the class with some sample data. Here's how we do it: graphql ``` # Define dimensions and initialize multi-head attention num_hiddens, num_heads = 100, 5 attention = MultiHeadAttention(num_hiddens, num_heads, dropout=0.5, bias=False) ``` We initialize the `MultiHeadAttention` class with 100 hidden units and 5 attention heads. This sets up the necessary parameters and weight matrices for the multi-head attention mechanism. python ``` # Define sample data batch_size, num_queries, num_kvpairs = 2, 4, 6 valid_lens = np.array([3, 2]) X = np.random.rand(batch_size, num_queries, num_hiddens) # Use random data to simulate input queries Y = np.random.rand(batch_size, num_kvpairs, num_hiddens) # Use random data to simulate key-value pairs ``` We create random data to simulate input queries (`X`) and key-value pairs (`Y`). The batch size is 2, the number of queries is 4, and the number of key-value pairs is 6. We also define valid lengths (`valid_lens`) to indicate the valid positions within the sequences. python ``` # Apply multi-head attention output = attention.forward(X, Y, Y, valid_lens) ``` We pass the sample data through the multi-head attention mechanism using the `forward` method. This processes the input queries, keys, and values, applying the multi-head attention calculations. python ``` print("Output shape:", output.shape) # Output should be: (2, 4, 100) print("Output data:", output) ``` We print the shape and content of the output. The expected output shape ensures that the output dimensions match the original input dimensions. Then, we print the output data after computing Multi-Head Attention. Now that you have an understanding of how the Multi-Head Attention mechanism works, try tweaking it. For instance, change the number of heads, try adding multiple FFN before and after Multi-Head Attention. Also, you could try to implement it in a machine translation task, and see it in action. Let me know if you would like me to do that in a next article. ## Conclusion Transformers have transformed deep learning, especially in NLP, by using self-attention mechanisms that allow for parallel processing of input sequences. This approach not only speeds up computation but also handles long-range dependencies more effectively than traditional recurrent neural networks). In this article, we've gained a comprehensive understanding of multi-head attention in Transformers, from it's math theory to a practical code implementation. Maybe for now the concepts we will still be abstracts, as you can't really do anything with the outputs from Multi-Head Attention, but soon we will see how they play a key role in the transformer architecture, which is the base of the well-known LLMs around (Claude, ChatGPT, ...). Stay tuned for future articles, where we'll explore the remaining components of the Transformer architecture, offering deeper insights into this powerful model. ## References * Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). [Attention is All You Need](https://arxiv.org/abs/1706.03762). In Advances in Neural Information Processing Systems (NeurIPS). * [Alammar, J. (2018). The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/). jalammar.github.io. ## Register Now to Continue Reading ### Unlock Towards Data Science with a Free Account All articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more. Register for freeSign in to your account * * * ![](../images/body-8f73ac5a40c91a7b503f.jpeg) Written By ## [Cristian Leo](https://towardsdatascience.com/author/cristianleo120/) [See all from Cristian Leo](https://towardsdatascience.com/author/cristianleo120/) [Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Artificial Intelligence](https://towardsdatascience.com/tag/artificial-intelligence/)[Machine Learning](https://towardsdatascience.com/tag/machine-learning/)[Math](https://towardsdatascience.com/tag/math/)[Programming](https://towardsdatascience.com/tag/programming/) Share This Article * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fthe-math-behind-multi-head-attention-in-transformers-c26cba15f625%2F&title=The%20Math%20Behind%20Multi-Head%20Attention%20in%20Transformers) * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fthe-math-behind-multi-head-attention-in-transformers-c26cba15f625%2F&title=The%20Math%20Behind%20Multi-Head%20Attention%20in%20Transformers) * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fthe-math-behind-multi-head-attention-in-transformers-c26cba15f625%2F&text=The%20Math%20Behind%20Multi-Head%20Attention%20in%20Transformers) Towards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program. [Write for TDS](https://towardsdatascience.com/submissions/) ## Related Articles * [![](../images/014.gif)](https://towardsdatascience.com/demystifying-artificial-intelligence-d2887879aaf1/) ## [Demystifying Artificial Intelligence!](https://towardsdatascience.com/demystifying-artificial-intelligence-d2887879aaf1/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) Breaking down the field of Artificial Intelligence and what the future holds for it. [Bharath K](https://towardsdatascience.com/author/bharath-k1297/) November 2, 20208 min read * [![](../images/015.jpg)](https://towardsdatascience.com/demystifying-the-lottery-ticket-hypothesis-in-deep-learning-158570b62674/) ## [Demystifying the Lottery Ticket Hypothesis in Deep Learning](https://towardsdatascience.com/demystifying-the-lottery-ticket-hypothesis-in-deep-learning-158570b62674/) [Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/) Why lottery tickets are the next big thing in training neural networks [David Yastremsky](https://towardsdatascience.com/author/davidyastremsky/) March 3, 20224 min read * [![](../images/016.jpeg)](https://towardsdatascience.com/fitbit-google-ai-5g-e6764d39e44e/) ## [Fitbit + Google AI + 5G = ?](https://towardsdatascience.com/fitbit-google-ai-5g-e6764d39e44e/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) Google Buys Fitbit for $2.1 billion [Alex Moltzau](https://towardsdatascience.com/author/alexmoltzau/) November 2, 20195 min read * [![](../images/017.jpg)](https://towardsdatascience.com/transformers-for-tabular-data-tabtransformer-deep-dive-5fb2438da820/) ## [Transformers for Tabular Data: TabTransformer Deep Dive](https://towardsdatascience.com/transformers-for-tabular-data-tabtransformer-deep-dive-5fb2438da820/) [Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/) Making sense out of TabTransformer and learning to apply it in TensorFlow [Antons Tocilins-Ruberts](https://towardsdatascience.com/author/antonsruberts/) September 13, 20227 min read * [![](../images/018.jpeg)](https://towardsdatascience.com/openai-api-intro-11-practical-implementation-examples-of-the-models-behind-chatgpt-18601f68b51b/) ## [OpenAI API - Intro & 11 Practical Implementation Examples of the Models Behind ChatGPT](https://towardsdatascience.com/openai-api-intro-11-practical-implementation-examples-of-the-models-behind-chatgpt-18601f68b51b/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) A programmatic approach to use models behind ChatGPT. [Farzad Nobar](https://towardsdatascience.com/author/fmnobar/) November 7, 202319 min read * [![](../images/019.jpg)](https://towardsdatascience.com/these-are-the-top-6-things-people-get-wrong-when-teaching-themselves-data-science-ae8a0bd72fe7/) ## [These Are the Top 6 Things People Get Wrong When Teaching Themselves Data Science](https://towardsdatascience.com/these-are-the-top-6-things-people-get-wrong-when-teaching-themselves-data-science-ae8a0bd72fe7/) [Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/) And how to avoid making the same mistakes yourself. [Madison Hunter](https://towardsdatascience.com/author/madison13/) September 6, 202110 min read * [![](../images/020.jpg)](https://towardsdatascience.com/7-handy-built-in-python-functions-3a32d2aa0301/) ## [7 Handy Built-In Python Functions](https://towardsdatascience.com/7-handy-built-in-python-functions-3a32d2aa0301/) [Programming](https://towardsdatascience.com/category/programming/) Do you know how to properly use these functions already built into Python? Awareness of... [Vivienne DiFrancesco](https://towardsdatascience.com/author/heythatsviv/) May 2, 20215 min read * [![](../images/body-63beb50f664a055ed739.jpg)](https://towardsdatascience.com/deploying-a-serverless-inference-service-with-amazon-sagemaker-pipelines-2d2f3cc96c39/) ## [Deploying a Serverless Inference Service with Amazon SageMaker Pipelines](https://towardsdatascience.com/deploying-a-serverless-inference-service-with-amazon-sagemaker-pipelines-2d2f3cc96c39/) [Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/) Deploying a serverless inference service with Amazon SageMaker Pipelines, AWS Lambda, Amazon API Gateway, and CDK [Sofian Hamiti](https://towardsdatascience.com/author/sofian-hamiti/) March 15, 20216 min read * [![](../images/022.jpg)](https://towardsdatascience.com/how-to-effectively-work-with-frontend-and-backend-code/) ## [How to Work Effectively with Frontend and Backend Code](https://towardsdatascience.com/how-to-effectively-work-with-frontend-and-backend-code/) [LLM Applications](https://towardsdatascience.com/category/artificial-intelligence/llm-applications/) Learn how to be an effective full-stack engineer with Claude Code [Eivind Kjosbakken](https://towardsdatascience.com/author/oieivind/) February 4, 20266 min read * [YouTube](https://www.youtube.com/c/TowardsDataScience) * [X](https://x.com/TDataScience) * [LinkedIn](https://www.linkedin.com/company/towards-data-science/) * [Threads](https://www.threads.net/@towardsdatascience) * [Bluesky](https://bsky.app/profile/towardsdatascience.com) [](https://towardsdatascience.com/) Your home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals. Subscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/) © Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network) #### Recommended Articles Close * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607686/image?etag=1782161464448&width=500&height=500) How Relevance Models Foreshadowed Transformers for NLP ](https://towardsdatascience.com/how-relevance-models-foreshadowed-transformers-for-nlp/ "How Relevance Models Foreshadowed Transformers for NLP") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ "Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ "The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs") * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ "Does Your Company Have a Data Strategy?")