# Why do we use square root of key dimension for scaling?

**URL:** <https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800>\
**Category:** Attention in Transformers: Concepts and Code in Py\
**Created:** [February 26, 2025, 7:31am UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800 "2025-02-26T07:31:01Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![rr11n](https://avatars.discourse-cdn.com/v4/letter/r/eb8c5e/32.png) [@rr11n](https://community.deeplearning.ai/u/rr11n)\
**Post date:** [February 26, 2025, 7:31am UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800/1 "2025-02-26T07:31:01Z")

</div>

Hello Folks,

I have a query around a part of the attention equation. The softmax function computes the dot product and then scales it by dividing with the square root of the dimensions in K. Why cant we use cosine similarity instead, which will give us a scaled dot product? Am I missing something here?

Thanks,  
RR.

---

<div class="post-metadata">

**Author:** ![carlosrl](https://avatars.discourse-cdn.com/v4/letter/c/2acd7d/32.png) [@carlosrl](https://community.deeplearning.ai/u/carlosrl)\
**Post date:** [February 26, 2025, 1:49pm UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800/2 "2025-02-26T13:49:00Z")

</div>

> [@rr11n](#):
>
> I have a query around a part of the attention equation. The softmax function computes the dot product and then scales it by dividing with the square root of the dimensions in K. Why cant we use cosine similarity instead, which will give us a scaled dot product? Am I missing something here?

Why not use cosine?  
The attention mechanism computes attention scores using the following formula:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d\_k}}\right)V

So, why not use cosine similarity?  
Cosine similarity measures the angle between two vectors, defined as:

\text{cosine\_similarity}(Q, K) = \frac{Q \cdot K}{\|Q\| \|K\|}

At first view, cosine similarity might seem like a natural choice because it normalizes the dot product by the magnitudes of the vectors, ensuring that the similarity score lies in the range[−1,1]. However, there are several reasons why cosine similarity is not typically used in attention mechanisms:

- Magnitude Information
- Softmax requires unbounded Inputs
- Computational Efficiency
- Empirical Evidence

To summarize:

- The scaled dot product is preferred in attention mechanisms because it preserves magnitude information, provides unbounded inputs for softmax, and is computationally efficient.
- Cosine similarity, while theoretically appealing, loses magnitude information, constrains the input range for softmax, and introduces additional computational overhead.
- Empirical evidence supports the effectiveness of the scaled dot product in a wide range of applications.

In fact, cosine similarity is not used in attention mechanisms because it loses magnitude information, constrains softmax inputs, and introduces computational overhead.

---

<div class="post-metadata">

**Author:** ![rr11n](https://avatars.discourse-cdn.com/v4/letter/r/eb8c5e/32.png) [@rr11n](https://community.deeplearning.ai/u/rr11n)\
**Post date:** [February 27, 2025, 8:11am UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800/3 "2025-02-27T08:11:45Z")

</div>

Thank you Carlos! That helps!

---

<div class="post-metadata">

**Author:** ![slopez](https://sea1.discourse-cdn.com/dlai/user_avatar/community.deeplearning.ai/slopez/32/102103_2.png) [@slopez](https://community.deeplearning.ai/u/slopez)\
**Post date:** [July 11, 2025, 9:40am UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800/4 "2025-07-11T09:40:57Z")

</div>

I think the question is still unanswered. Ok, don’t use cosine similarity to keep computation simple and preserve magnitude information. However, you could achieve the same without scaling _at all_. Why then pick a constant \sqrt{d\_k}?

Copy-paste from ChatGPT to this specific question:

> [@](#):
>
> Each element of QK^T is a dot product q\_i \cdot k\_j, where both q\_i and k\_j are vectors in \mathbb{R}^{d\_k}.
> 
> If these vectors are roughly unit-variance and uncorrelated (as is often the case early in training), the **expected magnitude of the dot product** grows with the dimension:
> 
> \mathbb{E}[q\_i \cdot k\_j] \sim \mathcal{O}(d\_k)
> 
> This means:
> 
> - As d\_k increases, the dot products can become **very large in magnitude** , especially when Q and K are not normalized.
> 
> * * *
> 
> ## **😬 Why is this a problem?**
> 
> You pass these dot products into a **softmax** :
> 
> \text{softmax}(q\_i \cdot k\_j)
> 
> If the inputs to softmax are **large in magnitude** , the softmax becomes **very peaky** :
> 
> - One token gets almost all the weight (close to 1)
> - All other tokens get near-zero weights
> 
> This causes:
> 
> - **Hard attention** instead of soft attention (bad for training early on)
> - **Vanishing gradients** , because softmax becomes too flat (gradient near-zero for almost all elements)
> - **Unstable training** , especially with larger models
> 
> * * *
> 
> ## **✅ The Solution: Scaling by** \sqrt{d\_k}
> 
> By dividing the dot products by \sqrt{d\_k}, we reduce their magnitude and keep the softmax in a **healthy, moderate range** :
> 
> \frac{q\_i \cdot k\_j}{\sqrt{d\_k}} \sim \mathcal{O}(1)
> 
> This ensures:
> 
> - The softmax doesn’t saturate
> - Gradients flow better
> - The attention weights stay meaningfully distributed

---

<div class="post-metadata">

**Author:** ![slopez](https://sea1.discourse-cdn.com/dlai/user_avatar/community.deeplearning.ai/slopez/32/102103_2.png) [@slopez](https://community.deeplearning.ai/u/slopez)\
**Post date:** [July 12, 2025, 11:01am UTC](https://community.deeplearning.ai/t/why-do-we-use-square-root-of-key-dimension-for-scaling/776800/5 "2025-07-12T11:01:50Z")

</div>

Edit: Even better than my previous answer (which I can’t edit), directly from [the paper],([https://arxiv.org/pdf/1706.03762](https://arxiv.org/pdf/1706.03762)) at the end of section 3.2.1:

> We suspect that for large values of d\_k , the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients. To counteract this effect, we scale the dot products by 1/\sqrt{d\_k}.

.