In Exercise 5 it is written that :
I am finding it difficult to understand point 2 , why we scale the embedding by multiplying it with square root of embedding dimension.
Any help is appreciated ! Thanks
Here’s the explanation from attention is all you need paper.
While for small values of d_k
the two mechanisms perform similarly, additive attention outperforms dot product attention without scaling for larger values of d_k
[3]. 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 1. To counteract this effect, we scale the dot products by
1 / \sqrt d_k
The paper quote @balaji.ambresh shared is about a related but different scaling, the one inside scaled dot-product attention (dividing QK^T by √d_k to stabilize softmax gradients). What the exercise is asking about is the scaling applied to the embedding itself before positional encoding is added.
The reason: embedding vectors learned by tf.keras.layers.Embedding tend to have small magnitudes (they’re initialized roughly in a unit-variance range, so individual values sit close to zero). Positional encoding values, on the other hand, use sin/cos and range between -1 and +1, so they’re relatively large compared to unscaled embeddings.
If you add positional encodings directly to these small embedding values, the positional signal dominates and the semantic content of the embedding gets drowned out. Multiplying the embedding by √d_model scales the embedding values up so they’re on a comparable magnitude to the positional encoding. That way, when you add them, neither signal overwhelms the other.
This comes from Section 3.4 of the original Transformer paper (“Attention Is All You Need”, Vaswani et al.), where they state they multiply the embedding weights by √d_model. The lecture on “Transformer Network” in Week 4 also walks through this step.
Think of it this way: if d_model = 512, embedding values might average around ±0.01–0.05, while positional encodings sit at ±0.5–1.0. Multiplying by √512 ≈ 22.6 brings the embedding into the same ballpark.
My addition by expanding this line:
The weights (in the embedding) are used for the pre-softmax linear transformation, so the linear transformation is the sum of d_{\text{model}} terms: \sum_i w_ih_i . Practically, d_{\text{model}} is large, and because h_i is ranged like +/- 1 (it is normalized - see discussion about LayerNorm in the paper), the sum can be large if w_i is not kept small.
Being too large is no good for softmax, so w_i should stay small.
And because w_i is shared between the pre-softmax linear transformation and the input/output embeddings, we need to compensate with a multiplying factor for the input/output embeddings as @arman101 explained.
A comment about initialization.
By default, tf.keras.layers.Embedding initializes weights using RandomUniform which is ranged between -0.05 and 0.05. tf.keras.layers.MultiHeadAttention is more adaptive because it uses GlorotUniform which is ranged between -l and l where l depends on d_{\text{model}} and is typically on the order of 0.01. This shows that the default std value for Llama is 0.02. So, either way, they are all on the order of 0.01.
But why same weight matrix is shared between the two embedding layers and the pre-softmax linear transformation as both are doing different tasks?
Hello @Dhruv_A,
The fact is that even though weight sharing (commonly known as Weight Tying) is common, there are also inapplicable cases like when the input and output have different sets of vocabulary (translation, for example). Even for the applicable case, there are also critiques. I think it would be the best for me to go through some references with you, so you can find out more yourself.
This paper (Press and Wolf, 2017) reported that “weight tying can reduce the size of neural translation models to less than half of their original size without harming their performance”.
I think reducing the size without performance loss is already a good reason. Reducing trainable weights don’t just reduce memory footprint, and traditionally, as we have learned in the DLS, it is also a way to regularize or counteract overfitting. However, I would keep in mind that it was back in 2017. In 2026, we have more computational resources and data.
This work (Inan et al., 2016) presented theory and empirical validation of weight tying under their loss framework. You may also find it interesting to read the discussions between the authors and reviewers on OpenReview.
This one (Bertolotti and Cazzola, 2024) is relatively new, and they claimed that their "findings indicate that words … with similar semantics tend to be encoded in similar input embeddings, while words that appear in similar contexts are encoded in similar output embeddings … As a consequence of these findings, the tying of the input and output embeddings is encouraged only when the distributional hypothesis holds for the underlying data. " One of the authors actually made a very interesting overview about Weight Tying and his idea on Reddit, and it is a nice read!
One more… This (Lopado et al., 2026) was published a few months ago and they found that “weight tying optimizes the embedding matrix for output prediction, compromising its role in input representation.” and I think it can be linked back to the previous paper.
In fact, if we look at this Llama codebase again, we find that the output projection layer has its own set of weights.
Cheers,
Raymond