Calibrating Generative Models to Feature Distributions with MMD Finetuning

Resumo Técnico: Calibrando Modelos Generativos para Distribuições de Características com Ajuste Fino por MMD (kCGM)

Declaração do Problema
Modelos generativos frequentemente produzem amostras que são individualmente plausíveis, mas exibem um desalinhamento distributivo em relação a características-chave de interesse. Por exemplo, um modelo pré-treinado em um espaço químico amplo pode gerar moléculas que diferem significamente da distribuição de características de uma classe terapêutica específica, como antibióticos. Corrigir isso via ajuste fino (finetuning) direto no conjunto alvo é problemático: corre-se o risco de sobreajuste (overfitting) a conjuntos de dados pequenos, exige amostras completas do alvo (que podem ser escassas) e falha em especificar quais características específicas devem ser correspondidas. Além além disso, métodos existentes frequentemente têm dificuldade com características de valores discretos ou de caixa-preta (ex: impressões digitais de Morgan) que não são diferenciáveis. O artigo identifica uma lacuna em algoritmos capazes de calibrar modelos generativos para corresponder a uma distribuição de características alvo usando apenas supervisão ao nível de características, sem exigir a diferenciabilidade do mapa de características ou do procedimento de amostragem.

Metodologia: kCGM
Os autores introduzem os Modelos Generativos de Calibração por Kernel (kCGM), um método de ajuste fino que minimiza a Discrepância de Média Máxima (MMD) entre a distribuição de características gerada e uma distribuição de características alvo, regularizada pela divergência de Kullback-Leibler (KL) em relação ao modelo pré-treinado.

• Função Objetivo: O método otimiza uma função objetivo $L_\theta$ composta por dois termos:

  1. Perda MMD: Um estimador não viesado de amostra finita do quadrado da MMD entre a distribuição de características induzida pelo modelo $P$ (onde $y = h(x)$ e $x \sim p_\theta$) e a distribuição empírica alvo $Q$ (proveniente de características alvo $y^{targ}_{1:N}$).
  2. Regularização KL: Um termo $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ para garantir que o modelo ajustado permaneça próximo à distribuição pré-treinada, evitando o esquecimento catastrófico ou o colapso de modo (mode collapse).

• Seleção de Kernel: Diferente de abordagens anteriores que podem depender de kernels de produto escalar (que apenas correspondem médias), o kCGM utiliza kernels que definem uma métrica em distribuições de probabilidade (zero apenas para distribuições idênticas).

  • Para características euclidianas contínuas, um kernel de distância de energia é utilizado.
  • Para características de vetores de bits discretos (ex: impressões digitais de Morgan), um kernel de Tanimoto/Jaccard é empregado.
  • Essas escolhas eviam a introdução de hiperparâmetros adicionais de largura de banda (bandwidth).

• Estimativa de Gradiente: Como o mapa de características $h$ e o processo de amostragem podem não ser diferenciáveis, o kCGM emprega gradientes de função de pontuação (score-function gradients) (estilo REINFORCE).

  • Baseline Leave-One-Out (LOO): Uma contribuição técnica crítica é a derivação de um novo baseline LOO para a perda MMD. Baselines LOO padrão fazem a média dos coeficientes de outras amostras em um lote, mas os coeficientes de MMD são acoplados entre as amostras. Os autores derivam um estimador "MMD-LOO" específico que aproxima o coeficiente médio caso uma amostra específica não tivesse aparecido no mini-lote, reduzindo a variância sem introduzir viés.
  • Algoritmo: O processo envolve a amostragem do modelo atual (sem rastreamento de gradiente), o cálculo dos vetores de características, o cálculo dos coeficientes de MMD e KL com baselines LOO, e a atualização dos pesos do modelo via gradientes de função de pontuação.

Principais Contribuições

1. Formulação: O artigo formula a correspondência de distribuição de características como ajuste fino de MMD regularizado por KL usando apenas amostras de características alvo, eliminando a necessidade de amostras completas do alvo no espaço de amostra original.
2. Estimador Não Viesado: Os autores derivam um estimador de gradiente de função de pontuação não viesado para o objetivo MMD, permitindo o uso de mapas de características não diferenciáveis e amostradores não reparametrizáveis (ex: modelos autorregressivos, difusão discreta).
3. Generalização: O método generaliza o Calibrating Generative Models (CGM). Enquanto o CGM faz a correspondência das médias das características alvo, o kCGM faz a correspondência da distribuição completa via MMD. O artigo observa que o uso de um kernel de produto escalar no kCGM recupera o objetivo CGM-relax.

Resultados Experimentais
Os autores avaliam o kCGM em três domínios: geração de pequenas moléculas, geração de estruturas de proteínas e geração de DNA regulatório.

• Pequenas Moléculas Antibióticas:

  • Tarefa: Calibrar um modelo de grafo autorregressivo G2PT para corresponder à distribuição de características de 174 antibióticos conhecidos.
  • Características: Impressões digitais de Morgan (subestrutura local), andaimes (scaffolds) de Murcko (topologia) e descritores físico-químicos.
  • Descobertas: O kCGM alcança um tradeoff superior entre a correspondência com a distribuição de características alvo (medida por MMD) e a distância do modelo pré-treinado em comparação ao CGM e ao ajuste fino direto. Crucialmente, enquanto o ajuste fino direto no pequeno conjunto de antibióticos levou a uma alta proporção de moléculas quimicamente inválidas, o kCGM melhorou a correspondência de características ao mesmo tempo em que aumentou a validade química. O kCGM também demonstrou transferência de características cruzadas, melhorando a correspondência de andaimes e descritores mesmo quando treinado apenas em impressões digitais.

• Diversidade de Estrutura de Proteínas:

  • Tarefa: Ajustar a proteína Genie 2 (um modelo de difusão de proteína) para corresponder à distribuição de estrutura secundária (alfa-hélices vs. folhas-beta) do conjunto de dados CATH.
  • Descobertas: O modelo pré-treinado produzia majoritariamente alfa-hélices. O kCGM reduziu o desacordo com a distribuição alvo em aproximadamente 2x comparado ao CGM. O kCGM mostrou uma curva de tradeoff suave com o peso KL ($\lambda$), enquanto o desempenho do CGM era menos sensível a $\lambda$, tornando a otimização difícil.

• Geração de DNA Regulatório:

  • Tarefa: Ajustar um modelo de difusão discreta para gerar sequências de DNA com perfis de atividade celular-específicos realistas (previstos pelo AlphaGenome).
  • Descobertas: O kCGM conseguiu corresponder à distribuição completa dos perfis de atividade para 47 diferentes tópicos cis-regulatórios, superando o CGM (que apenas correspondia às médias). O modelo ajustado gerou sequências com perfis de atividade similares a enhancers reais, enquanto o modelo pré-treinado gerava sequências com baixa atividade.

Significância e Alegações
O artigo afirma que o kCGM fornece uma alternativa robusta ao ajuste fino direto, particularmente quando os conjuntos de dados alvo são pequenos ou quando apenas aspectos específicos de uma distribuição precisam ser preservados. Sua principal significância reside na sua capacidade de:

• Alinhar modelos generativos a distribuições de características especificadas pelo usuário usando apenas supervisão ao nível de características.
• Lidar com características de caixa-preta e não diferenciáveis (ex: impressões digitais, saídas de preditores) que são comuns em domínios científicos, mas incompatíveis com o ajuste fino baseado em gradiente padrão.
• Manter ou melhorar a qualidade das amostras (validade) enquanto corrige a má calibração distributiva, um modo de falha comum do ajuste fino direto.
• Ser aplicável a diversas arquiteturas, incluindo modelos autorregressivos, difusão em espaço contínuo e modelos de difusão discreta.

Os autores observam limitações, como o potencial de ineficácia para características de altíssima dimensão devido à maldição da dimensionalidade em testes de hipótese MMD, e a exigência de log-probabilidades de trajetórias de amostragem tratáveis (excluindo modelos de fluxo/flow matching). Eles sugerem que trabalhos futuros poderiam estender este framework para cenários onde apenas restrições distributivas existem, sem amostras completas no espaço generativo original.