Resumo Técnico: Grafos de Conhecimento e LLMs de Raciocínio para Encontrar Preditores de Perturbação Transcriptômica Simples e Eficazes

1. Declaração do Problema

Prever os efeitos transcriptômicos de perturbações de nocaute de genes não observados permanece um desafio significativo para modelos de células virtuais. Embora abordagens recentes utilizem grafos de conhecimento biológico para melhorar a extrapolação além das perturbações de treinamento, modelos complexos de aprendizagem profunda frequentemente lutam para superar baselines simples (ex: a média das amostras de treinamento). A dificuldade central reside em generalizar para novas distribuições intervencionais, $P(x|do(i))$, onde cada perturbação induz uma distribuição única restrita por processos biológicos complexos. Métodos existentes que utilizam Redes Neurais de Grafos (GNNs) ou pré-treinamento em larga escala mostraram promessa, mas os mecanismos específicos pelos quais o conhecimento biológico permite a generalização para novas intervenções permanecem obscuros. Além disso, otimizar a estrutura causal subjacente diretamente é computacionalmente intratável devido ao espaço de busca exponencial de possíveis grafos de genes.

2. Metodologia

2.1 Fundamentação Teórica: Um Modelo Hierárquico Esparso

Os autores postulam que as respostas celulares a nocautes de genes seguem um grafo causal hierárquico esparso. Neste modelo, genes formam complexos proteicos, que formam vias, levando eventualmente a organelas e respostas celulares. Um nocaute interrompe a produção de uma proteína inicial, propagando a falha para cima na hierarquia.

• Implicação: Perturbações que afetam o mesmo complexo de baixo nível devem gerar perfis de expressão semelhantes.
• Predição: A expressão média de uma perturbação alvo pode ser aproximada por uma média simples de vizinhos mais próximos sobre o vizinhança local definida por este grafo causal.

2.2 Baseline: O Preditor de Vizinho de Grafo

Os autores primeiro avaliam um preditor "Graph Neighbour" (Vizinho de Grafo), que utiliza grafos de conhecimento biológico existentes (ex: StringDB, GO) como priors.

• Mecanismo: Para um gene $g$, o preditor identifica seus vizinhos $N(g)$ no grafo de conhecimento. O delta da perturbação previsto ($\hat{\delta}_g$) é calculado como a média dos deltas de treinamento desses vizinhos:
  $$\hat{y}_g = \bar{x}_g + \frac{1}{|N_{Tr}(g)|} \sum_{\tilde{g} \in N_{Tr}(g)} \bar{\delta}_{\tilde{g}}$$
  onde $\bar{x}_g$ é a expressão de controle correspondente e $N_{Tr}(g)$ representa os vizinhos presentes no conjunto de treinamento.
• Achado: Esta abordagem simples de $K$-vizinhos mais próximos supera modelos complexos de aprendizagem profunda (ex: GEARS, scLAMBDA) e aproxima-se do desempenho de modelos de estado da arte (SOTA) como o TxPert, particularmente em configurações fora da distribuição (OOD).

2.3 Refinamento: Aprendizagem por Reforço com LLMs de Raciocínio

Reconhecendo que os grafos de conhecimento são incompletos e contêm arestas subótimas, os autores propõem o uso de Grandes Modelos de Linguagem (LLMs) para refinar essas vizinhanças.

• Redução do Espaço de Busca: Em vez de buscar no espaço intratável de todos os grafos possíveis ($O(2^{\#genes^2})$), a busca é restrita a arestas candidatas plausíveis sob as crenças a priori do LLM.
• Formulação de RL:
  • Política: Um LLM ($\pi_\theta$) atua como uma política para selecionar ações (Adicionar $A$, Remover $R$) a um conjunto inicial de vizinhos do grafo de conhecimento ($N_{KG}$).
  • Espaço de Ação: O conjunto refinado é $S_{LLM} = (N_{KG} \setminus R) \cup A$.
  • Função de Recompensa: O modelo é treinado via Otimização de Política Relativa de Grupo (GRPO) para maximizar uma recompensa $R(s, a)$ composta por:
    1. Recompensa de Formatação ($R_{Fmt}$): Garante que a saída seja analisável (parsable).
    2. Recompensa de Validade ($R_{Valid}$): Incentiva a adição de genes do conjunto de treinamento e a remoção de genes que não estavam na vizinhança original.
    3. Recompensa de Desempenho ($R_\Delta$): Incentiva melhorias na métrica Pearson Delta em comparação com o preditor baseline do grafo de conhecimento.
• Treinamento: O LLM é treinado com dados de perturbação de linhagens celulares de Replogle et al. (2022), com uma divisão rigorosa de treino/teste para garantir a avaliação em perturbações não observadas.

3. Principais Contribuições

1. Demonstração de Baselines Simples: O artigo estabelece que um preditor simples de $K$-vizinhos mais próximos baseado em grafos de conhecimento biológico é um baseline altamente competitivo para a predição de perturbação OOD, superando a maioria dos métodos complexos de aprendizagem profunda.
2. LLM como um Prior de Refinamento: Os autores demonstram que LLMs, quando otimizados via Aprendizagem por Reforço, podem efetivamente editar vizinhanças de grafos de conhecimento (adicionando relações ausentes e podando arestas irrelevantes) para melhorar o desempenho preditivo.
3. Desempenho de Estado da Arte: O LLM ajustado por RL atinge desempenho equivalente aos métodos SOTA atuais (especificamente o TxPert) em linhagens celulares de Replogle et al. (2022), apesar de partir de uma arquitetura simples de vizinho mais próximo.
4. Generalização para Tarefas Downstream: O estudo mostra que o treinamento por RL melhora o desempenho do LLM em tarefas subsequentes (especificamente a predição de expressão diferencial no benchmark PerturbQA), sugerindo que o modelo aprende padrões de raciocínio generalizáveis em vez de sofrer overfitting à tarefa de predição específica.

4. Resultados

• Desempenho em Benchmarks: Em quatro linhagens celulares (K562, RPE1, HEPG2, Jurkat) de Replogle et al. (2022), o LLM treinado por RL ($LLM + KG + RL$) alcançou pontuações de Pearson Delta comparáveis ao TxPert (ex: 0.625 vs 0.626 em K562) e superou significativamente o LLM não treinado e o baseline bruto do grafo de conhecimento.
• Generalização OOD: O método generaliza para perturbações não observadas, mantendo alto desempenho mesmo quando os nocaute de genes específicos não estavam presentes no treinamento.
• Transferência de Aprendizado: Quando treinado em apenas três das quatro linhagens celulares, o modelo ainda melhorou o desempenho na linhagem mantida fora (held-out) (K562) em comparação com o modelo base, indicando generalização da política de raciocínio entre linhagens celulares.
• Tarefa Downstream: No dataset PerturbQA (Wu et al.), o modelo pós-treinado superou o LLM base na predição de expressão diferencial e direção de mudança, mesmo sem fine-tuning específico para esta tarefa.

5. Significância e Alegações

O artigo alega que grafos de conhecimento biológico servem como eficazes priors de modelo, e que o modelo mais simples que utiliza esses priors (vizinho mais próximo) é surpreendentemente robusto. A principal significância reside em mostrar que a Aprendizagem por Reforço pode refinar LLMs em ferramentas generalizáveis para predição biológica complexa.

Os autores argumentam que, embora o treinamento por RL seja computacionalmente caro e não seja necessariamente o método de implantação mais prático para predição imediata, ele serve como um prova de conceito:

• Ele valida que os LLMs possuem sinais preditivos significativos sobre relações gênicas além do que é explicitamente codificado em grafos de conhecimento estáticos.
• Demonstra que o RL pode ajustar um LLM para agir como uma posterior sobre mudanças de vizinhança de genes, aprendendo efetivamente uma política para selecionar conjuntos de genes preditivos.
• Destaca que melhorias no desempenho preditivo podem se transferir para tarefas de raciocínio subsequentes, sugerindo um caminho para modelos de "célula virtual" mais robustos.

Os autores mantêm a modéstia, apontando limitações como o custo computacional do RL e a dependência de ter um conjunto de treinamento para a linhagem celular específica, o que impede predições puramente "zero-shot" em novas linhagens celulares. Eles não alegam ter resolvido inteiramente o problema da generalização, mas mostraram um caminho viável para refinar modelos simples e interpretáveis usando raciocínio avançado de modelos de linguagem.