Training Diffusion Language Models for Black-Box Optimization

1. O Problema: Otimização de Caixa-Preta Offline (BBO)

O artigo aborda o problema da Otimização de Caixa-Preta Offline (BBO), comum em áreas como design robótico, síntese de DNA e descoberta de materiais.

• Objetivo: Descobrir novos designs ($x$) que maximizem uma função objetivo desconhecida $f(x)$ (ex: afinidade de ligação), utilizando apenas um conjunto de dados offline limitado de pares (design, rótulo/valor).
• Desafios Principais:
  1. Escassez de Dados: Em aplicações reais, obter rótulos (resultados experimentais) é caro e demorado, resultando em conjuntos de dados pequenos (ex: $N=500$).
  2. Incerteza Epistêmica: Métodos tradicionais (como modelos substitutos ou GPs) sofrem com incerteza quando exploram regiões fora da distribuição dos dados de treinamento.
  3. Limitação dos Modelos Autoregressivos (AR): Trabalhos recentes usam Grandes Modelos de Linguagem (LLMs) autoregressivos (leitura da esquerda para a direita). No entanto, muitos problemas de design (como sequências de DNA) possuem dependências bidirecionais fortes (um elemento depende tanto do prefixo quanto do sufixo), o que modelos AR não capturam bem.

2. Metodologia Proposta: DiBO

Os autores propõem o DiBO (Diffusion Black-Box Optimization), um framework que adapta Modelos de Linguagem Difusivos (Diffusion LLMs) para o domínio de BBO. A abordagem consiste em três pilares principais:

A. Adaptação de Domínio (Domain Adaptation - DA)

Existe uma lacuna entre o pré-treinamento dos LLMs (texto natural) e os sinais heterogêneos do BBO (instruções em texto, designs estruturados e rótulos numéricos).

• Corpus Unificado: Os autores constroem um corpus unificado de "Prompt-Resposta".
• Tokens Delimitadores: Para marcar explicitamente as fronteiras semânticas, são introduzidos tokens especiais:
  • |design-start| e |design-end| para envolver os designs.
  • |label-start| e |label-end| para envolver os rótulos.
• Loss Conjunta: O modelo é treinado para prever tokens mascarados tanto no prompt quanto na resposta simultaneamente, aprendendo a relação entre a descrição da tarefa, os exemplos offline e o design otimizado.

B. Pós-Treinamento em Duas Etapas

Para alinhar a geração do modelo com designs de alto valor (rótulos altos), propõe-se um pipeline de duas etapas:

1. Ajuste Fino Supervisionado (SFT):
   • O modelo é ajustado no corpus unificado, focando apenas na reconstrução dos tokens da resposta (o novo design).
   • Isso introduz um viés indutivo para gerar designs que superam os exemplos fornecidos no prompt.
2. Aprendizado por Reforço (RL):
   • Um conjunto de dados de RL é construído onde a recompensa ($r$) é definida pela melhoria do rótulo: $r(q, o) = y(o) - y(q)$, onde $y(o)$ é o rótulo do design gerado e $y(q)$ é o melhor rótulo no prompt.
   • Otimização Eficiente: Em vez de usar o processo de denoising iterativo completo (que é lento), os autores utilizam uma aproximação de log-probabilidade em um único passo para calcular a perda de RL, tornando o treinamento viável.
   • O modelo aprende a maximizar a melhoria incremental em relação ao contexto do prompt.

3. Contribuições Chave

1. Adaptação de Diffusion LLMs para BBO: Primeira aplicação de modelos difusivos (com capacidade de modelagem bidirecional) para otimização de caixa-preta offline, superando as limitações direcionais dos LLMs autoregressivos.
2. Framework de Adaptação de Domínio: Introdução de tokens delimitadores e um corpus unificado para integrar texto, designs e rótulos heterogêneos, permitindo que o modelo entenda a estrutura do problema de otimização.
3. Pipeline de Pós-Treinamento Híbrido: Combinação de SFT (para alinhamento inicial e estabilidade) e RL (para integração de feedback de recompensa granular), demonstrando que ambas as etapas são complementares e necessárias.
4. Desempenho SOTA: Demonstração de resultados state-of-the-art em configurações de dados pequenos no benchmark Design-Bench.

4. Resultados Experimentais

Os experimentos foram realizados no Design-Bench, cobrindo tarefas discretas (sequências de DNA: TF Bind 8 e TF Bind 10) e contínuas (morfologia de robôs: Ant e D'Kitty).

• Desempenho Geral: O DiBO alcançou o 1º lugar na classificação média e mediana entre todos os métodos comparados (incluindo métodos baseados em surrogate, VAEs, GANs, LLMs autoregressivos e outros métodos difusivos).
• Comparação Específica:
  • No tarefa TF Bind 8, o DiBO superou o melhor baseline por uma margem significativa (>0.108 na pontuação normalizada).
  • No Ant Morphology, obteve a melhor pontuação absoluta (0.944).
  • Em TF Bind 10, obteve a melhor pontuação (0.755).
• Análise de Ablação:
  • A remoção dos tokens delimitadores causou queda de desempenho em todas as etapas, especialmente na adaptação de domínio.
  • A remoção de qualquer uma das três etapas (DA, SFT ou RL) resultou em queda de desempenho, confirmando que o pipeline completo é necessário. O RL sozinho falhou em gerar saídas válidas em tarefas complexas sem o pré-alinhamento do SFT.
  • A construção do contexto baseada em similaridade (escolher exemplos de prompt próximos ao design alvo) foi crucial para o sucesso, especialmente em tarefas de DNA.

5. Significado e Conclusão

O trabalho estabelece os Diffusion LLMs como uma nova e robusta paradigma para otimização offline em espaços de design de alta dimensão.

• Vantagem Principal: A capacidade de modelagem bidirecional permite capturar dependências estruturais complexas que modelos autoregressivos ignoram.
• Eficiência: O método é eficaz mesmo com conjuntos de dados muito pequenos (500 amostras), reduzindo a necessidade de experimentos físicos caros e iterativos.
• Impacto: Oferece uma ferramenta geral para descoberta científica em biologia, robótica e ciência de materiais, embora os autores alertem para a necessidade de supervisão humana para evitar o uso indevido (ex: design de agentes biológicos nocivos).

Em resumo, o DiBO preenche a lacuna entre a linguagem natural e a otimização estruturada, utilizando a flexibilidade dos modelos difusivos e um treinamento cuidadoso em duas etapas para superar o estado da arte em otimização de caixa-preta.