from itertools import chain from functools import partial import torch import torch.nn as nn import torch.nn.functional as F from .acm_gin import ACM_GIN_model from .acm_gineconv import ACM_GINEConv_model from torch_geometric.utils import dropout_edge from torch_geometric.utils import add_self_loops from torch_geometric.nn import global_mean_pool def sce_loss(x, y, alpha=3): x = F.normalize(x, p=2, dim=-1) y = F.normalize(y, p=2, dim=-1) loss = (1 - (x * y).sum(dim=-1)).pow_(alpha) loss = loss.mean() return loss def setup_module( m_type, in_dim, out_dim, num_hidden, num_layers, activation, batchnorm, edge_in_dim, norm_type="none", input_norm="none", edge_distance_in_proj=True, ) -> nn.Module: if m_type == "acm_gin": mod = ACM_GIN_model( int(in_dim), int(out_dim), num_layers, int(num_hidden), edge_in_dim, batchnorm, activation=activation, ) elif m_type == "acm_gineconv": mod = ACM_GINEConv_model( int(in_dim), int(out_dim), num_layers, int(num_hidden), edge_in_dim, batchnorm, activation=activation, norm_type=norm_type, input_norm=input_norm, edge_distance_in_proj=edge_distance_in_proj, ) else: raise NotImplementedError return mod class PreModel(nn.Module): def __init__( self, in_dim: int, edge_in_dim: int, # Dimensionality of raw edge features (e.g. 74) num_hidden: int, num_layers: int, activation: str, mask_rate: float = 0.3, encoder_type: str = "gat", decoder_type: str = "gat", loss_fn: str = "sce", drop_edge_rate: float = 0.0, replace_rate: float = 0.1, alpha_l: float = 2, concat_hidden: bool = False, batchnorm=False, encoder_norm="none", encoder_input_norm="none", edge_distance_in_proj=True, vicreg_var_weight: float = 0.0, vicreg_cov_weight: float = 0.0, vicreg_gamma: float = 1.0, ): super(PreModel, self).__init__() self._mask_rate = mask_rate self._encoder_type = encoder_type self._decoder_type = decoder_type self._drop_edge_rate = drop_edge_rate self._output_hidden_size = num_hidden self._concat_hidden = concat_hidden # VICReg-style anti-collapse regularizer (0 = off); see _vicreg_terms. self._vicreg_var_weight = vicreg_var_weight self._vicreg_cov_weight = vicreg_cov_weight self._vicreg_gamma = vicreg_gamma self.last_loss_components = {} self._replace_rate = replace_rate self._mask_token_rate = 1 - self._replace_rate enc_num_hidden = num_hidden dec_in_dim = num_hidden dec_num_hidden = num_hidden # Build encoder self.encoder = setup_module( m_type=encoder_type, in_dim=in_dim, out_dim=enc_num_hidden, num_hidden=enc_num_hidden, num_layers=num_layers, activation=activation, batchnorm=batchnorm, edge_in_dim=edge_in_dim, norm_type=encoder_norm, input_norm=encoder_input_norm, edge_distance_in_proj=edge_distance_in_proj, ) # Build decoder for attribute prediction self.decoder = setup_module( m_type=decoder_type, in_dim=dec_in_dim, out_dim=in_dim, num_hidden=dec_num_hidden, num_layers=1, activation=activation, batchnorm=batchnorm, edge_in_dim=edge_in_dim, edge_distance_in_proj=edge_distance_in_proj, ) self.enc_mask_token = nn.Parameter(torch.zeros(1, in_dim)) if concat_hidden: self.encoder_to_decoder = nn.Linear( dec_in_dim * num_layers, dec_in_dim, bias=False ) else: self.encoder_to_decoder = nn.Linear(dec_in_dim, dec_in_dim, bias=False) # Setup loss function self.criterion = self.setup_loss_fn(loss_fn, alpha_l) @property def output_hidden_dim(self): return self._output_hidden_size def setup_loss_fn(self, loss_fn, alpha_l): if loss_fn == "mse": criterion = nn.MSELoss() elif loss_fn == "sce": criterion = partial(sce_loss, alpha=alpha_l) else: raise NotImplementedError return criterion def encoding_mask_noise(self, x, mask_rate=0.3, virtual_node_index=None): num_nodes = x.shape[0] all_indices = torch.arange(num_nodes, device=x.device) # Remove virtual node index from masking candidates if virtual_node_index is not None: all_indices = all_indices[~torch.isin(all_indices, virtual_node_index)] perm = all_indices[torch.randperm(len(all_indices), device=x.device)] # random masking num_mask_nodes = int(mask_rate * len(perm)) mask_nodes = perm[:num_mask_nodes] keep_nodes = perm[num_mask_nodes:] out_x = x.clone() if self._replace_rate > 0: num_noise_nodes = int(self._replace_rate * num_mask_nodes) perm_mask = torch.randperm(num_mask_nodes, device=x.device) token_nodes = mask_nodes[ perm_mask[: int(self._mask_token_rate * num_mask_nodes)] ] noise_nodes = mask_nodes[ perm_mask[-int(self._replace_rate * num_mask_nodes) :] ] noise_to_be_chosen = torch.randperm(len(perm), device=x.device)[ :num_noise_nodes ] noise_to_be_chosen = all_indices[noise_to_be_chosen] out_x[token_nodes] = 0.0 out_x[noise_nodes] = x[noise_to_be_chosen] else: token_nodes = mask_nodes out_x[mask_nodes] = 0.0 out_x[token_nodes] += self.enc_mask_token return out_x, (mask_nodes, keep_nodes) def forward(self, batch): # ---- attribute reconstruction ---- x, edge_index, edge_attr, virtual_node_index, batch = ( batch.x, batch.edge_index, batch.edge_attr, getattr(batch, "virtual_node_index", None), batch.batch, ) loss = self.mask_attr_prediction( x, edge_index, edge_attr, batch, virtual_node_index ) return loss def mask_attr_prediction(self, x, edge_index, edge_attr, batch, virtual_node_index): use_x, (mask_nodes, keep_nodes) = self.encoding_mask_noise( x, self._mask_rate, virtual_node_index, ) if self._drop_edge_rate > 0: use_edge_index, masked_edges = dropout_edge( edge_index, self._drop_edge_rate ) use_edge_attr = edge_attr[masked_edges] use_edge_index, use_edge_attr = add_self_loops( use_edge_index, use_edge_attr, fill_value="min" ) else: use_edge_index = edge_index use_edge_attr = edge_attr enc_rep, all_hidden = self.encoder( use_x, use_edge_index, use_edge_attr, batch=batch, return_hidden=True ) if self._concat_hidden: enc_rep = torch.cat(all_hidden, dim=1) # ---- attribute reconstruction ---- rep = self.encoder_to_decoder(enc_rep) # VICReg anti-collapse regularization on the graph-level pooled embedding, # computed BEFORE the decoder remask below, and only while training so that # val_loss stays pure reconstruction and the collapse metric is independent. if self.training and ( self._vicreg_var_weight > 0 or self._vicreg_cov_weight > 0 ): reg_loss, reg_comps = self._vicreg_terms(rep, batch) else: reg_loss, reg_comps = rep.new_zeros(()), {} if self._decoder_type not in ("mlp", "linear"): # * remask, re-mask rep[mask_nodes] = 0 if self._decoder_type in ("mlp", "linear"): recon = self.decoder(rep) else: recon = self.decoder(rep, use_edge_index, use_edge_attr) x_init = x[mask_nodes] x_rec = recon[mask_nodes] loss = self.criterion(x_rec, x_init) self.last_loss_components = {"sce": float(loss.detach()), **reg_comps} return loss + reg_loss def _vicreg_terms(self, rep, batch): """VICReg-style anti-collapse terms on the graph-level pooled embedding. Returns ``(reg_loss, components)``. Combines a variance hinge (keep each embedding dim's per-batch std >= gamma) and a covariance penalty (decorrelate dims). Applied on ``global_mean_pool(rep)`` -- the same representation the PCA-collapse metric is computed on -- so it directly opposes the dimensional collapse observed once reconstruction saturates. """ z = global_mean_pool(rep, batch) # [G, D] G = z.size(0) D = z.size(1) reg = rep.new_zeros(()) comps = {} if self._vicreg_var_weight > 0: std = torch.sqrt(z.var(dim=0, unbiased=False) + 1e-4) var_term = torch.clamp(self._vicreg_gamma - std, min=0).mean() reg = reg + self._vicreg_var_weight * var_term comps["vic_var"] = float(var_term.detach()) if self._vicreg_cov_weight > 0 and G > 1: zc = z - z.mean(dim=0, keepdim=True) cov = (zc.t() @ zc) / (G - 1) # [D, D] off_diag_sq = cov.pow(2).sum() - cov.diagonal().pow(2).sum() cov_term = off_diag_sq / D reg = reg + self._vicreg_cov_weight * cov_term comps["vic_cov"] = float(cov_term.detach()) return reg, comps def embed(self, x, edge_index, edge_attr, batch): if self._concat_hidden: enc_rep, all_hidden = self.encoder( x, edge_index, edge_attr, batch=batch, return_hidden=True ) enc_rep = torch.cat(all_hidden, dim=1) else: enc_rep = self.encoder(x, edge_index, edge_attr, batch=batch) rep = self.encoder_to_decoder(enc_rep) return rep @property def enc_params(self): return self.encoder.parameters() @property def dec_params(self): return chain(*[self.encoder_to_decoder.parameters(), self.decoder.parameters()])