Transformer · Encoder: 6 layers

Encoder: 6 layers

learnedtransformer_encoder(encoder_input, residual_fn, encoder_self_attention_bias, hparams) · transformer.py 129–167

6 identical layers with independent weights over the source positions: self-attention, then a feed-forward network, each added to its input and normalized.

w=x'_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{ffn}})

x_{n+1}=\operatorname{LayerNorm}(w)

All 23 steps

\text{Multi-head self-attention:}

p^{QKV}=xW^{QKV}+b^{QKV}

{\htmlClass{vl-stream-q}{q}}=p^{QKV}_{1:512}

{\htmlClass{vl-stream-k}{k}}=p^{QKV}_{513:1024}

{\htmlClass{vl-stream-v}{v}}=p^{QKV}_{1025:1536}

{\htmlClass{vl-stream-q}{Q}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q}})

\tilde {\htmlClass{vl-stream-q}{Q}}={\htmlClass{vl-stream-q}{Q}}/\sqrt{64}

{\htmlClass{vl-stream-k}{K}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k}})

{\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})

S=\tilde QK^{\mathsf T}

\tilde S=S+{\htmlClass{vl-stream-context}{M}}

A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{S-1}\exp(\tilde S[b,i,t,u])}

O=AV

o=\operatorname{reshape}(O)

\Delta=oW^{O}+b^{O}

u=x_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{attn}})

x'_n=\operatorname{LayerNorm}(u)

\text{Position-wise feed-forward network:}

h=xW_1+b_1

a=\max(0,\,h)

\Delta=aW_2+b_2

w=x'_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{ffn}})

x_{n+1}=\operatorname{LayerNorm}(w)

in: $D_{\mathrm{P_{\mathrm{drop}}}}(x_0)$ [B batch, S source positions, 512 d] $M$ [B batch, 1, 1, S keys]
out: $x_{n+1}$ [B batch, S source positions, 512 d]
parameters: 3,152,384 · × 6 layers = 18,914,304

Parameters

3,152,384 parameters per layer, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
Multi-head self-attention	1,050,624 across its components	1,050,624
LayerNorm	$512 + 512$	1,024
Position-wise feed-forward network	2,099,712 across its components	2,099,712
LayerNorm	$512 + 512$	1,024
Per-instance total		3,152,384

Across 6 independent layers: $3{,}152{,}384 \times 6 = 18{,}914{,}304$ parameters.

Derived from the modules the pinned tensor2tensor v1.0.14 code creates, with a vocabulary of about 37,000 (the paper’s words); no trained weights were published to count.

Explanation

The output of each sublayer is LayerNorm $(x+\text{Sublayer}(x))$ ; all sublayers and embeddings produce outputs of $d_{\text{model}}=512$ (§3.1).

Codetransformer.py 59–61, 129–167

transformer.py 129–167

def transformer_encoder(encoder_input,
                        residual_fn,
                        encoder_self_attention_bias,
                        hparams,
                        name="encoder"):
  """A stack of transformer layers.

  Args:
    encoder_input: a Tensor
    residual_fn: a function from (layer_input, layer_output) -> combined_output
    encoder_self_attention_bias: bias Tensor for self-attention
       (see common_attention.attention_bias())
    hparams: hyperparameters for model
    name: a string

  Returns:
    y: a Tensors
  """
  x = encoder_input
  # Summaries don't work in multi-problem setting yet.
  summaries = "problems" not in hparams.values() or len(hparams.problems) == 1
  with tf.variable_scope(name):
    for layer in xrange(hparams.num_hidden_layers):
      with tf.variable_scope("layer_%d" % layer):
        x = residual_fn(
            x,
            common_attention.multihead_attention(
                x,
                None,
                encoder_self_attention_bias,
                hparams.attention_key_channels or hparams.hidden_size,
                hparams.attention_value_channels or hparams.hidden_size,
                hparams.hidden_size,
                hparams.num_heads,
                hparams.attention_dropout,
                summaries=summaries,
                name="encoder_self_attention"))
        x = residual_fn(x, transformer_ffn_layer(x, hparams))
  return x

transformer.py 59–61

    def residual_fn(x, y):
      return common_layers.layer_norm(x + tf.nn.dropout(
          y, 1.0 - hparams.residual_dropout))

Sources2

Direct evidence

Each of num_hidden_layers encoder layers: self-attention over the encoder input with the padding bias, then the feed-forward network, each through residual_fn.

encoder · tensorflow/tensor2tensor / tensor2tensor/models/transformer.py
Pinned revision 0c66117d4297 ↗ · lines 129–167

Read code excerpt

def transformer_encoder(encoder_input,
                        residual_fn,
                        encoder_self_attention_bias,
                        hparams,
                        name="encoder"):
  """A stack of transformer layers.

  Args:
    encoder_input: a Tensor
    residual_fn: a function from (layer_input, layer_output) -> combined_output
    encoder_self_attention_bias: bias Tensor for self-attention
       (see common_attention.attention_bias())
    hparams: hyperparameters for model
    name: a string

  Returns:
    y: a Tensors
  """
  x = encoder_input
  # Summaries don't work in multi-problem setting yet.
  summaries = "problems" not in hparams.values() or len(hparams.problems) == 1
  with tf.variable_scope(name):
    for layer in xrange(hparams.num_hidden_layers):
      with tf.variable_scope("layer_%d" % layer):
        x = residual_fn(
            x,
            common_attention.multihead_attention(
                x,
                None,
                encoder_self_attention_bias,
                hparams.attention_key_channels or hparams.hidden_size,
                hparams.attention_value_channels or hparams.hidden_size,
                hparams.hidden_size,
                hparams.num_heads,
                hparams.attention_dropout,
                summaries=summaries,
                name="encoder_self_attention"))
        x = residual_fn(x, transformer_ffn_layer(x, hparams))
  return x

Direct evidence

Post-norm: each sublayer output is dropped out, added to its input and layer-normalized.

residual-fn · tensorflow/tensor2tensor / tensor2tensor/models/transformer.py
Pinned revision 0c66117d4297 ↗ · lines 59–61

Read code excerpt

    def residual_fn(x, y):
      return common_layers.layer_norm(x + tf.nn.dropout(
          y, 1.0 - hparams.residual_dropout))