Distribution Alignment: A Unified Framework for Long-tail Visual Recognition

In plain words

Real image datasets are lopsided. A few classes have many photos and many classes have only a handful. A network trained on such data learns to guess the common classes, so its accuracy on rare classes is poor, even though a balanced test set rewards getting those rare classes right.

A common fix trains in two steps. First the network learns its features on the lopsided data. Then the final layer that makes the class decision is retrained or adjusted to favour rare classes. The authors found that the features from the first step are good, and the decision layer is the weak part. Retraining it on balanced data would give much better results.

This method adjusts that decision layer with a small amount of extra learning. Each class gets its own learned scale and offset on the score. A small learned gate decides, for each image, how much to trust the adjusted score versus the original one. Training also weights the rare classes more heavily. The same recipe is applied to classification, segmentation and object detection, with task-specific details listed below.

TL;DR

Distribution Alignment is a second-stage method for long-tailed recognition. It freezes the backbone and the original classifier, then learns per-class scale and offset parameters plus an input-dependent confidence gate. It trains with a cross-entropy whose class weights are inverse class frequency raised to a power rho. The authors report state-of-the-art results on image classification, semantic segmentation and LVIS detection and instance segmentation §1.

Problem

Long-tailed label distributions make the balanced evaluation metrics hard to reach: a model trained on imbalanced data is biased toward head classes §1 fig. 1. Two-stage methods learn features on the imbalanced data and then re-balance the classifier, but their adjustments rely on hand-designed rules and tuning §1.

The authors estimate the ceiling of the second stage: they keep the first-stage features fixed and retrain the classifier head on a balanced ImageNet training set, excluding the ImageNet-LT validation images §3.1. The gap between this ceiling and the existing decoupling methods is large §3.1 fig. 2. The values behind the ceiling are shown only in the figure, not stated in the text.

The method assumes a backbone trained on the imbalanced set with instance-balanced sampling in the joint stage §3.2, and a balanced test metric (mean per-class accuracy, mIoU, LVIS AP) §3.1 §4.1.

Method

Setup. The network is a feature extractor f followed by a classifier head h. The head outputs one score per class and the prediction is the argmax §3.1 eq. 1. The head is either linear (score = w_j · x) or cosine with a scale s (score = s · cosine(w_j, x)) §3.1 eq. 2 eq. 3. The scale s is taken from a prior work and may be learned or constant; the paper does not fix its value for classification §3.1.

Stage 1, joint learning. f and the original head h_o are trained together on the imbalanced data with instance-balanced sampling §3.2.

Stage 2, distribution calibration. f is frozen. Only the calibration parameters are trained §3.2 app. A.2.

1. Class-specific transform. Each original score z_j^o is mapped to s_j = alpha_j · z_j^o + beta_j, with one alpha and one beta per class, learned from data §3.2.1 eq. 4.
2. Input-dependent blend. The output is z_hat_j = sigma(x) · s_j + (1 − sigma(x)) · z_j^o, which equals (1 + sigma(x) · alpha_j) · z_j^o + sigma(x) · beta_j §3.2.1 eq. 5 eq. 6. The confidence sigma(x) is a learned linear layer on x followed by a non-linearity; the paper names a sigmoid as an example §3.2.1. Sigma controls how much calibration an example gets.
3. Prediction. Softmax over the calibrated scores z_hat §3.2.1 eq. 7.

Loss. The method minimises the expected KL divergence from a reference class distribution p_r to the model prediction p_m §3.2.2 eq. 8. Using the empirical average over the training set, this is a cross-entropy with the reference as target, up to a constant §3.2.2 eq. 9.

Generalized re-weighting. The reference is a class-weighted empirical distribution: p_r(y=c | x_i) = w_c · delta_c(y_i) §3.2.2 eq. 10. The weight is w_c proportional to (1/r_c)^rho, normalised over classes, where r_c is the training-set frequency of class c and rho is a scale hyper-parameter §3.2.2 eq. 11. With rho = 0 the weights are instance-balanced, with rho = 1 they are class-balanced §3.2.2.

Relation to earlier methods. The authors state that LWS and tau-normalisation learn a magnitude and drop the margin; Logit Adjustment and Deconfound use a hand-set margin with magnitude 1.0; DisAlign has both as learnable terms §3.3 tab. 2. Resampling is hard to apply to detection and segmentation, while the generalized re-weight can be used there §3.3.

Tasks. The same formula is used for image classification, semantic segmentation (per pixel), and detection and instance segmentation (per box, with the class weight on the foreground classes) §3.1 §4.2 §4.3 app. C.2.

# Distribution Alignment, stage 2  [§3.2, eq. 4-6]
# Inputs:  features (B, d), original_scores (B, K) from frozen h_o
# Learned: alpha (K,), beta (K,), confidence_layer (d -> 1)
def calibrated_scores(features, original_scores, alpha, beta,
                      confidence_layer):
    # Class-specific linear transform of the original scores  [eq. 4]
    transformed_scores = alpha * original_scores + beta      # (B, K)
    # Per-example confidence in (0, 1)  [§3.2.1]
    confidence = sigmoid(confidence_layer(features))         # (B, 1)
    # Blend: confident examples use the transformed scores  [eq. 5]
    return (confidence * transformed_scores
            + (1.0 - confidence) * original_scores)          # (B, K)

# Generalized re-weighting and loss  [§3.2.2, eq. 8-11]
# Inputs: class_counts (K,) training images per class, the r_c values
#         rho: scale hyper-parameter (see Hyperparameters)
def reweight_vector(class_counts, rho):
    # Inverse frequency raised to rho  [eq. 11]
    powered = (1.0 / class_counts) ** rho                    # (K,)
    # Normalise so the weights sum to one  [eq. 11]
    return powered / sum(powered)                            # (K,)

def alignment_loss(calibrated, labels, class_weights):
    # Weight of the true class of each example  [eq. 10]
    example_weights = class_weights[labels]                  # (B,)
    # Log-softmax over the calibrated scores  [eq. 7, eq. 9]
    log_probs = log_softmax(calibrated)                      # (B, K)
    picked = log_probs[arange(len(labels)), labels]          # (B,)
    # Batch mean of the weighted negative log-likelihood  [eq. 9]
    return mean(-example_weights * picked)

The pseudocode follows the paper's equations. The code's reduction differs; see Paper vs code.

Hyperparameters

Classification settings are from the appendix, which follows the decoupling protocol of Kang et al. §4.1 app. A.2. Segmentation and detection settings are from the appendix sections B.2 and C.2.

Results

Author-reported. Each table below uses the paper's own numbers. Setup: classification on ImageNet-LT (1000 classes), Places-LT (365 classes) and iNaturalist 2018 (8,142 species); metric is class-balanced top-1 accuracy, reported as average and as Many, Medium and Few subsets §4.1 app. A.1. Segmentation on ADE20K, metric mIoU and mean accuracy §4.2. Detection and instance segmentation on LVIS v0.5 and v1.0, COCO-style AP §4.3. No seeds or variance are reported in the active tables.

ImageNet-LT, top-1 average tab. 1. R-50 = ResNet-50, X-50 = ResNeXt-50. * = cosine classifier.

DisAlign many / medium / few (R-50): 59.9 / 49.9 / 31.8 (DisAlign), 61.3 / 52.2 / 31.4 (DisAlign*). For X-50: 61.5 / 50.7 / 33.1 and 62.7 / 52.1 / 31.4 tab. 1. The text reports gains of 2.5 (R-50) and 1.6 (X-50) over the best prior row §4.1.

ImageNet-LT, backbone sweep (top-1 average) app. A.

Component ablation, ImageNet-LT tab. 5. Backbone not named in the table caption; the text says ResNet-50 was used for ablations §4.1.

iNaturalist 2018, top-1 average tab. 3. Columns: ResNet-50 at 90 and 200 epochs, ResNet-152 at 90 and 200 epochs.

The text reports 69.5% with ResNet-50 at 90 epochs, 1.5 points above LDAM §4.1.

Places-LT, ResNet-152 tab. 4. Average / Many / Medium / Few.

Other rows in the same table: Focal Loss 34.6, Range Loss 35.1, OLTR 35.9, Feature Aug 36.4 (average) tab. 4.

ADE20K semantic segmentation, mIoU / mean accuracy tab. 6. Average over head, body and tail classes.

The text reports 47.8 mIoU as a new state of the art on ADE20K with DeepLabV3+ and ResNeSt-101 §4.2.

LVIS v0.5, Mask R-CNN FPN with ResNet-50, box AP / mask AP tab. 8.

Gains across backbones and with Cascade R-CNN: box AP +2.1 to +3.6 over the matching baseline tab. 7. LVIS v1.0 with Cascade R-CNN, box AP / mask AP: ResNet-50 Baseline* 26.5 / 23.5 vs DisAlign* 30.5 / 27.0; ResNet-101 28.9 / 25.6 vs 32.7 / 28.9; ResNeXt-101 30.7 / 27.3 vs 33.7 / 29.7 tab. 9.

Paper vs code

Repository: https://github.com/Megvii-BaseDetection/DisAlign at commit a2fc3500a108cb83e3942293a5675c97ab3a2c6e, Apache-2.0. The paper's footnote and abstract link it. The owner is the Megvii-BaseDetection organisation; two authors are listed with Megvii affiliations on the title page. The repo has three task folders: classification (a single ImageNet-LT ResNeXt-50 config pair), semantic segmentation (mmseg) and instance segmentation (cvpods). Many building blocks are in external packages that are not in this repo.

• The classification layer DisAlignLinear and the classification loss GRWCrossEntropyLoss come from the external cvpods package, not from this repo. Their formulas cannot be checked here net.py net.py net.py.
• The classification category frequencies are loaded from an external datasets folder net.py. The LVIS frequencies are loaded from num_shots_v0.5.npy, also external fast_rcnn.py.
• Stage-2 freezing matches the paper for classification: all parameters are frozen except alpha, beta and the confidence layer net.py. Segmentation does the same in the training script train.py.
• The segmentation head uses the same formula as the paper: confidence times (scale times score plus bias), plus (1 minus confidence) times the original score fcn_head_disalign.py. The confidence layer is a 1×1 ConvModule built from the head's norm_cfg and act_cfg before the sigmoid, so it has normalisation and an activation that the paper's linear-plus-non-linearity description does not name fcn_head_disalign.py fcn_head_disalign.py. The head defaults come from mmseg, which is not in the repo, so the exact layers are not verified.
• Segmentation rho is 0.2 in the config disalign_fcn_r50-d8_512x512_160k_ade20k.py. The paper states 0.3 for all experiments app. B.2.
• Segmentation stage-2 length is 4000 iterations in the config disalign_fcn_r50-d8_512x512_160k_ade20k.py. The paper states 8k app. B.2.
• LVIS stage-2 length is 4500 iterations in the config, with the LR step at 4000 config.py. The paper states 9k at LR 0.02 app. C.2.
• The LVIS weight vector gives the background class weight 1 and rescales the foreground weights to sum to the number of foreground classes fast_rcnn.py. This matches the paper's statement that the re-weight is used on foreground classes only app. C.2. The paper's equation 11 normalises to one instead.
• The LVIS loss uses PyTorch's weighted mean, which divides by the sum of the batch weights fast_rcnn.py. The paper's equation 9 reads as a plain batch mean. The two differ in scale whenever the batch weights change, so the effective step size changes too.
• The segmentation weights are rescaled to sum to the number of classes grw_cross_entropy_loss.py and read from the ratio column of objectInfo150.txt grw_cross_entropy_loss.py. The paper's weights sum to one eq. 11. The segmentation loss is inherited from mmseg, so its reduction is not verified here.
• The classification stage-2 config uses base LR 0.1, 10 epochs, cosine decay with no warmup, batch 256 and weight decay 5e-4 config.py. The paper gives the batch size and the 10-epoch length but not the stage-2 LR or the weight decay app. A.2.
• The repo's model-zoo table lists ResNeXt-50 DisAlign at 52.438 top-1 on the test split README.md, row for resx50 DisAlign. The paper's non-cosine ResNeXt-50 DisAlign row is 52.6 tab. 1. The same table lists the ResNeXt-50 baseline at 46.918, while the paper's baseline is 44.4 tab. 1. The repo does not say how the baseline was trained.
• Only the dot-product classification configs are in the repo; the cosine classification config (rho 1.5) is absent . The LVIS and segmentation cosine variants are in the repo.

Pitfalls

• Text and table disagree. The backbone ablation text gives 54.9% for ResNet-152 §4.1. The appendix table gives 54.8 for the cosine variant and 53.7 for the dot-product variant app. A. The text does not say which variant 54.9 belongs to.
• Equation numbers. The numbers in this card were counted in order from the LaTeX source, where the align blocks give one number per row. Check them against the PDF before citing one for a specific claim.
• Equation 10 is not normalised. The reference is written as w_c times delta, which sums to w_{y_i} for each example rather than to one eq. 10. The code treats it as a weighted cross-entropy. An implementer has to choose the reduction; see Paper vs code.
• Reduction and scale. The paper's loss is a batch mean of the weighted negative log-likelihood eq. 9. The LVIS code uses a weighted mean fast_rcnn.py. This changes the effective step size, so either copy the code's reduction or retune the learning rate.
• rho is task-specific and its selection is not stated. The paper sweeps rho on the ImageNet-LT validation split fig. 5. The values used (1.2 and 1.5 for classification, 0.3 for segmentation, 0.8 for LVIS) are given per task, and the paper does not say how they were chosen for the other tasks app. A.2 app. B.2 app. C.2.
• The calibration is cheap but its size is not stated. The paper says the method has far fewer free parameters than cRT §3.2.1, but gives no count. Per-class alpha and beta add 2K parameters for K classes. The confidence layer is per example for classification and per pixel for segmentation fcn_head_disalign.py.
• Stage-2 learning rate is missing for classification and segmentation. The paper says the LR is restarted app. A.2 and follows stage 1 app. B.2 without giving values. The code gives 0.1 with cosine for classification.
• Starting checkpoint for segmentation. The run script starts stage 2 from a released mmseg FCN checkpoint, passed with --load-from, and adds --disalign disalign_fcn_r50-d8_512x512_160k_ade20k.sh. The paper does not say which checkpoint its stage 2 started from.
• One gate for all classes. The confidence is a single value per example (per pixel for segmentation) and multiplies the whole class vector fcn_head_disalign.py. The code comment says only the foreground scores are adjusted, but the formula applies to every class.
• Results are single-run. No seeds or variance are reported in the active tables, so gaps of a few tenths of a point, such as the 0.2 difference between the repo README and the paper's ResNeXt-50 DisAlign row, cannot be judged as meaningful.
• Benchmarks use different backbones and training lengths. Many comparisons mix 90 and 200 epochs, cosine and dot-product heads, and different pre-training (ImageNet vs COCO for LVIS). Compare only within one row of a table.

Lineage

• Builds on: two-stage decoupling of representation and classifier learning, with the cRT, LWS and tau-Norm variants [Kang et al. 2019, ICLR, as cited in §3.2] tab. 2. No arXiv id is given in the source bibliography, so no link is added.
• Builds on: generalized re-weighting extends class-balanced re-weighting with a tunable exponent rho [Cui et al. 2019, CVPR, as cited in §3.2.2] eq. 11.
• Builds on: the cosine classifier scale s, cited in §3.1 as qi2018low.
• Compared against: Logit Adjustment [menon2020long], Deconfound [tang2020long], LDAM [cao2019learning], BBN [zhou2020bbn], OLTR [liu2019large], Focal Loss, Range Loss, Feature Aug, NCM, LWS and tau-Norm (classification); LDAM and BBN (iNaturalist); EQL, RFS, SimCal, LST and GroupSoftmax (LVIS). No arXiv ids are given in the source bibliography.
• Superseded by: not stated.