RobustNet: Improving Domain Generalization in Urban-Scene Segmentation via Instance Selective Whitening

In plain words

A segmentation model trained on one set of street scenes often fails on scenes that look different: other lighting, rain, fog or a new layout. Domain generalization means making the model robust to such scenes without having seen them during training.

The paper looks at how the channels inside a network's feature maps move together. Some of that movement carries the look of a scene, such as color, brightness and blur. Other movement carries what is actually in the scene, such as road, cars and sidewalks. Removing all of it would hurt recognition.

The method finds the movements that change most when the colors or blur of a training image are altered, and penalizes only those during training. The rest is left alone. The extra work happens only during training; the network used at test time has the same cost as before.

TL;DR

• Instance selective whitening (ISW) is a training loss for domain generalization in urban-scene segmentation.
• It penalizes only the feature-covariance entries that move under photometric change, so style-like correlations are reduced and content-like ones are kept.
• It adds no parameters and no inference cost in the authors' measurements.
• Gains are author-reported, averaged over three runs, and some text claims are not supported by the paper's own tables (see Pitfalls).

Problem

Deep networks trained on a source dataset often degrade on target images with a different distribution, such as other illumination, weather or cities §1. Domain adaptation reduces this gap but needs target-domain images, which cannot be collected in advance for an open-ended real world §1§2. Domain generalization (DG) needs no target data during training. Many DG methods, however, need several labelled source domains, which are costly to collect §1.

A normalization-based route needs only one source domain. Instance normalization (IN) standardizes each channel of each image, but it ignores correlations between channels. Those correlations are reported to carry style such as color and texture §1§2. Full whitening removes both style and content correlations, and its exact form (eigen-decomposition) is slow and blocks gradients §3. The paper's problem is to remove only the style part §1§3.

Method

Input: a DeepLabV3+ segmentation network with ResNet-50 (output stride 16) §5.1.1tab. 1. The loss is used only during training, and it is attached to three instance-normalization layers §4.4app. A.3.

1. Instance standardization. For each image and channel, subtract the spatial mean and divide by the spatial standard deviation over the H×W map, with no learned affine parameters §4.1eq. 8§4.4. This sets the diagonal of the covariance close to one, so the loss only has to shrink the off-diagonal terms. Plain whitening losses, in contrast, push diagonal and off-diagonal terms in conflicting directions §4.1.
2. Covariance of the standardized map. Σ_s = (1/HW) X_s X_sᵀ, computed per instance eq. 9§4.1.
3. Instance whitening (IW) loss. The mean absolute value of the off-diagonal entries of Σ_s, with a strictly upper-triangular mask M eq. 10eq. 11§4.1. This is the all-entries variant, used as an ablation.
4. Margin relaxation (IRW), ablation only. The IW loss is replaced by max(E‖Σ_s ⊙ M‖₁ − δ, 0), so the covariance may keep a margin of values eq. 12§4.2. The best margin in the paper is δ = 1/64 tab. 8.
5. Selecting style-sensitive entries. This is done once, after a warm-up §4.3:
   • Build pairs of each training image x and a photometrically altered copy τ(x). The paper's transform is color jitter and Gaussian blur §4.3app. A.3.
   • Train with the standardization layers only, without the whitening loss, for n = 5 initial epochs §4.3.
   • For each pair, compute Σ_s of x and of τ(x). The variance matrix V averages, over images, the variance of each entry across the pair eq. 13eq. 14eq. 15§4.3.
   • Cluster the strict upper-triangular entries of V into k groups with k-means, with k = 3 §4.3tab. 7.
   • Call the lowest cluster insensitive and the rest style-sensitive (G_high). Set m = 1, meaning only the lowest group is dropped §4.3.
   • The mask M̃ is 1 for entries in G_high and 0 otherwise eq. 16§4.3.
6. ISW loss. L_ISW = E‖Σ_s ⊙ M̃‖₁, with the same masked form as step 3 but with the style-sensitive mask eq. 17§4.3. Training continues with it for the remaining epochs §4.3.
7. Total loss. L_total = L_task + λ · (1/L) Σ_i L_ISW^i, with λ = 0.6 and L = 3 layers eq. 18§4.4. The supplement's total loss also includes the auxiliary segmentation loss with weight 0.4, so the main-text equation is incomplete for the reproduced setup app. A.3.
8. Placement. Three IN layers are used: the input convolution (where IN with ISW replaces BN), and the last residual block of ResNet stage 1 and stage 2 §4.4app. A.3. The supplement says these early positions follow IBN-Net's observation that earlier layers carry style app. A.3.

# Instance selective whitening loss for one IN layer  [§4.1] [eq. 8-9]
# Input:  feats (B, C, H, W) from one IN layer; mask (C, C), 1 = penalize
# Output: scalar loss

def isw_layer_loss(feats, mask):
    # Standardize each channel of each image, no affine  [eq. 8]
    standardized = instance_norm(feats)                 # (B, C, H, W)
    flat = standardized.reshape(B, C, H * W)            # (B, C, HW)
    # Covariance of the standardized map, scaled by 1/HW  [eq. 9]
    cov = flat @ flat.transpose(1, 2) / (H * W)         # (B, C, C)
    # Penalize only masked entries  [eq. 17]
    penalty = (cov * mask).abs().sum(dim=(1, 2))        # (B,)
    # Divide by the masked count, as the code does
    # [code: network/instance_whitening.py#L23-L25]
    penalty = penalty / mask.sum()
    return penalty.mean()

# Style-sensitive mask for one IN layer  [§4.3] [eq. 13-16]
# Input:  images (N tensors); tau (photometric transform); k (default 3)
# Output: mask (C, C), 1 for style-sensitive entries

def select_mask(images, tau, k):
    # Variance of each entry over (x, tau(x)) pairs, averaged over images
    variance = zeros(C, C)
    for image in images:
        cov_plain = covariance(instance_norm(feats(image)))
        cov_photo = covariance(instance_norm(feats(tau(image))))
        variance += pair_variance(cov_plain, cov_photo)  # [eq. 13-15]
    variance = variance / len(images)
    # Cluster the strict upper-triangle values into k groups  [§4.3]
    upper = variance[upper_triangle_indices]
    labels = kmeans_1d(upper, k)                 # label 0 = lowest group
    # Keep every group except the lowest (m = 1)  [§4.3] [tab. 7]
    mask = zeros(C, C)
    mask[upper_triangle_indices] = labels != 0   # [eq. 16]
    return mask

# Training with the ISW term  [§4.3] [eq. 18] [app. A.3]
# Input:  loader of (images, labels); n = 5 warm-up epochs

def train(network, loader, num_epochs, n=5):
    masks = None
    for epoch in range(num_epochs):
        for images, labels in loader:
            outputs = network(images)  # main, aux, per-layer IN features
            # Main and auxiliary cross-entropy  [app. A.3]
            loss = ce(outputs.main, labels) + 0.4 * ce(outputs.aux, labels)
            if masks is not None:
                # Average over the three layers, weight 0.6  [eq. 18]
                terms = [isw_layer_loss(f, m)
                         for f, m in zip(outputs.feats, masks)]
                loss = loss + 0.6 * mean(terms)
            optimizer_step(loss)
        if epoch == n - 1:
            # Fix one mask per IN layer from the warm-up network  [§4.3]
            masks = select_masks(loader, tau, k)

Inference uses the network without these terms; the paper says the whitening operations are used only during training fig. 3.

Hyperparameters

Results

Author-reported. Setup: DeepLabV3+ with ResNet-50 (output stride 16), trained on GTAV unless noted. Columns C, B, M, S and G are the evaluation sets, scored as mIoU in percent §5.1.1tab. 1. The † models are the authors' re-implementations on their baseline §5.1.1. The text does not define the Baseline row beyond its label. Results are averages over three runs; no spread is reported §5.2.1. The paper states that ISW outperforms other methods on every target set except the source §5.2.1.

GTAV-trained, ResNet-50 tab. 1

Cityscapes-trained, ResNet-50 tab. 2

Other backbones, GTAV-trained tab. 3

Multiple synthetic sources (GTAV + SYNTHIA), ResNet-50 tab. 4

Comparison with DRPC and IBN-Net, GTAV-trained tab. 5

The table lists each block's own baseline with its gain. The DRPC and IBN-Net rows are taken from those papers §5.2.2. IBN-Net reports no B or M values.

Computational cost on 2048×1024, NVIDIA A100, averaged over 500 trials tab. 6

Ablations on GTAV-trained, ResNet-50

Number of clusters k tab. 7:

Margin δ in IRW tab. 8:

Loss weight tab. 9 (called γ there):

Qualitatively, the paper shows covariance matrices from the first and second convolution layers of an IBN-Net model and of the ISW model. The ISW matrices are mostly whitened, with a small number of large entries left fig. 6. Reconstructions from ISW features keep object content, while illumination and color are suppressed fig. 7. The paper also shows that the most sensitive covariance tends to track illumination and the most insensitive one tends to track scene complexity, on BDD-100K app. A.4.

Paper vs code

Official code: https://github.com/shachoi/RobustNet, licensed BSD-3-Clause, copyright Sungha Choi, at commit 1ff5c047d7abaf9ab1a7cca23aff812b2f09de15. The repository is linked from the abstract. Its owner matches the first author, so no further search was needed. Covariance, mask and loss code is in network/instance_whitening.py, network/cov_settings.py and network/deepv3.py; training is in train.py.

The following differences were checked against the paper:

• Covariance normalization. The paper normalizes by 1/HW eq. 3eq. 9. The code divides by HW−1 and adds 1e-5 to the diagonal deepv3.pyinstance_whitening.py. The effect is small but not zero for small maps.
• Loss normalization. Eq. 17 writes E‖Σ_s ⊙ M̃‖₁ without saying what E averages over. The code sums the masked absolute entries, divides by the count of masked entries, then averages over the batch and over the three layers instance_whitening.pydeepv3.py. The loss values are therefore per-entry means, not sums.
• Variance statistic. The paper averages, over images, the per-pair variance eq. 13-15. The code computes the variance over a batch of two maps (one plain, one photometric) with the unbiased estimator, then averages over calls deepv3.pytrain.py[code: datasets/init.py#L268]. This ranks entries the same way up to a constant factor. This equivalence is derived from the code, not stated by the authors.
• Trials. The code runs 10 trials, each over up to 500 random pairs, builds a mask per trial and intersects the masks train.pytrain.pycov_settings.py. The paper describes one pass.
• Warm-up timing. The paper uses n = 5 epochs without the loss §4.3. The code's epoch counter starts at 0; the loss is off for indices 0 to 5 and the mask is computed after index 5 train.pytrain.pytrain.pytrain.py. That is six warm-up epochs, not five.
• Clusters per setting. The paper states k = 3 §4.3. The repo's scripts use k = 2 for the Cityscapes ResNet-50 run train_r50os16_cty_isw.sh, k = 2 for ResNet-101 train_r101os8_cty_isw.sh, k = 2 for ShuffleNetV2 with weight 0.8 and batch multiplier 8 train_shuffle_gtav_isw.shtrain_shuffle_gtav_isw.shtrain_shuffle_gtav_isw.sh, and k = 20 with weight 1.0 for MobileNetV2 train_mobile_gtav_isw.shtrain_mobile_gtav_isw.sh. The paper does not state these per-setting values.
• Auxiliary loss. Main eq. 18 has no auxiliary term. The code and the supplement add 0.4 times the auxiliary cross-entropy train.pyapp. A.3.
• Dynamic masks. The code has a --dynamic option that recomputes masks periodically, but the released scripts do not use it train.pytrain.py.
• Weight decay. The optimizer hard-codes 5e-4 and ignores the command-line argument optimizer.pytrain.py.
• Masking on the full matrix. The code clusters all C×C variance values, with the lower triangle and diagonal set to zero. Those zeros fall in the lowest cluster, so the outcome matches clustering only the upper triangle with m = 1 cov_settings.py.
• IRW margin. The code sets the margin to the upper-triangle count divided by the denominator and divides the hinge by the same count cov_settings.pyinstance_whitening.py. So δ = 1/64 is a per-entry mean threshold. The paper does not give units for δ.
• Photometric blur. The paper does not give the Gaussian sigma. The code samples it in [0.15, 1.30] transforms.py. The color-jitter values match the supplement [code: datasets/init.py#L108].
• Multi-source. The paper reports 110K iterations for multi-source training §5.1.1. No multi-source script is in the repo, so that setting is not reproduced by a script.
• Layer placement. Consistent with the paper: IN without affine, applied after the input convolution and after the last block of stages 1 and 2 Resnet.pytrain_r50os16_gtav_isw.sh.

Pitfalls

• Mapillary claim. The paper says ISW beats the others on BDD-100K, Mapillary and SYNTHIA when trained on Cityscapes §5.2.1. On Mapillary, ISW (58.64) is below IRW (59.20) and IW (58.90) in tab. 2.
• IBN-Net claim. The supplement says ISW beats IBN-Net for every k app. A.2. On SYNTHIA, IBN-Net scores 27.90 tab. 1, and only k = 3 (28.30) exceeds it; k = 2, 5, 10 and 20 score 27.70, 27.52, 27.68 and 27.47 tab. 7.
• Compute. The "no additional computational cost" claim refers to inference, parameters and FLOPs §5.2.3tab. 6fig. 3. Training adds the warm-up pass, the variance statistics over 10 trials of up to 500 pairs, and the loss on each of three layers. The paper does not measure training time.
• Variance across runs. Tables average three runs but report no spread §5.2.1. Differences of a few tenths of mIoU, for example between ISW and IRW on G, cannot be judged as significant from the paper. Seed and random sampling for the mask are not stated.
• Mask sensitivity. The mask depends on which images are drawn for the statistics. The paper does not give N, the photometric sampling, or the number of trials. The code uses 10 trials and intersects them train.py.
• Equation and notation drift. λ (main text), γ (Table 9) and γ₁, γ₂ (supplement) name the same weights. The supplement's total loss has an auxiliary term that the main-text equation lacks eq. 18app. A.3.
• IRW δ. The unit of δ is not stated in the text. Table 8 uses fractions, and the code uses a per-entry mean cov_settings.py.
• Epsilon and normalization. Small details such as 1/HW versus 1/(HW−1), the 1e-5 diagonal term and the masked-count division change the loss scale. Keep the code's convention when reproducing reported numbers deepv3.py.
• Baseline definition. The text does not define "Baseline" beyond the table label. The repo has scripts/train_r50os16_gtav_base.sh for it, which was not compared line by line against the paper.
• Table and figure numbers. These were counted from the LaTeX source in order, not from the compiled PDF.

Lineage

• Builds on: IBN-Net (instance normalization inside a ResNet for domain generalization, named in §2); the group-wise deep whitening loss of GDWCT, which is the source of the masked whitening loss §3.
• Compared against: IBN-Net, Switchable Whitening (SW), IterNorm, DRPC and the authors' own IW and IRW variants tab. 1tab. 5.
• Related: domain adaptation methods for segmentation. The supplement compares only their reported gains over their baselines, since DA needs target data that DG does not use app. A.1.