modeling/models/edcoder.py

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()])