Enhancing Neural Subset Selection: Integrating Background Information into Set Representations

TL;DR

This work addresses neural subset selection by incorporating background information from the superset into subset representations. It develops INSET, an information-aggregation module built on invariant-sufficient representations that are stable under hierarchical permutation symmetries of (S,V). Theoretical results connect functional and probabilistic symmetries, guiding network design to model $P(Y|S,V)$ via $M(S,V)$, and empirically demonstrate superior performance across product recommendation, anomaly detection, and drug-discovery tasks. The approach improves both prediction quality and training efficiency, validating the practical value of superset-aware, permutation-invariant learning for set-valued tasks.

Abstract

Learning neural subset selection tasks, such as compound selection in AI-aided drug discovery, have become increasingly pivotal across diverse applications. The existing methodologies in the field primarily concentrate on constructing models that capture the relationship between utility function values and subsets within their respective supersets. However, these approaches tend to overlook the valuable information contained within the superset when utilizing neural networks to model set functions. In this work, we address this oversight by adopting a probabilistic perspective. Our theoretical findings demonstrate that when the target value is conditioned on both the input set and subset, it is essential to incorporate an \textit{invariant sufficient statistic} of the superset into the subset of interest for effective learning. This ensures that the output value remains invariant to permutations of the subset and its corresponding superset, enabling identification of the specific superset from which the subset originated. Motivated by these insights, we propose a simple yet effective information aggregation module designed to merge the representations of subsets and supersets from a permutation invariance perspective. Comprehensive empirical evaluations across diverse tasks and datasets validate the enhanced efficacy of our approach over conventional methods, underscoring the practicality and potency of our proposed strategies in real-world contexts.

Figures (4)

• Figure 1: (Left): The DeepSet-style models only focus on processing the subset $S$. (Right): In contrast, INSET not only identifies the subset $S$ but also takes the identification of $V$ into account, which parameterizes the information that $S$ is a subset of $V$ during the training process.
• Figure 2: Left: A sample from the Double MNIST dataset, comprising $|S^*|$ images displaying the same digit (indicated by the red box). Right: The two figures on the right display the validation performance plotted against the number of epochs for Toys and Diaper datasets, respectively. The x-axis represents the epochs.
• Figure 3: Here is an illustration of the CelebA dataset from the work by ou2022learning. Each row in the dataset represents a sample, containing a combination of $|S^*|$ anomaly images (highlighted in red boxes) and $8-|S^*|$ normal images. Notably, within each sample, normal images possess two specific attributes, which are indicated in the rightmost column. In contrast, anomalies lack both of these attributes. This clear distinction between normal images and anomalies allows for a comprehensive analysis and understanding of the dataset's characteristics.
• Figure 4: Sensitivity analysis of INSET performance under varying numbers of Monte Carlo (MC) sampling.

Theorems & Definitions (18)

• Definition 3.2
• Definition 3.3
• Definition 3.3
• Theorem 3.4
• Corollary 3.4
• Proposition 3.4
• Proposition 3.5
• Lemma A.1: Conditional independence and randomization
• Lemma A.2
• Lemma A.3