Sequence-Wise Feature Normalization with Affine Transform

In modern deep learning architectures, particularly those handling sequential data like Transformers, normalization techniques play a crucial role in stabilizing training and improving convergence. Unlike batch normalization which operates across samples in a batch, layer normalization standardizes activations across the feature dimension independently for each position in a sequence.

Given a 3D input tensor X with dimensions (batch_size, sequence_length, d_model) representing a batch of sequences, your task is to normalize each feature vector at every sequence position. For each position in the sequence, you will:

1. Compute the mean (μ) across the feature dimension (d_model)
2. Compute the variance (σ²) across the feature dimension
3. Standardize the features using: $\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}$
4. Apply affine transformation using learnable parameters: $y = \gamma \cdot \hat{x} + \beta$

Mathematical Formulation:

For each position in the sequence, given a feature vector $\mathbf{x} \in \mathbb{R}^{d}$:

$$\mu = \frac{1}{d} \sum_{i=1}^{d} x_i$$

$$\sigma^2 = \frac{1}{d} \sum_{i=1}^{d} (x_i - \mu)^2$$

$$\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$$

$$y_i = \gamma_i \cdot \hat{x}_i + \beta_i$$

Where:

• d is the dimension of the feature space (d_model)
• ε (epsilon) is a small constant (typically 1e-5) added for numerical stability
• γ (gamma) is the scale parameter (shape: 1 × 1 × d_model)
• β (beta) is the shift parameter (shape: 1 × 1 × d_model)

Your Task: Implement a function that performs layer normalization on a 3D tensor representing batch sequences. The function should normalize across the last dimension (feature dimension) independently for each batch and sequence position, then apply the affine transformation using the provided gamma and beta parameters.