Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebysheff Scalarization

This paper introduces STOMP, a novel offline reinforcement learning algorithm that employs smooth Tchebysheff scalarization to effectively align large language models with multiple conflicting objectives, demonstrating superior performance over state-of-the-art baselines in multi-objective protein engineering tasks.

The Gist

Imagine you are a chef trying to create the perfect dish. You have a recipe book (your AI model), but you want to tweak it to satisfy three different critics at once:

1. The Flavor Critic: Wants it to taste amazing.
2. The Health Critic: Wants it to be low in calories.
3. The Budget Critic: Wants it to be cheap to make.

The problem? These goals often fight each other. A dish that tastes amazing might be full of butter (bad for health) or use expensive truffles (bad for budget). In the world of Artificial Intelligence, this is called Multi-Objective Reinforcement Learning. The goal isn't to find one "perfect" dish, but to find the Pareto Front: the set of all the best possible compromises where you can't improve one thing without making another worse.

The Old Way: The "Weighted Average" Mistake

For a long time, AI researchers tried to solve this by giving the critics a single score. They'd say, "Okay, let's give Flavor 50% weight, Health 30%, and Budget 20%, and just add them up."

Think of this like a linear ruler. If you try to measure a curved mountain range with a straight ruler, you miss all the valleys and peaks. In math terms, this "linear ruler" (called linear scalarization) fails to find the best solutions in the "non-convex" areas—the tricky, curved parts of the compromise curve where the most interesting solutions live. It forces the AI to pick a boring, middle-of-the-road solution and ignores the truly unique, high-performing compromises.

The New Solution: STOMP

The authors of this paper introduce a new algorithm called STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences).

Here is the simple analogy for how STOMP works:

Instead of adding the scores together, STOMP acts like a smart, flexible rubber band.

1. The Rubber Band: Imagine a rubber band stretched around a group of balloons (the different goals). The rubber band naturally hugs the outer edge of the balloons, finding the best shape that touches all of them, even if the shape is weird or curved.
2. The "Smooth" Part: The math behind this (Tchebysheff scalarization) is usually very sensitive to the scale of the numbers. If one critic uses a score of 1–10 and another uses 1–1,000,000, the math breaks. STOMP fixes this by standardizing the scores. It looks at the history of every dish made in the kitchen, figures out what a "normal" score looks like for each critic, and adjusts the rubber band accordingly. It doesn't just look at the raw numbers; it looks at how rare or common a score is.

The Experiment: Cooking with Proteins

To prove this works, the researchers didn't just talk about food; they cooked with proteins.

Proteins are the building blocks of life. Scientists want to design proteins that do two or three things at once, like:

• DHFR: A protein that fights bacteria (Activity) but doesn't get stopped by a common drug (Stability).
• PbrR: A protein that grabs lead (good for cleaning water) but ignores zinc (so it doesn't get confused).
• Alpha-Amylase: A protein that works fast, is stable, and is easy to produce.

They took three different AI "chefs" (Protein Language Models) and taught them to cook using:

1. The old "Linear Ruler" method.
2. A slightly better "Z-score" method.
3. The new STOMP method.

The Results

When they tested the results:

• The Old Methods: Often got stuck in the middle, creating proteins that were "okay" at everything but great at nothing. They missed the "curved" solutions.
• STOMP: Successfully found the "sweet spots." It generated proteins that were significantly better at balancing these conflicting goals. In 8 out of 9 tests, STOMP created the best possible set of compromises (the highest "hypervolume").

Why This Matters

This isn't just about proteins. This is a new way of teaching AI to handle conflicting desires.

• Chatbots: Making them helpful and harmless and concise.
• Self-driving cars: Making them fast and safe and energy-efficient.
• Image Generators: Making images high-quality and accurate to the prompt and diverse.

In a nutshell: The paper says, "Stop trying to average out your goals with a straight line. Use a flexible, smart rubber band (STOMP) that understands the unique shape of your problems, so you can find the best possible solutions that everyone can agree on."

Technical Summary

1. Problem Statement

The paper addresses the challenge of Multi-Objective Offline Reinforcement Learning (RL) for aligning Large Language Models (LLMs), specifically Protein Language Models (PLMs).

• The Core Issue: Real-world applications often require optimizing multiple conflicting objectives simultaneously (e.g., protein activity vs. stability, or chatbot helpfulness vs. harmlessness).
• Limitation of Current Methods: The standard approach, Linear Reward Scalarization (taking a weighted average of rewards), is mathematically proven to fail in recovering solutions from non-convex regions of the Pareto front. These non-convex regions often contain the most valuable trade-offs between conflicting objectives.
• The Gap: While Smooth Tchebysheff Scalarization (STS) exists to solve non-convex multi-objective optimization, applying it directly to reward vectors in RL is difficult because STS is highly sensitive to the scale of individual rewards, requiring difficult-to-tune hyperparameters.

2. Methodology: STOMP

The authors propose Smooth Tchebysheff Optimization of Multi-Objective Preferences (STOMP), a novel offline RL algorithm. Instead of scalarizing the rewards directly, they frame the multi-objective RL problem itself as an optimization problem to be scalarized.

Key Technical Innovations:

1. Reframing the Optimization Problem:

   • Standard Multi-Objective RL seeks to maximize a vector of expected rewards $[E[r_1], \dots, E[r_k]]$ subject to a KL-divergence constraint from a reference policy $\pi_0$.
   • The authors apply Smooth Tchebysheff Scalarization (STS) directly to this optimization objective. This transforms the problem into minimizing a scalarized upper bound of the KL divergence from multiple optimal policies (one for each reward).

2. Distribution-Relative Reward Standardization:

   • A critical insight is that STS requires rewards to be on comparable scales. The authors derive a policy-independent reward scalarization ($R_{ST}^\lambda$) that dynamically standardizes rewards based on their observed distributions in the offline dataset.
   • Instead of simple min-max or mean-variance scaling, they utilize partition functions ($Z_i(x)$) derived from the training data to create "distribution-relative rewards" ($\rho_i$).
   • This approach automatically handles skewed reward distributions (common in biological data) without manual hyperparameter tuning, ensuring the scalarization covers the full Pareto front.

3. The STOMP Loss Function:

   • Building on Direct Preference Optimization (DPO) and OffsetDPO, the authors define a new loss function.
   • Preference Pairing: They use the policy-independent scalarized reward ($R_{ST}^\lambda$) to identify "winner" ($y_w$) and "loser" ($y_l$) pairs to ensure stability during training.
   • Optimization Objective: The loss minimizes the difference between the log-likelihood ratio of the winner/loser and the difference in their scalarized rewards, similar to OffsetDPO but with the STS-derived reward.
   • Regularization: Includes a length-averaged negative log-likelihood term for winners to prevent likelihood degradation, a common issue in offline RL.

3. Key Contributions

• Theoretical Formulation: The paper derives a principled way to apply Smooth Tchebysheff Scalarization to Multi-Objective RL by framing the RL problem itself as the target for scalarization, rather than just the rewards.
• Algorithm (STOMP): Introduction of a robust offline RL algorithm that extends Direct Preference Optimization to the multi-objective setting. It uniquely handles reward scaling via distribution-relative standardization, eliminating the need for manual reward scaling hyperparameters.
• Empirical Validation: Comprehensive evaluation on protein engineering tasks, demonstrating that STOMP outperforms linear scalarization baselines and naive STS implementations.

4. Experimental Results

The authors evaluated STOMP on three protein engineering datasets using three different autoregressive Protein Language Models (ProGen3-3B, ProGen-RA-3B, and ProGen-RA-10B).

• Datasets:
  • DHFR: Maximizing activity with and without an inhibitor (uncorrelated objectives).
  • PbrR: Maximizing lead binding while minimizing zinc binding (strongly negatively correlated).
  • $\alpha$-Amylase: Maximizing activity, expression, and thermostability (positively correlated).
• Metrics: The primary metric was the Hypervolume of the Pareto front (the volume of space dominated by the set of solutions relative to a reference point).
• Offline Off-Policy Evaluation:
  • STOMP achieved the highest hypervolume in 8 out of 9 settings.
  • In the 9th setting, it was second-best, achieving 98.7% of the best performance.
  • Baselines (DPO-Lin, ODPO-Lin, ODPO-STZ) showed inconsistent performance across different datasets and models.
• Generative Evaluation:
  • Models aligned with STOMP generated sequences with the highest expected hypervolumes in 8 out of 9 settings when sampled and evaluated via Gaussian Process reward models.
  • Notably, STOMP showed superior performance in generating diverse, high-quality proteins in the challenging PbrR dataset (highly negative correlation), where linear methods failed to find the optimal trade-offs.

5. Significance and Impact

• Solving Non-Convexity: STOMP successfully recovers solutions in non-convex regions of the Pareto front, which linear scalarization provably cannot do. This is crucial for applications where the best solutions are complex compromises between conflicting goals.
• Robustness: The method is domain-agnostic. While tested on proteins, the framework applies to any multi-objective alignment task, including chatbots (helpfulness vs. safety) and text-to-image generation (quality vs. prompt adherence).
• Practicality: By deriving a reward scalarization that standardizes rewards based on data distributions, STOMP removes the need for tedious manual tuning of reward scales, making multi-objective RL more accessible and stable for practitioners.
• Future Directions: The authors suggest this approach can be extended to online RL and other KL-constrained RL formulations, potentially revolutionizing how AI models are aligned with complex, multi-faceted human or biological preferences.