7
$\begingroup$

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?
$\endgroup$
9
$\begingroup$

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

enter image description here

Transcription:

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).

$\endgroup$
2
  • $\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 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 at 6:17
3
$\begingroup$

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.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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