Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebysheff Scalarization

Este artigo apresenta o STOMP, um novo algoritmo de aprendizado por reforço offline multiobjetivo que utiliza a escalarização Tchebysheff suave para superar as limitações da escalarização linear, demonstrando desempenho superior na otimização de proteínas com múltiplos atributos.

O essencial

Imagine que você é um chef de cozinha talentoso, mas está tentando criar o prato perfeito para um cliente muito exigente. O problema é que o cliente quer coisas que, na culinária, normalmente se excluem: ele quer que o prato seja extremamente saboroso (como um bolo de chocolate), mas também super saudável (como uma salada de folhas verdes) e muito barato de fazer.

Se você tentar apenas somar esses desejos (digamos, 50% sabor + 50% saúde), você provavelmente vai acabar com um prato medíocre: um bolo sem graça ou uma salada sem graça. É assim que a maioria dos modelos de Inteligência Artificial (especialmente os que geram texto ou proteínas) funciona hoje: eles tentam "somar" objetivos conflitantes e acabam perdendo o equilíbrio.

Este artigo apresenta uma nova receita chamada STOMP (uma sigla engraçada, mas que significa "Otimização Suave de Preferências Múltiplas de Tchebysheff") para resolver esse problema.

Aqui está a explicação passo a passo, usando analogias do dia a dia:

1. O Problema: A "Soma" não funciona

Antes, os cientistas usavam uma técnica chamada "escalarização linear". Imagine que você tem um controle remoto com vários botões de volume para diferentes músicas tocando ao mesmo tempo. Se você apenas aumentar o volume geral, algumas músicas ficam abafadas e outras estridentes.
No mundo das proteínas (que são como "instruções" biológicas para o corpo), isso significava tentar criar uma proteína que fosse boa em várias coisas ao mesmo tempo (ex: matar vírus, mas não fazer mal ao corpo). A abordagem antiga falhava em encontrar as soluções "perfeitas" onde você consegue o melhor de dois mundos sem sacrificar nada demais. Elas só encontravam soluções óbvias e chatas.

2. A Solução: O "Chef de Equilíbrio" (STOMP)

Os autores criaram o STOMP. Em vez de apenas somar os desejos do cliente, o STOMP age como um chef mestre que entende de equilíbrio.

• A Analogia do "Mapa de Tesouros": Pense nos objetivos (sabor, saúde, preço) como montanhas em um mapa. O objetivo é encontrar todos os picos onde você não pode subir mais em uma montanha sem descer na outra. Isso é chamado de "Frente de Pareto".
• O Truque do STOMP: Em vez de tentar escalar uma montanha única (somando tudo), o STOMP usa uma técnica matemática inteligente (chamada Tchebysheff Suave) que olha para o "pior" objetivo em cada tentativa e tenta melhorá-lo, sem deixar os outros caírem. É como se o chef dissesse: "Ok, o prato está ótimo no sabor, mas a saúde está baixa. Vamos ajustar a receita para subir a saúde, mas sem derrubar o sabor".

3. A "Calibragem" Inteligente

Um dos maiores problemas ao tentar equilibrar coisas diferentes é que elas têm escalas diferentes.

• Exemplo: Se você quer equilibrar "número de erros de digitação" (que vai de 0 a 10) com "tempo de leitura" (que vai de 1 a 1000 minutos), você não pode simplesmente somar 5 erros com 500 minutos. O tempo vai dominar tudo.
• O que o STOMP faz: Ele olha para os dados que já tem (proteínas que já foram testadas em laboratório) e "padroniza" tudo. Ele transforma os números em uma escala comum, como se estivesse convertendo Dólar, Euro e Real para uma moeda única antes de somar. Isso permite que o modelo compare "maçãs com laranjas" de forma justa.

4. O Teste: Cozinhar Proteínas

Os autores testaram essa nova técnica em proteínas (que são como blocos de construção da vida). Eles usaram modelos de linguagem (como o ChatGPT, mas treinado em biologia) para criar novas proteínas.

• O Desafio: Criar uma proteína que seja muito ativa (faça seu trabalho rápido) mas também muito estável (não quebre fácil).
• O Resultado: O STOMP conseguiu criar proteínas que eram melhores em 8 de 9 cenários diferentes comparado aos métodos antigos. Foi como se o chef tivesse descoberto um novo tipo de bolo que é ao mesmo tempo delicioso, saudável e barato, algo que os outros chefs diziam ser impossível.

Resumo em uma frase

O STOMP é um novo algoritmo que ensina a Inteligência Artificial a fazer "negócios" perfeitos entre objetivos conflitantes (como criar uma proteína que seja forte e segura ao mesmo tempo), evitando os erros de cálculo que os métodos antigos cometiam ao tentar apenas "somar" as qualidades.

Por que isso importa?
Isso pode acelerar a descoberta de novos medicamentos, enzimas para limpar o meio ambiente e materiais mais fortes, porque permite que os cientistas usem a IA para encontrar soluções criativas e equilibradas que humanos ou métodos antigos não conseguiam ver.

Resumo técnico

Título: Aprendizado por Reforço Offline Pareto-Ótimo via Escalarização Suave de Tchebysheff (STOMP)

1. O Problema

O alinhamento de Grandes Modelos de Linguagem (LLMs) e modelos de linguagem de proteínas (PLMs) com preferências humanas ou biológicas é frequentemente realizado através de Aprendizado por Reforço (RL) offline.

• Limitação Atual: A maioria das abordagens existentes utiliza escalarização linear (uma média ponderada de múltiplas recompensas) para transformar um problema multi-objetivo em um único objetivo.
• O Desafio: Em cenários do mundo real (ex: engenharia de proteínas onde se busca alta atividade e especificidade simultaneamente, ou chatbots que devem ser úteis e seguros), os objetivos são frequentemente conflitantes. A escalarização linear é matematicamente incapaz de recuperar soluções em regiões não convexas da frente de Pareto. Isso significa que as melhores compensações entre objetivos conflitantes (que muitas vezes são as mais desejáveis) são perdidas.
• O Objetivo: Desenvolver um algoritmo de RL offline que possa otimizar múltiplos objetivos conflitantes simultaneamente, garantindo a cobertura de toda a frente de Pareto, incluindo regiões não convexas.

2. Metodologia

Os autores propõem uma nova formulação que trata o próprio problema de RL multi-objetivo como um problema de otimização a ser escalarizado, em vez de escalarizar as recompensas diretamente de forma heurística.

• Escalarização Suave de Tchebysheff (STS): O trabalho baseia-se na técnica STS, introduzida por Lin et al., que supera as limitações da escalarização linear. A STS é definida como:
  $$\min_z \tau \log \sum_{i=1}^k \exp\left(\frac{\lambda_i (L_i(z) - L_i^*)}{\tau}\right)$$
  onde $\tau$ é um parâmetro de suavização. Isso permite encontrar todas as soluções na frente de Pareto via otimização baseada em gradiente.

• Desafio da Escala: A aplicação direta da STS é sensível às escalas individuais das recompensas. Se as recompensas tiverem distribuições diferentes (ex: uma com média 0 e outra com média 1000), a escalarização falha.

• A Inovação (STOMP): Os autores derivam o algoritmo STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences) com as seguintes características principais:

  1. Formulação via Otimização Multi-Objetivo: Eles reencenam o problema de RL com restrição de KL (ou MaxEnt) como um problema de minimização de múltiplas divergências KL em relação a políticas ótimas individuais.
  2. Escalarização de Recompensa Baseada em Distribuição: Em vez de usar hiperparâmetros manuais para escalonar as recompensas, o STOMP padroniza dinamicamente as recompensas com base em suas distribuições observadas no conjunto de dados offline.
  3. Funções de Partição ($Z_i$): O método utiliza funções de partição estimadas a partir dos dados para criar recompensas relativas à distribuição ($\rho_i$). Isso garante que recompensas raras (outliers) sejam penalizadas adequadamente sem distorcer o aprendizado, superando métodos de normalização simples (como min-max ou z-score).
  4. Extensão do DPO: O algoritmo estende o Direct Preference Optimization (DPO) e o OffsetDPO para o cenário multi-objetivo, definindo uma nova função de perda que incorpora a escalarização suave de Tchebysheff.

3. Contribuições Chave

1. Novo Algoritmo (STOMP): Introdução de um algoritmo de RL offline multi-objetivo que deriva a escalarização de recompensa a partir da formulação de otimização, garantindo estabilidade e cobertura teórica da frente de Pareto.
2. Padronização Dinâmica: Uma técnica para normalizar recompensas heterogêneas baseada em suas distribuições empíricas, eliminando a necessidade de ajuste manual de escalas e permitindo que o algoritmo navegue por trade-offs complexos.
3. Validação em Engenharia de Proteínas: Aplicação prática em três modelos de linguagem de proteínas (ProGen3-3B, ProGen-RA-3B, ProGen-RA-10B) em três conjuntos de dados laboratoriais distintos (DHFR, PbrR, $\alpha$-Amilase), demonstrando a eficácia em cenários biológicos reais.

4. Resultados Experimentais

Os autores avaliaram o STOMP comparando-o com baselines de última geração (DPO-Lin, ODPO-Lin e ODPO-STZ) em nove configurações (combinações de modelos e datasets).

• Métrica Principal: O volume hipervolume (hypervolume) das soluções não dominadas, que mede a qualidade e a cobertura da frente de Pareto.
• Desempenho Offline (Off-Policy): O STOMP alcançou o maior volume hipervolume em 8 de 9 configurações. Na nona configuração, ficou em segundo lugar, mas com desempenho muito próximo ao melhor método (98,7%).
• Desempenho Generativo: Ao gerar novas sequências de proteínas e avaliá-las com modelos de recompensa (Gaussian Process), o STOMP novamente superou ou empatou com os baselines na maioria dos cenários, especialmente em tarefas de geração com amostras pequenas.
• Robustez: Enquanto os métodos baseados em escalarização linear ou z-score simples (ODPO-STZ) mostraram desempenho inconsistente dependendo do dataset e do modelo base, o STOMP demonstrou robustez consistente em todos os cenários.

5. Significado e Impacto

• Avanço Teórico: O trabalho demonstra que a escalarização linear é insuficiente para RL multi-objetivo complexo e que a escalarização suave de Tchebysheff, quando aplicada corretamente à estrutura de RL (e não apenas às recompensas brutas), resolve o problema da cobertura de regiões não convexas.
• Aplicabilidade Prática: O STOMP oferece uma ferramenta poderosa para otimização de proteínas multi-atributo, permitindo que pesquisadores projetem moléculas que equilibrem propriedades conflitantes (ex: estabilidade vs. atividade) de forma mais eficaz do que métodos anteriores.
• Generalidade: Embora testado em proteínas, o método é agnóstico ao domínio e pode ser aplicado a qualquer tarefa de alinhamento multi-objetivo, incluindo chatbots (utilidade vs. segurança) e geração de imagens (qualidade vs. aderência ao prompt).

Em resumo, o STOMP representa um avanço significativo na capacidade de alinhar modelos generativos a múltiplos objetivos simultaneamente, superando as limitações fundamentais dos métodos de escalarização linear tradicionais.