Table of Contents
Fetching ...

Mitigating over-exploration in latent space optimization using LES

Omer Ronen, Ahmed Imtiaz Humayun, Richard Baraniuk, Randall Balestriero, Bin Yu

TL;DR

Latent Space Optimization for discrete, structured problems often yields invalid or impractical solutions due to over-exploration. The paper introduces Latent Exploration Score ($LES$), a differentiable, decoder-centric density-based constraint derived from the push-forward sequence density through a CPA-represented decoder, enabling seamless integration into LSO without retraining. Empirical results across five benchmark tasks and 22 VAEs show that $LES$ consistently improves solution validity and objective quality, with AUROC analyses confirming reliable identification of valid latent regions and with competitive or superior performance to multiple baselines. The method is computationally intensive due to determinant calculations but demonstrates strong robustness and portability across architectures, offering a practical path to more realistic generative-discovery in discrete spaces; open-source code is provided.

Abstract

We develop Latent Exploration Score (LES) to mitigate over-exploration in Latent Space Optimization (LSO), a popular method for solving black-box discrete optimization problems. LSO utilizes continuous optimization within the latent space of a Variational Autoencoder (VAE) and is known to be susceptible to over-exploration, which manifests in unrealistic solutions that reduce its practicality. LES leverages the trained decoder's approximation of the data distribution, and can be employed with any VAE decoder - including pretrained ones - without additional training, architectural changes or access to the training data. Our evaluation across five LSO benchmark tasks and twenty-two VAE models demonstrates that LES always enhances the quality of the solutions while maintaining high objective values, leading to improvements over existing solutions in most cases. We believe that new avenues to LSO will be opened by LES' ability to identify out of distribution areas, differentiability, and computational tractability. Open source code for LES is available at https://github.com/OmerRonen/les.

Mitigating over-exploration in latent space optimization using LES

TL;DR

Latent Space Optimization for discrete, structured problems often yields invalid or impractical solutions due to over-exploration. The paper introduces Latent Exploration Score (), a differentiable, decoder-centric density-based constraint derived from the push-forward sequence density through a CPA-represented decoder, enabling seamless integration into LSO without retraining. Empirical results across five benchmark tasks and 22 VAEs show that consistently improves solution validity and objective quality, with AUROC analyses confirming reliable identification of valid latent regions and with competitive or superior performance to multiple baselines. The method is computationally intensive due to determinant calculations but demonstrates strong robustness and portability across architectures, offering a practical path to more realistic generative-discovery in discrete spaces; open-source code is provided.

Abstract

We develop Latent Exploration Score (LES) to mitigate over-exploration in Latent Space Optimization (LSO), a popular method for solving black-box discrete optimization problems. LSO utilizes continuous optimization within the latent space of a Variational Autoencoder (VAE) and is known to be susceptible to over-exploration, which manifests in unrealistic solutions that reduce its practicality. LES leverages the trained decoder's approximation of the data distribution, and can be employed with any VAE decoder - including pretrained ones - without additional training, architectural changes or access to the training data. Our evaluation across five LSO benchmark tasks and twenty-two VAE models demonstrates that LES always enhances the quality of the solutions while maintaining high objective values, leading to improvements over existing solutions in most cases. We believe that new avenues to LSO will be opened by LES' ability to identify out of distribution areas, differentiability, and computational tractability. Open source code for LES is available at https://github.com/OmerRonen/les.
Paper Structure (26 sections, 3 theorems, 27 equations, 4 figures, 12 tables, 1 algorithm)

This paper contains 26 sections, 3 theorems, 27 equations, 4 figures, 12 tables, 1 algorithm.

Key Result

Theorem 3.1

Let where $\boldsymbol{p}^{(i)}_{\boldsymbol{z}} = \text{Softmax}(\mathbf{L}_\theta(\boldsymbol{z}))_{.i}$ and $c^{(i)}_{\boldsymbol{z}} = \sum_{j=1}^D e^{\mathbf{L}_\theta(\boldsymbol{z})_{ji}}$. Assume that $\mathbf{L}_\theta$ is bijective and can be expressed as a CPA (eq:nn_cpas), and that $\boldsym for where $\left(\boldsymbol{A}^{(1)}_{\omega}, \dots, \boldsymbol{A}^{(L)}_{\omega}\right)^{

Figures (4)

  • Figure 1: Incorporating $\text{LES}$ promotes valid solutions. We consider the task of approximating the expression 1/3 + x + sin(x * x), using LSO. Optimization trajectories with (blue) and without (red) $\text{LES}$ constraint in the latent space of a VAE are projected onto a two-dimensional subspace that contains the starting point and the end-points obtained after 10 gradient ascent steps. In the left panel, we show the $\text{LES}$ score for latent vectors on the two-dimensional subspace, with darker shades corresponding to lower $\text{LES}$. In the right panel, we show the validity of the decoder outputs for each latent vector, with orange denoting invalid generations. High $\text{LES}$ values correlate with valid areas, and incorporating $\text{LES}$ in LSO produce an expression that adheres to the grammatical rules of \ref{['ex:expressions']}.
  • Figure 2: Derivation of $\text{LES}$. The decoder network ($\mathbf{G}_\theta$), which maps from the latent space to the output space, is assumed to be the composition of a softmax operation over a continuous piecewise affine (CPA) spline operator. $\text{LES}$ is the density of a random variable ($\boldsymbol{z}$) in the latent space, under the decoder transformation. Calculating $\text{LES}$ only requires a pre-trained decoder.
  • Figure 3: Cumulative objective for the top-20 solutions found during Bayesian optimization with the pre-trained SELFIES-VAE (maus2022local). Each method is shown in a distinct color. Solid lines represent solutions passing quality filters, while dashed lines include all evaluations. $\text{LES}$ outperforms all baselines on $\text{Ranolazine MPO}$ and $\text{Perindopril MPO}$, achieving competitive results on $\text{Zaleplon MPO}$.
  • Figure 4: Cumulative true objective for the best solution found during Bayesian optimization with the pre-trained SELFIES-VAE (maus2022local). Each method is shown in a distinct color. Solid lines represent solutions passing quality filters, while dashed lines include all evaluations. $\text{LES}$ achieves the best performance on $\text{Perindopril MPO}$ and is comptetitive on $\text{Zaleplon MPO}$ and $\text{Ranolazine MPO}$.

Theorems & Definitions (10)

  • Example 1.1: Arithmetic expressions
  • Example 1.2: Simplified molecular-input line-entry system (SMILES)
  • Example 1.3: Quality filters for molecules
  • Theorem 3.1: DGN sequence density
  • Remark 3.2
  • Lemma A.1
  • proof
  • proof : Proof of \ref{['thm:img_den']}
  • Lemma A.2
  • proof