PaperlayerView exact revision on arXiv

Source-linked AI summary

Decorrelated Batch Normalization

Lei Huang, Dawei Yang, Bo Lang, Jia Deng

arXiv:1804.08450v1cs.CVcs.LGstat.ML

TL;DR

Batch Normalization standardizes activations but does not decorrelate them, motivating a method that can capture the benefits of full whitening during deep-network training. DBN whitens mini-batch activations with differentiable backpropagation, uses ZCA whitening to avoid PCA’s stochastic axis swapping, and reports improved optimization and generalization across architectures and datasets.

  • Problem

    Batch Normalization centers and scales activations but does not decorrelate them, leaving open whether full whitening can further improve training and generalization.

  • Method

    DBN whitens layer activations over mini-batches, differentiates through the whitening transform, uses ZCA whitening, and can whiten smaller activation groups.

  • Results

    DBN improves BN’s optimization efficiency and generalization, with consistent residual-network accuracy improvements on CIFAR-10, CIFAR-100, and ILSVRC-2012.

  • Takeaways & Limitations

    ZCA whitening is critical for successful DBN because it avoids the stochastic axis swapping associated with PCA whitening while retaining BN’s advantages.

  • Takeaways & Limitations

    The current un-optimized DBN implementation has high computational cost, so residual networks use only a subset of DBN modules, including one before the first residual block.

Abstract

from arXiv · show

Batch Normalization (BN) is capable of accelerating the training of deep models by centering and scaling activations within mini-batches. In this work, we propose Decorrelated Batch Normalization (DBN), which not just centers and scales activations but whitens them. We explore multiple whitening techniques, and find that PCA whitening causes a problem we call stochastic axis swapping, which is detrimental to learning. We show that ZCA whitening does not suffer from this problem, permitting successful learning. DBN retains the desirable qualities of BN and further improves BN's optimization efficiency and generalization ability. We design comprehensive experiments to show that DBN can improve the performance of BN on multilayer perceptrons and convolutional neural networks. Furthermore, we consistently improve the accuracy of residual networks on CIFAR-10, CIFAR-100, and ImageNet.

1. Introduction

Decorrelated Batch Normalization extends Batch Normalization by whitening mini-batch activations, addressing the limits of standardization while preserving differentiable training. The paper identifies PCA whitening’s stochastic axis swapping and finds that ZCA whitening and grouped decorrelation enable effective improvements in training and generalization.

  • Batch Normalization standardizes mini-batch activations to have zero mean and unit variance, accelerating deep network training.
  • Full whitening can improve conditioning when activations are highly correlated, whereas standardization does not decorrelate them.
  • DBN whitens each layer’s activations within a mini-batch and back-propagates through the whitening transform using differentiable eigen-decomposition.
  • PCA whitening can randomly permute layer neurons across batches through stochastic axis swapping, potentially preventing convergence.
  • ZCA whitening avoids stochastic axis swapping by rotating PCA-whitened activations back toward the original axes, where the benefits of decorrelation are observed.
  • Grouped whitening limits each operation to groups of size kG < d, reducing cost from O(d^2 max(m, d)) to O(mdkG).
  • Experiments report that DBN improves BN’s training speed and generalization, including consistent residual-network improvements on CIFAR-10, CIFAR-100, and ILSVRC-2012.

2. Related Work

DBN extends normalization research by decorrelating activations, complementing methods that address batch-size limitations, parameterized whitening, regularization, or implicit normalization.

  • Layer Normalization addresses BN’s need for reasonable batch sizes by normalizing each example using statistics from its layer.
  • DBN’s activation decorrelation is described as orthogonal to prior efforts targeting small batch sizes, recurrent networks, or related normalization changes.
  • Natural Neural Networks whiten activations with periodically estimated whitening matrices, but treating statistics as model parameters can cause instability.
  • DeCov and related regularizers encourage non-redundant or group-wise decorrelated representations but are not designed to speed training and often slow it.
  • Other approaches implicitly normalize activations through weight re-parameterization, Riemannian optimization, weight regularization, or scaling and bias choices.

3. Decorrelated Batch Normalization

DBN whitens layer activations within mini-batches, while addressing the instability and optimization difficulties caused by some whitening choices. The paper identifies stochastic axis swapping in PCA whitening and uses ZCA whitening and grouped decorrelation to enable effective training.

  • Whitening transformation: DBN defines a mini-batch whitening transformation that centers layer inputs and produces decorrelated, scaled activations.The transformation uses the batch mean and covariance, with a positive stability constant; the transformed covariance is the identity.
  • Stochastic axis swapping: Stochastic axis swapping occurs when batch-dependent PCA rotations permute activation dimensions across iterations, changing the data representation between batches.The phenomenon arises because PCA whitening rotates and scales activations using statistics that vary with the batch.
  • Stochastic axis swapping: PCA whitening can severely impair learning: on a 4-layer MNIST MLP, DBN-PCA behaves similarly to random guessing and performs significantly worse than the plain network.Different batches can reorder PCA axes when eigenvalue order changes, producing the swapping illustrated in Figure 1.
  • ZCA whitening: ZCA whitening rotates PCA-whitened activations back to the original axis system, minimizing distortion and avoiding the harmful PCA axis-swapping behavior.Experiments on MNIST MLPs report that DBN-ZCA improves training performance over both no whitening and DBN-PCA.
  • Optimization and inference: DBN performs differentiable back-propagation through the batch-dependent mean, covariance, and whitening transform, while maintaining running statistics for inference.The method derives gradients through eigen-decomposition and updates expected mean and whitening matrix estimates by running averages.
  • Optimization and inference: DBN extends BN with mini-batch decorrelation while retaining BN’s efficient-training properties and further targeting better dynamical isometry and conditioning.The paper presents these as benefits over standard BN.

4. Experiments

Experiments evaluate DBN on MLPs, CNNs, and residual networks, focusing on conditioning, convergence, generalization, depth, learning rate, and group size. Across these settings, DBN generally improves optimization and test performance relative to BN, while computational cost motivates applying it selectively in residual networks.

  • MLP experiments: DBN achieves the best conditioning among normalization methods and significantly speeds convergence on MLPs, despite being 2× slower than BN per iteration.The Yale-B experiments measure relative FIM conditioning and training loss over wall-clock time.
  • MLP experiments: Intermediate group sizes such as G16 converge faster than both G1 and G128, while smaller groups reduce eigendecomposition cost.The experiments motivate group whitening because full whitening estimates may be poor with mini-batch samples.
  • CNN experiments: As network depth increases, BN becomes more difficult to optimize than DBN on the S-plain architecture.The authors conjecture that approximate dynamical isometry alleviates the degradation problem.
  • CNN experiments: With a 4× higher learning rate of 0.4, DBN attains significantly better training accuracy than BN, which the authors relate to improved conditioning.This comparison uses the S-plain architecture on CIFAR-10.
  • Residual-network experiments: DBN improves wide residual networks on CIFAR-10 and CIFAR-100, reducing test error by 3.74% and 18.27%, respectively.The method also improves residual-network test errors across the evaluated networks, with larger gains for deeper models.

5. Conclusions

DBN extends Batch Normalization with whitening, but successful training requires ZCA rather than PCA whitening. The paper reports improved optimization efficiency and generalization abilities through decorrelated representations.

  • DBN extends Batch Normalization by whitening activations over mini-batch data.
  • PCA whitening can harm training through stochastic axis swapping, whereas ZCA whitening avoids this issue.The swapping effectively randomly permutes a layer’s neurons between batches.
  • DBN improves models’ optimization efficiency and generalization abilities while retaining Batch Normalization’s advantages.The paper attributes these properties to approximate dynamical isometry and improved Fisher Information Matrix conditioning.

A. Derivation for Back-propagation

The derivation uses column-vector matrix notation while representing gradient vectors as row vectors. It presents the forward pass before deriving the backward pass.

  • The derivation first presents the forward pass and then derives the backward pass.
  • The notation treats all vectors as column vectors.
  • Gradient vectors are represented as row vectors.

A.1. Forward Pass

The forward pass computes ZCA-whitened mini-batch inputs from the batch mean, covariance eigendecomposition, and whitening transformations. The intermediate variables correspond to PCA-whitened outputs, which the paper says perform poorly for deep networks.

  • ZCA whitening maps each mini-batch input through a whitening matrix and subsequent rotation.The displayed sequence is U = Λ^-1/2D^T, x̃_i = U(x_i − µ), and x̂_i = Dx̃_i.
  • The batch mean and covariance define the eigendecomposition used to construct the whitening transformation.D^TD = I, and Λ is diagonal with eigenvalues on its diagonal.
  • The auxiliary variable x̃_i is the output of PCA whitening, which the paper says hardly works for deep networks.

A.2. Back-propagation

The backward-pass derivation applies the chain rule and prior matrix-differentiation results to propagate gradients through the whitening transformation. It introduces a matrix K whose off-diagonal entries depend on eigenvalue differences.

  • The backward pass derives gradients using the chain rule and a prior result on differentiating the whitening transformation.
  • The matrix K is d×d, has zero diagonal, and uses K_ij = 1/(σ_i − σ_j) for i ≠ j.
  • The displayed derivative formulation is attributed to prior work on back-propagating through whitening transformations.A similar formulation was derived for learning an orthogonal weight matrix.

A.3. Derivation for Simplified Formulation

This section presents a simplified formulation and outlines its derivation from earlier equations. The derivation proceeds by combining specified equations to obtain intermediate results and the final formulation.

  • The section introduces a simplified formulation for more efficient computation.
  • The derivation of Eqn. A.14 is developed through a sequence of intermediate steps.
  • One intermediate result combines Eqns. A.4, A.5, A.8, and A.9.
  • Another intermediate result follows from Eqns. A.4 and A.10.
  • A further intermediate result is obtained from Eqn. A.11.

B. Computational Cost of DBN

DBN adds computational overhead whose magnitude depends on convolutional dimensions, group size, implementation, and hardware configuration. The current implementation is substantially less efficient than BN, but selective placement and single-GPU execution reduce the practical gap.

  • Adding DBN with group size K incurs overhead O(dKmhw + dK^2) for a convolutional layer.The relative overhead is negligible when K is small, such as 16.
  • 71ms versus 32ms: unoptimized full whitening costs more than an optimized 3 × 3 cuDNN convolution for the same input.These are averaged forward-plus-backward costs over 10 runs with a 64 × 32 × 32 input and batch size 64.
  • Small groups can cost more time than larger groups because the implementation whitens groups sequentially rather than in parallel.
  • 20%-40% versus 95%+: DBN has much lower average GPU utilization than BN in the current implementation.The authors identify substantial room for a more efficient implementation.
  • On a single GPU with batch size 32, the time-cost difference between BN and DBN is much smaller than in the multi-GPU ImageNet experiments.This suggests room for optimizing DBN for multi-GPU training and inference.
Loading 1804.08450v1…