Transformer · Token embedding

Token embedding

learnedinput_modality.bottom(inputs) · modalities.py 86–90

Look up a learned feature vector for each token ID.

{\htmlData{light=source-embedded source-embedding}{x_e}}=\sqrt{512}\,E[{\htmlData{light=source-ids}{I^{src}}}]

in: $I^{src}$ [B batch, S source positions]
out: $x_e$ [B batch, S source positions, 512 d]
parameters: 18,944,000

Parameters

18,944,000 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Token embedding	$37{,}000 \times 512$	18,944,000
Per-instance total		18,944,000

About 37,000 × 512 values, counted here once for all three uses.

Derived from the configured table dimensions; checkpoint tensors have not been independently counted.

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

One table $E$ of about 37,000 rows serves the source embedding, the target embedding and the output projection (§3.4, shared_embedding_and_softmax_weights). The paper gives the shared byte-pair vocabulary of WMT 2014 English–German as «about 37000 tokens» (§5.1), so the table’s size and every total that holds it are approximate. The code stores the table in 16 shards (common_hparams.py 71).

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,