Set Interdependence Transformer: Set-to-Sequence Neural Networks for Permutation Learning and Structure Prediction

1 Introduction

There is a wide range of challenges where the objective is to find an optimal mapping from an unordered collection of objects to a permutation. This group of set-to-sequence tasks covers combinatorial optimization and structure prediction problems where exhaustive search is often not tractable, lending itself to neural network (NN) approaches.

Set-to-sequence challenges arise in many areas of application. Examples include natural language processing, in the form of sentence and paragraph ordering [Wang and Wan, 2019; Pandey and Chowdary, 2020], text comprehension

[Li and Gao, 2020] and discourse coherence maximization [Farag, 2021]; computer vision for relative attribute learning [Santa Cruz and Fernando, 2017] and rigid point cloud registration [Yew and Lee, 2020]; reinforcement learning for managing the combinatorial action space of an agent [Vinyals et al., 2019]. For an overview, see Jurewicz and Derczynski (2021b). We adapt our method to a novel application domain in the form of predicting the structure of digital catalogs.

Set-to-sequence models can be thought of as consisting of two distinct parts (see figure 1). Firstly, the set encoder obtains representations of both the elements of a set individually and of the set in its entirety. Secondly, a permutation learning module uses these two representations to predict a reordered sequence. Each stage presents unique challenges. Set-input methods are required to handle inputs of any dimensionality due to the examples being sets of varying cardinalities [Lee and Lee, 2019]. Additionally, the learned representation of a given set of cardinality n must be identical for all n! possible permutations of the vector that represents it. This permutation invariance can be achieved through a variety of symmetric functions, from simple summation [Zaheer et al., 2017], through self-attention [Lee and Lee, 2019] to the output of a bipartite matching algorithm [Skianis and Konstantinos, 2020]. Solving challenging set-to-sequence problems requires a degree of relational reasoning, thus these functions benefit from being capable of encoding higher-order interactions between set elements [Huang et al., 2020].

Predicting a permutation is also a challenge by itself. One of the main difficulties in dealing with combinatorial objects is that the number of possible output sequences increases factorially in the cardinality of the set. Additionally, the space of all possible permutations is not smooth, preventing direct use of gradient-based methods without a relaxation of the concept [Diallo et al., 2020]. Finally, when the target is a complex structure represented by a permutation, it can be difficult to obtain evaluation and loss functions that are only sensitive to meaningful alterations of this structure.

To address these challenges, we propose a novel set encoding method which, unlike its predecessors, jointly learns the permutation invariant representation of the entire set and the permutation equivariant representations of individual set elements. Whilst methods that rely entirely on a pooled representation of the set have been applied to set-to-sequence problems [Vinyals et al., 2016], their performance drops sharply as the cardinality of the input set increases, compared to methods that obtain both types of representations [Wang and Wan, 2019; Yin et al., 2020]. To our knowledge, no other NN set encoding method learns the representations of set elements and the entire set jointly through an adjusted attention mechanism. Instead, they obtain the permutation equivariant representations of elements and then pool them through various symmetric operations to derive the encoding of the set proper. In this paper, we show empirically that learning them jointly is advantageous for the purposes of permutation learning and structure prediction.

We propose a complete, 3-part neural network architecture designed for performing the set-to-sequence mapping on inputs of any cardinality, consisting of an initial set encoder, an interdependence encoder and a permutation module. We then showcase its usefulness on a number of challenges, ranging from toy problems such as the Travelling Salesman Problem, through learning context-free and context-sensitive grammars to robust tasks such as sentence ordering and the novel task of catalog structure prediction. We also demonstrate our model's ability to generalize to unseen lengths and empirically prove its ability to learn higher-order interaction rules on a number of easily customized, synthetic structure datasets and via corresponding evaluation functions. All code, hyperparameters and datasets required for repeated experiments are provided in the supplement.

Our main contributions are summarized as follows:

• A novel, fully differentiable set encoding method designed specifically for permutation learning and structure prediction challenges, capable of learning higherorder interactions within sets of any cardinality in a single layer of the proposed transformations.
• A complete set-to-sequence model outperforming stateof-the-art methods on established datasets and within the novel application domain of catalog structure prediction.
• An expanded library for generating synthetic set-tosequence structure datasets. In addition, we provide easy-to-use tools for obtaining detailed performance reports through customizable metrics, which allow researchers to measure and empirically confirm a model's relational reasoning capabilities.

2 Model

The goal is to transform an input set of any cardinality into the permuted sequence of its elements. To do this, the proposed set-to-sequence model consists of three core components: (i) a basic set encoder, (ii) a novel interdependence encoder and (iii) a permutation decoder. The initial encoder uses a learned pooling function to obtain (a) the permutation equivariant representations of individual set elements, and subsequently (b) the permutation invariant representation of the entire set. These two representations are then transformed together in the interdependence encoder, such that higher order interactions between both individual set elements and the set in its entirety can be learned in a single step. Finally, the permutation decoder sequentially selects the elements to form the output sequence by using these two representation via an enhanced pointer attention mechanism. An overview of our complete interdependence architecture is shown in Figure 2.

From a formal standpoint, the model is given a set X of any cardinality n, consisting of individual set elements represented by fixed-length vectors xi of dimensionality d, in the form of a matrix Xπ arbitrarily ordered according to some permutation π:

The task is then to sequentially select individual set elements X in the target order, represented by a vector of indices y ∈ N1×n referencing the original placement of set elements in the Xπ matrix. This process continues until there are no remaining candidate elements.

2.1 Basic Set Encoder

The initial encoder necessarily consists of learned function fe for transforming set elements in a permutation equivariant way, and learned function fs for pooling those element representations (Eπ) into a permutation invariant embedding of the entire set (s), such that:

In our proposed model this basic set encoder takes the form of a simplified Set Transformer [Lee and Lee, 2019], which we chose for its ability to explicitly represent inter-item interactions. We denote initial transformer-style attention toward permutation equivariant element representations Eπ as:

where Q, K, V are the projected query, key and value matrices obtained from Xπ via weight matrices WQ, WK and WV ; √dk is the standard transformer normalizing factor.

This operation is repeated in a multi-head fashion for each of the m heads, whose outputs are concatenated and further transformed via learned weight matrix WO , without positional encoding or dropout:

To obtain a permutation invariant representation (s) of the set we apply simplified Pooling by Multihead Attention [Lee and Lee, 2019], which performs an attention transformation between a learned seed vector k ∈ R1×d as the query and Eπ as the key and value for each of the m heads:

2.2 Set Interdependence Transformer

At this stage, we could rely on the transformations that lead to Eπ to encode the dependencies between elements of X. However, a single layer of such transformations can only explicitly capture pairwise relations, as it computes attention between pairs of elements [Lee and Lee, 2019]. Thus it would require up to n stacks to explicitly encode higher-order interactions among an entire set of cardinality n. Our proposed set encoder allows efficient capture of dependencies between set elements and the set in its entirety through adjusted transformer-style attention. We refer to it as the Set Interdependence Transformer, or SIT.

The SIT performs attention-based transformations between both individual set elements and the permutation-invariant representation of the set itself in the form of an augmented matrix Sπ. The initial permutation invariant representation of the entire set s is treated as an element of a new set:

where ds is a scaling factor equal to the length of the permutation invariant vector representing the entire set, σ is a nonlinear activation function andWmatrices are learned parameters. The above operation can be repeated over any number of heads as described under the inital basic encoder (subsection 2.1.

This method enables encoding of dependencies between individual set elements and the set in its entirety. Even a single SIT layer is capable of modelling higher-order interactions. This is beneficial for tasks over large sets, such as point cloud challenges, where the cardinality is prohibitively high for pairwise transformations.

To obtain the separate set and set-element representations required by a permutation module, we reverse the augmentation and retrieve the transformed s′ vector by its index. Note that the required permutation invariance of s′ is maintained during the permutation equivariant SITi transformations.

2.3 Permutation Decoder

Finally, we use the E′π and s′ representations in a permutation decoder employing enhanced pointer-style attention [Yin et al., 2020]. This takes the form of an LSTM-based pointer network with two additional mechanisms for pairwise ordering predictions towards improved global and local coherence of the output sequence. Formally, we calculate the conditional probability of a predicted order ˆy as:

where v, W1 and W2 are model parameters, n is the set's cardinality, tanh is the hyperbolic tangent nonlinearity, e′i-1 is the embedding of the set element selected at previous step i -1 and hdi is the hidden state of the permutation module at current step i. The first hidden state hd0 is initialized from the permutation invariant set representation s′ . The Mi matrix provides additional context consisting of two types of information. The first is global orientation relating all remaining unordered set elements to each other and the second is local coherence between previously selected elements and remaining candidates. This context is obtained through "history" and "future" sub-modules from E′π . The sub-modules make pairwise ordering predictions in relation to each candidate element, which are then concatenated to formMi. For specific implementation details, see Yin et al. (2020).

During training, given a batch B of n examples of the form (Xi, yi), we minimize the loss function:

where θ is the set of all model parameters and λ = 0.1 is a hyperparameter that balances the first term of the loss with Ls, a cross-entropy loss of the future and history sub-modules.

3 Experiments

All datasets, hyperparameters and code are freely available and described in detail in the linked supplementary material. The provided code includes any and all data pre-processing and generation steps, where applicable. Full requirements are explicitly listed. Best performance is reported in bold, second best is underlined. Reported results are averaged over 3 full training runs, standard deviation is reported after the ± sign.

3.1 Setup

We compare all models on the planar symmetric Travelling Salesman Problem (TSP) [Vinyals et al., 2015] and 3 types of context-free and context-sensitive grammars, as well as the ROCStory sentence ordering dataset [Mostafazadeh and Chambers, 2016] and the PROCAT product catalog dataset [Jurewicz and Derczynski, 2021a], following provided train, test and validation splits. Additionally, we expand upon the synthetic structure dataset from PROCAT and report performance per n-th order rule of interaction required.

We use Kendall's Rank Correlation Coefficient (τ ) and Perfect Match Ratio (PMR) as primary metrics, scaled by a factor of a hundred for readability, following convention [Wang and Wan, 2019; Yin et al., 2020; Pandey and Chowdary, 2020]. For TSP we report average tour length (shorter is better) and in both grammars and synthetic structure experiments the percentage of valid predictions per ruleset. For the representations of natural language entities in ROCStory and PROCAT we use the concatenated and averaged output of the last 4 layers of the cased large version of BERT [Devlin et al., 2019], frozen during training to isolate the effect of different set encoding methods on the permutation task performance. The models' exact layer dimensions are given in the supplement, with the number of learnable parameters of each model varying by less than 5% per task. The proposed model consisted of 3-5 layers of SIT transformations over 4 attention heads, with the hidden size of 256. The AdamW [Loshchilov and Hutter, 2017] optimizer was used with weight decay coefficient 1e-2, learning rate (α) 1e-4, dropout rate of 0.1 and batch size 32, for 50 epochs. All experiments were performed on cloud-based GPU instances, with NVIDIA Quadro P6000 graphics cards (24 GB) and 8 CPU cores.

3.2 Baselines

Our proposed model is compared with state-of-the-art set encoders combined with the same permutation component:

• DeepSets: Zaheer et al. [2017] encodes set elements independently and identically via a fully connected layer and then sums them.
• Set Transformer: Lee et al. [2019] uses transformerstyle attention in a learned pooling by multihead attention function.
• AttSets: Yang et al. [2020] uses convolutional weighted attention and sum pooling.
• RepSet: Skianis et al. [2020] obtains permutation invariance through bipartite graph matching.

and with several complete set-to-sequence models:

• Pointer Network: Vinyals et al. [2015] consists of an RNN set encoder and basic pointer-style attention without pairwise predictions.
• Read-Process-Write: Vinyals et al. [2016] obtains permutation invariance via an adjusted LSTM network, which performs a number of passes over the input set.
• ATTOrderNet: Vinyals et al. [2018] combines transformer-style attention for set encoding with layer normalization, average pooling and a basic pointer-style permutation module.
• Enhanced PtrNet: Yin et al. [2020] introduces the global and local context sub-modules, adding pairwise judgments to the permutation mechanism.

3.3 Core Results

As shown in Table 2, our proposed model outperforms baselines on the TSP by predicting shorter average paths and with smaller standard deviation. The improvement is more pronounced as input set cardinalities increasingly differ from those seen during training, unseen-cardinality sets being marked with an asterisk. Specifically, the proposed model predicts paths that are shorter than the second-best by 0.02 and 0.06 when n = 15 and n = 20 respectively, exhibiting increased ability to generalize to unseen path node counts, even for sets of twice the cardinality.

Overall results on ROCStory and PROCAT are shown in Table 1. Our proposed model, SIT, outperforms the state-ofthe-art on both datasets. PMR scores are increased by 2.68 and 2.89, and τ by 3.43 and 3.67, over second-best performances on the ROCStory and PROCAT datasets respectively.

We also compare all methods on synthetically-generated structures of length 30, following the default n-th order rulesets. Compared to the best benchmark, our method offers an improvement of 2.32 (PMR) and 1.89 (τ). Table 3 presents the average percentage of predicted sequences that adhered to the rules underlying two context sensitive and one context free grammar, in the form of the Dyck language [Yu et al., 2019] consisting of n ∈ [2, 100] pairs of {} and (). The simplest grammar (an bn cn , n sampled uniformly from [1..100]) was able to be fully learned by most models using the enhanced pointer network as their permutation module, our proposed method among them. With regards to the more challenging an bk cnk context-sensitive grammar (n and k sampled uniformly from [1..25]) our method outperformed the second best by 1.24 % and by 1.56 % in the case of the Dyck language. Performance in terms of learning higher-order interaction rulesets is discussed in Section 3.4.

The main improvement of the proposed model architecture stems from the ability to relate the input elements to the set in its entirety, supplementing its capacity to encode higher-order interactions in the set representation, fromwhich the first hidden state of the permutation decoder is initialized. Compared with the baseline set encoders and full set-to-sequence models, the combination of the SIT set encoder and the enhanced pointer network consistently outperforms on a plethora of tasks where the target permutation depends on the input set's composition. Additionally, in order to explicitly confirm the proposed model's ability to learn higher-than-pairwise interactions between set elements in a single layer of SIT transformation, we designed a comparative experiment on the expanded synthetic structure prediction task. Results are presented in Table 4 and discussed in the next section.

3.4 Relational Reasoning Study

To empirically compare a model's capacity to learn n-th order interactions between set elements, we expanded the synthetic structure library provided by Jurewicz and Derczynski (2021a). This task requires ordering a set of tokens into a target structure, defined by customizable rulesets which depend on n-th order relations between input set elements. Each benchmark model consisted of 5 stacked layers of its respective attention transformations (except the Set Transformer, which was tested with both 4 and 5 layers). The hypothesis is that even with learning only pairwise interactions per-layer, benchmark models should be able to learn the highest-order ruleset (5th). By comparison, the tested versions of our proposed set-to-sequence model had 2, 3 and 4 attention-based encoding layers, requiring them to learn higher-order interactions in a single SIT layer. Further details are given in the supplementary material.

As shown in Table 4, the 3-layer version of the proposed model has the highest score on the 3rd-order ruleset by a margin of 4.07. However, even the 2-layer model outperforms all but two benchmark methods, each consisting of 5 stacked layers. On the 4th-order ruleset the best result is obtained through the proposed model's 4-layer version; second-best performance is attained by the 3-layer model, showing its ability to learn higher-than-pairwise interactions in a single SIT layer. Similarly, on the 5th-order ruleset (n = 5 in the table) the 4-layer model outperforms the best benchmark by 3.07 percentage points, with a 5-layer Set Transformer followed by the same permutation component as our proposed set-to-sequence model achieving second best performance. As an ablation study, we include results for a 4-layer Set Transformer, the only difference between it and our 4-layer set-to-sequence model being the SIT matrix augmentation. Its presence accounts for a 4.36 performance increase.

4 RelatedWork

The earliest neural set-to-sequence model was proposed by Hopfield and Tank (1985). It introduced a constrained version of the set-to-sequence problem, in which the input set must be of a fixed cardinality. However, sets of varying cardinality require different neural architectures, making it limited in application. Kernels between sets have also been proposed to allow kernel methods, such as Support Vector Machines, to tackle set-input problems [Lyu, 2005]. These require a twostep approach where representation learning is separate.

Set-input NN methods, such as DeepSets [Zaheer et al., 2017], obtain a permutation equivariant representation of the input set by encoding each element in an identical and independent way and summing them into a permutation invariant representation. The main limitation is an inability to explicitly model higher order interactions, which are lost during pooling. This is addressed by the Set Transformer [Lee and Lee, 2019], using self-attention; the AttSets [Yang and Wang, 2020], though weighted pooling with learned, feature-specific weights; and RepSet [Skianis and Konstantinos, 2020], by solving a series of network flow problems.

Complete set-to-sequence NN models include the Pointer Network [Vinyals et al., 2015], later improved by its authors to obtain a permutation invariant representation of the set [Vinyals et al., 2016], ATTOrderNet [Cui et al., 2018] and the Enhanced Pointer Network, which requires the model to continuously make pairwise order predictions, resulting in increased sequence coherence [Yin et al., 2020]. There also exists a plethora of other NN order prediction methods, ranging from listwise ranking approaches [Ai et al., 2018] in the field of information retrieval (which, however, presume an existence of a query for which the relevant order is predicted), to permutation matrix approaches [Zhang et al., 2019], which do not handle sets of varying cardinalities. Finally, reinforcement learning has also been applied to combinatorial optimization problems in conjunction with pointer-based attention [Kool et al., 2018]. However, many interesting tasks, such as sentence ordering and complex structure prediction, do not provide the kind of fine-grained signal ideally required by these methods.

5 Conclusion

In this paper, we propose a novel set encoding method for permutation learning and structure prediction. The Set Interdependence Transformer (SIT) is capable of effectively learning higher-order interactions between input set elements, which is crucial for tasks requiring a degree of relational reasoning over large sets of varying cardinality. Our method can take input sets of any size and generalizes well to unseen lengths of the target output sequence. It is easily modularized and can be combined with an attention-based permutation decoder to form a complete set-to-sequence model. Experiments show that this architecture achieves state-of-the-art results on combinatorial optimization challenges, on NLP tasks such as sentence ordering, and in the novel domain of complex structure prediction in the context of product catalogs.

Acknowledgements

This work was partly supported by an Innovation Fund Denmark research grant (number 9065-00017B) and Tjek A/S.