Why OPT's token embeddings are not scaled by sqrt(dim) as in the original OPT implementation?

rodrigo-nogueira · February 7, 2023, 1:19pm

In the original OPT implementation from metaseq, token embeddings are multiplied by sqrt(embedding_dim):

facebookresearch/metaseq/blob/59403be56179e2ad4bad1be339b741a78e684cd1/metaseq/models/transformer_decoder.py#L327


      
          # see IncrementalDecoder for important information about
          # incremental state
          if incremental_state:
              tokens = tokens[:, -1:]
              if positions is not None:
                  positions = positions[:, -1:]
          
          
if token_embedding is None:
              token_embedding = self.embed_tokens(tokens)
          
          
x = embed = self.embed_scale * token_embedding
          
          
if positions is not None:
              x += positions
          
          
if self.dropout_module is not None:
              x = self.dropout_module(x)
          
          
# Returning in T x B x C format as that makes integrating sequence parallelism easier.
          x = x.transpose(0, 1).contiguous()
          return x, embed, positions

However, I couldn’t find this scaling in HF’s implementation. Is there a reason for not having it?

nielsr · February 7, 2023, 2:18pm

Hi,

The self.embed_scale attribute is set to 1.0, hence no scaling is happening (args.no_scale_embedding is set to True). See:

github.com

facebookresearch/metaseq/blob/59403be56179e2ad4bad1be339b741a78e684cd1/metaseq/models/transformer_decoder.py#LL83C50-L83C50


      
          self.embed_scale = 1.0 if args.no_scale_embedding else math.sqrt(self.embed_dim)

rodrigo-nogueira · February 7, 2023, 4:54pm

It seems that no_scale_embeddings is set to False here:

github.com

facebookresearch/metaseq/blob/51871bd73cd04c038f239ea2a26db1d7f6b37927/metaseq/models/transformer_lm.py#L76


      
          decoder_learned_pos: bool = field(
              default=False,
              metadata={"help": "use learned positional embeddings in the decoder"},
          )
          decoder_learned_sinusoidal: bool = field(
              default=False,
              metadata={
                  "help": "use learned positional embeddings init with sinusoidal in the decoder"
              },
          )
          no_scale_embedding: bool = field(
              default=False, metadata={"help": "if True, dont scale embeddings"}
          )
          checkpoint_activations: bool = field(
              default=False, metadata={"help": "checkpoint activations at each layer"}
          )
          offload_activations: bool = field(
              default=False,
              metadata={"help": "move checkpointed activations to CPU after they are used."},
          )
          # config for Fully Sharded Data Parallel (FSDP) training

and here:

github.com

facebookresearch/metaseq/blob/51871bd73cd04c038f239ea2a26db1d7f6b37927/metaseq/models/transformer_lm.py#L348


      
              args.attention_dropout = getattr(args, "attention_dropout", 0.0)
              args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 512)
              args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 2048)
              args.decoder_layers = getattr(args, "decoder_layers", 6)
              args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 8)
              args.share_decoder_input_output_embed = getattr(
                  args, "share_decoder_input_output_embed", False
              )
              args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
              args.decoder_learned_sinusoidal = getattr(args, "decoder_learned_sinusoidal", False)
              args.no_scale_embedding = getattr(args, "no_scale_embedding", False)
              args.add_bos_token = getattr(args, "add_bos_token", False)
          
          

          
@register_model_architecture("model_parallel_transformer_lm", "transformer_lm_megatron")
          def transformer_lm_megatron(args):
              args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 3072)
              args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 3072 * 4)
              args.decoder_layers = getattr(args, "decoder_layers", 72)
              args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 32)
              args.dropout = getattr(args, "dropout", 0.1)

Or is it overwritten by a config file?

nielsr · February 8, 2023, 9:41am

Pinging @ybelkada here

Topic		Replies	Views
Understanding FLOPs-per-token estimates from OpenAI's scaling laws Research	6	9182	September 20, 2023
Add new tokens and learn the embeddings of the new tokens and keeping all the other parametes frozen 🤗Tokenizers	0	380	April 30, 2021
Layoutlmv3 sequence_length vs token_sequnce_length size mismatch Models	2	503	November 19, 2022
Is resize_token_embeddings available to the FlaxPreTrainedModel? Flax/JAX Projects	1	1419	August 25, 2022
Why are my special tokens not appearing as predictions? 🤗Transformers	0	343	July 29, 2021

Why OPT's token embeddings are not scaled by sqrt(dim) as in the original OPT implementation?

Related Topics