Implementation:Tensorflow Tfjs Normalization Layers

Knowledge Sources	Tensorflow_Tfjs TensorFlow.js API
Domains	Deep_Learning, Layers_API
Last Updated	2026-02-10 06:00 GMT

Overview

The normalization.ts module implements normalization layers for the TensorFlow.js Layers API: BatchNormalization and LayerNormalization. These layers normalize activations to stabilize training, accelerate convergence, and act as regularizers.

The module also provides standalone normalization functions: batchNormalization() (applies the formula (x - mean) / sqrt(var + epsilon) * gamma + beta), normalizeBatchInTraining() (computes mean/variance from the batch and normalizes), and internal helpers for regular vs. broadcasting normalization paths.

Code Reference

Source Location

tfjs-layers/src/layers/normalization.ts (View on GitHub)

Key Signatures

// Standalone functions
export function batchNormalization(
    x: Tensor, mean: Tensor, variance: Tensor,
    beta?: Tensor, gamma?: Tensor, epsilon?: number): Tensor;

export function normalizeBatchInTraining(
    x: Tensor, gamma: Tensor, beta: Tensor,
    reductionAxes: number[], epsilon?: number): [Tensor, Tensor, Tensor];

// Layer classes
export class BatchNormalization extends Layer {
  static className = 'BatchNormalization';
  constructor(args?: BatchNormalizationLayerArgs);
  build(inputShape: Shape | Shape[]): void;
  call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
  getConfig(): serialization.ConfigDict;
}

export class LayerNormalization extends Layer {
  static className = 'LayerNormalization';
  constructor(args?: LayerNormalizationLayerArgs);
  build(inputShape: Shape | Shape[]): void;
  call(inputs: Tensor | Tensor[], kwargs: Kwargs): Tensor | Tensor[];
  getConfig(): serialization.ConfigDict;
}

Import

import { BatchNormalization, LayerNormalization } from './layers/normalization';

Layer Details

BatchNormalization

Parameter	Default	Description
axis	-1	Axis to normalize (typically features axis)
momentum	0.99	Moving average momentum for mean/variance
epsilon	1e-3	Small constant to avoid division by zero
center	true	Whether to add learnable offset (beta)
scale	true	Whether to add learnable scale (gamma)
betaInitializer	'zeros'	Initializer for beta
gammaInitializer	'ones'	Initializer for gamma
movingMeanInitializer	'zeros'	Initializer for moving mean
movingVarianceInitializer	'ones'	Initializer for moving variance

During training, normalizes using batch statistics and updates moving averages. During inference, uses stored moving averages.

LayerNormalization

Parameter	Default	Description
axis	-1	Axis or axes to normalize
epsilon	1e-3	Small constant to avoid division by zero
center	true	Whether to add learnable offset (beta)
scale	true	Whether to add learnable scale (gamma)

Normalizes across the specified axes independently for each sample (no batch statistics). Unlike BatchNormalization, behavior is the same during training and inference.

I/O Contract

Input

• BatchNormalization: Tensors of rank 2-4. The normalization axis defaults to -1 (last axis). Supports channelsFirst and channelsLast data formats.
• LayerNormalization: Tensors of any rank. The axis parameter specifies which axes to normalize. Supports single axis or array of axes.

Output

Both layers output tensors of the same shape as the input, normalized and optionally scaled/shifted.

Build Phase

• BatchNormalization: Creates gamma, beta (if center/scale=true), movingMean, and movingVariance weights. Moving statistics are non-trainable.
• LayerNormalization: Creates gamma and beta weights shaped by the specified axes dimensions.

Usage Examples

import * as tf from '@tensorflow/tfjs';

// BatchNormalization after convolution
const model = tf.sequential();
model.add(tf.layers.conv2d({
  inputShape: [28, 28, 1], filters: 32, kernelSize: 3
}));
model.add(tf.layers.batchNormalization());
model.add(tf.layers.activation({ activation: 'relu' }));

// LayerNormalization in a transformer-style architecture
const lnModel = tf.sequential();
lnModel.add(tf.layers.dense({ inputShape: [512], units: 512 }));
lnModel.add(tf.layers.layerNormalization({ axis: -1 }));

// BatchNormalization with channels-first format
model.add(tf.layers.batchNormalization({
  axis: 1,  // normalize along channels axis
  momentum: 0.9,
  epsilon: 1e-5
}));

Related Pages

• Tensorflow_Tfjs_Convolutional_Layers - Convolutional layers commonly followed by batch normalization
• Tensorflow_Tfjs_Exports_Layers - Public API factory functions tf.layers.batchNormalization() and tf.layers.layerNormalization()
• Tensorflow_Tfjs_Layer_Variables - LayerVariable used for gamma, beta, movingMean, and movingVariance weights
• Tensorflow_Tfjs_Backend_Operations - Backend operations used in the normalization computation