Foundations

BatchNorm training, evaluation and functional running statistics

MediumRaw PyTorchAbout 20 min2 public examples plus private tests

scVIscANVItotalVI

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Normalize valid observations channel by channel in training or evaluation mode. Return output features, the statistics used for normalization, and newly computed running statistics without mutating the supplied ones.

Function signature

batch_normalization(
    x: torch.Tensor,
    valid: torch.Tensor,
    scale: torch.Tensor,
    bias: torch.Tensor,
    running_mean: torch.Tensor,
    running_var: torch.Tensor,
    eps: float = 0.001,
    momentum: float = 0.01,
    mode: str = 'train',
) -> dict[str, torch.Tensor]

Worked example

Three valid observations across two batch items

The first channel contains valid values [1, 3, 5], giving mean 3 and population variance $(4+0+4)/3=8/3$ . The second channel contains [2, 4, 8], giving mean $14/3$ . Only the running-variance update multiplies by $3/2$ ; the masked NaN row returns zeros.

import torch

x = torch.tensor([[[1.0, 2.0], [float('nan'), float('nan')]], [[3.0, 4.0], [5.0, 8.0]]], dtype=torch.float64, requires_grad=True)
valid = torch.tensor([[True, False], [True, True]], dtype=torch.bool)
scale = torch.tensor([0.8561484843061281, 0.919331295013565], dtype=torch.float64, requires_grad=True)
bias = torch.tensor([-0.5966353626151273, 0.18203648506130554], dtype=torch.float64, requires_grad=True)
running_mean = torch.tensor([-0.8566745932963743, 1.100604170903427], dtype=torch.float64, requires_grad=True)
running_var = torch.tensor([1.6474423668847282, 0.5150555940922124], dtype=torch.float64, requires_grad=True)
eps = 0.001
momentum = 0.25
mode = 'train'

result = batch_normalization(x=x, valid=valid, scale=scale, bias=bias, running_mean=running_mean, running_var=running_var, eps=eps, momentum=momentum, mode=mode)

Expected return

{
    'output': torch.tensor([[[-1.6450022775470605, -0.8006910434720396], [0.0, 0.0]], [[-0.5966353626151273, -0.06364539707203079], [0.45173155231680595, 1.4104458957279866]]], dtype=torch.float64),
    'mean': torch.tensor([3.0, 4.666666666666667], dtype=torch.float64),
    'variance': torch.tensor([2.6666666666666665, 6.222222222222221], dtype=torch.float64),
    'running_mean': torch.tensor([0.10749405502771925, 1.992119794844237], dtype=torch.float64),
    'running_var': torch.tensor([2.235581775163546, 2.719625028902492], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`x`torch.Tensor	$[B, N, C]$ Minibatch observations and channels.
`valid`torch.Tensor	$[B, N]$ `True` means include observation.
`scale`torch.Tensor	$[C]$ Floating learned affine multiplier $[C]$ .
`bias`torch.Tensor	$[C]$ Floating learned affine offset $[C]$ .
`running_mean`torch.Tensor	$[C]$ Floating supplied previous `mean` $[C]$ .
`running_var`torch.Tensor	$[C]$ Floating supplied previous nonnegative `variance` $[C]$ .
`eps`floatDefault `0.001`	Positive stabilizer inside sqrt(used channel variance + eps), during both train and eval normalization.
`momentum`floatDefault `0.01`	Value m in [0,1] for next_running=(1-m)old_running + mmeasured_statistic. The returned running variance uses the specified unbiased correction; the normalization variance remains population variance.
`mode`strDefault `'train'`	Choose "train" to measure statistics from valid observations and return running-statistic updates, or "eval" to normalize with the supplied running statistics.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

$x [B,N,C]$ has $B,C\ge 1$ and $N\ge 0$ ; valid $[B,N]$ is True for included observations. Reduce jointly across $B$ and $N$ , never $C$ . Every floating input shares dtype/device, learned scale/bias and running_mean/running_var have shape $[C]$ , and $\mathrm{running\_var}\ge 0$ .
Active $x$ is finite; invalid $x$ may be NaN/Inf. Replace invalid $x$ by zero before statistics. Config mode='train' (default) or 'eval', $\varepsilon>0$ default $10^{-3}$ , momentum $m$ in $[0,1]$ default $0.01$ (the supplied scvi-tools setting). Let $n=\sum(\mathrm{valid})$ .
Training: $\mu=\sum_{\mathrm{valid}} x/\operatorname{max}(n,1)$ , $v=\sum_{\mathrm{valid}} (x-\mu)^{2}/\operatorname{max}(n,1)$ , using POPULATION variance in $y=(x-\mu)/\sqrt{v+\varepsilon}\cdot \mathrm{scale}+\mathrm{bias}$ . Zero $y$ where invalid.
If $n>0$ , $\mathrm{next\_mean}=(1-m)\cdot \mathrm{running\_mean}+m\cdot \mu$ ; otherwise preserve running_mean. If $n>1$ , $\mathrm{next\_var}=(1-m)\cdot \mathrm{running\_var}+m\cdot v\cdot n/(n-1)$ , using the UNBIASED correction only for the running statistic; if $n\le 1$ preserve running_var.
Thus a singleton trains with zero variance and output bias; empty batches have output zero, $\mu=v=0$ and unchanged running states. These $n\le 1$ rules are explicit teaching extensions to ordinary BatchNorm training rejection.
Evaluation uses $\mu=\mathrm{running\_mean}$ and $v=\mathrm{running\_var}$ for $y$ and returns both states unchanged, regardless of valid count. Return output, the mean/variance USED for normalization, and the two next running states.
No running-state mutation or detach, cumulative-momentum counter, $\epsilon$ outside sqrt, affine-only shortcut, per-example statistics, or hidden batch state. The returned running-state calculation is differentiable for this functional exercise, even though production BatchNorm usually maintains nondifferentiable buffers.
Numerical domain: $\varepsilon$ must be finite and remain strictly positive when represented in the input dtype.
Active inputs and parameters must be bounded so the required arithmetic intermediates are finite and within that dtype's representable range: normalization sums/means, centered differences, squares, variance-plus- $\varepsilon$ , projections, residual/statistic updates, and intermediates for the graded first-order derivatives.
Finite inputs alone do not guarantee finite arithmetic; overflow-sized reductions or squares are outside the contract.

Why this matters

Training statistics describe the current observations; evaluation uses saved statistics. The variance used for normalization differs from the running-variance update.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def batch_normalization(
    x: torch.Tensor,  # (B, N, C)
    valid: torch.Tensor,  # (B, N)
    scale: torch.Tensor,  # (C,)
    bias: torch.Tensor,  # (C,)
    running_mean: torch.Tensor,  # (C,)
    running_var: torch.Tensor,  # (C,)
    eps: float = 0.001,
    momentum: float = 0.01,
    mode: str = 'train',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    output: (B, N, C).
    mean: (C,).
    variance: (C,).
    running_mean: (C,).
    running_var: (C,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement batch_normalization")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.