Possible error in DMoNPooling implementation and documentation (and example).

[chester-tan]

🐛 Describe the bug

I think there might be an error in the implementation and documentation of DMoNPooling https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/dense/dmon_pool.py .

The documentation writes

    based on dense learned assignments :math:`\mathbf{S} \in \mathbb{R}^{B
    \times N \times C}`.
    Returns the learned cluster assignment matrix, the pooled node feature
    matrix, the coarsened symmetrically normalized adjacency matrix, and three
    auxiliary objectives: (1) The spectral loss

    .. math::
        \mathcal{L}_s = - \frac{1}{2m}
        \cdot{\mathrm{Tr}(\mathbf{S}^{\top} \mathbf{B} \mathbf{S})}

    where :math:`\mathbf{B}` is the modularity matrix, (2) the orthogonality
    loss

    .. math::
        \mathcal{L}_o = {\left\| \frac{\mathbf{S}^{\top} \mathbf{S}}
        {{\|\mathbf{S}^{\top} \mathbf{S}\|}_F} -\frac{\mathbf{I}_C}{\sqrt{C}}
        \right\|}_F

    where :math:`C` is the number of clusters, and (3) the cluster loss

    .. math::
        \mathcal{L}_c = \frac{\sqrt{C}}{n}
        {\left\|\sum_i\mathbf{C_i}^{\top}\right\|}_F - 1.

As I understand, for it to match the original paper https://arxiv.org/pdf/2006.16904 , there should be only 2 losses, and the "spectral" loss should be called a "modularity" loss

    based on dense learned assignments :math:`\mathbf{C} \in \mathbb{R}^{B
    \times N \times C}`.
    Returns the learned cluster assignment matrix, the pooled node feature
    matrix, the coarsened symmetrically normalized adjacency matrix, and two
    auxiliary objectives: (1) The modularity loss

    .. math::
        \mathcal{L}_s = - \frac{1}{2m}
        \cdot{\mathrm{Tr}(\mathbf{C}^{\top} \mathbf{B} \mathbf{C})}

    where :math:`\mathbf{B}` is the modularity matrix, and (2) the cluster loss

    .. math::
        \mathcal{L}_c = \frac{\sqrt{C}}{n}
        {\left\|\sum_i\mathbf{C_i}^{\top}\right\|}_F - 1.

     where :math:`C` is the number of clusters.

I think the orthogonality loss in the current implementation is actually from MinCutPool. The DMoNPooling example in https://github.com/pyg-team/pytorch_geometric/blob/master/examples/proteins_dmon_pool.py should also only sum the modularity loss and cluster loss (ignoring the orthogonality loss).

The implementation might not have to change since I think the calculation of the modularity loss and cluster (or "collapse regularization") loss is correct, though the calculation of the additional orthogonality loss might be both unnecessary computationally and confusing for users.

Versions

Not applicable. Referring to the "latest" documentation and master branch as of 10 June 2024.

[rusty1s]

I think this is correct. Looking at the original implementation here, it uses three loss regularizations.

[chester-tan]

I think this might be the correct implementation (in the same repo) with only 2 losses https://github.com/google-research/google-research/blob/master/graph_embedding/dmon/dmon.py