Losses API

NTXentLoss

NTXentLoss(tau: float = 0.5, chunk_size: Optional[int] = 4096)

Bases: Module

Symmetric contrastive loss over two augmented views.

Positive pair: (z1_i, z2_i). All other pairs in the batch are negatives.

Parameters:

Name	Type	Description	Default
tau	float	Temperature. Lower = harder negatives. Default 0.5 (GraphCL paper).	0.5
chunk_size	Optional[int]	Max rows of the [2N, 2N] similarity matrix materialized at once. When 2N exceeds it, the loss is computed in row chunks with gradient checkpointing: exactly the same value and gradients, but peak memory O(chunk_size * 2N) instead of O((2N)^2). Full-batch PubMed (2N = 39,434) needs well over 16 GB unchunked. None always materializes the full matrix.	4096

Source code in src/graphssl/losses/nt_xent.py

def __init__(self, tau: float = 0.5, chunk_size: Optional[int] = 4096):
    super().__init__()
    if chunk_size is not None and chunk_size <= 0:
        raise ValueError(f"chunk_size must be > 0 or None, got {chunk_size}")
    self.tau = tau
    self.chunk_size = chunk_size

VICRegLoss

VICRegLoss(invariance: float = 25.0, variance: float = 25.0, covariance: float = 1.0, gamma: float = 1.0)

Bases: Module

Triple loss: invariance + variance + covariance.

Parameters:

Name	Type	Description	Default
invariance	float	Weight for the MSE invariance term (default 25.0).	25.0
variance	float	Weight for the per-dimension variance term (default 25.0).	25.0
covariance	float	Weight for the off-diagonal covariance term (default 1.0).	1.0
gamma	float	Target std per dimension for the variance term (default 1.0).	1.0

Source code in src/graphssl/losses/vicreg.py

def __init__(
    self,
    invariance: float = 25.0,
    variance: float = 25.0,
    covariance: float = 1.0,
    gamma: float = 1.0,
):
    super().__init__()
    self.inv_w = invariance
    self.var_w = variance
    self.cov_w = covariance
    self.gamma = gamma

BarlowTwinsLoss

BarlowTwinsLoss(lambda_param: float | None = None, out_dim: int = 256)

Bases: Module

Cross-correlation matrix loss.

On-diagonal terms are pushed to 1; off-diagonal terms are pushed to 0.

Parameters:

Name	Type	Description	Default
lambda_param	float | None	Weight for off-diagonal terms. Defaults to 1/out_dim.	None
out_dim	int	Projection dimension used to compute the default lambda.	256

Source code in src/graphssl/losses/barlow.py

def __init__(self, lambda_param: float | None = None, out_dim: int = 256):
    super().__init__()
    self.lambda_param = lambda_param if lambda_param is not None else 1.0 / out_dim

CosineRegressionLoss

CosineRegressionLoss(symmetric: bool = True)

Bases: Module

Cosine similarity regression loss.

Symmetric (BGRL): loss = 2 - cos_sim(p1, t2) - cos_sim(p2, t1)

Asymmetric / single-sided (AFGRL): loss = 1 - cos_sim(p, t).mean()

Parameters:

Name	Type	Description	Default
symmetric	bool	If True (default) expects all four tensors (p1, t1, p2, t2).	True

Source code in src/graphssl/losses/regression.py

def __init__(self, symmetric: bool = True):
    super().__init__()
    self.symmetric = symmetric

DINOLoss

Bases: Module

Cross-entropy between teacher softmax and student log-softmax over N views.

Inputs are pre-processed by DINOHead.forward(): - student_out: log-softmax with student_temp, shape [n_views * batch, D] - teacher_out: softmax with center subtraction and teacher_temp, shape [n_global_views * batch, D]

For every (teacher_i, student_j) pair with i != j, computes H(teacher_i, student_j) and averages over all such pairs. When n_global_views == n_views == 2 this reduces to −0.5 * ((t2 * s1).sum(−1).mean() + (t1 * s2).sum(−1).mean()).

CombinedLoss

CombinedLoss(losses: List[Tuple[Module, float]])

Bases: Module

Weighted sum of multiple loss functions.

Weights are plain scalars and need not sum to 1. All component losses must share the same call signature.

Source code in src/graphssl/losses/combined.py

def __init__(self, losses: List[Tuple[nn.Module, float]]):
    super().__init__()
    if not losses:
        raise ValueError("CombinedLoss requires at least one (loss, weight) pair.")
    self._loss_modules = nn.ModuleList([m for m, _ in losses])
    self._weights = [w for _, w in losses]

from_config classmethod

from_config(config_list: list) -> 'CombinedLoss'

Build from a list of dicts with 'name', 'weight', and loss-specific kwargs.

Source code in src/graphssl/losses/combined.py

@classmethod
def from_config(cls, config_list: list) -> "CombinedLoss":
    """Build from a list of dicts with 'name', 'weight', and loss-specific kwargs."""
    pairs: List[Tuple[nn.Module, float]] = []
    for entry in config_list:
        entry = dict(entry)
        name = entry.pop("name")
        weight = float(entry.pop("weight", 1.0))
        pairs.append((LOSSES.build(name, **entry), weight))
    return cls(pairs)