Closing the Loop: Machine Learning for Optimization and Discovery

Closing the Loop: Machine Learning for Optimization and Discovery

Speakers: Andreas Krause

Conference: ICML 2025

YouTube: https://slideslive.com/39043349

Overview

In this insightful talk at ICML 2025, Andreas Krause, a leading researcher in machine learning, presented a compelling vision for "Closing the Loop: Machine Learning for Optimization and Discovery." The presentation transcends the traditional linear machine learning pipeline—collecting vast datasets, fitting models, and generalizing—to advocate for a dynamic, iterative approach where ML actively drives scientific experimentation and discovery. Krause argues that while current ML advancements, from image classification to large generative models, are profound, their full potential in science remains untapped due to the inherent challenges of data scarcity and the high cost of real-world experiments.

The core of Krause's argument lies in the necessity of integrating machine learning into the experimental process itself, using it not just to analyze data but to intelligently generate the most informative data. This paradigm shift, rooted in principles of optimal experimental design and information theory, aims to accelerate progress in fields ranging from protein engineering and materials discovery to precision agriculture and robotics. By systematically quantifying uncertainty and leveraging that uncertainty to guide subsequent experiments, ML can enable scientists to efficiently navigate vast search spaces, discover novel solutions, and gain deeper insights into complex systems, ultimately moving towards the vision of "self-driving labs."

Krause meticulously unpacks the methodological and practical advancements required to realize this closed-loop vision, highlighting the evolution from Bayesian optimization to Bayesian model-based reinforcement learning (RL). He demonstrates how these frameworks, when augmented with sophisticated techniques like physics-aware priors and active learning for large language models, can tackle high-dimensional, real-world problems that are otherwise intractable. The talk not only celebrates the foundational work in this community but also lays out critical open challenges, such as uncertainty quantification in modern foundation models and the orchestration of human-AI scientific workflows, positioning this area as a frontier for future ML research.

Background

▶ Watch: Introduction: ML's impact on science and protein design (0:00)

The prevailing practice in machine learning often involves a linear pipeline: gather an abundance of data (often from the internet), train a model, and hope it generalizes well to new, similar examples. While effective for many applications, this approach falters in scientific domains where data acquisition is inherently expensive and time-consuming. As eloquently captured by Richard Feynman, "it doesn't matter how beautiful your theory is, it doesn't matter how smart you are if it doesn't agree with experiment it's wrong." This underscores the indispensable role of experimentation in scientific progress, and thus, the critical need for algorithmic support to make experiments smarter and more efficient.

The quest for intelligent information gathering has a rich history, drawing from diverse fields. In statistics, optimal experimental design dates back to the 19th century with pioneering work by Fisher and Lindley. Economics contributed ideas like the value of information, while information theory and approximation theory provided foundational concepts. Within AI and machine learning, these ideas have crystallized into frameworks such as active learning, multi-armed bandits, Bayesian optimization, and reinforcement learning (RL), all aimed at formalizing sequential decision-making under uncertainty to acquire the most useful information.

Key methodological questions underpin this field:

1. Quantifying Epistemic Uncertainty: Before designing experiments, one must understand what is unknown. This involves developing robust methods for uncertainty quantification, leveraging techniques like Gaussian process models, Bayesian deep learning, and causal models.
2. Objective Functions for Experimentation: How should potential experiments be scored? Objectives might include reducing uncertainty, improving decision-making, or maximizing information gain.
3. Combinatorial Design and Planning: Given an objective, how does one design a diverse and informative set of experiments, or create policies for sequential information gathering?
4. Fundamental Theoretical Questions: What are the sample complexity guarantees for active learning? How much regret is suffered? Can these methods compete with optimal designs, and do they introduce biases by actively selecting data?

These questions are motivated by real-world applications across various disciplines. Early work included optimizing sensor placement for environmental monitoring (e.g., water quality) and guiding mobile robots. More recently, these techniques have been applied to tuning complex scientific instruments like the SwissFEL Free-Electron Laser, finding sustainable soil management practices in precision agriculture, and a major focus of Krause's talk: protein engineering and design. The increasing trend towards lab automation and the vision of "self-driving labs" further amplify the relevance of these closed-loop ML approaches.

Key Findings

▶ Watch: Closing the loop: Algorithmic support for smart experimentation (2:30)

Andreas Krause's talk highlights several pivotal findings and contributions in the realm of machine learning for optimization and discovery, demonstrating a clear progression from foundational techniques to advanced, scalable solutions:

1. Bayesian Optimization (BO) as a Foundational Tool: BO is presented as a highly effective method for optimizing expensive, black-box objective functions. Its strength lies in building a probabilistic model of the unknown landscape, quantifying epistemic uncertainty, and using this uncertainty to guide exploration and exploitation, leading to the discovery of novel and superior solutions (e.g., more stable proteins).
2. Scaling to Complex Domains via Bayesian Model-Based Reinforcement Learning (RL): Recognizing the limitations of BO in high-dimensional or sequential decision-making tasks, Krause advocates for a shift to Bayesian model-based RL. This approach explicitly models the uncertainty in system dynamics (transition kernels and rewards) and uses this model for introspective planning, effectively addressing the "sim-to-real" gap prevalent in real-world applications where perfect simulators are unavailable.
3. Enhanced Sample Efficiency with Physics-Aware Priors: A significant finding is that incorporating domain knowledge—either from simplified first-principle models or extensive simulations—into the Bayesian model as a physics-aware prior drastically improves sample efficiency. This allows real-world agents (like an RC car) to learn complex tasks with far fewer physical experiments, transforming raw learning into informed adaptation.
4. Active Learning for Large Language Models (LLMs): Even with powerful pre-trained LLMs, active learning at test-time proves crucial for sharpening predictions and guiding reasoning on specific, high-stakes tasks. By selectively querying relevant data sources to minimize uncertainty, this approach significantly enhances predictive performance (e.g., cardiovascular risk prediction) and enables more robust reasoning, even allowing smaller LLMs to achieve performance comparable to much larger ones with additional computational investment.
5. Guiding Generative Models for De Novo Design: The talk demonstrates that modern generative models, particularly diffusion models used for de novo protein design, can be effectively "steered" using principles from optimal control and RL. This allows for targeted optimization of desired properties (e.g., thermostability) or for more open-ended, unsupervised exploration (e.g., entropy maximization) of novel designs within the generated manifold.

These findings collectively underscore the power of uncertainty-aware, closed-loop machine learning to drive scientific progress, offering practical solutions to challenges of data scarcity, simulation inaccuracy, and the efficient exploration of vast design spaces.

Technical Deep Dive

▶ Watch: Key methodological questions in active information gathering (4:00)

The technical core of Krause's talk revolves around the evolution and scaling of uncertainty-aware machine learning techniques for real-world optimization and discovery.

Bayesian Optimization (BO)

At its foundation, BO operates by maintaining a probabilistic model of an unknown, expensive-to-evaluate objective function, often using Gaussian processes (GPs). A GP models the function as a distribution over functions, providing not only a mean prediction but also a measure of epistemic uncertainty (variance) at any untried point. This uncertainty is crucial for guiding exploration.

The BO loop typically involves:

1. Prior Data: Start with an initial set of experimentally characterized points.
2. Model Fitting: Fit a probabilistic model (e.g., GP) to this data, yielding a posterior distribution over the objective function.
3. Acquisition Function: Use an acquisition function to decide which point to evaluate next. A prominent example is the Upper Confidence Bound (UCB), which selects the point maximizing (posterior mean + kappa * posterior standard deviation). This function naturally balances exploitation (choosing points with high predicted mean) and exploration (choosing points with high uncertainty). Other acquisition functions include Expected Improvement (EI) or Probability of Improvement (PI).
4. Experimentation: Evaluate the chosen point in the real world (e.g., lab experiment).
5. Update Model: Add the new observation to the dataset and update the probabilistic model.

This iterative process efficiently navigates the search space, especially when experiments are costly. However, BO's traditional strength lies in relatively low-dimensional, black-box problems, with scaling to high-dimensional or sequential decision-making tasks being a recognized limitation.

Bayesian Model-Based Reinforcement Learning (RL)

To scale beyond the limitations of classical BO, Krause introduces Bayesian model-based RL. Here, the goal is not just to optimize a static objective but to learn a policy for sequential decision-making in an environment where the dynamics are unknown and uncertain.

The key idea is to build a Bayesian model of the environment's transition kernel P(S_t+1 | S_t, A_t) and potentially its rewards. This means quantifying uncertainty not just in the value of a single point, but in how actions change the state of the world.

The loop for Bayesian model-based RL is:

1. Prior Belief: Start with a prior over the dynamics (possibly uninformative or based on rough assumptions).
2. Iterative Refinement: In each round (experiment):

• Bayesian Inference: Update the posterior belief about the environment dynamics using new experimental data.
• Introspective Planning: Given the posterior belief, plan a new behavioral policy. This is the most challenging step. The naive approach would be to find a policy that maximizes expected reward across all plausible dynamics (a Bayes-optimal policy). This is computationally intractable and prone to "dynamics hacking" (exploiting model inaccuracies rather than true environmental properties).
• A more robust approach, inspired by the optimism in the face of uncertainty principle, seeks a policy that performs well under the best plausible dynamics. This encourages exploration by preferring actions that could lead to high rewards if the world turns out to be favorable.
• Algorithms like PILCO (using GPs for dynamics) and PETS (using deep ensembles) embody this principle. Krause describes a relaxation to a tractable optimistic value function, allowing the use of standard policy gradient methods to solve it. This approach is proven to achieve sublinear regret under certain regularity assumptions and is compatible with state-of-the-art model-based RL algorithms like MBPO and Dreamer.
• Execution: Roll out the new policy in the real world (the only real-world interaction).
• Learning: The outcome of the experiment updates the posterior belief, shrinking uncertainty and informing future planning.

Physics-Aware Priors for RL

The sample efficiency of Bayesian model-based RL can be dramatically improved by incorporating prior knowledge. This is achieved through hierarchical Bayesian modeling, where a data-driven prior is learned from similar domains or simulations.

• Domain Randomization (DR) is a common technique in robotics where a simulator's parameters (e.g., mass, friction) are varied to create diverse training environments. Instead of just training a robust policy, Krause proposes distilling this varied simulated data into a stochastic process prior.
• This "physics-aware prior" is essentially a large neural network (a "foundation model for physics") that captures the higher-order marginals and regularity of sample paths consistent with the simulations. It's trained using techniques like score matching, similar to diffusion models.
• By initializing the real-world learning with such a prior, the agent avoids "spiraling wildly out of control" and learns the task much faster, as demonstrated with the RC car experiment. It quickly focuses on relevant actions, bridging the "sim-to-real" gap more effectively.

Active Learning for Large Language Models (LLMs)

Even for powerful, pre-trained LLMs, active learning can significantly enhance performance for specific tasks at test time.

• The concept is to "sharpen" the pre-trained model's prediction for a particular prompt X_star by selectively conditioning it on relevant training data. This is akin to transductive learning.
• The acquisition function here aims to minimize uncertainty about the model's response f(X_star).
• A major challenge is quantifying epistemic uncertainty in LLMs, which remains an open research problem. A pragmatic approach involves extracting relevant features from the LLM (or an embedding model) and using Gaussian processes to calculate entropy, providing a measure of uncertainty.
• Crucially, the selection process must be diversity-aware to avoid selecting redundant information, which can lead to overfitting and performance degradation, especially in retrieval-augmented generation (RAG) systems where corpora might contain many similar examples.

Steering Generative Models

Modern generative models, especially diffusion models (e.g., RFdiffusion for proteins), have revolutionized de novo design by transforming simple noise into complex, structured outputs. Krause demonstrates how these models can be actively guided.

• The generation process can be viewed through the lens of optimal control and reinforcement learning.
• Adjoint methods or adjoint matching can be used to fine-tune diffusion models to optimize specific objectives (e.g., thermostability, binding affinity). This allows incorporating prior knowledge about desired properties directly into the generation.
• For more open-ended exploration, one can steer generation towards maximizing expected improvement (finding better-than-current solutions) or maximizing entropy over the generated manifold (unsupervised exploration to discover diverse, novel structures).
• The Flow Density Control framework provides flexibility in defining these optimization objectives for generative models, with applications spanning image and molecular generation.

Experimental Setup & Results

▶ Watch: Diverse applications: from sensors to free-electron lasers (6:00)

The talk illustrates these technical concepts with compelling experimental results across several domains, showcasing the practical impact of closing the loop with ML.

Protein Engineering and Design

• Early Work (with Frances Arnold): In collaboration with Nobel laureate Frances Arnold and Phil Romero, early applications focused on optimizing protein properties. Starting with wild-type proteins, mutations were introduced, and an initial library of candidates was created. Bayesian optimization was then used in an iterative, "novelty-seeking" manner to design subsequent experimental batches.
• Results: This approach successfully identified protein sequences that were considerably more stable in the lab than those found by previous methods, demonstrating the power of uncertainty-driven exploration.
• Recent Work (NCCR Catalysis): More recent collaborations, such as with the NCCR Catalysis consortium in Switzerland, applied Bayesian optimization to discover artificial metalloenzymes for new-to-nature reactions, aiming for sustainable alternatives to fossil fuel-dependent catalysts.
• Setup: Leveraging advances in lab automation (parallel experiments), computational atomic descriptors, and protein language models (like ESM-1 for feature extraction), the complexity of characterizing dynamic catalytic reactions could be tackled.
• Results: With only a few rounds of experimentation, the approach identified pools of designs that were much more active than the original wild type and yielded solutions with significantly higher yields.

RC Car Reverse Parking (Bayesian Model-Based RL)

• Task: Training an RC car to reverse park into a designated spot.
• Challenge: This task requires modeling complex physical phenomena like tire properties, slippage, and aerodynamic effects, making it difficult to build an accurate simulator (sim-to-real gap).
• Baseline: Estimating a simplified first-principle model and training a fixed policy against it failed to solve the task due to model mismatch.
• Bayesian Model-Based RL (Uninformed Prior): Learning from scratch, the agent initially drifted randomly but learned to solve the task after approximately 15-20 episodes of real-world interaction. This highlights the sample efficiency challenge in real-world RL.
• Bayesian Model-Based RL (Physics-Aware Prior): By incorporating a prior learned from domain randomization in simulation, the agent's initial behavior was much more controlled.
• Results: It learned to turn around after about 5 episodes and solved the task after approximately 10 episodes, demonstrating a 2x improvement in sample efficiency compared to the uninformed prior, and significantly faster and more stable learning than the baseline.

Active Learning for LLMs

• Cardiovascular Risk Prediction: Fine-tuning LLMs on the UK Biobank data for cardiovascular risk prediction.
• Results: Even with relatively few patients, test-time tuning yielded substantial improvements in predictive performance.
• Language Modeling Performance (Pile Benchmark): Experiments on open language models evaluated on the Pile benchmark for language modeling performance (bits per byte).
• Baseline: Nearest-neighbor prediction (common in retrieval-augmented generation).
• Challenge: Duplication in corpora can lead to selecting similar information repeatedly, causing overfitting and degradation of predictions when conditioning on more information.
• Results: By employing a diversity-aware active learning strategy for fine-tuning, this overfitting was avoided, leading to improved and stable predictive performance across various language models. The work also showed that smaller language models, when combined with more test-time compute for active learning, can achieve similar performance to significantly larger models.

Subgoal Selection for RL (AntMaze)

• Task: A simulated ant navigating a maze and reaching target states.
• General Goal Exploration: An approach that generally explores all reachable subgoals effectively covers the state space.
• Target-Oriented Subgoal Selection: By incorporating relevance to a specific target task (using ideas from transductive learning), the agent's exploration becomes much more focused.
• Results: In complex environments with obstacles, the target-oriented approach was able to solve tasks that were otherwise difficult or impossible, by providing a sense of direction and effectively breaking down the problem into relevant subgoals.

Steering Generative Models (Flow Density Control)

• Experiment: Demonstrating the Flow Density Control framework on a dataset with two distinct modes (one with many samples, one with few).
• Objective: Entropy maximization (unsupervised exploration).
• Results: By optimizing for entropy, the model was able to spread the generated mass more broadly, including the boundaries of the known data, showcasing its flexibility for discovering novel structures beyond the most populated modes of the training data. This was applied to both image and molecular generation.

Practical Implications

▶ Watch: Focus: Self-driving labs and protein engineering (7:00)

The advancements presented by Andreas Krause have profound practical implications for a wide array of fields, transforming how machine learning interacts with real-world systems and scientific discovery.

For Practitioners and Infra Teams:

• Accelerated Scientific Discovery: The closed-loop approach directly addresses the bottleneck of expensive and time-consuming experimentation in fields like materials science, chemistry, and biology. By intelligently guiding experiments, the discovery cycle for new proteins, catalysts, or materials can be drastically shortened, leading to faster innovation.
• Resource Efficiency: For industries where experiments are costly (e.g., drug discovery, advanced manufacturing, large-scale scientific facilities), methods like Bayesian optimization and sample-efficient RL minimize the number of real-world trials needed. This translates into significant savings in time, materials, and labor.
• Overcoming the Sim-to-Real Gap: The development of Bayesian model-based RL offers a robust framework for deploying ML agents in real-world environments where perfect simulators are elusive or impossible to build. This is critical for robotics, autonomous systems, and any application requiring interaction with complex, uncertain physical systems.
• Leveraging Existing Knowledge and Data: The concept of physics-aware priors provides a principled way to integrate domain expertise, simplified models, or data from simulations into real-world learning. This "grey-box modeling" approach prevents learning from scratch, leading to faster convergence and more robust performance, making ML more accessible and effective for domain experts.
• Enhanced LLM Utility for High-Stakes Applications: Active learning for LLMs at test-time means these powerful models can be fine-tuned for specific, critical tasks (e.g., medical diagnosis, targeted molecule generation) with higher accuracy and reliability. This also opens up possibilities for deploying smaller, more efficient LLMs by investing compute at inference time, potentially reducing the massive computational demands of large models.

For Model Builders and Deployers:

• Uncertainty Quantification as a First-Class Citizen: The emphasis on quantifying epistemic uncertainty is a paradigm shift. Model builders must prioritize not just predictive accuracy but also the model's awareness of what it doesn't know. This is crucial for safe and reliable deployment, especially in high-stakes environments.
• Designing for Iterative Interaction: Instead of building static predictive models, the focus shifts to designing "agentic" systems that can actively query, experiment, and learn. This requires new architectures and evaluation metrics that account for iterative improvement and information gain.
• Trade-offs and Limitations:
• Computational Cost vs. Experimental Cost: While these methods reduce real-world experiments, they often increase computational demands (e.g., for planning in Bayesian model-based RL, test-time tuning for LLMs). Deployers need to balance these costs based on the application.
• Epistemic Uncertainty in LLMs: Acknowledged as a significant open challenge, the lack of robust uncertainty quantification in large foundation models limits their full potential in closed-loop systems. Continued research is needed to make these models truly "aware of what they don't know."
• Orchestrating Human-AI Workflows: The vision of "self-driving labs" requires careful consideration of how humans and AI systems collaborate in scientific discovery. Designing effective interfaces and decision-making protocols for this joint effort is a complex, unsolved problem.
• Causality and Safety: For systems that drive actions in the real world, understanding causal effects and ensuring safety during exploration are paramount. This involves incorporating constraints into the learning process and moving beyond purely observational data.
• Multi-Objective Optimization: Real-world problems rarely have a single objective. Balancing competing goals (e.g., maximizing reward while minimizing risk or cost) requires sophisticated multi-objective optimization techniques, an area that needs more dedicated research in this context.
• Building Trust: For ML models to be adopted by domain experts, they must consistently provide useful and reliable suggestions. Balancing aggressive exploration with delivering tangible, positive results is crucial for fostering trust in these intelligent systems.

In essence, Krause's talk provides a roadmap for moving ML from a passive analytical tool to an active, intelligent partner in scientific discovery and real-world decision-making, albeit with significant technical and practical challenges still to overcome.

Key Takeaways

• Closing the Loop is Essential for Scientific Discovery: Traditional linear ML pipelines are insufficient for science due to expensive experiments and data scarcity. Iterative, closed-loop ML, where algorithms actively inform and drive experimentation, is critical for accelerating discovery.
• Uncertainty Quantification Guides Exploration: The ability to quantify epistemic uncertainty is fundamental. Methods like Bayesian optimization and Bayesian model-based RL leverage this uncertainty to intelligently balance exploration (gathering information) and exploitation (optimizing known good solutions).
• Bayesian Model-Based RL Scales to Complex Real-World Tasks: By explicitly modeling and quantifying uncertainty in environmental dynamics, Bayesian model-based RL can learn effectively in real-world settings where perfect simulators are unavailable, addressing the "sim-to-real" gap.
• Physics-Aware Priors Boost Sample Efficiency: Incorporating domain knowledge (from simplified models or simulations) as physics-aware priors significantly improves the sample efficiency of real-world learning, enabling complex tasks to be mastered with far fewer physical experiments.
• Active Learning Enhances LLM Performance for Specific Tasks: Even powerful pre-trained Large Language Models (LLMs) benefit from active learning at test-time. By selectively querying data to minimize uncertainty, LLMs can achieve higher accuracy and guide reasoning for specific, high-stakes applications.
• Generative Models Can Be Intelligently Steered: Modern generative models like diffusion models can be guided using principles from optimal control and RL to achieve specific objectives (e.g., optimizing protein properties) or to perform open-ended, unsupervised exploration for novel designs.

About the Speaker(s)

Andreas Krause is a distinguished researcher in machine learning, currently serving as a Professor at ETH Zurich. His extensive research career has consistently focused on the fundamental challenges of information gathering, active learning, Bayesian optimization, and reinforcement learning, particularly in the context of scientific discovery and real-world applications. He has a notable history of impactful collaborations, including pioneering work with Nobel laureate Frances Arnold during his time at Caltech, focusing on protein engineering and directed evolution. Krause's work spans diverse applications, from environmental monitoring and precision agriculture to the tuning of advanced scientific instruments and the design of novel molecules, reflecting his commitment to applying machine learning to drive tangible progress in science and engineering.