from typing import List, Optional, Union
import torch
from torch import Tensor
from torch.distributions import Beta
from torch_geometric.typing import Adj
from tgp.connect import DenseConnect
from tgp.lift import BaseLift
from tgp.reduce import BaseReduce
from tgp.select import DPSelect, SelectOutput
from tgp.src import DenseSRCPooling, PoolingOutput
from tgp.utils import (
batched_negative_edge_sampling,
connectivity_to_edge_index,
negative_edge_sampling,
)
from tgp.utils.losses import (
cluster_connectivity_prior_loss,
kl_loss,
sparse_bce_reconstruction_loss,
weighted_bce_reconstruction_loss,
)
from tgp.utils.typing import LiftType, SinvType
[docs]
class BNPool(DenseSRCPooling):
r"""The BN-Pool operator from the paper `"BN-Pool: Bayesian Nonparametric Graph Pooling" <https://arxiv.org/abs/2501.09821>`_
(Castellana & Bianchi, 2025).
BN-Pool implements a Bayesian nonparametric approach to graph pooling using a Dirichlet Process
with stick-breaking construction for cluster assignment. The method learns both the number of clusters
and their assignments through variational inference.
+ The :math:`\texttt{select}` operator is implemented with :class:`~tgp.select.DPSelect` to perform variational inference of the stick-breaking process.
+ The :math:`\texttt{reduce}` operator is implemented with :class:`~tgp.reduce.BaseReduce`.
+ The :math:`\texttt{connect}` operator is implemented with :class:`~tgp.connect.DenseConnect`.
+ The :math:`\texttt{lift}` operator is implemented with :class:`~tgp.lift.BaseLift`.
The method uses a truncated stick-breaking representation of the Dirichlet Process:
.. math::
v_{ik} \sim \text{Beta}(\alpha_{ik}, \beta_{ik}), \quad i = 1, \ldots, N \quad k = 1, \ldots, K-1
.. math::
\pi_{ik} = v_{ik} \prod_{j=1}^{k-1} (1 - v_{ij})
where :math:`\pi_{ik}` represents the probability of assigning node :math:`i` to cluster :math:`k`.
The coefficients :math:`\alpha_{ik}` and :math:`\beta_{ik}` are computed by an MLP
from node features :math:`\mathbf{x}_i`.
The cluster connectivity is modeled through a learnable matrix :math:`\mathbf{K} \in \mathbb{R}^{K \times K}`
and the pooled adjacency matrix is computed as:
.. math::
\mathbf{A}_{\text{rec}} = \mathbf{S} \mathbf{K} \mathbf{S}^{\top}
where :math:`S_{ik} = \pi_{ik}`.
This layer optimizes three auxiliary losses:
+ **Reconstruction loss** (:func:`~tgp.utils.losses.weighted_bce_reconstruction_loss`): Binary cross-entropy loss between the true and reconstructed adjacency matrix :math:`\mathbf{A}_{\text{rec}}`.
+ **KL divergence loss** (:func:`~tgp.utils.losses.kl_loss`): KL divergence between the prior and posterior variational approximation of the stick-breaking variables.
+ **Cluster connectivity prior loss** (:func:`~tgp.utils.losses.cluster_connectivity_prior_loss`): Prior regularization on the cluster connectivity matrix :math:`\mathbf{K}`.
Args:
in_channels (Union[int, List[int]]): The number of input node feature channels.
If a list is provided, it specifies the architecture of the MLP in :class:`~tgp.select.DPSelect`.
k (int): The maximum number of clusters :math:`K` to be used in the pooling mechanism.
The actual number of active clusters is learned through the stick-breaking process.
alpha_DP (float, optional): Prior concentration parameter :math:`\alpha` of the Dirichlet Process.
Controls the expected number of clusters. Higher values encourage more clusters.
(default: ``1.0``)
K_var (float, optional): Variance :math:`\sigma^2` of the Gaussian prior on the cluster connectivity matrix :math:`\mathbf{K}`.
(default: ``1.0``)
K_mu (float, optional): Mean parameter for the cluster connectivity prior. The prior mean matrix is constructed as
:math:`\mathbf{K}_{\mu} = \mu \mathbf{I} - \mu (\mathbf{1}\mathbf{1}^{\top} - \mathbf{I})`.
(default: ``10.0``)
K_init (float, optional): Initial value for the cluster connectivity matrix :math:`\mathbf{K}`.
(default: ``1.0``)
eta (float, optional): Weights the KL divergence loss term.
(default: ``1.0``)
train_K (bool, optional): If :obj:`True`, the cluster connectivity matrix :math:`\mathbf{K}` is learnable.
If :obj:`False`, :math:`\mathbf{K}` is fixed to its initial value.
(default: :obj:`True`)
act (str, optional): Activation function for the MLP in :class:`~tgp.select.DPSelect`.
(default: :obj:`None`)
dropout (float, optional): Dropout rate in the MLP of :class:`~tgp.select.DPSelect`.
(default: ``0.0``)
remove_self_loops (bool, optional):
If :obj:`True`, the self-loops will be removed from the adjacency matrix.
(default: :obj:`True`)
degree_norm (bool, optional):
If :obj:`True`, the adjacency matrix will be symmetrically normalized.
(default: :obj:`True`)
edge_weight_norm (bool, optional):
Whether to normalize the edge weights by dividing by the maximum absolute value per graph.
(default: :obj:`False`)
adj_transpose (bool, optional):
If :obj:`True`, the preprocessing step in :class:`tgp.src.DenseSRCPooling` and
the :class:`tgp.connect.DenseConnect` operation returns transposed
adjacency matrices, so that they could be passed "as is" to the dense
message-passing layers.
(default: :obj:`True`)
lift (~tgp.utils.typing.LiftType, optional):
Defines how to compute the matrix :math:`\mathbf{S}_\text{inv}` to lift the pooled node features.
- ``"precomputed"`` (default): Use as :math:`\mathbf{S}_\text{inv}` what is
already stored in the ``"s_inv"`` attribute of the :class:`tgp.select.SelectOutput`.
- ``"transpose"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Recomputes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
s_inv_op (~tgp.utils.typing.SinvType, optional):
The operation used to compute :math:`\mathbf{S}_\text{inv}` from the select matrix
:math:`\mathbf{S}`. :math:`\mathbf{S}_\text{inv}` is stored in the ``"s_inv"`` attribute of
the :class:`tgp.select.SelectOutput`. It can be one of:
- ``"transpose"`` (default): Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^\top`,
the transpose of :math:`\mathbf{S}`.
- ``"inverse"``: Computes :math:`\mathbf{S}_\text{inv}` as :math:`\mathbf{S}^+`,
the Moore-Penrose pseudoinverse of :math:`\mathbf{S}`.
cache_preprocessing (bool, optional):
If :obj:`True`, caches the dense adjacency produced during preprocessing.
This should only be enabled when the same graph is reused across iterations.
(default: :obj:`False`)
batched (bool, optional):
If :obj:`True`, uses the batched dense representation of the input.
If :obj:`False`, uses an unbatched representation without padding.
(default: :obj:`True`)
sparse_output (bool, optional):
If :obj:`True`, returns block-diagonal sparse outputs. If :obj:`False`,
returns batched dense outputs. (default: :obj:`False`)
num_neg_samples (int, optional):
Cap on the number of negative edges sampled **per graph** in the unbatched
(sparse-loss) path. If :obj:`None`, defaults to matching the number of
positive edges. (default: :obj:`None`)
"""
def __init__(
self,
in_channels: Union[int, List[int]],
k: int, # hyperparameters of the method
alpha_DP=1.0,
K_var=1.0,
K_mu=10.0,
K_init=1.0,
eta=1.0,
train_K=True, # hyperparameters of the selector
act: str = None,
dropout: float = 0.0,
remove_self_loops: bool = True,
degree_norm: bool = True,
edge_weight_norm: bool = False,
adj_transpose: bool = True,
lift: LiftType = "precomputed",
s_inv_op: SinvType = "transpose",
batched: bool = True,
sparse_output: bool = False,
cache_preprocessing: bool = False,
num_neg_samples: Optional[int] = None,
):
if alpha_DP <= 0:
raise ValueError("alpha_DP must be positive")
if K_var <= 0:
raise ValueError("K_var must be positive")
if eta <= 0:
raise ValueError("eta must be positive")
if k <= 0:
raise ValueError("max_k must be positive")
super(BNPool, self).__init__(
selector=DPSelect(
in_channels,
k,
batched_representation=batched,
act=act,
dropout=dropout,
s_inv_op=s_inv_op,
),
reducer=BaseReduce(),
lifter=BaseLift(matrix_op=lift),
connector=DenseConnect(
remove_self_loops=remove_self_loops,
degree_norm=degree_norm,
adj_transpose=adj_transpose,
edge_weight_norm=edge_weight_norm,
sparse_output=sparse_output,
),
adj_transpose=adj_transpose,
cache_preprocessing=cache_preprocessing,
batched=batched,
sparse_output=sparse_output,
)
self.k = k
self.K_init_val = K_init
self.alpha_DP = alpha_DP
self.K_var_val = K_var
self.K_mu_val = K_mu
self.train_K = train_K
self.eta = eta # coefficient for the kl_loss
self.num_neg_samples = num_neg_samples
# prior of the Stick Breaking Process
self.register_buffer("alpha_prior", torch.ones(self.k - 1))
self.register_buffer("beta_prior", torch.ones(self.k - 1) * alpha_DP)
# prior of the cluster-cluster prob. matrix
self.register_buffer("K_var", torch.tensor(K_var))
self.register_buffer(
"K_mu",
K_mu * torch.eye(self.k, self.k) - K_mu * (1 - torch.eye(self.k, self.k)),
)
# cluster-cluster prob matrix
self.K = torch.nn.Parameter(
K_init * torch.eye(self.k, self.k)
- K_init * (1 - torch.eye(self.k, self.k)),
requires_grad=train_K,
)
def reset_parameters(self):
super().reset_parameters()
self.K.data = self.K_init_val * torch.eye(
self.k, self.k, device=self.K.device
) - self.K_init_val * (1 - torch.eye(self.k, self.k, device=self.K.device))
[docs]
def forward(
self,
x: Tensor,
adj: Optional[Adj] = None,
edge_weight: Optional[Tensor] = None,
so: Optional[SelectOutput] = None,
batch: Optional[Tensor] = None,
batch_pooled: Optional[Tensor] = None,
lifting: bool = False,
mask: Optional[Tensor] = None,
**kwargs,
) -> PoolingOutput:
r"""Forward pass.
Args:
x (~torch.Tensor): Node feature tensor
:math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
batch-size :math:`B`, (maximum) number of nodes :math:`N` for
each graph, and feature dimension :math:`F`.
adj (~torch_geometric.typing.Adj, optional): The connectivity matrix.
In batched mode, this accepts sparse connectivity
(``edge_index``, ``torch_sparse.SparseTensor``, or torch COO),
which is internally converted to a dense padded tensor
:math:`\mathbf{A} \in \mathbb{R}^{B \times N \times N}`, or an
already dense adjacency tensor with the same shape.
(default: :obj:`None`)
edge_weight (~torch.Tensor, optional): Edge weights associated with
``adj`` when sparse connectivity is provided.
(default: :obj:`None`)
so (~tgp.select.SelectOutput, optional): The output of the :math:`\texttt{select}` operator.
(default: :obj:`None`)
batch (~torch.Tensor, optional): Batch assignment vector for input
nodes. Required in sparse mode and optional in dense mode.
(default: :obj:`None`)
batch_pooled (~torch.Tensor, optional): Optional precomputed batch
assignment for pooled nodes, used when ``lifting=True``.
(default: :obj:`None`)
lifting (bool, optional): If set to :obj:`True`, the :math:`\texttt{lift}` operation is performed.
(default: :obj:`False`)
mask (~torch.Tensor, optional): Input-node validity mask
:math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}`, where :obj:`True`
marks real (non-padded) nodes. Only used when inputs are already
dense/padded.
(default: :obj:`None`)
Returns:
~tgp.src.PoolingOutput: The output of the pooling operator.
"""
if lifting:
# Lift
x_lifted = self.lift(
x_pool=x, so=so, batch=batch, batch_pooled=batch_pooled
)
return x_lifted
# === Batched path ===
if self.batched:
x, adj, mask = self._ensure_batched_inputs(
x=x,
edge_index=adj,
edge_weight=edge_weight,
batch=batch,
mask=mask,
)
# Select
so = self.select(x=x, mask=mask)
# Reduce
x_pooled, batch_pooled = self.reduce(x=x, so=so, batch=batch)
# Connect
adj_pool, _ = self.connect(edge_index=adj, so=so)
loss = self.compute_loss(adj, mask, so)
if self.sparse_output:
x_pooled, edge_index_pooled, edge_weight_pooled, batch_pooled = (
self._finalize_sparse_output(
x_pool=x_pooled,
adj_pool=adj_pool,
batch=batch,
batch_pooled=batch_pooled,
so=so,
)
)
return PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
loss=loss,
)
return PoolingOutput(x=x_pooled, edge_index=adj_pool, so=so, loss=loss)
# === Unbatched (sparse-loss) path ===
# Select
so = self.select(x=x, batch=batch)
loss = self.compute_sparse_loss(adj, batch, so)
# Reduce
return_batched = not self.sparse_output
x_pooled, batch_pooled = self.reduce(
x=x, so=so, batch=batch, return_batched=return_batched
)
# Connect
edge_index_pooled, edge_weight_pooled = self.connect(
edge_index=adj,
so=so,
edge_weight=edge_weight,
batch=batch,
batch_pooled=batch_pooled,
)
return PoolingOutput(
x=x_pooled,
edge_index=edge_index_pooled,
edge_weight=edge_weight_pooled,
batch=batch_pooled,
so=so,
loss=loss,
)
def compute_loss(self, adj, mask, so) -> dict:
r"""Computes the loss components for BN-Pool training.
This method calculates three loss components that guide the learning of cluster assignments
and connectivity patterns:
1. **Reconstruction Loss**: Measures how well the learned cluster connectivity matrix :math:`\mathbf{K}`
can reconstruct the original adjacency matrix through :math:`\mathbf{A}_{\text{rec}} = \mathbf{S} \mathbf{K} \mathbf{S}^{\top}`.
2. **KL Divergence Loss**: Regularizes the posterior stick-breaking variables :math:`q(v_k)` towards
the Dirichlet Process prior :math:`p(v_k)`.
3. **Cluster Connectivity Prior Loss**: Regularizes the learned connectivity matrix :math:`\mathbf{K}`
towards the specified prior distribution.
All losses are normalized by :math:`N^2` (number of node pairs) to ensure consistent scaling
across different graph sizes.
Args:
adj (~torch.Tensor): True adjacency matrix of shape :math:`(B, N, N)` to reconstruct.
mask (~torch.Tensor): Input-node validity mask of shape :math:`(B, N)`,
where :obj:`True` marks real (non-padded) nodes.
so (:class:`~tgp.select.SelectOutput`): Selection output containing:
- :attr:`s`: Soft assignment matrix :math:`\mathbf{S} \in \mathbb{R}^{B \times N \times K}`
- :attr:`q_z`: Posterior Beta distributions for stick-breaking variables
Returns:
dict: Dictionary containing three loss components:
- :obj:`'quality'`: Reconstruction loss :math:`\mathcal{L}_{\text{rec}}`
(see :func:`~tgp.utils.losses.weighted_bce_reconstruction_loss`)
- :obj:`'kl'`: KL divergence loss :math:`\eta \cdot \mathcal{L}_{\text{KL}}` weighted by :attr:`eta`
(see :func:`~tgp.utils.losses.kl_loss`)
- :obj:`'K_prior'`: Cluster connectivity prior loss :math:`\mathcal{L}_{\mathbf{K}}`
(see :func:`~tgp.utils.losses.cluster_connectivity_prior_loss`).
Set to ``0.0`` if :attr:`train_K=False`.
Note:
The total training loss is typically computed as:
:math:`\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{rec}} + \mathcal{L}_{\text{KL}} + \mathcal{L}_{\mathbf{K}}`
"""
s, q_z = so.s, so.q_z
rec_adj = self.get_rec_adj(s)
if mask is not None:
N = mask.sum(-1) # has shape B x 1
else:
N = torch.tensor(adj.shape[-1], device=adj.device) # N
N_squared = N**2
# Reconstruction loss
rec_loss = weighted_bce_reconstruction_loss(
rec_adj,
adj,
mask,
balance_links=True,
normalizing_const=N_squared,
batch_reduction="mean",
)
# KL loss
alpha_prior = self.get_buffer("alpha_prior")
beta_prior = self.get_buffer("beta_prior")
prior_dist = Beta(alpha_prior, beta_prior)
kl_loss_value = kl_loss(
q_z,
prior_dist,
mask=mask,
normalizing_const=N_squared,
batch_reduction="mean",
)
# K prior loss
if self.train_K:
K_prior_loss = cluster_connectivity_prior_loss(
self.K,
self.get_buffer("K_mu"),
self.get_buffer("K_var"),
normalizing_const=N_squared,
batch_reduction="mean",
)
else:
K_prior_loss = torch.tensor(0.0)
# build the output dictionary
return {
"quality": rec_loss,
"kl": self.eta * kl_loss_value,
"K_prior": K_prior_loss,
}
[docs]
def compute_sparse_loss(
self, adj: Adj, batch: Optional[Tensor], so: SelectOutput
) -> dict:
"""Compute BNPool auxiliary losses in unbatched sparse mode."""
node_assignment, q_z = so.s, so.q_z
if batch is not None:
batch_size = int(batch.max()) + 1
else:
batch_size = 1
rec_loss, norm_const = self.get_sparse_rec_loss(
node_assignment, adj, batch, batch_size
)
alpha_prior = self.get_buffer("alpha_prior")
beta_prior = self.get_buffer("beta_prior")
prior_dist = Beta(alpha_prior, beta_prior)
kl_loss_value = kl_loss(
q_z,
prior_dist,
batch=batch,
batch_size=batch_size,
normalizing_const=norm_const,
batch_reduction="mean",
)
if self.train_K:
K_prior_loss = cluster_connectivity_prior_loss(
self.K,
self.get_buffer("K_mu"),
self.get_buffer("K_var"),
normalizing_const=norm_const,
batch_reduction="mean",
)
else:
K_prior_loss = torch.tensor(0.0)
return {
"quality": rec_loss,
"kl": self.eta * kl_loss_value,
"K_prior": K_prior_loss,
}
def extra_repr_args(self) -> dict:
return {
"batched": self.batched,
"alpha_DP": self.alpha_DP,
"k_prior_variance": self.K_var_val,
"k_prior_mean": self.K_mu_val,
"k_init_value": self.K_init_val,
"eta": self.eta,
"train_K": self.train_K,
"num_neg_samples": self.num_neg_samples,
}
[docs]
def get_rec_adj(self, S):
"""Return the reconstructed dense adjacency logits from assignments."""
return S @ self.K @ S.transpose(-1, -2)
[docs]
def get_sparse_rec_loss(self, node_assignment, adj, batch, batch_size):
"""Compute sparse reconstruction loss using sampled positive/negative edges."""
edge_index, _ = connectivity_to_edge_index(adj)
dev = edge_index.device
if batch is None:
neg_edge_index = negative_edge_sampling(
edge_index,
num_neg_samples=self.num_neg_samples,
force_undirected=True,
)
else:
neg_edge_index = batched_negative_edge_sampling(
edge_index,
batch,
num_neg_samples=self.num_neg_samples,
force_undirected=True,
)
num_edges = edge_index.size(1)
num_neg_edges = neg_edge_index.size(1)
all_edges = torch.cat([edge_index, neg_edge_index], dim=1)
edges_batch_id = None
if batch is not None:
edges_batch_id = batch[all_edges[0]]
link_prob_logit = self.get_prob_link_logit(node_assignment, all_edges)
pred_y = torch.cat(
[torch.ones(num_edges, device=dev), torch.zeros(num_neg_edges, device=dev)],
dim=0,
)
return sparse_bce_reconstruction_loss(
link_prob_logit,
pred_y,
edges_batch_id=edges_batch_id,
batch_size=batch_size,
)
[docs]
def get_prob_link_logit(self, node_assignment, edges_list):
"""Score candidate edges from assignment probabilities and matrix :math:`K`."""
left = node_assignment[edges_list[0]]
right = node_assignment[edges_list[1]]
aux = left @ self.K
return (aux * right).sum(-1)
