# Learning Neural Set Functions Under the Optimal Subset Oracle

arXiv: https://arxiv.org/abs/2203.01693 · HTML: https://mlsanity.com/paper/2203.01693/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 68 references that are in the map.
- First posted on arXiv 3 Mar 2022.

## Builds on (most cited first, 30 of 34)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [Advances in Variational Inference](https://mlsanity.com/paper/1711.05597/) (2017)
- [GLISTER: Generalization based Data Subset Selection for Efficient and Robust Learning](https://mlsanity.com/paper/2012.10630/) (2020)
- [Conditional Random Fields as Recurrent Neural Networks](https://mlsanity.com/paper/1502.03240/) (2015)
- [DeepDTA: Deep Drug-Target Binding Affinity Prediction](https://mlsanity.com/paper/1801.10193/) (2018)
- [Learning Gradient Fields for Molecular Conformation Generation](https://mlsanity.com/paper/2105.03902/) (2021)
- [Learning Non-Convergent Non-Persistent Short-Run MCMC Toward Energy-Based Model](https://mlsanity.com/paper/1904.09770/) (2019)
- [Set Functions for Time Series](https://mlsanity.com/paper/1909.12064/) (2019)
- [Inference Suboptimality in Variational Autoencoders](https://mlsanity.com/paper/1801.03558/) (2018)
- [Discriminative Embeddings of Latent Variable Models for Structured Data](https://mlsanity.com/paper/1603.05629/) (2016)
- [Oops I Took A Gradient: Scalable Sampling for Discrete Distributions](https://mlsanity.com/paper/2102.04509/) (2021)
- [SIMILAR: Submodular Information Measures Based Active Learning In Realistic Scenarios](https://mlsanity.com/paper/2107.00717/) (2021)
- [Structure-aware Interactive Graph Neural Networks for the Prediction of Protein-Ligand Binding Affinity](https://mlsanity.com/paper/2107.10670/) (2021)
- [DrugOOD: Out-of-Distribution (OOD) Dataset Curator and Benchmark for AI-aided Drug Discovery -- A Focus on Affinity Prediction Problems with Noise Annotations](https://mlsanity.com/paper/2201.09637/) (2022)
- [AtomNet: A Deep Convolutional Neural Network for Bioactivity Prediction in Structure-based Drug Discovery](https://mlsanity.com/paper/1510.02855/) (2015)
- [Guarantees for Greedy Maximization of Non-submodular Functions with Applications](https://mlsanity.com/paper/1703.02100/) (2017)
- [Conditional Set Generation with Transformers](https://mlsanity.com/paper/2006.16841/) (2020)
- [Atomic Convolutional Networks for Predicting Protein-Ligand Binding Affinity](https://mlsanity.com/paper/1703.10603/) (2017)
- [Expectation-Maximization for Learning Determinantal Point Processes](https://mlsanity.com/paper/1411.1088/) (2014)
- [Differentiable Submodular Maximization](https://mlsanity.com/paper/1803.01785/) (2018)
- [DeepSetNet: Predicting Sets with Deep Neural Networks](https://mlsanity.com/paper/1611.08998/) (2016)
- [Deep Submodular Functions](https://mlsanity.com/paper/1701.08939/) (2017)
- [Set Prediction without Imposing Structure as Conditional Density Estimation](https://mlsanity.com/paper/2010.04109/) (2020)

## Built on by (most cited first, 10 of 10)

- [Learning Energy-Based Models by Cooperative Diffusion Recovery Likelihood](https://mlsanity.com/paper/2309.05153/) (2023)
- [On permutation-invariant neural networks](https://mlsanity.com/paper/2403.17410/) (2024)
- [Energy Discrepancies: A Score-Independent Loss for Energy-Based Models](https://mlsanity.com/paper/2307.06431/) (2023)
- [Learning Exceptional Subgroups by End-to-End Maximizing KL-divergence](https://mlsanity.com/paper/2402.12930/) (2024)
- [Energy-Based Modelling for Discrete and Mixed Data via Heat Equations on Structured Spaces](https://mlsanity.com/paper/2412.01019/) (2024)
- [Advances in Set Function Learning: A Survey of Techniques and Applications](https://mlsanity.com/paper/2501.14991/) (2025)
- [CheMixHub: Datasets and Benchmarks for Chemical Mixture Property Prediction](https://mlsanity.com/paper/2506.12231/) (2025)
- [Enhancing Neural Subset Selection: Integrating Background Information into Set Representations](https://mlsanity.com/paper/2402.03139/) (2024)
- [Learning Set Functions with Implicit Differentiation](https://mlsanity.com/paper/2412.11239/) (2024)
- [Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions](https://mlsanity.com/paper/2607.11555/) (2026)

## Abstract (arXiv)

Learning neural set functions becomes increasingly more important in many applications like product recommendation and compound selection in AI-aided drug discovery. The majority of existing works study methodologies of set function learning under the function value oracle, which, however, requires expensive supervision signals. This renders it impractical for applications with only weak supervisions under the Optimal Subset (OS) oracle, the study of which is surprisingly overlooked. In this work, we present a principled yet practical maximum likelihood learning framework, termed as EquiVSet, that simultaneously meets the following desiderata of learning set functions under the OS oracle: i) permutation invariance of the set mass function being modeled; ii) permission of varying ground set; iii) minimum prior; and iv) scalability. The main components of our framework involve: an energy-based treatment of the set mass function, DeepSet-style architectures to handle permutation invariance, mean-field variational inference, and its amortized variants. Thanks to the elegant combination of these advanced architectures, empirical studies on three real-world applications (including Amazon product recommendation, set anomaly detection, and compound selection for virtual screening) demonstrate that EquiVSet outperforms the baselines by a large margin.
