Resumo Técnico: GRAFT (Tempo Acelerado de Falta com Resíduo Portão)

Declaração do Problema

A análise de sobrevivência enfrenta um compromisso persistente entre interpretabilidade/calibração e poder discriminativo. Modelos clássicos, como o Modelo de Riscos Proporcionais de Cox (CoxPH) e modelos paramétricos de Tempo Acelerado de Falta (AFT), oferecem estimativas de sobrevivência bem calibradas e interpretabilidade, mas são restritos por suposições rígidas (por exemplo, riscos proporcionais ou formas distribucionais específicas) e relações lineares entre características. Por outro lado, modelos de aprendizado profundo (por exemplo, DeepSurv, DeepHit) capturam interações não lineares complexas e alcançam forte discriminação, mas frequentemente sofrem com calibração deficiente e sensibilidade a ruído de alta dimensão. Além disso, conjuntos de dados clínicos do mundo real são frequentemente de alta dimensão com características irrelevantes, contudo muitos modelos carecem de mecanismos robustos e integrados para seleção automática de características, levando ao sobreajuste e à degradação do desempenho em ambientes ruidosos.

Metodologia: A Arquitetura GRAFT

Os autores propõem o GRAFT (Tempo Acelerado de Falta com Resíduo Portão), um modelo híbrido projetado para desacoplar a classificação prognóstica da calibração de sobrevivência. A arquitetura integra três componentes principais:

1. Arquitetura Híbrida (Linear + Não Linear)

Em vez de modelar diretamente o tempo de sobrevivência logarítmico por meio de uma única função linear ou não linear, o GRAFT aprende uma pontuação prognóstica $s_i$ que é monótona em relação ao tempo de sobrevivência logarítmico. A pontuação é construída como:
$$s_i = \beta^T \phi_i + \mu$$
onde a representação $\phi_i$ combina:

• Componente Linear Portão: As características de entrada $x_i$ são mascaradas por um vetor de porta estocástico $g_{\eta,i}$ (produto elemento a elemento), efetivamente realizando a seleção de características.
• MLP Residual: As características portadas são passadas por uma Rede Neural Multicamada (MLP) residual não linear, $f_\theta(\tilde{x}_i)$.
  Este design permite que o modelo capture simultaneamente efeitos lineares simples e interações não lineares complexas.

2. Seleção Estocástica de Características (STG)

Para lidar com ruído de alta dimensão, o GRAFT integra Portas Estocásticas (STG) diretamente na arquitetura.

• Mecanismo: Para cada característica $j$, um valor de porta $g_{\eta,j}$ é amostrado a partir de um relaxamento contínuo de uma distribuição Bernoulli (baseado em uma distribuição Gaussiana com média aprendível $\eta_j$).
• Treinamento: Durante o treinamento, as portas são amostradas para criar uma máscara binária probabilística. Uma aproximação diferenciável da norma $\ell_0$ é adicionada à função de perda para penalizar o número esperado de portas ativas, incentivando a esparsidade.
• Inferência: No momento do teste, a amostragem estocástica é substituída pelo valor esperado determinístico da porta para garantir previsões reproduzíveis.

3. Objetivo de Treinamento: Classificação e Calibração Desacopladas

O GRAFT separa explicitamente a otimização da classificação (discriminação) da estimativa de probabilidades de sobrevivência (calibração).

• Imputação Condicional Estocástica: Para lidar com dados censurados sem assumir uma distribuição paramétrica para erros, o GRAFT emprega estimadores locais de Kaplan-Meier (KM). Para sujeitos censurados, o modelo constrói uma vizinhança local de $k$ vizinhos mais próximos (incluindo sujeitos censurados e com evento) para ajustar uma curva KM condicional. Os tempos de evento para sujeitos censurados são imputados amostrando dessa distribuição condicional (imputação múltipla) em vez de usar uma única estimativa pontual.
• Função de Perda de Classificação Diferenciável: O modelo é treinado para otimizar o Índice de Concordância (C-index) minimizando uma aproximação diferenciável da correlação de posto negativa de Spearman entre as pontuações previstas e os tempos de evento logarítmicos imputados. Isso utiliza um operador de classificação suave (baseado em regressão isotônica) para permitir otimização baseada em gradiente. A média de Monte Carlo sobre múltiplas amostras de imputação ($M=5$) é usada para contabilizar a incerteza da imputação.
• Calibração Pós-Treinamento: O modelo não gera curvas de sobrevivência diretamente. Em vez disso, as pontuações prognósticas aprendidas $s_i$ são usadas como entrada para um modelo CoxPH padrão (ajustado pós-treinamento) para estimar o risco basal e gerar funções de sobrevivência calibradas $S(t|x)$. Esta etapa garante a calibração sem modificar os parâmetros de classificação aprendidos.

Contribuições Principais

1. Desacoplamento de Classificação e Calibração: O artigo propõe uma separação explícita onde o modelo é treinado exclusivamente para otimizar a discriminação por meio de um objetivo de classificação diferenciável, enquanto a calibração é alcançada por meio de uma etapa leve pós-treinamento. Isso evita os compromissos inerentes às abordagens de aprendizado profundo de objetivo único.
2. Arquitetura AFT Híbrida: O modelo combina uma estrutura AFT linear com uma MLP residual não linear, permitindo que ele aprenda tanto efeitos lineares quanto não lineares de características, mantendo a interpretabilidade da estrutura AFT.
3. Seleção Integrada de Características Estocástica: Ao incorporar Portas Estocásticas diretamente na arquitetura, o GRAFT realiza seleção automática e diferenciável de características durante o treinamento. Isso fornece robustez contra características irrelevantes e ruído de alta dimensão, uma capacidade frequentemente ausente em modelos padrão de sobrevivência profunda.

Resultados Experimentais

Os autores avaliaram o GRAFT em seis conjuntos de dados de referência públicos (GBSG, METABRIC, SUPPORT, NWTCO, FLCHAIN, AIDS) cobrindo taxas de censura variadas (32% a 92%) e domínios (oncologia, cuidados críticos, HIV).

• Desempenho vs. Baselines: O GRAFT superou consistentemente as linhas de base clássicas (CoxPH, AFT Weibull) e as linhas de base modernas de aprendizado profundo (DeepSurv, DeepHit, Floresta de Sobrevivência Aleatória) tanto em discriminação (C-index) quanto em calibração (Pontuação Brier Integrada) na maioria dos conjuntos de dados.
  • Exemplo: No conjunto de dados SUPPORT, o GRAFT alcançou um C-index de 0,6143, superando o DeepSurv (0,6071) e a RSF (0,6057).
  • Alta Censura: No conjunto de dados FLCHAIN (72% de censura), o GRAFT alcançou o melhor IBS (0,0919) e um C-index de 0,7965.
• Robustez ao Ruído: Em estudos de ablação onde características de ruído foram injetadas (até 10 vezes o número de características originais), o GRAFT demonstrou robustez superior. Enquanto o DeepSurv e o DeepHit sofreram degradação catastrófica de desempenho, o GRAFT manteve pontuações de C-index estáveis. O estudo de ablação confirmou que a variante "No STG" (removendo a seleção de características) performou significativamente pior sob ruído do que o modelo completo, validando a necessidade do mecanismo de porta.
• Limitações: O modelo mostrou desempenho reduzido no conjunto de dados AIDS (92% de censura), onde as vizinhanças locais de KM continham informações de evento insuficientes para imputação estável. Em tais cenários de censura extrema, os autores observam que modelos paramétricos mais simples (como CoxPH) podem ser preferíveis.

Significado e Alegações

O artigo alega que desacoplar classificação e calibração é um princípio de design crítico para análise de sobrevivência. Ao otimizar a classificação durante o treinamento e calibrar post-hoc, o GRAFT alcança desempenho superior em ambas as métricas simultaneamente, abordando uma falha comum de modelos de sobrevivência de aprendizado profundo.

Além disso, os autores enfatizam que a seleção integrada de características é essencial para robustez em cenários de saúde caracterizados por dados de alta dimensão e ruidosos. A capacidade do GRAFT de suprimir automaticamente características não informativas permite que ele mantenha o poder discriminativo onde outros modelos degradam. O trabalho sugere que, para aplicações biomédicas de alta dimensão, uma abordagem híbrida que combina a flexibilidade do aprendizado profundo com as restrições estruturais e propriedades de calibração de modelos clássicos oferece uma solução prática e eficaz.