swag

probly.method.swag(base: Predictor[In, Out], max_rank: int = 20, scale: float = 0.5, rngs: Rngs | RngStream | int = 0) → SWAGPredictor[In, Out][source]

Create a SWAG predictor from a base predictor based on [MIG+19].

SWAG (SWA-Gaussian) fits a Gaussian distribution to the weights visited by SGD: the mean is the stochastic weight average (SWA) and the covariance is the sum of a diagonal term and a low-rank term formed from the last max_rank weight snapshots. Train the returned predictor as usual and call collect_swag() periodically to record snapshots. At prediction time, drawing repeated predictions (e.g. via representer(model, num_samples=...)) samples a fresh weight vector from the fitted Gaussian for every forward pass.

Parameters:

  • base – The base model to be used for SWAG.

  • max_rank – Maximum number of columns of the low-rank deviation matrix. Set to 0 for a diagonal-only (SWAG-Diag) posterior. Default is 20.

  • scale – Scaling factor applied to the sampled weight perturbations. The default of 0.5 is the 1/2 covariance scaling the paper uses for the full posterior; its diagonal-only variant samples without this factor, which corresponds to a scale of 1.0.

  • rngs – Optional rngs for the sampling randomness of the flax backend; the torch backend uses the global torch generator instead and ignores this. Default is 0.

Returns:

The SWAG predictor.

Examples using probly.method.swag

SWAG on Two Moons

SWAG on Two Moons

SWAG on MNIST

SWAG on MNIST