Collaborative Contextual Bayesian Optimization

1. Problem Statement

The paper addresses the challenge of Contextual Bayesian Optimization (CBO) in a distributed, multi-client setting.

• Contextual BO (CBO): Unlike standard BO, which seeks a single global optimal design $x^*$, CBO aims to learn a mapping $x^*(c) = \arg\max_x f(x, c)$ that provides the optimal design $x$ for any given context $c$ (e.g., patient weight, product specifications). This requires simultaneous exploration across the context space and exploitation within each context.
• The Challenge: CBO is significantly more resource-intensive than standard BO because it must approximate an entire function rather than a single point. In modern applications (e.g., manufacturing, personalized medicine), multiple related but potentially heterogeneous clients (e.g., different machines or hospitals) face similar optimization tasks.
• The Gap: Existing Collaborative BO methods focus on finding a single optimal design per client. There is no established framework for collaboratively learning context-dependent mappings ($x^*(c)$) across multiple clients while handling heterogeneity and privacy constraints.

2. Methodology: CCBO Framework

The authors propose CCBO (Collaborative Contextual Bayesian Optimization), a unified framework enabling multiple clients to jointly learn context-specific optimal designs. The core components are:

A. Surrogate Modeling

Each client $k$ maintains a local Gaussian Process (GP) surrogate model $GP(D_{k,t})$ based on their local dataset $D_{k,t}$. The response function is modeled as $f_k(x, c)$.

B. Disagreement-Driven Switching Mechanism

The algorithm dynamically balances between collaborative learning (using global knowledge) and independent learning (refining local models) via a probability $p_t$ that decays over time.

1. Collaborative Mode (High $p_t$):

   • Clients share their local posterior mean functions $\mu_{k,t-1}$.
   • A global operational model is constructed by averaging these means: $\bar{\mu}_{t-1} = \frac{1}{K}\sum \mu_{k,t-1}$.
   • Selection Strategy: The algorithm identifies contexts where the local client's optimal design differs significantly from the global model's recommendation. It defines a "collaborative improvement potential" $\Delta^{collab}$ and selects the context $c$ that maximizes this discrepancy. The design is then chosen based on the global model.
   • Goal: Correct inaccurate local estimates early in the process by leveraging "strength from others."

2. Independent Mode (Low $1-p_t$):

   • Clients operate independently using Thompson Sampling (TS).
   • A client samples a function $\tilde{f}_{k,t-1}$ from its local posterior GP.
   • It selects the context and design that maximize the potential improvement within this sampled function.
   • Goal: Refine the specific client's model as local data accumulates, reducing bias from the global average.

C. Privacy-Preserving Communication

To address privacy concerns, the paper introduces a mechanism using Random Fourier Features (RFF).

• Instead of sharing raw data or full posterior mean functions, clients approximate their posterior means using a linear model in a randomized feature space: $\mu_k(z) \approx \phi(z)^\top w_k$.
• Clients only share the coefficient vector $w_k$ (a compressed representation) with a central coordinator.
• The coordinator aggregates these weights to form the global model $\bar{w}$, which is broadcast back. This prevents the reconstruction of raw historical data while enabling collaboration.

D. Offline Initialization

The framework supports an "offline" mode where a new client can initialize its search using the historical posterior means of peers without requiring new coordinated experiments, facilitating immediate transfer learning.

3. Key Contributions

1. Unified Framework: The first collaborative approach explicitly designed for CBO, enabling multiple heterogeneous clients to jointly learn the mapping $x^*(c)$.
2. Adaptive Switching: A novel decision mechanism that uses "disagreement" between local and global models to guide information sharing, transitioning smoothly from collaborative to independent operation as data accumulates.
3. Privacy Preservation: A practical communication protocol using RFF approximations that reduces communication costs and protects sensitive client data.
4. Theoretical Guarantees: The authors prove sublinear regret ($O(\sqrt{T})$), demonstrating that the average regret vanishes as the number of iterations increases, even under client heterogeneity.

4. Experimental Results

The authors evaluated CCBO against baselines including Random Sampling (RS), Independent Multi-Task Thompson Sampling (MTS), and Federated Thompson Sampling (FTS).

• Synthetic Benchmarks:

  • Tested on Ackley, Levy, and Hartmann functions with varying context and design dimensions.
  • Homogeneous Settings: CCBO significantly outperformed all baselines, reducing regret in early iterations by effectively pooling information.
  • Heterogeneous Settings: Even with clients having shifted response functions, CCBO consistently outperformed MTS and FTS, demonstrating robustness to client variation.
  • Scalability: Performance improved as the number of clients ($K$) increased, showing the benefit of collaboration.
  • Privacy: The RFF-based approximation achieved performance comparable to exact sharing in later stages, validating the privacy-preserving approach.

• Real-World Application (Hot Rolling):

  • Scenario: Optimizing grain size in a metal hot rolling process.
  • Context: Final sheet thickness ($h_f$).
  • Heterogeneity: Caused by variations in roller radius ($R$) across different machines (clients).
  • Result: CCBO achieved the lowest regret compared to FTS, MTS, and RS, successfully learning context-specific optimal parameters (speed, load, flux) faster than independent methods.

5. Significance and Impact

• Efficiency: CCBO drastically reduces the experimental budget required to learn context-specific optimal designs by leveraging cross-client data, which is critical in expensive real-world applications like manufacturing and clinical trials.
• Scalability: The framework handles client heterogeneity effectively, making it suitable for distributed systems where machines or users have slightly different characteristics.
• Privacy: By integrating RFF approximations, the method offers a viable path for collaborative optimization in privacy-sensitive domains (e.g., healthcare, proprietary manufacturing) without compromising data security.
• Theoretical Foundation: The sublinear regret guarantee provides a rigorous mathematical basis for the method's long-term performance, bridging the gap between practical heuristics and theoretical optimality in continuous contextual settings.

In conclusion, CCBO represents a significant advancement in Bayesian Optimization, moving from single-point optimization to collaborative, context-aware learning, thereby enabling more efficient and scalable decision-making in complex, distributed environments.