Few-Shot Design Optimization by Exploiting Auxiliary Information

Many real-world design problems involve optimizing an expensive black-box function $f(x)$, such as hardware design or drug discovery. Bayesian Optimization has emerged as a sample-efficient framework for this problem. However, the basic setting considered by these methods is simplified compared to real-world experimental setups, where experiments often generate a wealth of useful information. We introduce a new setting where an experiment generates high-dimensional auxiliary information $h(x)$ along with the performance measure $f(x)$; moreover, a history of previously solved tasks from the same task family is available for accelerating optimization. A key challenge of our setting is learning how to represent and utilize $h(x)$ for efficiently solving new optimization tasks beyond the task history. We develop a novel approach for this setting based on a neural model which predicts $f(x)$ for unseen designs given a few-shot context containing observations of $h(x)$. We evaluate our method on two challenging domains, robotic hardware design and neural network hyperparameter tuning, and introduce a novel design problem and large-scale benchmark for the former. On both domains, our method utilizes auxiliary feedback effectively to achieve more accurate few-shot prediction and faster optimization of design tasks, significantly outperforming several methods for multi-task optimization.

💡 Research Summary

This paper addresses a realistic gap in current Bayesian optimization (BO) practice: real‑world experiments often produce rich, high‑dimensional auxiliary data (denoted h(x)) in addition to the scalar objective f(x). Moreover, designers typically have access to a repository of previously solved, related tasks that share the same type of auxiliary observations. The authors formalize a new “few‑shot design optimization with auxiliary information” setting in which (i) each trial of a design x returns a pair F(x) = (f(x), h(x)), and (ii) a multi‑task history {T₁,…,T_N} with corresponding datasets D_i is available for meta‑training. The central challenge is to learn a representation of h(x) that generalizes across tasks so that, given only a few observed (x, f(x), h(x)) tuples for a new task, the algorithm can accurately predict f(x) for unseen designs and guide BO efficiently.

To meet this challenge, the authors propose a transformer‑based neural model P_θ that takes a few‑shot context C = {(x_k, f(x_k), h(x_k))}_{k=1}^K and a target design x′ as input and outputs a probabilistic prediction P_θ(f | C, x′) in the form of a Gaussian (mean μ and variance σ²). The model is trained in a meta‑learning fashion on the entire task collection. Training optimizes a joint objective consisting of (a) the log‑likelihood of the observed scalar rewards and (b) a reconstruction loss for the auxiliary data h(x). The reconstruction term forces the transformer to learn a compact latent embedding of h(x) that preserves information useful for predicting f(x). Task identifiers are injected as special tokens, allowing the model to capture task‑specific nuances while sharing a common representation of h.

During optimization of a new task, the learned P_θ serves as a surrogate model in a standard BO loop. Because P_θ produces a predictive variance, classic acquisition functions such as Expected Improvement (EI) or Probability of Improvement (PI) can be applied directly. At each iteration the acquisition function selects a promising design x_t, the true experiment is run to obtain F(x_t) = (f(x_t), h(x_t)), and the new observation is appended to the context C. No fine‑tuning of P_θ is required; the model generalizes from the meta‑training tasks to the unseen test task.

The authors evaluate the approach on two challenging domains. First, they construct a large‑scale robotic gripper benchmark: 1,000 distinct objects, 4.3 million evaluated gripper designs, and high‑dimensional tactile sensor streams (≈2,000 dimensions) as h(x). The goal is to maximize grasp stability f(x). In this setting, the proposed method reaches near‑optimal grasp stability after only 5–10 evaluations, outperforming multi‑task Gaussian‑process baselines by roughly 30 % fewer evaluations and discovering creative design strategies (e.g., exploiting object gaps). Second, they apply the method to neural network hyper‑parameter tuning using the LCBench suite, treating per‑epoch learning curves as h(x). With just five initial hyper‑parameter configurations, the model predicts the final validation accuracy accurately and identifies high‑performing configurations with less than half the evaluations required by state‑of‑the‑art multi‑task GP, meta‑RL, or ensemble methods.

Key insights from the experiments include: (1) auxiliary information provides a strong signal about the underlying dynamics of a design, enabling rapid inference of the objective even with very few samples; (2) a transformer‑based meta‑learner can capture the shared structure of heterogeneous, high‑dimensional h(x) across many tasks, achieving better scalability than GP‑centric approaches; (3) integrating the learned surrogate into a conventional BO loop yields a simple yet powerful pipeline that retains the exploration‑exploitation balance of BO while benefiting from rich side‑information.

In summary, the paper makes three major contributions: (i) defining a new, practically motivated optimization setting that incorporates auxiliary data and multi‑task histories; (ii) introducing a novel transformer‑based few‑shot predictor that learns to encode auxiliary information for accurate reward prediction; and (iii) providing a new robotic gripper benchmark together with extensive empirical validation showing substantial gains over existing multi‑task BO methods. The work pushes AI‑driven design toward more realistic, information‑rich environments and opens avenues for future research on other modalities of auxiliary data (e.g., images, graphs) and online, real‑time feedback loops.