CleanSurvival: Automated data preprocessing for time-to-event models using reinforcement learning

Resumo Técnico: CleanSurvival

Declaração do Problema
O pré-processamento de dados é um componente crítico, porém frequentemente negligenciado, em pipelines de aprendizado de máquina, impactando significativamente o desempenho do modelo. Embora frameworks de Aprendizado de Máquina Automatizado (AutoML) tenham começado a integrar pré-processamento automatizado para tarefas de classificação e regressão, há uma falta distinta de soluções automatizadas e sob medida para análise de sobrevivência (modelagem de tempo até evento). A análise de sobrevivência apresenta desafios únicos devido a dados censurados, onde os tempos de evento são observados apenas parcialmente. Ferramentas existentes de AutoML ou falham em lidar com dados censurados nativamente, dependem de pré-processamento manual ou convertem tarefas de sobrevivência em problemas de classificação (por exemplo, via discretização de tempo), levando à perda de informação. Além disso, as ferramentas atuais de pré-processamento automatizado frequentemente carecem de suporte para as estruturas de recompensa e tipos de modelo específicos exigidos pela análise de sobrevivência, como o modelo de riscos proporcionais de Cox ou Florestas Aleatórias de Sobrevivência.

Metodologia
O artigo introduz o CleanSurvival, um framework de pré-processamento de dados automatizado projetado especificamente para análise de sobrevivência, construído sobre uma abordagem de Aprendizado por Reforço (RL) utilizando Q-learning.

• Arquitetura: O framework trata a seleção de etapas de pré-processamento como um problema de tomada de decisão sequencial. Um agente navega em um espaço de estados definido pelas características do conjunto de dados para selecionar uma sequência de ações (técnicas de pré-processamento) que maximize uma função de recompensa.
• Ações de Pré-processamento: O framework suporta uma variedade de técnicas para:
  • Imputação de Valores Faltantes: Análise de Casos Completos (CCA), Imputação por Média/Mediana, Imputação Múltipla por Equações Cadeadas (MICE) e k-Vizinhos Mais Próximos (KNN).
  • Detecção de Outliers: Envelope Elíptico (baseado na distância de Mahalanobis) e resíduos de Martingala.
  • Seleção/Extração de Características: Seleção Univariada de Riscos Proporcionais de Cox, LASSO, Eliminação Recursiva de Características (RFE) e Ganho de Informação.
• Modelos de Sobrevivência: O sistema é compatível com vários modelos de tempo até evento, incluindo o modelo de Riscos Proporcionais de Cox, Florestas Aleatórias de Sobrevivência (RSF), DeepHit (uma abordagem de aprendizado profundo) e modelos fornecidos pelo usuário.
• Estrutura de Recompensa: O agente de Q-learning é guiado por um sinal de recompensa baseado no Índice de Concordância (C-index) avaliado em um conjunto de retenção (hold-out). Isso garante que o pipeline de pré-processamento selecionado seja otimizado especificamente para o desempenho preditivo do modelo de sobrevivência a jusante.
• Modos de Funcionamento: A ferramenta oferece quatro modos:
  1. Algoritmo Principal: Otimização por Q-learning.
  2. Limpeza Aleatória: Busca aleatória para comparação de linha de base.
  3. Pipeline Personalizado: Pipelines fixos definidos pelo usuário via arquivos de configuração.
  4. Sem Preparação: Entrada de dados brutos para desempenho de linha de base.

Principais Contribuições

1. Primeiro Pré-processamento Baseado em RL para Sobrevivência: O CleanSurvival estende o framework Learn2Clean (originalmente para regressão/classificação) para lidar com as restrições específicas de dados de tempo até evento censurados.
2. Otimização Específica do Modelo: Diferentemente do pré-processamento estático, o framework seleciona dinamicamente pipelines que são ótimos para classes específicas de modelos (por exemplo, um pipeline ótimo para Cox pode diferir de um para RSF).
3. Implementação de Código Aberto: Os autores fornecem um pacote Python disponível no GitHub, integrando-se a bibliotecas padrão como scikit-survival.
4. Tratamento de Censura: O framework incorpora explicitamente o ponderamento pela probabilidade inversa de censura (IPCW) em métricas como o Escore Brier Integrado (IBS) para contabilizar observações censuradas durante a avaliação.

Resultados Experimentais
O framework foi avaliado em três conjuntos de dados do mundo real: Rotterdam (câncer de mama), Flchain (cadeia leve livre sérica) e GBSG (Grupo de Estudo Alemão de Câncer de Mama). Ausências sintéticas (MCAR, MAR, MNAR) foram introduzidas para testar a robustez.

• Desempenho: O CleanSurvival superou consistentemente linhas de base simples (Análise de Casos Completos e Imputação por Média) e frequentemente igualou ou superou buscas aleatórias não direcionadas e Optuna (otimização bayesiana) em termos de C-index.
  • No conjunto de dados Flchain, o CleanSurvival alcançou um C-index de 0,67, comparado a 0,39 para CCA e 0,62 para Imputação por Média.
  • No conjunto de dados Rotterdam com 50% de ausências, o CleanSurvival alcançou C-indices entre 0,82 e 0,83, superando significativamente a CCA (0,77–0,79) e a Imputação por Média (0,78–0,79).
• Seleção de Características: A política de Q-learning selecionou frequentemente operadores de seleção de características, com o LASSO aparecendo em ~68% dos pipelines otimizados e a RFE em ~42%.
• Calibração: O Escore Brier Integrado (IBS) indicou que a calibração foi geralmente estável entre os mecanismos de ausência, com exceção das condições de 40% MCAR, que mostraram maior variância.
• Tempo de Execução: O comportamento do tempo de execução foi dependente da condição. Embora o framework tenha identificado com sucesso pipelines de alto desempenho, o tempo para atingir 95% do melhor C-index observado variou conforme o conjunto de dados e o mecanismo de ausência.

Significado e Afirmações
O artigo afirma que o CleanSurvival aborda uma lacuna crítica no ecossistema AutoML ao generalizar o pré-processamento automatizado para análise de sobrevivência. Seu significado reside em:

• Robustez: Demonstrar que o pré-processamento automatizado, guiado por RL, pode produzir modelos de sobrevivência mais robustos através de níveis variados de ausência e ruído em comparação com abordagens manuais ou estáticas.
• Acessibilidade: Tornar estudos de análise de sobrevivência "mais fáceis e rápidos de realizar" para praticantes que podem não possuir a expertise de domínio para ajustar manualmente pipelines de pré-processamento complexos para dados censurados.
• Integração: Fornecer uma ponte fluida entre a limpeza de dados e a modelagem de sobrevivência a jusante, garantindo que as decisões de limpeza sejam guiadas diretamente pelo seu impacto nas previsões de sobrevivência, em vez de métricas genéricas de qualidade de dados.

Os autores mantêm um tom modesto, reconhecendo que, embora os resultados sejam promissores, a generalização para conjuntos de dados clínicos de alta dimensão e multimodais requer verificação adicional, e a sobrecarga computacional do loop de Q-learning (especialmente com modelos KNN e de aprendizado profundo) permanece uma limitação. Trabalhos futuros são propostos para estender o framework para tarefas adicionais de pré-processamento e melhorar a escalabilidade.