Test-Time Adaptation API¶

Complete reference for torchregress.test_time. This package provides model-agnostic test-time adaptation and uncertainty utilities for regression: Bayesian linear heads, label-shift correction, OT-based conformal reweighting, feature alignment, EMA ensembling, and confidence-based sample selection. Every primitive operates on torch.Tensor or numpy.ndarray inputs and returns a PredictiveBatch or self-describing dataclass (see the Predictive containers section in the Utilities API).

For background, see Bayesian Linear Regression and Optimal Transport Conformal. The shared PredictiveBatch container is documented in the Utilities API — Predictive containers section.

Adaptation interfaces (test_time.base)¶

from torchregress.test_time import (
    AdaptationBatch, SupportsPredictiveBatch, flatten_adaptation_parameters
)
assert isinstance(my_model, SupportsPredictiveBatch)
batch = AdaptationBatch(x=x_test, predictions=predictive_batch)
# Any dict of name -> iterable of parameters works
flat = flatten_adaptation_parameters(
    {"head": my_model.head.parameters(), "norm": my_model.norm.parameters()}
)

Bayesian linear heads (test_time.bayes)¶

Conjugate Gaussian linear regression on fixed features (closed-form posterior, online updates).

Reference: Bishop, Pattern Recognition and Machine Learning (2006), §3.3 (conjugate Bayesian linear regression).

import torch
from torchregress.test_time import BayesianLinearHead

phi = torch.randn(64, 5)        # fixed features
y   = torch.randn(64, 1)
head = BayesianLinearHead(in_features=5, out_features=1, noise_variance=0.1)
head.fit(phi, y)
batch = head.predictive_batch(phi[:8])   # mean / std / epistemic / aleatoric
samples = head.sample_weights(n_samples=16)   # \[16, 1, 6\]  (with intercept)

Label-shift correction (test_time.label_shift)¶

EM-based target-prior estimation and Gaussian predictive adjustment under marginal label shift \(p_{\text{target}}(y) \neq p_{\text{source}}(y)\).

EM iterates the Saerens et al. fixed point: with source posteriors \(p_{ik} = p_s(y = k \mid x_i)\) and source prior \(\pi\),

which increases the target log-likelihood \(\frac{1}{n}\sum_i \log \sum_k p_{ik}\, q_k / \pi_k\) at every step.

import numpy as np
from torchregress.test_time import correct_gaussian_predictions_for_label_shift

mean, std, meta = correct_gaussian_predictions_for_label_shift(
    mean=test_mean, std=test_std, source_targets=train_y,
    features=test_x,                           # optional, enables local consistency weights
)
print(meta["estimate_converged"], meta["selected_rows"])

COSA Output Adaptation (test_time.cosa)¶

COSA-style (Conformal Output Space Adaptation) residual adaptation for point predictions and uncertainty under delayed label observations.

from torchregress.test_time import DelayedLabelResidualAdapter

adapter = DelayedLabelResidualAdapter(base_model, ema_beta=0.2)
# Under streaming protocol:
# 1. Predict first
adapted_batch = adapter.predict_distribution(X_t)
# 2. Observe delayed labels later
adapter.partial_fit(X_old, y_old)

Optimal-transport conformal reweighting (test_time.ot_conformal)¶

Score-CDF matching reweighting for classification-style nonconformity scores under non-exchangeable target shift, and weighted split-conformal regression.

import torch
from torchregress.test_time import WeightedConformalRegressionAdapter

adapter = WeightedConformalRegressionAdapter(alpha=0.1)
# Fit density ratio estimator and calibrate scores
adapter.calibrate(y_pred_cal, y_cal, X_cal, X_tgt)
# Predict intervals under covariate shift
lower, upper = adapter.predict_interval(y_pred_test, X_test)

Density-ratio weights & joint TTA (test_time.shift_weights, test_time.joint_tta)¶

JointDistributionalTTA head matching minimizes diagonal Gaussian \(D_{\mathrm{KL}}(\mathcal{N}(\mu_t,v_t)\Vert\mathcal{N}(\mu_s,v_s))\) with frozen teacher moments. Its per-coordinate score is \(\tfrac12[\log(v_s/v_t)+(v_t+(\mu_s-\mu_t)^2)/v_s-1]\). An unperturbed model matched to its own moments has zero gradient; no variance-sharpening signal is invented. With alignment disabled and identical teacher/student views, use pseudo_label_rounds=0 for an honest weighted-CP control. Meaningful self-training requires distinct justified views or teachers. The conformal statement concerns observed targets and the required sampling/ weight assumptions; it does not establish calibrated physical truth.

Streaming evaluation (test_time.benchmark)¶

Shift calibration primitives (test_time.calibration)¶

Backward-compatible re-export of the variance-inflating temperature calibrator used by ShiftFactoredPredictiveTransport and any other shift-aware pipeline.

from torchregress.test_time import RepresentationShiftInflator

cal = RepresentationShiftInflator(base_temperature=1.0, slope=0.2,
                                    max_temperature=2.0, clip_quantile=0.05)
cal.fit(source_features)                # learns per-source calibration
std_calibrated = cal.calibrate_std(test_std, test_features)

Confidence and selection (test_time.selection)¶

import numpy as np
from torchregress.test_time import local_consistency_weights, select_high_confidence

w = local_consistency_weights(features, probabilities)         # [B] weights
mask = select_high_confidence(probabilities, top_fraction=0.5, min_count=32)

Feature / subspace alignment (test_time.subspace)¶

Reference: Fernando, Habrard, Sebban, Tuytelaars, "Unsupervised Visual Domain Adaptation Using Subspace Alignment" (ICCV 2013).

import numpy as np
from torchregress.test_time import WeightedSubspaceMomentAligner

aligner = WeightedSubspaceMomentAligner(rank=10, variance_threshold=0.95)
aligner.fit(X_train, y_train)
X_test_aligned = aligner.transform(X_test)

Dynamic ensembling (test_time.dynamic)¶

from torchregress.test_time import ParameterEMA
ema = ParameterEMA(decay=0.99)
ema.initialize(model)
for x, y in loader:
    ...
    ema.update(model)
ema.copy_to(model)   # evaluate with EMA weights

Shift-factored predictive transport (test_time.transport)¶

End-to-end orchestrator: combines prior transport, feature alignment, uncertainty inflation, and conformal calibration for predictive batches under target shift.

fit_source discretises the source predictive distributions on the support grid (n_support points spanning the source labels plus support_margin) and uses their mean as the source prior — the marginal the predictive posteriors integrate to, which the EM prior estimate in adapt_unlabeled_target needs (see the warning above). EM runs on the top_fraction most confident target rows and stops once the mean target log-likelihood gains less than 1e-4 nats per row per iteration; the estimate is then shrunk toward the source prior and ratio-clipped (prior_transport_strength, prior_ratio_clip). Metadata reports target_prior_raw (EM), target_prior (stabilised), prior_source_target_tv, estimate_converged and transport_applied.

For a Gaussian predictive the prior correction is applied as a moment change: the native mean moves by \(w_i\,(m^{\text{transported}}_i - m^{\text{grid}}_i)\) and the std is scaled by \(1 + w_i\,(s^{\text{transported}}_i / s^{\text{grid}}_i - 1)\), where \(m, s\) are the moments of the row's binned predictive before and after the prior reweighting and \(w_i\) is the row's predictive mass inside the support grid. A prior equal to the source prior returns the input predictive, and a row whose predictive lies past the grid passes through unchanged. (Before 0.3.0 the output was the moments of the transported grid itself, which added the discretisation and the truncation at the grid edge: target predictives beyond the source label range collapsed onto the edge bin, std of the order of the grid step and NLL in the thousands.) Metadata prior_mass_outside_source_range is the raw EM prior mass on the support bins outside [min, max] of the source labels.

The transport supports PPI (mean / quantile / ols estimands) and dispatches conformal calibration to cqr / cti / interval / split based on the predictive family. The conformal threshold is the exact split-conformal order statistic k = ceil((n + 1)(1 - alpha)) of the calibration scores (+inf, i.e. an unbounded interval, when k > n). Outputs keep the floating dtype of the predictions (float64 predictions stay float64).

import numpy as np
from torchregress.test_time import ShiftFactoredPredictiveTransport, ShiftFactoredTransportConfig
from torchregress.prediction import PredictiveBatch

src_pb = PredictiveBatch(mean=src_mean, std=src_std)
config = ShiftFactoredTransportConfig(alpha=0.1, enable_alignment=True)
transport = ShiftFactoredPredictiveTransport(config)
transport.fit_source(src_pb, source_targets=src_y, source_inputs=src_x)
adapted = transport.adapt_unlabeled_target(target_predictions=tgt_pb, target_inputs=tgt_x)
calibrated = transport.calibrate_target(cal_pb, cal_y)  # fits conformal threshold
final = transport.predict(target_predictions=tgt_pb, target_inputs=tgt_x)

Quick example¶

import numpy as np
import torch
from torchregress.test_time import (
    BayesianLinearHead, ScoreCDFReweighter, WeightedSplitConformalAdapter,
    WeightedSubspaceMomentAligner, ParameterEMA,
)

# 1. Closed-form Bayesian linear regression on fixed features
head = BayesianLinearHead(in_features=8, out_features=1, noise_variance=0.1)
head.fit(features, labels)
batch = head.predictive_batch(query_features)   # mean / std / epistemic / aleatoric

# 2. OT-style conformal reweighting under target shift
rw = ScoreCDFReweighter().fit(cal_scores, tgt_scores)
adapter = WeightedSplitConformalAdapter(alpha=0.1).calibrate(cal_scores, rw.weights_)
mask = adapter.predict_from_test_scores(test_scores)

# 3. Feature alignment
aligner = WeightedSubspaceMomentAligner(rank=8)
aligner.fit(X_train, y_train)
X_test_aligned = aligner.transform(X_test)

# 4. EMA over model parameters
ema = ParameterEMA(decay=0.999)
ema.initialize(model)
for x, y in loader: ... ; ema.update(model)
ema.copy_to(model)

Next steps¶

• Bayesian Linear Regression
• Optimal Transport Conformal

Test-Time Adaptation Details¶

ScoreCDFReweighter¶

Learns target weights by minimising the Wasserstein/CDF gap between source and target predictions.

WeightedSplitConformalAdapter¶

Weighted split-conformal calibration using per-sample weights (e.g. from ScoreCDFReweighter).

BayesianLinearHead¶

Conjugate Gaussian Bayesian linear head for closed-form posterior updates and predictive uncertainty at test time.

ShiftFactoredPredictiveTransport¶

Orchestrates covariate shift correction, feature alignment, and conformal calibration.