Representation-Level Imbalanced Regression: PRIME & ConR¶

← Imbalanced targets | Rank-N-Contrast → Balanced MSE

Long-tailed regression is not only a loss-weighting problem. Even with a perfectly reweighted output loss, a neural network's hidden representation can collapse in sparse target regions: the few metal-poor stars, young ages, or deep-depth pixels are absorbed into the feature statistics of the well-populated bulk. Output-level methods (LDS, Balanced MSE, density weighting) re-weight errors; representation-level methods shape features.

torchregress already provides two representation-level tools:

• FDS — FeatureDistributionSmoother: calibrates feature first/second moments across target bins by kernel smoothing. This is a feature-statistics layer, and it is not duplicated here.
• Rank-N-Contrast — RankNContrastLoss: orders features by pairwise label distance.

This page adds two complementary losses, both implemented in torchregress.losses.prime:

• PRIME (PrimeLoss) — learnable proxies whose targets tile the label range uniformly; every mini-batch, even one containing only bulk targets, still pulls features toward the under-populated proxies.
• ConR (ConR) — a contrastive regularizer that penalises feature-space collapse: it selects label-dissimilar but prediction-similar pairs as negatives and pushes them apart proportionally to their label distance and inverse label density.

1. PRIME: proxy-based representation learning¶

PRIME (Lim et al., ICML 2025) introduces \(C\) learnable proxies \(\{(\mathbf z_j^p, y_j^p)\}_{j=1}^{C}\), with feature points \(\mathbf z_j^p\) and targets \(y_j^p\) distributed uniformly over the training target range \([y_{\min}, y_{\max}]\) (the \((C+1)\)-quantiles for scalar targets). The proxy features are optimized jointly with the model, so the proxy set is a persistent, balanced reference frame for the feature space.

For a sample feature \(\mathbf z\) with label \(y\), PRIME aligns the sample to the proxies with a soft-label cross-entropy:

where \(c_j\) are optional class-balanced effective-number weights (class_balanced=True, paper Eq. 10, \(\beta = 0.99\)). As \(\tau_t \to 0\) this becomes an InfoNCE-style alignment to the single nearest proxy; with a finite \(\tau_t\) it is a smooth soft-label version that shares information between neighbouring target values.

To keep the proxies ordered and separated, the proxy geometry adds a t-SNE-style KL term and a spread term:

The total regularizer is \(\lambda_p\mathcal L_{\text{proxy}} + \lambda_a\,\overline{\mathcal L_{\text{align}}}\), added to the task loss. Because the proxies are parameters, the whole target range receives gradient on every step — a mini-batch that happens to contain no tail objects still updates the tail proxies, which is exactly what Rank-N-Contrast and batch-contrastive methods lack.

import torch
import torch.nn as nn
from torchregress.losses import PrimeLoss

torch.manual_seed(0)
n = 2000
# Long-tailed target: many objects near 0, few near the edges of [-1, 1].
y = torch.cat([0.15 * torch.randn(1600), torch.rand(400) * 2 - 1]).clamp(-1.0, 1.0)
x = torch.stack([y, torch.sin(3 * y)], dim=1) + 0.05 * torch.randn(n, 2)

feature_dim = 32
encoder = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, feature_dim))
head = nn.Linear(feature_dim, 1)
prime = PrimeLoss(
    feature_dim=feature_dim, n_proxies=50, label_range=(-1.0, 1.0), class_balanced=True
)
optimizer = torch.optim.Adam(
    list(encoder.parameters()) + list(head.parameters()) + list(prime.parameters()),
    lr=1e-3,
)
mse = nn.MSELoss()
for _ in range(50):
    z = encoder(x)
    pred = head(z).squeeze(-1)
    loss = mse(pred, y) + 0.5 * prime(z, y)  # both terms, one backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

print("tail loss:", float((head(encoder(x[y > 0.8])).squeeze(-1) - y[y > 0.8]).abs().mean()))

Call prime.fit(y_train) instead of label_range to estimate the proxy targets and class-balanced weights from the training labels.

2. ConR: contrastive regularizer for collapse¶

Uncertainty-aware ConR¶

UncertaintyAwareConR wraps ConR for labels with known measurement uncertainty. Instead of forcing each pair into a hard positive/negative class, it uses the Gaussian probability that two noisy labels fall within label_threshold to soften the positive and negative pair weights. The observed-label distances still enter the pushing weight; only pair membership is softened. For scalar Gaussian errors the exact positive probability is \(\Phi((\omega-|y_i-y_j|)/s_{ij})-\Phi((-\omega-|y_i-y_j|)/s_{ij})\), where \(s_{ij}^2=\sigma_i^2+\sigma_j^2\). The weighted mean of positive similarities uses these memberships and reduces exactly to ordinary ConR at zero error. Soft noisy memberships currently accept scalar labels only; a multivariate Euclidean distance is not a scalar Gaussian random variable. For gross label contamination (rather than measurement error), use a contamination model or a robust task loss instead.

from torchregress.losses import UncertaintyAwareConR

conr = UncertaintyAwareConR(label_sigma=0.02, label_threshold=0.1)
loss = mse(pred, y) + 4.0 * conr(features, y, prediction=pred)

ConR (Keramati et al., ICLR 2024) selects pairs using both label and prediction similarity. For anchor \(i\) with feature \(\mathbf z_i\), label \(y_i\) and prediction \(p_i\):

with \(\mathcal L_i=0\) when the anchor has no positive or no negative partner. The intuition: a negative pair is a pair with very different labels that the model has nevertheless mapped to the same region of feature space — that is precisely the collapse that hurts minority targets. The pushing weight \(S_{iq}\) grows with the label disagreement and with the anchor's inverse-label-density weight \(\eta_i\) (ConR.fit(train_targets)), so minority anchors push harder.

from torchregress.losses import ConR

# Reuse the long-tailed data and encoder defined above.
conr = ConR(label_threshold=0.2, temperature=0.07, push_scale=0.2)
conr.fit(y)   # inverse-label-density pushing weights from the training targets
for _ in range(50):
    z = encoder(x)
    pred = head(z).squeeze(-1)
    loss = mse(pred, y) + 4.0 * conr(z, y, prediction=pred)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Comparison: which representation method?¶

Caveats¶

References¶

See also¶

• Imbalanced targets — LDS, FDS, Balanced MSE, Focal-R.
• Rank-N-Contrast — pairwise label-order regularizer.
• Post-hoc calibration — re-calibrate after changing the objective.
• Imbalanced / rare-target regression — task-level selection matrix.