RankSim: Ranking Similarity Regularization¶

← Imbalanced targets | PRIME & ConR →

Long-tailed regression is not only an output-space reweighting problem. A network trained with a pointwise loss can learn a feature space in which rare targets are absorbed by the well-populated bulk: the few metal-poor stars, young ages, or extreme-depth pixels end up close to majority features even though their labels are far away.

RankSim (Gong, Mori & Tung, ICML 2022) attacks this directly. For every anchor it compares two sorted lists — the other samples ordered by label distance and the same samples ordered by feature similarity — and penalises any disagreement between the two rankings. The regularizer is parameter-free, needs no model predictions, and is orthogonal to reweighting (LDS, density weighting, Balanced MSE) and to distribution smoothing (FDS).

Pipeline¶

graph LR
    A["features z: B x D"] --> B["cosine similarity matrix"]
    C["targets y: B"] --> D["label similarity: negative distance"]
    B --> E["rank every row"]
    D --> F["rank every row"]
    E --> G["MSE between the two rank vectors"]
    F --> G
    G --> H["gamma * L_RankSim"]

The term is added to a task regression loss; the encoder receives gradients from both terms in a single backward pass.

Mathematical definition¶

Let \(z_i = f_\theta(x_i) \in \mathbb{R}^D\) be the feature vector of sample \(i\) and \(y_i \in \mathbb{R}\) its scalar target. For an anchor \(i\), define the competition rank of a score vector \(\mathbf a \in \mathbb{R}^M\) as

so the largest score receives rank \(1\) and equal scores receive equal ranks. The label-space and feature-space similarities are

where the cosine similarity is the paper's default and an L2 option is also available. The per-anchor rank-matching loss and the regularizer are

which is Eq. (4) of Gong et al. (2022) with the ranks normalised by \(M\) (as in the reference implementation). Minimising \(\ell_i\) maximises the Spearman correlation between the label and feature rankings of anchor \(i\).

Differentiating through the ranking¶

Ranking is piecewise constant, so its true gradient is zero almost everywhere. RankSim uses the continuous interpolation of Vlastelica et al. (2020). With \(\lambda>0\) the interpolation strength and \(g = \partial \mathcal{L}/\partial \operatorname{rk}(S^z)\) the incoming gradient,

The backward pass therefore costs one extra ranking (one extra sort per row). interpolation_strength is \(\lambda\); the paper uses \(\lambda = 2\).

Composing with a task loss¶

The total objective is the task loss plus the balancing weight \(\gamma\) times the regularizer:

RankSimLoss folds \(\gamma\) into its return value, so the composition is

rank_sim = RankSimLoss(gamma=1.0)
loss = task_loss(prediction, target) + rank_sim(features, target)   # gamma folded in

Equivalently, keep the balancing weight at the call site and use the functional form:

gamma = 1.0
loss = task_loss(prediction, target) + gamma * rank_sim_regularizer(features, target)

Both lines describe the same objective. Do not multiply by gamma twice when using the module.

Runnable example¶

The example trains the same small MLP twice on a long-tailed half-circle: once with MSE only, once with MSE + RankSimLoss(gamma=1.0). Only 8% of the training targets are in the tail. Evaluation uses a balanced 101-point probe, following the balanced validation/test protocol of the deep imbalanced regression benchmarks.

import copy

import torch
import torch.nn as nn
import torch.nn.functional as F

from torchregress.losses.rank_sim import RankSimLoss

torch.manual_seed(0)

# 1) Long-tailed half-circle: only 8% of the training targets are in the tail.
n_train, n_tail = 240, 19
generator = torch.Generator().manual_seed(0)
y = torch.cat(
    [
        0.45 * torch.rand(n_train - n_tail, generator=generator) ** 1.2,
        0.65 + 0.35 * torch.rand(n_tail, generator=generator),
    ]
)
angle = torch.pi * y
x = torch.stack([torch.cos(angle), torch.sin(angle)], dim=1)

# 2) Balanced probe set for evaluation (101 uniformly spaced targets).
y_probe = torch.linspace(0.0, 1.0, 101)
probe_angle = torch.pi * y_probe
x_probe = torch.stack([torch.cos(probe_angle), torch.sin(probe_angle)], dim=1)


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(2, 32), nn.SiLU(), nn.Linear(32, 4), nn.SiLU()
        )
        self.head = nn.Linear(4, 1)

    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        features = self.encoder(inputs)
        return self.head(features).squeeze(-1), features


def train(model: Model, gamma: float) -> tuple[torch.Tensor, torch.Tensor]:
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
    task_loss = nn.MSELoss()
    rank_sim = RankSimLoss(gamma=gamma)
    for _ in range(200):
        optimizer.zero_grad()
        prediction, features = model(x)
        loss = task_loss(prediction, y)
        if gamma > 0.0:
            loss = loss + rank_sim(features, y)  # task_loss + gamma * L_RankSim
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        prediction, features = model(x_probe)
    return prediction, features


def ordering(features: torch.Tensor) -> float:
    """Spearman correlation between feature distance and label distance."""
    unit = F.normalize(features, dim=-1)
    feature_distance = 1.0 - unit @ unit.T
    label_distance = (y_probe[:, None] - y_probe[None, :]).abs()
    upper = torch.triu_indices(y_probe.numel(), y_probe.numel(), offset=1)
    left = feature_distance[upper[0], upper[1]]
    right = label_distance[upper[0], upper[1]]
    left_rank = left.argsort().argsort().float()
    right_rank = right.argsort().argsort().float()
    return float(torch.corrcoef(torch.stack([left_rank, right_rank]))[0, 1])


baseline_model = Model()
ranksim_model = copy.deepcopy(baseline_model)  # identical initialisation
baseline_prediction, baseline_features = train(baseline_model, gamma=0.0)
ranksim_prediction, ranksim_features = train(ranksim_model, gamma=1.0)

tail = y_probe > 0.6
for name, prediction, features in (
    ("baseline", baseline_prediction, baseline_features),
    ("RankSim ", ranksim_prediction, ranksim_features),
):
    tail_mae = (prediction[tail] - y_probe[tail]).abs().mean()
    print(f"{name}  tail MAE = {float(tail_mae):.4f}   ordering = {ordering(features):.3f}")

Example output on CPU (exact values can vary slightly with the PyTorch build):

baseline  tail MAE = 0.0120   ordering = 0.726
RankSim   tail MAE = 0.0054   ordering = 0.997

The baseline fits the crowded region but leaves the tail poorly ordered. RankSim raises the feature/label rank correlation from 0.73 to nearly 1.0 and cuts the tail error by more than half. The improvement is a property of the representation, so it survives a change of regression head; validate it on tail-slice metrics (few-shot MAE/RMSE, worst-bin error), not only on the overall MSE.

API¶

The implementation lives in torchregress.losses.rank_sim and is re-exported from the package:

from torchregress.losses import RankSimLoss, rank_sim_regularizer

mask and weights follow the shared BaseLoss contract: masked samples are removed from the anchor set and from every ranking, and reduction="mean" returns \(\sum_i w_i \ell_i / \sum_i w_i\). Under a deduplicating policy, a representative inherits the mean weight of its label group.

Implementation choices and simplifications¶

• Duplicate labels. The paper samples the batch subset \(\mathcal{M}\) so that each label occurs at most once, choosing one instance per label at random. Randomness would break the library's deterministic contract, so duplicate_policy="mean" (default) uses the group-mean feature of all samples sharing a label — deterministic and invariant to sample order — while "first" deterministically takes the first occurrence in the batch. "keep" disables the reduction entirely.
• Ties. The paper defines the rank by Eq. (1), which gives equal scores equal ranks. The reference code instead breaks ties arbitrarily through argsort. This implementation follows Eq. (1): deterministic and invariant to the order of equal-label samples.
• Rank normalisation. Ranks are divided by the number \(M\) of retained anchors, and the per-anchor losses are averaged under the default reduction. The paper's Eq. (4) is the "sum" reduction; see the warning above for the \(\gamma\) conversion.
• Scalar targets. The paper defines \(\sigma^{y}\) for scalar labels. Multi-output targets are rejected rather than silently reduced to a scalar; use one RankSim term per target dimension if needed.
• Similarity. Cosine is the paper's default. "l2" (negative Euclidean) is offered for parity with RankNContrastLoss, but the paper's ablations and tuned hyperparameters assume cosine.
• No fitting, no parameters. Unlike PRIME (learnable proxies) or FDS (running statistics), RankSim has no state; it can be added to an existing training loop without a fit call.

Comparison: which representation method?¶

graph TD
    A["Rare tail targets"] --> B{"Do tail labels appear in most batches?"}
    B -->|"No"| C["PRIME or FDS: persistent signal"]
    B -->|"Yes"| D{"Are model predictions available inside the loss?"}
    D -->|"Yes"| E["ConR: collapse penalty"]
    D -->|"No"| F{"Full ordering or pairwise contrast?"}
    F -->|"Full ordering"| G["RankSim"]
    F -->|"Pairwise"| H["Rank-N-Contrast"]

Choosing the hyperparameters¶

Caveats¶

References¶

See also¶

• Imbalanced targets — LDS, FDS, Balanced MSE, density weighting.
• PRIME & ConR — proxy and contrastive representation losses.
• Base classes — reduction, mask, and weight contracts.
• Metric selection — tail-slice evaluation.
• Post-hoc calibration — re-calibrate after changing the objective.