Models API

Supervised

Supervised(config: Dict, in_channels: int, num_classes: int)

Bases: BaseSSLModel

Encoder + linear head, with CrossEntropyLoss or — task: regression — L1 loss.

For regression num_classes is the number of targets and data.y holds floats.

The labels the loss may see depend on the mode:

• Graph-level (encoder.pool=True): node embeddings are mean-pooled per graph and every graph in the batch contributes.
• Node-level mini-batch (batch.batch_size set by NeighborLoader): only the first batch_size seed nodes contribute. Build the loader with input_nodes=train_idx.
• Node-level full-batch: only data.train_mask nodes contribute. A full graph without train_mask raises rather than silently training on validation/test labels.

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against SupervisedConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required
num_classes	int	Number of output classes, resolved from the dataset.	required

Source code in src/graphssl/models/supervised.py

def __init__(self, config: Dict, in_channels: int, num_classes: int):
    super().__init__()
    cfg = SupervisedConfig.from_dict(config)
    self.encoder = cfg.encoder.build(in_channels)
    self.head = nn.Linear(cfg.encoder.hidden_dim, num_classes)
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout
    self.task: str = cfg.task

DGI

DGI(config: Dict, in_channels: int)

Bases: BaseSSLModel

Mutual information maximisation between node embeddings and global summary.

Discriminates real node embeddings from a corrupted-graph counterpart using a learnable bilinear discriminator W (DGI, Velickovic et al. 2019).

The summary vector s is computed as sigmoid(mean(h_pos)). For graph-level mode (encoder.pool=True) one summary per graph is computed and broadcast to the corresponding nodes via the batch index. On a NeighborLoader batch only the seed nodes (the first batch_size) enter the summary and the loss.

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against DGIConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

Source code in src/graphssl/models/dgi.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = DGIConfig.from_dict(config)
    self.encoder = cfg.encoder.build(in_channels)
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout
    self.corruption: str = cfg.corruption
    self.shuffle_ratio: float = cfg.shuffle_ratio
    hidden_dim = cfg.encoder.hidden_dim
    self.W = nn.Parameter(torch.empty(hidden_dim, hidden_dim))
    nn.init.xavier_uniform_(self.W)

GraphCL

GraphCL(config: Dict, in_channels: int)

Bases: BaseSSLModel

NT-Xent contrastive loss over two independently augmented views.

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against GraphCLConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

Source code in src/graphssl/models/graphcl.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = GraphCLConfig.from_dict(config)
    self.encoder = cfg.encoder.build(in_channels)
    hidden_dim = cfg.encoder.hidden_dim
    self.projector = Projector(hidden_dim, hidden_dim, cfg.proj_dim)
    self.loss_fn = NTXentLoss(tau=cfg.tau, chunk_size=cfg.loss_chunk_size)
    self.aug_list: List[Tuple[str, dict]] = [(a.name, a.kwargs) for a in cfg.augment]
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout

BGRL

BGRL(config: Dict, in_channels: int)

Bases: BaseSSLModel

Teacher-student SSL with EMA target encoder and cosine regression loss.

Key invariant: the target encoder is deepcopied then reset via reset_parameters(), giving it intentionally different initial weights. This asymmetry is critical for convergence (App. B, BGRL paper).

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against BGRLConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

Source code in src/graphssl/models/bgrl.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = BGRLConfig.from_dict(config)
    self.online_enc = cfg.encoder.build(in_channels)
    hidden_dim = cfg.encoder.hidden_dim
    self.online_pred = Predictor(
        hidden_dim,
        cfg.pred_hidden,
        hidden_dim,
        norm=cfg.pred_norm == "batch",
        activation=cfg.pred_activation,
    )

    # Target encoder: deepcopy + reset — intentionally different from online.
    self.target_enc = copy.deepcopy(self.online_enc)
    reset_fn = getattr(self.target_enc, "reset_parameters", None)
    if callable(reset_fn):
        reset_fn()
    for p in self.target_enc.parameters():
        p.requires_grad = False

    self._ema_tau: float = cfg.ema_tau
    self._ema_scheduler: Optional[CosineEMAScheduler] = (
        CosineEMAScheduler(cfg.ema_tau, cfg.ema_tau_end, cfg.total_steps)
        if cfg.total_steps > 0
        else None
    )
    self._step: int = 0
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout
    self.loss_fn = CosineRegressionLoss(symmetric=True)
    self.aug_list: List[Tuple[str, dict]] = [(a.name, a.kwargs) for a in cfg.augment]

AFGRL

AFGRL(config: Dict, in_channels: int)

Bases: BaseSSLModel

No-augmentation SSL: positive pairs mined via local adjacency and global k-means.

Identical to BGRL in the student/teacher setup, but replaces augmented views with node pairs identified as positives through: - Local: top-k cosine-similarity neighbours that are also graph-adjacent (kNN ∩ adj). - Global: top-k neighbours sharing a k-means cluster in ANY of the independent runs.

Key difference from BGRL: the target encoder starts with the same weights as the online encoder (deepcopy without reset). Diversity comes from the miner, not from weight asymmetry.

Graph-level mode: the PositiveMiner is topology-aware and not meaningful across independent graphs. In this mode it is skipped; a simple teacher-student loss on pooled embeddings is used instead.

Requires faiss-cpu: pip install faiss-cpu

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against AFGRLConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

Source code in src/graphssl/models/afgrl.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = AFGRLConfig.from_dict(config)
    self.online_enc = cfg.encoder.build(in_channels)
    hidden_dim = cfg.encoder.hidden_dim
    self.online_pred = Predictor(
        hidden_dim,
        cfg.pred_hidden,
        hidden_dim,
        norm=cfg.pred_norm == "batch",
        activation=cfg.pred_activation,
    )

    # Target encoder: same initial weights as online (no reset — unlike BGRL).
    self.target_enc = copy.deepcopy(self.online_enc)
    for p in self.target_enc.parameters():
        p.requires_grad = False

    self._ema_tau: float = cfg.ema_tau
    self._ema_scheduler: Optional[CosineEMAScheduler] = (
        CosineEMAScheduler(cfg.ema_tau, cfg.ema_tau_end, cfg.total_steps)
        if cfg.total_steps > 0
        else None
    )
    self._step: int = 0
    self._topk: int = cfg.topk
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout
    self._graph_loss = CosineRegressionLoss(symmetric=False)
    self.positive_miner = PositiveMiner(
        num_centroids=cfg.num_centroids,
        num_kmeans=cfg.num_kmeans,
        clus_num_iters=cfg.clus_num_iters,
        kmeans_threads=cfg.kmeans_threads,
        knn_chunk_size=cfg.knn_chunk_size,
    )

VICReg

VICReg(config: Dict, in_channels: int)

Bases: BaseSSLModel

Symmetric augmentation with triple VICReg loss (invariance + variance + covariance).

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against VICRegConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

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

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = VICRegConfig.from_dict(config)
    self.encoder = cfg.encoder.build(in_channels)
    hidden_dim = cfg.encoder.hidden_dim
    self.projector = Projector(hidden_dim, hidden_dim * 2, cfg.proj_dim, num_layers=3)
    self.loss_fn = VICRegLoss(
        invariance=cfg.invariance,
        variance=cfg.variance,
        covariance=cfg.covariance,
    )
    self.aug_list: List[Tuple[str, dict]] = [(a.name, a.kwargs) for a in cfg.augment]
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout

BarlowTwins

BarlowTwins(config: Dict, in_channels: int)

Bases: BaseSSLModel

Barlow Twins cross-correlation loss with symmetric augmentation.

Parameters:

Name	Type	Description	Default
config	Dict	Config dict validated against BarlowTwinsConfig.	required
in_channels	int	Node feature dimensionality, resolved from the dataset.	required

Source code in src/graphssl/models/barlow_twins.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = BarlowTwinsConfig.from_dict(config)
    self.encoder = cfg.encoder.build(in_channels)
    hidden_dim = cfg.encoder.hidden_dim
    self.projector = Projector(hidden_dim, hidden_dim * 2, cfg.proj_dim, num_layers=3)
    self.loss_fn = BarlowTwinsLoss(lambda_param=cfg.lambda_param, out_dim=cfg.proj_dim)
    self.aug_list: List[Tuple[str, dict]] = [(a.name, a.kwargs) for a in cfg.augment]
    self.graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout

GraphDINO

GraphDINO(config: Dict, in_channels: int)

Bases: BaseSSLModel

Parameters:

Name	Type	Description	Default
config	Dict	Config dict (see graphdino.yaml); validated against GraphDINOConfig.	required
in_channels	int	Node feature size, resolved from the dataset at runtime.	required

Source code in src/graphssl/models/graphdino.py

def __init__(self, config: Dict, in_channels: int):
    super().__init__()
    cfg = GraphDINOConfig.from_dict(config)

    self.student_enc = cfg.encoder.build(in_channels)
    # _build_head() returns nn.Module (it goes through the HEADS registry, which is
    # deliberately generic — see registry pattern). GraphDINO itself only ever
    # requests the "dino" head, so this cast documents a real, narrower invariant
    # rather than papering over one; it's what unlocks DINOHead-specific access
    # below (.center, .set_epoch(), .update_center(), .cancel_last_layer_gradients()).
    self.student_head = cast(DINOHead, self._build_head(cfg.head, cfg.encoder.hidden_dim))

    # Build fresh instances rather than deepcopy: weight_norm creates non-leaf
    # tensors that break copy.deepcopy.
    self.teacher_enc = cfg.encoder.build(in_channels)
    self.teacher_head = cast(DINOHead, self._build_head(cfg.head, cfg.encoder.hidden_dim))
    self.teacher_enc.load_state_dict(self.student_enc.state_dict())
    self.teacher_head.load_state_dict(self.student_head.state_dict())
    for p in self.teacher_enc.parameters():
        p.requires_grad = False
    for p in self.teacher_head.parameters():
        p.requires_grad = False

    self.aug_list_teacher: List[Tuple[str, dict]] = [
        (a.name, a.kwargs) for a in cfg.augment_teacher
    ]
    self.aug_list_student: List[Tuple[str, dict]] = [
        (a.name, a.kwargs) for a in cfg.augment_student
    ]
    self._n_views: int = cfg.n_views
    self._n_global_views: int = cfg.n_global_views
    # pool=True means the encoder returns graph-level embeddings; no batch_size crop needed
    self._graph_level: bool = cfg.encoder.pool
    self.readout: str = cfg.encoder.readout

    self._ema_tau: float = cfg.ema_tau_base
    self._loss_fn = DINOLoss()

    self._freeze_last_layer_epochs: int = cfg.freeze_last_layer_epochs
    self._epoch: int = 0

    if cfg.total_steps > 0 and cfg.ema_tau_base < cfg.ema_tau:
        self._ema_scheduler: CosineEMAScheduler | None = CosineEMAScheduler(
            ema_base=cfg.ema_tau_base,
            ema_end=cfg.ema_tau,
            total_steps=cfg.total_steps,
        )
    else:
        self._ema_scheduler = None
    self._step: int = 0

    self._last_teacher_out: Tensor | None = None
    self._last_teacher_logits: Tensor | None = None