Resumo Técnico: Impulsionando as Fronteiras Utilidade-Diversidade Biomolecular com Otimização Relativa de Política de Supergroup

1. Declaração do Problema

A geração biomolecular visa produzir candidatos (moléculas pequenas ou proteínas) que satisfaçam objetivos químicos ou biológicos específicos. Embora o Aprendizado por Reforço (RL) seja um framework padrão para o pós-treinamento de geradores pré-treinados a fim de maximizar a utilidade (por exemplo, similaridade a fármacos, escores de acoplamento, estabilidade), a otimização apenas para utilidade frequentemente leva ao colapso de modos, onde o gerador concentra a massa de probabilidade em uma família estreita de candidatos de alta pontuação. Por outro lado, manter a diversidade é desafiador porque a diversidade é uma propriedade de nível de conjunto (definida sobre coleções de amostras), enquanto a otimização padrão de política atualiza trajetórias individuais.

Métodos existentes de RL que promovem diversidade frequentemente dependem de proxies indiretos, como penalidades de novidade em relação a amostras históricas ou filtros baseados em memória. Essas abordagens podem penalizar excessivamente modos úteis de alta densidade ou induzir deriva distribucional. O desafio central abordado por este artigo é: Como podemos otimizar diretamente a diversidade de conjuntos de amostras como um objetivo de primeira classe, ao mesmo tempo em que atribuímos crédito útil a candidatos gerados individualmente?

2. Metodologia: Otimização Relativa de Política de Supergroup (SGRPO)

Os autores propõem a Otimização Relativa de Política de Supergroup (SGRPO), um framework flexível inspirado na Otimização Relativa de Política de Grupo (GRPO) que otimiza diretamente o trade-off entre utilidade em nível de trajetória e diversidade em nível de conjunto.

Mecanismo Central

O SGRPO opera amostrando múltiplos conjuntos de candidatos (grupos) sob a mesma condição e comparando-os para redistribuir recompensas de diversidade. O processo envolve quatro etapas-chave:

1. Amostragem de Supergroup na Mesma Condição:
   Para uma dada condição $C$ (por exemplo, um bolso proteico específico ou nenhuma condição), a política amostra $M$ grupos, cada um contendo $K$ trajetórias geradas independentemente. Esta coleção $S(C) = \{G_1, \dots, G_M\}$ forma um "supergroup" de $N=MK$ candidatos. As comparações são restritas a este supergroup para evitar confundir a qualidade da política com a dificuldade intrínseca de diferentes condições.

2. Pontuação de Diversidade em Nível de Grupo:
   Cada grupo $G_m$ recebe uma pontuação de diversidade $R_m = D(G_m)$ baseada em uma métrica especificada pelo usuário (por exemplo, distância interna de Tanimoto para moléculas, similaridade de Levenshtein para proteínas). O sinal relativo de diversidade do grupo é computado via uma comparação leave-one-out (deixar um de fora) contra outros grupos no supergroup:
   $$A_m^{grp} = R_m - \frac{1}{M-1}\sum_{h \neq m} R_h$$
   Um $A_m^{grp}$ positivo indica que o grupo é mais diverso do que suas alternativas.

3. Redistribuição Consciente do Conjunto (Contribuição Leave-One-Out):
   Para bridar a lacuna entre sinais em nível de conjunto e atualizações de política individuais, o SGRPO redistribui a recompensa de diversidade do grupo para trajetórias individuais com base em sua contribuição para a diversidade do grupo.

   • Para cada trajetória $x_{m,i}$, uma contribuição leave-one-out é calculada: $c_{m,i} = D(G_m) - D(G_m \setminus \{x_{m,i}\})$.
   • Essas contribuições são padronizadas e usadas para formar pesos softmax conscientes de sinal.
   • Se um grupo é altamente diverso ($A_m^{grp} > 0$), o sinal positivo é concentrado em candidatos com altas contribuições de diversidade. Se um grupo é menos diverso ($A_m^{grp} < 0$), o sinal negativo é concentrado em candidatos com baixas contribuições.
   • Isso garante que a recompensa média do grupo seja preservada, ao mesmo tempo em que atribui crédito às moléculas específicas que impulsionam a diversidade.

4. Atualização de Política Relativa ao Supergroup:
   A recompensa final para uma trajetória combina sua utilidade individual $r_{m,i}$ e a recompensa de diversidade redistribuída $\tilde{R}_{m,i}$:
   $$\bar{r}_{m,i} = (1-\lambda)r_{m,i} + \lambda \tilde{R}_{m,i}$$
   onde $\lambda$ controla o trade-off utilidade-diversidade. A função de vantagem é computada como uma linha de base leave-one-out sobre todas as trajetórias no supergroup, e a política é atualizada usando um objetivo estilo PPO com clipping e regularização KL.

3. Contribuições Principais

• Otimização Direta em Nível de Conjunto: O SGRPO é o primeiro framework a otimizar diretamente a diversidade em nível de conjunto como um objetivo primário na geração biomolecular, em vez de depender de proxies de novidade indiretos ou penalidades dependentes de histórico.
• Design Desacoplado: O método é desacoplado de arquiteturas específicas de geradores (autoregressivos ou de difusão), recompensas de utilidade ou métricas de diversidade, permitindo instanciação com vários otimizadores estilo GRPO (por exemplo, GRPO padrão para modelos autoregressivos, GRPO Acoplado para difusão discreta).
• Mecanismo de Atribuição de Crédito: A introdução de um mecanismo de contribuição de diversidade leave-one-out permite que o sistema atribua crédito a candidatos individuais que genuinamente apoiam a diversidade do conjunto, resolvendo o problema de otimizar uma propriedade de conjunto com atualizações individuais.
• Robustez ao Tamanho do Grupo: O método permanece eficaz mesmo com tamanhos de grupo pequenos, tornando-o computacionalmente viável para o design biomolecular em grande escala.

4. Resultados Experimentais

Os autores avaliaram o SGRPO em três configurações distintas de geração biomolecular:

1. Design de Moléculas Pequenas De Novo: Usando GenMol (difusão discreta).
2. Design de Moléculas Pequenas Baseado em Bolso: Usando GenMol-P (difusão condicionada a bolso).
3. Design de Proteínas De Novo: Usando ProGen2 (autoregressivo).

Métricas de Desempenho:
A avaliação focou na Fronteira de Pareto Utilidade-Diversidade, medida por:

• Hypervolume (HV): A área dominada pelos pontos operacionais não dominados.
• Distância ao Ponto Ideal (DIP): Distância para a melhor utilidade e diversidade teóricas.
• Indicador R2: Defasagem média ponderada do pior caso.

Descobertas:

• Expansão da Fronteira: Em todas as tarefas e varreduras de decodificação, o SGRPO expandiu consistentemente a fronteira de Pareto utilidade-diversidade alcançável em comparação com geradores pré-treinados, GRPO padrão e GRPO assistido por memória.
• Preservação da Diversidade: Em regimes de alta utilidade onde as linhas de base (especialmente o GRPO) sofreram colapso severo de diversidade, o SGRPO manteve diversidade significativamente maior. Por exemplo, no design baseado em bolso, o SGRPO deslocou a trajetória do ponto operacional para fora, retendo diversidade marcadamente maior em níveis de utilidade comparáveis.
• Estudos de Ablação: A remoção da recompensa de diversidade reverteu o desempenho para os níveis do GRPO padrão. A remoção da atribuição de crédito leave-one-out enfraqueceu o desempenho, confirmando a importância de redistribuir sinais de diversidade para contribuintes específicos.
• Dinâmicas Distribucionais: Visualizações das distribuições de sequências mostraram que, enquanto o GRPO e o GRPO assistido por memória contraíam-se em regiões estreitas ou derivavam, o SGRPO explorou e preservou múltiplos clusters, indicando uma cobertura mais ampla do espaço de sequências.

5. Significado e Alegações

O artigo alega que o SGRPO oferece um princípio de pós-treinamento amplamente aplicável para expandir os pontos operacionais disponíveis para modelos de geração biomolecular. Ao recompensar diretamente conjuntos gerados diversos, o método empurra a fronteira utilidade-diversidade para fora, permitindo que os modelos alcancem melhor utilidade em níveis de diversidade fixos (e vice-versa) sem sacrificar a cobertura.

Os autores enfatizam que essa abordagem aborda uma limitação fundamental no design biomolecular baseado em RL atual: a incapacidade de tratar a diversidade como um objetivo de primeira classe durante a otimização de política. Os resultados sugerem que otimizar diretamente a diversidade em nível de conjunto é uma estratégia prática e robusta para mitigar o colapso de modos em tarefas de geração tanto de moléculas pequenas quanto de proteínas. O código é disponibilizado para facilitar pesquisas futuras e instanciação em diferentes domínios.