P RUNING E DGES AND G RADIENTS TO L EARN H YPERGRAPHS FROM L ARGER S ETS
Anonymous authors
Paper under double-blind review
A BSTRACT
This paper aims for set-to-hypergraph prediction, where the goal is to infer the set of relations for a given set of entities.
This is a common abstraction for applications in particle physics, biological systems and combinatorial optimization.
We address two common scaling problems encountered in set-to-hypergraph tasks that limit the size of the input set: the exponentially growing number of hyperedges and the run-time complexity, both leading to higher memory requirements.
We make three contributions.
First, we propose to predict and supervise the positive edges only, which changes the asymptotic memory scaling from exponential to linear.
Second, we introduce a training method that encourages iterative reﬁnement of the predicted hypergraph, which allows us to skip iterations in the backward pass for improved efﬁciency and constant memory usage.
Third, we combine both contributions in a single set-to-hypergraph model that enables us to address problems with larger input set sizes.
We provide ablations for our main technical contributions and show that our model outperforms prior state-of-the-art, especially for larger sets.
I NTRODUCTION
Inferring the relational structure for a given set of entities is a common abstraction for many applications, including vertex reconstruction in particle physics (Shlomi et al., 2020a;
Serviansky et al., 2020), inferring higher-order interactions in biological and social systems (Brugere et al., 2018;
Battiston et al., 2020) or combinatorial optimization problems, such as ﬁnding the convex hull or
Delaunay triangulation (Vinyals et al., 2015; Serviansky et al., 2020).
The wide spectrum of different application areas underlines the expressivity of this abstract task, which is known in machine learning as set-to-hypergraph prediction (Serviansky et al., 2020).
Here, the hypergraph generalizes the pairwise relations in a graph to multi-way relations, a.k.a. hyperedges.
We distinguish this task from the related, but different, task of link prediction that aims to discover the missing edges in a graph, given the set of vertices and a subset of the edges (L ¨ u & Zhou, 2011).
For the set-to-hypergraph problem considered in this paper, we start with a set of nodes without any edges.
A common approach to set-to-hypergraph problems is to decide for every edge, whether it exists or not (Serviansky et al., 2020).
For a set of n nodes, the number of all possible hyperedges grows in O (2 n ) , which already becomes intractable for moderately sized n .
This is the scaling problem of set-to-hypergraph prediction that we will address in this paper.
Combinatorial optimization challenges, like set-to-hypergraph prediction, introduce the additional problem of complexity .
For example, convex hull ﬁnding in d dimensions has a run time complexity of O ( n log( n ) + n (cid:98) d 2 (cid:99) )
(Chazelle, 1993).
This means that larger input sets require more compute regardless of the quality of the hypergraph prediction algorithm.
Indeed, it was observed in (Serviansky et al., 2020) that for larger set sizes performance was worse.
In this paper we aim to address the scaling and complexity problems in order to predict hypergraphs from larger sets.
We make three contributions in this paper.
First, in Section 2 we improve the scalability of hypergraph representations for set-to-hypergraph tasks, by pruning the non-existing edges .
We prove that during training it sufﬁces to supervise the existing edges only, thus improving the asymptotic memory requirements from O (2 n ) to O ( mn ) , that is linear in the input set size.
Second, to address the complexity problem, we introduce in Section 3 a training method that encourages iterative reﬁnement of the predicted hypergraph with memory requirements scaling constant in the number of reﬁnement steps.
This addresses the need for more compute by complex problems in a scalable manner.
Third, we combine in Section 4 the efﬁcient representation from the ﬁrst contribution with the requirements of the scalable training method from the second contribution in a recurrent model that performs iterative reﬁnement on a pruned hypergraph.
Our model handles different input set sizes and varying numbers of edges, while respecting the permutation symmetry of both.
In our experiments in Section 5, we provide an in-depth ablation on each of our technical contributions and compare our model against prior work on common set-to-hypergraph benchmarks. 1. P RELIMINARY
Hypergraphs generalize normal graphs by replacing the normal edges that connect exactly two nodes with hyperedges that connect an arbitrary number of nodes. Since we only consider the general version, we shorten hyperedges to edges in the remainder of the paper. In set-to-hypergraph tasks, we treat the input set as the set of nodes of a hypergraph and aim to learn a function f that predicts the corresponding set of edges. For example, the input set could consist of objects in an image and the edges would represent their relations.
Next, we provide an overview of the Set2Graph neural network (Serviansky et al., 2020).
We focus on it because most previous networks follow a similar structure.
In Set2Graph, f is split into a collection of set-to- k -edge functions F k , where k -edges connect exactly k nodes.
All F k are composed of three steps: a set-to-set model maps the input set to a latent set, a broadcasting step forms all possible k -tuples from the latent set elements, and a ﬁnal graph-to-graph model that predicts for each k -edge whether it exists or not.
Serviansky et al.
(2020) show that this can approximate any continuous set-to- k -edge function, and by extension the family of F k functions can approximate any continuous set-to-hypergraph function.
Since the asymptotic memory scaling of F k is in O ( n k ) , modeling k -edges beyond k > 2 already becomes intractable in many settings and one has to apply heuristics to recover higher-order edges from pairwise edges (Serviansky et al., 2020).
S CALING BY PRUNING THE NON - EXISTING EDGES
In this section, we propose a solution for the memory scaling problem encountered in set-tohypergraph tasks.
The goal is to learn a model f ( X ) =
H that maps a set X of input vectors to the hypergraph H .
The choice on how to represent the hypergraph H can already drastically impact the asymptotic complexity of f , as we saw for Set2Graph.
In what follows, we explain how we represent the nodes and edges of H to learn a pruned incidence matrix, and we will motivate why that is necessary for scaling to problems with more nodes.
Nodes.
Each input element x ∈ X gets an “identity” as a node in the hypergraph, meaning if a subset of the nodes are connected by an edge then there exists a relation between the corresponding input elements.
We differentiate between the input elements and the nodes of the hypergraph, as we expect the latter to be represented as latent vectors v ∈ V of d V dimensions.
The importance of this becomes clear in the discussion on the training objective later on.
Edges.
The set of all possible edges can be expressed using the power set P ( V ) \ { ∅ } , that is the set of all subsets of V minus the empty set.
Different from the situation with the nodes, we do not know which edge will be part of the hypergraph, since this is what we want to predict.
Listing out all 2 | V | − 1 possible edges and deciding for each edge whether it exists or not, becomes intractable very quickly.
Furthermore, we observe that in many cases the number of existing edges is much smaller than the total number of possible edges.
We leverage this observation by keeping only a ﬁxed number of edges m that is sufﬁcient to cover all (or most) hypergraphs for a given task.
Thus, we improve the memory requirement from O (2 | V | d E ) to O ( md E ) , where d E is the vector size of the edge representations e ∈ E .
All possible edges that do not have an edge representation in E are implicitly predicted to not exist.
After specifying the training objective, we provide a more formal argument on why pruning all but m edges from E is sound.
Incidence matrix.
Since both the nodes and edges are represented by latent vectors, we require an additional component for specifying the connections.
Different from previous approaches, we use the incidence matrix over adjacency tensors (Serviansky et al., 2020; Ouvrard et al., 2017).
The two differ in that incidence describes whether an edge is connected to a node, while adjacency describes whether an edge between a subset of nodes exists.
The rows of the incidence matrix I correspond to the edges and the columns to the nodes.
Thus, an entry
I i,j ∈ [0 , 1] represents the probability of the i -th edge being incident to the j -th node.
Theoretically we can express any hypergraph in both representations, but pruning edges becomes especially simple in the incidence matrix, where we just remove the corresponding rows.
We interpret the pruned edges e / ∈ E that have no corresponding row in the pruned I as having zero probability of being incident to any node.
Loss function.
We apply the loss function only on the incidence probability matrix I . For efﬁciency purposes, we would like to train each incidence value separately as a binary classiﬁcation and apply a constant threshold ( > 0 . 5 ) on I at inference time to get the discrete incidence matrix I (cid:48) .
In probabilistic terms, this translates to a factorization of the joint distribution p ( I | X ) as, (cid:81) i,j p ( I i,j | X ) .
In order to still be able to model interrelations between different incidence probabilities, we impose a requirement on the model f : the probability I i,j produced by f depends on e i and v j .
This highlights the importance of the latent node and edge representations, which enable us to model the dependencies in the output while still being efﬁcient at training and inference time.
Furthermore, this changes our assumption on the incidence probabilities from that of independence to conditional independence on e i and v j , and we apply the binary cross-entropy loss on each I i,j .
The binary classiﬁcation over I i,j highlights yet another reason for picking the incidence representation over the adjacency.
When we prune all non-existing edges, learning a binary classiﬁer in the adjacency case would no longer work due to the lack of negative examples.
In contrast, an existing edge in the incidence matrix contains both ones and zeros (except for the edge connecting all nodes), ensuring that a binary incidence classiﬁer sees both positive and negative examples.
However, an adjacency tensor has the advantage that the order of the entries is fully decided by the order of the nodes, which are given by the input X in our case.
In the incidence matrix, the row order of the incidence matrix depends on the edges, which are orderless.
When comparing the predicted incidence matrix with a ground-truth matrix, we need to account for the orderless nature of the edges and the given order of the nodes.
Hence, we require a loss function that is invariant towards reordering over the rows of the incidence matrix, but equivariant to reordering over the columns.
We achieve this by matching every row in I with a row in the pruned ground-truth incidence matrix I ∗ (containing the existing edges), such that the binary cross-entropy loss H over all entries is minimal:
[Equation]
Finding a permutation π that minimizes the total loss is known as the linear assignment problem and can be solved with the Hungarian matching algorithm (Kuhn, 1955).
Having established the training objective in Equation 1, we can now offer a more formal argument on the soundness of supervising existing edges only while pruning the non-existing ones , where J can be understood as the full incidence matrix (proof in Appendix A):
Proposition 1 (Supervising only existing edges) .
Let J ∈ [ (cid:15), 1) (2 n − 1) × n be a matrix with at most m rows for which ∃ j : J ij > (cid:15) , with a small (cid:15) > 0 . Similarly, let J ∗ ∈ { 0 , 1 } (2 n − 1) × n be a matrix with at most m rows for which ∃ j : J ij = 1 .
Let prune ( · ) denote the function that maps from a (2 n − 1) × n matrix to a k × n matrix, by removing (2 n − 1) − k rows where all values are ≤ (cid:15) . Then, for a constant c = (2 n − 1 − k ) n · H ( (cid:15), 0) and all such J and J ∗ :
[Equation]
The matrix prune ( J ) can be understood as the pruned incidence matrix that we deﬁned earlier and prune ( J ∗ ) as the pruned ground-truth.
In practice, the (cid:15) corresponds to a lower bound on the log in the entropy computation, like -100 in PyTorch (Paszke et al., 2019).
Since the losses in Equationare equivalent up to an additive constant, the gradients of the parameters are exactly equal in both the pruned and non-pruned cases.
Thus, pruning the non-existing edges does not affect learning, while signiﬁcantly reducing the asymptotic memory requirements.
Summary.
In set-to-hypergraph tasks, the number of different edges that can be predicted grows exponentially with the input set size.
We address this computational limitation by representing the edge connections with the incidence matrix and pruning all non-existing edges before explicitly deciding for each edge whether it exists or not.
We show that pruning the edges is sound, when the loss function respects the permutation symmetry in the edges.
S CALING BY PRUNING THE NON - ESSENTIAL GRADIENTS
Next, we consider how to learn the pruned incidence matrix.
Some tasks may require more compute than others, which can result in worse performance or intractable models if not properly addressed.
A naive approach would increase the number of parameters, either by increasing the number of hidden dimensions or the depth of the neural network, which is clearly not scalable.
Furthermore, the memory requirement of backprop would also grow with greater depth.
Instead, we would like to increase the amount of sequential computation by reusing parameters.
That is, we want the model f to be recurrent, H t +1 = f
( X , H t ) .
Recurrent models are commonly applied to sequential data, where the input varies for each time step t (Lipton et al., 2015).
In our case, we use the same input X at every step.
Using a recurrent model, we can increase the total number of iterations – to scale the amount of sequential computation steps – without increasing the number of parameters.
However, the recurrency does not address the growing memory requirements of backprop, since the activations of each iteration still need to be kept in memory.
Iterative reﬁnement.
In the rest of this section, we present an efﬁcient training algorithm that can scale to any number of iterations at a constant memory cost.
We build on the idea that if each iteration applies a small reﬁnement, then it becomes unnecessary to supervise every iteration.
We can deﬁne an iterative reﬁnement as reducing the loss (by a little) after every iteration, L ( I t , I ∗ ) <
L ( I t − 1 , I ∗ ) .
Thus, the long-term dependencies between H t for t ’s that are far apart can also be ignored, since f only needs to improve the current H t .
We can encourage f to iteratively reﬁne the prediction H t , by applying the loss L after each iteration.
This has the effect that f learns to move the H t in the direction of the negative gradient of L , making it similar to a gradient descent update.
Input: X , I ∗ , S, B, N
H ← initialize( X ) for s in S : with no grad() : for t in range( s ) :
H ← f ( X , H ) l ←for t in range( B ) :
H ← f ( X , H ) l ← l + L ( H , I ∗ ) backward( l ) gradient step and reset()
Backprop with skips.
Similar to previous works that encourage iterative reﬁnement through (indirect) supervision on the intermediate steps (Jastrzebski et al., 2018), we expect the changes of each step to be small.
Thus, it stands to reason that supervising every step is unnecessary and redundant.
This leads us to a more efﬁcient training algorithm that skips iterations in the backward-pass of backprop.
Algorithm 1 describes the training procedure in pseudocode (in syntax similar to PyTorch (Paszke et al., 2019)).
We perform a total of N gradient updates per mini-batch.
Each gradient update consist of two phases, ﬁrst s iterations without gradient accumulation and then B iterations that add up the losses for backprop.
Through these hyperparameters we control the amount of resources used during training.
Increasing hyperparameter B allows for models that do not strictly decrease the loss after every step and require supervision over multiple subsequent steps.
Note that having the input X at every reﬁnement step is important so that the model does not forget the initial problem.
Summary.
Motivated by the need for more compute to address complex problems, we propose a method that increases the amount of sequential compute of the neural network without increasing the memory requirement at training time.
Our training algorithm requires the model f to perform iterative reﬁning of the hypergraph, for which we present a method in the next section.
S CALING THE SET - TO - HYPERGRAPH PREDICTION MODEL
In Section 2 and Section 3 we proposed two methods to overcome the memory scaling problems that appear for set-to-hypergraph tasks.
To put these methods into practice, we need to specify a model f that fulﬁlls the required properties.
In what follows, we propose a speciﬁc implementation for each such property.
Initialization.
As the starting point for the iterative reﬁnement, we initialize the nodes V 0 from the input set as v 0 i
=
W x i + b , where W ∈ R d V × d X , b ∈ R d V are learnable parameters.
The afﬁne map allows for hidden dimensions d V that are different from the input feature dimensions d X .
An informed initialization similar to the nodes is not available for the edges and the incidence matrix, since these are what we aim to predict.
Instead, we choose an initialization scheme that respects the permutation symmetry of a set of edges while also ensuring that each edge starts out differently.
The last point is necessary for permutation-equivariant operations to distinguish between different edges.
The random initialization e 0 i ∼ N ( µ , diag ( σ )) , with shared learnable parameters µ ∈ R d E and σ ∈ R d E fulﬁlls both these properties, as it is highly unlikely for two samples to be identical.
Conditional independence.
We want the incidence probabilities
I i,j to be conditionally independent of each other given e i and v j .
A straightforward way to model this is by concatenating both vectors (denoted with [ · ] ) and applying an MLP with a sigmoid activation on the scalar output:
[Equation]
The superscripts point out that we produce a new incidence matrix for step t based on the edge and node vectors from the previous step.
Note that we did not specify an initialization for the incidence matrix, since we directly replace it in the ﬁrst step.
Iterative reﬁnement.
The training algorithm in Section 3 assumes that f performs iterative reﬁnement on H t , but leaves open the question on how to design such a model.
Instead of iteratively reﬁning the incidence matrix, i.e., the only term that appears in the loss (Equation 1), we focus on reﬁning the edges and nodes.
A reﬁnement step for some node v i ∈ V needs to account for the rest of the hypergraph, which also changes with each iteration.
For this purpose we apply the permutation-equivariant DeepSets (Zaheer et al., 2017) to produce node updates dependent on the full set of nodes from the previous iteration V t − 1 .
The permutation-equivariance of DeepSets means that the output set retains the input order; thus it is sensible to refer to v ti as the updated version of the same node v t − 1 i from the previous step.
Furthermore, we concatenate the aggregated neighboring edges weighted by the incidence probabilities
ρ E → V ( j, t ) = (cid:80) ki =
I ti,j e t − 1 i , to incorporate the relational structure between the nodes.
This aggregation works akin to message passing in graph neural networks (Gilmer et al., 2017).
An indispensable input, required for adding skips in the backward pass during training, are the input features X .
Instead of directly concatenating the raw features x i , we use the initial nodes v 0 i .
Finally, we express the reﬁnement part for the nodes as:
[Equation]
The updates to the edges E t mirror that of the nodes, except for the injection of the input set.
Together with the aggregation function ρ V → E ( i, t ) =
{+(cid:80) nj =
+}
I ti,j v t − 1 j , we can update the edges as:
[Equation]
By sharing the parameters between different reﬁnement steps, we naturally obtain a recurrent model.
Previous works on recurrent models (Locatello et al., 2020) saw improvements in training convergence by including layer normalization (Ba et al., 2016) between each iteration.
Shortcut connections in ResNets (He et al., 2016) have been shown to encourage iterative reﬁnement of the latent vector (Jastrzebski et al., 2018).
We add both shortcut connections and layer normalization to the updates in Equation 4 and Equation 5.
Although we prune the negative edges, we still want to predict a variable number thereof.
To achieve that we simply extend the incidence model in Equationwith an existence indicator:
[Equation]
This can be seen as factorizing the probability into “ p ( e i incident to v j ) · p ( e i exists ) ” and replaces the aggregation weights in ρ E → V and ρ V → E . [Table]
Summary.
We propose a model that fulﬁlls the requirements of our scalable set-to-hypergraph training framework from Section 2 and Section 3.
By adding shortcut connections [Table] we encourage it to perform iterative reﬁnements on the hypergraph while being permutation equivariant with respect to both the nodes and the edges [Table]
E XPERIMENTS
In this section [Table] we evaluate our approach on multiple set-to-hypergraph tasks, in order to compare to prior work and examine the main design choices.
We refer to Appendix B for further details, results and ablations.
Code is included in the supplementary material.
5. S CALING S ET - TO -H YPERGRAPH P REDICTION
First, we compare our model from Section 4 on three different set-to-hypergraph tasks against the state-of-the-art model.
This allows us to see the difference between predicting the incidence matrix and predicting the adjacency tensors.
Baselines.
Our main comparison is against Set2Graph (Serviansky et al., 2020), which is a strong and representative baseline for approaches that predict the adjacency structure, which we generally refer to as adjacency-based approaches.
Serviansky et al.
(2020) modify the task in two of the benchmarks, to avoid storing an intractably large adjacency tensor.
We explain in Appendix B how this affects the comparison.
Additionally, we compare to Set Transformer (Lee et al., 2019) and Slot
Attention (Locatello et al., 2020), which we adapt to the set-to-hypergraph setting by treating the output as the pruned set of edges and producing the incidence matrix with the MLP from Equation 3.
We refer to these two as incidence-based approaches that also include our model.
Particle partitioning.
Particle colliders are an important tool for studying the fundamental particles of nature and their interactions.
During a collision, several particles are emanated and measured by nearby detectors, while some particles decay beforehand.
Identifying which measured particles share a common progenitor is an important subtask in the context of vertex reconstruction (Shlomi et al., 2020b).
We can treat this as a set-to-hypergraph task: the set of measured particles is the input set and the common progenitors are the edges of the hypergraph.
We use a simulated dataset of 0 .
9 M datasample with the default train/validation/test split (Serviansky et al., 2020; Shlomi et al., 2020b).
Each data-sample is generated from on one of three different distributions for which we report the results separately: bottom jets , charm jets and light jets .
The ground-truth target is the incidence matrix that can also be interpreted as a partitioning of the input elements, since every particle has exactly one progenitor (edge).
In Table 1 we report the performances on each type of jets as the F1 score and Adjusted Rand Index (ARI).
Our method outperforms all alternatives on bottom and charm jets, while being competitive on light jets.
Convex hull.
The convex hull of a ﬁnite set of d -dimensional points can be efﬁciently represented as the set of simplices that enclose all points.
In the 3 D case, each simplex consists of 3 points that together form a triangle.
For the general d -dimensional case, the valid incidence matrices are limited to those with d incident vertices per edge.
Finding the convex hull is an important and wellunderstood task in computational geometry, with optimal exact solutions (Chazelle, 1993; Preparata [Table] [Table] Shamos, 2012).
Nonetheless, predicting the convex hull for a given set of points poses a challenging problem for current machine learning methods, especially when the number of points increases (Vinyals et al., 2015; Serviansky et al., 2020).
We generate an input set by drawing n 3-dimensional vectors from one of two distributions: Gaussian or spherical.
For the Gaussian setting, points are sampled i.i.d.
from a standard normal distribution.
For the spherical setting, we additionally normalize each point to lie on the unit sphere.
Following Serviansky et al.
(2020), we use n =30 , n = and n ∈ [20 . .
100] , where in the latter case the input set size varies between 20 and 100.
Table 2 shows our results.
Our method consistently outperforms all the baselines by a considerable margin.
In contrast to Set2Graph, our model does not suffer from a drastic performance decline when increasing the input set size from 30 to 50.
Furthermore, based on the results in the Gaussian setting, we also observe that all the incidence-based approaches handle varying input sizes much better than the adjacency-based approach.
Delaunay triangulation.
A Delaunay triangulation of a ﬁnite set of 2D points is a set of triangles for which the circumcircles of all triangles have no point lying inside.
When there exists more than one such set, Delaunay triangulation aims to maximize the minimum angle of all triangles.
We consider the same learning task and setup as Serviansky et al.
(2020), who frame Delaunay triangulation as mapping from a set of 2D points to the set of Delaunay edges, represented by the adjacency matrix.
Since this is essentially a set-to- graph problem instead of set-to-hypergraph one, we adapt our method for efﬁciency reasons, as we describe in Appendix B.
We generate the input sets of size n , by sampling 2-dimensional vectors uniformly from the unit square, with n =50 or n ∈ [20 . . 80] .
In Table 3, we report the results for Set2Graph (Serviansky et al., 2020) and our adapted method.
Our method again outperforms Set2Graph on all metrics.
Summary.
By predicting the positive edges only, we can signiﬁcantly reduce the amount of required memory for set-to-hypergraph tasks.
On three different benchmarks, we observe performance improvements when using this incidence-based approach, compared to the adjacency-based baseline.
Interestingly, our method does not see a large discrepancy in performance between different input set sizes, both in convex hull ﬁnding and Delaunay triangulation.
We attribute this to the recurrence of our iterative reﬁnement scheme, which we look into next.
5. A BLATIONS
Effects of increasing (time) complexity.
The intrinsic complexity of ﬁnding a convex hull for a d -dimensional set of n points is in O ( n log( n ) + n (cid:98)
d 2 (cid:99) )
(Chazelle, 1993).
This scaling behavior offers an interesting opportunity to study the effects of increasing (time) complexity on model performance.
The time complexity implies that any algorithm for convex hull ﬁnding scales superlinearly with the input set size.
Since our learned model is not considered an algorithm that (exactly) solves the convex hull problem, the implications become less clear.
In order to assess the relevancy of the problem’s complexity for our approach, we examine the relation between the number of reﬁning steps and increases in the intrinsic resource requirement.
The following experiments are all performed with standard backprop, in order to not introduce additional hyperparameters that may affect the conclusions.
[Figure]
First, we examine the performance of our approach with 3 iterations, trained on increasing set sizes n ∈ [10 . .
In Figurewe observe a monotone drop in performance when training with the same number of iterations.
The negative correlation between the set size and the performance conﬁrms a relationship between the computational complexity and the difﬁculty of the learning task.
Next, we examine the performance for varying number of iterations and set sizes.
We refer to the setting, where the number of iterations is 3 and set size n =10 , as the base case.
All other set sizes and number of iterations are chosen such that the performance matches the base case as closely as possible.
In Figure 1, we observe that the required number of iterations increases with the input set size, further highlighting that an increase in the number of iterations actually sufﬁces in counteracting the performance decrease.
Furthermore, we observe that the number of reﬁnement steps scales sub-linearly with the set size, different from what we would expect based on the complexity of the problem.
We speculate this is due to the parallelization of our edge ﬁnding process, differing from incremental approaches that produce one edge at a time.
[Figure]
the same memory, except standard backprop T ∈{ 16 , 32 } , which require more. Efﬁciency of backprop with skips. To assess the efﬁciency of backprop with skips, we compare to truncated backpropagation through time (TBPTT)
(Williams & Peng, 1990).
We consider two variants of our training algorithm: 1. Skipping iterations at ﬁxed time steps and 2. Skipping randomly sampled time steps.
In both the ﬁxed and random skips versions, we skip half of the total iterations.
We train all models on convex hull ﬁnding in 3 -dimensions for 30 spherically distributed points.
In addition, we include baselines trained with standard backprop that contingently inform us about performance degradation incurred by our method or TBPTT.
Standard backprop increases the memory footprint linearly with the number of iterations T , inevitably exceeding the available memory at some threshold.
Hence, we deliberately choose a small set size in order to accommodate training with backprop for T ∈{ 4 , 16 , 32 } number of iterations.
We illustrate the differences between standard backprop,
TBPTT and our backprop with skips in Figure 4 in the Appendix.
The results in Figure 2 demonstrate that skipping half of the iterations in the backward-pass, signiﬁcantly decreases the training time without hurting predictive performance.
When the memory budget is constricted to 4 iterations in the backward-pass, both TBPTT and backprop with skips outperform standard backprop considerably.
Recurrent vs. stacked.
Recurrence plays a crucial role in enabling more computation without an increase in the number of parameters.
By training the recurrent model using backprop with skips, we can further reduce the memory cost during training to a constant amount.
Since our proposed training algorithm from Section 3 encourages iterative reﬁnement akin to gradient descent, it is natural to believe that the weight-tying aspect of recurrence is a good inductive bias for modelling this.
A reason for thinking so, is that the “gradient” should be the same for the same I , no matter at which iteration it is computed.
Hence, we compare the recurrent model against a non-weight-tied (stacked) version that applies different parameters at each iteration.
First, we compare the models trained for 3 to 9 reﬁnement steps.
In Figure 3a, we show that both cases beneﬁt from increasing the iterations.
Adding more iterations beyond 6 only slightly improves the performance of the stacked
[Figure]
model, while the recurrent version still beneﬁts, leading to an absolute difference of 0.03 in F score for 9 iterations.
Next, we train both versions with 15 iterations until they achieve a similar validation performance, by stopping training early on the recurrent model.
The results in Figure 3b show that the recurrent variant performs better when tested at larger set sizes than trained, indicating an improved generalization ability.
R ELATED W ORK
Adjacency prediction.
Set2Graph (Serviansky et al., 2020) is a family of maximally expressive permutation equivariant neural networks that map from an input set to (hyper)graphs.
They show that their method outperforms many popular alternatives, including Siamese networks (Zagoruyko & Komodakis, 2015) and graph neural networks (Morris et al., 2019) applied to a k -NN induced graph.
(Serviansky et al., 2020) extend the general idea, of applying a scalar-valued adjacency indicator function on all pairs of nodes (Kipf & Welling, 2016), to the l -edge case (edges that connect l nodes).
In Set2Graph, for each l the adjacency structure is modeled by an l -tensor, requiring memory in O ( n l ) .
This becomes intractable already for small l and moderate set sizes.
By pruning the negative edges, our approach scales in O ( nk ) , making it applicable even when l = n .
Set prediction.
Recent works on set prediction map a learned initial set (Zhang et al., 2019; Lee et al., 2019) or a randomly initialized set (Locatello et al., 2020; Carion et al., 2020; Zhang et al., 2021) to the target space.
Out of these, the closest one to our hypergraph reﬁning approach is Slot
Attention (Locatello et al., 2020), which recurrently applies the Sinkhorn operator (Adams & Zemel, 2011) in order to associate each element in the input set with a single slot (hyperedge).
None of the prior works on set prediction consider the set-to-hypergraph task, but some can be naturally extended by mapping the input set to the set of positive edges, an approach similar to ours.
C ONCLUSION AND FUTURE WORK
By representing and supervising the set of positive edges only, we substantially improve the asymptotic scaling and enable learning tasks with higher-order edges.
On common benchmarks, we have demonstrated that our method outperforms previous works, while offering a more favorable asymptotic scaling behavior.
In further evaluations, we have highlighted the importance of recurrence for addressing the intrinsic complexity of problems.
Efﬁcient set loss.
We identify the Hungarian matching (Kuhn, 1955) as the main computational bottleneck during training.
Replacing the Hungarian matched loss with a faster alternative, like a learned energy function (Zhang et al., 2021), would greatly speed up training for tasks with a large maximum number of edges.
Larger input dimensions.
Our empirical analysis is limited to datasets on with low dimensional inputs.
Learning on higher dimensional input data might require extensions to the model that can make larger changes to the latent hypergraph than is feasible with small iterative reﬁnement steps.
The idea here is similar to the observation from Jastrzebski et al.
(2018) for ResNets (He et al., 2016) that also encourage iterative reﬁnement: earlier residual blocks apply large changes to the intermediate features while later layers perform (small) iterative reﬁnements.
E THICS S TATEMENT
We propose methods to deal with the scaling issues encountered in set-to-hypergraph tasks and evaluate it largely on synthetic datasets.
Our contributions enable larger input set sizes due to better asymptotic memory scaling.
This may enable applications on data with many nodes and sparse connections, including amongst others social networks (e g., clique prediction for a new group of users) or scientiﬁc data (e g., vertex reconstruction for more particles).
In the context of applying machine learning methods in social networks, it is good to be mindful about biases relating to human ethnicity, gender or groups that experience frequent discrimination (Torralba & Efros, 2011; Mehrabi et al., 2021).
R EPRODUCIBILITY S TATEMENT
The proof is included in Appendix A. Additional results and experiment descriptions are in Appendix B.
We provide the code for reproducing the experiments in the supplementary material.
A P ROOF : S UPERVISING POSITIVE EDGES ONLY SUFFICES
Proposition 1 (Supervising only existing edges) .
Let J ∈ [ (cid:15), 1) (2 n − 1) × n be a matrix with at most m rows for which ∃ j : J ij > (cid:15) , with a small (cid:15) > 0 . Similarly, let J ∗ ∈ { 0 , 1 } (2 n − 1) × n be a matrix with at most m rows for which ∃ j : J ij = 1 .
Let prune ( · ) denote the function that maps from a (2 n − 1) × n matrix to a k × n matrix, by removing (2 n − 1) − k rows where all values are ≤ (cid:15) . Then, for a constant c = (2 n − 1 − k ) n · H ( (cid:15), 0) and all such J and J ∗ :
[Equation]
Proof.
We shorten the notation with I = prune ( J ) and I ∗ =
prune ( J ) ∗ , making the relation to the incidence matrix I deﬁned in Section 2 explicit. Since L is invariant to permutations over the rows of its input matrices, we can assume w.l.o.g. that the not-pruned rows are the ﬁrst k rows, J : k
= I and J ∗ : k = I ∗ .
For improved readability, let ˆ H
( J π ( i ) , J ∗ i ) =
(cid:80) j H ( J π ( i ) ,j , J ∗ i,j ) denote the element-wise binary cross-entropy, thus the loss in Equation 1 can be rewritten as L ( J , J ∗ )=
min π ∈ Π (cid:80) i ˆ H
( J π ( i ) , J ∗ i ) . First, we show that there exists an optimal assignment between J , J ∗ that assigns the ﬁrst k rows equally to an optimal assignment between I , I ∗ . More formally, for an optimal assignment π I ∈ arg min π ∈ Π (cid:80)
i ˆ H
( I π ( i ) , I ∗ i ) we show that there exists an optimal assignment π J ∈ arg min π ∈ Π (cid:80)
i ˆ H
( J π ( i ) , J ∗ i ) such that ∀ 1 ≤ i ≤ k : π J ( i )=
π I ( i ) . If π J ( i ) ≤ k for all 1 ≤ i ≤ k then the assignment for the ﬁrst k rows is also optimal for I , I ∗ .
So we only need to show that there exists a π J such that π J ( i ) ≤ k for all 1 ≤ i ≤ k .
Let π J be an optimal assignment that maps an i<k to π J >k . Since π J is a bijection, there also exists a j<k that π − 1 J ( j ) >k assigns to.
The corresponding loss terms are lower bounded as follows:
[Equation]
Equality of Equation 8 holds since all rows with index >k are (cid:15) -vectors in J and zero-vectors in
J ∗ .
The inequality in Equation 11 holds since all values in J are lower bounded by (cid:15) . Thus, we have shown that either there exists no optimal assignment π J that maps from a value ≤ k to a value >k (which is the case when ˆ H ( J i , J ∗ π J ( i ) )
+ ˆ H
( J π − 1 J ( j ) > ˆ H ( J i , J ∗ j ) + ˆ H ( (cid:15) , 0 ) ) or that there exists an equally good assignment that does not mix between the rows below and above k .
Since the pruned rows are all identical, any assignment between these result in the same value (2 n − 1 − k )
ˆ H ( (cid:15) , 0 )=(2 n − 1 − k ) n · H ( (cid:15), 0) that only depends on the number of pruned rows 2 n − 1 − k and number of columns n .
B E XPERIMENTAL DETAILS
In this section, we provide further details on the experimental setup and additional results.
B. P ARTICLE PARTITIONING
For the particle partitioning problem, we follow the experimental setup of (Serviansky et al., 2020).
The problem considers the case where particles are collided at high energy, resulting in multiple particles shooting out from it.
Each example in the dataset consists of the input set, which corresponds to some outgoing particles, and the ground truth partition of the input set.
Each element in the partition is a subset of the input set and corresponds to some intermediate particle that was not measured, because it decayed into multiple particles before it could reach the sensors.
The learning task consists of inferring which elements in the input set originated from the same intermediate particle.
The particle partitioning task bears resemblance to the classical clustering setting.
It can be understood as a meta-learning clustering task, where both the number of clusters and the similarity function depend on the context.
Thought, clustering algorithms such as k -means, require a priori speciﬁcation of the number of clusters and use a ﬁxed similarity function and are therefore unsuitable.
For more information on how this task ﬁts into the area of particle physics more broadly, we refer to Shlomi et al.
(2020a).
Dataset.
We use the publicly available dataset of 0 .
9 M data-sample with the default train/validation/test split (Serviansky et al., 2020; Shlomi et al., 2020b). Each data-sample is The input sets consist of 2 to 14 particles, with each particle represented by 10 features. The target partitioning indicate the common progenitors and restrict the valid incidence matrices to those with a single incident edge per node.
Setup.
While Set2Graph is only one instance of an adjacency-based approach, (Serviansky et al., 2020) show that it outperforms many popular alternatives: Siamese networks, graph neural networks and a non-learnable geometric-based baseline.
All adjacency-based approaches incur a prohibitively large memory cost when predicting edges with high connectivity.
In the case of particle partitioning, Set2Graph resorts to only predicting edges with at most 2 connecting nodes, followed by an additional heuristic to infer the partitions (Serviansky et al., 2020).
In contrast to that, all the incidence-based approaches do not require the additional post-processing step at the end.
We simplify the hyperparameter search by choosing the same number of hidden dimensions d for the latent vector representations of both the nodes d V and the edges d E .
In all runs dedicated to searching d , we set the number of total iterations T =3 and backpropagate through all iterations.
We start with d =32 and double it, until an increase yields no substantial performance gains on the validation set, resulting in d =128 .
In our reported runs, we use T =16 total iterations, B =4 backprop iterations, N =2 gradient updates per mini-batch, and a maximum of 10 edges.
We apply the same d =128 to both the Slot Attention and Set Transformer baselines.
Similar to the original version (Locatello et al., 2020), we train Slot Attention with 3 iterations.
Attempts with more than 3 iterations resulted in frequent divergences in the training losses.
We attribute this behavior to the recurrent sinkhorn operation, that acts as a contraction map, forcing all slots to the same vector in the limit.
We train all models using the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 0 .for
400 epochs and retain the parameters corresponding to the lowest validation loss.
All models additionally minimize a soft F1 score (Serviansky et al., 2020).
Since each particle can only be part of a single partition, we choose the one with the highest incidence probability at test time.
Our model has 268162 trainable parameters, similar to 251906 for the Slot Attention baseline, but less than 517250 for Set Transformer and 461289 for Set2Graph (Serviansky et al., 2020).
The total training time is less than 12 hours on a single GTX 1080 Ti and 10 CPU cores.
Further results.
For completeness, we also report the results for the rand index (RI) in Table 4.
B. C ONVEX HULL FINDING
On convex hull ﬁnding in 3D, we compare our method to the same baselines as on the particle partitioning task.
Setup.
Set2Graph learns to map the set of 3D points to the 3rd order adjacency tensor.
Since storing this tensor in memory is not feasible, they instead concentrate on a local version of the [Table] [Table] problem, which only considers the k -nearest neighbors for each point (Serviansky et al., 2020).
We train our method with T total =48 , T BPTT =4 , N BPTT =6 and set k equal to the highest number of triangles in the training data.
At test time, a prediction admits an edge e i if its existence indicator σ i > 0 .
Each edge is incident to the three nodes with the highest incidence probability.
We apply the same hyperparameters, architectures and optimizer as in the particle partitioning experiment, except for: T =48 , B =4 , N =6 .
We set the maximum number of edges k equal to the highest number of triangles in the training data.
Since we do not change the model, the number of parameters remains at 268162 for our model.
This notably differs to Set2Graph, which reports an increased parameter count of 1186689 (Serviansky et al., 2020).
We train our method until we observe no improvements on the F1 validation performance for 20 epochs, with a maximum of 1000 epochs.
The set-to-set baselines are trained for 4000 epochs, and we retain the parameters resulting in the highest f1 score on the validation set.
The total training time is between 14 and 50 hours on a single
GTX 1080 Ti and 10 CPU cores.
B.
[Equation] TRIANGULATION
The problem of Delaunay triangulation is, similar to convex hull ﬁnding a well-studied problem in computational geometry and has exact solutions in O ( n log ( n ))
(Rebay, 1993).
We consider the same learning task as Serviansky et al.
(2020), who frame Delaunay triangulation as mapping from a set of 2D points to the set of Delaunay edges, represented by the adjacency matrix.
Note that this differs from ﬁnding the set of triangles, as an edge no longer remembers which triangles it is part of.
Thus, this reduces to a set-to-graph task, instead of a set-to-hypergraph task.
Model adaptation.
The goal in this task is to predict the adjacency matrix of an ordinary graph – a graph consisting of edges that connect two nodes – where the number of edges are greater than the number of nodes.
One could recover the adjacency matrix based on the matrix product of I T I , by clipping all values above 1 back to 1 and setting the diagonal to 0.
This approach is inefﬁcient, since in this case the incidence matrix is actually larger than the adjacency matrix.
Instead of applying our method directly, we consider a simple adaptation of our approach to the graph setting.
We replace the initial set of edges with the (smaller) set of nodes and apply the same node reﬁnements on both sets.
This change results in E = V for the prediction and effectively reduces the incidence matrix to an adjacency matrix, since it is computed based on all pairwise combinations of E and V .
We further replace the concatenation for the MLP modelling the incidence probability with a sum, to ensure that the predicted adjacency matrix is symmetric and represents an undirected graph.
Two of the main design choices of our approach remain in this adaptation: Iterative reﬁning of the complete graph with a recurrent neural network and BPTT with gradient skips.
We train our model with
T =32 , B =4 and N =4 .
At test-time, an edge between two nodes exists if the adjacency value is greater than 0.5.
Setup.
We increase the latent dimensions to d =256 , resulting in 595201 trainable parameters.
This notably differs to Set2Graph, which increases the parameter count to 5918742 (Serviansky et al., 2020), an order of magnitude larger.
The total training time is less than 9 hours on a single
GTX 1080 Ti and 10 CPU cores.
B. L EARNING HIGHER ORDER EDGES
The particle partitioning experiment exempliﬁes a case where a single edge can connect up tovertices.
Set2Graph (Serviansky et al., 2020) demonstrates that in this speciﬁc case it is possible to approximate the hypergraph with a graph.
They leverage the fact that any vertex is incident to exactly one edge and apply a post-processing step that constructs the edges from noisy cliques.
Instead, we consider a task for which no straightforward graph based approximation exists.
Speciﬁcally, we consider convex hull ﬁnding in 10 -dimensional space for 13 standard normal distributed points.
We train with T =32 , N =4 and B =4 . The test performance reaches an F1 score of 0.75, clearly demonstrating that the model managed to learn. This result demonstrates the improved scaling behavior can be leveraged for tasks that are computationally out of reach for adjacency-based approaches. We demonstrated that the improved scaling behavior of our proposed method can be leveraged for tasks that are computationally out of reach for adjacency based approaches.
The number of points and dimensions were chosen in conjunction, such that the corresponding adjacency tensor would require more storage than is feasible with current GPUs (available to us).
For 13 points indimensions, explicitly storing the full adjacency tensor using 32-bit ﬂoating-point numbers would already require more than 500 GB.
We intentionally kept the number of points and dimensions low, to highlight that the asymptotic scaling issue cannot be met by hardware improvements, since small numbers already pose a problem.
Note that Set2Graph already struggles with convex hull ﬁnding in 3D, where the authors report that storing 3-rd order tensors in memory is not feasible.
Instead, they consider a local version of the problem and take the k -Nearest-Neighbors out of the set of points that are part of the convex hull, with k = 10 .
While we limited our calculation of the storage requirement to the adjacency tensor itself, a typical implementation of a neural network also requires storing the intermediate activations, further exacerbating the problem for adjacency based approaches.
B. B ACKPROP WITH SKIPS
We compare backprop with skips to TBPTT (Williams & Peng, 1990) with B =4 every 4 iterations, which is the setting that is most similar to ours with regard to training time.
In general, TBPTT allows for overlaps between subsequent BPTT applications, as we illustrate in Figure 4.
We constrict both TBPTT and backprop with skips to a ﬁxed memory budget, by limiting any backward pass to the most recent B =4 iterations, for T ∈{ 16 , 32 } .
The standard backprop results serve as a reference point to answer the question: “What if we apply backprop more frequently, resulting in a better approximation to the true gradients?”, without necessitating a grid search over all possible hyperparameter combinations for TBPTT.
The results on standard backprop appear to indicate that performance worsens when increasing the number of iterations from 16 to 32 .
We observe that applying backprop on many iterations leads to increasing gradient norms in the course of training, complicating the training process.
The memory limited versions did not exhibit a similar behavior, evident from the improved performance, when increasing the iterations from 16 to 32 .
[Figure]
with skips at iterations 1, 6, 7, 8, which effectively reduces the training time, while retaining the same number of reﬁnement steps.