VISReg#

class stable_pretraining.methods.VISReg(encoder_name: str = 'vit_base_patch16_224', projector: Module | None = None, num_projections: int = 256, lambda_scale: float = 1.0, lambda_shape: float = 1.0, lambda_center: float = 1.0, lamb: float = 0.02, pretrained: bool = False, drop_path_rate: float = 0.1)[source]#

Bases: Module

VISReg: multi-view invariance + centering + variance + sliced wasserstein distance.

Architecture:

Backbone: timm ViT (CLS-pooled, num_classes=0)
Projector: MLP projection head
Loss: (1 - λ) * invariance + (λ * VISReg)

Centers are computed from global-view projections only. The invariance term penalizes the MSE between each view’s projection and the center. The VISReg term is a sliced goodness-of-fit regularizer that pushes projected embeddings toward an isotropic standard Gaussian, applied per view over the batch dimension.

Parameters:

encoder_name – timm model name (e.g., "vit_base_patch16_224")
projector – Optional projection head. When None, a 3-layer BN+ReLU MLP (embed_dim → 2048 → 2048 → 512) is created.
num_projections – Random projection directions for the shape term (default: 256)
lambda_scale – Weight on the VISReg scale (unit-variance) term (default: 1.0)
lambda_shape – Weight on the VISReg shape (sliced-quantile) term (default: 1.0)
lambda_center – Weight on the VISReg center (zero-mean) term (default: 1.0)
lamb – Convex mixing weight λ ∈ [0, 1] between the invariance and VISReg terms: (1 - λ) * inv + λ * visreg (default: 0.02)
pretrained – Load pretrained timm weights
drop_path_rate – Stochastic depth rate for the backbone (default: 0.1)

Example:

model = VISReg("vit_base_patch16_224")
images = torch.randn(4, 3, 224, 224)

model.train()
output = model(
    global_views=[images, images],
    local_views=[images, images],
)
output.loss.backward()

model.eval()
output = model(images=images)
features = output.embedding  # [4, 768]

Example with Lightning:

import lightning as pl
from stable_pretraining.methods import VISReg


class VISRegLightning(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = VISReg("vit_base_patch16_224")

    def training_step(self, batch, batch_idx):
        views = [v["image"] for v in batch["views"]]
        output = self.model(global_views=views, local_views=views)
        self.log("loss", output.loss)
        return output.loss

    def configure_optimizers(self):
        return torch.optim.AdamW(self.parameters(), lr=1e-3)

forward(global_views: list[Tensor] | None = None, local_views: list[Tensor] | None = None, images: Tensor | None = None) → VISRegOutput[source]#

Same as torch.nn.Module.forward().

Parameters:

*args – Whatever you decide to pass into the forward method.
**kwargs – Keyword arguments are also possible.

Returns:

Your model’s output

VISReg

Contents

VISReg#