Variational Imbalanced Regression (VIR)¶

← Continuous Group DRO | Algorithms overview → Methods overview

Standard regression on a long-tailed target fails in two ways at once: the bulk of the labels dominates the objective (so the rare tails are under-fitted), and a deterministic model gives no honest uncertainty estimate exactly where it is least reliable. Variational Imbalanced Regression (VIR) (Wang & Wang, NeurIPS 2023) addresses both in a single model. It learns a probabilistic latent representation whose per-bin statistics are smoothed across neighbouring labels, and it predicts a full Normal–Inverse–Gamma (NIG) distribution whose pseudo-counts are reweighted by the smoothed label density.

This sits alongside output reweighting (LDS), feature-statistic smoothing (FDS), representation learning (PRIME & ConR), and worst-regime training (Continuous Group DRO). Its distinguishing feature is that the output is a predictive distribution, not only a point estimate.

Why both a representation and a distribution?¶

• Output reweighting alone (LDS, density weighting) changes the loss but leaves the hidden representation of a rare label dominated by the bulk. A mini-batch with no tail object still produces no tail gradient.
• Feature smoothing alone (FDS) borrows first/second moments across neighbouring label bins, but the model is still deterministic: there is no predictive distribution to inspect in the tail.
• Evidential regression alone (DER) predicts an NIG distribution but is designed for the balanced regime; on imbalanced data the tail evidence is unstable and the NIG parameters can be poorly estimated.
• VIR couples the two: the representation of a rare label borrows statistics from nearby labels, and the evidential pseudo-count is enlarged for rare labels so the likelihood focuses there. The result is a single-forward-pass predictive distribution.

Mathematical formulation¶

Notation and label bins¶

Let \(\{(x_i,y_i)\}_{i=1}^N\) be the training set and partition the target range into \(B\) equal-width bins with edges \(e_0<\dots<e_B\) and centres \(c_b=(e_{b-1}+e_b)/2\). Write \(b(i)\) for the bin of \(y_i\) and \(N_b\) for its population. Bins are half-open, \([e_b,e_{b+1})\); out-of-range values are clamped into the boundary bins.

Probabilistic representation¶

A feature encoder maps \(x_i\) to features, and a linear head emits a diagonal Gaussian

The N.I.D. assumption replaces the usual I.I.D. assumption: a representation may be calibrated using data with similar labels, not only its own input.

Statistics of the statistics¶

For each bin, VIR computes three quantities. The mean of the latent means and the variance of the latent means are

The expected within-bin covariance follows from the law of total variance, with the variance of the bin mean subtracted:

Neighbouring-label smoothing¶

A symmetric Gaussian kernel \(k(c_b,c_{b'})=\exp\!\big(-(c_b-c_{b'})^2/(2\sigma^2)\big)\) is row-normalised. Empty bins carry no information and are excluded from each row before normalisation; otherwise their global-mean imputation would drag the extreme bins toward the bulk. The smoothed statistics are

where \(\bar k^{\,2}\) is the squared kernel renormalised over occupied bins (the variance of a weighted mean is the sum of squared weights).

Probabilistic whitening and recolouring¶

Each representation is whitened with its own bin's statistics and recoloured with the smoothed ones:

The scale ratio is clamped to \([0.1,10]\) for numerical stability. A count-dependent shrinkage

protects singleton and very sparse bins: the paper assumes enough samples per bin for the within-bin covariance to be meaningful, while a catalog tail often has only a handful. Set smoothing_shrinkage=0 to recover the paper's full recolouring.

During training, \(z_i\sim\mathcal N(\tilde z_i^\mu,\operatorname{diag}(\tilde z_i^\Sigma))\) is reparameterised before the NIG head. The statistics are treated as constants and recomputed from the current encoder at the start of each epoch.

Reweighted NIG posterior¶

The smoothed label density is \(\tilde p_b=\sum_{b'}k_{bb'}\hat p_{b'}\), and the importance weight is

A linear head emits a raw pseudo-count \(n_i^{\text{raw}}\), a location \(\psi_i\), and a raw scale \(\phi_i^{\text{raw}}\). They are constrained as

with \(n_{\min}=2\) following the paper. Given the prior \(\mathrm{NIG}(\gamma_0,\nu_0,\alpha_0,\beta_0)\), the reweighted posterior parameters are

Because \(w_i\) grows as the smoothed density falls, a rare label receives a larger pseudo-count and therefore a larger weight in the likelihood. This is the probabilistic reweighting mechanism; it is the NIG analogue of inverse-density weighting, not a post-hoc multiplication of the loss.

Objective¶

Marginalising the Gaussian likelihood over the NIG prior gives a Student-t predictive

Its negative log-likelihood is

with \(\Omega_i=2\beta_i^*(1+\nu_i^*)\). The training objective adds the evidential regulariser of the paper,

and optionally a small \(\mathrm{KL}\big(q_\phi(z_i\mid x_i)\,\|\,\mathcal N(0,I)\big)=\tfrac12\sum_{j}\big(\sigma_{ij}^2+\mu_{ij}^2-1-\log\sigma_{ij}^2\big)\) term (summed over latent dimensions, averaged over the batch) scaled by kl_weight.

Prediction¶

The predictive mean is \(\hat y_i=\gamma_i^*\). The total predictive variance of the Student-t is

which is the sum of the aleatoric and evidential epistemic variances. predict(..., return_std=True) returns the square root of this quantity; predict_interval(...) uses the exact Student-t quantiles with \(2\alpha_i^*\) degrees of freedom.

Algorithm¶

1. Fit label statistics. Bin the training targets, estimate \(\hat p_b\), smooth the density, and store \(w_b=\tilde p_b^{-1/2}\) and the bin edges.
2. Warm-up epoch. Train with the raw representation (identity recolouring). The first per-bin statistics are collected at the start of the second epoch.
3. Epoch loop (from epoch 2).
4. Refresh \((\mu_b^\mu,\mu_b^\Sigma,\Sigma_b^\mu)\) and the smoothed counterparts from the current encoder on the full training set (no gradient).
5. For each mini-batch: encode, recolour with the start-of-epoch statistics, reparameterise, and evaluate the NIG head.
6. Minimise \(\mathcal L_i\) with Adam, adding the optional KL regulariser.
7. Inference. Encode the test point; a first pass gives an initial label, its bin selects the stored training statistics, and the recoloured representation gives the final NIG posterior. The density weight is evaluated at the model's own predicted label, so no test labels are used. n_refine controls how many times the bin is refreshed.

Usage¶

End-to-end regressor¶

import torch
import torch.nn as nn

from torchregress.algorithms import VIRRegressor

torch.manual_seed(0)
y = torch.cat([0.4 * torch.randn(500), 2.5 * torch.randn(50)])
X = torch.stack([torch.sin(y), torch.cos(y)], dim=1) + 0.05 * torch.randn(len(y), 2)

encoder = nn.Sequential(nn.Linear(2, 32), nn.SiLU(), nn.Linear(32, 32))

regressor = VIRRegressor(
    encoder,
    n_bins=12,
    kernel_sigma=1.0,          # bandwidth in bin-width units
    smoothing_shrinkage=10.0,  # protects sparse tail bins
    reg_weight=0.01,           # lambda in the evidential regulariser
    epochs=60,
    batch_size=64,
    lr=1e-3,
    seed=0,
)
regressor.fit(X, y)

mean, std = regressor.predict(X[:5], return_std=True)
lower, upper = regressor.predict_interval(X[:5], confidence=0.9)
print(mean.shape, std.shape, lower.shape, upper.shape)

The fitted attributes are:

The loss in a custom loop¶

VIRLoss consumes the raw output of a 3-unit head (it applies the softplus constraints itself). It is a DistributionLoss, so it follows the shared (y_pred, target, mask, weights) contract:

from torchregress.algorithms import VIRLoss

loss_fn = VIRLoss(
    prior=(0.0, 1.0, 2.0, 0.1),
    reg_weight=0.01,
    n_bins=20,
    kernel_sigma=1.0,
    max_importance_weight=10.0,
).fit(y_train)                      # estimates the smoothed label density

for x, y in loader:
    z, _ = encoder(x)
    raw = nig_head(z)               # (N, 3): raw_n, psi, raw_phi
    loss = loss_fn(raw, y)          # reweighted NIG NLL + regulariser
    loss.backward()
    optimizer.step()

The loss also exposes posterior_parameters(raw, importance_weight=None) and predictive_mean_std(posterior) for inspection; before fit() the importance weights are all one and the loss is a plain DER-style evidential NLL.

Full runnable comparison¶

The following example trains a plain-MSE baseline and VIR on a deliberately imbalanced, ordered-target regression problem. It compares the rare-tail error and the empirical coverage of 90 % intervals. Every call matches the signatures above.

import torch
import torch.nn as nn

from torchregress.algorithms import VIRRegressor


def make_data(n, seed, probabilities=(0.02, 0.08, 0.80, 0.08, 0.02)):
    """Five ordered target values, with an optional balanced test distribution."""
    generator = torch.Generator().manual_seed(seed)
    classes = torch.tensor([-2.0, -1.0, 0.0, 1.0, 2.0])
    probability = torch.tensor(probabilities)
    targets = classes[torch.multinomial(probability, n, replacement=True, generator=generator)]
    direction = torch.tensor([1.0, 0.2, -0.3, 0.4, 0.1, -0.2, 0.3, 0.5])
    features = 1.5 * targets.unsqueeze(1) * direction + 0.8 * torch.randn(
        n, direction.numel(), generator=generator
    )
    return features, targets.unsqueeze(1)


def make_encoder(n_features, hidden=32):
    return nn.Sequential(
        nn.Linear(n_features, hidden),
        nn.SiLU(),
        nn.Linear(hidden, hidden),
        nn.SiLU(),
    )


torch.manual_seed(0)
X_train, y_train = make_data(2000, seed=0)
X_test, y_test = make_data(1000, seed=99, probabilities=(0.2, 0.2, 0.2, 0.2, 0.2))
tail = y_test.abs().squeeze(1) >= 1.0

# Plain-MSE baseline: same encoder architecture, optimiser and training budget.
torch.manual_seed(0)
encoder = make_encoder(X_train.shape[1])
baseline = nn.Sequential(make_encoder(X_train.shape[1]), nn.Linear(32, 1))
with torch.no_grad():
    baseline[0].load_state_dict(encoder.state_dict())
optimizer = torch.optim.Adam(baseline.parameters(), lr=3e-3)
generator = torch.Generator().manual_seed(1)
for _ in range(120):
    baseline.train()
    order = torch.randperm(len(X_train), generator=generator)
    for start in range(0, len(X_train), 64):
        index = order[start : start + 64]
        optimizer.zero_grad(set_to_none=True)
        ((baseline(X_train[index]) - y_train[index]) ** 2).mean().backward()
        optimizer.step()
baseline.eval()
with torch.no_grad():
    mse_prediction = baseline(X_test)
    residual_std = (baseline(X_train) - y_train).std()
mse_tail = (mse_prediction - y_test).abs()[tail].mean()
mse_coverage = ((y_test - mse_prediction).abs() <= 1.645 * residual_std).float().mean()

# VIR: same encoder weights, probabilistic smoothing and reweighted NIG head.
vir = VIRRegressor(
    encoder,
    n_bins=10,
    kernel_sigma=1.0,
    smoothing_shrinkage=10.0,
    epochs=120,
    lr=3e-3,
    batch_size=64,
    seed=0,
)
vir.fit(X_train, y_train)
vir_prediction, vir_std = vir.predict(X_test, return_std=True)
lower, upper = vir.predict_interval(X_test, confidence=0.9)
vir_tail = (vir_prediction - y_test).abs()[tail].mean()
vir_coverage = ((y_test >= lower) & (y_test <= upper)).float().mean()

print(f"plain MSE : tail MAE={mse_tail:.3f}  90% coverage={mse_coverage:.3f}")
print(f"VIR       : tail MAE={vir_tail:.3f}  90% coverage={vir_coverage:.3f}")

Typical output on CPU with these seeds:

plain MSE : tail MAE≈0.50  90% coverage≈0.57
VIR       : tail MAE≈0.41  90% coverage≈0.75

Both models are deterministic for the given seeds. The absolute numbers are not the point: the ordered rare targets are systematically under-predicted by plain MSE, and the single global residual standard deviation is badly calibrated, while VIR's per-sample evidential scale is closer to nominal.

When to use VIR¶

graph TD
    A["Imbalanced continuous target?"] -->|No| B["Standard loss"]
    A -->|Yes| C["Need a predictive distribution?"]
    C -->|No| D["LDS / FDS / PRIME / ConR / Group DRO"]
    C -->|Yes| E["VIR"]
    E --> F["Check per-bin support"]
    F -->|"many samples per bin"| G["Full smoothing (shrinkage=0)"]
    F -->|"sparse tail bins"| H["Shrinkage > 0 + report bin counts"]

Caveats and failure modes¶

References¶

See also¶

• VIRRegressor / VIRLoss API
• Imbalanced regression losses — LDS, FDS, Balanced MSE, Focal-R, density weighting.
• Evidential regression — the NIG head and its uncertainty decomposition.
• PRIME & ConR — representation-level alternatives without predictive uncertainty.
• Continuous Group DRO — worst-regime rather than worst-sample robustness.
• Forward generative inversion — another route to tail-aware regression when features are noisier than labels.
• Method selection — task-level guidance.