Semi-Supervised Regression (teacher–student SSL)¶

← Causal Inference | PPI Inference →

The semi-supervised module provides modular workflows for training regression models with both labeled and unlabeled data. It implements teacher–student consistency with distributional pseudo-labeling for semi-supervised regression.

Research track (SAGE-Reg)

The SAGE-Reg paper track (benchmarks, manuscript claims, and committed result tables) lives in torchregress-research. The API here is stable enough for experimentation and demos, but it is not held to the same battle-tested bar as core point losses, conformal wrappers, and ensemble utilities in the main library.

When to use

Semi-supervised regression is valuable when you have abundant unlabeled data (\(N_{\text{unlabeled}} \gg N_{\text{labeled}}\)) and believe the unlabeled distribution carries information about the conditional target distribution \(p(y \mid x)\).

Core Concepts¶

Teacher-Student Consistency¶

The training loop alternates between:

Supervised step: Train the student model on labeled data using a standard regression loss
Unsupervised step: Generate multiple predictive views of unlabeled inputs using a teacher model (EMA of the student), build a consensus distribution, and penalize the student for disagreeing with the consensus

Predictive Views¶

A predictive view is a PredictiveBatch — a distributional prediction from one stochastic forward pass of the model. Multiple views are generated by:

Multiple forward passes of a stochastic model (MC Dropout, BNN)
Applying different augmentations to the input
Using ensemble members as different views

Agreement Scoring¶

Disagreement between views is measured as the average pairwise symmetric KL divergence between the predictive densities on a common support grid:

\[\text{disagreement} = \frac{2}{M(M-1)} \sum_{i < j} \frac{1}{2} \left( \text{KL}(p_i \parallel p_j) + \text{KL}(p_j \parallel p_i) \right)\]

Low disagreement → the model is consistent → the consensus pseudo-label is trusted with high weight.

Core API Components¶

`PredictiveBatch`¶

The universal container for distributional predictions. Supports multiple head types:

Head Type	Required Fields	Example
Gaussian	`mean`, `std`	`PredictiveBatch(mean=mu, std=sigma)`
Quantile	`quantiles`, `quantile_levels`	`PredictiveBatch(quantiles=q_vals, quantile_levels=[0.1,0.5,0.9])`
Binned PDF	`bar_logits`, `bin_edges`	`PredictiveBatch(bar_logits=logits, bin_edges=edges)`
Density	`support`, `density`	`PredictiveBatch(support=grid, density=pdf)`
Point	`point`	`PredictiveBatch(point=y_hat)`

`TeacherStudentTrainer`¶

The main training orchestrator. Injects custom sample weighting policies:

from torchregress.semi_supervised import TeacherStudentTrainer

trainer = TeacherStudentTrainer(
    optimizer=optimizer,
    supervised_loss_fn=supervised_loss,
    predictive_batch_fn=predictive_fn,
    augment_fn=augment_fn,       # Optional input augmentation
    n_views=4,                    # Number of stochastic views
    agreement_weight=1.0,         # Weight of unsupervised loss term
    ema_decay=0.99,               # Teacher EMA momentum
)

history = trainer.fit(
    model,
    labeled_loader=labeled_loader,
    unlabeled_loader=unlabeled_loader,
    epochs=50,
    lr_schedule="cosine",
)

Parameter	Type	Default	Description
`optimizer`	`Optimizer`	—	PyTorch optimizer
`supervised_loss_fn`	`callable`	—	`(model, x, y) -> loss`
`predictive_batch_fn`	`callable`	—	`(model, x) -> PredictiveBatch`
`augment_fn`	`callable` or `None`	`identity`	Input augmentation for view diversity
`n_views`	`int`	`4`	Number of predictive views per unlabeled batch
`agreement_weight`	`float`	`1.0`	Weight of unsupervised loss
`ema_decay`	`float` or `None`	`0.99`	Teacher EMA momentum (None → no teacher)
`n_support`	`int`	`128`	Grid points for density evaluation
`range_margin`	`float`	`0.05`	Margin beyond prediction range
`gaussian_std_span`	`float`	`4.0`	Std spans for Gaussian-based bounds
`min_scale`	`float`	`1e-4`	Minimum scale for numerical stability
`detach_weights`	`bool`	`True`	Detach sample weights from computation graph
`tau`	`float`	`0.2`	Temperature of the default trust weights \(w = \exp(-d/\tau)\)
`weight_power`	`float`	`1.0`	Exponent \(p\) applied to the trust weights
`hard_weight_threshold`	`float` or `None`	`None`	Zero out weights below this value
`batch_relative_mode`	`bool`	`False`	z-score the disagreement within each batch before weighting
`batch_trust_top_k`	`int` or `None`	`None`	Keep only the \(k\) lowest-disagreement samples per batch

The weighting arguments (tau … batch_trust_top_k) configure the default disagreement_to_weight policy and are ignored when a custom sample_weight_fn is passed. Unknown keyword arguments raise TypeError.

`SAGERegLoss`¶

Standalone loss module for manual training loops:

from torchregress.semi_supervised import SAGERegLoss

sage_loss = SAGERegLoss(
    tau=0.2,                    # Temperature for weight computation
    agreement_weight=1.0,
    n_support=128,
)

# In training loop:
output = sage_loss(supervised_loss, unlabeled_views)
total_loss = output.total_loss  # supervised_loss + weighted unsupervised
# output.agreement.disagreement → per-sample disagreement scores
# output.agreement.weights → per-sample trust weights

Weighting Functions¶

Flexible functions to convert prediction quality signals into sample weights:

Function	Input	Weight Logic	Use Case
`disagreement_to_weight`	Disagreement scores	\(w = \exp(-d / \tau)\)	Default SAGE-Reg weighting
`uncertainty_to_weight`	`PredictiveBatch` std	\(w = \exp(-\sigma / \tau)\)	Teacher uncertainty gating
`conformal_width_to_weight`	Conformal interval bounds	\(w = \exp(-\text{width} / \tau)\) or binary mask	Width-based pseudo-label gating

from torchregress.semi_supervised import disagreement_to_weight, conformal_width_to_weight

# Soft weighting: trust samples where views agree
weights = disagreement_to_weight(disagreement, tau=0.2, power=2.0)

# Hard gating: only use samples with narrow conformal intervals
mask = conformal_width_to_weight(lower, upper, threshold=1.5)

# Batch-relative: trust the top-k most consistent samples
weights = disagreement_to_weight(
    disagreement, tau=0.2,
    batch_trust_top_k=32,
)

`predictive_agreement_score`¶

Compute disagreement between any set of PredictiveBatch views:

from torchregress.semi_supervised import predictive_agreement_score

disagreement = predictive_agreement_score(
    predictive_views,
    n_support=128,
    reduction="none",  # Per-sample scores
)

`build_consensus_predictive_batch`¶

Build a single consensus distribution from multiple predictive views:

from torchregress.semi_supervised import build_consensus_predictive_batch

consensus = build_consensus_predictive_batch(views, n_support=128)
# consensus.mean → consensus mean
# consensus.std → consensus std
# consensus.density → consensus density on support grid

Complete Example¶

import torch
import torch.nn as nn
from torchregress.semi_supervised import (
    TeacherStudentTrainer,
    PredictiveBatch,
)
from torchregress.losses import GaussianNLLLoss

torch.manual_seed(42)

# 1. Heteroscedastic model
class HeteroModel(nn.Module):
    def __init__(self, in_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, 64), nn.ReLU(),
            nn.Linear(64, 32), nn.ReLU(),
        )
        self.mean_head = nn.Linear(32, 1)
        self.logvar_head = nn.Linear(32, 1)

    def forward(self, x):
        h = self.net(x)
        return self.mean_head(h), self.logvar_head(h)

# 2. Predictive batch function
def predictive_batch_fn(model, x):
    mean, logvar = model(x)
    return PredictiveBatch(
        mean=mean,
        std=torch.exp(0.5 * logvar),
    )

# 3. Supervised loss
def supervised_loss_fn(model, x, y):
    mean, logvar = model(x)
    loss_fn = GaussianNLLLoss()
    return loss_fn((mean, logvar), y)

# 4. Simple Gaussian noise augmentation
def augment_fn(x):
    return x + 0.01 * torch.randn_like(x)

# 5. Generate data
n_labeled, n_unlabeled = 200, 2000
X_all = torch.randn(n_labeled + n_unlabeled, 5)
y_labeled = 2.0 * X_all[:n_labeled, [0]] + 0.3 * torch.randn(n_labeled, 1)

labeled_dataset = torch.utils.data.TensorDataset(
    X_all[:n_labeled], y_labeled
)
unlabeled_dataset = torch.utils.data.TensorDataset(
    X_all[n_labeled:]
)

labeled_loader = torch.utils.data.DataLoader(labeled_dataset, batch_size=32)
unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=128)

# 6. Train
model = HeteroModel(5)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

trainer = TeacherStudentTrainer(
    optimizer=optimizer,
    supervised_loss_fn=supervised_loss_fn,
    predictive_batch_fn=predictive_batch_fn,
    augment_fn=augment_fn,
    n_views=4,
    agreement_weight=1.0,
    ema_decay=0.99,
)

history = trainer.fit(
    model,
    labeled_loader=labeled_loader,
    unlabeled_loader=unlabeled_loader,
    epochs=30,
    lr_schedule="cosine",
)

print(f"Final total loss: {history['total_loss'][-1]:.4f}")
print(f"Mean trust weight: {history['mean_weight'][-1]:.4f}")

When to Use¶

Semi-supervised regression helps when

You have many unlabeled samples (\(N_{\text{unlabeled}} \gg N_{\text{labeled}}\))
Your model is stochastic (MC Dropout, BNN) or you can create diverse views via augmentations
The unlabeled distribution is broadly similar to the labeled distribution (no extreme covariate shift)
You have a well-calibrated supervised baseline to compare against

Semi-supervised regression can hurt when

The labeled set is already large enough that the supervised model is near-optimal
Unlabeled data comes from a very different distribution (covariate shift without correction)
The model is deterministic with no augmentation diversity → views collapse to identical predictions
Hyperparameters (tau, agreement_weight, n_views) are not tuned

Limitations¶

Practical constraints

Confirmation bias: The teacher model (EMA of student) can reinforce the student's errors if not properly regularized. Use detach_weights=True and monitor disagreement scores.
View diversity is essential: If all predictive views are identical, disagreement = 0 everywhere and the unsupervised loss provides no signal. Ensure stochasticity via MC Dropout, BNN sampling, or meaningful augmentations.
Density estimation cost: Computing pairwise KL on a support grid of n_support=128 points requires numerical integration for every pair of views, per unlabeled batch. With \(M\) views, this is \(\mathcal{O}\binom{M}{2} \cdot \text{n\_support} \cdot \text{batch\_size}\) trapezoidal integrations — non-trivial overhead, especially with the default n_views=4 (6 pairs).
Hyperparameter sensitivity: tau (temperature) and agreement_weight heavily influence the tradeoff. Too small tau → only a few samples trusted; too large → noisy pseudo-labels dominate.
No coverage guarantees: Unlike conformal prediction, SAGE-Reg provides no finite-sample coverage guarantees. Validate prediction intervals empirically.
EMA teacher cost: The teacher is a copy.deepcopy(model), which is memory-intensive for large models. The default ema_decay=0.99 updates the teacher every step via update_ema_teacher_.

Recommendations¶

Default trainer: Use TeacherStudentTrainer with default settings (n_views=4, ema_decay=0.99, detach_weights=True). It provides a solid baseline for most semi-supervised regression problems.
View diversity is essential: Ensure your model or augmentation pipeline produces genuinely different predictive views. For deterministic models, add Gaussian input noise (augment_fn) or use MC Dropout. Monitor disagreement scores — if they are near zero, views are collapsing.
Tune tau and agreement_weight: Start with tau=0.2 and agreement_weight=1.0. If the unsupervised loss dominates, reduce agreement_weight. If too few samples are trusted, increase tau.
Validate with held-out labeled data: Never evaluate on pseudo-labeled samples. Maintain a clean labeled test set and compare against a supervised-only baseline trained on the same labeled data.
For simpler workflows: If you don't need custom weighting policies, use SAGERegLoss directly in your training loop — it handles consensus building, weighting, and pseudo-loss computation internally.
Combine with conformal calibration: SAGE-Reg provides no coverage guarantees. For prediction intervals with finite-sample coverage, wrap the trained model with SplitConformal or CQR.

Next steps¶

Noisy label losses — complementary approach for explicitly corrupted target values
Uncertain ground truth — pseudo-label strategies when labels are known-uncertain
Conformal prediction — coverage-guaranteed intervals that complement heuristic SSL consistency
PPI inference — valid statistical inference combining small labeled + large unlabeled sets

References¶

#	Reference
1	SAGE-Reg manuscript and reproduction: `torchregress-research`.
2	D.-H. Lee. "Pseudo-Label: The Simple and Efficient Semi-Supervised Learning Method for Deep Neural Networks." ICML Workshop, 2013.
3	A. Tarvainen, H. Valpola. "Mean Teachers are Better Role Models." NeurIPS, 2017.