HORSE: Hierarchical Representation for Large-Scale Neural Subset Selection

Authors: binghui  xie
yixuan  wang
yongqiang  chen
kaiwen  zhou
yu  li
wei  meng
james  cheng
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:52:30

Abstract: Subset selection tasks, such as anomaly detection and compound selection in AI-assisted drug discovery, are crucial for a wide range of applications. Learning subset-valued functions with neural networks has achieved great success by incorporating permutation invariance symmetry into the architecture. However, existing neural set architectures often struggle to either capture comprehensive information from the superset or address complex interactions within the input. Additionally, they often fail to perform in scenarios where superset sizes surpass available memory capacity. To address these challenges, we introduce the novel concept of the Identity Property, which requires models to integrate information from the originating set, resulting in the development of neural networks that excel at performing effective subset selection from large supersets. Moreover, we present the Hierarchical Representation of Neural Subset Selection (HORSE), an attention-based method that learns complex interactions and retains information from both the input set and the optimal subset supervision signal. Specifically, HORSE enables the partitioning of the input ground set into manageable chunks that can be processed independently and then aggregated, ensuring consistent outcomes across different partitions. Through extensive experimentation, we demonstrate that HORSE significantly enhances neural subset selection performance by capturing more complex information and surpasses state-of-the-art methods in handling large-scale inputs by a margin of up to 20%.

Main paper: 
 introduction: 
 Set-valued functions are of great importance to a wide range of real-world applications. For example, anomaly detection aims to identify a set of outliers from a larger dataset that could be users or financial transactions [Zhang et al., 2020] . Another example is the recommender system, where the objective is to identify a set of products that better satisfy customer preferences [Ou et al., 2022] . In these scenarios, there is a need for implicitly learning a set function [Rezatofighi et al., 2017 , Zaheer et al., 2017] that quantifies the usefulness of a given subset of the inputs. The set function assigns a utility value to each subset, and the subset with the highest utility corresponds to the most desired output.To illustrate the concept, let us consider the task of a recommender system. In this task, we aim to recommend a subset of items S from a larger item pool V , denoted as S ∈ V , that maximizes the utility of S with respect to the satisfaction of the customers. The utility can be captured by a parameterized utility function, denoted as F θ (S; V ), and our goal is to optimize the following criteria:S * = arg max S∈2 V F θ (S; V ).(1)One straightforward method [Balcan and Harvey, 2018] involves explicitly modeling the utility by learning the function U = F θ (S; V ) using supervised data. This data consists of pairs {(S i , V i ), U i } N i=1 , where U i represents the actual utility value of subset S i given the respective item pool V i . However, this training approach becomes challenging to implement due to the significant number of supervision signals required, which can be expensive and time-consuming to acquire [Ou et al., 2022] . To overcome this limitation, an alternative approach is to tackle Eq. 1 using an implicit learning method from a probabilistic perspective [Tschiatschek et al., 2018] . This approach requires utilizing data in the form of {(V i , S * i )} N i=1 , where S * i represents the optimal subset corresponding to V i . The objective is to estimate the parameters θ such that Eq. 1 holds for all possible (V i , S i ). In practical training, with limited data sampled from the underlying distribution P (S, V ), the empirical log-likelihood N i=1 [log p θ (S * |V )] is maximized for all the data pairs {S, V }, where p θ (S|V ) is proportional to F θ (S; V ) for all S ∈ 2 V [Ou et al., 2022 , Xie et al., 2024] . Additional details are available in Appendix D.3.The crux of the matter lies in determining the structure of neural networks that can effectively model F θ (S, V ) throughout the entire process. One commonly employed approach in the literature is to utilize an encoder to generate feature vectors for each element in V . These vectors are then inputted into DeepSets [Zaheer et al., 2017] , along with the corresponding supervised subset S, in order to learn the permutation invariant set function F (S). However, this methodology may neglect the interaction between S and V , leading to a reduction in the expressive capacity of the models. In the previous study, Xie et al. [2024] suggest incorporating the sum-pooling representation of V into S to enhance the performance. Yet, the simple integration in Xie et al. [2024] , Wang et al. [2024b] limits its capacity to effectively model interactions among elements or subsets within these sets. Furthermore, the approach struggles with high-cardinality sets V , as encoding the entire set into memory may not be feasible [Bruno et al., 2021] .To address these problems, and inspired by [Willette et al., 2023] , we introduce the notion of the Identity Property, a desirable concept for the effective functioning of the model F (S, V ). Identity Property requires F (S, V ) to accurately reflect which set V the information S originates from. In order to capture the interplay between S and V by adhering to the Identity Property, we propose a subset-based attentive set encoder. Additionally, this encoder facilitates the division of a large set V into smaller and manageable subsets. These subsets can be processed independently and later aggregated, ensuring no loss of information from V . Hence, our approach is able to efficiently handle large-scale subset representation learning. As depicted in Figure 1 , our method is capable of modeling more complex information and managing large-scale inputs more effectively than the two state-of-the-art approaches in the field of Neural Subset Selection tasks, EquiVSet [Ou et al., 2022] and INSET [Xie et al., 2024] .In this work, we make several contributions to the field of neural subset selection, which can be summarized as follows:• We introduce and rigorously define a critical concept termed as the Identity Property for neural subset selection. This property requires that models can reliably determine the source set V from which the information of the subset S is derived, which is a crucial requirement for neural subset selection tasks.• To adhere to the Identity Property and model complex interaction, we present a subset-based attention mechanism. This mechanism is crafted to learn the Hierarchical Representation of Neural Subset Selection, denoted as HORSE. Our theoretical analysis confirms that HORSE not only upholds the Identity Property but also maintains Permutation Invariance.• Through extensive empirical research, we validate the effectiveness of HORSE. Our experiments across a variety of datasets demonstrate the consistently superior performance of HORSE. Additionally, we specifically explore HORSE's capabilities in large set environments, further showcasing its practical applicability and efficiency compared with the baselines. 
related work: 
 2.1 Set Encoding.The exploration of network architectures tailored for set-structured inputs has become a vibrant area of research in recent years. A number of key studies [Ravanbakhsh et al., 2017 , Edwards and Storkey, 2017 , Zaheer et al., 2017 , Qi et al., 2017 , Horn et al., 2020 , Bloem-Reddy and Teh, 2020 , Wang et al., 2023] have laid the groundwork in this domain, primarily focusing on creating models that are permutation equivariant using conventional feed-forward neural networks. These foundational models have been successful in universally approximating continuous permutation-invariant functions, primarily through the application of set-pooling layers to aggregate information across different elements of a set regardless of their order.However, these methodologies have primarily concentrated on learning representations at the aggregate set level, paying less attention to more nuanced interactions occurring at elements. Recognizing this gap, more recent research efforts have aimed at introducing more sophisticated interaction modeling within invariant set functions for various applications. A notable example is the work by [Lee et al., 2019b] , which incorporates self-attention mechanisms to facilitate the processing of elements within sets, thereby effectively capturing element-wise interactions. Moreover, the concept of Janossy pooling, proposed by [Murphy et al., 2018] , introduced a novel approach to incorporate higher-order interactions within the pooling process. Since then, subsequent studies have built upon this advancement, leading to further refinements and innovations in the field, e.g., [Kim, 2021 , Li et al., 2020 , Bruno et al., 2021 , Willette et al., 2023] .The existing literature primarily concentrates on processing entire input sets, often overlooking the information provided by the sub-levels. Addressing this oversight, Maron et al. [2020] introduced an innovative approach that integrates the symmetry of elements to generate representations of an input set. This methodology was further expanded into a broader context by Wang et al. [2020] . Moreover, Bevilacqua et al. [2022] proposed a novel framework aimed at enhancing graph representations by including whole-graph representations to encode each subgraph. Along similar lines, Xie et al. [2024] developed an information aggregation module designed to learn F (S, V ) effectively.Table 1 : Properties of Various Methods: "Attn" indicates the use of the attention mechanism, "V" signifies the explicit utilization of information from V , and 'Large-scale' denotes the capability of the methods to generalize effortlessly to large-scale settings.Model Attn V Large-Scale DeepSets [Zaheer et al., 2017 ] ✗ ✗ ✓ Set Transformer [Lee et al., 2019a] ✓ ✓ ✗ EquiVSet [Ou et al., 2022] ✗ ✗ ✓ INSET [Xie et al., 2024] ✗ ✓ ✗Despite these advancements, a gap remains in the current research landscape. These methods tend to overlook more complex interactions between elements or subsets within sets. Furthermore, they often fall short in scenarios where the input set has a significantly large cardinality, indicating a need for more scalable and interaction-sensitive approaches in set processing. In Table 1, we compare our proposed method with importance baselines commonly used in subset selection tasks. Specifically, DeepSets can handle large-scale settings but may lose complex information due to its simple pooling-based structure. Set Transformer excels at modeling complex information within sets but faces challenges with large input set sizes. Methods tailored for subset selection tasks, like EquiVSet and INSET, may struggle with learning intricate interactions and often overlook large-scale settings. HORSE is designed to address these drawbacks.Recent work has focused on extracting subsets from training datasets to decrease cost and improve effectiveness [Wei et al., 2015 , Mirzasoleiman et al., 2020 , Yang et al., 2023] . This research also highlights the importance of modeling the relationship between the original dataset and its subsets. Unlike neural subset selection, these core subsets are unlabeled, and typically, more data in the core subset enhances the performance of the models. Our approach differs in that our optimal subset is labeled within the training set, and its size is constrained.3 Method 
preliminaries: 
 In this paper, we concentrate on the development of neural networks for the purpose of modeling the hierarchical set function F (S, V ), a critical component for tasks involving sets, such as neural subset selection. For every ground set V , assumed to consist of n elements represented as x i , that is, V = {x 1 , x 2 , ..., x n }, each element x i belonging to X is characterized by a d-dimensional tensor. Typically, the ground set V can be conceptualized as an assembly of multiple disjoint subsets, explicitlyV = S 1 ∪ S 2 ∪ • • • ∪ S m, where S i ∩ S j = ∅ for i ̸ = j and each S i is a subset in R ni×d . In this context, n i denotes the number of elements in subset S i . Generally, S ⊆ V acts as a supervisory signal in the form of a mask over V to indicate the elements to be selected. For the sake of clarity, we define S as the concrete subset derived from this mask.In the context of neural subset selection, the task entails the encoding of subsets S i into representative vectors to forecast the associated function value Y ∈ Y. Traditional approaches, such as those documented in Zaheer et al. [2017] and Ou et al. [2022] , involve directly selecting S i based on the encoding embeddings of all elements within V , subsequently feeding S i into feed-forward networks.Nonetheless, these methods model the function F (S i , V ) solely based on the explicit subsets S i , potentially leading to suboptimal results due to the omission of the broader context provided by the ground set V . This section introduces a novel attention-based method for encoding subset representations, which distinctively incorporates information from the entire input set V , thereby enhancing performance.To effectively model F (S, V ), Xie et al. [2024] have proposed to combine the representations of V and S. In practice, this involves utilizing two DeepSets architectures, as proposed by Zaheer et al. [2017] , to independently process S and V before merging their outputs, as presented by Figure 1 .Given that set pooling operations process each element independently, certain information about the interactions among elements is inevitably lost. This omission can render some problems more challenging than necessary. To address this issue and facilitate the learning of complex interactions within sets, we introduce the following principles: Property 3.1. Consider V ∈ R n×d and S ⊆ V where S ∈ R s×d , assuming that V is partitioned into a random collection of disjoint subsetsV = S 1 ∪ S 2 ∪ • • • ∪ S m .Here, m varies within the range [1, n], dependent on the chosen method of partitioning. The function F is said to satisfy the Identity Property if and only if there exist functions g and h such that F (h(S), h(V )) = F (h(S), g(h(S 1 ), . . . , h(S m ))),(2) where g serves as an aggregation function that effectively combines the encoded representations of the subsets, ensuring that F leverages both the specific subset S and the ground set V through the transformations applied by h and the aggregation by g.The method introduced by Xie et al. [2024] is notable for satisfying the Identity Property through its utilization of sum-pooling to simultaneously process all elements. However, this approach may not be practical for scenarios involving large inputs and may struggle to capture more complex information. In response to these limitations, we propose an attention-based method designed to fulfill the requirements of the Identity Property. Additionally, our interpretation of this property accommodates scenarios where S differs from the union of subsets {S 1 ∪ S 2 ∪ • • • ∪ S m }. In practice, S is often chosen to be S 1 for simplicity. This nuanced approach allows for greater flexibility and effectiveness in encoding set information, especially in complex or large-scale settings.In this section, we introduce a formulation for an attention-based set encoding function F , leveraging the concept of partitions (referred to as slots in [Bruno et al., 2021 , Willette et al., 2023] ). Given a ground set V ∈ R n×d , we randomly divide it into m subsets. For each subset, we allocate a unique embedding s i ∈ R ds . Furthermore, we establish ζ = [s 1 , . . . , s m ] T as a matrix in R m×ds . Similar to [Willette et al., 2023] , we initialize ζ by sampling m embeddings s i from a parameterized Gaussian distribution with random initialization. Following this setup, we calculate the unnormalized attention scores between ζ and V , facilitating a dynamic weighting of elements within V based on their relevance to each partition's embedding. This process aims to capture the nuanced interrelations within subsets and between elements and their corresponding subsets.EQUATIONEQUATIONIn this expression, "LN " represents Layer Normalization, and the matrices W q ∈ R ds×d h , W k ∈ R d×d h , and W v ∈ R d×d h are introduced. These matrices serve to project V and ζ into a shared dimensional space d h . Subsequently, we employ a dot product attention mechanism to assess the interactions between V and ζ. This process is governed by the specified formulation, strategically aligning elements of V with the embeddings in ζ through dimensional congruence, thus enabling a nuanced, attention-driven analysis of set elements in relation to their partitioned subsets.EQUATIONEQUATIONwhere σ denotes an element-wise activation function. Utilizing the unnormalized attention scores, denoted by Â, we proceed to define the following mapping operation:EQUATIONIn this context, h signifies a transformation function mapping from R n×d to R m×d h . The term "nl" represents a normalization operation, defined as follows:EQUATIONwhich normalizes the column of Â. Then, we can apply a pooling function (such as sum, mean, min, or max) across the columns of h(V ) and select the sigmoid function for σ, thereby establishing an attention mechanism akin to the SSE (Set Stream Embedding) method proposed by Bruno et al. [2021] . However, this approach has its limitations. Given that the attention score nl( Â) i,j is calculated independently of the other n-1 attention scores, it is not feasible for the rows of nl( Â) to form convex coefficients, unlike the softmax outputs typically observed in conventional attention mechanisms, as described by Willette et al. [2023] .To address this issue, we follow Willette et al. [2023] to aggregate information across all rows of nl( Â), thereby incorporating dependencies among different set elements into the attention mechanism. This is achieved through a specific normalization process, outlined as follows:M = diag(nl( Â)1 n ) -1 (9)where 1 n = (1, . . . , 1) ∈ R n represents a vector of ones of dimension n, and M 2 ∈ R m signifies a vector in an m-dimensional space. Subsequently, we can compute h(V ) by applying the normalization term M , as follows:EQUATION)Since S i is a subset of V based on a partition method, h V (S i ) can be derived from h(V ). Detailed steps are provided in Algorithm 1. For simplicity, we omit V and use h(S i ) going forward. This process ensures that h meets the criteria specified in Eq. 2. By constructing such an h function, we ensure that the model can recognize the input set V regardless of its partitioning, leading to the property that g(h(S 1 ), h(S 2 ), . . . , h(S m )) yields the same value for any partition of V.Furthermore, it facilitates the learning of interactions among the partitioned segments of V , essentially enabling the model to identify the characteristics of the input ground set V . Specifically, we concatenate h(S) with g(h(S 1 ), h(S 2 ), . . . , h(S m )). In practice, a Multilayer Perceptron (MLP) is utilized to process the concatenated tensor into a vector Z ∈ R do , achieving the following:EQUATION)Given that the aforementioned process delineates the entire calculation in a matrix format, which may be complex for some readers, we have taken steps to enhance comprehension. To better illustrate how our method establishes an attention map between subsets S i , we have detailed the procedural steps in the Appendix (see Algorithm 1), with a particular focus on the generation of h(S i ). This will not only clarify the operational details but also emphasize the underlying methodology and thought process involved in constructing h(S i ). 
theoretical results: 
 In the realm of machine learning, particularly within the scope of set-based tasks, a fundamental requirement is the invariance to the permutation of input set elements. This characteristic ensures that the computation or outcome of a task is unaffected by the order in which the set's elements are presented, a principle that is especially pertinent to neural subset selection tasks. To address and formalize this aspect within the context of our proposed method, we present a theorem that rigorously demonstrates the permutation invariance of our approach. Theorem 4.1. Let S n denote the set of all permutations of a given set V . Since V ∈ R n×d is represented by a matrix, let π V ∈ R n×n be a random permutation applied to V . Given that S ⊆ V represents a subset of V , the permutation π V naturally induces a corresponding permutation π S on S. Under these conditions, HORSE exhibits permutation invariance, which is defined as:EQUATIONTheorem 4.1 assures that irrespective of how the elements in the input set V are ordered, the output generated by HORSE remains consistent. This property is crucial for ensuring the reliability and applicability of our method across a wide spectrum of set-based tasks, where the inherent order of data points should not influence the task outcome. Furthermore, an illustration of the underlying concept is provided in Figure 2 . By satisfying this property, HORSE ensures that its encoding captures both the individual characteristics of subsets within V and the overarching structure of the entire set, thereby facilitating a more nuanced and comprehensive understanding of set information. The proofs is inspired by [Willette et al., 2023] and can be found in Appendix A. 
experiments: 
 In this section, we aim to demonstrate that HORSE significantly outperforms baseline models in a suite of benchmarks tailored to Neural Subset Selection tasks. Subsequently, we extend our investigation to scenarios involving large-scale input settings. Due to the page limitation, we have included additional experiments in Appendix D.Evaluations. To assess the performance of various methods, we employ the mean Jaccard coefficient (MJC) as the evaluation metric. This metric quantifies the similarity between the predicted subset S ′ and the true subset S * for each data sample (S * , V ). The Jaccard coefficient is calculated as follows:JC(S * , S ′ ) = |S * ∩S ′ | |S * ∪S ′ |, where the intersection and union operations determine the size of the overlap and the total unique elements in both sets, respectively. The MJC is then derived by averaging the Jaccard coefficient across all test set samples. Please note that all the following reported performance metrics are presented in percentages, with a default multiplication factor of 100%. Firstly, We validate the effectiveness of our models through experimental trials focused on learning set functions, using two synthetic datasets: the two-moons dataset [Pedregosa et al., 2011] with an added noise variance of σ 2 = 0.1, and and a Gaussian mixture represented as 1 2 N (µ 0 , Σ) + 1 2 N (µ 1 , Σ). Take the Gaussian mixture as an example, the data generation procedure as follows: i) Initially, we select an index, denoted as b, using a Bernoulli distribution with a probability of 1 2 . ii) Subsequently, we sample 10 points from the Gaussian distribution N (µ b , Σ) to construct the set S * . iii) Further, we sample 90 points for V \S * from the Gaussian distribution N (µ 1-b , Σ). We follow the procedure of Ou et al. [2022] to obtain 1,000 samples, subsequently divided into training, validation, and test sets. We effectively demonstrate the efficacy of our approach in mastering complex set functions. Detailed results can be found in Table 3 .The task involves recommending the most suitable subset of 30 products to a customer within a specific category. For this experiment, we utilize the dataset from the Amazon baby registry, sourced from Gillenwater et al. [2014a] . This dataset includes numerous product subsets chosen by various customers, with Amazon categorizing each item on a baby registry into specific categories such as "Bath", "Health" and "Feeding". Detailed information can be found in Appendix D.Table 2 presents the performance of all models across different categories. Notably, out of the twelve cases evaluated, HORSE outperforms other models in 11 of them. Even in the Diaper category, our method achieves results that are comparable to INSET, which is noteworthy. These significant improvements highlight the effectiveness and superiority of HORSE. Specifically, while EquiVSet and INSET struggle to surpass classical neural subset selection baselines in the Safety, Car Seats, and Furniture categories, HORSE consistently outperforms all baselines in a notable manner.In drug discovery, the screening of compounds with diverse biological activities and favorable ADME (absorption, distribution, metabolism, and excretion) properties is a critical step [Li et al., 2021 , Ji et al., 2022 , Gimeno et al., 2019] . Virtual screening typically involves a sequential filtering process that employs multiple essential filters. However, neural networks encounter challenges when learning the complete screening process. This difficulty arises from the absence of intermediate supervision signals, which can be costly or impossible to obtain due to pharmaceutical protection policies. Therefore, we implement a single filter, namely, high bioactivity, to obtain the optimal subset of compound selection, following the methodology in [Ou et al., 2022] . Our experiments are conducted on two datasets: PDBBind [Liu et al., 2015a] and BindingDB [Liu et al., 2007] . To be more practical, we further enhance our approach with the inclusion of two filters: the high bioactivity filter and the diversity filter. This extended analysis is denoted as PDBBind-2 and BindingDB-2, representing the two-stage filtering process, for a more practical perspective.Table 4 demonstrates that our method outperforms the baselines and significantly surpasses random guessing, especially on the BindingDB-2 and PDBBind-2 datasets. However, the improvement on PDBBind and BindingDB is less significant. This marginal enhancement is due to the informative structural characteristics of complexes (the elements within a set), which inherently provide substantial information for this task. Consequently, the model can effectively predict the activity values of complexes even without explicitly considering the interactions between the optimal subset and its complement in single filter scenarios. Nonetheless, our method still achieves superior results compared to other methods, confirming its effectiveness.Set encoding mechanisms, as introduced by Zaheer et al. [2017] and further explored by Lee et al. [2019a] , have fundamentally shifted the way neural networks perceive and process sets by emphasizing the importance of permutation invariance and the ability to handle variable-sized inputs. These models are designed to learn from the entire set in a single gradient step, ensuring that the learned representations encapsulate the holistic properties of the set. However, this approach encounters practical limitations when dealing with large-scale sets, where processing the entire set in a single step becomes computationally infeasible due to memory constraints or the sheer volume of data.To circumvent these challenges, an effective strategy involves training models on partitions of the set, sampled dynamically at each iteration of the optimization process [Lee et al., 2019a , Wang et al., 2024a] . This method allows for manageable subsets to be used for training, significantly reducing the computational load. However, this method will lose information [Bruno et al., 2021 , Willette et al., 2023] since it does not process all the elements from V. Therefore, we instead partition the set elements into mini-batches, independently encode each batch, and aggregate them to obtain a single set encoding. By applying this methodology across both baseline models and HORSE in scenarios characterized by large-scale input sets, we can highlight the efficiency and scalability of our proposed solution. We conducted experiments on the Two-Moons and Gaussian-Mixture datasets. To ensure consistency, we set the size of the optimal subset S * to be 10. Subsequently, we varied the size of the input ground set V within the range of {200, 400, 600, 800, 1000}. Notably, the memory capacity of the GeForce RTX 3090 is insufficient when the size reaches 600. The ground set was divided into 5 disjoint partitions, with each partition containing one-fifth of the elements in V . For the purpose of comparison, we selected INSER and Set Transformer as baselines alongside our proposed method, HORSE. INSET demonstrated the best performance among baselines, while Set Transformer is an alternative method that incorporates an attention mechanism. The results obtained from these experiments are presented in the left two subfigures of Figure 3 . It is evident that HORSE outperforms both INSER and Set Transformer by a significant margin, demonstrating its superior effectiveness in handling large-scale sets.Furthermore, to further enhance the practicality of our approach and investigate the potential impact of the partition numbers on the results, we conducted additional experiments on the BindingDB dataset. In this experiment, we set the size of the optimal subset S * to be 15, while the size of the ground set remained fixed at 1000. We partitioned the ground set into a range of 2 to 50 partitions. The results of these experiments are presented in the right subfigure of Figure 3 . Remarkably, it becomes evident that HORSE exhibits remarkable robustness with respect to the number of partitions considered. Regardless of the specific partitioning scheme employed, HORSE consistently delivers exceptional performance, which is more robust than our baselines. 
conclusion: 
 In this paper, we have addressed the limitations observed in existing methods for neural subset selection tasks. These methods often struggle to effectively model complex information and lack scalability when dealing with large-scale inputs. To overcome these challenges, we propose an innovative and scalable approach called HORSE, which leverages the power of the attention mechanism. Theoretically, we establish that HORSE satisfies the Identity Property and Permutation Invariance, ensuring its soundness and effectiveness. Empirically, we thoroughly evaluate the performance of HORSE against various baselines in both standard and large-scale settings.Limitation and Future Work. Our theoretical and empirical results demonstrate how the attention mechanism can enhance models for neural subset selection tasks in both standard and large-scale settings. However, in large-scale scenarios, our support is currently limited to a theoretical framework for partitioning the set into different groups within a synthetic distributed setting, rather than practical experimentation in a real distributed environment. Moving forward, we plan to implement and test our model in more practical, real-world scenarios to further validate its effectiveness. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Please refer to our Appendix, and we follow the settings of [Ou et al., 2022] .Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: The experiments in this paper are conducted with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: The focus and experimental finding of this paper do not relevant to societal impacts.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper does not have such a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All baselines used are properly cited. Sources of the datasets are also stated and all datasets are under CC-BY 4.0.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: This paper does not release new assets. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Proof of Theorem 4.1Proof. Consider a given set V = {x 1 , x 2 , . . . , x n } ∈ R n×d , where n represents the number of elements in V . Let S n represent the set of all permutations of V . Now, suppose π V denotes a random permutation applied to V . By utilizing this permutation π V , we can construct a permutation matrix M ∈ R n×n :M V =    -x ⊤ π V (1) - . . . -x ⊤ π V (n) -    .Assuming the use of an activation function σ that is strictly positive for each element,σ( √ d -1 • q(M V W K ) ⊤ ) = σ( √ d -1 • q(V W K ) ⊤ M ⊤ ) = σ( √ d -1 • qk ⊤ )M ⊤ = ÂM ⊤ .The normalized attention score nl( Â) can be computed using the given permutation π V , resulting inEQUATIONNow, we consider the matrix multiplication ofnl( Â)M ⊤ =    nl( Â) 1,π(1) • • • nl( Â) 1,π(n) . . . . . . . . . nl( Â) k,π(1) • • • nl( Â) k,π(n)    and M V W V =    x ⊤ π(1) W V . . . x ⊤ π(n) W V    . Since    nl( Â) 1,π(1) • • • nl( Â) 1,π(n) . . . . . . . . . nl( Â) k,π(1) • • • nl( Â) k,π(n)       x ⊤ π(1) W V . . . x ⊤ π(N ) W V    is equal to     N j=1 nl( Â) 1,π(j) x ⊤ π(j) W V . . . N j=1 nl( Â) m,π(j) x ⊤ π(j) W V     ,which can also be formulated as:EQUATIONTherefore, nl( Â)v is permutation invariant under the permutation group of S n . Sincenl( Â)1 n = n j=1 nl( Â) i,j = N j=1 nl( Â) i,π V (j) , thus diag nl( Â)1 n -1is also invariant with respect to the permutation of input V , which leads to the conclusion thath(M V ) = h(V ).Similarly, we can construct the permutation matrix M ∈ R s×s for a given S and permutatation π S , such that:M S =    -x ⊤ π S (1) - . . . -x ⊤ π S (s) -    .with the same process as Equation 13 and 14, we can have the following conclusion:F (h(S), h(V )) = F (h(π S • S), h(π V • V )).B Proof of Theorem 4.2Proof. Consider the input set V ∈ R n×d and letV = S 1 ∪ S 2 ∪ • • • ∪ S m represent a partition of V with |S i | = n i .In other words, V can be expressed as the union of all S i and each S i is disjoint from S j for i ̸ = j. Without loss of generality, we can make the assumption that.k = V W k =    S 1 W k . . . S m W k    , v = V W v =    S 1 W v . . . S m W v   where S i W k ∈ R ni×d h and S i W v ∈ R ni×d h for i = 1, 2 . . . , m. Then we can express the matrix nl( Â) as follows:nl( Â) = nl( Â(1) ) • • • nl( Â(m) ) ,whereÂ(i) = σ( √ d -1 • q(S i W k ) ⊤ ) ∈ R m×ni for i = 1, 2 . . . , m since nl( Â) i,j is independent to nl( Â) i,t for all t ̸ = j. Since h = nl( Â(1) ) • • • nl( Â(m) )    S 1 W v . . . S m W v    ,the following equality holdsEQUATIONSinceEQUATIONwhere1 ni = (1, . . . , 1) ∈ R ni .Combining Equation 15 and 16h(V ) = diag m i=1 nl( Â(i) )1 ni -1 m i=1 nl( Â(i) )S i W v . We define h 1 (S i ) = nl( Â(i) )1 ni and h 2 (S i ) = nl( Â(i) )S i W v . Moreover, h(S i ) = diag(h 1 (S i )) -1 h 2 (S i ) Now, we define a function, g({h(S 1 ), . . . , h(S m )}) := g 1 ({h 1 (S 1 ), . . . , h 1 (S m )}) • g 2 ({h 2 (S 1 ), . . . , h 2 (S m )}),whereg 1 ({h 1 (S 1 ), . . . h 1 (S m )} := diag m i=1 h 1 (S i ) -1 g 2 ({h 2 (S 1 ), . . . , h 2 (S m )}) := m i=1 h 2 (S i ).Then h(V ) = g({h(S 1 ), . . . , h(S m )}). Since the partition is arbitrary, h satisfies Property 3.1.In the main text, we present HORSE using matrix calculations, which may be challenging to comprehend. To improve understanding of how our method establishes an attention map between subsets S i , we detail the procedural steps in Algorithm 1, with a special emphasis on the generation of h(S i ). This approach is designed to elucidate the operational details and highlight the methodology involved in constructing h(S i ).Algorithm 1 HORSE. V = {S 1 , S 2 , . . . , S m } is the input set partitioned into m subsets. ξ ∈ R m×ds is the initialized embedding and g is the choice of aggregation function.1: Product Recommendation. The task involves recommending the most suitable subset of 30 products to a customer within a specific category. For this experiment, we utilize the dataset from the Amazon baby registry, sourced from Gillenwater et al. [2014a] . This dataset includes numerous product subsets chosen by various customers, with Amazon categorizing each item on a baby registry into specific categories such as "Bath", "Health" and "Feeding". Additionally, each product is represented by a 768-dimensional vector generated by a pre-trained BERT model, based on its textual description.Input: V = {S 1 , S 2 , . . . , S m }, S = S 1 , ζ ∈ R m×ds , g 2: Output: Z ∈ R do 3: Initialize ζ 4: q = LN (ζW q ) 5: for i = 1, 2, . . . , m do 6: k i = S i W k 7: v i = S i W v 8: end for 9: k = [k 1 , k 2 , . . . , k m ] T 10: v = [v 1 , v 2 , . . . , v m ] T 11: M = d -1 h • qk T 12: Â = σ(M 1 ) 13: A = nl( Â) = [A 1 , A 2 , . . . , A m ] T 14: M = diag(nl(A)1 n ) -1 = [M 1 , . . . , M m ] T 15: for i = 1, 2, . . . , m do 16: h(S i ) = M i A i S i W v 17: end for 18: Ŝi = g(h(S 1 ), . . . h(S m )) 19: Z = F (h(S), ŜiThe Amazon baby registry data [Gillenwater et al., 2014b] comprises various datasets collected from Amazon, encompassing different categories such as toys, furniture, and more. Within each category, there exist |V | sets of products that have been selected by different customers. To create a sample (S * , V ), we follow a specific procedure. Initially, we remove any subset with an optimal subset size |S * | greater than or equal to 30. The remaining subsets are then divided into training, validation, and test folds using a 1:1:1 ratio. Additionally, we randomly select an additional 30 -|S * | products from the same category to construct (S * , V ). This process allows us to create a data point (S * , V ). For comprehensive information, please refer to Table 5 in Ou et al. [2022] , which presents the statistics of the categories.Set Anomaly Detection. We tackle set anomaly detection tasks on four real-world datasets: double MNIST [Sun, 2019] , CelebA [Liu et al., 2015b] , F-MNIST [Xiao et al., 2017] , and CIFAR-10 [Krizhevsky, 2009] . Each dataset is partitioned into training, validation, and test sets, each comprising 10,000, 1,000, and 1,000 samples, respectively. In each dataset, we randomly select n images from the dataset to create the OS Oracle S * , where n can be either 2, 3, or 4. This setup aligns with the approach outlined in [Zaheer et al., 2017 , Ou et al., 2022 ].Let's take CelebA as an illustrative example. In this scenario, the goal is to identify anomalous faces solely through visual observation, without using any attribute values. The CelebA dataset consists of 202,599 face images, each annotated with 40 boolean attributes. When constructing sets, for each ground set V, we randomly select n images from the dataset to create the OS Oracle S * , ensuring that none of the selected images contain any of the two attributes. Additionally, we ensure that no individual person's face appears in both the training and test sets.Regarding the results presented in Table 6 , it is evident that our model exhibits a significant performance advantage over all the baseline methods. This substantial improvement underscores the superior capabilities of our model in addressing the given task.Compound Selection in AI-aided Drug Discovery. In drug discovery, the screening of compounds with diverse biological activities and favorable ADME (absorption, distribution, metabolism, and excretion) properties is a critical step [Li et al., 2021 , Ji et al., 2022 , Gimeno et al., 2019] . Virtual screening typically involves a sequential filtering process that employs multiple essential filters. These filters initially select diverse subsets from highly active compounds and subsequently eliminate compounds with unfavorable ADME characteristics. After passing through several filtering stages, an optimal subset of compounds is identified. However, neural networks encounter challenges when learning the complete screening process. This difficulty arises from the absence of intermediate supervision signals, which can be costly or impossible to obtain due to pharmaceutical protection policies. Consequently, models are expected to learn this intricate selection process in an end-to-end manner. In other words, models must predict S * based solely on the optimal subset supervision signals, without knowledge of the intermediate steps. Therefore, we simulate the optimal subset oracle of compound selection by applying one or two filters by uising PDBBind and BindingDB, as [Ou et al., 2022] .PDBBind offers an extensive compilation of experimentally measured binding affinity data for biomolecular complexes. We utilized the "refined" portion of the complete PDBBind dataset, which consists of 179 complexes, to construct our subsets. To create a data point (V, S * ), we randomly sampled 30 complexes from the dataset to form the ground set V . The subset S * was then generated by selecting the five most active complexes within V . We constructed separate training, validation, and test splits, comprising 1000, 100, and 100 data points, respectively.BindingDB is an openly accessible database that provides measured binding affinities for a collection of 52, 273 drug-target pairs involving small, drug-like molecules. Similar to PDBBind, we randomly selected 300 drug-target pairs from the BindingDB database to form the ground set V . From this dataset, we carefully identified the 15 most active drug-target pairs and designated them as S * . To ensure comprehensive evaluation and robust model training, we subsequently created distinct training, validation, and test sets. These comprised 1000, 100, and 100 data points respectively. Random. This represents the expected performance of a random guess, serving as a baseline to help us gauge the actual difficulty of the tasks.PGM [Tschiatschek et al., 2018] . PGM, which stands for Probabilistic Greedy Model, tackles the optimization Problem 1 by employing a differentiable extension of the greedy maximization algorithm. For a deeper understanding of this approach, please refer to the origianl paper or Appendix A in [Ou et al., 2022] .DeepSet [Zaheer et al., 2017] . Here, we employ DeepSet as a baseline model. We use DeepSet to predict the probabilities of including specific instances in S * , essentially learning an invariant permutation mapping from the power set 2 V to the interval [0, 1] of size |V |. This model was also used as a backbone for set function learning in EquiVSet. Moreover, it is suitable for subset selection tasks, as detailed in its original paper.Set Transformer [Lee et al., 2019a] . Set Transformer extends DeepSet's capabilities by integrating the self-attention mechanism. This addition allows the model to consider pairwise interactions between elements, enabling it to capture dependencies and relationships among different elements more effectively. It can be also utilized for subset selection tasks, similar to DeepSet.EquiVSet [Ou et al., 2022] . EquiVSet employs an energy-based model (EBM) to establish the set mass function, denoted as P (S|V ) from a probabilistic standpoint. Their primary objective lies in learning a distribution P (S|V ) that monotonically increases with respect to the utility function F (S, V ). It's worth noting that their framework focuses on approximating the symmetric function F (S) rather than the symmetric function F (S, V ), with DeepSet serving as the foundational component of their model to approximate the set function.INSET [Xie et al., 2024] . As discussed in the Introduction, EquiVSet faces limitations in incorporating information from the ground set V. In response to this challenge, Xie et al. [2024] present an innovative solution. They propose the generation of embeddings for V and subsequently concatenate these embeddings with the representations of S, which has been presented in Figure 1 .Among these baselines, DeepSets and Set Transformer are two crucial model structures widely used in set-based tasks, including subset selection tasks. On the other hand, PGM, EquiVSet, and INSET are methods specifically designed for neural subset selection tasks. Notably, both INSET and HORSE are implemented based on the EquiVSet framework, yet they significantly outperform EquiVSet.Our method, HORSE, is applicable for learning F (S, V ) across a range of tasks. In this paper, we primarily utilize the framework established in [Ou et al., 2022] and modify it with our approach to model F (S, V ). The optimization objective aims to solve Equation 1 by employing an implicit learning strategy based on probabilistic reasoning. This approach can be formulated concisely as follows: arg max θ E P(V,S) [log p θ (S * |V )] s. t. p θ (S|V ) ∝ F θ (S; V ), ∀S ∈ 2 V ,p θ (S|V ) = exp(F θ (S; V )) Z , Z := S ′ ⊆V exp(F θ (S ′ ; V )),In practice, we approximate the Energy-Based Model (EBM) through a variational approximation.Due to the scope of this paper, we omit the detailed explanation for the sake of simplicity. We kindly invite readers to refer to [Ou et al., 2022] for further information on this topic.In this subsection, we present the implementation details of HORSE. The setup closely follows that of [Ou et al., 2022] . The proposed models are trained using the Adam optimizer [Kingma and Ba, 2014] with a fixed learning rate of 1e-4 and a weight decay rate of 1e-5. To accommodate different model sizes across various datasets, we select the batch size from the set {4, 8, 16, 32, 64, 128}. Importantly, we choose the largest batch size that allows efficient training on a single GeForce RTX 3090 GPU.To enhance training efficiency and mitigate overfitting, we utilize an early stopping strategy for both the baselines and our proposed models. In this strategy, if there is no performance improvement over 10 consecutive epochs, we terminate the training process prematurely. For each dataset, the maximum number of epochs allowed for training is set to 80. At the end of each epoch, we assess the model's performance on the validation set and save the model with the best performance. Finally, we evaluate the saved models on the test set to determine their performance.In order to consider the effect of randomness and ensure the reliability of the findings, we conduct all experiments five times using different random seeds. The average performance metrics, along with their corresponding standard deviations, are reported as the final performance measures. This approach allows for a comprehensive evaluation of the models' performance while accounting for the influence of random variations.To further investigate the robustness of INSET, we have conducted ablation studies specifically focusing on the Monte-Carlo sample numbers k for each input pair (S * , V ). In our framework, the model θ is trained to accurately predict the optimal subset S * from a given ground set V . Following the Energy-Based Method (EBM) proposed in [Ou et al., 2022] , we incorporate a necessary hyperparameter. During the training process, we sample k subsets from V in order to optimize the model parameters θ, thereby maximizing the conditional probability distribution p θ (S * |V ) among all pairs of (S, V ) for a given V .To • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 


Decision: **Summary of the paper:**
The paper deals with the supervised learning of deep networks for subset selection. The paper identifies expressivity issues with the recent best-performing methods, particularly, in modeling the interaction within and between the original set elements and the optimal subset elements. It then defines a desirable formal property on set representations that their proposed model, HORSE, satisfies. HORSE uses an attention mechanism to model the desired interactions and furthermore uses a random partitioning of the full set for a more efficient handling of large sets. Several experiments including, two on synthetic data, few on proteins binding, and one on a product recommendation task, are conducted exhibiting consistent improvements over the baselines. 

**Summary of the reviews:**
The reviewers appreciate the importance of the problem, the relevance of “identity property” for dealing with large input sets and the simplicity of the attention mechanism in modeling the desired interactions. They further found the experiments substantiating the claim. Most reviewers found the presentation well done. 

On the other hand, the reviewers asked for a broader set of benchmarks and baselines, particularly large-scale sets. They further brought up the importance of ablation study for different components and asked for a discussion on the role of partitioning. Errors and complicatedness of the notation were also raised. 

**Summary of the rebuttal and discussions:**
The authors provided (and referred to) the ablation study of the various pooling functions and the different partitioning scenarios, and further discussed the role and the effect of partitioning. They further compared with more advanced baselines equipped with attention and introduced a couple of new benchmarks during the rebuttal. They acknowledge the possible improvements in the notation.

**Consolidation report:**
The AC agrees with the reviewers' positive points on the relevance and significance of the problem and the proposed approach. The AC believes the concerns are mitigated with the rebuttal regarding the benchmarks, baselines, the ablation studies and the choice of partitioning. The notation can be improved but is relatively minor and can be done in a camera-ready submission.

**Recommendation:**
The AC agrees with the unanimous recommendation by the reviewers and suggests acceptance.

Reviews: review: paper_summary: paper_summary: The paper introduces HORSE, a method for neural subset selection. It addresses the limitations of existing methods by introducing the concept of Identity Property and utilizing an attention-based mechanism to capture complex interactions within the input set. HORSE demonstrates superior performance on various tasks, including product recommendation and compound selection.
strength_weakness: strength_weakness: - While the paper includes several experiments, a more extensive evaluation with a wider range of datasets and tasks would strengthen the claims. It would be beneficial to compare the proposed method to a broader set of baselines, including recently proposed methods for set-based tasks.
- A more detailed ablation study to isolate the contributions of different components of the proposed method (e.g., attention mechanism, partitioning scheme) would provide deeper insights into the model's effectiveness.
- The paper could provide more details about the computational complexity of the proposed method and compare it to other approaches. This information would be valuable for practical applications.
- Although the paper mentions handling large-scale inputs, a more thorough evaluation on extremely large datasets would demonstrate the scalability of the proposed method.
- While the paper provides theoretical analysis of the Identity Property and permutation invariance, a deeper theoretical investigation into the properties of the attention mechanism and its impact on the model's performance could provide additional insights
questions: questions: - How does HORSE compare to other state-of-the-art methods for subset selection, such as those based on graph neural networks or deep sets with attention mechanisms?
- Could the authors provide a more detailed ablation study to analyze the impact of different components of the proposed method, such as the attention mechanism, the partitioning strategy?
- What is the computational complexity of HORSE compared to other methods, especially for large-scale datasets? How does the choice of partitioning strategy affect computational efficiency?
- Can the authors provide further theoretical insights into the properties of the attention mechanism and its relationship to the Identity Property?
- How does HORSE handle imbalanced datasets or noisy data? What are the potential challenges and limitations of applying HORSE to real-world applications?
limitations: limitations: # Limitations already addressed
- The paper acknowledges limitations like sensitivity to hyperparameter tuning (through potential need for careful hyperparameter search) and potential overfitting (although not extensively discussed).

# Limitations NOT addressed

- The impact of noise in the data is not explicitly discussed.
- The paper mentions handling large-scale inputs but lacks thorough evaluation on extremely large datasets, leaving scalability concerns unaddressed.
paper_summary: paper_summary: Paper is interested in subset selection. Given ground-set $V$, how to choose a subset $S \subseteq V$ that maximizes utility function $F(S, V)$. The core contribution of the method is in (randomly) partitioning $V$ into multiple subsets $S_1, S_2, \dots, S_m$, computing some matrices (based on this random partitioning), and using these matrices to compute information from $S$ and each of $S_i$, plugging into neural network, that should model $F$.
strength_weakness: strength_weakness: In my opinion, there are two main weaknesses in the paper:

1. The "subset formulation" seems to be introduced by force (I don't see a good fit here).
2. The notation is **unnecessarily** complex.

## Subset Formulation

Usually, for subset-selection (IMO), there are usually some conditions along:

1. Cardinality constraint. E.g., $|S| \le k$. In such applications, enlarging $S$ is usually good -- almost every item gives some positive utility, however, some give utility more than others  -- e.g., submodular.

2. The presence of some items makes other items less-needed. E.g., if a customer order contains diapers from company A then they get less if buying diapers from company B.

However, the experiments (in my understanding) seem to be classification-like or detection-like -- they are more fit for multi-class classification. This becomes ecident as the qualitative measure of Table 1 experiments is Jaccard similarity (it strictly grows when including correct elements and degrades when including incorrect elements).

NOTE: this does *not* degrade any of the modeling tricks they have invented in the paper, I am only complaining about advertising the wrong thing (per my understanding, and please correct my understanding if I am wrong).

## Unnecessary complexity

The reason for notation is conciseness and clarity. I think one should spend a lot of time on carefully-choosing notation to carry the message clearly. In cases where notation is complicating things, it is defeating its own purpose. I will point some specifics here.

* The calculation of $h(V)$ seems to be an algorithm -- specifically, detailed in section 3.3  -- roughly, compose 
matrix $\zeta$ based on random partitioning of rows of $V$, from which, matrices (q, v, k) are calculated which finally produce the output of $h(V)$. The notation in Equation (11) does not use $h(V)$ but instead uses $h(S)$ and $h(S_i)$. In that case, would the $\zeta$ be recalculated for every invocation of $h$ or is it computed from the ground-set $V$? If the latter, it might make sense to denote $h_V(S)$.

* matrix $A$ is double (L1) normalized: Eq.8 normalizes columns to sum to `1` and Eq.9 (and its application in Eq.10) normalizes rows to sum to `1` -- but why have completely different notations to the normalization?

* Note: the wording "utilizing unnormalized" scores made me think that it might be a good idea to do "sum" rather than "mean" (i.e., unnormalized pooling). However, the matrix $A$ is indeed double normalized.

The letter $n$ denotes number of nodes, yet $nl$ stands for normalization. I recommend you use \textrm{nl} for the function call.

* Set are permutation invariant by construction. This makes Theorem 4.1 unnecessary. However, I understand your intent: you are referring to the matrix version and not the set. Please correct the writing.

* Please be clear on the permutation's $\pi$. It seems that it is a binary matrix where every row and every column has exactly single 1.
questions: questions: Q1: What is the objective function that you used for training? While you are focusing on the model, it is especially important to note it, as the method is advertised as "subset selection" from the start.

Q2: How is row $i$ of $\zeta$ calculated from points clustered in $S_i$. Is it the average? This crucial information is missing.

Q3: Do h(V) and h(S) share the same $\zeta$?
limitations: limitations: Authors have added this as part of their conclusion
paper_summary: paper_summary: The paper studies an interesting and important problem that is general in machine learning, focusing on motivations (the coexistence of the set-to-set interaction and the large-scale setting issues) that are crucial in this field. While people may like to see the real-world practical analysis of this method, the theoretical analysis and experiments show that the proposed method is effective in the field of recommendation and anomaly detection. Also, the paper is well written and easy to understand.
strength_weakness: strength_weakness: 1. In Section 3.1, I cannot find citations for the problem setting. At first glance, it seems the problem setting does not reflect real cases. Why is $W$ split into $$ m"disjoint" subsets a given assumption? If the subsets are provided as candidates, what if they limit the results to suboptimal solutions or introduce bias?
2. More discussion is needed regarding the impacts of $n_i$ and $m$. Although they are given as fixed inputs, it is possible to merge some subsets, which $m$ decreases and $n_i$ increases, before running the algorithm.
3. The proposed method has been tested under anomaly detection and recommendation tasks. However, it has not been compared with popular approaches in these fields that do not use subset-valued functions. The method has only been compared with others within the set-function learning domain. While this is necessary, it might be a minor weakness given that the author has stated their model is theoretical rather than practical for real-world scenarios.
questions: questions: 1. Can the proposed framework be used in the image pixel grouping task, such as image segmentation? Does predicting a subset of pixels from the entire image pixels fit in the subset-valued function?
2. Why splitting V into "disjoint" subsets is necessary? Why not just randomly sampling subsets from V such that different subsets may overlap?
limitations: limitations: The authors addressed the limitations in the conclusion section.
paper_summary: paper_summary: The authors propose a novel approach for architecture for the task of subset selection that is based on developing layers (HORSE) which satisfy a notion of an identity property (leveraging learning representations that connect a set V with its subset S) in a permutationally-invariant way, leveraging key-query-based attention mechanisms. Extensive experiments and theoretical justification validate the proposed construction.
strength_weakness: strength_weakness: Unless I have missed it, could you let me know what is the effect of the random partitioning of the set V? Do you have ablations over that? How sensitive is your method to the partitioning?

nit: In Figure 2's caption you have a typo: if the ground set <- of the ground set.

nit: In Table 4 you may want to use '' instead of " for a closing quotation mark.
questions: questions: How do you define an "effective" aggregation function in Property 3.1? In principle, by this definition, $g$ can be a trivial function such as $g \equiv 0$, which will ignore all the subsets $S_1,\dots,Sm$.
limitations: limitations: Yes.

Comments: title: Thanks for the rebuttal.
Some Empirical Studies Complementing Our Rebuttal
comment: Thank you for your response and for dedicating your time to thoroughly review our paper and rebuttal.
The authors have addressed my comments, I have read the discussions, and maintain my score.
Dear Reviewer aVVy,

We sincerely appreciate your valuable and helpful suggestions. We are wondering whether you have any more suggestions or questions after our response. Specifically, do you have any questions on our training objective or do you have any reference about Corset Selection to help use differentiate Neural Subset Selection with Corset Selection. Moreover, if you find our response satisfactory, we kindly invite you to consider the possibility of improving your rating.

Sincerely,
Authors
***Ablation Studies on the pooling methods***

Our method employs pooling methods as our aggregation function. Following the common selection in set-related works, such as [1, 2], we use mean pooling, as it typically yields better performance. In our studies, we also find that these four methods exhibit similar overall performance, with sum pooling and mean pooling being slightly superior.

|   | mean pooling |  sum pooling          | min pooling          | max pooling          
| ----------- | -------------- | -------------- | -------------- | -------------- | 
| PDBBind      | **0.941 ± 0.007** | 0.939±0.010 | 0.934±0.013 | 0.935±0.012 | 
| BindingDB   | 0.742 ± 0.011 | **0.745±0.012** | 0.739±0.015 | 0.736±0.019 | 
| PDBBind-2      | 0.432 ± 0.006 | **0.435±0.008** | 0.423±0.009 | 0.425±0.012 |
| BindingDB-2      | **0.213 ± 0.009** | 0.210±0.007 | 0.211±0.011 | 0.206±0.014 |

***How does HORSE handle imbalanced datasets or noisy data?***

In neural subset selection tasks, the imbalance may arise from differences in the size of the optimal subsets and the size of the ground set. To provide empirical evidence, we conduct additional experiments that showcase INSET's consistent superiority over the baselines, even in scenarios with imbalanced ground set sizes. Specifically, we train the model on the two-moons dataset using fixed optimal subset sizes of 12 and evaluate its performance on various ground set sizes ranging from 200 to 1000.

|   | 200 |  400        | 600 | 800 | 1000|         
| ----------- | -------------- | -------------- | -------------- | -------------- | -------------- | 
| EquiVSet      | 0.486 ± 0.002 | 0.453 ± 0.003 | 0.432 ± 0.002 | 0.425 ± 0.005 | 0.396 ± 0.003  |
| INSET | 0.494 ± 0.002 | 0.478 ± 0.005|  0.447 ± 0.003 | 0.434 ± 0.002 | 0.425 ± 0.002 |
|HORSE | **0.527 ± 0.004** | **0.503 ± 0.006** |  **0.476 ± 0.005** | **0.466 ± 0.002** | **0.453 ± 0.003** |

For imbalanced datasets, we carry out additional experiments using the CelebA Dataset to underscore INSET's consistent superiority over the baseline models. Specifically, we will randomly insert one wrong image  in the optimal subset for anomaly detection tasks. From the results, the performance of all methods drops significantly, while HORSE still outperforms the others. This inspires further promising research on the topic to make neural subset selection methods more robust against noisy labels.

|   | PGM |  DeepSet        | Set-T  | EquiVSet | INSET |    HORSE     | 
| ----------- | -------------- | -------------- | -------------- |  -------------- | -------------- | -------------- |
| CelebA      | 0.371±0.027 | 0.436±0.031 | 0.417±0.028 | 0.437±0.035 |  0.465±0.037  | **0.469±0.034**  |