Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

O artigo introduz o CausalLongPFN, uma rede pré-ajustada congelada, pré-treinada em modelos causais temporais sintéticos, que prevê resultados longitudinais contrafatuais sem treinamento específico de domínio, demonstrando um desempenho competitivo em benchmarks como câncer, HIV e MIMIC-III, ao mesmo tempo em que oferece uma alternativa de custo-benefício superior à modelagem tradicional específica de coorte.

O essencial

Imagine que você é um médico tentando prever como um paciente responderá a um plano de tratamento específico ao longo das próximas semanas. O problema é que os pacientes são complexos: sua saúde muda dia após dia, seus tratamentos passados afetam seu estado atual e outros fatores (como dieta ou estresse) interferem nos resultados. Normalmente, para fazer essas previsões, você tem que construir um modelo de computador novo e altamente especializado para cada novo grupo de pacientes que encontra. Isso é como contratar um novo arquiteto para projetar uma casa toda vez que você se muda para um novo bairro. É lento, caro e exige muitos dados.

Este artigo apresenta uma nova ferramenta chamada CAUSALLONGPFN. Pense nela como um "Motor de Intuição Médica Universal" que já aprendeu as regras do jogo antes mesmo de ver um paciente real.

Veja como ela funciona, dividido em conceitos simples:

1. O "Campo de Treinamento" (Pré-treinamento Sintético)

Em vez de esperar que os pacientes reais cheguem, os criadores construíram um enorme "campo de treinamento" virtual. Neste campo, eles simularam milhões de pacientes falsos com milhões de diferentes tipos de corpos, doenças e reações a tratamentos. Eles programaram esses pacientes falsos para terem comportamentos complexos:

• Alguns melhoram lentamente; outros pioram rapidamente.
• Alguns tratamentos funcionam imediatamente; outros têm um efeito retardado.
• Alguns pacientes reagem de forma diferente com base em sua genética oculta.

O modelo de IA passou todo o seu tempo neste campo, aprendendo a prever desfechos para esses pacientes falsos. Ele não apenas memorizou respostas; ele aprendeu a lógica subjacente de como os tratamentos, o tempo e a biologia interagem.

2. O "Especialista Congelado" (Sem Necessidade de Novo Treinamento)

Aqui está o truque de mágica: assim que o modelo terminou seu campo de treinamento, eles o congelaram. Eles trancaram seu cérebro. Ele não pode aprender nada novo ou alterar suas configurações internas.

Quando um novo grupo de pacientes reais chega (como pacientes com câncer em um hospital), o modelo não começa do zero. Ele não precisa ser retreinado. Em vez disso, ele age como um superestagiário que lê o prontuário.

• As Trajetórias de Suporte: Você mostra ao modelo alguns exemplos de pacientes do hospital atual (o "suporte"). Estes são como dar ao estagiário alguns arquivos de casos para ler, para que ele entenda o estilo específico deste hospital.
• A Consulta (Query): Você pergunta: "O que acontecerá com este paciente específico se dermos o Tratamento A por 5 dias?"
• A Resposta: O modelo usa instantaneamente o que aprendeu no campo de treinamento, combinado com o contexto dos arquivos de casos que você mostrou a ele, para prever o resultado. Ele faz isso sem dar um único passo de "descida de gradiente" (um termo técnico para o processo usual de retreinamento).

3. O "Simulador de Viagem no Tempo"

O modelo é projetado para lidar com dados longitudinais, o que significa que ele entende o tempo. Ele não apenas adivinha o resultado de amanhã; ele simula o futuro passo a passo.

• Ele prevê o Dia 1.
• Ele pega essa previsão e a utiliza como ponto de partida para o Dia 2.
• Ele repete esse processo para ver o que acontece no Dia 5, 6 ou 7.

Isso é como um simulador de voo que não mostra apenas a decolagem, mas simula todo o trajeto de voo com base nas decisões do piloto, mesmo que o clima mude durante o percurso.

4. Por Que Isso Importa (Os Resultados)

Os autores testaram este "especialista congelado" contra a forma antiga de fazer as coisas (construir um novo modelo para cada hospital).

• O Testo: Eles usaram dados de câncer, HIV, anticoagulantes (varfarina) e registros reais de UTI.
• O Resultado: O modelo congelado teve um desempenho tão bom quanto, e às vezes melhor do que, os modelos que foram especialmente treinados para cada doença específica.
• A Grande Vitória: Funcionou especialmente bem em dados reais de UTI, onde não podiam testar cenários de "e se" (porque não é ético testar diferentes tratamentos em pacientes reais em uma simulação). O modelo conseguiu prever o que provavelmente aconteceria a seguir com base apenas nos dados observados.

A Conclusão

O artigo afirma que você nem sempre precisa construir um novo modelo customizado para cada novo conjunto de dados médicos. Em vez disso, você pode treinar um modelo massivo em uma ampla variedade de cenários de "e se" (dados sintéticos) e então usá-lo como uma ferramenta congelada e pronta para uso.

É como ter um mestre chef que praticou o preparo de todos os tipos de culinária em uma cozinha virtual. Quando você traz um novo conjunto de ingredientes (um novo grupo de pacientes), ele não precisa aprender a cozinhar novamente; ele apenas olha para os seus ingredientes e imediatamente sabe como preparar o prato, usando sua vasta intuição pré-aprendida.

Nota Importante: Os autores deixam claro que esta ferramenta é para previsão e pesquisa, não para tomar decisões finais de vida ou morte em uma clínica. Ela ajuda os médicos a entenderem potenciais desfechos, mas ainda depende das mesmas suposições médicas padrão de que qualquer outro modelo causal depende. É uma poderosa ferramenta de pesquisa, não um substituto para o julgamento médico.

Para quem deseja explorar a implementação ou usar o modelo, o código-fonte está disponível no GitHub e os pesos do modelo pré-treinado podem ser baixados no Hugging Face.

Resumo técnico

Resumo Técnico: Redes de Ajuste de Prior Longitudinal Causal para Predição de Resultados Contrafatuais

Formulação do Problema

O artigo aborda o desafio de prever resultados potenciais sob sequências de tratamentos futuros em dados observacionais longitudinais. A tarefa central é estimar o resultado potencial condicionado ao histórico $E[Y(\bar{a})_{t+\tau} | H_t]$, onde $H_t$ representa o histórico observado até o tempo $t$, e $\bar{a}$ é uma sequência de tratamento planejada.

Este problema é complicado por três fatores primários:

1. Confundimento variante no tempo: As atribuições de tratamento em cada etapa dependem de covariáveis que são, elas próprias, resultados de tratamentos anteriores.
2. Dinâmicas heterogêneas de pacientes: As trajetórias individuais exibem evolução de estado complexa e não linear, além de heterogeneidade latente.
3. Limitações de dados: Coortes observacionais são frequentemente pequenas demais para treinar modelos de sequência profundos confiáveis do zero para cada novo domínio ou simulador.

Estimadores causais longitudinais existentes (ex: RMSN, CRN, G-Net, Causal Transformer) tipicamente requerem uma execução de treinamento supervisionado separada, incluindo seleção de hiperparâmetros e modelagem de propensão, para cada nova coorte. Este pipeline é custoso e impraticável quando o treinamento repetido específico para o domínio é necessário.

Metodologia: CAUSALLONGPFN

Os autores introduzem o CAUSAL LONGITUDINAL PRIOR-FITTED NETWORKS (CAUSALLONGPFN), um preditor de contexto ajustado por prior (prior-fitted in-context predictor) projetado para amortizar a predição causal longitudinal através de um amplo prior sobre modelos causais estruturais temporais (TSCMs).

1. Pré-treinamento Sintético em um Prior Amplo

O modelo é pré-treinado inteiramente em episódios sintéticos amostrados de um amplo prior sobre TSCMs. Este prior é projetado para abranger uma vasta classe de dinâmicas causais longitudinais, em vez de replicar um único simulador. As principais características do prior sintético incluem:

• Grafos Temporais Causais: Dependências contemporâneas e defasadas esparsas com grafos instantâneos acíclicos.
• Mecanismos Não Lineares: As coordenadas de estado seguem atualizações autorregressivas esparsas e não lineares usando diversas não linearidades elementares (identidade, tanh, sinusoidal, ReLU, etc.) e várias distribuições de ruído.
• Motivos Dinâmicos: Motivos estruturados como memória de ação, saturação, homeostase, controle de feedback e canais de leitura suavizada são sobrepostos para capturar mecanismos qualitativos como efeitos retardados e feedback regulatório.
• Políticas de Comportamento Confundidas: Os tratamentos são amostrados de políticas estocásticas dependentes do estado, influenciadas pela heterogeneidade latente da unidade ($Z_i$), criando um feedback entre tratamento e confundidor variante no tempo.
• Modelos de Desfecho: Desfechos escalares são gerados via leituras autorregressivas com efeitos de tratamento diretos e cumulativos.

2. Arquitetura

O CAUSALLONGPFN emprega uma arquitetura de codificador duplo:

• Codificador de Histórico Causal: Um Transformer causal de trajetória (usando autoatenção mascarada) que mapeia sequências longitudinais para representações de histórico, garantindo que a representação no tempo $r$ dependa apenas das informações disponíveis até aquele tempo.
• Codificador de Contexto PFN: Um codificador de contexto baseado em Transformer que realiza a adaptação em contexto. Ele processa trajetórias de suporte (tratadas como um conjunto não ordenado) e um token de consulta conjuntamente via autoatenção. Nenhum codificação posicional é atribuída à ordenação das trajetórias de suporte.
• Cabeça de Predição de Mistura Gaussiana: A representação de consulta final parametriza uma distribuição de mistura gaussiana de 5 componentes para o desfecho normalizado, fornecendo tanto predições pontuais quanto estimativas de incerteza.

3. Predição em Contexto e Rollout

No momento do teste, o modelo permanece congelado. Ele recebe:

• Trajetórias de suporte: Exemplos do novo domínio/tarefa.
• Histórico de consulta: Observado até o tempo $t_{obs}$.
• Sequência de tratamento futura proposta: As intervenções planejadas.

O modelo retorna uma distribuição preditiva sem atualizações de gradiente, ajuste de modelo de propensão ou balanceamento adversário. Para predição de múltiplos passos ($\tau > 1$), o modelo realiza um rollout de plug-in autorregressivo: ele prevê a distribuição do desfecho de um passo, insere a média da mistura no histórico de consulta e repete o processo sob a sequência de tratamento especificada.

Principais Contribuições

1. Um Modelo de Ajuste de Prior para Predição Causal Longitudinal: O CAUSALLONGPFN é o primeiro modelo estilo PFN para predição de desfecho potencial condicionado ao histórico sob sequências de tratamento longitudinais planejadas. Ele opera como um modelo congelado que não requer adaptação no tempo de teste.
2. Um Prior Sintético sobre Tarefas Causais Longitudinais: Os autores projetam um prior de TSCM que gera tarefas diversas com feedback tratamento-confundidor, heterogeneidade latente, dinâmicas não lineares, efeitos retardados/cumulativos e mudanças de regime.
3. Arquitetura para Inferência em Contexto Longitudinal: Um novo codificador duplo que combina um codificador de histórico Transformer causal com um codificador de contexto PFN e uma cabeça de mistura gaussiana.
4. Rollout Contrafatual Autorregressivo: Uma extensão do preditor de um passo aprendido para predição de múltiplos passos via rollout de plug-in determinístico.
5. Avaliação Zero-Shot: Uma avaliação abrangente de um único modelo congelado contra baselines treinados por domínio (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) em benchmarks de contrafatuais ramificáveis e dados reais fatuais.

Resultados

O modelo foi avaliado em quatro benchmarks: Crescimento de tumor em câncer, PK/PD de Varfarina, Dinâmicas de tratamento de HIV e trajetórias de UTI MIMIC-III.

• Desempenho Balanceado por Domínio: O CAUSALLONGPFN alcançou o melhor RMSE normalizado de um passo balanceado por domínio (0.222), superando por margem estreita o MSM e o RMSN. Para predição de cinco passos, classificou-se em terceiro lugar geral, atrás do RMSN e G-Net, mas superou o MSM, CRN e os baselines baseados em transformer.
• Desempenho por Domínio:
  • MIMIC-III (Factual): O modelo ficou em primeiro lugar tanto na predição de um passo quanto na de cinco passos, demonstrando forte transferência para trajetórias clínicas do mundo real sem treinamento específico de domínio.
  • Benchmarks Contrafatuais (Câncer, HIV, Varfarina): O modelo manteve-se competitivo, ocupando o segundo ou terceiro lugar em tarefas de um passo. No entanto, em tarefas contrafatuais de horizonte mais longo (ex: predição de câncer de 5 passos), modelos recorrentes especializados e treinados por domínio (RMSN, CRN) alcançaram erros menores, sugerindo uma vantagem quando existem dados suficientes do domínio alvo para ajuste específico.
• Calibração de Incerteza: A cabeça de mistura gaussiana forneceu informações distribuicionais úteis. A calibração variou por domínio, com a Varfarina apresentando a melhor calibração e o MIMIC-III apresentando intervalos mais amplos devido à maior heterogeneidade.

Disponibilidade de Código e Modelos

A implementação do CAUSALLONGPFN está disponível no GitHub em https://github.com/Amirhossein-Zare/causal-long-pfn. Os pesos do modelo pré-treinado estão disponíveis no Hugging Face em https://huggingface.co/Amirhossein-Zare/causal-long-pfn .

Significância e Alegações

O artigo alega que o pré-treinamento em um prior sintético amplo pode fornecer uma alternativa congelada útil ao treinamento repetido específico de domínio quando:

• O retreinamento de modelos especializados é custoso ou impraticável.
• A adaptação rápida a uma nova coorte é necessária.
• A supervisão contrafatual não está disponível (como em tarefas de predição factual do mundo real como o MIMIC-III).

Os autores enfatem que o CAUSALLONGPFN não remove as suposições causais padrão (consistência, positividade, trocaabilidade sequencial) necessárias para interpretar dados observacionais. Em vez disso, ele amortiza o problema de estimação. Os resultados sugerem que um prior sintético suficientemente amplo pode capturar estruturas reutilizáveis através de tarefas de resposta ao tratamento, tornando o modelo um forte preditor de contexto de uso geral. No entanto, o artigo nota modestamente que o treinamento específico de domínio continua sendo valioso quando há sinais de validação e dados suficientes do domínio alvo, particularmente para predições de longo horizonte em domínios específicos.

O trabalho posiciona-se como uma ferramenta de pesquisa para modelagem de sequências causais e geração de hipóteses, em vez de um sistema de decisão clínica autônomo, alertando contra a confiança excessiva em predições quando as suposições causais ou o suporte do prior são inadequados.