Long-Tailed Recognition via Weight Balancing

In plain words

Train a classifier on data where a few classes have many examples and most have only a few, and it learns to favour the common classes. The authors saw that the final layer of such a model gives the common classes larger weights, so their scores win more often. Their fix leaves the data and the loss alone and keeps those weights in check with standard tools. Weight decay penalises large weights, and a MaxNorm limit caps them at a set size.

Training runs in two stages. First the feature extractor is trained normally, with weight decay tuned for accuracy. Then only the final classifier is retrained, using a loss that gives rare classes more influence, plus the same penalty and cap. The authors report that this simple recipe beats several more complicated methods on standard benchmarks, and they offer it as a baseline for the field.

TL;DR

On CIFAR100-LT with imbalance factor 100, the full recipe reaches 53.35% class-averaged accuracy, against 38.38% for plain cross-entropy training tab. 2. Tuning weight decay alone in stage 1 gives 46.08% tab. 2. The released code reproduces only the CIFAR100-LT ablation, and it differs from the paper in several details listed under Paper vs code. Relevant if you train long-tailed image classifiers and can afford a short second pass over the last layer.

Problem

Real data is long-tailed: a few classes are frequent and most are rare. A network trained naively on such data scores well on the frequent classes and poorly on the rare ones §1. The authors measure the imbalance as the ratio of the largest to the smallest class count, and score models by accuracy averaged over classes §3 Preliminaries. Existing remedies rebalance the data, reweight losses or gradients, transfer features from common to rare classes, or ensemble experts and use heavy augmentation or self-supervised pretraining §1 §2.

The paper's starting point is an observation about the classifier. After naive training, the per-class weight norms of the last layer are larger for common classes, because they receive more gradient signal §1 fig. 1. The paper assumes the target is class-averaged accuracy on a class-balanced test set §4.1.

Method

The pipeline has two stages §3.3. Stage 1 learns features with cross-entropy and tuned weight decay. Stage 2 freezes the features and retrains the classifier with a class-balanced loss, weight decay and MaxNorm §3.3.

Weight-balancing techniques §3.1. All three act on the per-class weight vectors \theta_k of the classifier.

1. L2-normalisation constrains each \theta_k to unit norm eq. 2. The paper also applies it post hoc, by dividing each trained class vector by its norm, and reports that this helps accuracy but can hurt it on iNaturalist §3.1 tab. 3.
2. Weight decay adds a penalty \lambda \sum_k \lVert \theta_k \rVert_2^2 to the cross-entropy objective eq. 3. The paper says it penalises large weights more and keeps small ones small §3.1 fig. 2c.
3. MaxNorm keeps each \theta_k inside a ball of radius \delta eq. 4. Training uses projected gradient descent, which rescales any row whose norm exceeds \delta back onto the ball eq. 5 §3.1. Small weights may still grow inside the ball fig. 2d.

The paper also discusses a Lagrangian view in which weight decay is one step of solving the MaxNorm problem eq. 6 §3.2. It reports that weight decay and MaxNorm together work better than either alone §3.2.

Stage 1 (feature learning). Train the network with cross-entropy and a single weight decay \lambda applied to all network parameters §3.3 §4.2. The paper says tuning \lambda is enough for a good backbone, and that MaxNorm in this stage needs per-layer thresholds that are hard to tune §3.3.

Stage 2 (classifier learning). Train only the classifier on the frozen features, with a class-balanced loss §3.3. The class-balanced loss weights each class by the inverse of its effective number of samples, (1-\beta)/(1-\beta^{n_c}), rescaled to sum to the number of classes §3.3 §4.2. Weight decay and MaxNorm are added on top §3.3. The paper tunes these hyperparameters with Bayesian optimisation §3.3.

Pseudocode. Stage 1 and stage 2 with the MaxNorm projection, as implemented in the released code (see Paper vs code).

# Stage 1: feature learning with tuned weight decay  [§3.3] [§4.1]
# Input:  train_set of (image, label); weight decay lam; epochs E
# Output: backbone and linear head, trained with plain cross-entropy

def train_stage_one(train_set, lam, epochs):
    # Backbone plus a linear head with one output per class  [§4.1]
    model = build_resnet_with_linear_head()
    # SGD, momentum 0.9, base learning rate 0.01  [§4.1] [code: demo1#L222]
    optimizer = sgd(model.parameters(), lr=0.01, momentum=0.9,
                    weight_decay=lam)
    # Cosine decay of the learning rate from 0.01 to 0  [§4.1]
    scheduler = cosine_decay(optimizer, epochs, eta_min=0.0)
    for epoch in range(epochs):
        # Random crop with 4-pixel padding and horizontal flip  [§4.1]
        for images, labels in train_set:
            logits = model(images)
            loss = cross_entropy(logits, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        scheduler.step()
    return model

# Stage 2: classifier relearning with class-balanced loss, weight decay
# and MaxNorm  [§3.3] [code: demo2#L316-L356]
# Input:  stage-1 model, class counts n_c, weight decay lam, thresh
# Output: model whose classifier is retrained

def train_stage_two(model, train_set, class_counts, lam, thresh):
    # Freeze all layers except the classifier weight and bias
    # [code: demo2#L324-L329]
    trainable = [model.fc.weight, model.fc.bias]
    # Class-balanced weights from effective numbers  [§3.3]
    class_weights = class_balanced_weights(class_counts, beta=0.9999)
    # One MaxNorm radius for the whole classifier, from its initial norms
    radius = radius_from_row_norms(model.fc.weight, thresh)
    optimizer = sgd(trainable, lr=0.01, momentum=0.9, weight_decay=lam)
    for epoch in range(10):  # [code: demo2#L336]
        for images, labels in train_set:
            logits = model(images)
            loss = class_balanced_loss(logits, labels, class_weights)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        # Projection once per epoch, not per batch  [code: trainval#L140-L142]
        model.fc.weight.data = project_max_norm(model.fc.weight.data, radius)
    return model


def class_balanced_weights(class_counts, beta):
    # Effective number of samples per class, as in the class-balanced loss
    effective = 1 - beta ** class_counts            # (K,)
    weights = (1 - beta) / effective                # (K,)
    # Rescale so the weights sum to the number of classes  [code: cb#L73-L75]
    return weights / weights.sum() * len(class_counts)


def radius_from_row_norms(weight, thresh):
    # L2 norm of each class row of the classifier  [code: regularizers#L21-L27]
    norms = row_l2_norms(weight)                    # (K,)
    # Sit a fraction thresh of the way from the smallest to the largest norm
    return norms.min() + thresh * (norms.max() - norms.min())


def project_max_norm(weight, radius):
    # Rescale any class row whose norm exceeds the radius  [eq. 5]
    norms = row_l2_norms(weight)                    # (K,)
    scale = minimum(1, radius / norms)              # (K,)
    return weight * scale[:, None]                  # (K, d)


def class_balanced_loss(logits, labels, class_weights):
    # Softmax probabilities against one-hot targets  [code: cb#L92-L94]
    probs = softmax(logits)                         # (B, K)
    targets = one_hot(labels, num_classes)          # (B, K)
    # Every entry carries the weight of its sample's true class
    sample_weights = class_weights[labels][:, None] # (B, 1)
    bce = -(targets * log(probs) + (1 - targets) * log(1 - probs))
    # Mean over all B*K entries, as the released code does
    return mean(bce * sample_weights)

Pipeline comparisons. The paper's table of variants covers, on the classifier: naive; post-hoc L2-normalisation; post-hoc \tau-normalisation; cross-entropy with L2-normalisation; class-balanced loss with L2-normalisation, weight decay, MaxNorm, or combinations tab. 1 tab. 4. The \tau-normalisation is a post-hoc rescaling of the trained classifier §4.2. The code divides each class row by its norm to the power \tau regularizers.py, with \tau = 1.9 in the reported runs app. A demo2_second-stage-training.py.

Hyperparameters

Values in the paper and in the released code are shown separately where they differ. "not stated" means the paper text gives no value.

Results

Author-reported. Setup: top-1 accuracy, averaged over classes, on balanced test sets. CIFAR-100 with imbalance factors 100, 50 and 10 tab. 2. ImageNet-LT and iNaturalist2018 tab. 3. The paper's ablation is CIFAR100-LT with imbalance factor 100 tab. 1. Competitor numbers come from their respective papers tab. 2 tab. 3.

CIFAR-100-LT, top-1 accuracy (%) tab. 2

The row labels follow the paper's table; see Pitfalls for the label conflict in its "+" notation.

CIFAR100-LT ablation, IF 100, top-1 accuracy (%) tab. 1

ImageNet-LT and iNaturalist2018, top-1 accuracy (%), all-class column tab. 3. Many, Medium and Few splits are in the paper.

The paper says these methods rival prior art but trail the ensemble and self-supervised methods on the large datasets §4.3 tab. 3. The released code reproduces only the CIFAR100-LT demos, so these large-dataset numbers cannot be rerun from it.

Paper vs code

Official code: https://github.com/ShadeAlsha/LTR-weight-balancing, linked in the paper's author block and in the open-source appendix §4.1 app. B. The repository owner matches the first author. Pinned at e89f29273d35bb02d58204e7e450a390920c264d, MIT licence (LICENSE). The release contains two demo scripts for CIFAR100-LT with imbalance factor 100 app. B. Discrepancies, each with both citations:

• Backbone. The paper uses ResNet32 on CIFAR100-LT §4.1. The code builds ResNet34 for both stages demo1_first-stage-training.py demo2_second-stage-training.py. The video demo also uses a modified ResNet34 app. C.
• Stage-1 epochs. The paper trains for 200 epochs §4.1. The code runs 320 demo1_first-stage-training.py. The release's README says the runtime can be reduced by setting this to 100.
• MaxNorm timing. The paper applies the renormalisation after each batch update §3.1. The code calls the projection once per epoch, after the batch loop trainval.py.
• MaxNorm radius. The paper writes a single \delta in eq. 4 and a per-layer \delta_l only in a commented-out appendix figure eq. 4 app. A. The code computes a radius from the initial classifier row norms, \min + 0.1 \cdot (\max - \min) regularizers.py. It applies MaxNorm only to the classifier weight and bias regularizers.py.
• Stage-2 scope. The paper's pipeline trains the classifier and does not state which layers are trainable §3.3. The code freezes everything except fc.weight and fc.bias demo2_second-stage-training.py. Weight decay is applied to the bias too, because the optimiser receives both tensors demo2_second-stage-training.py.
• Class-balanced loss variant. The paper names the class-balanced loss but does not give its form §3.3. The code uses the softmax variant: binary cross-entropy on softmax probabilities against one-hot targets, with each entry weighted by its sample's class weight class_balanced_loss.py. The gamma argument goes unused in that branch.
• Model selection. The paper says it tunes on a validation set and reports on a test set §4.1. The CIFAR code has only train and test loaders demo1_first-stage-training.py. The training loop picks the best epoch by per-class accuracy on the test loader trainval.py. The reported best-epoch numbers are therefore selected on the evaluation split.
• Test-time augmentation. The paper does not mention it. print_accuracy averages logits over horizontal flips by default eval_funcs.py eval_funcs.py. It is not in the training-loop selection path, so the reported numbers may or may not include flip averaging. The paper does not say.
• Weight-decay scale. Eq. 3 penalises \lambda \sum_k \lVert\theta_k\rVert^2 eq. 3. PyTorch's weight_decay argument adds weight_decay * w to the gradient, which corresponds to a penalty of half that size. Inferred from the optimiser's documented behaviour, not stated in the repository: the paper's \lambda is about half the code's weight_decay for the same objective. Weight decay in stage 1 applies to all parameters, including normalisation layers demo1_first-stage-training.py.
• Headline number. The paper reports 46.08% for weight decay alone and 53.35% for the full recipe tab. 2 tab. 1. The released logs record 47.18% best-epoch accuracy for weight decay alone with_WD_model_note_bestModel.log. The stage-2 log records 51.98% best-epoch and 51.83% final accuracy MaxNorm_WD_note_bestModel.log MaxNorm_WD_train.log. Below 53.35%, and below the ">52%" the release states app. B readme.md. The naive log records 39.62% best-epoch, against 38.38% in the paper naive_note_bestModel.log.
• Scope. The paper gives CIFAR100-LT for imbalance factors 50 and 10, ImageNet-LT and iNaturalist2018 §4.1 tab. 2 tab. 3. The release does not include those experiments.

Pitfalls

• Test-split selection inflates the numbers. Best-epoch selection on the evaluation split is in the code (see Paper vs code). Budget for lower accuracy when you select on a held-out split.
• "+" labels are inconsistent in the paper. The caption of tab. 2 says "+" methods use the class-balanced loss. Its "+ L2norm" (49.60) and "+ τ-norm" (51.31) rows match the post-hoc numbers in tab. 1, which use no class-balanced loss. The main text also gives weight decay alone as 46.09% in one place §4.3 and 46.08% in the tables.
• Post-hoc L2-normalisation can hurt. The paper says it favours rare classes at the cost of common ones, and on iNaturalist it drops to 51.3% from 65.4% §3.1 tab. 3. Check per dataset before adopting it.
• λ is sensitive. The paper says large λ makes training unstable and that the optimal λ varies with dataset size §3.2 §4.2 fig. 4. No λ value for the general case is given. The released value of 5\times10^{-3} is for CIFAR100-LT with imbalance factor 100 only demo1_first-stage-training.py.
• MaxNorm needs a rule for the radius. The paper says per-layer thresholds take time to tune §3.3. The code's rule (min + 0.1 × range of initial row norms) is one data-dependent choice; the paper gives no reason for it.
• The class-balanced loss scale. The code averages over all B×K entries, not over the batch class_balanced_loss.py. Its loss scale therefore differs from a batch-mean cross-entropy, and the learning rate is not retuned for it in the release. Check the scale when you swap the loss.
• Second-stage hyperparameters are not searched in the release. The paper says Bayesian optimisation tunes them §3.3. The code hard-codes the chosen values and includes no search code.
• Cost. Stage 1 is a full training run. In the release, stage 1 takes about two hours on an RTX 3090 app. B. Stage 2 trains one layer and takes a few minutes on a GPU readme.md, but needs a stage-1 checkpoint.
• Imbalanced intermediate layers. The paper finds filter norms are also imbalanced in hidden layers and says weight decay balances them too §3.2 fig. 3. It says MaxNorm could be applied to all layers but that it found this non-trivial, so it did not §3.2. The release applies MaxNorm only to the classifier.
• Balanced evaluation is assumed. The paper evaluates on balanced test or validation sets §4.1 and does not discuss other test distributions.

Lineage

• Builds on: the two-stage decoupling of feature learning and classifier learning from Kang et al. §1 §3.3. The class-balanced loss §3.3 and the \tau-normalisation §4.2 are used as components; the paper cites them but gives no arXiv ids in the text.
• Compared against: tab. 2 lists CE, CE+CB, KD, LDAM-DRW, BBN, LogitAdjust, LDAM+SSP, Focal, De-confound, \tau-norm, SSD, DiVE, DRO-LT, PaCo, ACE and RIDE. tab. 3 lists the large-dataset comparisons.
• Superseded by: not stated.