On Deep Set Learning and the Choice of Aggregations

TL;DR

This paper explores the impact of different aggregation functions in Deep Set learning, demonstrating that learnable aggregations improve performance, reduce hyper-parameter sensitivity, and enhance out-of-distribution generalization.

Contribution

It introduces and empirically evaluates learnable aggregation functions as alternatives to standard ones in Deep Set architectures.

Findings

Learnable aggregations outperform fixed ones in accuracy.

Deep Set networks are highly sensitive to aggregation choice.

Learnable aggregations improve out-of-distribution generalization.

Abstract

Recently, it has been shown that many functions on sets can be represented by sum decompositions. These decompositons easily lend themselves to neural approximations, extending the applicability of neural nets to set-valued inputs---Deep Set learning. This work investigates a core component of Deep Set architecture: aggregation functions. We suggest and examine alternatives to commonly used aggregation functions, including learnable recurrent aggregation functions. Empirically, we show that the Deep Set networks are highly sensitive to the choice of aggregation functions: beyond improved performance, we find that learnable aggregations lower hyper-parameter sensitivity and generalize better to out-of-distribution input size.

Figures

Tables

Table 1. Table 1 : Test set accuracy on ModelNet40 classification.

	Equivariant layer type & aggregation type
$\left|𝒳\right|$	$\max$ $\max$	$\max$ r-$L\Sigma E$	$\max$ r-$\mathrm{sum}$	$\max$ q-$\max$	$\max$ q-$\mathrm{sum}$	r-$\mathrm{sum}$ r-$\mathrm{sum}$	$\max$ r-$\max$	r-$\max$ r-$\max$	r-$L\Sigma E$ r-$L\Sigma E$	q-sum q-sum
1000	87.3	85.8	85.7	83.8	83.5	82.0	81.7	81.2	78.0	77.5
100	66.5	75.3	73.0	69.5	68.4	71.9	45.3	22.0	64.0	60.3
50	47.0	62.8	58.4	52.4	51.3	61.0	35.5	14.6	51.9	46.8

Equations

$$f\left(\left\{{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\right\}\right)=f\left(\left\{{\mathbf{x}}_{\pi (1)},\dots ,{\mathbf{x}}_{\pi (N)}\right\}\right).$$

$$f({\mathbf{X}}_{\pi })=(f(\mathbf{X}){)}_{\pi }.$$

$$f(\mathcal{X})=\rho \left(\sum _{\mathbf{x}\in \mathcal{X}}\varphi \left(\mathbf{x}\right)\right),$$

$$\mathrm{equivariant}(\mathbf{X})=\sigma \left(\mathbf{X}-1\alpha \left(\mathbf{X}\right)\right),$$

$${\mathbf{m}}_i$$

$$\mathbf{C}$$

$$\mathbf{a}$$

$$\mathbf{r}$$

$${\mathbf{q}}_t$$

$${\hat{w}}_{i,t}$$

$${\mathbf{w}}_t$$

$${\mathbf{a}}_t$$

$$\mathbf{a}$$

$${\mathbf{q}}_t$$

$${\hat{w}}_{i,t}$$

$${\mathbf{w}}_t$$

$${\mathbf{a}}_t$$

$$\mathbf{a}$$

Full text

1 Introduction

Machine learning algorithms make implicit assumptions on the data set encoding. For instance, feed-forward neural networks assume that data is encoded in a unique vector representation, e. g. by one-hot encoding categorical variables. Yet, many interesting learning tasks revolve around data sets consisting of sets: depth vision with 3D point clouds, probability distributions represented by finite samples, or operations on unstructured sets of tags [16, 26, 21].

Naively, a population111 Disambiguating terms like set and sample, we discuss data sets of populations of particles. is embedded by ordering and concatenating particle vectors into a matrix. While standard neural networks can learn to imitate order-invariant behavior, e.g. by random input permutation at each gradient step, such architectures are no true set functions. Further, they cannot easily handle varying population sizes. This motivated research into order-invariant neural architectures [24, 6, 4, 20]. From this, the Deep Set framework emerged, proving that many interesting invariant functions allow for a sum decomposition [29, 18, 25]. It allows for straightforward application of neural networks that are order-invariant by design, and can handle varying population sizes.

In this work, we study aggregations—the component of a Deep Set architecture that induces order invariance by mapping a variable-sized population to a fixed-sized description. After discussing desirable properties and extending the theory around aggregation functions, we suggest multiple alternatives, including learnable recurrent aggregation functions. Studying them in several experimental settings, we find that the choice of aggregation impacts not only the performance, but also hyper-parameter sensitivity and robustness to varying population sizes. In the light of these findings, we argue for new evaluation techniques for neural set functions.

2 Order-Invariant Deep Architectures

We discuss populations $\mathcal{{X}}$ of particles $\bm{\mathbf{x}}$ from a particle space $\mathscr{{X}}\subset\mathbb{R}^{d}$, i. e. $\bm{\mathbf{x}}\in\mathcal{{X}}$ and $\mathcal{{X}}\subset\mathscr{{X}}\subset\mathbb{R}^{d}$. We are further interested in representations $\mathbf{{X}}\in\mathbb{R}^{p\times d},p={\left\lvert\mathcal{{X}}\right\rvert},$ achieved by concatenating the particles of $\mathcal{{X}}$. A permutation of the particle axis with a permutation $\pi$ is denoted by $\mathbf{{X}}_{\pi}$, i. e. $\mathbf{{X}}\neq\mathbf{{X}}_{\pi}$ but $\mathbf{{X}}\equiv\mathcal{{X}}\equiv\mathbf{{X}}_{\pi}$. Data sets $\mathcal{D}$ consist of finite populations $\mathcal{{X}}_{i}$ of potentially varying size.

2.1 Invariance, Equivariance, and Decomposition of Invariant Functions

We study invariant functions according to

Definition 1 (Invariance)

A function $f$ on the power set $\mathcal{P}\left(\mathscr{{X}}\right)$ is order-invariant if for any permutation $\pi$ and input ${\left\{\bm{\mathbf{x}}_{1},\dots,\bm{\mathbf{x}}_{N}\right\}}\in\mathcal{P}\left(\mathscr{{X}}\right)$

[TABLE]

If it is clear from the context, we will call such functions invariant. When the input is embedded as a matrix, definition 1 can be formulated as $f(\mathbf{{X}})=f(\mathbf{{X}}_{\pi}).$ A related, important notion is that of equivariant functions:

Definition 2 (Equivariance)

A function $f$ is equivariant if input permutation results in equivalent output permutation, i. e. for any $\mathbf{{X}}$ and $\mathbf{{X}}_{\pi}$

[TABLE]

In [29], a defining structural property of order-invariant functions was proven:

Theorem 2.1 (Deep Sets, [29])

A function $f$ on populations $\mathcal{{X}}$ from countable particle space $\mathscr{{X}}$ is invariant if and only if there exists a decomposition,

[TABLE]

with appropriate functions $\phi$ and $\rho$.

We call such functions sum-decomposable; this follows [25], where severe pathologies for uncountable input spaces are pointed out:

There exist invariant functions that have no sum decomposition. 2. 2.

There exist sum decompositions that are everywhere-discontinuous. 3. 3.

Even relevant functions such as $\max(\mathcal{{X}})$ cannot be continuously decomposed when the image space of the embedding $\phi$ is smaller than the population size $|\mathcal{{X}}|$.

As a consequence they refine theorem 2.1 to

Theorem 2.2 (Uncountable Particle Spaces, [25])

A continuous function $f$ on finite populations $\mathcal{{X}}$, $|\mathcal{{X}}|\leq p$, is invariant if and only if it is sum-decomposable via $\mathbb{R}^{p}$.

That is, for arbitrary $f$, the image space of $\phi$ has to have at least dimension $p$, which is both necessary and sufficient. More restrictive in scope than theorem 2.1, it is more applicable in practice where most function approximators—neural networks, Gaussian processes—are continuous.

2.2 Deep Sets

A generic invariant neural architecture emerges from theorems 2.1 and 2.2 by using neural networks for $\rho$ and $\phi$, respectively. In practice, to allow for higher-level particle interaction during the embedding $\phi$, equivariant neural layers are introduced [29],

[TABLE]

where $\sigma(\cdot)$ denotes a per-particle feed-forward layer, and $\alpha(\cdot)$ denotes an aggregation. Aggregations—our object of study—induce invariance by mapping a population to a fixed-size description, typically e. g. sum, mean, or $\max$. The full architecture is

[TABLE]

with $\phi$ implemented by a per-particle embedding followed by an equivariant combination function consisting of equivariant layers. Summation is replaced by a generic aggregation operation. In [29, 18], the $\max$ operation is suggested as an alternative summation. Lastly, $\rho$ can be implemented by arbitrary functions, since the aggregation in eq. 4 is already invariant. This framework is depicted in fig. 1(a).

2.3 Order Matters

Recurrent neural networks can handle set-valued input by feeding one particle at a time. However, it has been shown that the result is sensitive to order, and an invariant read-process-write architecture has been suggested as a remedy [24]:

[TABLE]

An embedded memory is queried. The invariant result $\bm{\mathbf{a}}_{t}$ is iteratively used to refine subsequent queries with an LSTM [8]. It is not obvious how to cast the recurrent structure into the setting of eqs. 4, 3, 5, 2, 2.1 and 2.2. To the best of our knowledge, this model has only been discussed in its sequence-to-sequence context. We will revisit and refine this architecture in section 3.3.

2.4 Further Related Work

Several papers introduce and discuss a Deep Set framework for dealing with set-valued inputs [18, 29]. A driving force behind research into order-invariant neural networks are point clouds [19, 17, 18], where such architectures are used to perform classification and semantic segmentation of objects and scenes represented as point clouds in $\mathbb{R}^{3}$. It is further shown that a $\max$ decomposition allows for arbitrarily close approximation [18].

Generative models of sets have been investigated: in an extension of variational auto-encoders [11, 22], the inference of latent population statistics resembles a Deep Sets architecture [4]. Generative models of point clouds are proposed by [1] and [28].

Permutation-invariant neural networks have been used for predicting dynamics of interacting objects [6]. The authors propose to embed the individual object positions in pairs using a feed-forward neural network. Similar pairwise approaches have been investigated by [3, 2], and applied to relational reasoning in [23].

Weighted averages based on attention have been proposed and applied to multi-instance learning [10]. Several works have focused on higher-order particle interaction, suggesting computationally efficient approximations of Janossy pooling [15], or propose set attention blocks as an alternative to equivariant layers [14].

3 The Choice of Aggregation

The invariance of the Deep Set architecture emerges from invariance of the aggregation function—eq. 4. Theorem 2.1 theoretically justifies summing the embeddings $\phi\left(\bm{\mathbf{x}}_{i}\right)$. In practice, mean or max-pooling operations are used. Equally simple and invariant, they are numerically favorable for varying population sizes, controlling input magnitude to downstream layers. This section discusses alternatives and their properties.

3.1 Alternative Aggregations

We start by justifying alternative choices with an extension of theorems 2.1 and 2.2:

Corollary 1 (Sum Isomorphism)

Theorems 2.1* and 2.2 can be extended to aggregations of the form ${\alpha}_{g}={g}\circ{\sum}\circ{g^{-1}}$, i. e. summations in an isomorphic space.*

Proof

From ${\rho}\circ{\sum}\circ{\phi}=({\rho}\circ{g^{-1}})\circ{g}\circ{\sum}\circ{g^{-1}}\circ({g}\circ{\phi})$, sum decompositions can be constructed from ${\alpha}_{g}$-decompositions and vice versa.

This class includes, e. g., mean (with $g((x_{1},\dots,x_{n+1}))=(x_{1},\dots,x_{n})/x_{n+1}$ and $g^{-1}(\bm{\mathbf{x}})=(\bm{\mathbf{x}}^{\top},1)^{\top}$) and $\operatorname{logsumexp}$ ($\operatorname{L\Sigma E}$) (with $g=\ln$).

In that light, there is an interesting case to be made for $\operatorname{L\Sigma E}$: depending on the input magnitudes, $\operatorname{L\Sigma E}$ can behave akin to $\max$ (cf. figs. 2(a), 2(b) and 2(c)) or like a linear function akin to summation (cf. figs. 2(f), 2(d) and 2(e)). Operating in log space, $\operatorname{L\Sigma E}$ further exhibits diminishing returns: $N$ identical scalar particles $x_{i}$ yield $\operatorname{L\Sigma E}({\left\{x_{i}\right\}})=\ln(N)+x_{1}$. The larger $N$, the smaller the output change from additional particles. Beyond making $\operatorname{L\Sigma E}$ a numerically useful aggregation, diminishing returns are a desirable property from a statistical perspective, where we would like to have asymptotically consistent results.

Divide and Conquer

Commutative and associative binary operations like addition and multiplication yield invariant aggregations. Widening this perspective, we see that divide-and-conquer style operations yield invariant aggregations: order invariance is equivalent to conquering being invariant to division. Examples beyond the previously mentioned operations are logical operators such as any or all, but also sorting (generalizing $\max$ and $\min$, and any percentile, e. g. median). While impractical for typical first-order optimization, we note that aggregations can be of very sophisticated nature.

3.2 Learnable Aggregation Functions

In [29], cf. eqs. 2, 3, 4 and 5, the aggregation is the only non-learnable component. We will now investigate ways to render the aggregations learnable. In section 2.3, we have seen that due to the structure of theorem 2.1, recurrent architectures as suggested by [24] had been overlooked as it is not straightforward to cast them into the Deep Sets framework. Inspired by the read-process-write architecture, we suggest recurrent aggregations:

Definition 3 (Recurrent and Query Aggregation)

A recurrent aggregation is a function $f(\mathcal{{X}})=\bm{\mathbf{a}}$ that can be written recursively as:

[TABLE]

where $\bm{\mathbf{m}}_{i}=\phi\left(\bm{\mathbf{x}}_{i}\right)$ is an embedding of the input population ${\left\{\bm{\mathbf{x}}_{i}\right\}}$ and $\bm{\mathbf{q}}_{1}$ is a constant. We further call the special case $T=1$ (i. e. a single query $\bm{\mathbf{q}}\equiv\bm{\mathbf{q}}_{1}$) a query aggregation.

As long as $\operatorname{reduce}$ is invariant and $\operatorname{normalize}$ is equivariant, recurrent and query aggregations are invariant. This architectural block is depicted in fig. 1(b).

Building upon eqs. 6, 7, 8, 9 and 10, recurrent aggregations introduce two modifications: firstly, we replace a weighted sum by a general weighted aggregation—giving us a rich combinatorial toolbox on the basis of simple invariant functions such as those mentioned in section 3.1. Secondly, we add post-processing of the step-wise results $\bm{\mathbf{a}}_{1:T}$. In practice, we use another recurrent network layer that processes $\bm{\mathbf{a}}_{1:T}$ in reversed order. Without this modification, later queries tend to be more important, as their result is not as easily forgotten by the forward recurrence. The backward processing reverses this effect, so that the first queries tend to be more important, and the overall architecture is more robust to common fallacies of recurrent architectures, in particular unstable gradients.

Observing eq. 14, we note that our learnable aggregation functions wrap around the previously discussed simpler non-learnable aggregations. A major benefit is that the inputs are weighted—sum becomes weighted average, for instance. This also allows the model to effectively exploit non-linearities as discussed with $\operatorname{L\Sigma E}$ (cf. fig. 2).

3.3 A Note on Universal Approximation

The key promise of universal approximation is that a family of approximators (e. g. neural nets, or neural sum decompositions) is dense within a wider family of interesting functions [12, 7, 9]. The universality granted by theorems 2.1 and 2.2, through constructive proofs, hinges on sum aggregation. Corollary 1 grants flexibility, but does not apply to arbitrary aggregations, like $\max$ or the suggested learnable aggregations. (Note that $\max$ allows for arbitrary approximation [18].) It remains open to what extent the sum can be replaced. As such, the suggested architectures might not grant universal approximators. As we will see in section 4, however, they provide useful inductive biases in practical settings, much like feed-forward neural nets are usually replaced with architectures targeted towards the task. It is worth noting that the embedding dimension constraint of theorem 2.2 is rarely met, trading theoretical guarantees for test-time performance.

4 Experiments

We consider three simple aggregations: mean (or weighted sum), $\max$, and $\operatorname{L\Sigma E}$. These are used in equivariant layers and final aggregations, and may be be wrapped into a recurrent aggregation. This combinatorially large space of configurations is tested in four experiments described in the following sections.

4.1 Mininmal Enclosing Circle

In this supervised experiment, we are trying to predict the minimal enclosing circle of a population of size $20$ from a Gaussian mixture model (GMM). A sample population with target circle is depicted in fig. 4. The sample mean does not approximate the center of the minimal enclosing circle well, and the correct solution is defined by at least three particles. The models are trained by minimizing the mean squared error (MSE) towards the center and radius of the true circle (computable in linear time [27]).

Results are given in fig. 4. Each row shows the best result out of 180 runs (20 runs for each of the 9 combinations of aggregations). We can see that both recurrent equivariant layers and recurrent aggregations improve the performance, with equivariant layers granting the larger performance boost. The challenge lies mostly in a better approximation of the center.

The top row indicates that an entirely non-recurrent model performs better than its counterpart with recurrent aggregation (second row). To test for a performance outlier, we compute a bootstrap estimate of the expected peak performance when only performing 20 experiments: we subsample all available experiments (with replacement) into several sets of 20 experiments, recording the best performance in each batch. The last column in fig. 4 reports the median of these best batch performances. The result shows increased robustness to hyper-parameters, despite having more hyper-parameters.

4.2 GMM Mixture Weights

In this experiment, our goal is to estimate the mixture weights of a Gaussian mixture model directly from particles. The GMM populations of size $100$ in our data set are sampled as follows: each mixture consists of two components; the mixture weights are sampled from $[.05,.95]$; the means span a diameter of the unit circle, their position is drawn uniformly at random; component variances are a fixed to the same diagonal value such that the clusters are not linearly separable. An example population is shown in fig. 5(a). The model outputs concentrations $a$ and $b$ of a Beta distribution. We train to maximize the log-likelihood of the smaller ground truth weight under this Beta distribution. At training time, for every gradient step the batch population size $N$ is chosen randomly, with $p\left(N=n\right)\propto n$. In fig. 5(a), we show how an estimator based on the learned model behaves with growing population size.

We were again interested in the robustness of the models. We compare to expectation maximization (EM)—the classic estimation technique for mixture weights—as a baseline by gathering 100 estimates each from EM and the model for each population size by subsampling (with replacement) the original population. Then we compare the likelihood of the true weight under a kernel density estimate (KDE) of these estimates. The final metric is the log ratio of the scores under the two KDEs. Then, as in the previous section, we compute the peak performance for batches of 5 experiments in order to see which configurations of models consistently perform well.

The results of this analysis are shown in fig. 5(b). The top row indicates that learnable equivariant layers lead to a significant performance boost across all reduction operations. Note that the y-axis is in log scale, indicating multiples of improvements over the EM baseline. We note that $\operatorname{L\Sigma E}$ benefits most drastically from learnable inputs. Notably, the middle column, which depicts $\max$-type aggregations, indicates that this type of aggregation significantly falls behind the alternatives. Notice that we had to scale the y-axes to even show the violins, and that a significant amount of peak performances perform worse than EM (indicated by sign flip of the metric).

4.3 Point Clouds

The previous experiment extensively tested the effect of aggregations in controlled scenarios. To test the effect of aggregations on a more realistic data set, we tackle classification of point clouds derived from the ModelNet40 benchmark data set [30]. The data set consists of CAD models describing the surfaces of objects from 40 classes. We sample point cloud populations uniformly from the surface. The training is performed on 1000 particles. For this experiment, we fixed all hyper-parameters—including optimizer parameters and learning rate schedules—as described in [29], and only exchanged the aggregation functions in the equivariant layers and the final aggregation.

The results for the 10 best configurations are summarized in table 1. The original model ($\max$/$\max$ column) performs best in the training scenario (${\left\lvert\mathcal{{X}}\right\rvert}=1000$, first row)—as expected on hyper-parameters that were optimized for the model. Otherwise, learnable final aggregations outperform all non-learnable aggregations. We further observe that $\max$-type aggregations in equivariant layers seem crucial for good final performance. This contrasts the findings from section 4.2. We believe this to be a result of either (i) the hyper-parameters being optimized for $\max$-type equivariant layers, or (ii) the classification task (as opposed to a regression task), favoring $\max$-normalized embeddings that amplify discriminative features.

The second and third row highlight an insufficiently investigated problem with invariant neural architectures: the top-performing model overfits to the training population size. Despite sharing all hyper-parameters except the aggregations, the test scenarios with fewer particles show that learnable aggregation functions generalize favorably. Compare the first two columns: both drops for the original model are comparable to the total drop for the learnable model.

4.4 Spatial Attention

In the previous experiments, we investigated models trained in isolation on supervised tasks. Here, we will test the performance as a building block of a larger model, trained end-to-end and unsupervised. The data consists of canvases containing multiple MNIST digits, cf. fig. 6(a). In [5], an unsupervised algorithm for scene understanding of such canvases was introduced. We plug an invariant model as the localization module, which repeatedly attends to the input image, at each step returning the bounding box of an object. To turn a canvas into a population, we interpret the gray-scale image as a two-dimensional density and create populations by sampling 200 particles proportional to the pixel intensities. Remarkably, the set-based approach requires an order of magnitude fewer weights, and consequently has a significantly lower memory footprint compared to the original model, which repeatedly processes the entire image.

The task is challenging in several ways: the loss is a lower bound to the likelihood of the input canvas, devoid of localization information. The intended localization behavior needs to emerge from interaction with downstream components of the overall model. As with enclosing circles, the bounding box center is correlated with the sample mean of isolated particles from one digit. However, depending on the digit, this can be inaccurate.

As fig. 6(b) indicates, the order-invariant architecture on 200 particles (as in training, vertical line) can serve as a drop-in replacement, performing on a par or slightly improved compared to the original model baseline, indicated by the vertical line. This is remarkable, with the original model being notoriously hard to train [13].

We investigate the performance of the model when the population size varies. We observe that the effect on performance varies with different aggregation functions. Learnable aggregation functions exhibit strictly monotonic performance improvements. This is reflected by tightening bounding boxes for increasing population sizes, fig. 6(a). Similar behavior cannot be found reliably for non-learnable aggregations. Note that we can trade off performance and inference speed at test time by varying the population size.

Lastly, we note that in both this and the point cloud experiment, section 4.3, learnable $\operatorname{L\Sigma E}$-aggregations performed well. We attribute this to the properties of diminishing returns and sum-$\max$-interpolation amplified by weighted inputs, cf. section 4.

5 Discussion and Conclusion

We investigated aggregation functions for order-invariant neural architectures. We discussed alternatives to previously used aggregations. Introducing recurrent aggregations, we showed that each component of the Deep Set framework can be learnable. Establishing the notion of sum isomorphism, we created ground for future aggregation models.

Our empirical studies showed that aggregation functions are indeed an orthogonal research axis within the Deep Set framework worth studying. The right choice of aggregation function may depend on the type of task (e. g. regression vs. classification). It affects not only training performance, but also model sensitivity to hyper-parameters and test time performance on out-of-distribution population sizes. We showed that the learnable aggregation functions introduced in this work are more robust in their performance and more consistent in their estimates with growing population sizes. Lastly, we showed how to exploit these features in larger architectures by using neural set architectures as drop-in replacements. In the light of our experimental results, we strongly encourage emphasizing desirable properties of invariant functions, and in particular actively challenge models in non-training scenarios in future research.