Transformer

Published
Vaswani et al. · Jun 2017
Line
Encoder–decoder
View
Architecture
Source

tensorflow/tensor2tensor v1.0.14 @ 0c66117 · arXiv:1706.03762v7 · WMT 2014 English–German · vocabulary of about 37,000 · no trained weights were publishedHow to read

Decoder: 6 layers

Open inside →×

learnedtransformer_decoder(decoder_input, encoder_output, residual_fn, …) · transformer.py 170–226

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

w=cn+DPdrop(Δffn)w=c_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{ffn}})
yn+1=LayerNorm⁡(w)y_{n+1}=\operatorname{LayerNorm}(w)
All 41 steps
Masked multi-head self-attention:\text{Masked multi-head self-attention:}
Mts={0s≤t−109s>t{\htmlClass{vl-stream-context}{M_{ts}}}=\begin{cases}0 & s\le t\\ -10^{9} & s>t\end{cases}
pQKV=xWQKV+bQKVp^{QKV}=xW^{QKV}+b^{QKV}
q=p1:512QKV{\htmlClass{vl-stream-q}{q}}=p^{QKV}_{1:512}
k=p513:1024QKV{\htmlClass{vl-stream-k}{k}}=p^{QKV}_{513:1024}
v=p1025:1536QKV{\htmlClass{vl-stream-v}{v}}=p^{QKV}_{1025:1536}
Q=reshape⁡(q){\htmlClass{vl-stream-q}{Q}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q}})
Q~=Q/64\tilde {\htmlClass{vl-stream-q}{Q}}={\htmlClass{vl-stream-q}{Q}}/\sqrt{64}
K=reshape⁡(k){\htmlClass{vl-stream-k}{K}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k}})
V=reshape⁡(v){\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})
S=Q~KTS=\tilde QK^{\mathsf T}
S~=S+M\tilde S=S+{\htmlClass{vl-stream-context}{M}}
A[b,i,t,s]=exp⁡(S~[b,i,t,s])∑u=0T−1exp⁡(S~[b,i,t,u])A[b,i,t,s]=\frac{\exp(\tilde S[b,i,t,s])}{\sum_{u=0}^{T-1}\exp(\tilde S[b,i,t,u])}
O=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δ=oWO+bO\Delta=oW^{O}+b^{O}
u=yn+DPdrop(Δself)u=y_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{self}})
sn=LayerNorm⁡(u)s_n=\operatorname{LayerNorm}(u)
Encoder–decoder attention:\text{Encoder–decoder attention:}
q=xWQ+bQ{\htmlClass{vl-stream-q}{q}}=xW^{{\htmlClass{vl-stream-q}{Q}}}+b^{{\htmlClass{vl-stream-q}{Q}}}
pKV=zWKV+bKVp^{KV}=zW^{KV}+b^{KV}
k=p1:512KV{\htmlClass{vl-stream-k}{k}}=p^{KV}_{1:512}
v=p513:1024KV{\htmlClass{vl-stream-v}{v}}=p^{KV}_{513:1024}
Q=reshape⁡(q){\htmlClass{vl-stream-q}{Q}}=\operatorname{reshape}({\htmlClass{vl-stream-q}{q}})
Q~=Q/64\tilde {\htmlClass{vl-stream-q}{Q}}={\htmlClass{vl-stream-q}{Q}}/\sqrt{64}
K=reshape⁡(k){\htmlClass{vl-stream-k}{K}}=\operatorname{reshape}({\htmlClass{vl-stream-k}{k}})
V=reshape⁡(v){\htmlClass{vl-stream-v}{V}}=\operatorname{reshape}({\htmlClass{vl-stream-v}{v}})
S=Q~KTS=\tilde QK^{\mathsf T}
S~=S+M\tilde S=S+{\htmlClass{vl-stream-context}{M}}
A[b,i,t,s]=exp⁡(S~[b,i,t,s])∑u=0S−1exp⁡(S~[b,i,t,u])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=AVO=AV
o=reshape⁡(O)o=\operatorname{reshape}(O)
Δ=oWO+bO\Delta=oW^{O}+b^{O}
v=sn+DPdrop(Δcross){\htmlClass{vl-stream-v}{v}}=s_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{cross}})
cn=LayerNorm⁡(v)c_n=\operatorname{LayerNorm}({\htmlClass{vl-stream-v}{v}})
Position-wise feed-forward network:\text{Position-wise feed-forward network:}
h=xW1+b1h=xW_1+b_1
a=max⁡(0, h)a=\max(0,\,h)
Δ=aW2+b2\Delta=aW_2+b_2
w=cn+DPdrop(Δffn)w=c_n+D_{\mathrm{P_{\mathrm{drop}}}}(\Delta_{\mathrm{ffn}})
yn+1=LayerNorm⁡(w)y_{n+1}=\operatorname{LayerNorm}(w)
in
DPdrop(y0)D_{\mathrm{P_{\mathrm{drop}}}}(y_0) [B batch, T target positions, 512 d]zz [B batch, S source positions, 512 d]MM [B batch, 1, 1, S keys]
out
yn+1y_{n+1} [B batch, T target positions, 512 d]
parameters
4,204,032 · × 6 layers = 25,224,192
Parameters

4,204,032 parameters per layer, across its learned components

Exact per-instance accounting
ComponentCalculationParameters
Masked multi-head self-attention1,050,624 across its components1,050,624
LayerNorm512+512512 + 5121,024
Encoder–decoder attention1,050,624 across its components1,050,624
LayerNorm512+512512 + 5121,024
Position-wise feed-forward network2,099,712 across its components2,099,712
LayerNorm512+512512 + 5121,024
Per-instance total4,204,032

Across 6 independent layers: 4,204,032×6=25,224,1924{,}204{,}032 \times 6 = 25{,}224{,}192 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

In addition to the two sublayers of an encoder layer, the decoder inserts a third, which performs multi-head attention over the output of the encoder stack (§3.1).

Codetransformer.py 59–61, 170–226
transformer.py 170–226
def transformer_decoder(decoder_input,
                        encoder_output,
                        residual_fn,
                        decoder_self_attention_bias,
                        encoder_decoder_attention_bias,
                        hparams,
                        name="decoder"):
  """A stack of transformer layers.

  Args:
    decoder_input: a Tensor
    encoder_output: a Tensor
    residual_fn: a function from (layer_input, layer_output) -> combined_output
    decoder_self_attention_bias: bias Tensor for self-attention
      (see common_attention.attention_bias())
    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
      (see common_attention.attention_bias())
    hparams: hyperparameters for model
    name: a string

  Returns:
    y: a Tensors
  """
  x = decoder_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,
                decoder_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="decoder_self_attention"))
        x = residual_fn(
            x,
            common_attention.multihead_attention(
                x,
                encoder_output,
                encoder_decoder_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="encdec_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 decoder layer: masked self-attention, attention over the encoder output with the encoder padding bias, then the feed-forward network, each through residual_fn.

decoder · tensorflow/tensor2tensor / tensor2tensor/models/transformer.py
Pinned revision 0c66117d4297 ↗ · lines 170–226

Read code excerpt
def transformer_decoder(decoder_input,
                        encoder_output,
                        residual_fn,
                        decoder_self_attention_bias,
                        encoder_decoder_attention_bias,
                        hparams,
                        name="decoder"):
  """A stack of transformer layers.

  Args:
    decoder_input: a Tensor
    encoder_output: a Tensor
    residual_fn: a function from (layer_input, layer_output) -> combined_output
    decoder_self_attention_bias: bias Tensor for self-attention
      (see common_attention.attention_bias())
    encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
      (see common_attention.attention_bias())
    hparams: hyperparameters for model
    name: a string

  Returns:
    y: a Tensors
  """
  x = decoder_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,
                decoder_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="decoder_self_attention"))
        x = residual_fn(
            x,
            common_attention.multihead_attention(
                x,
                encoder_output,
                encoder_decoder_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="encdec_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))