While going over a Tensorflow tutorial for the Transformer model I realized that their implementation of the Encoder layer (and the Decoder) scales word embeddings by sqrt of embedding dimension before adding positional encodings. Notice that this is different from scaling the dot product attention.

I'm referring to the 3rd line of the call method of the Encoder class here: https://www.tensorflow.org/tutorials/text/transformer#encoder

def call(self, x, training, mask):

  seq_len = tf.shape(x)[1]
  # adding embedding and position encoding.
  x = self.embedding(x) # (batch_size, input_seq_len, d_model)
  x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
  x += self.pos_encoding[:, :seq_len, :]
  x = self.dropout(x, training=training)

  for i in range(self.num_layers):
    x = self.enc_layers[i](x, training, mask)

  return x # (batch_size, input_seq_len, d_model)

I could not find any mention of this scaling in the papers I've read so far. People always show the input to the encoder as WE + PE, that is word embedding plus positional encoding. But this implementation seems to use sqrt(d_model) * WE + PE.

My questions:

  1. Have you ever seen this extra scaling step mentioned in a paper? I didn't find it in "Attention is all you need" (Vaswani et. al.).
  2. What is this additional scaling trying to achieve?

4 Answers 4


This is specified in the original Transformer paper, at the end of section 3.4:

enter image description here


3.4 Embeddings and Softmax

Similarly to other sequence transduction models, we use learned embeddings to convert the input tokens and output tokens to vectors of dimension 𝑑model. We also use the usual learned linear transformation and softmax function to convert the decoder output to predicted next-token probabilities. In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to [24]. In the embedding layers, we multiply those weights by √𝑑model

This aspect is not justified by the authors, either on the paper or anywhere else. It was specifically asked as an issue in the original implementation by Google with no response.

Other implementations of the Transformer have also wondered if this was actually needed (see this, this and this).

Some hypothesithed arguments (source) are:

  • It is for the sharing weight between the decoder embedding and the decoder pre-softmax linear weights.
  • It is not actually needed.
  • It is to make the positional encoding relatively smaller. This means the original meaning in the embedding vector won’t be lost when we add them together.

For reference, there are other StackExchange questions discussing this (see this and this).

  • $\begingroup$ noe, doesn't this scaling factor cancel out with the scaling factor used in the scaled dot-product attention of the first layer of the encoder/decoder? Isn't this scaling factor there exactly because the authors don't want the choice of d_model to impact the ability of the first softmax function to focus? Please see arxiv.org/pdf/1801.07704v2, figure 4. $\endgroup$
    – toliveira
    Jun 11, 2021 at 1:21
  • $\begingroup$ This scaling is probably there to place the positional embeddings in the same scale as the normal embeddings, so no, I don't think the effect of this scaling is canceling out with the scaling in the dot product, because the latter affects the already combined embeddings. $\endgroup$
    – noe
    Jun 11, 2021 at 6:17

Thank-you!! I'd also missed that multiply in my (fairseq transformer) code study, and it helps clear up a mystery that I'd noted: the (sinusoidal, non-learned) positional embeddings are initialized with a range of -1.0 to +1.0, but the word-embeddings are initialized with a mean of 0.0 and s.d. of embedding_dim ** -0.5 (0.044 for 512, 0.03125 for 1024).

So, on the face of it, the positional embeddings would overwhelm any signal coming from the word embeddings.

But now I can see word embeddings are scaled by math.sqrt(embed_dim) (22.6 for 512, 32 for 1024), it makes sense again.

Following the links in the other answer, it seems it is done this way because the same embeddings can be used in other parts of the transformer model, and that has decided the initialization values.

  • $\begingroup$ Can you explain why the word-embeddings are initialized with a mean of 0.0 and s.d. of dimension of model and where can i check this in the paper or source code?? Thank you $\endgroup$ Aug 16, 2023 at 0:02
  • $\begingroup$ @SeungyeopHyun See the quoted text in other answer for the paper, and also the links there. $\endgroup$ Aug 16, 2023 at 6:52

I agree with Noe's third explanation "It is to make the positional encoding relatively smaller. This means the original meaning in the embedding vector won’t be lost when we add them together.".

I find a reference book (Dive into Deep Learning-Section 11.7.4) mentions this as "In the following Transformer encoder implementation, we stack num_blks instances of the above TransformerEncoderBlock classes. Since we use the fixed positional encoding whose values are always between -1 and 1 , we multiply values of the learnable input embeddings by the square root of the embedding dimension to rescale before summing up the input embedding and the positional encoding."


Multiplying Weights by √dmodel

In the embedding layers, the weights are multiplied by the square root of dmodel. This is done to scale the weights appropriately and ensure that the dot product between the input tokens and the weights falls within a reasonable range. Scaling the weights helps in preventing the gradients from becoming too large or too small during training, which can lead to unstable learning dynamics.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.