Transformer · Token embedding

Token embedding

reuses weightstarget_modality.targets_bottom(targets) · modalities.py 92–96

Look up a learned feature vector for each token ID.

{\htmlData{light=target-embedded target-embedding}{y_e}}=\sqrt{512}\,E[{\htmlData{light=shifted-ids shift-right}{I^{in}}}]

in: $I^{in}$ [B batch, T target positions]
out: $y_e$ [B batch, T target positions, 512 d]

Parameters

0 new parameters.

Looks up the table that source-embedding owns; its rows are counted once, there.

Explanation

Multiplied by √512 ≈ 22.63 after the lookup (§3.4; multiply_embedding_mode = "sqrt_depth", common_hparams.py 68), and positions holding padding id 0 get a zero vector (modalities.py 83).

It reads the same table $E$ as the source embedding (modalities.py 92–96).

Codemodalities.py 51–73, 75–96 · common_hparams.py 68–71

modalities.py 75–96

  def bottom_simple(self, x, name, reuse):
    with tf.variable_scope(name, reuse=reuse):
      # Squeeze out the channels dimension.
      x = tf.squeeze(x, axis=3)
      var = self._get_weights()
      ret = tf.gather(var, x)
      if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
        ret *= self._body_input_depth**0.5
      ret *= tf.expand_dims(tf.to_float(tf.not_equal(x, 0)), -1)
      return ret

  def bottom(self, x):
    if self._model_hparams.shared_embedding_and_softmax_weights:
      return self.bottom_simple(x, "shared", reuse=None)
    else:
      return self.bottom_simple(x, "input_emb", reuse=None)

  def targets_bottom(self, x):
    if self._model_hparams.shared_embedding_and_softmax_weights:
      return self.bottom_simple(x, "shared", reuse=True)
    else:
      return self.bottom_simple(x, "target_emb", reuse=None)

modalities.py 51–73

  def _get_weights(self):
    """Create or get concatenated embedding or softmax variable.

    Returns:
       a list of self._num_shards Tensors.
    """
    num_shards = self._model_hparams.symbol_modality_num_shards
    shards = []
    for i in xrange(num_shards):
      shard_size = (self._vocab_size // num_shards) + (
          1 if i < self._vocab_size % num_shards else 0)
      var_name = "weights_%d" % i
      shards.append(
          tf.get_variable(
              var_name, [shard_size, self._body_input_depth],
              initializer=tf.random_normal_initializer(
                  0.0, self._body_input_depth**-0.5)))
    if num_shards == 1:
      ret = shards[0]
    else:
      ret = tf.concat(shards, 0)
    ret = eu.ConvertGradientToTensor(ret)
    return ret

common_hparams.py 68–71

      multiply_embedding_mode="sqrt_depth",
      norm_type="none",  # "batch", layer", "noam", "none".
      layer_norm_epsilon=1e-6,
      symbol_modality_num_shards=16,

Sources3

Direct evidence

Embedding: gather rows of the table, multiply by √depth, zero the padding id 0; with shared weights, inputs and targets read the same table.

symbol-bottom · tensorflow/tensor2tensor / tensor2tensor/models/modalities.py
Pinned revision 0c66117d4297 ↗ · lines 75–96

Read code excerpt

  def bottom_simple(self, x, name, reuse):
    with tf.variable_scope(name, reuse=reuse):
      # Squeeze out the channels dimension.
      x = tf.squeeze(x, axis=3)
      var = self._get_weights()
      ret = tf.gather(var, x)
      if self._model_hparams.multiply_embedding_mode == "sqrt_depth":
        ret *= self._body_input_depth**0.5
      ret *= tf.expand_dims(tf.to_float(tf.not_equal(x, 0)), -1)
      return ret

  def bottom(self, x):
    if self._model_hparams.shared_embedding_and_softmax_weights:
      return self.bottom_simple(x, "shared", reuse=None)
    else:
      return self.bottom_simple(x, "input_emb", reuse=None)

  def targets_bottom(self, x):
    if self._model_hparams.shared_embedding_and_softmax_weights:
      return self.bottom_simple(x, "shared", reuse=True)
    else:
      return self.bottom_simple(x, "target_emb", reuse=None)

Direct evidence

The symbol table: vocab_size × depth, built from num_shards shards and concatenated.

symbol-weights · tensorflow/tensor2tensor / tensor2tensor/models/modalities.py
Pinned revision 0c66117d4297 ↗ · lines 51–73

Read code excerpt

  def _get_weights(self):
    """Create or get concatenated embedding or softmax variable.

    Returns:
       a list of self._num_shards Tensors.
    """
    num_shards = self._model_hparams.symbol_modality_num_shards
    shards = []
    for i in xrange(num_shards):
      shard_size = (self._vocab_size // num_shards) + (
          1 if i < self._vocab_size % num_shards else 0)
      var_name = "weights_%d" % i
      shards.append(
          tf.get_variable(
              var_name, [shard_size, self._body_input_depth],
              initializer=tf.random_normal_initializer(
                  0.0, self._body_input_depth**-0.5)))
    if num_shards == 1:
      ret = shards[0]
    else:
      ret = tf.concat(shards, 0)
    ret = eu.ConvertGradientToTensor(ret)
    return ret

Direct evidence

Defaults the models inherit: embeddings are multiplied by √depth; the symbol table is stored in 16 shards.

hparams-common · tensorflow/tensor2tensor / tensor2tensor/models/common_hparams.py
Pinned revision 0c66117d4297 ↗ · lines 68–71

Read code excerpt

      multiply_embedding_mode="sqrt_depth",
      norm_type="none",  # "batch", layer", "noam", "none".
      layer_norm_epsilon=1e-6,
      symbol_modality_num_shards=16,