Feature whitening

Overview

Feature whitening removes correlations between the channels of a feature map, so that channels stop carrying the same information. The paper that covers it here applies this to domain generalization: a segmentation network trained on one source domain of street scenes should degrade less on unseen ones 2103.15597 §1.

The paper's argument is that some channel correlations carry style, such as colour, texture and blur, while others carry the content needed to recognise objects. The aim is to remove the first kind and keep the second 2103.15597 §1 2103.15597 §3. Full whitening removes both kinds, and its exact form (eigen-decomposition) is slow and blocks gradients 2103.15597 §3.

The variants share one invariant. The penalty is computed on the covariance of each instance after per-channel standardization, and it is added to the training loss only. The network used at test time is unchanged 2103.15597 §4.1 2103.15597 fig. 3.

Variants

All three variants come from one paper, which compares them as ablations on the same setup 2103.15597 tab. 1 2103.15597 tab. 7 2103.15597 tab. 8.

Instance whitening, all entries (IW)

Penalises the mean absolute value of every off-diagonal covariance entry after standardization 2103.15597 §4.1 2103.15597 eq. 10. It treats all channel pairs alike. The paper's GTAV-trained ResNet-50 run scores 33.21 on Cityscapes, against 36.58 for ISW 2103.15597 tab. 1. Card: RobustNet.

Instance relaxed whitening (IRW)

Replaces the penalty with a hinge, so that entries may keep a margin of values before any cost is paid 2103.15597 §4.2 [arxiv:2103.15597 eq. 12]. The paper's best margin is δ = 1/64, which scores 33.57 on Cityscapes 2103.15597 tab. 8. It is an ablation, not the main method. Card: RobustNet.

Instance selective whitening (ISW)

The paper's main method. It penalises only the covariance entries that change most under photometric change, so style-like correlations are reduced and content-like ones are left alone 2103.15597 §4.3.

The mask is fixed once, after a warm-up:

1. Train for n = 5 epochs with the standardization layers but without the whitening loss 2103.15597 §4.3.
2. For each training image x and a photometrically altered copy τ(x), compare the covariance of x with the covariance of τ(x). The paper's τ is colour jitter and Gaussian blur 2103.15597 §4.3 2103.15597 app. A.3.
3. Average the per-entry variance over images, cluster the upper-triangle values with k-means (k = 3), and drop the lowest cluster. The remaining entries form the penalty set 2103.15597 §4.3 [arxiv:2103.15597 tab. 7].

The penalty is applied to three instance-normalization layers: after the input convolution and after the last block of ResNet stages 1 and 2 2103.15597 §4.4 2103.15597 app. A.3. Its weight is λ = 0.6 2103.15597 §4.4 2103.15597 eq. 18. The paper reports the GTAV-trained ResNet-50 scores as 36.58 on Cityscapes, 35.20 on BDD-100K and 40.33 on Mapillary, against a baseline of 28.95, 25.14 and 28.18 2103.15597 tab. 1. These are author-reported averages of three runs, with no spread given 2103.15597 §5.2.1. Card: RobustNet.

When to use which

The paper's data are tied to one kind of scene, street driving, and to photometric change as the source of style. The cited results do not show how the method behaves for other kinds of shift.

Implementation sketch

Python-flavoured pseudocode of the common core. The variants differ in the mask only.

# Covariance penalty for one instance-normalization layer  [arxiv:2103.15597 §4.1]
# Input:  feats (B, C, H, W); mask (C, C), 1 = penalise this entry
# Output: scalar loss

def whitening_penalty(feats, mask):
    # Standardize each channel of each image, no affine  [eq. 8]
    z = instance_norm_no_affine(feats).reshape(B, C, H * W)
    # Covariance of the standardized map  [eq. 9]
    cov = z @ z.transpose(1, 2) / (H * W)           # (B, C, C)
    # Mean absolute value over the masked entries  [eq. 17]
    # The paper does not say what the mean is over; this divides by the
    # number of masked entries.
    return (cov * mask).abs().sum(dim=(1, 2)).mean() / mask.sum()

# Select the ISW penalty set once, after a warm-up  [arxiv:2103.15597 §4.3]
# Input:  images; photometric transform tau; k = 3
# Output: mask (C, C), 1 for style-sensitive entries

def select_mask(images, tau, k):
    var = zeros(C, C)
    for x in images:
        # Covariance of the plain image and of its photometric copy  [eq. 13]
        a = covariance(instance_norm(feats(x)))
        b = covariance(instance_norm(feats(tau(x))))
        # Per-entry variance across the pair  [eq. 14-15]
        var += pair_variance(a, b)
    var = var / len(images)
    # Cluster the upper-triangle values, then drop the lowest cluster  [eq. 16]
    labels = kmeans_1d(var[upper], k)
    mask = zeros(C, C)
    mask[upper] = labels != lowest_cluster(labels)
    return mask

# Training with the ISW term  [arxiv:2103.15597 §4.3] [eq. 18]
# Input: loader; n = 5 warm-up epochs

masks = None
for epoch in range(num_epochs):
    for x, y in loader:
        loss = task_loss(network(x), y)
        if masks is not None:
            # Average over the three layers, weight 0.6  [eq. 18]
            penalties = [whitening_penalty(f, m)
                         for f, m in zip(network.in_feats(x), masks)]
            loss = loss + 0.6 * mean(penalties)
        step(loss)
    if epoch == n - 1:
        # Fix one mask per layer from the warm-up network  [§4.3]
        masks = [select_mask(loader, tau, k) for each layer]

Contradictions

• Mapillary claim. The text says ISW beats the other methods on BDD-100K, Mapillary and SYNTHIA when trained on Cityscapes 2103.15597 §5.2.1. In the Cityscapes table, ISW scores 58.64 on Mapillary, below IRW (59.20) and IW (58.90) 2103.15597 tab. 2.
• IBN-Net claim. The appendix says ISW beats IBN-Net for every k 2103.15597 app. A.2. On SYNTHIA, IBN-Net scores 27.90 2103.15597 tab. 1. Only k = 3 (28.30) is above it. k = 2, 5, 10 and 20 score 27.70, 27.52, 27.68 and 27.47 2103.15597 tab. 7.
• Warm-up length. The paper uses n = 5 epochs without the loss 2103.15597 §4.3. The released training script's epoch counter is zero-based, which gives six warm-up epochs train.py.
• Cluster count. The paper states k = 3 2103.15597 §4.3. The released Cityscapes scripts use k = 2 for ResNet-50 [code: scripts/train_r50os16_cty_isw.sh#L25], and other backbones use k and loss weights the paper does not give.
• Auxiliary loss. The main equation for the total loss has no auxiliary term 2103.15597 eq. 18. The appendix adds one with weight 0.4 2103.15597 app. A.3.
• Unreported spread. Tables average three runs with no spread, so differences of a few tenths of mIoU, such as ISW against IRW on G, cannot be judged as significant from the paper 2103.15597 §5.2.1.
• Training cost. The paper's "no additional cost" claim refers to inference, parameters and FLOPs 2103.15597 tab. 6. Training time is not measured 2103.15597 §5.2.3.

Papers

• RobustNet: Improving Domain Generalization in Urban-Scene Segmentation via Instance Selective Whitening: introduces ISW, with IW and IRW as ablations, for single-source domain generalization in segmentation.