BatchNorm folding

Overview

BatchNorm folding (also called BN merging or fusion) removes BatchNorm layers from a trained network before deployment by absorbing them into an adjacent convolution. The network then computes the same function with fewer layers and fewer memory passes. FasterNet uses it as the reason for choosing BatchNorm over LayerNorm, GroupNorm or InstanceNorm, and its benchmarks fold BN wherever that applies 2303.03667 §3.4, §4. The two quantization whitepapers use the same folding in their int8 pipelines. The 2018 whitepaper reports that folded weights have a wider per-filter range, which hurts per-layer quantization, and that BN needs a correction during quantization-aware training 1806.08342 §3.1.3, §3.2.2. The 2021 white paper treats folding as the default for per-tensor QAT and gives a second option for per-channel QAT 2106.08295 §4.2.

All variants rely on the same fact: once training is done, BN's running mean and variance are fixed, so BN is a per-channel scale and shift. The 2021 white paper states that BN is folded into the preceding linear layer 2106.08295 §2.3.1. The AdaRound paper also absorbs BN into adjacent weights before it quantizes, and runs its experiments on that basis 2004.10568 §5. A scale and shift applied right after a linear layer can be absorbed into that layer's weights and bias 1806.08342 eq. 19–21. LayerNorm, GroupNorm and InstanceNorm compute their statistics from each input at run time, so they cannot be folded this way 2303.03667 §3.4.

Variants

Conv → BN folding as a design constraint (FasterNet)

FasterNet places BN in three spots: after the 4×4 stride-4 embedding conv, after each 2×2 stride-2 merging conv, and after the middle 1×1 conv of each block, which is followed by the activation 2303.03667 §3.4, tab. 8. Every BN directly follows a conv, so all of them can be folded and the deployed network has no normalization layers. The block also has no norm after the PConv or the final 1×1 conv. The authors' reasons are feature diversity and latency 2303.03667 §3.4, and that choice also means there are fewer BN layers to fold.

The paper's only measured evidence is the T0 normalization ablation. With BN (folded at inference) and LayerNorm, top-1 accuracy is the same, and BN is faster on all three devices 2303.03667 tab. 5. This is a single comparison on one variant, not a sweep over norm types:

That is about 19% lower GPU throughput with LN, and about 16% and 11% more CPU and ARM latency. Card: Run, Don't Walk: Chasing Higher FLOPS for Faster Neural Networks.

Folding in a hybrid transformer (EfficientFormer)

EfficientFormer uses the same placement in its early stages. Its 4D blocks put BN after each 1×1 conv, and its stage-transition convs also use BN, so those layers fold at deploy time 2206.01191 §4.1. Its late 3D blocks put LayerNorm before the attention and before the MLP 2206.01191 eq. 5, and LN cannot be folded. The paper's latency analysis says LN adds roughly 10–20% of latency, and that convolution with BN is faster at a small accuracy cost 2206.01191 §3, obs. 3. Its ablation puts numbers on this. With the same architecture and no searched attention, BN gives 80.3% top-1 at 2.5 ms on iPhone 12, LN gives 80.3% at 3.7 ms, and GroupNorm gives 80.6% at 3.0 ms 2206.01191 tab. 3, V1, V2-LN, V2-GN. The ablation is one architecture on one device with CoreML, so read it as an indication, not a general rate. Card: EfficientFormer: Vision Transformers at MobileNet Speed.

Folding for int8 deployment (2018 quantization whitepaper)

For inference, the paper gives the folded weight as the original weight times γ/σ, and the folded bias as β − γμ/σ, where μ and σ are the moving mean and standard deviation 1806.08342 eq. 20–21. The 2021 white paper gives the same fold with the stability term ε inside the square root 2106.08295 §2.3.1, eq. 9–12. The folded conv then replaces the conv and BN pair.

Folding makes per-filter ranges uneven. Per-layer weight quantization fails mainly because the folded BN makes kernel ranges inside one layer differ a lot. Per-channel quantization gives much higher per-kernel SQNR, and folded weights have much longer tails than unfolded ones 1806.08342 §3.1.3, app. A, fig. 18–20. The 2021 white paper says per-channel weight scales are supported by many accelerators and are needed for some networks, but its card does not tie per-layer failure to folding 2106.08295 §2.2.3. Under 8-bit weight-and-activation post-training quantization, per-channel weights with per-layer activations stay within about 2 points of float top-1 on the nine networks in the table 1806.08342 tab. 3. Per-layer 8-bit weights alone collapse the MobileNets to between 0.1% and 0.4% top-1 1806.08342 tab. 2.

Folding during quantization-aware training

The 2018 whitepaper finds that folding during training needs more care. If the folded weights use batch statistics, the quantized weights change from batch to batch and accuracy suffers. Keeping moving averages throughout training removes that effect but destabilises training 1806.08342 §3.2.2, fig. 8. Its fix has two phases 1806.08342 eq. 22–31, fig. 9:

• Before freezing: scale the folded weights by a correction factor, the ratio of batch σ to moving σ, so the quantized weights always use the long-term σ. Undo that scale on the output and keep the bias from batch statistics, so the output matches ordinary BN.
• After the freeze step: stop undoing the scale, freeze the moving mean and variance, and add a bias correction term. This matches what inference computes.

On MobileNet-v1, this correction plus freezing gives higher and steadier evaluation accuracy than naive folding, batch renormalization, or moving-average weights without freezing 1806.08342 §4, fig. 14. The freeze step used for MobileNet-v1 is about 200000 steps 1806.08342 fig. 14. For MobileNet-v2 the card gives about 300000 steps in §3.2.2 and 400000 in §4 and the fig. 15 caption, so the value is not settled there 1806.08342 §3.2.2, §4, fig. 15.

The 2021 white paper uses a different default. It folds BN statically before training, which it calls simpler than the double forward pass of the 2018 whitepaper's method. In its runs, static folding is on par or better 2106.08295 §4.2, tab. 7. For per-channel quantization it can instead keep BN unfolded during training and merge BN's scale into the per-channel weight scales afterwards, which the paper also reports as on par or better 2106.08295 §4.2, eq. 41. The per-tensor QAT pipeline also needs cross-layer equalization for MobileNetV2 to train at all 2106.08295 §4.3, tab. 9.

Card: Quantizing deep convolutional networks for efficient inference: A whitepaper. Card: A White Paper on Neural Network Quantization.

Folding before post-training rounding (AdaRound)

AdaRound is post-training, so there is no fold inside a training loop. It takes the BN-absorbed network as its starting point 2004.10568 §5. Its contribution is the rounding choice for each weight, not the fold. With per-layer symmetric 4-bit weights and FP32 activations on ResNet18, nearest rounding gives 23.99% top-1 and AdaRound gives 68.60%, against 69.68% for FP32 2004.10568 tab. 6, tab. 7. These are author-reported numbers from 1024 calibration images and 10k iterations 2004.10568 §5. Bias correction on the same network gives 38.87% 2004.10568 tab. 8, so most of the gain is not from a bias fix alone. The method is covered on its own page, adaptive rounding.

The released code folds the soft rounding values back into the weights at export. The weights may therefore sit off the integer grid unless every rounding variable has reached 0 or 1 during optimisation. The card's "Paper vs code" notes list this as unchecked, so verify that exported weights lie on the grid before deploying. Card: Up or Down? Adaptive Rounding for Post-Training Quantization.

When to use which

Implementation sketch

Inference-time fold for a conv (or linear) layer followed by BN in inference mode. The formulas are from 1806.08342 eq. 20–21 and 2106.08295 eq. 9–12; the 2021 paper's folded formula includes the stability term ε.

# Trained conv: weight W [c_out, c_in, k, k], bias b [c_out] (zeros if absent)
# Trained BN on c_out channels: gamma, beta, running_mean mu, running_var var, eps

fold_conv_bn(W, b, gamma, beta, mu, var, eps):
    s      = gamma / sqrt(var + eps)          # per output channel, shape [c_out]
    W_fold = W * s[:, None, None, None]       # scale each output filter
    b_fold = (b - mu) * s + beta
    return W_fold, b_fold                     # replace conv+BN with conv(W_fold, b_fold)

deploy(model):
    model.eval()                              # use running stats, not batch stats
    for each (conv, bn) where bn's input is exactly conv's output:
        conv.W, conv.b = fold_conv_bn(...)
        replace bn with identity
    check: max |model_folded(x) - model(x)| is at float rounding level on a few inputs

For FasterNet, the pairs are embedding conv → BN, each merging conv → BN, and each block's middle 1×1 conv → BN, folded before the activation 2303.03667 §3.4, tab. 8.

For quantization-aware training, the fold happens inside the training forward pass. Compute batch statistics for the conv output, scale the folded weights by the ratio of batch σ to moving σ, fake-quantize those weights, and apply the bias and phase rules from the Overview's QAT variant 1806.08342 §3.2.2. The static-folding option of the 2021 paper skips the batch-statistics step and folds once before training 2106.08295 §4.2. Check the paper's pseudocode in the card before implementing, because the bias terms are split across several equations.

Contradictions

No conflicting numbers between the speed and int8 results. They measure different things: FasterNet's speed numbers are float-precision latency and throughput with BN folded 2303.03667 §4, while the quantization papers' results are int8 accuracy with per-layer and per-channel weight scales 1806.08342 tab. 2, tab. 3 2106.08295 tab. 6. A model that folds cleanly and runs fast in float can still lose accuracy under per-layer int8, so the two results should not be combined into one claim. The two whitepapers disagree on the best QAT treatment of folded BN, or at least do not agree on a single answer. The 2018 paper reports that keeping moving averages throughout training destabilises training, and that its correction plus freezing beats naive folding on MobileNet-v1 1806.08342 §3.2.2, §4, fig. 14. The 2021 paper uses static folding as its per-tensor default and reports it on par with or better than the double forward pass, on ResNet-18 and MobileNetV2 2106.08295 §4.2, tab. 7. The 2021 card does not say which statistics static folding uses during training, and the two papers test different networks and baselines, so this is not a direct contradiction. It does mean a reader cannot take either result as the one correct QAT choice. The 2018 paper's MobileNet-v2 freeze step is also not settled 1806.08342 §3.2.2, §4, fig. 15.

The 2018 paper ties per-layer weight failure partly to folded BN 1806.08342 §3.1.3. The AdaRound paper reports a per-layer 4-bit failure with nearest rounding on BN-absorbed networks 2004.10568 tab. 6, and attributes its fix to rounding, not to the fold. The two results use different bit-widths and methods and do not compare folded and unfolded weights, so neither settles how much of the loss comes from folding.

Papers

• Quantizing deep convolutional networks for efficient inference: A whitepaper: gives the fold formulas, shows that folding widens per-filter ranges and hurts per-layer int8, and proposes the BN correction and freezing for quantization-aware training.
• A White Paper on Neural Network Quantization: gives the fold formula with ε, and reports static folding and per-channel BN handling for QAT on several networks.
• Run, Don't Walk: Chasing Higher FLOPS for Faster Neural Networks: picks BN over LN in FasterNet because BN folds into adjacent convs; on T0 it matches LN's accuracy with higher throughput and lower latency.
• EfficientFormer: Vision Transformers at MobileNet Speed: puts BN in the convolution-style early stages and LayerNorm only in the late attention blocks; its ablation measures the latency cost of LN against BN and GN.
• Up or Down? Adaptive Rounding for Post-Training Quantization: assumes BN is absorbed before quantization and optimises per-layer weight rounding; reports 4-bit results on ResNet18 and others, against nearest rounding and bias correction.