Transformer · Position-wise feed-forward network

in: $x$ [B batch, S source positions, 512 d]
out: $\Delta$ [B batch, S source positions, 512 d]
parameters: 2,099,712 · × 6 layers = 12,598,272

Parameters

2,099,712 parameters per instance, across its learned components

Exact per-instance accounting
Component	Calculation	Parameters
First linear map	$512 \times 2{,}048 + 2{,}048$	1,050,624
Second linear map	$2{,}048 \times 512 + 512$	1,049,088
Per-instance total		2,099,712

Across 6 independent layers: $2{,}099{,}712 \times 6 = 12{,}598{,}272$ 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

FFN $(x)=\max(0,\,xW_1+b_1)W_2+b_2$ (§3.3, eq. 2): the same at every position, with different parameters from layer to layer.

Codetransformer.py 229–244 · common_layers.py 354–360, 1160–1194

transformer.py 229–244

def transformer_ffn_layer(x, hparams):
  """Feed-forward layer in the transformer.

  Args:
    x: a Tensor of shape [batch_size, length, hparams.hidden_size]
    hparams: hyperparmeters for model

  Returns:
    a Tensor of shape [batch_size, length, hparams.hidden_size]
  """
  if hparams.ffn_layer == "conv_hidden_relu":
    return common_layers.conv_hidden_relu(
        x,
        hparams.filter_size,
        hparams.hidden_size,
        dropout=hparams.relu_dropout)

common_layers.py 1160–1194

def conv_hidden_relu(inputs,
                     hidden_size,
                     output_size,
                     kernel_size=(1, 1),
                     second_kernel_size=(1, 1),
                     summaries=True,
                     dropout=0.0,
                     **kwargs):
  """Hidden layer with RELU activation followed by linear projection."""
  name = kwargs.pop("name") if "name" in kwargs else None
  with tf.variable_scope(name, "conv_hidden_relu", [inputs]):
    if inputs.get_shape().ndims == 3:
      is_3d = True
      inputs = tf.expand_dims(inputs, 2)
    else:
      is_3d = False
    conv_f1 = conv if kernel_size == (1, 1) else separable_conv
    h = conv_f1(
        inputs,
        hidden_size,
        kernel_size,
        activation=tf.nn.relu,
        name="conv1",
        **kwargs)
    if dropout != 0.0:
      h = tf.nn.dropout(h, 1.0 - dropout)
    if summaries and not tf.get_variable_scope().reuse:
      tf.summary.histogram("hidden_density_logit",
                           relu_density_logit(
                               h, list(range(inputs.shape.ndims - 1))))
    conv_f2 = conv if second_kernel_size == (1, 1) else separable_conv
    ret = conv_f2(h, output_size, second_kernel_size, name="conv2", **kwargs)
    if is_3d:
      ret = tf.squeeze(ret, 2)
    return ret

common_layers.py 354–360

def conv(inputs, filters, kernel_size, **kwargs):
  return conv_internal(tf.layers.conv2d, inputs, filters, kernel_size, **kwargs)


def conv1d(inputs, filters, kernel_size, **kwargs):
  return tf.squeeze(
      conv(tf.expand_dims(inputs, 2), filters, (kernel_size, 1), **kwargs), 2)

Sources3

Direct evidence

The feed-forward layer is conv_hidden_relu with filter_size hidden units and hidden_size outputs.

ffn · tensorflow/tensor2tensor / tensor2tensor/models/transformer.py
Pinned revision 0c66117d4297 ↗ · lines 229–244

Read code excerpt

def transformer_ffn_layer(x, hparams):
  """Feed-forward layer in the transformer.

  Args:
    x: a Tensor of shape [batch_size, length, hparams.hidden_size]
    hparams: hyperparmeters for model

  Returns:
    a Tensor of shape [batch_size, length, hparams.hidden_size]
  """
  if hparams.ffn_layer == "conv_hidden_relu":
    return common_layers.conv_hidden_relu(
        x,
        hparams.filter_size,
        hparams.hidden_size,
        dropout=hparams.relu_dropout)

Direct evidence

conv_hidden_relu: a 1×1 convolution to the hidden size with ReLU, optional dropout, then a 1×1 convolution to the output size.

conv-relu · tensorflow/tensor2tensor / tensor2tensor/models/common_layers.py
Pinned revision 0c66117d4297 ↗ · lines 1160–1194

Read code excerpt

def conv_hidden_relu(inputs,
                     hidden_size,
                     output_size,
                     kernel_size=(1, 1),
                     second_kernel_size=(1, 1),
                     summaries=True,
                     dropout=0.0,
                     **kwargs):
  """Hidden layer with RELU activation followed by linear projection."""
  name = kwargs.pop("name") if "name" in kwargs else None
  with tf.variable_scope(name, "conv_hidden_relu", [inputs]):
    if inputs.get_shape().ndims == 3:
      is_3d = True
      inputs = tf.expand_dims(inputs, 2)
    else:
      is_3d = False
    conv_f1 = conv if kernel_size == (1, 1) else separable_conv
    h = conv_f1(
        inputs,
        hidden_size,
        kernel_size,
        activation=tf.nn.relu,
        name="conv1",
        **kwargs)
    if dropout != 0.0:
      h = tf.nn.dropout(h, 1.0 - dropout)
    if summaries and not tf.get_variable_scope().reuse:
      tf.summary.histogram("hidden_density_logit",
                           relu_density_logit(
                               h, list(range(inputs.shape.ndims - 1))))
    conv_f2 = conv if second_kernel_size == (1, 1) else separable_conv
    ret = conv_f2(h, output_size, second_kernel_size, name="conv2", **kwargs)
    if is_3d:
      ret = tf.squeeze(ret, 2)
    return ret

Direct evidence

conv1d is a 1×1 tf.layers.conv2d, which owns a kernel and a bias.

conv1d · tensorflow/tensor2tensor / tensor2tensor/models/common_layers.py
Pinned revision 0c66117d4297 ↗ · lines 354–360

Read code excerpt

def conv(inputs, filters, kernel_size, **kwargs):
  return conv_internal(tf.layers.conv2d, inputs, filters, kernel_size, **kwargs)


def conv1d(inputs, filters, kernel_size, **kwargs):
  return tf.squeeze(
      conv(tf.expand_dims(inputs, 2), filters, (kernel_size, 1), **kwargs), 2)