Multiset Transformer: Advancing Representation Learning in Persistence Diagrams

TL;DR

This is the first neural network that utilizes attention mechanisms specifically designed for multisets as inputs and offers rigorous theoretical guarantees of permutation invariance.

Abstract

To improve persistence diagram representation learning, we propose Multiset Transformer. This is the first neural network that utilizes attention mechanisms specifically designed for multisets as inputs and offers rigorous theoretical guarantees of permutation invariance. The architecture integrates multiset-enhanced attentions with a pool-decomposition scheme, allowing multiplicities to be preserved across equivariant layers. This capability enables full leverage of multiplicities while significantly reducing both computational and spatial complexity compared to the Set Transformer. Additionally, our method can greatly benefit from clustering as a preprocessing step to further minimize complexity, an advantage not possessed by the Set Transformer. Experimental results demonstrate that the Multiset Transformer outperforms existing neural network methods in the realm of persistence diagram representation learning.

Figures (5)

• Figure 1: Persistence diagram examples. PDs are represented as point sets in $\mathbb{R}^2$ above the diagonal. Each point's size indicates its multiplicity, and its color reflects the distance from the diagonal. The left PD contains 184 points, with 183 being distinct. After applying DBSCAN clustering, the diagram is reduced to 54 distinct points, as depicted in the right figure. Both PDs are characterized as multisets.
• Figure 2: MST architecture. Base set $X$ with multiplicities $M_X$ is processed by equivariant layers, preserving permutation order. Representation output $R$ is generated, with multiplicities $M_X$ used as input to an invariant layer.
• Figure 3: Sythetic data classification pipeline. A multiset sample $(X, M_X)$ is processed by a Multiset Transformer (MST) to generate its representation $R$. This representation is then used by a fully connected (FC) classifier to make predictions $Y$. In the context of problem setup (Section \ref{['sec:problem']}), MST corresponds to the representation function $f_\theta$ and FC corresponds to the task-specific function $g_\phi$.
• Figure 4: Graph classification architecture. Given a graph $G$, it's encoded into various ordinary or extended PDs, i.e., $(X_i, M_{X_i})$. Here, we have $i \in \{1,2,3\}$ as demonstration. In experimentation, we can have $i \in \{1,2,\ldots, n\}$ for some finite integer $n$. Each diagram is processed by an independent instance of the Multiset Transformer (MST $i$), yielding its representation $R_i$. These representations are concatenated to form the complete feature set of the graph. The classifier, represented by a fully connected layer (FC), then makes predictions based on this comprehensive representation.
• Figure 5: Multiset Transformer architecture with clustering preprocessing. In this architecture, 'CL.' represents the clustering preprocessing step. Initially, the input multiset $(X, M_X)$ undergoes clustering to approximate its structure. The output of this preprocessing stage is a clustered multiset, denoted by $(X', M_{X'})$, which then serves as the input for the Multiset Transformer.

Theorems & Definitions (6)

• Definition 3.1: Permutation Invariance
• Definition 3.2: Permutation Equivariance
• Theorem 5.1
• Theorem 5.2
• proof
• proof