STRIDE: Post-Training LLMs to Reason and Refine Bio-Sequences via Edit Trajectories

Resumo Técnico: STRIDE

1. O Problema

A otimização de sequências biológicas discretas (como proteínas e pequenas moléculas) é um desafio central na biologia computacional. O objetivo não é apenas a geração de novo, mas o refinamento orientado a metas: partir de uma sequência viável existente e aplicar um número limitado de edições para melhorar uma propriedade específica (ex: fluorescência, estabilidade, propriedades farmacocinéticas), mantendo restrições sintáticas e estruturais rigorosas (ex: gramática SMILES para moléculas ou restrições de aminoácidos para proteínas).

Os métodos existentes apresentam limitações:

• Modelos de Difusão: Oferecem refinamento progressivo, mas não expõem naturalmente operações de edição discreta controláveis (como inserção/deleção) e têm dificuldade em garantir validade sintética em cada passo intermediário.
• LLMs Autoregressivos Padrão: Embora gerem sequências sintaticamente válidas, tendem a ser miopes (focam apenas no próximo token) e carecem de planejamento explícito de longo prazo para navegar em paisagens de aptidão complexas sob orçamentos de edição estritos.

2. Metodologia: STRIDE

O STRIDE (Sequence Trajectory Refinement via Internalized Denoising Emulation) é um framework de pós-treinamento que reformula a otimização de sequências como um planejamento de trajetória no espaço de edição. Em vez de amostrar estados intermediários ruidosos, o modelo aprende a emitir uma trajetória explícita e executável de edições atômicas (INSERT, DELETE, REPLACE) que transforma a sequência fonte na candidata otimizada.

O treinamento segue um pipeline de duas etapas:

• Etapa I: Ajuste Fino Supervisionado (SFT) com Rastreamento Levenshtein

  • Objetivo: Internalizar a validade e o viés de edição mínima.
  • Processo: Utiliza um algoritmo determinístico de programação dinâmica (alinhamento de Levenshtein) para calcular o caminho de edição mais curto entre uma sequência fonte e um alvo.
  • Backtracing: O caminho ótimo é convertido em um script de edição executável com índices dinâmicos (reavaliados após cada operação).
  • Formato de Treino: O modelo é treinado para primeiro emitir a trajetória de raciocínio (o script de edição) e, em seguida, a sequência final. Isso induz um prior de validade e encoraja trajetórias concisas e não redundantes.

• Etapa II: Alinhamento Funcional via Otimização de Política (GRPO)

  • Objetivo: Alinhar as trajetórias de edição com recompensas de tarefas específicas (ex: aumento de fluorescência, melhoria de QED).
  • Método: Utiliza Group Relative Policy Optimization (GRPO) (ou variantes como GSPO/CISPO) sem um crítico aprendido.
  • Mecanismo: Para cada prompt, o modelo gera um grupo de saídas. A vantagem relativa é calculada com base nas recompensas das sequências finais.
  • Regularização KL: Uma penalidade KL é aplicada em relação à política de referência (SFT) para garantir que o comportamento de edição coerente e válido aprendido na Etapa I não seja perdido (evitando "hacking" de recompensa).
  • Verificação: Se a trajetória emitida não for executável ou não reproduzir a sequência final, a recompensa é zerada, forçando a consistência entre o raciocínio e a ação.

3. Principais Contribuições

1. Otimização de Trajetória de Edição Executável: Formulação do refinamento de sequências biológicas como planejamento sobre trajetórias atômicas (INSERT/DELETE/REPLACE), permitindo edição de comprimento variável consistente com índices.
2. Supervisão via Rastreamento Levenshtein: Proposta de um pipeline determinístico que converte pares alinhados em scripts de edição fundamentados em índices para SFT, induzindo um prior de edição mínima e preservação de validade.
3. Pós-treinamento Alinhado a Recompensas: Uso de otimização de política baseada em grupos com regularização KL para alinhar a geração de trajetórias com objetivos de tarefa, garantindo consistência de análise e execução.
4. Interface de Controle Interpretável: O modelo fornece um "rastro de raciocínio" transparente (o script de edição), permitindo que humanos verifiquem e controlem o processo de otimização passo a passo.

4. Resultados Experimentais

O STRIDE foi avaliado em dois domínios principais:

• Otimização de Fluorescência de Proteínas (TAPE):

  • Cenário de Edição Completa (Inserção/Deleção/Substituição): O STRIDE melhorou a taxa de sucesso de edição variável de 42% (SFT puro) para 89%.
  • Diversidade e Novidade: Aumentou a novidade das sequências geradas de 47% para 97%, superando significativamente modelos de difusão discreta (como EvoDiff) e abordagens Vanilla SFT, que tendiam a gerar mutações repetitivas ou falhar em edições complexas de comprimento variável.

• Otimização Molecular Condicionada a Instruções (MEGA-MolEdit / DrugAssist):

  • Controle e Estabilidade: Em 14 condições de otimização (ex: solubilidade, permeabilidade), o STRIDE (especificamente a variante com GSPO) alcançou a melhor controllabilidade global, reduzindo significativamente o desvio de propriedades não-alvo (off-target drift) em comparação com geração direta ou RL sem trajetórias.
  • Validade: Aumentou a taxa de validade sintática (SMILES válidos) de 75% para 90,9%.
  • Comparação: Embora a geração direta (Vanilla SFT) tenha alta taxa de sucesso inicial, ela sofre de grandes desvios indesejados. O STRIDE equilibra sucesso na tarefa com estabilidade estrutural.

5. Significado e Impacto

O trabalho STRIDE demonstra que raciocínio explícito sobre trajetórias de edição é uma interface prática e superior para a otimização de sequências biológicas discretas.

• Superioridade em Longo Prazo: A abordagem é particularmente eficaz em cenários de "longo horizonte" onde múltiplas edições (especialmente inserções e deleções) alteram os índices subsequentes, exigindo um controle de estado consistente que modelos autoregressivos padrão ou difusão não conseguem manter facilmente.
• Transparência Científica: Ao forçar o modelo a emitir um script de edição executável, o processo de design torna-se interpretável, permitindo validação humana e correção de erros antes da síntese experimental.
• Aplicabilidade: O framework tem potencial para acelerar pipelines de descoberta de fármacos e engenharia de proteínas, reduzindo custos experimentais ao gerar candidatos mais válidos, diversos e otimizados.

O código do projeto está disponível publicamente, facilitando a reprodução e extensão para contextos genômicos mais longos e oráculos de validação experimental.