batchNorm1dMath function

Tensor<Vector> batchNorm1dMath(Tensor<Vector> x, Tensor<Vector> gamma, Tensor<Vector> beta, Vector runningMean, Vector runningVariance, int numFeatures, bool isTraining, double momentum, double epsilon,)

Applies Batch Normalization over a 1D vector.

During training (isTraining == true), this normalizes the input using the computed mean and variance of the current input, and updates the runningMean and runningVariance using the specified momentum. During inference, it normalizes using the provided running statistics. The normalized output is then scaled by gamma and shifted by beta. Includes an automatic differentiation node for backpropagation.

• x - Input 1D tensor.
• gamma - Learnable scale parameter vector.
• beta - Learnable shift (bias) parameter vector.
• runningMean - Persistent state vector storing the moving average of means.
• runningVariance - Persistent state vector storing the moving average of variances.
• numFeatures - The number of features in the input vector.
• isTraining - Flag dictating whether to compute batch statistics or use running stats.
• momentum - Weight used to update the running statistics (e.g., 0.9).
• epsilon - A small value added to the variance to prevent division by zero.

Returns a Tensor<Vector> containing the normalized, scaled, and shifted result.

Implementation

Tensor<Vector> batchNorm1dMath(
  Tensor<Vector> x,
  Tensor<Vector> gamma,
  Tensor<Vector> beta,
  Vector runningMean,
  Vector runningVariance,
  int numFeatures,
  bool isTraining,
  double momentum,
  double epsilon,
) {
  Vector xHat = [];
  Vector currentMean = [];
  Vector currentVariance = [];

  if (isTraining) {
    for (int i = 0; i < numFeatures; i = i + 1) {
      currentMean.add(x.data[i]);
      currentVariance.add(0.0);
    }

    for (int i = 0; i < numFeatures; i = i + 1) {
      runningMean[i] =
          momentum * runningMean[i] + (1.0 - momentum) * currentMean[i];
      runningVariance[i] =
          momentum * runningVariance[i] + (1.0 - momentum) * currentVariance[i];
    }
  } else {
    for (int i = 0; i < numFeatures; i = i + 1) {
      currentMean.add(runningMean[i]);
      currentVariance.add(runningVariance[i]);
    }
  }

  Vector varianceToUse = isTraining ? runningVariance : currentVariance;
  Vector meanToUse = isTraining ? runningMean : currentMean;

  for (int i = 0; i < numFeatures; i = i + 1) {
    xHat.add((x.data[i] - meanToUse[i]) / sqrt(varianceToUse[i] + epsilon));
  }

  Vector outValue = [];
  for (int i = 0; i < numFeatures; i = i + 1) {
    outValue.add(gamma.data[i] * xHat[i] + beta.data[i]);
  }

  Tensor<Vector> out = Tensor<Vector>(outValue);
  out.creator = Node(
    [x, gamma, beta],
    () {
      for (int i = 0; i < numFeatures; i = i + 1) {
        double invStd = 1.0 / sqrt(varianceToUse[i] + epsilon);
        gamma.grad[i] = gamma.grad[i] + out.grad[i] * xHat[i];
        beta.grad[i] = beta.grad[i] + out.grad[i];
        x.grad[i] = x.grad[i] + out.grad[i] * gamma.data[i] * invStd;
      }
    },
    opName: 'batch_norm_1d',
    cost: numFeatures * 4,
  );

  return out;
}