Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions

Learning neural set functions is pivotal in a wide range of important applications, such as compound selection in AI-driven drug discovery and product recommendation. In recent work, optimal subset oracles have been introduced to implicitly learn set functions under practical weakly supervised settings, where the parameters are optimized via mean-field variational inference. However, such frameworks rely on Monte Carlo sampling to estimate the gradients of evidence lower bound when updating the variational distribution. Extensive sampling across iterations results in inefficient inference, while the induced stochasticity also perturbs the convergence trajectory. In this work, we propose treating the evidence lower bound as a relaxation of the set function and directly learning a surrogate objective to approximate it. In this way, we can efficiently and stably obtain the gradient at any point in the continuous domain. Furthermore, we derive the approximation ratio of the proposed framework under submodular maximization and analyze its connection with variational free energy. Experimental results show the proposed method consistently outperforms baselines across various real-world tasks.