# Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 61 references that are in the map.
- First posted on arXiv 13 Jul 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Variational Inference: A Review for Statisticians](https://mlsanity.com/paper/1601.00670/) (2016)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [DeepDTA: Deep Drug-Target Binding Affinity Prediction](https://mlsanity.com/paper/1801.10193/) (2018)
- [On Learning Sets of Symmetric Elements](https://mlsanity.com/paper/2002.08599/) (2020)
- [Structure-aware Interactive Graph Neural Networks for the Prediction of Protein-Ligand Binding Affinity](https://mlsanity.com/paper/2107.10670/) (2021)
- [AtomNet: A Deep Convolutional Neural Network for Bioactivity Prediction in Structure-based Drug Discovery](https://mlsanity.com/paper/1510.02855/) (2015)
- [Learning Iterative Reasoning through Energy Diffusion](https://mlsanity.com/paper/2406.11179/) (2024)
- [Gradient Methods for Submodular Maximization](https://mlsanity.com/paper/1708.03949/) (2017)
- [Learning Iterative Reasoning through Energy Minimization](https://mlsanity.com/paper/2206.15448/) (2022)
- [Energy-Based Transformers are Scalable Learners and Thinkers](https://mlsanity.com/paper/2507.02092/) (2025)
- [Atomic Convolutional Networks for Predicting Protein-Ligand Binding Affinity](https://mlsanity.com/paper/1703.10603/) (2017)
- [Equilibrium Matching: Generative Modeling with Implicit Energy-Based Models](https://mlsanity.com/paper/2510.02300/) (2025)
- [Expectation-Maximization for Learning Determinantal Point Processes](https://mlsanity.com/paper/1411.1088/) (2014)
- [Differentiable Submodular Maximization](https://mlsanity.com/paper/1803.01785/) (2018)
- [FEAT: Free energy Estimators with Adaptive Transport](https://mlsanity.com/paper/2504.11516/) (2025)
- [DeepSetNet: Predicting Sets with Deep Neural Networks](https://mlsanity.com/paper/1611.08998/) (2016)
- [Neural Set Function Extensions: Learning with Discrete Functions in High Dimensions](https://mlsanity.com/paper/2208.04055/) (2022)
- [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)
- [Relevance meets Diversity: A User-Centric Framework for Knowledge Exploration through Recommendations](https://mlsanity.com/paper/2408.03772/) (2024)
- [Learning Neural Set Functions Under the Optimal Subset Oracle](https://mlsanity.com/paper/2203.01693/) (2022)
- [A Unified Approach for Maximizing Continuous DR-submodular Functions](https://mlsanity.com/paper/2305.16671/) (2023)
- [Learning Graphical Model Parameters with Approximate Marginal Inference](https://mlsanity.com/paper/1301.3193/) (2013)
- [Deep Submodular Networks for Extractive Data Summarization](https://mlsanity.com/paper/2010.08593/) (2020)

## Abstract (arXiv)

Learning neural set functions is pivotal to a wide range of important applications, including compound selection in AI-driven drug discovery and product recommendation. Recent work has introduced optimal subset oracles to implicitly learn set functions under practical weakly supervised settings, where model parameters are optimized through mean-field variational inference. However, these frameworks rely on Monte Carlo sampling to estimate gradients of the evidence lower bound when updating the variational distribution. Repeated sampling across iterations incurs substantial computational overhead, while the resulting stochasticity can destabilize the optimization trajectory. In this work, we reinterpret the evidence lower bound as a continuous relaxation of the set function and learn a surrogate objective that replaces sampling-based ELBO gradient estimation during variational optimization. The learned surrogate provides stable and efficient gradients throughout the continuous domain, thereby reducing computational overhead and accelerating inference. Furthermore, we establish an approximation guarantee for the proposed framework under submodular maximization and characterize its connection to variational free energy. Experiments on a variety of real-world tasks demonstrate consistent improvements over existing baselines.