In-Context Multi-Objective Optimization

Technical Summary: In-Context Multi-Objective Optimization (TAMO)

Problem Statement
Multi-objective optimization (MOO) is a pervasive challenge in scientific and engineering domains, such as drug discovery and autonomous systems, where practitioners must balance competing objectives (e.g., efficacy vs. toxicity, accuracy vs. cost). The standard approach, Multi-objective Bayesian Optimization (MOBO), relies on fitting probabilistic surrogates (typically Gaussian Processes) and optimizing acquisition functions at each step. However, the authors identify three critical limitations in real-world deployment:

1. Computational Overhead: Each new problem requires training surrogates from scratch and repeatedly optimizing the acquisition function, creating bottlenecks in parallel or time-sensitive loops.
2. Lack of Transferability: Performance is highly sensitive to modeling choices (kernels, likelihoods) and rarely transfers to new problems without re-tuning.
3. Myopia: Most acquisition functions are myopic, optimizing for one-step gains rather than long-horizon Pareto frontier discovery, which often requires multi-step planning.

Methodology: TAMO
The paper introduces TAMO (Task-Agnostic Multi-Objective), a fully amortized, universal policy for multi-objective black-box optimization. Unlike traditional MOBO, TAMO eliminates per-task surrogate fitting and acquisition engineering by mapping the entire observation history directly to the next query via a single neural network forward pass.

• Architecture: TAMO utilizes a transformer backbone designed to be dimension-agnostic on both inputs and outputs. It employs a novel "dimension-aggregating embedder" that maps scalar observations to vectors regardless of the input ($d_x$) or output ($d_y$) dimensions. This allows the model to be pretrained on heterogeneous tasks with varying dimensionalities and transferred to new problems without retraining.
• Training Strategy: The model is trained offline using a two-phase approach:
  1. Prediction Warm-up: The model is first trained on an auxiliary in-context regression task to learn the function landscape and stabilize representations.
  2. Reinforcement Learning (RL): The policy is optimized via RL to maximize the cumulative hypervolume improvement over full trajectories. The state includes the history of observations, the current time step, and the total budget. The reward is the normalized hypervolume, encouraging long-horizon planning rather than myopic gains.
• Inference: At test time, the pretrained model proposes the next design by scoring a candidate set in a single forward pass, conditioned on the entire query history.

Key Contributions

1. Fully Amortized Policy: TAMO is the first end-to-end, dimension-agnostic architecture for multi-objective black-box optimization that operates entirely in-context, removing the need for iterative surrogate fitting.
2. Dimension Agnosticism: The proposed transformer architecture handles varying input and output dimensions, enabling pretraining on diverse synthetic corpora and transfer to real-world tasks with unseen dimensionalities.
3. Long-Horizon Optimization: By optimizing for cumulative hypervolume over full trajectories via RL, TAMO addresses the myopia inherent in standard acquisition functions.

Experimental Results
The authors evaluated TAMO on synthetic GP tasks, standard analytic benchmarks (e.g., Ackley-Rastrigin, Branin-Currin), and real-world problems (Oil Sorbent design, Laser-Plasma acceleration).

• Performance: TAMO achieves competitive or superior Pareto quality (measured by simple regret and hypervolume) compared to strong GP-based baselines (qNEHVI, qNParEGO, qHVKG) and amortized baselines like BOFormer.
• Efficiency: The most significant gain is in inference speed. TAMO reduces proposal time by 50× to 1000× compared to GP-based methods. This is because TAMO avoids the overhead of surrogate refitting and acquisition optimization, performing only a single forward pass per proposal.
• Generalization: The model demonstrates robust generalization to out-of-distribution (OOD) tasks, including those with unseen input/output dimensions (e.g., training on $d_x \in \{1,2\}$ and testing on $d_x=4$) and decoupled observation settings where objectives are measured independently.

Significance and Claims
The paper claims that TAMO establishes a path toward foundation-style, plug-and-play optimizers for scientific discovery. By demonstrating that transformers can perform multi-objective optimization entirely in-context, the work suggests that the field can move away from task-specific surrogate engineering toward universal policies trained on diverse data. The authors position TAMO as a universal engine capable of handling heterogeneous legacy datasets and accelerating closed-loop laboratories, high-throughput campaigns, and time-sensitive workflows where evaluation costs are low or parallelization is high. The work remains modest regarding limitations, noting that performance depends on the composition of the synthetic pretraining corpus and that current inference assumes a discrete candidate pool, which may restrict application in continuous or combinatorial generative design spaces without further architectural adaptation.