On Learning Sets of Symmetric Elements (Extended Abstract)

1 Introduction

Learning with data that consists of unordered sets of elements is an important problem with numerous applications, from classification and segmentation of 3D data [Zaheer et al., 2017; Qi et al., 2017; Su et al., 2015; Kalogerakis et al., 2017] to image deblurring [Aittala and Durand, 2018]. In this setting, each data point consists of a set of elements, and the task is independent of element order. This independence induces symmetry structure, which can be used to design deep models with improved efficiency and generalization. Indeed, models that respect set symmetries, e.g. [Zaheer et al., 2017;

∗The original version of this paper was presented at the International Conference on Machine Learning (ICML), 2020.

†Contact Author

Importantly, the elements of the set themselves often adhere to certain symmetries, as happens when learning with sets of images, sets of point clouds, and sets of graphs. It is still unknown what is the best way to utilize these additional symmetries.

A common approach to handle per-element symmetries is based on processing elements individually. First, one processes each set-element independently into a feature vector using a Siamese architecture [Bromley et al., 1994], and only then fuses information across all feature vectors. When following this process, the interaction between the elements of the set only occurs after each element has already been processed, possibly omitting low-level details. Indeed, it has been recently shown that for certain visual tasks [Aittala and Durand, 2018; Sridhar et al., 2019; Liu et al., 2019], significant gain can be achieved with intermediate informationsharing layers.

Here, we present a principled approach to learning sets of symmetric elements. First, we describe the symmetry group of these sets, and then fully characterize the space of linear layers that are equivariant to this group. Notably, this characterization implies that information between set elements should be shared in all layers. For example, Figure 1 illustrates a DSS layer for sets of images. DSS layers provide a unified framework that generalizes several previously-described architectures for a variety of data types. In particular, it directly generalizes DeepSets [Zaheer et al., 2017]. Moreover, other recent works can also be viewed as special cases of our approach [Hartford et al., 2018; Aittala and Durand, 2018; Sridhar et al., 2019].

A potential concern with equivariant architectures is that restricting layers to be equivariant to some group of symmetries may reduce the expressive power of the model [Maron et al., 2019; Morris et al., 2018; Xu et al., 2019]. We address this concern by proving two universal-approximation theorems for invariant and equivariant DSS networks. Simply put, these theorems state that if invariant (equivariant) networks for the elements of interest are universal, then the corresponding invariant (equivariant) DSS networks on sets of such elements are also universal. One important corollary of these results is that DSS networks are strictly more expressive than Siamese networks.

2 Preliminaries

2.1 Notation and Definitions

Let x ∈ Rℓ represent an input that adheres to a group of symmetries G ≤ Sℓ, the symmetric group on ℓ elements. G captures those transformations that our task-of-interest is invariant (or equivariant) to. The action of G on Rℓ is defined by (g · x)i = xg-1 (i). For example, when inputs are images of size h × w, we have ℓ = hw and G can be a group that applies cyclic translations, or left-right reflections to an image. A function is called G-equivariant if f(g · x) = g · f(x) for all g ∈ G. Similarly, a function f is called G-invariant if f(g · x) = f(x) for all g ∈ G.

2.2 G-invariant and G-equivariant Networks

G-equivariant networks are a popular way to model Gequivariant functions. These networks are composed of several linear G-equivariant layers, interleaved with activation functions like ReLU, and have the following form:

where Li : Rℓ×di →Rℓ×di+1 are linear G-equivariant layers, di are their feature dimensions and σ is a point-wise activation function. It is straightforward to show that this architecture results in a G-equivariant function. G-invariant networks are defined by adding an invariant layer on top of a G-equivariant function followed by a multilayer Perceptron.

2.3 Characterizing Equivariant Layers

The main building block of G-invariant/equivariant networks are linear G-invariant/equivariant layers. To implement these networks, one has to characterize the space of linear Ginvariant/equivariant layers. For example, it is well known that for images with the group G of circular 2D translations, the space of linear G-equivariant layers is simply the space of all 2D convolutions operators [Puschel and Moura, 2008]. Unfortunately, such elegant characterizations are not available for most permutation groups. See [Wood and ShaweTaylor, 1996; Ravanbakhsh et al., 2017] for further details.

2.4 Deep Sets

The current paper generalizes DeepSets [Zaheer et al., 2017] and we summarize their main results for completeness. Let {x1, . . . xn} ⊂ Rd be a set, which we represent in arbitrary order as a matrix X ∈ Rn×d . [Zaheer et al., 2017] characterized all Sn-equivariant linear layers for this case, which take the form:

where L1,L2 : Rd → Rd are general linear functions and the subscript represents the i-th row of the output. [Zaheer et al., 2017; Qi et al., 2017] established the universality of invariant networks that are composed of DeepSets Layers and [Segol and Lipman, 2019] extended this result to the equivariant case.

3 DSS Layers

Our main goal is to design deep models for sets of elements with non-trivial per-element symmetries. In this section, we first formulate the symmetry group G of such sets. The deep models we advocate are composed of linear G-equivariant layers (DSS layers); therefore, our next step is to find a simple and practical characterization of the space of these layers.

3.1 Sets with Symmetric Elements

Let {x1, . . . xn} ⊂ Rd be a set of elements with symmetry group H ≤ Sd. We wish to characterize the space of linear maps L : Rn×d → Rn×d that are equivariant to both the natural symmetries of the elements, represented by the elements of the group H, as well as to the order of the n elements, represented by Sn.

In our setup, H operates on all elements xi in the same way. More formally, the symmetry group is defined by G = Sn × H, where Sn is the symmetric group on n elements. This group operates on X ∈ Rn×d by applying the permutation q ∈ Sn to the first dimension and the same element h ∈ H to the second dimension, namely ((q, h) · X)ij = Xq-1 (i)h-1 (j). Figure 2 illustrates this setup. Notably, this setup generalizes several popular learning setups: (1) DeepSets, where H = {Id} is the trivial group. (2) Tabular data [Hartford et al., 2018], where H = Sd. (3) Sets of images, where H is the group of cyclic translations [Aittala and Durand, 2018].

Considering the symmetry in the form of a direct product suggests implicitly that the elements {xi} are aligned, as they are in many popular tasks like image deblurring [Aittala and Durand, 2018] (also see figure 1). For other datasets and tasks, this assumption might not hold. In that case, it is natural to consider another setup, where the members of H applied to each element may differ 1 . Unfortunately, as we show in the full paper [Maron et al., 2020], the corresponding equivariant layers are essentially Siamese layers. See also [Wang et al., 2020].

3.2 Characterization of Equivariant Layers

This subsection provides a practical characterization of linear G-equivariant layers for G = Sn ×H. Our result generalizes DeepSets (equation 2) whose layers are tailored for H = {Id}, by replacing the linear operators L1, L2 with linear H-equivariant operators. This result is summarized in the following theorem:

Theorem 1. Any linear G-equivariant layer L : Rn×d → Rn×d is of the form

where LH1 ,LH2 are linear H-equivariant functions

Note that this is equivalent to the following formulations L(X)i = LH1 (xi) + LH2 (∑nj=1 xj) = LH1 (xi) + ∑nj=1 LH2 (xj) due to linearity. Figure 1 illustrates Theorem 1 for sets of images. In this case, applying a DSS layer amounts to: (i) Applying the same convolutional layer L1 to all images in the set (blue); (ii) Applying another convolutional layer L2 to the sum of all images (orange); and (iii) summing the outputs of these two layers.

Relation to previous work. In the case of a set of images and translation equivariance, LHi are convolutions. In this setting, [Aittala and Durand, 2018; Sridhar et al., 2019] have previously proposed using similar set-aggregation layers after convolutional blocks. This work provides a theoretical analysis of their setup and generalizes it all permutation groups.

4 AUniversal Approximation Theorem

When restricting a network to be invariant (equivariant) to some group action, one may worry that these restrictions could reduce the network expressive power. We now show that networks that are constructed from DSS layers do not suffer from loss of expressivity. Specifically, we show that for any group H that induces a universal H-invariant (equivariant) network, its corresponding G-invariant (equivariant) network has high expressive power: we prove universal approximation for invariant and equivariant functions defined on any compact set with zero intersection with a specific lowdimensional set E ⊂ Rn×d . The precise definition of E is given in the supplementary of the full paper [Maron et al., 2020]. The following theorem summarizes these results:

1 Here, the symmetry group is the wreath product of Sn and H.

Theorem 2. Let K ⊂ Rn×d be a compact domain such that K = ∪g∈GgK and K ∩ E = ∅. G-equivariant networks are universal approximators (in ‖ · ‖∞ sense) of continuous Rn×d → Rn×d G-equivariant functions on K if Hequivariant networks are universal.

The invariant case is a direct corollary of Theorem 2.

5 Experiments

This section investigates the effectiveness of DSS layers in practice by comparing them to previously suggested architectures and different aggregation schemes. We use the experiments to answer two basic questions: (1) Early or late aggregation? Can early aggregation architectures like DSS and its variants improve learning compared to Late aggregation architectures, which fuse the set information at the end of the data processing pipeline? and (2) How to aggregate? What is the preferred early aggregation scheme?

Tasks. In the full paper, [Maron et al., 2020], we evaluated DSS in a series of six experiments spanning a wide range of tasks: from classification (Rn×d → R), through selection (Rn×d → Rn ) and burst image deblurring (Rn×d → Rd ) to general equivariant tasks (Rn×d → Rn×d ). The experiments also demonstrate the applicability of DSS to a range of data types, including point clouds, images and graphs.

Competing methods. We compare DSS to a narchitecture with a Siamese network followed by DeepSets (Siamese+DS). See full paper for more comparisons. We also compare several variants of our DSS layers: (1) DSS(sum): our basic DSS layer from Theorem 1 (2) DSS(max): DSS with max-aggregation instead of sum-aggregation (3) DSS(Aittala): DSS with the aggregation proposed in [Aittala and Durand, 2018], namely, L(x)i ↦→ [LH (xi),maxnj=1 LH (xj)] where [] denotes feature concatenation and LH is a linear H-equivariant layer (4) DSS(Sridhar): DSS layers with the aggregation proposed in [Sridhar et al., 2019] ,i.e., L(x)i ↦→ LH (xi) - 1n ∑nj=1 LH (xj ).

Evaluation protocol. All models have roughly the same number of parameters for each particular task. In all experiments, we report the mean and standard deviation over five random initializations. Experiments were conducted using NVIDIA DGX with V100 GPUs.

5.1 Selection Tasks

In a selection task, we are given a set and wish to choose one element from the set that obeys a predefined property. Formally, each task is modeled as a G-equivariant function f : Rn×d → Rn , where the output vector represents the probability of selecting each element. The architecture comprises three convolutional blocks employing Siamese or DSS variants, followed by a DeepSets block. We note that the Siamese+DS model was suggested for similar selection tasks in [Zaheer et al., 2017].

Frame selection in images and shapes. The first selection task is to find a particular frame within an unordered set of frames extracted from a video/shape sequence. For videos, we used the UCF101 dataset [Soomro et al., 2012]. Each set contains n = 8 frames generated by randomly drawing a video, a starting position, and frame order. The task is to select the "first" frame, namely, the one that appeared earliest in the video. Table 1 details the accuracy of all compared methods in this task, showing that DSS(sum) and DSS(Aittala) outperform Siamese+DS and DSS(Sridhar) by a large margin.

In a second frame selection task, we demonstrate that DSS can handle multiple data types. Specifically, we showcase how DSS operates on point clouds and graphs. Given a short sequence of 3D human shapes performing various activities, the task is to identify which frame was the center frame in the original non-shuffled sequence. These human shapes are represented as point clouds in the first experiment and as graphs (point clouds + connectivity) in the second experiment.

To generate the data, we cropped 7-frame-long sequences from the Dynamic Faust dataset [Bogo et al., 2017] in which the shapes are given as triangular meshes. In [Bogo et al., 2017] the points of each mesh are ordered consistently, providing point-to-point correspondence across frames. When this correspondence is not available, a shape matching algorithm like [Litany et al., 2017; Maron and Lipman, 2018] can be used as preprocessing. See Figure 3 for an illustration of this task. Results are summarized in Table 1, comparing DSS variants to a late-aggregation baseline (Siamese +DS) and to random choice.

Highest quality image selection. Given a set of n = 20 degraded images of the same scene, the task is to select the highest-quality image. We generate data for this task from the Places dataset [Zhou et al., 2017], by adding noise and Gaussian blur to each image. The target image is defined to be the image that is the most similar in L1 norm sense to the original image. Notably, DSS consistently improves over Siamese+DS with a margin of 1% to 3%. See Table 2.

5.2 Summary of Experiments

The above experiments demonstrate that applying early aggregation using DSS layers improves learning in various tasks and data types, compared with earlier architectures like Siamese+DS. This improvement can be attributed to the provably higher expressive power of DSS networks. More specifically, the basic DSS layer, DSS(sum), performs well on all tasks, and DSS(Aittala) has also yielded strong results. DSS(Sridhar) performs well on some tasks but fails on others. See the full paper for additional experiments: signal classification, image deblurring, and multi-view reconstruction.

6 Conclusion

In this short paper, we have summarized the main results of [Maron et al., 2020]. Specifically, we presented a principled approach for designing deep networks for sets of elements with symmetries. We characterized the space of equivariant maps for such sets, analyzed its expressive power, exemplified its benefits over standard set learning approaches over various tasks and data types, and shown that our approach generalizes several successful previous works. Using our framework as it is, one implicitly assumes that the set elements are aligned. A worthwhile direction for future work is to suggest approaches to circumvent this limitation.

Acknowledgments

This research was supported by an Israel science foundation grant 737/18. We thank Srinath Sridhar and Davis Rempe for valuable discussions.