Wasserstein Hypergraph Neural Network

## Abstract

Abstract The ability to model relational information using machine learning has driven advancements across various domains, from medicine to social science. While graph representation learning has become mainstream over the past decade, representing higher-order relationships through hypergraphs is rapidly gaining momentum.
In the last few years, numerous hypergraph neural networks have emerged, most of them falling under a two-stage, set-based framework. The messages are sent from nodes to edges and then from edges to nodes. However, most of the advancement still takes inspiration from the graph counterpart, often simplifying the aggregations to basic pooling operations.
In this paper we are introducing Wasserstein Hypergraph Neural Network, a model that treats the nodes and hyperedge neighbourhood as distributions and aggregate the information using Sliced Wasserstein Pooling. Unlike conventional aggregators such as mean or sum, which only capture first-order statistics, our approach has the ability to preserve geometric properties like the shape and spread of distributions. This enables the learned embeddings to reflect how easily one hyperedge distribution can be transformed into another, following principles of optimal transport.
Experimental results demonstrate that applying Wasserstein pooling in a hypergraph setting significantly benefits node classification tasks, achieving top performance on several real-world datasets.

## 1 Introduction

The potential to learn from relational data has substantially broadened the applicability of machine learning, extending its reach to a wide range of fields from medicine journals/corr/abs-2310-13767; pred_pat_outcome , to physics gonzalez20a_battaglia_simulate; lam2023graphcastlearningskillfulmediumrange, social science monti2019fake_news and chemistry pmlr-v70-gilmer17a; huang2020skipgnn_molecular_interactions. The flexibility of graph structures makes them well-suited for representing complex natural phenomena involving various types of interactions. As a result, graphs quickly became synonymous with modelling interactions. However, while graphs are restricted to model pairwise connections, many real-world interactions involve more than two entities. To fill this gap, a generalisation of graphs called hypergraphs were introduced, allowing for the representation of higher-order relationships among multiple elements.

More precisely, a hypergraph is characterized by a set of edges, where each edge connects a set of nodes, potentially of varying cardinality. The challenge of designing hypergraph neural networks becomes the challenge of properly modelling these sets. Many approaches allset; wang2022equivariant; UniGNN tackle this using a two-step process: first, the model aggregates information from the nodes within each hyperedge to compute a representation for that hyperedge. Then, in the second step, it updates each node’s representation using information from the hyperedges it belongs to. Both steps rely on methods designed to handle sets of elements.

Although set representation learning has seen significant progress in recent years set_survey, hypergraph networks still largely rely on sum-based aggregation methods such as Deep Sets DeepSets and Set Transformers set_transformer. Despite their strong theoretical foundation, these methods can struggle to effectively capture the full geometry of set-structured inputs Wass_pooling.

In this work we are introducing Wasserstein Hypergraph Neural Networks (WHNN), a class of hypergraph models that uses Sliced Wasserstein Pooling (SWP)  (Wass_pooling,) as a node and hyperedge aggregator. This pooling is based on the Wasserstein distance - an optimal transport metric which measures the distance between two distributions based on the cost of transporting mass from one to another. SWP treats the set elements as samples from an underlying distribution and generates a vector representation that captures geometric relationships between inputs, such as shape, spread, and density. This information is often lost when aggregating using summation.

We argue that this geometric information is highly relevant for hypergraph learning. Our experimental results support this claim, showing that WHNN not only outperforms traditional set-based aggregation methods used in previous hypergraph models, but also achieves superior performance compared to several strong hypergraph methods across a range of real-world datasets.

Our main contributions are summarised as follow:

- 1.
We propose a novel hypergraph architecture, the Wasserstein Hypergraph Neural Network (WHNN), which leverages Sliced Wasserstein Pooling for both node and hyperedge aggregation to more effectively capture the geometric structure of the feature space.
- 2.
We empirically show that Wasserstein aggregation is highly effective for hypergraph representation, consistently outperforming traditional sum-based methods such as Deep Sets DeepSets and Set Transformers set_transformer, regardless of the encoder used to process the nodes.
- 3.
Wasserstein Hypergraph Neural Network achieves top results on multiple real-world datasets, highlighting the advantages of incorporating optimal transport into hypergraph processing.

The paper is structured as follows: Sections [2](https://arxiv.org/html/2506.09682v1#S2) and  [3](https://arxiv.org/html/2506.09682v1#S3) describe similar efforts and introduce background concepts. Sections [4](https://arxiv.org/html/2506.09682v1#S4) introduces the architectures and innovative aspects of the methodology. Then, Section [5](https://arxiv.org/html/2506.09682v1#S5) contains and discusses experiments supporting the claimed contributions.

## 2 Related Work

#### Hypergraph representation learning.

Hypergraphs represent a versatile structure for modeling group-wise interactions, which allows us to capture interactions between various number of elements. This flexibility, combined with the widespread presence of higher-order interactions in real-world scenarios, has led to a growing interest in developing machine learning architectures for modeling hypergraph data. Some methods HyperNN; SHNN reduce the hypergraph to a clique-expansion graph that can be further processed with standard graph neural networks. A more popular approach is based on a two-stage framework allset; UniGNN, which sends the information from node to hyperedges and then from hyperedges back to nodes. Depending on how these stages are instantiated, several architectures emerged. HCHA and HERALD (HCHA,; HERALD,) use an attention mechanism to combine the information, AllDeepSets (allset,) uses Deep Set model, while AllSetTransformer (allset,) is using a PMA-like set_transformer pooling.

In all of these methods the information sent from the node is independent of the target hyperedge. Recently, models that create edge-dependent node representation have gained traction. ED-HNN wang2022equivariant uses as messages a concatenation of node and hyperedge information, while MultiSetMixer telyatnikov2025hypergraph uses MLP-Mixer tolstikhin2021mlpmixer to combine the information. Similar to our node encoder, CoNHDzheng2025corepresentation incorporates pairwise propagation at the hyperedge-level using self-attention blocks (SAB set_transformer) to create edge-dependent representations. However, similar to whatsnet, the model is only tested on hyperedge-dependent node classification tasks, where each node is assigned multiple labels corresponding to the number of hyperedges it participates in. A complementary line of work thnn is representing uniform hypergraphs as high-dimensional tensors and applies tensorial operators to propagate the information.

In contrast, we are interpreting the hyperedges as samples from a set of probability distributions, and uses Sliced Wasserstein Pooling to aggregate the information such that we preserve geometric information. In terms of node encoders, we are experimenting with both edge-dependent and edge-independent modules.

#### Set representation learning.

The core operation in set representation learning is the permutation-invariant operator that aggregate the information without imposing an order among elements. Popular examples of such operator include summation, mean or maximum. More recently, learnable version of permutation-invariant poolings were introduced. Among these, Deep Sets DeepSets is using element-wise encoding of the elements followed by summation and is proved to be universal approximator for permutation-invariant functions. Janossy Pooling murphy2018janossy extends this model by explicitly aggregating pairs of elements. On the other hand, Set Transformer set_transformer and Skianis2019RepTS uses an anchor set as a reference and compute the similarity against this set as a representation, while FSPool Zhang2019FSPool sorts the elements feature-wise to create a canonical order. Recently,  equiv_vs_invariant shows empirically that combining an equivariant backbone with an invariant pooling layer creates powerful set representation learning. Inspired by optimal transport literature, Sliced Wasserstein Pooling was introduced in  Wass_pooling as a geometrically-interpretable set representation technique.

#### Wasserstein embeddings.

In recent years, Wasserstein distance has attracted significant attention in deep learning, demonstrating success in areas such as generative modeling WGAN; nguyen2021distributional, natural language processing frogner2019learningembeddingsentropicwasserstein and point cloud processing point_cloud_Wass. In graph representation learning, Wasserstein distance was used to define a similarity kernel between pair of graphs Togninalli19. While recognised as a powerful tool, computing this distance for each pair of compared graphs is extremely inefficient. More recent works WGRL; OTK; courty2018learning try to reduce this cost by introducing Wasserstein embeddings. The purpose of a Wasserstein embedding is to infer a vector representation such that the $L_{2}$ distance in the vector space approximates the Wasserstein distance in the input space. Particularly important for us is the work of  Wass_pooling which produces set representations using efficient Wasserstein embeddings.

To more effectively capture the internal structure of node and hyperedge neighborhoods, we employ Sliced Wasserstein Pooling as the aggregation operator in hypergraph message passing, demonstrating its advantages for hypergraph representation learning.

## 3 Background

### 3.1 Hypergraph Representation Learning

A hypergraph is a tuple $\mathcal{H}=(V,E)$ where $V=\{v_{1},v_{2}\dots v_{N}\}$ is a set of nodes, and $E=\{e_{1},e_{2}\dots e_{M}\}$ is a set of hyperedges. Different than the graph structure, where each edge contains exactly two nodes, in a hypergraph each hyperedge contains a set of nodes which can vary in cardinality. Each node $v_{i}$ is characterize by a feature vector $x_{i}\in\mathbb{R}^{d}$. We denote by neighbourhood of hyperedge $e_{i}$ the set of nodes that are part of that hyperedge $\{v_{j}|v_{j}\in e_{i}\}$. Similarly, the neighbourhood of a node $v_{i}$ is the set of all hyperedges containing that node $\mathcal{N}_{v_{i}}=\{e_{j}|v_{i}\in e_{j}\}$.

Several architectures were developed for hypergraph-structured input HyperNN; wang2022equivariant; ijcai21-UniGNN; allset. However, the most general pipeline follow a two-stage framework, inspired by the bipartite representation of the hypergraphs.
First, the information is sent from nodes to the hyperedges using a permutation-invariant operator $z_{j}=f_{V\rightarrow E}(\{x_{i}|v_{i}\in e_{j}\})$. Secondly, the messages are sent back from hyperedge to nodes $\tilde{x}_{i}=f_{E\rightarrow V}(\{z_{j}|v_{i}\in e_{j}\})$.

**Table 1: The update rules used as aggregation steps in various hypergraph neural networks from the literature. As shown by the equations, all of them relies either on (weighted) summation or mean to combine the information. While theoretically powerful, summing can easily destroy all the geometric relationships between points. $\mathcal{N}_{v}(i)$ is the neighbourhood of node $v_{i}$ of cardinality $d_{i}$, $\mathcal{N}_{e}(j)$ is the neighbourhood of edge $e_{j}$ of cardinality $d_{j}$ and $\epsilon$, $W_{*}$, $\tilde{W}_{*}$ are learnable parameters.**
| Model | Hyperedge aggregation | Node aggregation |
| --- | --- | --- |
| HGNN | $h_{e}\leftarrow\sum_{i\in\mathcal{N}_{e}(e)}\frac{1}{\sqrt{d_{i}}}x_{i}W$ | $x_{i}\leftarrow\frac{1}{\sqrt{d_{i}}}\sum_{e\in\mathcal{N}_{v}(i)}\frac{1}{d_{<br>e}}h_{e}$ |
| HCHA(^1^11The coefficients $\alpha_{e,i}$ used in summations are scalars predicted as $MLP(x_{i}||h_{e})$) | $h_{e}\leftarrow\sum_{i\in\mathcal{N}_{e}(e)}\alpha_{e,i}x_{i}W$ | $x_{i}\leftarrow\sum_{e\in\mathcal{N}_{v}(i)}\tilde{\alpha}_{i,e}h_{e}\tilde{W}$ |
| UniGIN | $h_{e}\leftarrow\sum_{i\in\mathcal{N}_{e}(e)}x_{i}$ | $x_{i}\leftarrow\sum_{e\in\mathcal{N}_{v}(i)}h_{e}W+(1+\epsilon)x_{i}W$ |
| ED-HNN | $h_{e}\leftarrow\sum_{i\in\mathcal{N}_{e}(e)}\text{MLP}(x_{u})$ | $x_{i}\leftarrow\sum_{e\in\mathcal{N}_{v}(i)}\text{MLP}(x_{i}\|h_{e})$ |
| AllDeepSets | $h_{e}\leftarrow\text{MLP}(\sum_{i\in\mathcal{N}_{e}(e)}\text{MLP}(x_{u}))$ | $x_{i}\leftarrow\text{MLP}(\sum_{e\in\mathcal{N}_{v}(i)}\text{MLP}(h_{e}))$ |
| AllSetTransformer(^2^22The function $\sigma$ is a combination of residual connections and layer normalisations, while $\alpha_{i}=(\theta W_{q})(x_{i}W_{k})^{T}$ with $\theta$, $W_{q}$ and $W_{k}$ as learnable parameters.) | $h_{e}\leftarrow\sigma(\sum_{i\in\mathcal{N}_{e}(e)}(\alpha_{i}x_{i}W_{v})$ | $x_{i}\leftarrow\sigma(\sum_{e\in\mathcal{N}_{v}(i)}(\tilde{\alpha}_{e}h_{e}<br>\tilde{W}_{v})$ |

While aggregators like Deep Sets DeepSets were theoretically capable of approximating any permutation-invariant function on sets, it relies on the initial encoder (such as MLPs) to reshape the feature space in a way in which the sum pooling is not losing important information. In other words, it moves the complexity of the representation from the pooling to the initial encoding. This is in line with the empirical results shown in  equiv_vs_invariant where, in order to preserve good performance, mean pooling requires more complex encoders compared to more sophisticated pooling methods.

In this work we are following the standard two-stage framework. Compared to existing methods, we take advantage of the success demonstrated by Sliced Wasserstein Pooling in capturing and retaining the geometric structure of sets and proposed the first hypergraph model that uses optimal transport techniques to perform the node and hyperedge aggregation.

### 3.2 Sliced Wasserstein Pooling (SWP)

To ensure the method’s readability, this section introduces all the key concepts underlying our Wasserstein Hypergraph Neural Network.
First, we will define the 2-Wasserstein metric, approximate it using the tractable Sliced-Wasserstein distance and finally present the algorithm to compute the SWP used as an aggregator in our model.

###### Definition 1 .

The 2-Wasserstein distance between two distributions $p_{i}$ and $p_{j}$ over $\mathbb{R}^{d}$ is defined as:

$$ $\mathcal{W}_{2}(p_{i},p_{j})=\Big{(}\inf_{\gamma\in\Gamma(p_{i},p_{j})}\int_{ \mathbb{R}^{n}\times\mathbb{R}^{n}}||x-y||^{2}d\gamma(x,y)\Big{)}^{\frac{1}{2}},$ (1) $$

where $\Gamma(p_{i},p_{j})$ represent the collection of all the transport plans with marginals $p_{i}$ and $p_{j}$.

In simpler terms, the 2-Wasserstein distance quantifies the cost of transforming one distribution into another.

Unfortunately, computing the infimum over all possible transport maps is generally untractable. However, in the one dimensional case (when $d=1$) ), a closed-form solution exists that avoids expensive optimization. Specifically, when $p_{i}$ and $p_{j}$ are probability distributions over $\mathbb{R}$, the 2-Wasserstein distance is given by $\mathcal{W}_{2}(p_{i},p_{j})=\Big{(}\int_{0}^{1}|F_{p_{i}}^{-1}(t)-F_{p_{j}}^{
-1}(t)|^{2}dt\Big{)}^{\frac{1}{2}}$, where $F_{p_{i}}^{-1}$ and $F_{p_{j}}^{-1}$ denote the inverse cumulative distribution functions of $p_{i}$ and $p_{j}$. A key practical benefit of this formulation is that this inner integral can be empirically estimated using a discrete sum over sorted samples from the distribution.

Building on this observation, Sliced Wasserstein distance sliced_wass was introduce to approximate the Wasserstein distance, by projecting the high-dimensional probabilities into 1D lines using all possible directions on the unit sphere.

###### Definition 2 .

The Sliced Wasserstein distance between two distributions $p_{i}$ and $p_{j}$ over $\mathbb{R}^{d}$ is defined as:

$$ $\mathcal{SW}_{2}(p_{i},p_{j})=\Big{(}\int_{S^{d-1}}\mathcal{W}_{2}(P_{\theta}p _{i},P_{\theta}p_{j})d\theta\Big{)}^{\frac{1}{2}}\approx\Big{(}\frac{1}{L}\sum _{l=1}^{L}\underbrace{{\mathcal{W}_{2}(P_{\theta_{l}}p_{i},P_{\theta_{l}}p_{j} )}}_{\text{{\color[rgb]{.5,.5,.5}\definecolor[named]{pgfstrokecolor}{rgb}{ .5,.5,.5}\pgfsys@color@gray@stroke{.5}\pgfsys@color@gray@fill{.5}1D Wasserstein distance}}}\Big{)}^{\frac{1}{2}},$ (2) $$

where $S^{d-1}$ is the unit sphere in $\mathbb{R}^{d}$, $P_{\theta}p_{i}$ represent the projection (pushforward) of $p_{i}$ onto the line direction $\theta$ and $\{\theta_{l}\}_{l=1}^{L}$ represents the set of $L$ directions used to empirically approximate the expectation.

To avoid the computational cost of calculating distances between every pair of probability distributions, the Sliced Wasserstein embedding Wass_pooling was introduced. It maps a probability distributions $p_{i}$ to a vector $\phi(p_{i})$ in such a way that the Euclidean distance between the vectors (which is inexpensive to compute) approximates the Sliced Wasserstein distance between the original distributions $||\phi(p_{i})-\phi(p_{j})||_{2}\approx\mathcal{SW}_{2}(p_{i},p_{j})$. In other words, it provides a vectorial representation that captures the geometric structure of distributions, preserving information about how costly it is to transform one distribution into another. This geometric encoding reflects characteristics such as shape, spread, and density. This proves useful in our context, as it allows us to quantify the cost of transforming one hyperedge into another — a measure we argue effectively captures the similarity between group interactions (hyperedges).

Since our nodes and hyperedges are sets rather than distributions, we use a variant of this embedding called Sliced Wasserstein Pooling Wass_pooling, which is designed not as an embedding of probability distributions themselves, but rather as an embedding of sets sampled from those distributions.
In short, Sliced Wasserstein Pooling encodes a set of points by measuring, in an efficient way, how different they are positioned compared to a set of reference points. The complete algorithm as used in our model is described in the following section.

Figure: Figure 1: (A) One stage (node-to-hyperedge) of Wasserstein Hypergraph Neural Network pipeline designed to be more sensitive to the geometric structure of the hyperedge compared to the traditional aggregators. First, a node encoder processes the nodes using a simple MLP or an edge-dependent self-attention block (SAB). The hypergraphs is than viewed as a collection of probability distributions $\{p_{i}\}$, one for each hyperedge, with the observed nodes treated as samples drawn from it. An additional distribution $q$ is picked as a reference. Finally the Sliced Wasserstein Pooling is adopted as an aggregation method: each hyperedge is represented by its Sliced Wasserstein distance to a reference distribution. (B) Compared to the standard mean pooling which only capture the difference between the mean of the distributions (visualised as a cross), the euclidean distance between the obtained hyperedges quantify the cost of transforming one group into another.
Refer to caption: extracted/6532951/figures/main_tall.png

###### Definition 1 .

###### Definition 2 .

## 4 Wasserstein Hypergraph Neural Network

Taking inspiration from the success of Wasserstein embeddings in set representation learning  Wass_pooling; OTK,
we are introducing Wasserstein Hypergraph Neural Network (WHNN), a neural network for processing hypergraph structured data which replace the standard (weighted) mean aggregator with SWP, thus better capturing the internal structure of the neighbourhoods.

The model follows the two-stage framework introduced in Section [3.1](https://arxiv.org/html/2506.09682v1#S3.SS1), by sending information from nodes to hyperedges and vice-versa. For simplicity this section only describes the nodes to hyperedges mechanism, as the hyperedge-to-node operation is entirely symmetrical. The entire pipeline is depicted in Figure [1](https://arxiv.org/html/2506.09682v1#S3.F1) and Algorithm [1](https://arxiv.org/html/2506.09682v1#alg1). For readability, the algorithm is presented sequentially for each hyperedge. However our implementation processes all hyperedges in parallel.

First we will project the node features into a more expressive representation. Each hyperedge is then associated with a probability distribution, with its constituent nodes treated as samples. These distributions are embedded using a Wasserstein-based aggregator to obtain the final hyperedge representations. These hyperedge representations are fed into the hyperedges-to-nodes stage.

#### Node encoder.

The goal of this module is to enhance the representation of node features by projecting them into a more informative space. We are experimenting with two types of encoders: an edge-independent one where the node is carrying the same representation in each hyperedge it is contained, and an edge-dependent one which takes into account pairwise interactions.

The edge-independent encoder is a simple MLP, which is applied in parallel for each node. This way a node $i$ is characterized by the same feature vector in each hyperedge $e$ it is part of.

$$ $\tilde{x}_{i}^{e}=MLP(x_{i})$ $$

On the other hand, for the edge-dependent encoder, each node has a different representation in each hyperedge it is part of. To achieve this, for each hyperedge, we are using a Set Attention Block layer (SAB) as introduce in set_transformer which propagates the information between each pair of two nodes contained in that hyperedge. The full version of the block acts as follow:

$$ $\displaystyle z_{i}^{e}$ $\displaystyle=\sigma(x_{i}+\sum_{j\in e}(x_{i}W_{q})(x_{j}W_{k})^{T}(x_{j}W_{v }))$ $\displaystyle\tilde{x}_{i}^{e}$ $\displaystyle=\sigma(z_{i}^{e}+MLP(z_{i}^{e})),$ $$

where $\sigma$ denote layer normalisation and $W_{k}$, $W_{q}$ and $W_{v}\in\mathbb{R}^{d\times d}$ are learnable parameters.

Figure: Algorithm 1 One Layer of Wasserstein Hypergraph Neural Network^∗

#### Hyperedges as probability distributions.

Unlike traditional hypergraph approaches that treat a hyperedge as a set of nodes, we model a hyperedge as a probability distribution, with its constituent nodes being samples drawn from that distribution. This way the hyperedges are not only characterized by the combination of its elements, but by the regions of the space where its elements are situated. The nodes became prototypes of the hyperedge behaviour. For example, a hyperedge in which nodes have similar representations ( homophilic behaviour) indicates a low-variance distribution while a hyperedge with diverse nodes suggests a more uniform distribution.

Lets consider $p_{i}$ the probability distribution where the elements of the hyperedge $e_{i}$ are sampled from. In other words, we assume each node $v_{j}\in e_{i}$ is sampled as $\tilde{x}^{i}_{j}\in\mathbb{R}^{d}\sim p_{i}$. The goal is to obtain hyperedge embeddings that preserve the geometric information of this underlying distribution, such as spreading, shape etc. See Figure [1](https://arxiv.org/html/2506.09682v1#S3.F1) for a visual representation of this data structure.

Note that, by treating nodes as sampled from an underlying distribution, we make the assumption that other unobserved nodes drawn from the same distribution are also likely to belong to the same hyperedge. This probabilistic interpretation proved to be powerful for set representation learning Wass_pooling and our experiments demonstrates that hypergraph models can benefit from it as well.

#### Wasserstein aggregator.

Interpreting hypergraphs as a collection of probability distributions enable us to derive more powerful similarity metrics between hyperedges. As showed in the previous section, most of the current hypergraph architectures rely on mean pooling to create hyperedge embeddings from node representations. However, from a probabilistic perspective, averaging compares distributions only based on their means. For complex data distributions, this approach fails to capture the full underlying geometry. While models relying on summation such as Deep Sets DeepSets where proved to be universal approximators, they heavily rely on the internal node encoder (an MLP) to map input features into a space where first-order statistics like the mean effectively approximate the distribution. In the hypergraph setting, where multiple sets interact in complex ways, this is hard to achieve.

This motivates us to adopt Sliced Wasserstein Pooling Wass_pooling to encode the hyperedge distributions. Concretely, for each hyperedge $e$, given the node embeddings of all the nodes in the hyperedge $\{~{}\tilde{x}_{i}^{e}\}_{i\in e}$, we are aggregating them using the Sliced Wasserstein Pooling described in Section [3.2](https://arxiv.org/html/2506.09682v1#S3.SS2), to obtain a vectorial hyperedge representation: $h_{e}=\mathcal{SWP}(\{~{}\tilde{x}_{i}^{e}\}_{i\in e}$). The algorithm works as follow:

- 1.
Step 1: Select a reference hyperedge distribution $q$ and sample $N$ points $\{y_{i}\}_{i=1}^{N}\sim q$. Choose a set of directions $\{\theta_{l}\}_{l=1}^{L}$ with $\theta_{l}\in\mathbb{R}^{d\times 1}$ used as projection slices in the pooling process. Note that, in order to obtain comparable embeddings across the entire hypergraph, we share the same reference distribution and the same set of slices for all hyperedges.
- 2.
Step 2: Project each node representation $\tilde{x}_{i}^{e}$ into each slice $\theta_{l}$ as follow: $z_{i}^{e,\theta_{l}}=(\tilde{x}_{i}^{e})^{T}\theta_{l}\in\mathbb{R}$. Since the algorithm requires the same number of sampled nodes from both the hyperedge distribution and the reference, when the cardinality of the hyperedge $|e|\neq N$, we increase/decrease the number of nodes in $e$ using linear interpolation. $z_{i}^{e,\theta_{l}}\leftarrow\text{interp}(z_{i}^{e,\theta_{l}},N)$
- 3.
Step 3: For each hyperedge, for each slice, compute the distance between the node representations and the reference points. $h_{e}^{\theta_{l}}=||z_{\pi(i)}^{e,\theta_{l}}-y_{\tilde{\pi}(i)}||$, where $z^{e,\theta_{l}}_{\pi}$ and $y_{\tilde{\pi}}$ represent the vectors in sorted order. The final hyperedge embedding is obtain as a weighted mean of these embeddings: $h_{e}=\sum_{l=1}^{L}(w_{l}h_{e}^{\theta_{l}})$, where $w_{l}$ are learnable scalars combining the slices.

Intuitively, each hyperedge is represented by a vector which measure how difficult it is to transform the hyperedge distribution into the reference distribution (^3^33As defined above, by hyperedge distribution we denote the distribution of nodes in the hyperedge.). Following the theoretical properties of Sliced Wasserstein Pooling Wass_pooling, the Euclidean distance between two hyperedge representations measures the cost of transforming one hyperedge distribution into another. Hyperedges that are similar in shape or spreading should be closer in this space compared to hyperedges that have completely different distributions. The algorithm is also described in Algorithm [2](https://arxiv.org/html/2506.09682v1#alg2). The directions $\theta_{l}$ and the reference distribution can be either fixed or learnable.

#### Edge to node step.

For simplicity, we only described in details the first stage of the framework which sends messages from nodes to hyperedges. The second stage of the framework which create node representation by aggregating the information from neighbouring hyperedges is done in a similar way, only with different parameters. In conclusion, we not only capture structural relationship between hyperedges, but also structural relationship between nodes’ neighbourhood.

Figure: Figure 2: Ablation study on the importance of Wasserstein aggregator for hypergraph representation learning on Citeseer and NTU2012 datasets. We are testing two versions of the Sliced Wasserstein Pooling: with fixed (FPSWE) or learnable (LPSWE) reference distribution. Regardless of the encoder used to project the nodes and hyperedges, the Wasserstein aggregators outperform both the Deep Sets and PMA commonly used inside hypergraph models.
Refer to caption: extracted/6532951/figures/citeseer_ntu.png

## 5 Experiments

Our main goal is to understand to what extent Wasserstein aggregation is beneficial for hypergraph neural networks. Additionally, we investigate how the choice of node encoder—whether edge-dependent or edge-independent—affects overall performance. Finally, we compare our model against a range of strong baseline methods from the existing literature.

#### Datasets.

We evaluate our model on the node-classification task.
We select seven real-world datasets that vary in domain and scale. These include Cora, Citeseer, Cora-CA, DBLP-CA HyperGCN, ModelNet40 modelnet, NTU2012 ntu2012 and 20News 20News. Among the datasets that are usually used for benchmarking hypergraph models HyperGCN, we omitted Pubmed due to the high percentage of isolated nodes ($80.5\%$) which makes the relational processing unnecessary. Senate and House are two other datasets used to test hypergraphs in the heterophilic regime. However, they do not provide node features which are a key component for our geometric and probabilistic interpretation.

For a fair comparison with the other methods, we follow the training procedures employed by wang2022equivariant. We randomly split the data into $50\%$ training samples, $25\%$ validation samples and $25\%$ test samples.

#### Importance of Wasserstein aggregator.

Our main contribution consists of adopting Sliced Wasserstein Pooling as a powerful aggregator inside the hypergraph networks. As described in the previous section, while most of the existing methods used variations of the sum pooling to aggregate the information from each node and each hyperedge neighbourhoods, our Wasserstein aggregator presents a more in-depth understanding of the neighbourhood distribution, being capable of capturing subtle differences such as the difference in shape or spread.

To understand to what extent this is contributing to a better hypergraph representation for real-world scenarios, we are designing an ablation study in which we keep the underlying architecture fixed and only modify the aggregator used in both the nodes-to-hyperedges and hyperdges-to-nodes stages. Concretely we are using as aggeregators either Deep Set module (as used by AllDeepSet and ED-HNN models) or the PMA module (as used by AllSetTransformer model). For our Wasserstein aggregator, we are experimenting with both a fixed-reference distribution (a model denoted as FPSWE) or with learnable reference distribution (a model denoted as LPSWE). For a robust evaluation we are comparing this aggregators using both the edge-independent encoder (MLP) and the edge-dependent encoder (SAB). The results on Citeseer and NTU2012 datasets are reported in Figure [2](https://arxiv.org/html/2506.09682v1#S4.F2).

Regardless of the encoder and the dataset we are testing on, both Wasserstein aggregators are consistently outperforming both the Deep Sets and the PMA aggregators by a significant margin. A learnable reference seems to be beneficial, however the improvement is generally marginal.
Additional experiments on other datasets show a similar trend and are provided in the Supplementary Material.

#### Importance of edge-dependent encoder.

The node and hyperedge encoder transforms features into a space where their distribution within each hyperedge captures meaningful information about the group. As stated in the model description, we equipped our model with two types of encoders. An edge-independent module represented by an MLP, and an edge-dependent encoder represented by a self-attention block (SAB). While the MLP is processing information independently for each node/hyperedge, SAB is capturing pairwise interactions between nodes/hyperedges sharing a neighbourhood.
The results in Figure [2](https://arxiv.org/html/2506.09682v1#S4.F2) and Table [2](https://arxiv.org/html/2506.09682v1#S5.T2) show similar results among the encoder, with the edge-dependent one being slightly more powerful. However, this comes with the cost of a more expensive model, as the edge-dependent encoder requires more memory to store the representation for all incident pairs (node, hyperedge). To alleviate that on the larger datasets (20News and DBLP), we replace the SAB block with the ISAB low-rank approximation introduced by  set_transformer.

**Table 2: Performance on a collection of hypergraph datasets. Our model using SWP as a node and hyperedge aggregator shows superior results, proving the advantage of moving beyond the standard sum pooling employed by most of the existing works. We test our model in both its variants: with edge-independent (MLP) and edge-dependent encoder (SAB). Both options are exhibiting competitive performance. We mark the first, second and third best performing models for each dataset.**
| Name | Cora | Citeseer | Cora_CA | DBLP_CA | ModelNet40 | NTU2012 | 20News |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| HCHA | 79.14 $\pm$ 1.02 | 72.42 $\pm$ 1.42 | 82.55 $\pm$ 0.97 | 90.92 $\pm$ 0.22 | 94.48 $\pm$ 0.28 | 87.48 $\pm$ 1.87 | 80.33 $\pm$ 0.80 |  |
| HNHN | 76.36 $\pm$ 1.92 | 72.64 $\pm$ 1.57 | 77.19 $\pm$ 1.49 | 86.78 $\pm$ 0.29 | 97.84 $\pm$ 0.25 | 89.11 $\pm$ 1.44 | 81.35 $\pm$ 0.61 |  |
| HyperGCN | 78.45 $\pm$ 1.26 | 71.28 $\pm$ 0.82 | 79.48 $\pm$ 2.08 | 89.38 $\pm$ 0.25 | 75.89 $\pm$ 5.26 | 56.36 $\pm$ 4.86 | 81.05 $\pm$ 0.59 |  |
| HyperGNN | 79.39 $\pm$ 1.36 | 72.45 $\pm$ 1.16 | 82.64 $\pm$ 1.65 | 91.03 $\pm$ 0.20 | 95.44 $\pm$ 0.33 | 87.72 $\pm$ 1.35 | 80.33 $\pm$ 0.42 |  |
| AllDeepSets | 76.88 $\pm$ 1.80 | 70.83 $\pm$ 1.63 | 81.97 $\pm$ 1.50 | 91.27 $\pm$ 0.27 | 96.98 $\pm$ 0.26 | 88.09 $\pm$ 1.52 | 81.06 $\pm$ 0.54 |  |
| AllSetTransformers | 78.58 $\pm$ 1.47 | 73.08 $\pm$ 1.20 | 83.63 $\pm$ 1.47 | 91.53 $\pm$ 0.23 | 98.20 $\pm$ 0.20 | 88.69 $\pm$ 1.24 | 81.38 $\pm$ 0.58 |  |
| UniGCNII | 78.81 $\pm$ 1.05 | 73.05 $\pm$ 2.21 | 83.60 $\pm$ 1.14 | 91.69 $\pm$ 0.19 | 98.07 $\pm$ 0.23 | 89.30 $\pm$ 1.33 | 81.12 $\pm$ 0.67 |  |
| ED-HNN | 80.31 $\pm$ 1.35 | 73.70 $\pm$ 1.38 | 83.97 $\pm$ 1.55 | 91.90 $\pm$ 0.19 | 97.75 $\pm$ 0.17 | 89.48 $\pm$ 1.87 | $81.36\pm 0.55$ |  |
| WHNN_MLP | 79.84 $\pm$ 1.56 | 74.79 $\pm$ 1.19 | 84.12 $\pm$ 1.94 | 91.73 $\pm$ 0.24 | 98.47 $\pm$ 0.19 | 90.87 $\pm$ 1.59 | 81.83 $\pm$ 0.68 |  |
| WHNN_(I)SAB | 80.72 $\pm$ 1.96 | 74.92 $\pm$ 1.60 | 84.62 $\pm$ 1.77 | 91.99 $\pm$ 0.33 | 98.54 $\pm$ 0.21 | 90.68 $\pm$ 1.68 | 81.42 $\pm$ 0.60 |  |

Comparison with baselines. In Table [2](https://arxiv.org/html/2506.09682v1#S5.T2) we are comparing against a series of hypergraph networks from the literature. With respect to aggregation strategies, HNHN hnhn, HyperGNN HyperNN, AllDeepSets allset, UniGCNII UniGNN and ED-HNN wang2022equivariant use variations of Deep Sets to aggregate the information, HyperGCN HyperGCN uses a max aggregator, while HCHA HCHA and AllSetTransformer allset are using an attention-based weighted summation. Regardless of the encoder used, our model consistently obtain top results, outperforming the other methods on all datasets. This demonstrates the advantages of using Wasserstein aggregators for higher-order processing. Note that, while we integrated this aggregator into a standard instantiation of the two-stage framework, many existing models from the literature can be adopted to take advantage of this type of geometric-inspired aggregation.

Implementation details. In all experiments, we train our models using Adam optimizer for $500$ epochs, on a single GPU NVIDIA Quadro RTX 8000 with 48GB of memory. For comparing against other models in the literature, each model is trained $10$ times with different random splits and different initialization. For each experiment we report average accuracy along with the standard deviation. The results represent the best performance obtained by each architecture using hyper-parameter optimisation with random search. For the ablation study the results are averaged across $5$ runs and the architecture is fixed to ensure a fair comparison. For all experiments we use a number of Wasserstein slices equal to the hidden dimension and we experiment with both learning the reference set or not. Details on all the model choices and hyper-parameters can be found in the Supplementary Material.

These experimental results show that aggregating node and hyperedge neighborhoods using Sliced Wasserstein Pooling is highly effective for hypergraph processing, the Wasserstein aggregator consistently outperforming standard methods like Deep Sets and PMA.

## 6 Conclusion

In this work we introduce Wasserstein Hypergraph Neural Networks (WHNN), a model for processing hypergraph structures. The model relies on Sliced Wasserstein Pooling to aggregate the nodes into hyperedge representations and vice versa. This design choice inspired by optimal transport literature enable us to capture more information about the internal structure of the neighbourhoods, preserving more geometric relation between elements. The experimental results on various datasets demonstrates that this Wasserstein aggregator is effective for modeling higher-order interactions, outperforming traditional aggregators, making WHNN a promising tool for hypergraph representation learning.

#### Acknowledgment.

The authors would like to thank Daniel McFadyen and Alex Norcliffe for fruitful discussions and constructive suggestions during the development of the paper. I.D. additionally acknowledges the support provided by Robinson College during this period.

## References

- [1]
Ruth Johnson, Michelle M. Li, Ayush Noori, Owen Queen, and Marinka Zitnik.
Graph ai in medicine.
CoRR, abs/2310.13767, 2023.
- [2]
Catherine Tong, Emma Rocheteau, Petar Veličković, Nicholas Lane, and Pietro Lio.
Predicting Patient Outcomes with Graph Representation Learning, pages 281–293.
01 2022.
- [3]
Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia.
Learning to simulate complex physics with graph networks.
In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119, pages 8459–8468, 2020.
- [4]
Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Ferran Alet, Suman Ravuri, Timo Ewalds, Zach Eaton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Oriol Vinyals, Jacklynn Stott, Alexander Pritzel, Shakir Mohamed, and Peter Battaglia.
Graphcast: Learning skillful medium-range global weather forecasting, 2023.
- [5]
Federico Monti, Fabrizio Frasca, Davide Eynard, Damon Mannion, and Michael M Bronstein.
Fake news detection on social media using geometric deep learning.
arXiv preprint arXiv:1902.06673, 2019.
- [6]
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl.
Neural message passing for quantum chemistry.
In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1263–1272, 2017.
- [7]
Kexin Huang, Cao Xiao, Lucas M Glass, Marinka Zitnik, and Jimeng Sun.
Skipgnn: predicting molecular interactions with skip-graph networks.
Scientific reports, 10(1):1–16, 2020.
- [8]
Eli Chien, Chao Pan, Jianhao Peng, and Olgica Milenkovic.
You are allset: A multiset function framework for hypergraph neural networks.
In International Conference on Learning Representations, 2022.
- [9]
Peihao Wang, Shenghao Yang, Yunyu Liu, Zhangyang Wang, and Pan Li.
Equivariant hypergraph diffusion neural operators.
arXiv preprint arXiv:2207.06680, 2022.
- [10]
Jing Huang and Jie Yang.
Unignn: a unified framework for graph and hypergraph neural networks.
In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21, 2021.
- [11]
Jiahao Xie and Guangmo Tong.
Advances in set function learning: A survey of techniques and applications, 2025.
- [12]
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola.
Deep sets.
In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
- [13]
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh.
Set transformer: A framework for attention-based permutation-invariant neural networks.
In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3744–3753. PMLR, 09–15 Jun 2019.
- [14]
Navid Naderializadeh, Joseph F Comer, Reed Andrews, Heiko Hoffmann, and Soheil Kolouri.
Pooling by sliced-wasserstein embedding.
In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 3389–3400. Curran Associates, Inc., 2021.
- [15]
Yifan Feng, Haoxuan You, Zizhao Zhang, Rongrong Ji, and Yue Gao.
Hypergraph neural networks.
Proc. Conf. AAAI Artif. Intell., 33(01):3558–3565, July 2019.
- [16]
Bohan Tang, Zexi Liu, Keyue Jiang, Siheng Chen, and Xiaowen Dong.
Hypergraph node classification with graph neural networks.
CoRR, abs/2402.05569, 2024.
- [17]
Song Bai, Feihu Zhang, and Philip H.S. Torr.
Hypergraph convolution and hypergraph attention.
Pattern Recognition, 110:107637, 2021.
- [18]
Jiying Zhang, Yuzhao Chen, Xiong Xiao, Runiu Lu, and Shutao Xia.
Learnable hypergraph laplacian for hypergraph learning.
In ICASSP, 2022.
- [19]
Lev Telyatnikov, Maria Sofia Bucarelli, Guillermo Bernardez, Olga Zaghen, Simone Scardapane, and Pietro Lio.
Hypergraph neural networks through the lens of message passing: A common perspective to homophily and architecture design.
Transactions on Machine Learning Research, 2025.
- [20]
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Peter Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy.
MLP-mixer: An all-MLP architecture for vision.
In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
- [21]
Yijia Zheng and Marcel Worring.
Co-representation neural hypergraph diffusion for edge-dependent node classification, 2025.
- [22]
Minyoung Choe, Sunwoo Kim, Jaemin Yoo, and Kijung Shin.
Classification of edge-dependent labels of nodes in hypergraphs.
In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, page 298–309. ACM, August 2023.
- [23]
Maolin Wang, Yaoming Zhen, Yu Pan, Yao Zhao, Chenyi Zhuang, Zenglin Xu, Ruocheng Guo, and Xiangyu Zhao.
Tensorized hypergraph neural networks, 2024.
- [24]
Ryan L. Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro.
Janossy pooling: Learning deep permutation-invariant functions for variable-size inputs.
In International Conference on Learning Representations, 2019.
- [25]
Konstantinos Skianis, Giannis Nikolentzos, Stratis Limnios, and Michalis Vazirgiannis.
Rep the set: Neural networks for learning set representations.
ArXiv, abs/1904.01962, 2019.
- [26]
Yan Zhang, Jonathon Hare, and Adam Prügel-Bennett.
FSPool: Learning set representations with featurewise sort pooling.
2019.
- [27]
Abihith Kothapalli, Ashkan Shahbazi, Xinran Liu, Robert Sheng, and Soheil Kolouri.
Equivariant vs. invariant layers: A comparison of backbone and pooling for point cloud classification, 2024.
- [28]
Martin Arjovsky, Soumith Chintala, and Léon Bottou.
Wasserstein generative adversarial networks.
In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 214–223. PMLR, 06–11 Aug 2017.
- [29]
Khai Nguyen, Nhat Ho, Tung Pham, and Hung Bui.
Distributional sliced-wasserstein and applications to generative modeling.
In International Conference on Learning Representations, 2021.
- [30]
Charlie Frogner, Farzaneh Mirzazadeh, and Justin Solomon.
Learning embeddings into entropic wasserstein spaces, 2019.
- [31]
Trung Nguyen, Quang-Hieu Pham, Tam Le, Tung Pham, Nhat Ho, and Binh-Son Hua.
Point-set distances for learning representations of 3d point clouds, 2021.
- [32]
Matteo Togninalli, Elisabetta Ghisu, Felipe Llinares-López, Bastian Rieck, and Karsten Borgwardt.
Wasserstein weisfeiler–lehman graph kernels.
In H. Wallach, H. Larochelle, A. Beygelzimer, F. d’Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32 (NeurIPS), pages 6436–6446. Curran Associates, Inc., 2019.
- [33]
Soheil Kolouri, Navid Naderializadeh, Gustavo K. Rohde, and Heiko Hoffmann.
Wasserstein embedding for graph learning.
In International Conference on Learning Representations, 2021.
- [34]
Grégoire Mialon, Dexiong Chen, Alexandre d’Aspremont, and Julien Mairal.
A trainable optimal transport embedding for feature aggregation and its relationship to attention, 2021.
- [35]
Nicolas Courty, Rémi Flamary, and Mélanie Ducoffe.
Learning wasserstein embeddings.
In International Conference on Learning Representations, 2018.
- [36]
Jing Huang and Jie Yang.
Unignn: a unified framework for graph and hypergraph neural networks.
In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21, 2021.
- [37]
Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister.
Sliced and radon wasserstein barycenters of measures.
51(1):22–45, January 2015.
- [38]
Naganand Yadati, Madhav Nimishakavi, Prateek Yadav, Vikram Nitin, Anand Louis, and Partha Talukdar.
Hypergcn: A new method for training graph convolutional networks on hypergraphs.
In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
- [39]
Zhirong Wu, Shuran Song, Aditya Khosla, Fisher Yu, Linguang Zhang, Xiaoou Tang, and Jianxiong Xiao.
3d shapenets: A deep representation for volumetric shapes, 2015.
- [40]
Ding-Yun Chen, Xiao-Pei Tian, Yu-Te Shen, and Ming Ouhyoung.
On visual similarity based 3d model retrieval.
Comput. Graph. Forum, 22:223–232, 09 2003.
- [41]
Tom Mitchell.
Twenty Newsgroups.
UCI Machine Learning Repository, 1997.
DOI: https://doi.org/10.24432/C5C323.
- [42]
Yihe Dong, Will Sawin, and Yoshua Bengio.
Hnhn: Hypergraph networks with hyperedge neurons.
In Graph Representation Learning and Beyond Workshop at ICML 2020, June 2020.
Code available: https://github.com/twistedcubic/HNHN.

## Appendix: Wasserstein Hypergraph Neural Network

This appendix contains details related to our model, including potential limitations and future work, additional datasets for the ablation experiments, details on the hyperparameters used in our experiments and derivation of the computational complexity. The appendix is structured as follow:

- •
Section A highlights a series of potential limitations that can be address to improve the current work, together with a discussion on potential future work.
- •
Section B presents additional experiments used as ablation for our model.
- •
Section C presents the list of hyperparameters used in our experiments.
- •
Section D derives the computational complexity of our model.

## Appendix A Limitations and Future work

As discussed in the main paper, we treat the neighborhood of each node as a sample from an underlying probability distribution. This approach assumes that any additional nodes drawn from this distribution should belong to the same neighborhood as the observed ones. This aligns with the intuition that elements within a group should share common characteristics. While the datasets we used support this assumption, there may be real-world scenarios where it does not hold. Our model relies solely on the node encoder to project features into a space where the assumption is approximately valid.

Moreover, due to this continuous view of the neighbourhood (as a distribution of probability) together with the interpolation step , the current model may lose information about the exact cardinality of the neighborhoods. In situations where neighborhood size is important, we recommend encoding it as an explicit feature. However, we mention that this is an issue we share with the mean-based pooling algorithms.

The main goal of this paper is to highlight the benefits of using geometrically-inspired poolings for aggregating neighbourhood information in hypergraphs. While we focused entirely on hypergraphs, similar idea can be apply on graph neural networks to aggregate messages coming from each node’s neighbourhood. As a future work, it would be interesting to see to what extent GNNs can benefit from Wasserstein aggregators.

Moreover, while the proposed model integrate the Wasserstein aggregator into a standard two-stage pipeline, several other architectures such as ED-HNN that uses summation as an aggregator might benefit from adopting it. We are leaving this investigation as future work.

## Appendix B Additional experiments

Due to space constraints, in the main paper we only included ablation studies on Citeseer and NTU datasets. Here we report additional results for Cora_CA and ModelNet40 datasets.

For each experiment, we kept the architecture fixed and modify the aggregator used in the two stages to be either Deep Set, PMA, and the learnable (LPSWE) or fixed (FPSWE) Wasserstein aggregator. The results are similar across the datasets, with Wasserstein Pooling proving to be beneficial compared to Deep Sets and PMA. In terms of encoder type, we noticed that, in some cases, for a fixed architecture, SAB tends to model the distribution better than MLPs.

Figure: Figure 3: Additional results for the ablation study on the importance of Wasserstein aggregator for hypergraph representation learning Cora_CA and ModelNet datasets. FPSWE denotes the Wasserstein aggregator with fixed reference while LPSWE denotes the Wasserstein aggergator with learnable reference distribution. Regardless of the encoder used to project the nodes and hyperedges, the Wasserstein aggregators outperform both the Deep Sets and PMA commonly used inside hypergraph models.
Refer to caption: extracted/6532951/figures/cora_ca_modelnet.png

## Appendix C Implementation details

The results reported in Table 2 of the main paper are obtained using random hyperparameter tuning. We report here the range of parameters that we searched for. Table [3](https://arxiv.org/html/2506.09682v1#A3.T3) and Table [4](https://arxiv.org/html/2506.09682v1#A3.T4) contains the best hyperparameter configuration for the WHNN_MLP model and WHNN_SAB. We depict in bold the parameters specific to the Wasserstein aggregator, in italic the parameters specific to the SAB encoder, while the rest of them are the standard parameters used in the two-stage hypergraph models. In our experiment we search for the following hyperparameters:

- •
num_ref: number of elements sampled from the reference distribution $\{5,10,25,50\}$
- •
learnable_W: choose between learning or not the reference distribution $\{\text{True},\text{False}\}$
- •
heads: number of heads used by the SAB block $\{1,2,4\}$
- •
MLP_layers: number of layers in all MLPs used $\{0,1,2\}$
- •
MLP_hid: number of hidden units in all MLPs used. This is also the number of slices used by Wasserstein aggregator. $\{128,256,512\}$
- •
MLP2_layers: using or not an additional linear projection after the residual connection of each stage $\{0,1\}$
- •
Cls_layers: number of layers in the final classifier MLP $\{1,2\}$
- •
Cls_hid: number of hidden units in the final classifier MLP $\{96,128,256\}$
- •
self_loops: using or not self loops $\{\text{True},\text{False}\}$
- •
dropout: dropout used inside the model $\{0.5,0.6,0.7\}$
- •
in_dropout: dropout used in the begining of the model $\{0.2,0.5,0.6,0.7\}$
- •
fixed hyperparameters: All models use 1 layer of WHNN, LayerNorm normalisation, the residual coefficient $\alpha$ fixed to $0.5$ and they are trained for $500$ epochs with a learning rate of $0.001$.

**Table 3: The best configuration of hyperparameters used by our model WHNN_MLP on all tested datasets. We mark with bold the parameters that are specific to the Wasserstein aggregator.**
| Parameter | Cora | Citeseer | Cora_CA | DBLP_CA | ModelNet40 | NTU2012 | 20News |
| --- | --- | --- | --- | --- | --- | --- | --- |
| num_ref | 25 | 10 | 25 | 5 | 50 | 25 | 25 |
| learnable_W | True | False | True | True | False | False | False |
| MLP_layers | 1 | 2 | 2 | 2 | 1 | 1 | 0 |
| MLP2_layers | 0 | 0 | 1 | 0 | 0 | 1 | 0 |
| MLP_hid | 128 | 256 | 256 | 512 | 256 | 512 | 512 |
| Cls_layers | 1 | 1 | 1 | 2 | 2 | 2 | 2 |
| Cls_hid | 256 | 128 | 96 | 96 | 96 | 96 | 96 |
| self_loops | True | True | True | True | True | False | False |
| dropout | 0.7 | 0.5 | 0.6 | 0.7 | 0.5 | 0.5 | 0.5 |
| in_dropout | 0.7 | 0.5 | 0.6 | 0.7 | 0.2 | 0.2 | 0.2 |

**Table 4: The best configuration of hyperparameters used by our model WHNN_SAB on all tested datasets. We mark with bold the parameters that are specific to the Wasserstein aggregator and with italic the parameters that are specific to the SAB encoder.**
| Parameter | Cora | Citeseer | Cora_CA | DBLP_CA | ModelNet40 | NTU2012 | 20News |
| --- | --- | --- | --- | --- | --- | --- | --- |
| num_ref | 10 | 5 | 50 | 5 | 25 | 25 | 5 |
| learnable_W | True | False | False | False | False | False | True |
| heads | 2 | 4 | 1 | 4 | 1 | 2 | 2 |
| MLP_layers | 2 | 2 | 2 | 1 | 1 | 2 | 2 |
| MLP2_layers | 0 | 0 | 1 | 1 | 0 | 0 | 0 |
| MLP_hid | 128 | 256 | 128 | 256 | 256 | 512 | 512 |
| Cls_layers | 1 | 1 | 1 | 2 | 2 | 2 | 2 |
| Cls_hid | 128 | 256 | 128 | 96 | 96 | 96 | 96 |
| self_loops | True | False | True | True | True | True | False |
| dropout | 0.7 | 0.7 | 0.5 | 0.7 | 0.5 | 0.5 | 0.5 |
| in_dropout | 0.7 | 0.7 | 0.5 | 0.7 | 0.2 | 0.2 | 0.2 |

## Appendix D Computational complexity

We derive the computational complexity for both versions of our Wasserstein Hypergraph Neural Network: using the edge-independent encoder (WHNN_MPN) and using the edge-dependent encoder (WHNN_SAB). We present the complexity for a hypergraph with $N$ nodes, $M$ hyperedges, $K_{e}$ the maximum cardinality of a hyperedge, $K_{v}$ the maximum number of hyperedges a node is part of and $R$ the number of reference points sampled.

For node encoders, the MLP encoder has a complexity of $O(N)$ while the SAB encoder has complexity $O(M\times K^{2})$ due to the pairwise exchange of messages ($K^{2}$) inside each hyperedge ($M$).

For the Wasserstein aggregator, we derive the complexity both for the node-to-hyperedge and hyperedge-to-node stages. For node-to-hyperedge the complexity for interpolation is $O(M\times(R\log K_{e}))$ and the complexity for sorting each neighbourhood is $O(M\times(R\log R))$. Symmetrically, for hyperedge-to-node the complexity for interpolation is $O(N\times(R\log K_{e}))$ and the complexity for sorting each neighbourhood is $O(N\times(R\log R)$. The overall complexity becomes $O(M\times(R\log K_{e})+M\times(R\log R))+N\times(R\log K_{e})+N\times(R\log R))$. In our experiments R is maximum $50$.

For comparison, the complexity of a Deep Set pooling is $O(M\times K_{e}+N\times K_{v})$

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 4
presentation: 3
questions: * Are the nodes of the hyper-edges ordered (i.e. order of nodes matter, as a generalization of directed graphs)? The math does not seem ordered. But what if the data is ordered? How can one extend the model to handle this data? E.g., (<Student> studied <program> at <university>)
* Please answer the questions in weaknesses, importantly, by updating the paper.
soundness: 3
strengths: * Paper is well-written. I learned a bunch, as perhaps apparent from the above summary. It provides both good mathematical understanding (step-by-step walkthrough), as well as demonstration figures in main paper and Appendix.
* The paper utilizes different grounds-up concepts, which is "refreshing" in LLM era.
* Good experimental results
summary: Paper is interested in modeling neural networks for Hypergraphs: a generalization of graphs where any edge can connect zero-or-more nodes. Prior neural networks on hypergraphs follow message-passing paradigm, where each edge updates itself as an *average* of (a transformation of) its incident (endpoint) node features, then each node updates itself with an *average* of its incident hyperedges. This *average* (or sum) pooling captures only a simple statistic: the average. The paper pictorially shows why average is insufficient for representing a node. Suppose a graduate course with student average of 80, where most students receive a grade around the average. Suppose course with bimodal grades (centered at 90, and at 70). While both courses have a similar average, i.e. modeled the same with most (earlier) models, they should be represented differently. Paper rather tries to represent nodes "using distribution of edges" (and vice-versa: represent edges using distribution of nodes). Distance between 2 distributions is quantified using Wasserstein (W) distance. While W is intensive to compute for distributions over multiple dimensions, it is much easier to compute it for probability distributions over one (scalar) variable. (as sum of distances of inverse of CDF), allowing one to approximate the distance between 2 distributions (over higher dimension) by taking many random direction lines, and projecting the distributions upon these lines i.e. break-down the problem into many 1D W-distances, and average them, to approximate the distance in higher dimensions. Further, rather than calculating the approximate distance (even in the 1D case), paper mentions that they can project onto a latent space where euclidean-distances are trained to approximate W-distances.

Sorry for the long description :)

All-in-all -- I like the paper a lot. However, there are two main things that need to be addressed for acceptance consideration (IMO):
1- Clarity around the math / algorithm (see Weaknesses).
2- Novelty. There are Wasserstein GNNs (https://arxiv.org/abs/2102.03450)
weaknesses: ## main weakness
The main weakness is in the clarity of the method. Not in terms of the write-up/motivation/high-level description (I think those are clear), but in terms of actually implementing or replicating this method. The math/algorithm section has loose ends, and personally, while I appreciate this work, I am unable to implement it from reading this paper alone. Great papers, on the other hand, I am able to implement them by reading only those papers. Let's try to improve this paper from "Good" to "Great" in order to have it appropriate for ICLR.

Specifically,

* what is the "interpolation" mentioned in Step 2 (Page 7) or line 12 in Alg.2. It seems that it will discard entries if there are too many, or invent new entries if there are only a few of them. I would guess that inventing could imply taking average of existing points?

* Line 5 of Alg2: Does it expand the dimensionality of X?
* Line 7 of Alg2: what are the vectors (reference points) sorted by? by their CDF value?
* Line8 onwards of Alg2: what is the lowercase symbol $s$? Of course, upper-case $S$ is the set of neighbors -- does lower-case $s$ correspond to the center (node or edge) for which its neighbros are recorded in upper-case $S$? If so, perhaps consider looping $(s, S) \in \mathcal{N}.items()$ -- or something similar. We are lucky that most ML folks are able to read pythonic-pseudocode.
* Line 18 of Alg2: is $W$ vector of scalars? In general, **please add sizes of all tensors**. Does this line sum-over a dimension? (weighted sum, with weights present in vector $W$)
* Alg1 line 7 or Step 1 on page 7 -- what is the reference distribution $q$? does it contain some anchor (randomly-sampled) nodes/edges? Do they have to go through the encoder (Line 11)?


## Novelty

Have you seen: https://arxiv.org/pdf/2102.03450 ? In what way does your method differ? Is it only in the application (Hypergraphs VS graphs)? or is the entire construction different?


## minor weakness:
* Typo: Perhaps replace "inner integral" with "integral" on line 201? I get that this becomes "inner" once plugged-into Eq.2 but perhaps the text is confusing as word "inner" shows up before Eq.2

---

confidence: 3
contribution: 1
initial_rating: 2
presentation: 2
questions: See weakness
soundness: 3
strengths: The key strengths are two-fold:
- The key idea is simple, intuitive, and solid. Existing set pooling functions may struggle to capture the full geometry of set-structured inputs, and using a more expressive pooling function to improve HNNs is a natural approach. 
- The approach is also somewhat original to hypergraph learning. Intuitively, node feature distribution within each hyperedge could be distinct and informative for hypergraph learning.
summary: This paper introduces Wasserstein Hypergraph Neural Network (HNN). Whereas most existing HNNs use sum-based pooling to smooth information during message passing, the proposed method uses Wasserstein aggregation to preserve geometric information of node feature distribution within each hyperedge. The proposed method outperforms the baseline HNNs in several node classification benchmark datasets.
weaknesses: I have five primary concerns:
- [limited novelty]: The key technical innovation is adapting Sliced Wasserstein Pooling (SWP), an existing module published in 2021, for HNNs. While the adaptation makes sense and seems to be performant, without a new technical innovation, the paper does not meet the standards of a top-tier conference.
- [thin content]: The experiments and analyses are thin. The authors used only 7 benchmark datasets with 8 (old) baseline HNNs. The method is evaluated only based on node classification performance. All the benchmark datasets have high label-homophily. No large dataset was tested. No formal or empirical analysis of the proposed method was provided. In fact, in a 9-page paper, the authors spent 3 pages solely on related work and background, which I do not think is generally acceptable.
- [weak method justification]: I do not agree with the authors' claim that the superior performance of the proposed method is due to 'capturing geometric information' of node features within each hyperedge. In most of the benchmark datasets used, the average hyperedge size is very small, e.g., 3-4 in Cora and Citeseer. Moreover, their input feature vectors are mostly sparse, multi-hot vectors. That is, I do not think such informative 'geometry' can be robustly identified in the benchmark datasets used. In fact, the authors did not provide any empirical analysis of the 'geometry' in the benchmark datasets. This drawback leaves the authors' key claim 'that this geometric information is highly relevant for hypergraph learning' (line 069) only speculative.
- [unfair hyperparameters]: The authors used self-loop addition as part of the tuned hyperparameters for each dataset. This is atypical and unfair. All the baseline methods, while they could, did not have the self-loop addition as their hyperparameter. This makes the empirical superiority of the proposed method suspicious.
- [copyright]: The baseline performance table is copied from the AllSet paper. This is okay, as long as the authors clearly mention that in the paper. However, I could not find it anywhere. This is a research ethics issue. Please revise properly. 

I have other, less significant concerns:
- [missing details]: Details about the synthetic case in Figure 1 are missing, making it unclear whether the case is illustrative or based on formal analysis. 
- [unclear writings]: Some writings are unclear. e.g., I do not understand 'This way, the hyperedges are not only characterised by the combination of their elements, but by the regions of the space where their elements are situated. The nodes became prototypes of the hyperedge behaviour.' Moreover, the authors write 'more sensitive to the geometric structure of the hyperedge' (line 228-229). However, if I understood correctly, they are not sensitive to the structure 'of hyperedge' but of 'node feature distribution within a hyperedge'. Please clarify.
- [unclear figures]: Some figures are quite confusing. E.g., in Fig. 2, what do node colors mean? What does the bar plot aim to depict? What do the contours aim to visualize? Adding legends, labels, and annotations may help.
- [minor errors]: Algorithm 1's line 4 is empty

Overall, in the current version, the paper is more like a workshop paper, with a simple, new, and promising idea and preliminary empirical outcomes.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: 1. This part feels vague in terms of its advantages—could it be presented in a more systematic way? 

“For example, a hypergraph containing two clusters of nodes suggests a bimodal underlying distribution. On the other hand, a hyperedge where nodes are close in the feature space denotes a unimodal probability distribution, suggesting a homophilic behaviour. A hyperedge in which nodes have similar representations indicates a low-variance distribution, while a hyperedge with diverse nodes suggests a more uniform distribution (see Figure 1).”

2. Why was optimal transport chosen to address this problem? Could other methods for measuring distances between distributions not work as well?

3. Please include an analysis of time complexity and a comparison of the algorithm’s runtime efficiency.

P.S. While I am generally favorable toward this paper, as I find the idea of modeling hyperedges as distributions very appealing, I would be more cautious and might consider a weak accept if the authors cannot demonstrate acceptable computational efficiency.
soundness: 4
strengths: 1. Although the paper addresses a traditional GNN problem, it takes a very interesting perspective. Instead of using conventional aggregators like mean or sum, the authors represent the nodes within a hyperedge as a distribution. This approach better captures the internal structure and higher-order relationships of hyperedges.
2. The writing of the paper is very logical and well-structured. Although it involves some theoretical concepts, the explanations are clear and easy to follow.
3. This paper presents an innovative application of optimal transport theory in the context of GNNs.
summary: This paper introduces the Wasserstein Hypergraph Neural Network, which models nodes and hyperedge neighborhoods as distributions and aggregates them using Sliced Wasserstein Pooling. Unlike traditional mean or sum aggregators, this approach preserves geometric properties of distributions, capturing richer higher-order relationships. Experiments show that it significantly improves node classification performance on real-world hypergraph datasets.
weaknesses: The main drawback of the paper is the lack of complexity analysis and runtime comparisons. Since the optimal transport algorithm is generally not cheap computationally, this raises some concerns about the efficiency of the proposed method.

---

confidence: 4
contribution: 1
initial_rating: 2
presentation: 2
questions: 1. Check the first point of Weaknesses, some theoretical proofs should be shown in the paper.
2. The Wasserstein aggregator appears significantly more computationally intensive than the simple summation used in baselines like ED-HNN. Could the authors please provide empirical performance metrics for WHNN? This would allow for a fair assessment of the accuracy-efficiency trade-off.
3. Given that the performance gains are sometimes modest (e.g., $\approx 0.4\%$ absolute improvement over ED-HNN on Cora), how do the authors justify the additional complexity of SWP for practical applications where inference speed or training cost are critical constraints?
4. The ablation studies show that both the baseline aggregators (Deep Set, PMA) and the proposed SWP aggregator benefit from a more powerful encoder. This seems to suggest that the initial motivation—that simpler aggregators uniquely "move the complexity... to the initial encoding"-may be an oversimplification. Could the authors please clarify or provide a more nuanced interpretation of these results? It appears a more general conclusion is that better encoders are complementary to better aggregators.
5. Wasserstein distance is extensively used in AI and machine learning. Why this WHNN adds new values?
soundness: 3
strengths: Principled Aggregator Choice: By adopting the distributional perspective, the choice of Sliced Wasserstein Pooling as an aggregator is well-justified. It endows the model with a powerful inductive bias to capture the shape, spread, and density of feature distributions within hyperedges, a capability that mean-based aggregators lack.   

Thorough Empirical Validation: The authors conduct a rigorous evaluation across ten datasets, comparing against eight strong baselines. The inclusion of comprehensive ablation studies effectively isolates the contribution of the Wasserstein aggregator, providing convincing evidence for the paper's central claims
summary: The paper introduces the Wasserstein Hypergraph Neural Network, an architecture designed to overcome a key limitation in existing hypergraph neural networks. WHNN proposes a conceptual shift by viewing these neighborhoods as empirical samples from underlying probability distributions. To capture their geometric properties, the model employs Sliced Wasserstein Pooling as its core aggregation mechanism. This design creates embeddings where the Euclidean distance approximates the Sliced Wasserstein distance between the original distributions, thus preserving structural information like shape, spread, and density. Implemented within a standard two-stage message-passing framework, the authors demonstrate that WHNN outperforms traditional aggregation methods and achieves state-of-the-art performance on multiple node classification benchmark datasets.
weaknesses: 1.	Theoretical Grounding not thoroughly discussed: the paper critiques sum-based aggregators like Deep Sets. Deep Sets is known to be a universal approximator for permutation-invariant functions defined on sets. The authors provide no such theoretical analysis for WHNN. They rely on the properties of SWP established in prior work, stating that their embedding approximates the SW distance, but they do not analyze what this approximation implies for the expressive power of the full HGNN model. This leaves several important theoretical questions unanswered. Does WHNN retain universal approximation capabilities? How does the error introduced by the SWP approximation propagate through multiple layers of message passing? Is the model guaranteed to distinguish between any two non-isomorphic hypergraph neighborhoods - is the aggregation function injective? The literature on SWP embeddings has explored properties like injectivity and bi-Lipschitz continuity, but the authors do not connect these foundational concepts to their hypergraph model. Without this theoretical grounding, the paper's contribution feels more like an engineering result rather than a fundamental advance in understanding hypergraph representation learning.

2.	Potentially Overstated Claims and Interpretation of Results: The paper claims to achieve "top performance" and "top results" across all datasets. While technically true based on the mean accuracy reported in Table 2, this language can be misleading when the margin of victory is small and falls within the standard deviation of a competing method. For example, on the NTU2012 dataset, WHNN-MLP scores 90.87 $\pm$ 1.59, which is statistically very close to ED-HNN's 89.48 $\pm$ 1.87. A more measured tone would be appropriate. More concerning is the interpretation of the ablation studies. The authors initially motivate SWP as a way to reduce the learning burden on the encoder. However, their own results in Figure 3 show that SWP's performance also significantly improves when paired with the more complex SAB encoder. This suggests the initial premise may be a false dichotomy. A more nuanced conclusion, better supported by the data, is that powerful encoders and expressive aggregators are complementary, and their combination is what yields state-of-the-art results.

3.	Omission of Computational Cost and Practicality Analysis: The paper's central claim revolves around achieving superior classification accuracy, but it entirely neglects to discuss the associated computational costs. The literature on SWP confirms its computational overhead, with a complexity of $O(LM \log M)$ for a set of size $M$ and $L$ slices, and often requires a large number of slices ($L$) to be effective. The paper's own complexity analysis in Appendix D confirms this super-linear dependency on neighborhood size.


Decision: Reject