netcl wiki
api

netcl.ssl: Self-Supervised Learning & Deep Clustering

netcl.ssl: Self-Supervised Learning & Deep Clustering

The netcl.ssl module provides end-to-end Self-Supervised Learning (SSL) training architectures and deep clustering algorithms. It couples backbone encoders with device-side image augmentations, contrastive or non-contrastive loss objectives, and evaluation probes.

[!TIP] Theory & Curriculum Link: For the first-principles mathematical derivation of InfoNCE alignment, representation collapse, and BYOL stop-gradient dynamics, see Curriculum: Self-Supervised Representation Learning.


1. Quick Example: SimCLR Pretraining

import netcl.autograd as ag
import netcl.optim as optim
from netcl import ssl
from netcl.core.device import manager

q = manager.default("auto").queue

# 1. GPU view generation: 2 random augmented views per image
views = ssl.MultiView.two_crop(out_size=(32, 32), seed=0)

# 2. Wrap backbone encoder with SimCLR model
method = ssl.SimCLR(encoder, feature_dim=128, projection_dim=64, queue=q)
opt = optim.AdamW(method.parameters(), lr=1e-3)

# 3. Training iteration
for batch in dataloader:
    v1, v2 = views(batch)
    with ag.Tape() as tape:
        loss = method.training_step([v1, v2])
    tape.backward(loss)
    opt.step()
    opt.zero_grad()
    method.on_step_end()

2. Method Selection Guide

Class Category Requires Negative Pairs? Target Encoder? Target Application
SimCLR Contrastive Yes (In-Batch) No High-memory setups supporting large batches ($B \ge 128$)
MoCo Contrastive Yes (Memory Queue) Yes (EMA) Small-batch setups ($B \le 64$) maintaining negative FIFO bank
BYOL Non-Contrastive No Yes (EMA) Asymmetric prediction head avoiding representation collapse
SimSiam Non-Contrastive No No (Stop-Grad) Minimalist bootstrap architecture without EMA overhead
BarlowTwins Redundancy Reduction No No Forces cross-correlation matrix toward identity
VICReg Variance Regularization No No Explicit variance, invariance, and covariance bounds
DeepCluster Iterative Clustering No No Couples GPU KMeans pseudo-labeling with supervised training
SCAN Graph Clustering No No Enforces nearest neighbors in feature space to share clusters

Self-Supervised Learning: Contrastive Hyperspheres vs Target Network EMA


3. Contrastive Methods

SimCLR

Optimizes the Normalized Temperature-scaled Cross Entropy (nt_xent) loss across positive and negative pairs:

simclr = ssl.SimCLR(
    encoder,
    feature_dim=512,          # Output dimension of backbone encoder
    projection_dim=128,       # Dimension of MLP projection head
    temperature=0.1,          # Softmax temperature scale
    queue=q,
)

MoCo (Momentum Contrast)

Maintains a dynamic FIFO dictionary queue of negative sample representations, decoupling dictionary size from mini-batch size:

moco = ssl.MoCo(
    encoder,
    target_encoder,           # Identical architecture updated via EMA
    feature_dim=512,
    projection_dim=128,
    queue_size=4096,          # Capacity of negative sample queue
    momentum=0.999,           # EMA update rate
    queue=q,
)

4. Non-Contrastive Methods

BYOL (Bootstrap Your Own Latent)

An online network learns to predict the representation of a target network updated via Exponential Moving Average:

byol = ssl.BYOL(
    encoder,
    feature_dim=512,
    projection_dim=256,
    pred_dim=256,
    momentum=0.996,
    queue=q,
)

BarlowTwins

Forces the empirical cross-correlation matrix between dual view representations toward the identity matrix: - Diagonal elements equal 1: Invariance to transformations. - Off-diagonal elements equal 0: Redundancy reduction, ensuring feature coordinates encode independent information.

barlow = ssl.BarlowTwins(encoder, feature_dim=512, proj_dim=2048, lambd=0.005, queue=q)

VICReg

Prevents collapse using three explicit mathematical bounds: - Variance: Maintains feature standard deviation $\ge 1$ across all coordinates. - Invariance: Minimizes mean squared distance between augmented views. - Covariance: Minimizes off-diagonal cross-correlation between different coordinates.

vicreg = ssl.VICReg(encoder, feature_dim=512, proj_dim=2048, queue=q)

5. Deep Clustering

DeepCluster

Iteratively clusters feature representations using GPU KMeans and trains the encoder using Cross-Entropy on these pseudo-labels:

dc = ssl.DeepCluster(
    encoder,
    feature_dim=512,
    n_clusters=100,
    recluster_every=1,        # Frequency of KMeans re-clustering in epochs
    balanced=True,            # Employs Sinkhorn equipartition against collapse
    queue=q,
)

6. Device-Side Image Augmentations (netcl.ssl.views)

All image transformations execute directly inside OpenCL GPU buffers without CPU memory transfer bottlenecks:

from netcl.ssl import views

pipeline = views.AugmentationPipeline([
    lambda x: views.random_resized_crop(x, size=(32, 32), scale=(0.2, 1.0)),
    lambda x: views.random_flip(x, p=0.5),
    lambda x: views.color_jitter(x, brightness=0.4, contrast=0.4, p=0.8),
    lambda x: views.random_grayscale(x, p=0.2),
])

mv = views.MultiView(pipelines=[pipeline, pipeline])
v1, v2 = mv(batch)  # Yields two transformed GPU tensors directly in VRAM

7. Diagnostic Evaluation Probes (netcl.ssl.evaluate)

from netcl.ssl import embed_dataset, knn_monitor, linear_probe

# 1. Extract feature representations
X_train, y_train = embed_dataset(encoder, train_loader)
X_test,  y_test  = embed_dataset(encoder, test_loader)

# 2. Online kNN Monitor: evaluates feature quality without training
acc = knn_monitor(X_train, y_train, X_test, y_test, k=20)
print(f"k-NN Test Accuracy: {acc * 100:.2f}%")

# 3. Linear Probe: trains linear classifier on frozen representations
probe_acc = linear_probe(X_train, y_train, X_test, y_test, epochs=15)
print(f"Linear Probe Accuracy: {probe_acc * 100:.2f}%")