Transformer · Tied vocabulary projection

Tied vocabulary projection

reuses weightstf.matmul(body_output, var, transpose_b=True) · modalities.py 117

Project each normalized feature vector into vocabulary logits.

{\htmlData{light=logits projection}{L}}={\htmlData{light=decoder-output decoder}{y_{6}}}E^{\mathsf T}

in: $y_{6}$ [B batch, T target positions, 512 d]
out: $L$ [B batch, T target positions, 37000 vocabulary]

Parameters

Reuses Token embedding: weights are counted there, with no additional parameter ownership here.

Explanation

Each entry scores one vocabulary token as the next token for its position. The softmax and the choice of token happen outside the model.

The pre-softmax linear transformation shares its matrix with the two embeddings (§3.4); without the √d scale, which only the embeddings apply. The paper’s Figure 1 ends in a softmax; the code returns the logits and the loss applies the softmax.

Codemodalities.py 98–120

modalities.py 98–120

  def top(self, body_output, targets):
    """Generate logits.

    Args:
      body_output: A Tensor with shape [batch, p0, p1, body_input_depth]
      targets: A Tensor with shape [batch, p0, p1, 1]
    Returns:
      logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
    """
    if self._model_hparams.shared_embedding_and_softmax_weights:
      scope_name = "shared"
      reuse = True
    else:
      scope_name = "softmax"
      reuse = False
    with tf.variable_scope(scope_name, reuse=reuse):
      var = self._get_weights()
      shape = tf.shape(body_output)[:-1]
      body_output = tf.reshape(body_output, [-1, self._body_input_depth])
      logits = tf.matmul(body_output, var, transpose_b=True)
      logits = tf.reshape(logits, tf.concat([shape, [self._vocab_size]], 0))
      # insert a channels dimension
      return tf.expand_dims(logits, 3)

Sources1

Direct evidence

With shared weights the logits are the decoder output times the transposed table.

symbol-top · tensorflow/tensor2tensor / tensor2tensor/models/modalities.py
Pinned revision 0c66117d4297 ↗ · lines 98–120

Read code excerpt

  def top(self, body_output, targets):
    """Generate logits.

    Args:
      body_output: A Tensor with shape [batch, p0, p1, body_input_depth]
      targets: A Tensor with shape [batch, p0, p1, 1]
    Returns:
      logits: A Tensor with shape  [batch, p0, p1, ?, vocab_size].
    """
    if self._model_hparams.shared_embedding_and_softmax_weights:
      scope_name = "shared"
      reuse = True
    else:
      scope_name = "softmax"
      reuse = False
    with tf.variable_scope(scope_name, reuse=reuse):
      var = self._get_weights()
      shape = tf.shape(body_output)[:-1]
      body_output = tf.reshape(body_output, [-1, self._body_input_depth])
      logits = tf.matmul(body_output, var, transpose_b=True)
      logits = tf.reshape(logits, tf.concat([shape, [self._vocab_size]], 0))
      # insert a channels dimension
      return tf.expand_dims(logits, 3)