Losses¶

mincut_loss(adj: Tensor, S: Tensor, adj_pooled: Tensor, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary mincut loss used by MinCutPooling operator from the paper “Spectral Clustering in Graph Neural Networks for Graph Pooling” (Bianchi et al., ICML 2020).

The loss is computed as

\[\mathcal{L}_\text{CUT} = - \frac{\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{A} \mathbf{S})} {\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{D} \mathbf{S})},\]

where

\(\mathbf{A}\) is the adjacency matrix,
\(\mathbf{S}\) is the dense cluster assignment matrix,
\(\mathbf{D} = \mathrm{diag}(\mathbf{A}^{\top}\mathbf{1})\) is the degree matrix.

Parameters:

adj (Tensor) – The adjacency matrix of shape \((B, N, N)\), where \(B\) is the batch size, \(N\) is the number of nodes, used to compute \(\mathbf{D}\).
S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(K\) is the number of clusters.
adj_pooled (Tensor) – The pooled adjacency matrix \(\mathbf{S}^{\top} \mathbf{A}\mathbf{S}\) of shape \((B, K, K)\).
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The mincut loss.

Return type:

Tensor

orthogonality_loss(S: Tensor, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary orthogonality loss used by MinCutPooling operator from the paper “Spectral Clustering in Graph Neural Networks for Graph Pooling” (Bianchi et al., ICML 2020).

The loss is computed as

\[\mathcal{L}_O = {\left\| \frac{\mathbf{S}^{\top} \mathbf{S}} {{\|\mathbf{S}^{\top} \mathbf{S}\|}_F} -\frac{\mathbf{I}_K}{\sqrt{K}} \right\|}_F,\]

where

\(\mathbf{S}\) is the dense cluster assignment matrix,
\(\mathbf{I}_K\) is the identity matrix of size \(K\),
\(K\) is the number of clusters.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The orthogonality loss.

Return type:

Tensor

sparse_mincut_loss(edge_index: Tensor, S: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean', num_per_graph: Tensor | None = None) → Tensor[source]¶

Sparse auxiliary mincut loss for unbatched graph pooling.

This is the sparse/unbatched version of mincut_loss() used by MinCutPooling in unbatched mode. It operates on sparse adjacency matrices and unbatched dense assignment matrices.

The loss is computed as

\[\mathcal{L}_\text{CUT} = - \frac{\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{A} \mathbf{S})} {\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{D} \mathbf{S})},\]

where

\(\mathbf{A}\) is the adjacency matrix (sparse),
\(\mathbf{S}\) is the dense cluster assignment matrix,
\(\mathbf{D} = \mathrm{diag}(\mathbf{A}^{\top}\mathbf{1})\) is the degree matrix.

Parameters:

edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\), where \(E\) is the number of edges.
S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\), where \(N\) is the total number of nodes and \(K\) is the number of clusters.
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\). If None, all edges have weight 1.0. (default: None)
batch (Tensor, optional) – Batch vector of shape \((N,)\) indicating which graph each node belongs to. If None, assumes single graph. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")
num_per_graph (Tensor, optional) – Pre-computed numerator \(\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{A}_{\mathrm{pool}} \mathbf{S})\) per graph of shape \((B,)\). If provided, it overrides the edge-wise numerator computation, which lets callers feed a postprocessed pooled adjacency while the denominator is still derived from the raw edge_index / edge_weight. (default: None)

Returns:

The mincut loss.

Return type:

Tensor

sparse_ho_mincut_loss(edge_index: Tensor, S: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Sparse higher-order (motif) mincut loss for unbatched graph pooling.

Computes the same scalar as applying mincut_loss() to the dense motif adjacency \(\mathbf{M} = \mathbf{A}^3\), without materializing \(\mathbf{M}\) or dense \((N, N)\) tensors.

For each graph \(g\), it computes:

\[\mathcal{L}_\text{CUT}^{(g)} = -\frac{\mathrm{Tr}(\mathbf{S}_g^\top \mathbf{M}_g \mathbf{S}_g)} {\mathrm{Tr}(\mathbf{S}_g^\top \mathbf{D}_g \mathbf{S}_g)},\]

where \(\mathbf{M}_g = \mathbf{A}_g^3\) and \(\mathbf{D}_g = \mathrm{diag}(\mathbf{M}_g \mathbf{1})\).

Implementation details: - Numerator uses \(\mathrm{Tr}(S^\top M S) = \sum_{i,k} S_{ik} (MS)_{ik}\)

with \(MS = A(A(AS))\).

Denominator uses \(d = M \mathbf{1} = A(A(A\mathbf{1}))\).

This keeps memory closer to \(O(E + NK)\) (plus sparse storage), though runtime can still grow with graph density/3-hop walk proliferation.

unbatched_orthogonality_loss(S: Tensor, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Unbatched auxiliary orthogonality loss for unbatched graph pooling.

This is the unbatched version of orthogonality_loss() used by MinCutPooling in unbatched mode. It operates on unbatched dense assignment matrices.

The loss is computed as

\[\mathcal{L}_O = {\left\| \frac{\mathbf{S}^{\top} \mathbf{S}} {{\|\mathbf{S}^{\top} \mathbf{S}\|}_F} -\frac{\mathbf{I}_K}{\sqrt{K}} \right\|}_F,\]

where

\(\mathbf{S}\) is the dense cluster assignment matrix,
\(\mathbf{I}_K\) is the identity matrix of size \(K\),
\(K\) is the number of clusters.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\), where \(N\) is the total number of nodes and \(K\) is the number of clusters.
batch (Tensor, optional) – Batch vector of shape \((N,)\) indicating which graph each node belongs to. If None, assumes single graph. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The orthogonality loss.

Return type:

Tensor

unbatched_hosc_orthogonality_loss(S: Tensor, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Unbatched HOSC orthogonality loss for unbatched graph pooling.

This is the unbatched version of hosc_orthogonality_loss() used by HOSCPooling in unbatched mode.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The HOSC orthogonality loss.

Return type:

Tensor

unbatched_cluster_loss(S: Tensor, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Unbatched cluster regularization loss for unbatched graph pooling.

This is the unbatched version of cluster_loss() used by DMoNPooling in unbatched mode.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The cluster regularization loss.

Return type:

Tensor

unbatched_entropy_loss(S: Tensor, num_nodes: int | None = None) → Tensor[source]¶

Unbatched entropy regularization loss for unbatched graph pooling.

This is the unbatched version of entropy_loss() used by DiffPool in unbatched mode. Matches the batched semantics: mean entropy per node over the batch (total entropy sum / total number of nodes), then optional reduction over graphs.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
num_nodes (int, optional) – The number of nodes in the graph. If not provided, it is inferred from the shape of \(\mathbf{S}\). (default: None)

Returns:

The entropy regularization loss.

Return type:

Tensor

unbatched_asym_norm_loss(S: Tensor, k: int, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Unbatched asymmetric norm loss for unbatched graph pooling.

This is the unbatched version of asym_norm_loss() used by AsymCheegerCutPooling in unbatched mode.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
k (int) – The number of clusters.
batch (Tensor, optional) – Batch vector of shape \((N,)\).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The asymmetrical norm regularization loss.

Return type:

Tensor

unbatched_just_balance_loss(S: Tensor, batch: Tensor | None = None, normalize_loss: bool = True, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Unbatched balance regularization loss for unbatched graph pooling.

This is the unbatched version of just_balance_loss() used by JustBalancePooling in unbatched mode.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
normalize_loss (bool, optional) – Whether to normalize by sqrt(N*K).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The balance regularization loss.

Return type:

Tensor

hosc_orthogonality_loss(S: Tensor, mask: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary orthogonality loss used by HOSCPooling operator from the paper “Higher-order Clustering and Pooling for Graph Neural Networks” (Duval & Malliaros, CIKM 2022).

The loss is computed as

\[\mathcal{L}_\text{HO} = \frac{1}{\sqrt{K}-1} \bigg( \sqrt{K} - \frac{1}{\sqrt{N}}\sum_{j=1}^K ||\mathbf{S}_{*j}||_F\bigg),\]

where

\(N\) is the number of nodes,
\(K\) is the number of clusters,
\(\mathbf{S}_{*j}\) is the \(j\)-th column of the cluster assignment matrix \(\mathbf{S}\).

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The orthogonality loss.

Return type:

Tensor

link_pred_loss(S: Tensor, adj: Tensor, normalize_loss: bool = True) → Tensor[source]¶

Auxiliary link prediction loss used by DiffPool operator from the paper “Hierarchical Graph Representation Learning with Differentiable Pooling” (Ying et al., NeurIPS 2018).

The loss is computed as

\[\mathcal{L}_{LP} = {\| \mathbf{A} - \mathrm{softmax}(\mathbf{S}) {\mathrm{softmax}(\mathbf{S})}^{\top} \|}_F,\]

where

\(\mathbf{A}\) is the adjacency matrix,
\(\mathbf{S}\) is the dense cluster assignment matrix.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
adj (Tensor) – The adjacency matrix of shape \((B, N, N)\).
normalize_loss (bool, optional) – If set to True, the loss will be normalized by the number of elements in the adjacency matrix. (default: True)

Returns:

The link prediction loss.

Return type:

Tensor

entropy_loss(S: Tensor, num_nodes: int) → Tensor[source]¶

Auxiliary entropy regularization loss used by DiffPool operator from the paper “Hierarchical Graph Representation Learning with Differentiable Pooling” (Ying et al., NeurIPS 2018).

The loss is computed as

\[\mathcal{L}_E = \frac{1}{N} \sum_{n=1}^N H(\mathbf{S}_n),\]

where

\(\mathbf{S}\) is the dense cluster assignment matrix,
\(N\) is the number of nodes,
\(H(\cdot)\) is the entropy function.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\) where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
num_nodes (int) – The number of nodes in the graph.

Returns:

The entropy regularization loss.

Return type:

Tensor

sparse_link_pred_loss(S: Tensor, edge_index: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, normalize_loss: bool = True) → Tensor[source]¶

Sparse link prediction loss giving the same scalar as batched link_pred_loss() (global Frobenius norm over the batch), without materializing \((B, N, N)\).

Uses \(\|\mathbf{A} - \mathbf{S}\mathbf{S}^\top\|_F^2 = \sum_{e} (w_e - (\mathbf{S}\mathbf{S}^\top)_{e})^2 + \sum_g \|\mathbf{S}_g \mathbf{S}_g^\top\|_F^2 - \sum_{e} (\mathbf{S}\mathbf{S}^\top)_{e}^2\).

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
normalize_loss (bool, optional) – If True, divide by total number of entries (sum over graphs of \(n_g^2\)). (default: True)

Returns:

The link prediction loss (scalar, matches batched).

Return type:

Tensor

totvar_loss(S: Tensor, adj: Tensor, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

The total variation regularization loss used by AsymCheegerCutPooling operator from the paper “Total Variation Graph Neural Networks” (Hansen & Bianchi, ICML 2023).

The loss is computed as

\[\mathcal{L}_\text{GTV} = \frac{\mathcal{L}_\text{GTV}^*}{2E} \in [0, 1],\]

with the total variation regularization loss defined as

\[\mathcal{L}_\text{GTV}^* = \displaystyle\sum_{k=1}^K\sum_{i=1}^N \sum_{j=i}^N a_{i,j} |s_{i,k} - s_{j,k}|.\]

where

\(N\) is the number of vertices,
\(K\) is the number of clusters,
\(a_{i,j}\) is the entry \((i,j)\) of the adjacency matrix,
\(s_{i,k}\) is the assignment of vertex \(i\) to cluster \(k\),
\(E\) is the number of edges.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\) where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
adj (Tensor) – The adjacency matrix of shape \((B, N, N)\).
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The total variation regularization loss.

Return type:

Tensor

sparse_totvar_loss(edge_index: Tensor, S: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Sparse total variation loss for unbatched graph pooling.

This is the sparse version of totvar_loss() used by AsymCheegerCutPooling in unbatched mode.

Parameters:

edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\).
S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The total variation regularization loss.

Return type:

Tensor

asym_norm_loss(S: Tensor, k: int, mask: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary asymmetrical norm term used by AsymCheegerCutPooling operator from the paper “Total Variation Graph Neural Networks” (Hansen & Bianchi, ICML 2023).

This term, \(\mathcal{L}_{\text{AN}}\), encourages balanced partitions of the graph by penalizing large deviations between each assignment vector and its \(\rho\)-quantile. It is defined as

\[\mathcal{L}_{\text{AN}} = \frac{\beta - \mathcal{L}^*_{\text{AN}}}{\beta} \in [0, 1],\]

where

\[\mathcal{L}^*_{\text{AN}} = \sum_{k=1}^{K} \bigl\|\mathbf{S}_{:,k} \;-\; \mathrm{quant}_\rho\bigl(\mathbf{S}_{:,k}\bigr)\bigr\|_{1,\rho}.\]

In this formulation:

\(\mathbf{S}\) is the cluster dense assignment matrix and \(\mathbf{S}_{:,k}\) denotes the \(k\)-th column of \(\mathbf{S}\), i.e., the assignments for cluster \(k\) across all nodes.
\(\mathrm{quant}_\rho(\mathbf{S}_{:,k})\) extracts the \(\rho\)-quantile of \(\mathbf{S}_{:,k}\), where \(\rho\) is a balancing parameter typically set to \(K-1\).
\(\|\cdot\|_{1,\rho}\) is the asymmetric \(\ell_1\) norm: \(\|\mathbf{x}\|_{1,\rho} = \sum_{i=1}^N |x_i|_{\rho},\, |x_i|_{\rho} = \rho x_i \,\text{if } x_i \ge 0,\text{ and } -x_i \text{ if } x_i < 0.\)
\(\beta\) is a normalization term ensuring that \(\mathcal{L}_{\text{AN}}\) stays in \([0,1]\). When \(\rho = K-1\), \(\beta = N\rho\). For other values of \(\rho\), \(\beta = N\rho \min\!\bigl(1, \frac{K}{\rho+1}\bigr)\).

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\) where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
k (int) – The number of clusters (\(K\)). This is used internally to set \(\rho = K - 1\) if no other value of \(\rho\) is explicitly chosen.
mask (Optional[Tensor]) – Input-node validity mask of shape \((B, N)\) with True for real (non-padded) nodes. If None, all nodes are used. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The asymmetrical norm regularization loss.

Return type:

Tensor

just_balance_loss(S: Tensor, mask: Tensor | None = None, normalize_loss: bool = True, num_nodes: int | None = None, num_supernodes: int | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary balance regularization loss used by JustBalancePooling operator from the paper “Simplifying Clustering with Graph Neural Networks” (Bianchi, NLDL 2023).

The loss is computed as

\[\mathcal{L}_{B} = - \mathrm{Tr}(\sqrt{\mathbf{S}^{\top} \mathbf{S}}),\]

where

\(\mathbf{S}\) is the dense cluster assignment matrix,
\(\mathrm{Tr}(\cdot)\) is the trace operator.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
normalize_loss (bool, optional) – If set to True, the loss is normalized by the number of nodes \(N\) and the number of clusters \(K\). (default: True)
num_nodes (Optional[int]) – The number of nodes in the graph. If not provided, it is inferred from the shape of \(\mathbf{S}\). (default: None)
num_supernodes (Optional[int]) – The number of clusters in the graph. If not provided, it is inferred from the shape of \(\mathbf{S}\). (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The balance regularization loss.

Return type:

Tensor

spectral_loss(adj: Tensor, S: Tensor, adj_pooled: Tensor, mask: Tensor | None = None, num_supernodes: int | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary spectral regularization loss used by DMoNPooling operator from the paper “Graph Clustering with Graph Neural Networks” (Tsitsulin et al., JMLR 2023).

The loss is computed as

\[\mathcal{L}_S = - \frac{1}{2m} \cdot{\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{B} \mathbf{S})},\]

where

\(\mathbf{B} = \mathbf{A} - \frac{\mathbf{d} \mathbf{d}^{\top}}{2m}\) is the modularity matrix,
\(\mathbf{A}\) is the adjacency matrix,
\(\mathbf{d}\) is the degree vector,
\(m = \frac{1}{2} \sum_{i,j} A_{i,j}\) is the total number of edges in the graph.

Parameters:

adj (Tensor) – The adjacency matrix.
S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
adj_pooled (Tensor) – The pooled adjacency matrix.
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
num_supernodes (Optional[int]) – The number of clusters in the graph. If not provided, it is inferred from the shape of \(\mathbf{S}\). (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The spectral regularization loss.

Return type:

Tensor

sparse_spectral_loss(edge_index: Tensor, S: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Sparse spectral regularization loss for unbatched graph pooling.

This is the sparse version of spectral_loss() used by DMoNPooling in unbatched mode.

Parameters:

edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\).
S (Tensor) – The dense cluster assignment matrix of shape \((N, K)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\).
batch (Tensor, optional) – Batch vector of shape \((N,)\).
batch_reduction (str, optional) – Reduction over the batch dimension.

Returns:

The spectral regularization loss.

Return type:

Tensor

cluster_loss(S: Tensor, mask: Tensor | None = None, num_supernodes: int | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary cluster regularization loss used by DMoNPooling operator from the paper “Graph Clustering with Graph Neural Networks” (Tsitsulin et al., JMLR 2023).

The loss is computed as

\[\mathcal{L}_C = \frac{\sqrt{K}}{N} {\left\|\sum_{i=1}^{N} \mathbf{S}_i^{\top} \right\|}_F - 1,\]

where

\(\mathbf{S}\) is the dense cluster assignment matrix,
\(N\) is the number of nodes,
\(K\) is the number of clusters.

Parameters:

S (Tensor) – The dense cluster assignment matrix of shape \((B, N, K)\), where \(B\) is the batch size, \(N\) is the number of nodes, and \(K\) is the number of clusters.
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
num_supernodes (Optional[int]) – The number of clusters in the graph. If not provided, it is inferred from the shape of \(\mathbf{S}\). (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The cluster regularization loss.

Return type:

Tensor

weighted_bce_reconstruction_loss(rec_adj: Tensor, adj: Tensor, mask: Tensor | None = None, balance_links: bool = True, normalizing_const: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Weighted binary cross-entropy reconstruction loss for adjacency matrices.

This function computes the binary cross-entropy loss between a reconstructed adjacency matrix and the true adjacency matrix. When balance_links is True, it applies class-balancing weights to handle the imbalance between edges and non-edges in sparse graphs.

The weighted BCE loss is computed as:

\[\mathcal{L}_{\text{BCE}} = \text{BCE}(\mathbf{A}_{\text{rec}}, \mathbf{A}, \mathbf{W})\]

where the weight matrix \(\mathbf{W}\) is computed to balance positive and negative samples:

\[W_{ij} = \frac{N^2}{n_{\text{edges}}} \cdot A_{ij} + \frac{N^2}{n_{\text{non-edges}}} \cdot (1 - A_{ij})\]

with \(n_{\text{edges}} = \sum_{i,j} A_{ij}\) and \(n_{\text{non-edges}} = N^2 - n_{\text{edges}}\).

When normalizing_const \(\gamma\) is not None, the loss is normalized by \(\gamma\):

\[\mathcal{L}_{\text{normalized}} = \frac{\mathcal{L}_{\text{BCE}}}{\gamma}\]

Note that \(\gamma\) can be a vector to specify a different constant for each graph in the batch.

Parameters:

rec_adj (Tensor) – The reconstructed adjacency matrix (logits) of shape \((B, N, N)\), where \(B\) is the batch size and \(N\) is the number of nodes. Contains the predicted edge probabilities.
adj (Tensor) – The true adjacency matrix of shape \((B, N, N)\).
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
balance_links (bool, optional) – Whether to apply class-balancing weights to handle edge/non-edge imbalance. (default: True)
normalizing_const (Optional[Tensor]) – The normalizing constant used to scale the loss. It allows batch computation to ensure consistent scaling across graphs of different sizes. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The weighted BCE reconstruction loss.

Return type:

Tensor

kl_loss(q: Distribution, p: Distribution, mask: Tensor | None = None, batch: Tensor | None = None, batch_size: int | None = None, normalizing_const: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Compute KL divergence between two distributions with flexible axis control.

This function computes the KL divergence \(D_{KL}(q \parallel p)\) between two distributions. It is possible to specify either a mask or a batch vector to allow correct computations on batched graphs.

\[D_{KL}(q \parallel p) = \mathbb{E}_{x \sim q}[\log q(x) - \log p(x)]\]

When normalizing_const \(\gamma\) is not None, the loss is normalized by \(\gamma\):

\[D_{KL,\text{normalized}} = \frac{D_{KL}(q \parallel p)}{\gamma}\]

Parameters:

q (Distribution) – The approximate posterior distribution.
p (Distribution) – The prior distribution.
mask (Optional[Tensor]) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)
batch_size (int, optional) – The batch size.
normalizing_const (Optional[Tensor]) – The normalizing constant used to scale the loss. It allows batch computation to ensure consistent scaling across graphs of different sizes. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The KL divergence loss.

Return type:

Tensor

Examples

>>> import torch
>>> from torch.distributions import Beta
>>> from tgp.utils.losses import kl_loss
>>> # Example: Stick-breaking process in BNPool
>>> # Shape: (B=2, N=4, K-1=3) for 4 nodes, max 4 clusters
>>> alpha_sb = torch.ones(2, 4, 3) + 0.5  # Posterior Alpha parameters
>>> beta_sb = torch.ones(2, 4, 3) + 1.0  # Posterior Beta parameters
>>> q_sb = Beta(alpha_sb, beta_sb)  # Posterior distributions
>>> # Prior: Beta(1, alpha_DP) for each stick-breaking fraction
>>> alpha_prior = torch.ones(3)
>>> beta_prior = torch.ones(3) * 2.0  # alpha_DP = 2.0
>>> p_sb = Beta(alpha_prior, beta_prior)
>>> # Node mask for variable-sized graphs
>>> mask = torch.tensor(
...     [[True, True, True, False], [True, True, True, True]], dtype=torch.bool
... )
>>> # Compute KL loss: sum over K-1 components, then over nodes
>>> loss = kl_loss(q_sb, p_sb, mask=mask)

cluster_connectivity_prior_loss(K: Tensor, K_mu: Tensor, K_var: Tensor, normalizing_const: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Prior loss for cluster connectivity matrix in BNPool.

This function computes the prior loss for the cluster connectivity matrix \(\mathbf{K}\), which regularizes the learned cluster-cluster connectivity probabilities towards a prior distribution. The prior loss is computed as the negative log-likelihood of a Gaussian prior:

\[\mathcal{L}_{\mathbf{K}} = \frac{1}{2} \sum_{i,j} \frac{(K_{ij} - \mu_{ij})^2}{\sigma^2}\]

where \(\mathbf{K} \in \mathbb{R}^{C \times C}\) is the cluster connectivity matrix, \(\boldsymbol{\mu} \in \mathbb{R}^{C \times C}\) is the prior mean matrix, and \(\sigma^2\) is the prior variance.

The prior mean \(\boldsymbol{\mu}\) typically has the structure:

\[\begin{split}\mu_{ij} = \begin{cases} \mu_{\text{diag}} & \text{if } i = j \text{ (within-cluster connectivity)} \\ \mu_{\text{off}} & \text{if } i \neq j \text{ (between-cluster connectivity)} \end{cases}\end{split}\]

This structure encourages block-diagonal patterns in the reconstructed adjacency matrix \(\mathbf{A}_{\text{rec}} = \mathbf{S} \mathbf{K} \mathbf{S}^{\top}\), promoting well-separated clusters.

When normalizing_const \(\gamma\) is not None, the loss is normalized by \(\gamma\):

\[\mathcal{L}_{\text{normalized}} = \frac{\mathcal{L}_{\mathbf{K}}}{\gamma}\]

Parameters:

K (Tensor) – The learnable cluster connectivity matrix of shape \((C, C)\), where \(C\) is the maximum number of clusters. This matrix models the expected connectivity patterns between different clusters.
K_mu (Tensor) – Prior mean matrix of shape \((C, C)\) specifying the expected values for the connectivity matrix. Usually designed to encourage higher within-cluster than between-cluster connectivity.
K_var (Tensor) – Prior variance parameter \(\sigma^2\) (scalar tensor). Controls the strength of the regularization - smaller values impose stronger constraints towards the prior mean.
normalizing_const (Optional[Tensor]) – The normalizing constant used to scale the loss. It allows batch computation to ensure consistent scaling across graphs of different sizes. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The cluster connectivity prior loss.

Return type:

Tensor

Note

Typically used with \(\mu_{\text{diag}} > 0\) and \(\mu_{\text{off}} < 0\)
The loss strength can be controlled through K_var

sparse_bce_reconstruction_loss(link_prob_loigit, true_y, edges_batch_id: Tensor | None = None, batch_size=None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tuple[Tensor, Tensor][source]¶

Sparse weighted binary cross-entropy reconstruction loss for sampled edges.

Parameters:

link_prob_loigit (Tensor) – Logits for sampled edges of shape \([E]\).
true_y (Tensor) – Ground-truth labels for sampled edges of shape \([E]\).
edges_batch_id (Tensor, optional) – Batch assignment for each sampled edge. (default: None)
batch_size (int, optional) – Number of graphs in the batch.
batch_reduction (str, optional) – Reduction applied across graphs. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The loss value and the number of sampled edges (per-graph counts if edges_batch_id is provided).

Return type:

Tuple[Tensor, Tensor | int]

maxcut_loss(scores: Tensor, edge_index: Tensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, batch_reduction: Literal['mean', 'sum'] = 'mean') → Tensor[source]¶

Auxiliary MaxCut loss used by MaxCutPooling operator from the paper “MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks” (Abate & Bianchi, ICLR 2025).

The MaxCut objective aims to maximize the sum of edge weights crossing a graph partition. For differentiable optimization, the loss minimizes the negative normalized MaxCut value:

\[\mathcal{L}_{\text{MaxCut}} = -\frac{1}{V} \sum_{(i,j) \in E} w_{ij} \cdot z_i \cdot z_j\]

where:

\(z_i \in [-1, 1]\) are the node scores/assignments,
\(w_{ij}\) are the edge weights,
\(V = \sum_{(i,j) \in E} w_{ij}\) is the graph volume (total edge weight),
\(E\) is the edge set.

The computation is performed efficiently using sparse matrix operations:

\[\mathcal{L}_{\text{MaxCut}} = -\frac{\mathbf{z}^{\top} \mathbf{A} \mathbf{z}}{V}\]

where \(\mathbf{A}\) is the weighted adjacency matrix and \(\mathbf{z}\) contains node scores.

Implementation Details:

Node scores are normalized via \(\tanh\) to \([-1, 1]\) range
Sparse matrix multiplication \(\mathbf{A} \mathbf{z}\) is computed efficiently
Volume normalization ensures loss comparability across different graph sizes
Batch processing handles multiple graphs simultaneously

Parameters:

scores (Tensor) – Node scores/assignments of shape \((N,)\) or \((N, 1)\). Typically normalized to \([-1, 1]\) via tanh activation.
edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\). If None, all edges have weight 1.0. (default: None)
batch (Tensor, optional) – Batch assignments for each node of shape \((N,)\). If None, assumes single graph. (default: None)
batch_reduction (str, optional) – Reduction method applied to the batch dimension. Can be 'mean' or 'sum'. (default: "mean")

Returns:

The MaxCut loss value (scalar for single graph, or reduced across batch).

Return type:

Tensor

Note

The volume normalization \(V = \sum_{(i,j) \in E} w_{ij}\) ensures that the loss magnitude is comparable across graphs of different sizes and densities, making it suitable for batched training scenarios.