Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions

2026-07-13

AI Summary

This work addresses the high computational cost and optimization instability of existing Monte Carlo sampling–based variational inference methods for learning neural set functions under weak supervision. The authors propose a sampling-free, continuous relaxation learning framework that reinterprets the evidence lower bound (ELBO) as a continuous relaxation of set functions and introduces a learnable surrogate objective to yield stable and efficient gradients. The approach enjoys approximation guarantees under submodular maximization and reveals a theoretical connection to variational free energy. Experimental results demonstrate that the method significantly outperforms current baselines across multiple real-world tasks, achieving faster convergence and substantially reduced computational overhead.

Abstract

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.

Problem

• neural set functions
• optimal subset oracle
• evidence lower bound
• variational inference
• gradient estimation

Innovation

• continuous relaxation
• neural set functions
• variational inference
• surrogate objective
• submodular maximization