Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions

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.

• The paper introduces ReSet, a method that replaces MC-based ELBO gradient estimation with a learned continuous relaxation for more stable set function optimization.
• The paper demonstrates that ReSet achieves up to a 70% improvement in mean Jaccard coefficient and significant computational speedups on both synthetic and real-world tasks.
• The paper provides theoretical guarantees by establishing a constant-factor approximation ratio and first-order convergence under standard smoothness assumptions.

Summary of "Advancing Optimal Subset Oracle via Learning Relaxation of Neural Set Functions" (2607.11555)

Context and Motivation

Neural set function learning undergirds numerous applications in ML, including AI-driven drug discovery, recommendation systems, and set anomaly detection. The challenge lies in efficiently learning set functions that guide optimal subset selection, especially from only weak or implicit supervision, such as selections made by users (optimal subset, OS, oracles) rather than explicit utility annotations (function-value, FV, oracles).

Recent OS oracle approaches embed set function learning in a variational inference framework, optimizing model parameters via the evidence lower bound (ELBO) and fitting a variational distribution over subsets. However, these methods typically rely on computationally intensive Monte Carlo (MC) sampling for unbiased gradient estimation, which introduces optimization noise and significant training and inference costs.

Methodology

This paper proposes a novel paradigm, "ReSet", that replaces MC-based ELBO gradient estimation in neural set function learning with a learned, continuous relaxation—a differentiable surrogate objective—of the set function. The central insight is that the ELBO itself can be viewed as a relaxation of the underlying discrete combinatorial optimization problem, and that a neural surrogate of this objective can be trained end-to-end to enable more stable and efficient optimization.

Key features of ReSet include:

• Relaxation of Set Functions: Learn a neural surrogate $g_\theta : [0,1]^{|V|} \to \mathbb{R}$ of the multilinear extension of the discrete set function $F_\theta(S;V)$, circumventing the need for repeated MC sampling.
• End-to-End Variational Optimization: Optimize both the variational distribution (parameterized as independent Bernoulli probabilities) and set function parameters via projected/gradient ascent directly on the learned relaxation, backpropagating through the inner optimization trajectory via Hessian-vector products.
• Theoretical Guarantees: The relaxation is (weakly) DR-submodular under mild architectural conditions. The paper establishes a provable constant-factor approximation ratio for stationary points and proves first-order convergence guarantees under standard smoothness assumptions.
• ELBO and Variational Free Energy Connection: The method is further interpreted as adaptive free energy minimization, naturally balancing expected utility (negative energy) and entropy.
• Scalable and Permutation-Invariant Architecture: All models leverage DeepSets for permutation invariance, auxiliary recognition networks for inference amortization, and Langevin dynamics for robust energy landscape exploration.

Empirical Results

ReSet is evaluated extensively on synthetic and real-world tasks:

• Synthetic Benchmarks: On structurally challenging datasets like two-moons and Gaussian mixtures, ReSet substantially outperforms MC-based OS oracle baselines, achieving up to 70% improvement in mean Jaccard coefficient (MJC).
• Product Recommendation (Amazon Registry): Across multiple product categories, ReSet variants consistently surpass DeepSet, Set Transformer, and prior OS oracles, with pronounced gains on outlier and sparse categories.
• Set Anomaly Detection: On Double MNIST, CelebA, F-MNIST, and CIFAR-10, ReSet shows robust improvements over baselines. For example, on Double MNIST, EquiVSet-R increases accuracy by nearly 84% vs. the original.
• Compound Selection in Drug Discovery: On PDBBind and BindingDB, ReSet-based oracles achieve the best MJC on the benchmarking tasks, especially excelling on large and compositional candidate sets (e.g., 300 elements in BindingDB).
• Computational Efficiency: ReSet achieves up to a 3.8x speedup and 2.5x lower memory usage compared to MC-based baselines for the same prediction quality, due to eliminating repeated sampling.

Ablation studies and sensitivity analyses confirm that observed gains stem from the learned relaxation, not just from architectural advances (e.g., attention, Langevin dynamics).

Theoretical Implications

This work demonstrates that adaptive learnable relaxations of combinatorial objectives can reliably replace MC-based estimators in variational set function inference. It solidifies connections between variational free energy minimization, continuous DR-submodular optimization, and neural permutation-invariant architectures. The constant-factor approximation and guaranteed convergence address concerns about tractability and reliability in set function learning, even as set sizes grow.

Practical Implications

By deploying ReSet as a plug-in module, practitioners can efficiently train OS oracle models in domains where explicit utility annotation is infeasible, such as in large-scale recommendation, automated chemical screening, and set-based data mining. The marked gains in speed, sample-efficiency, and predictive accuracy broaden the applicability of neural set function learning in practical weakly supervised regimes.

Outlook

Future research directions include:

• Extending learned relaxations to richer distributions (beyond mean-field, e.g., structured/discrete variational families).
• Tightening the correspondence between surrogate landscape optimization and combinatorial exact solutions, potentially with task-dependent regularization for stricter submodularity.
• Exploring relaxation-based training for other energy-based and combinatorial objectives beyond OS oracles, e.g., generalized set cover, clustering, and diverse subset selection.

Conclusion

The proposed ReSet method advances the state-of-the-art in weakly supervised neural set function learning by substituting MC-based ELBO gradient estimation with a trainable continuous relaxation. This leads to faster, more stable inference, theoretically justified by approximation and convergence guarantees, and validated empirically on diverse, large-scale benchmarks. The approach both illuminates the structure of variational set function optimization and provides practical tools for deployment in high-impact application domains.