sngp

probly.method.sngp(base: Predictor[In, Out], name: str = 'weight', n_power_iterations: int = 1, norm_multiplier: float = 6.0, eps: float = 1e-12, num_random_features: int = 1024, ridge_penalty: float = 1.0, momentum: float = -1.0, random_feature_init_std: float = 1.0) → SNGPPredictor[In, Out][source]

Wrap a model with SNGP (Spectral-normalized Neural Gaussian Process).

Replaces the last nn.Linear with an SNGPLayer (random Fourier features + Laplace-approximated Gaussian process) and registers a spectral-norm parametrization on every preceding nn.Linear and nn.Conv2d. Defaults match the ImageNet ResNet-50 baseline at google/uncertainty-baselines/baselines/imagenet/sngp.py.

Parameters:

  • base – The model to wrap.

  • name – The name of the weight parameter to spectrally normalize on non-output layers. Defaults to "weight".

  • n_power_iterations – Power iterations per training step for the spectral-norm estimate. Defaults to 1. (A 15-iteration warmup runs once at construction time, independent of this value.)

  • norm_multiplier – Upper bound on each non-output layer’s spectral norm. Defaults to 6.0 (high because ResNet’s BatchNorm already applies its own Lipschitz scaling).

  • eps – Small constant to stabilize the spectral-norm denominator. Defaults to 1e-12.

  • num_random_features – Dimensionality of the random Fourier feature map. Defaults to 1024.

  • ridge_penalty – Ridge factor used inside the covariance inversion inv(ridge * I + precision). Defaults to 1.0.

  • momentum – Discount factor for the precision-matrix update. Default -1.0 triggers accumulate mode (paper Algorithm 1; imagenet gp_cov_discount_factor=-1); the user must call reset_precision_matrix() at the start of each training epoch in this mode. Pass momentum > 0 for EMA mode (no reset needed; matches the CLINC reference’s gp_cov_momentum=0.999).

  • random_feature_init_std – Standard deviation of the Gaussian used to initialize the frozen random projection W_L. Defaults to 1.0 (paper / imagenet / Edward2; full RFF kernel approximation, expects from-scratch training). Set to 0.05 (matching untangle.wrappers.sngp_wrapper) when fine-tuning from a pretrained backbone: keeps W_L^T h in the near-linear regime of cos so pretrained-feature signal flows through the RFF, at the cost of a longer effective kernel lengthscale and weaker distance-aware uncertainty.

Returns:

An SNGPPredictor whose predict(...) returns a GaussianDistribution over logits.

Examples using probly.method.sngp

SNGP Distance Awareness on 2D Toys

SNGP Distance Awareness on 2D Toys

SNGP on MNIST

SNGP on MNIST