Time-Aware Latent Space Bayesian Optimization

This paper proposes Time-Aware Latent-space Bayesian Optimization (TALBO), a novel framework that integrates temporal dynamics into both the surrogate model and the generative latent space to effectively address objective drift in structured design tasks like molecular discovery, outperforming existing methods while maintaining robustness across varying drift conditions.

The Gist

Imagine you are a master chef trying to invent the perfect new recipe. But there's a catch: you can't just taste the dish while you're cooking it. You have to send the recipe to a distant, expensive lab, wait three days for them to cook it, taste it, and send you a score. This is how scientists design new molecules (like medicines) today. It's slow, expensive, and you can only try a few recipes before you run out of money.

To speed this up, scientists use a "smart assistant" called Bayesian Optimization. Think of this assistant as a mapmaker. Instead of tasting every possible recipe, it builds a map of the "flavor landscape" based on the few samples it has, guessing where the best-tasting dishes might be hidden.

However, there are two big problems with the old way of doing this:

1. The Map is Wrong: The "flavor landscape" for molecules is incredibly complex and messy (like a jungle). The old assistants tried to flatten this jungle into a simple, 2D map. But often, the map didn't match the real jungle, leading the chef to search in the wrong places.
2. The Goalposts Keep Moving: In the real world, what makes a "good" recipe changes. Maybe yesterday you wanted a spicy dish, but today you need something sweet. Or maybe a new virus emerged, and you need an antibody that fights that specific strain, not the old one. The old assistants kept searching for the "spicy" peak even after you told them you wanted "sweet." They were stuck in the past.

Enter TALBO: The Time-Aware Chef's Assistant

The paper introduces TALBO (Time-Aware Latent-space Bayesian Optimization). It's a new kind of assistant that solves both problems by understanding that time changes everything.

Here is how TALBO works, using simple analogies:

1. The "Living" Map (The Latent Space)

Imagine the old map was a static piece of paper. If the terrain changed (e.g., a river flooded), the paper map was useless.
TALBO uses a dynamic, living map. It doesn't just learn what molecules look like; it learns how they change over time.

• The Analogy: Think of a GPS app like Google Maps. A static map shows roads. A living GPS knows that traffic is heavy at 5 PM but empty at 2 AM. It adjusts the "best route" in real-time based on the current time. TALBO does this for molecules. It knows that the "shape" of the best molecule changes as the goal changes.

2. The "Dual-Brain" Approach

TALBO has two brains working together, both aware of time:

• Brain A (The Predictor): This brain guesses the score of a new recipe. It knows that a recipe that was great yesterday might be mediocre today. It constantly updates its predictions based on the current time.
• Brain B (The Translator): This brain translates the messy, complex molecule into the simple "map coordinates." In the old days, this translator was frozen. In TALBO, the translator also changes with time. It learns that "spicy" molecules look different from "sweet" molecules, and it redraws the map to make the "sweet" area easier to find.

3. The "Drifting" Objective

The paper tests this by simulating a scenario where the chef's preferences drift slowly.

• The Scenario: Imagine you are designing a drug. At first, you want it to be very strong. Then, you realize it needs to be less toxic. Then, you need it to dissolve faster.
• The Old Way: The old assistant would keep hunting for the "strongest" drug, ignoring the new need for "safety." By the time it realized the goal changed, it had wasted months.
• The TALBO Way: As soon as the goal shifts, TALBO's "living map" and "predictor" shift with it. It instantly reorients its search to find the new "sweet spot."

Why This Matters

The researchers tested TALBO on designing new molecules (specifically, finding molecules similar to existing medicines). They compared it to the best existing methods.

• The Result: TALBO consistently found better molecules faster. Even when the goals changed rapidly, TALBO adapted quickly. The old methods got stuck, chasing ghosts of the past.
• The "Beyond the Dataset" Magic: The researchers also found that TALBO could find molecules that were better than anything in its initial training library. It didn't just copy-paste old recipes; it invented new ones that fit the current needs perfectly.

The Bottom Line

Think of TALBO as a chef who doesn't just have a recipe book, but has a time-traveling intuition.

• If the world changes, the chef's intuition changes with it.
• It doesn't just look at what worked yesterday; it understands that what works today requires a different kind of search.

In a world where problems (like new viruses or changing market trends) evolve constantly, TALBO is the tool that helps us stop looking for the "perfect solution of the past" and start finding the "perfect solution for right now."

Technical Summary

1. Problem Statement

The paper addresses a critical limitation in Latent Space Bayesian Optimization (LSBO). While LSBO effectively optimizes structured, discrete domains (e.g., molecular design, protein engineering) by mapping them to a continuous latent space via generative models (like VAEs), standard LSBO assumes a static objective function.

In real-world design campaigns (e.g., antibody discovery against mutating pathogens, human-in-the-loop optimization, or shifting multi-objective preferences), the objective function drifts over time.

• The Challenge: When the objective changes, two things happen:
  1. The surrogate model (predicting performance) becomes outdated.
  2. The latent search space geometry itself may become misaligned with the new objective because the generative model was trained on static data or does not account for temporal shifts.
• The Gap: Existing time-varying BO methods operate in continuous input spaces and do not account for the complexity of a learned latent representation that may also need to evolve.

2. Methodology: TALBO

The authors propose Time-Aware Latent Space Bayesian Optimization (TALBO), a framework that incorporates temporal dynamics into both the surrogate model and the generative representation.

A. Dual Temporal Modeling

TALBO treats time ($t$) as a first-class citizen in two distinct components:

1. Time-Aware Surrogate: Models the objective function $g(z, t)$ in the latent space using a Spatio-Temporal Gaussian Process (GP). The kernel is defined over both the latent code $z$ and time $t$, allowing the model to weigh past observations based on their temporal relevance.
2. Time-Aware Generative Representation: Instead of a fixed VAE, TALBO uses a GP-Prior Variational Autoencoder (specifically DGBFGP).
   • Covariates: The latent code $z$ is conditioned on covariates $c(x, t)$, which include molecular descriptors and the current time step $t$.
   • Mechanism: The latent space is induced by an additive GP prior over covariate-dependent latent functions. As time evolves, the mapping from the design space $X$ to the latent space $Z$ shifts, ensuring the latent geometry remains aligned with the current objective landscape.

B. The DGBFGP Architecture

To make GP-prior VAEs scalable, TALBO utilizes Deep Generative Basis Function GP (DGBFGP):

• Basis Function Approximation: It approximates the GP prior using Hilbert-space basis functions, converting the infinite-dimensional GP into a finite linear form $z(c) \approx A \phi(c)$.
• Variational Inference: Global parameters (weights $A$ and kernel hyperparameters) are learned via stochastic variational inference (SVI).
• Alignment & Inversion: To prevent misalignment between the latent code and the decoded molecule (a common issue in LSBO), TALBO integrates:
  • Latent Alignment Regularizers: Encourages smoothness between latent distance and objective value (inspired by CoBO).
  • Latent Inversion: Refines latent codes via gradient descent to ensure they decode back to the evaluated molecule (inspired by InvBO).

C. Optimization Loop

1. Embed: Encode evaluated molecules into the time-conditioned latent space.
2. Surrogate Update: Fit a spatio-temporal GP on the latent observations $(z, t, y)$.
3. Acquisition: Select the next latent query using an acquisition function (e.g., Thompson Sampling) that balances exploration and exploitation under the current time $t$.
4. Decode & Evaluate: Decode the latent point to a molecule, evaluate it under the current time-varying objective $f_t$.
5. Update: Update the surrogate and, periodically, the generative model (re-embedding all data to maintain alignment).

3. Key Contributions

• Conceptual: Formalizes the problem of Time-Varying LSBO, identifying that temporal drift affects both the surrogate and the latent geometry.
• Methodological: Introduces TALBO, the first framework to couple a spatio-temporal surrogate with a conditional, time-aware generative representation (GP-Prior VAE) for BO.
• Empirical:
  • Adapts standard molecular design benchmarks (GuacaMol) to a drifting multi-property objective setting.
  • Proposes evaluation metrics suitable for time-varying settings (Current-objective best-so-far, Cumulative Regret).
  • Demonstrates robustness across different drift speeds and design choices.

4. Experimental Results

The authors evaluated TALBO against strong baselines (InvBO, CoBO, LOL-BO, standard LSBO, and Random Search) on four molecular design tasks with drifting priorities.

• Performance: TALBO consistently outperformed all baselines in terms of best-so-far value and cumulative regret. It adapted faster to shifts in objective weights, maintaining higher performance when priorities changed.
• Ablation Studies:
  • Dual Modeling: Removing time from either the surrogate or the representation degraded performance, confirming that both components are necessary.
  • Drift Speed: TALBO remained robust across varying rates of objective drift (controlled by GP length scales).
  • Static Regime: When the objective was fixed (time-invariant), TALBO remained competitive, showing no systematic penalty for modeling time when it wasn't strictly necessary.
• Efficiency: TALBO could identify molecules outperforming the top 0.1% of the pre-training dataset (1.27M molecules) using only ~6,000 queries, demonstrating "beyond-dataset" capability even in dynamic settings.

5. Significance and Impact

• Bridging a Gap: This work bridges the gap between time-varying optimization and latent space optimization, a combination previously unexplored despite its relevance to real-world applications like drug discovery and materials science.
• Dynamic Adaptation: It provides a principled mechanism for optimization campaigns where "what is good" changes over time, preventing the optimizer from getting stuck in stale local optima defined by outdated objectives.
• Scalability: By leveraging basis-function approximations and variational inference, the method remains computationally feasible for high-dimensional structured data.
• Future Directions: The paper highlights the need for better stale-data management in LSBO (e.g., forgetting mechanisms) and more expressive temporal kernels to handle abrupt regime changes.

In summary, TALBO represents a significant step forward in making Bayesian Optimization robust to the dynamic, non-stationary nature of real-world scientific design problems.