Domain Adaptation for Semantic Segmentation with Maximum Squares Loss

In plain words

Teaching a model to label every pixel of a street photo usually takes many photos with hand-drawn labels, and drawing those labels is slow and expensive. A common shortcut is to train on labelled images from a video game and then adapt the model to real photos that have no labels. The two look different, so a model trained only on the game images does poorly on real ones.

One popular way to adapt is entropy minimization: push the model to be sure of its answer on the unlabelled real photos. Its pull is strongest on the pixels the model already sees clearly, so the hard pixels get little learning signal. This paper replaces that penalty with the maximum squares loss, a sum of squared probabilities whose pull grows steadily instead of blowing up, so confident pixels stop dominating the update.

Two extras address other problems. Each photo gets class weights based on which classes the model predicts in that photo, which counters the imbalance between common and rare classes. A second output head is trained on confident labels built from both heads' predictions. The method needs no extra discriminator network. It changes the loss in an existing adaptation pipeline and was tested on game-to-city and city-to-city transfer.

TL;DR

For unsupervised domain adaptation of semantic segmentation (labelled synthetic source, unlabelled real target), replace the entropy target loss with a sum of squared probabilities, weighted per image by predicted class counts §3.2 §3.3. Optionally add a low-level head trained on confident self-produced labels §3.4. Authors report 46.4 mIoU on GTA5-to-Cityscapes with DeepLabv2 and ResNet-101, against 42.2 for their entropy-minimization run tab. 2. The official code is MIT-licensed and differs from the paper in several normalisations and training details; see Paper vs code.

Problem

Pixel-level labels are costly: the paper cites 90 minutes to annotate one Cityscapes image accurately §1. Synthetic data (GTA5, SYNTHIA) is cheaper to label, but a model trained on it does not transfer directly to real images because of the appearance gap §1. Unsupervised domain adaptation (UDA) uses the unlabelled target images to narrow that gap §1.

Entropy minimization, one semi-supervised approach used for UDA, gives a gradient that is largest on target pixels the model already predicts with high probability §1 §3.2. The paper assumes those pixels are the more accurate ones §1. The effect is that easy-to-transfer pixels dominate the update and hard ones are under-trained, which the paper calls probability imbalance §1 fig. 1. The target labels are also class-imbalanced, and there is no target label set from which to compute class frequencies §3.3.

Method

Setup §3.1. Source images carry pixel labels and target images do not. The objective is the source cross-entropy plus a target term weighted by λ_T eq. 1 eq. 2. Entropy minimization uses the Shannon entropy of the target prediction as that term eq. 3.

Maximum squares loss §3.2. The paper replaces the entropy term with the negative sum of squared target probabilities eq. 6. For the two-class case the gradient magnitude is |4p − 2| eq. 8, which grows linearly in p, while the entropy gradient |log p − log(1 − p)| eq. 5 grows without bound as p approaches 0 §3.2. The paper also reads the loss as maximising a Pearson χ² divergence to the uniform distribution eq. 10. In that reading, the loss pushes target features away from the decision boundary towards the source feature distribution §3.2.2 fig. 3.

Image-wise class weights §3.3. Class counts are taken per target image from the argmax prediction eq. 11 eq. 12. Each term is divided by 2 (N^c)^α N^(1−α), where N^c is that image's pixel count for class c, N the number of pixels and α a hyperparameter eq. 13. This is the loss the headline "MaxSquare+IW" result uses tab. 2.

Multi-level self-produced guidance §3.4. An extra low-level head is applied to the conv4 features of ResNet with an ASPP module §3.4. The guidance label of a pixel is the argmax of the averaged output of the two heads, kept only when either head gives that class probability above δ eq. 15. The low head is trained with cross-entropy on those labels, weighted by λ_low, and this term is added to the target loss eq. 14 eq. 16. The paper says this needs no extra pseudo-label pass §1.

Pseudocode. The paper has no algorithm box. The block below is assembled from the equations above, using the paper's own formulation. Helpers are defined after the main function.

# Adaptation objective with image-wise maximum squares  [eq. 1] [eq. 14]
# Input:  source batch (x_s, y_s); target batch x_t
# Output: scalar loss
def uda_loss(x_s, y_s, x_t, lambda_t, lambda_low, alpha, delta):
    # Source pixel cross-entropy on the final head  [eq. 2]
    p_s = softmax(final_head(x_s))                 # (B, C, H, W)
    loss_source = cross_entropy(p_s, y_s)

    # Target probabilities from both heads  [§3.4]
    p_t = softmax(final_head(x_t))                 # (B, C, H, W)
    p_t_low = softmax(low_head(x_t))               # (B, C, H, W)

    # Image-wise maximum squares on the final head  [eq. 13]
    loss_t_final = image_wise_max_squares(p_t, alpha)

    # Guidance from the averaged heads, for the low head  [eq. 15]
    guidance = self_produced_guidance(p_t, p_t_low, delta)
    # Cross-entropy on guided pixels only  [eq. 16]
    loss_t_low = cross_entropy(p_t_low, guidance, ignore=IGNORE)

    # Low-head term enters the target loss with weight lambda_low  [eq. 14]
    loss_t = loss_t_final + lambda_low * loss_t_low

    # Target term enters the total loss with weight lambda_t  [eq. 1]
    return loss_source + lambda_t * loss_t


def image_wise_max_squares(p, alpha):
    # Hard class per pixel: the argmax of the prediction  [eq. 11]
    labels = argmax(p, axis=class)                 # (B, H, W)
    # Pixels per class in each image, N^c  [eq. 12]
    counts = count_per_class(labels)               # (B, C)
    n_pixels = H * W                               # N
    # Denominator 2 (N^c)^alpha N^(1 - alpha); undefined when N^c = 0
    denom = 2 * counts ** alpha * n_pixels ** (1 - alpha)   # (B, C)
    # Negative sum of squared probabilities  [eq. 13]
    return -sum(p ** 2 / denom)                    # denom broadcast over H, W


def self_produced_guidance(p_final, p_low, delta):
    # Ensemble is the mean of the two heads  [eq. 15] [§3.4]
    p_ens = (p_final + p_low) / 2                  # (B, C, H, W)
    labels = argmax(p_ens, axis=class)             # (B, H, W)
    # Probability of the ensemble's class under each head  [eq. 15]
    conf_final = take(p_final, labels, axis=class) # (B, H, W)
    conf_low = take(p_low, labels, axis=class)     # (B, H, W)
    # Keep a pixel if either head is above delta  [eq. 15]
    keep = (conf_final > delta) | (conf_low > delta)
    # Pixels below delta get no target  [eq. 15]
    return where(keep, labels, IGNORE)             # IGNORE = -1

Key equations (numbers counted from the LaTeX source, see Pitfalls):

Total objective eq. 1:

\mathcal{L} = \mathcal{L}_{CE}(p_s, y_s) + \lambda_T \mathcal{L}_T(x_t) eq. 1

Source cross-entropy eq. 2 and entropy target term eq. 3:

\mathcal{L}_{CE}(p_s, y_s) = -\frac{1}{N} \sum_{n} \sum_{c} y_s^{n,c} \log p_s^{n,c} eq. 2

\mathcal{L}_T(x_t) = -\frac{1}{N} \sum_{n} \sum_{c} p_t^{n,c} \log p_t^{n,c} eq. 3

Binary entropy gradient eq. 5 and binary maximum squares eq. 7 with gradient eq. 8:

\left| \frac{dH}{dp} \right| = \left| \log p - \log(1-p) \right| eq. 5

MS(p) = -p^2 - (1-p)^2, \qquad \left| \frac{dMS}{dp} \right| = |4p - 2| eq. 7 eq. 8

Maximum squares loss eq. 6:

\mathcal{L}_T(x_t) = -\frac{1}{2N} \sum_{n=1}^{N} \sum_{c=1}^{C} \left( p_t^{n,c} \right)^2 eq. 6

Image-wise weighted version eq. 11–eq. 13:

m^{n,c^*} = \begin{cases} 1 & c^* = \arg\max_c p^{n,c} \\ 0 & \text{otherwise} \end{cases} \qquad N^c = \sum_n m^{n,c} eq. 11 eq. 12

\mathcal{L}_T(x_t) = -\sum_{n=1}^{N} \sum_{c=1}^{C} \frac{1}{2 (N^c)^{\alpha} N^{1-\alpha}} \left( p_t^{n,c} \right)^2 eq. 13

Two-head target loss eq. 14, guidance eq. 15 and low-head term eq. 16:

\mathcal{L}_T = \mathcal{L}_T^{final} + \lambda_{low} \mathcal{L}_T^{low} eq. 14

\bar{y}^{n,c^*} = \begin{cases} 1 & c^* = \arg\max_c p_{ens}^{n,c}, \; p_{final}^{n,c^*} > \delta \text{ or } p_{low}^{n,c^*} > \delta \\ 0 & \text{otherwise} \end{cases} eq. 15

\lambda_{low} \mathcal{L}_T^{low} = \lambda_{low} \mathcal{L}_{CE}(p_{low}, \bar{y}) eq. 16

Hyperparameters

Segmentation experiments use DeepLabv2 with a ResNet-101 backbone for every setting §4.2. Classification uses ResNet-50 §4.2.

Results

Author-reported. Setup: DeepLabv2 with ResNet-101 for segmentation, ResNet-50 for classification; trained on the stated source and target splits §4.1. Per-class IoU tables for the segmentation settings are in tab. 2, tab. 3 and tab. 6 and are not repeated here.

GTA5 to Cityscapes, mIoU over 19 classes tab. 2.

SYNTHIA to Cityscapes, mIoU over 16 classes and mIoU over 13 classes tab. 3.*

Cityscapes to NTHU cross-city, mIoU over 13 shared classes tab. 6.

Office-31 classification, average accuracy over six tasks tab. 1.

The paper calls these classification results uncompetitive with the state of the art, and its main comparison is with entropy minimization and DANN §4.3.

Ablation, GTA5 to Cityscapes mIoU tab. 4.

Parameter sensitivity, GTA5 to Cityscapes mIoU tab. 5.

Diagnostic on A to W fig. 5 §4.3. On the 30% most confident target samples ("top set"), EntMin and MaxSquare reach the same accuracy. On the 30% least confident ("bottom set"), MaxSquare is clearly higher. The paper reads this as evidence that its gain comes from hard samples.

Paper vs code

Code: official repository, ZJULearning/MaxSquareLoss, at commit 2fbc18fc4ee0a618bfb6a9c0e1318261c80541db, MIT licence. The paper links it in its abstract. Citations below are to that commit. Loss scale comparisons come from reading the code; nothing was run.

• Loss normalisation differs from the equations. Eq. 6 divides by 2N with N = H·W per image eq. 6. The code's MaxSquareloss takes the mean over all B·C·H·W entries and divides by 2 loss.py. For one image that is 1/C of eq. 6, with C = 19. The image-wise version uses weights 1/(max(hist^α · sum^(1−α), 1)) with no factor 2, and divides by B·C loss.py loss.py. For one image that is 2/C of eq. 13. The λ_T values in the paper are therefore tied to the code's scale, not to the literal equations.
• Per-pixel weights, not per-class. Eq. 13 indexes the weight by the class c of each term. The code takes the weight from each pixel's own predicted class and applies it to all of that pixel's class terms loss.py loss.py. The weight is detached from the graph loss.py.
• Source low-head term is not in the paper. The code adds λ_seg times the cross-entropy of the low head on source batches to the source loss solve_gta5.py. Eq. 1 and §3.4 contain no such term. λ_seg defaults to 0.1 train_source.py.
• λ_T in the released commands. The paper uses 0.1 for all segmentation experiments §4.2. The README MaxSquare+IW+Multi command uses --lambda_target 0.09 and a checkpoint directory named ST=0.09 README.md. The code's own default is 1 solve_gta5.py.
• Guidance uses per-head maxima and ignores unconfident pixels. Eq. 15 tests the probability of the ensemble's argmax class, p^{c*}, in each head. The code tests each head's maximum over classes solve_gta5.py. Unconfident pixels get the ignore label −1, which the cross-entropy skips solve_gta5.py, rather than a 0 target. The paper's "0 otherwise" is an indicator per class, so both forms exclude these pixels from the sum. They can differ in the mean's denominator, which the paper does not state.
• Two forward passes, so BatchNorm is per domain. Source and target batches are passed through the model separately solve_gta5.py solve_gta5.py. The paper does not say how BatchNorm statistics are computed. Concatenating the two domains into one batch would change them.
• Source pre-training length. The paper uses 70k steps §4.2. The README stops source pre-training at 80000 iterations (--iter_stop 80000) README.md.
• Adaptation length is set by code defaults. Iterations per epoch are (N + B) // B over the target loader cityscapes_Dataset.py, capped at ITER_MAX = 5000 train_source.py train_source.py (the cap does not bind for the 2,975-image Cityscapes loader). Total iterations are that count times epoch_each_round times round_num solve_gta5.py. The defaults are 2 epochs per round and 1 round solve_gta5.py, and the README uses round_num 5 README.md. For 2,975 target images this comes to about 29,760 iterations (derived, not stated in the paper).
• Classifier layers get 10× the learning rate. The optimiser has two parameter groups; the second is the classifier at 10× the base rate deeplab_multi.py train_source.py. The paper does not mention it.
• Input size is set by resizing, not stated in the paper. GTA5 is resized to 1280×720 and Cityscapes to 1024×512 by default, with random crop off train_source.py train_source.py. The README uses 1280×640 for the multi-level Cityscapes run README.md.
• Augmentation. The paper names random mirror and Gaussian blur §4.2. The target (Cityscapes) loader mirrors with probability 0.5 and blurs with probability 0.5, using a radius drawn uniformly from [0, 1) cityscapes_Dataset.py cityscapes_Dataset.py. The GTA5 loader file contains no Gaussian blur call.
• Test-time flip. The paper's body does not mention test-time augmentation; a commented-out sentence in the source mentions a flipped ensemble §4.2. The README evaluation command uses --flip True README.md, and the evaluator has a flip option evaluate.py.
• BatchNorm freezing. freeze_bn defaults to False train_source.py, and the README sets it to False. The paper does not say.

Pitfalls

• Eq. 13 is undefined when a class has no pixels. With α > 0, a class with N^c = 0 gives a division by zero for every term of that class. The literal equation therefore has no finite value whenever a class is absent from an image's argmax. The code avoids this by weighting each pixel with its own predicted class count, clamped to at least 1 loss.py. Either way the guard is an implementation choice that the paper does not state.
• The code's IW path only broadcasts for one image per step. As read, IW_MaxSquareloss multiplies a (B, C, H, W) tensor by a (B, H, W) weight loss.py loss.py. That product broadcasts only when B = 1, which the released single-GPU commands use train_source.py. Scaling the batch will break it unless the weights are reshaped. Not executed.
• Chi-squared reading is inconsistent. §3.2.2 says the divergence is to the uniform distribution U §3.2.2 eq. 10. The conclusion says "with the normal distribution" §5. Implement from eq. 10 and §3.2.2.
• Binary analysis drops the ½. Eq. 7 and eq. 8 are written without the ½ of eq. 6 eq. 6 eq. 7 eq. 8. The full loss's gradient is half the binary gradient shown in eq. 8, so a λ_T tuned against eq. 8 does not transfer to eq. 6 without a factor of 2.
• The paper's own statements overreach its tables. §4.4 says the multi-level head improves results "significantly", but Table 4 shows +0.9 for MaxSquare (44.3 to 45.2) and +1.2 with IW (45.2 to 46.4) tab. 4. §4.6 says MaxSquare+IW beats the other methods "by about 1 point". The margins over the strongest reported competitor per city are about +0.7 (Rome, AdaptSegNet), +1.1 (Rio, CBST), +0.6 (Tokyo, AdaptSegNet) and +0.3 (Taipei, CBST) tab. 6. Quote the table, not the prose.
• α "always" beats MaxSquare is not quite right. At α = 0 the IW weight is constant and the table reads 44.3, equal to MaxSquare tab. 5. Only α ≥ 0.1 is clearly above it.
• Wrong table reference in §4.3. The sentence that introduces the Office-31 results cites its table with a label that is defined on a figure (fig_004) in the source, so the reference resolves to Figure 4, not to the table §4.3. The Office-31 table is Table 1 tab. 1.
• Dataset size claim needs checking. §4.1 says Office-31 has 4,652 images in 13 categories. Office-31 is usually described with 31 categories; check the dataset before quoting the count.
• Adaptation length and input resolution are unstated. Neither the number of adaptation steps nor the training resolution is given in the paper §4.2. Use the code defaults listed under Paper vs code, or state your own choice.
• Hidden cost of the multi-level variant. §1 says the method needs no additional structure or discriminator and no extra pseudo-label computation §1. The multi-level variant adds a second ASPP head on conv4 §3.4 deeplab_multi.py. The no-extra-structure claim holds for the base method, not the full one.
• Equation and table numbers were counted from source. The PDF was not compiled. Numbers follow amsmath row order for align blocks and float order for tables and figures. If the compiled numbers differ, cite sections instead.
• Guidance threshold and λ_low interact with the loss scale. λ_low multiplies the low-head term inside λ_T eq. 14 eq. 1. The effective weight on the low head is therefore λ_T × λ_low, which is 0.01 with the paper's values. Changing the loss normalisation (see Paper vs code) changes that product too.

Lineage

• Builds on: entropy minimization for semi-supervised learning and its use in segmentation UDA (ADVENT) §2; class-balanced self-training (CBST) §2; self-produced guidance from weakly-supervised learning §3.4; the low-level loss weight from AdaptSegNet §3.4 §4.2. None of these has a card here.
• Compared against: AdaptSegNet, ADVENT (MinEnt and AdvEnt+MinEnt), CBST and CBST-SP on the segmentation benchmarks tab. 2 tab. 3 tab. 6; the NTHU cross-city baseline tab. 6.
• Related mechanism: the paper compares its gradient weighting to focal loss §4.4.