Transformer · LayerNorm · OpenModelAtlas

Encoder 6 layers

LayerNorm

learnedcommon_layers.layer_norm(…) · 60

Center and scale each position’s features, then apply learned gain and bias.

\begin{aligned}\mu_d(x)&=\frac{1}{d}\sum_{j=1}^d x_j\\\sigma_d^2(x)&=\frac{1}{d}\sum_{j=1}^d(x_j-\mu_d(x))^2\\y&=\gamma_2\odot\frac{x-\mu_d(x)}{\sqrt{\sigma_d^2(x)+\varepsilon}}+\beta_2\\\varepsilon&=10^{-6}\end{aligned}

in: $w$ [B batch, S source positions, 512 d]
out: $x_{n+1}$ [B batch, S source positions, 512 d]
parameters: 1,024 · × 6 layers = 6,144

Post-norm: the sum of the sublayer’s input and its output is normalized, as the paper’s LayerNorm(x + Sublayer(x)).

Parameters

1,024 parameters per instance.

Exact per-instance accounting
Component	Calculation	Parameters
Learned gain	$512$	512
Learned bias	$512$	512
Per-instance total		1,024

Each normalization instance owns its learned gain and bias vectors.

Across 6 independent layers: $1{,}024 \times 6 = 6{,}144$ parameters.

Derived from the configured feature width; checkpoint tensors have not been independently counted.

Codecommon_layers.py 406–443 · transformer.py 59–61

common_layers.py 406–443

def layer_norm_compute_python(x, epsilon, scale, bias):
  """Layer norm raw computation."""
  mean = tf.reduce_mean(x, axis=[-1], keep_dims=True)
  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keep_dims=True)
  norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
  return norm_x * scale + bias


@function.Defun(compiled=True)
def layer_norm_compute_grad(x, epsilon, scale, bias, dy):
  y = layer_norm_compute_python(x, epsilon, scale, bias)
  dx = tf.gradients(ys=[y], xs=[x, epsilon, scale, bias], grad_ys=[dy])
  return dx


@function.Defun(
    compiled=True,
    separate_compiled_gradients=True,
    grad_func=layer_norm_compute_grad)
def layer_norm_compute(x, epsilon, scale, bias):
  return layer_norm_compute_python(x, epsilon, scale, bias)


def layer_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
  """Layer normalize the tensor x, averaging over the last dimension."""
  if filters is None:
    filters = x.get_shape()[-1]
  with tf.variable_scope(
      name, default_name="layer_norm", values=[x], reuse=reuse):
    scale = tf.get_variable(
        "layer_norm_scale", [filters], initializer=tf.ones_initializer())
    bias = tf.get_variable(
        "layer_norm_bias", [filters], initializer=tf.zeros_initializer())
    if allow_defun:
      result = layer_norm_compute(x, tf.constant(epsilon), scale, bias)
      result.set_shape(x.get_shape())
    else:
      result = layer_norm_compute_python(x, epsilon, scale, bias)

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

LayerNorm over the last axis with a learned scale and bias and ε = 1e-6.

layer-norm · tensorflow/tensor2tensor / tensor2tensor/models/common_layers.py
Pinned revision 0c66117d4297 ↗ · lines 406–443

Read code excerpt

def layer_norm_compute_python(x, epsilon, scale, bias):
  """Layer norm raw computation."""
  mean = tf.reduce_mean(x, axis=[-1], keep_dims=True)
  variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keep_dims=True)
  norm_x = (x - mean) * tf.rsqrt(variance + epsilon)
  return norm_x * scale + bias


@function.Defun(compiled=True)
def layer_norm_compute_grad(x, epsilon, scale, bias, dy):
  y = layer_norm_compute_python(x, epsilon, scale, bias)
  dx = tf.gradients(ys=[y], xs=[x, epsilon, scale, bias], grad_ys=[dy])
  return dx


@function.Defun(
    compiled=True,
    separate_compiled_gradients=True,
    grad_func=layer_norm_compute_grad)
def layer_norm_compute(x, epsilon, scale, bias):
  return layer_norm_compute_python(x, epsilon, scale, bias)


def layer_norm(x, filters=None, epsilon=1e-6, name=None, reuse=None):
  """Layer normalize the tensor x, averaging over the last dimension."""
  if filters is None:
    filters = x.get_shape()[-1]
  with tf.variable_scope(
      name, default_name="layer_norm", values=[x], reuse=reuse):
    scale = tf.get_variable(
        "layer_norm_scale", [filters], initializer=tf.ones_initializer())
    bias = tf.get_variable(
        "layer_norm_bias", [filters], initializer=tf.zeros_initializer())
    if allow_defun:
      result = layer_norm_compute(x, tf.constant(epsilon), scale, bias)
      result.set_shape(x.get_shape())
    else:
      result = layer_norm_compute_python(x, epsilon, scale, bias)

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