xgbse: Análise de Sobrevivência robusta e eficiente em Python

Davi Vieira
Loft
Published in
13 min readFeb 11, 2021

xgbse: XGBoost Survival Embeddings, nossa biblioteca de análise de sobrevivência, agora open-source!

“Existem duas culturas no uso da modelagem estatística para chegar a conclusões a partir dos dados. Uma assume que os dados são gerados por um determinado modelo de dados estocástico. A outra utiliza modelos algorítmicos e considera o mecanismo dos dados algo desconhecido.”

Leo Breiman, Statistical Modeling: The Two Cultures
(em tradução livre)

Na Loft, aplicamos ciência de dados em diferentes etapas da jornada de compra e venda de um imóvel, buscando trazer conforto e segurança às pessoas envolvidas neste processo que exige bastante cuidado e transparência. Modelar o tempo até determinados eventos nessa jornada é parte fundamental da nossa estratégia como companhia (ex.: o tempo até venda de um imóvel). Nossos modelos de liquidez têm uma vasta gama de aplicações, desde servir como um termômetro para a demanda do marketplace e ajudar pessoas vendedoras a melhorar seus anúncios até projetar o fluxo de caixa esperado para os nossos fundos imobiliários.

Para construir esses modelos utilizamos um conjunto de técnicas chamado Análise de Sobrevivência, a área da estatística interessada em estudar a duração entre eventos. Particularmente, a análise de sobrevivência nos permite lidar com um tipo especial de informação faltante chamado censura. Por exemplo: um imóvel está no mercado há 90 dias e ainda não foi vendido. Sabemos que o tempo até a venda dele é de pelo menos 90 dias, mas não podemos afirmar muito mais do que isso. Como incluir essa informação nos nossos modelos? A análise de sobrevivência tem a resposta para essa pergunta.

Caso você não tenha familiaridade com este tema, preparamos este post onde explicamos o porquê de usá-la e qual seu valor no mercado imobiliário e em outros mercados também.

O principal resultado neste tipo de modelagem é uma curva de sobrevivência, que representa a probabilidade do evento de interesse não acontecer até determinado tempo. No nosso caso, isso significa o imóvel “sobreviver” (não vender) ao longo do tempo, como mostra o gráfico abaixo. Algumas técnicas de modelagem fazem simplificações desta curva, disponibilizando somente o seu valor esperado ou um fator de decaimento (risco).

Curva de sobrevivência — no eixo X o tempo e no eixo Y a probabilidade de sobrevivência — P(T>t)

Devido à criticidade das aplicações de liquidez na Loft, principalmente no que diz respeito ao nosso negócio de compra, reforma e venda (onde temos alto capital alocado), precisamos de uma solução de análise de sobrevivência que entregue alta capacidade preditiva e expressividade sem abrir mão do rigor estatístico, explicabilidade e calibração de probabilidades.

Contudo, ao investigar ferramentas disponíveis no mercado, não conseguimos encontrar uma solução completa. Em geral, encontramos dois cenários bem definidos:

  1. A cultura “da estatística tradicional”: modelos com curvas de sobrevivência calibradas (probabilidade que o modelo produz é igual à fração do evento na vida real) e explicabilidade, porém deixando a desejar em eficiência computacional e capacidade preditiva.
  2. A cultura “do machine learning”: modelos com alta eficiência computacional e capacidade preditiva, porém com baixa explicabilidade e falta de rigor estatístico: curvas de sobrevivência descalibradas ou apenas a possibilidade de estimar riscos relativos que não conversam com estimativas de tempo.

Dado o nosso objetivo, foi ficando cada vez mais claro que precisávamos unir esses dois mundos em um só.

Apresentando o XGBoost Survival Embeddings

No incrível artigo, citado no início desse post, Statistical Modeling: The Two Cultures (2001), Leo Breiman, criador dos algoritmos Classification and Regression Trees (CART) e Random Forests, apresenta duas culturas de modelagem estatística: modelagem de dados (Data Modeling Culture) e modelagem algorítmica (Algorithmic Modeling Culture).

Duas culturas de modelagem estatística: dados e algorítmica (adaptado de Breiman, 2001)

A primeira, cultura de modelagem de dados, é representada pelo grupo de modelos tradicionais estatísticos, calibrados e com alta explicabilidade. Em geral, esta cultura produz modelos que são ótimos “estimadores”, com garantias de que em média as estimativas de tempo vão coincidir com a realidade, mas pecam no papel de “ordenadores”: saber que um indivíduo é mais propenso a sobreviver que outro.

A segunda, cultura de modelagem algorítmica, mais representada por modelos de machine learning, de alta eficiência computacional e alta capacidade preditiva, segue um padrão inverso: são ótimos “ordenadores” através de seus modelos expressivos e não-lineares, mas pecam no papel de “estimadores”, muitas vezes entregando previsões descalibradas ou adimensionais, que não têm interpretação além de somente um “score” de sobrevivência.

Motivados e com o foco em atingir o objetivo de unir essas duas culturas, criamos o XGBoost Survival Embeddings — xgbse, nosso pacote de análise de sobrevivência em Python. Durante o último ano viemos utilizando esta ferramenta internamente na Loft e agora estamos disponibilizando ela para a comunidade em formato open source!

Devido ao seu alto poder preditivo e performance, o XGBoost foi usado como base para esta construção. Sobre essa base, implementamos módulos que adicionam uma camada de rigor estatístico ao modelo, e permitem cobrir o restante das brechas que a sua implementação padrão de análise de sobrevivência não cobria.

A seguir mostramos e exemplificamos esses problemas.

Análise de Sobrevivência no XGBoost

O XGBoost disponibiliza duas adaptações feitas a partir de modelos tradicionais da estatística:

  • Modelo de Cox (sem estimativa do risco basal): XGBoost Cox;
  • Modelo AFT (Accelerated Failure Time): XGBoost AFT.

Ambos são modelos que mantém a alta performance computacional característica da lib XGBoost, com rápido treinamento e predição mesmo para grandes conjuntos de dados. Além disso, notamos que, em validação, ambos apresentam um alto poder de discriminação, próximo ao estado da arte, mensurados pelo C-index (métrica equivalente ao ROC AUC considerando dados censurados).

Contudo, O XGBoost Cox disponibiliza predições para o Hazard Ratio (HR) das amostras, o que se aproxima mais ao “score” de sobrevivência, ou seja, discrimina bem as amostras somente quanto a ordem. Não é disponibilizada a taxa de risco basal do modelo, o que inviabiliza recuperar a curva completa de sobrevivência, primordial para, por exemplo, nossos controles de risco de liquidez na Loft.

Partindo para a segunda opção, XGBoost AFT, este nos fornece uma implementação flexível, lidando com diferentes tipos de censuras e entregando previsões pontuais de tempo de sobrevivência. Não era exatamente o que desejávamos (curvas de sobrevivência), porém estávamos mais próximos dos nossos objetivos listados previamente.

Porém, observando as estimativas de tempo esperado de sobrevivência obtidas pelo XGBoost AFT, notamos que estas eram muito sensíveis aos hiperparâmetros utilizados. A título de exemplificação (acesso ao script aqui), geramos três cenários variando um de seus hiperparâmetros, o aft_loss_distribution_scale, em um dataset padrão de análise de sobrevivência (METABRIC). Para um mesmo conjunto de amostras, medimos o C-index e calculamos a média global do tempo esperado de sobrevivência previsto pelo modelo, chegando no seguinte resultado:

aft_loss_distribution_scale: 1.5
C-index: 0.645
Tempo médio estimado de sobrevivência: 203 dias
----
aft_loss_distribution_scale: 1.0
C-index: 0.648
Tempo médio estimado de sobrevivência: 165 dias
----
aft_loss_distribution_scale: 0.5
C-index: 0.646
Tempo médio estimado de sobrevivência: 125 dias
----

Conseguimos ver que, apesar da capacidade de ordenação das amostras ser mantida, avaliada pelo C-index, tivemos uma variância significativa para a média global dos tempos previstos pelo modelo (78 dias). Se estivéssemos analisando uma única amostra isso até seria compreensível, mas essa variação diz respeito ao conjunto completo de validação, onde a previsão média não deveria variar tanto independentemente do modelo utilizado. Nos questionamos: em que modelo confiar?

Comparando, sobre a mesma amostra, estes 3 cenários com a curva de sobrevivência estimada usando o estimador de Kaplan-Meier (estimador não-viesado da curva de sobrevivência) chegamos no seguinte resultado:

XGBoost AFT — 3 cenários e suas estimativas versus o estimativas retiradas pelo Kaplan-Meier (estimador não-viesado da curva de sobrevivência)

Observamos que para cada variação de 0.5 no nosso hiperparâmetro geramos previsões em diferentes decis na curva de sobrevivência, ou seja, representando riscos absolutos muito diferentes. Além disso, não contamos com intervalos de confiança para as estimativas geradas. Essa falta de estabilidade nas estimativas é um fator desmotivador para a adesão e apoio dos stakeholders no uso deste modelo em processos de alto valor e risco para o negócio como, por exemplo, alocar capital em um portfolio de imóveis, ou quem sabe, construir um cronograma de manutenção em uma frota de caminhões.

Os pontos anteriores nos trouxeram questionamentos sobre como poderíamos utilizar a performance do XGBoost e sua alta capacidade de discriminação e ainda assim obter estimativas para curvas de sobrevivência completas e intervalos de confiança, permitindo o uso de tais modelos para nossos objetivos de negócio. A seguir, ilustraremos a ideia base: utilizar o XGBoost como um transformador de features.

XGBoost como um transformador de features

Modelos de árvores são essencialmente conjuntos de decisões/quebras que dividem nossa base de dados sobre as features que utilizamos no treinamento. Estas divisões são definidas com um simples objetivo: encontrar subconjuntos de amostras que discriminem bem a variável alvo. No caso de árvores de sobrevivência essa discriminação é feita a fim de obter grupos de amostras com sobrevivências distintas ao fim de cada árvore, ou seja, em suas folhas.

Ensemble de árvores e o caminho percorrido (em laranja) por uma amostra. Podemos construir um vetor que representa em que folhas essa amostra esteve (1) e não esteve (0) por estimador/árvore, ou seja, um one-hot encoding das folhas

A informação sobre em que nó terminal/folha cada amostra da nossa base se encontra é incrivelmente valiosa. Ela representa grande parte da informação usada pelo ensemble para predizer o evento de interesse.

Afinal, amostras que passam pelas mesmas decisões nas árvores e terminam em folhas iguais devem ser bastante parecidas sobre o que pode ser esperado da target e features, certo?

Sendo assim, em um ensemble de árvores, a representação obtida através da ocorrência em diferentes folhas (e.g. um vetor binário, como na figura) é bastante poderosa. Ao usá-la, saímos do espaço original das features X para um espaço (embedding de árvores) que tem algumas propriedades especiais:

  • Esparso e de alta dimensionalidade: ao deixarmos o trabalho de lidar com as não linearidades dos dados para o ensemble de árvores, temos como resultado um espaço de folhas onde mesmo modelos lineares ao serem treinados sobre esse espaço obtém resultados competitivos versus o próprio ensemble de árvores e com benefícios como a calibração de suas estimativas.
  • Supervisionado: as folhas trazem amostras que são similares sobre a variável alvo quando condicionadas nas features (quebras das árvores), ou seja, o XGBoost serve de filtro escolhendo features com alto potencial preditivo, eliminando ruído e features redundantes. Sendo assim, nosso embedding está muito menos sujeito à maldição da dimensionalidade e funciona muito bem com modelos de vizinhos próximos, que, por sua vez, nos permitem a flexibilidade de usar métodos robustos (como o Kaplan-Meier).
Características desejáveis do embedding criado à partir das folhas de um ensemble de árvores

Os princípios listados acima formam a base para a construção dos modelos atualmente existentes no XGBoost Survival Embeddings, que serão introduzidos a seguir.

Vizinhos próximos e Kaplan-Meier

Ao comparar os vetores de folhas terminais para cada amostra podemos ter uma noção de similaridade no embedding construído. Como mostramos anteriormente, amostras que se encontram nas mesmas folhas em diferentes árvores devem ser bastante parecidas. O mesmo podemos dizer do inverso, caso ao comparar amostras estas não tenham nenhuma folha igual, temos uma menor semelhança. Assim, podemos comparar diferentes amostras e computar distâncias entre elas com base na coocorrência de folhas, usando por exemplo a distância de Hamming.

Com essa base, aplicamos algoritmos de vizinhos próximos (Nearest Neighbors) usando essa distância, obtendo para novas amostras seus N vizinhos no embedding. Por fim, utilizamos o estimador de Kaplan-Meier sobre esses vizinhos, obtendo uma estimativa para a curva de sobrevivência para a nova amostra. Então, temos o XGBSEKaplanNeighbors.

PS.: recomendamos o uso do booster dart a fim de evitar que alguma das árvores domine a variância do ensemble e consequentemente quebre a lógica de distâncias.

Construímos uma forma vetorizada de ajustar o estimador de Kaplan-Meier e calcular seus intervalos de confiança usando a fórmula Exponencial de Greenwood, deixando esse último estágio bastante performático. A busca de vizinhos em si é uma etapa custosa do ponto de vista computacional, o que implica em uma performance mais lenta deste algoritmo para dados em alta escala, porém ainda superior versus alguns benchmarks.

Também é possível obter estimativas de sobrevivência para intervalos de tempo pré-definidos pela pessoa usuária.

Múltiplas regressões logísticas treinadas no embedding

Este modelo (XGBSEDebiasedBCE) é considerado um modelo de tempos discretos (discrete time survival analysis model), isto significa que dividimos o tempo avaliado em diferentes intervalos discretos, o que torna necessário a pessoa usuária disponibilizar as janelas de tempo em que os eventos serão avaliados.

Sua construção segue os seguintes passos:

  • Utilizamos o XGBoost como transformador de features e geramos o embedding de árvores em um esquema de one-hot encoding.
  • O embedding é usado como entrada em múltiplas regressões logísticas, onde cada regressão tem o objetivo de prever a probabilidade de um evento acontecer em determinada janela de tempo;
  • À medida que avançamos ao longo das janelas, os indivíduos em risco, eventos e censuras são computados e avaliados, como também são removidas as censuras ao longo do tempo. A regressão logística busca estimar a proporção de falhas sobre indivíduos em risco no intervalo de tempo avaliado, ou seja, o termo di/ni encontrado no estimador de Kaplan-Meier (para mais informações sobre esse estimador, veja o nosso post — seção Kaplan-Meier);
  • Utilizar apenas as probabilidades de falha (termo di/ni) computadas pelas regressões logísticas nos traria uma estimativa viesada da curva de sobrevivência. Assim, usamos a fórmula do estimador de Kaplan-Meier para transformar esses diferentes termos avaliados nas diferentes janelas de tempo em uma estimativa não-viesada da curva de sobrevivência (utilização do complementar de cada termo 1-di/ni e é retirado o produtório para todos os tempos avaliados).

Conseguimos escalar o modelo para treinar e predizer utilizando grandes bases de dados (milhões de exemplos) de maneira eficiente através de paralelização das regressões logísticas utilizando o joblib .

Por fim, também conseguimos computar intervalos de sobrevivência para as nossas estimativas através do uso da técnica de Bootstrap. Implementamos um meta-estimador chamado XGBSEBootstrapEstimator para executar este cálculo para qualquer estimador base do xgbse.

Outras features

Além dos módulos apresentados neste post, temos outras features que valem ser destacadas:

  • XGBSEBootstrapEstimator: meta-estimador construído a partir de reamostragens com reposição;
  • O modelo XGBSEKaplanTree: versão similar ao XGBSEKaplanNeighbors porém onde estimamos apenas uma árvore, retirando a necessidade de executar buscas de vizinhos. O estimador de Kaplan-Meier utiliza todas as amostras da folha para estimar a curva de sobrevivência. Sua implementação é altamente eficiente, porém por ser apenas um estimador, sua capacidade preditiva é reduzida. Recomendamos seu uso em conjunto com o meta estimador citado acima XGBSEBootstrapEstimator;
  • Explicabilidade através de protótipos: é possível fazer a query dos vizinhos utilizados nas predições dos módulos XGBSEKaplanNeighbors e do XGBSEKaplanTree. Assim, pode-se entender que tipos de amostras foram usadas para compor a construção da curva de sobrevivência de determinada amostra, utilizando-as para avaliar e explicar previsões;
  • Extrapolação de curvas de sobrevivência;
  • Construção de intervalos de confiança paramétricos e não-paramétricos .

Esse conjunto de features nos trouxe soluções para alguns dos problemas listados no início desse post, permitindo que tenhamos curvas de sobrevivência individualizadas, intervalos de confiança para nossas estimativas e melhor mensuração do risco de predição.

E o problema original do XGBoost para Análise de Sobrevivência?

Voltando ao exemplo no início do post, vamos refazer a análise de sensibilidade aos hiperparâmetros do XGBoost AFT, mas agora usando-o apenas como transformador de features através do XGBSEKaplanNeighbors.

Logo abaixo, ilustramos o código exemplificando a interface do xgbse (aqui para vê-lo completo):

E os resultados são mostrados a seguir:

aft_loss_distribution_scale: 1.5
C-index: 0.640
Probabilidade média de sobrevivência nos [30, 90, 150] dias: [0.9109, 0.6854, 0.528]
----
aft_loss_distribution_scale: 1.0
C-index: 0.644
Probabilidade média de sobrevivência nos [30, 90, 150] dias: [0.9111, 0.6889, 0.5333]
----
aft_loss_distribution_scale: 0.5
C-index: 0.650
Probabilidade média de sobrevivência nos [30, 90, 150] dias: [0.913, 0.6904, 0.5289]

Podemos perceber que mantivemos a estabilidade do C-index para os diferentes hiperparâmetros e agora conseguimos também estimar em conjunto as curvas de sobrevivência, que se mantém estáveis para as diferentes configurações!

Tal como foi feito anteriormente, podemos comparar nossas predições com o estimador de Kaplan-Meier:

Ao comparar o XGBSEKaplanNeighbors com o Kaplan-Meier (estimador não-viesado da curva de sobrevivência) podemos notar a estabilidade das predições, com redução significativa do viés e estimativas da curva de sobrevivência para todos os tempos avaliados!

Comece a usar o XGBoost Survival Embeddings

Esperamos que a nossa visão de como a Análise de Sobrevivência pode trazer impacto em diferentes indústrias seja fortalecida para fronteiras além das da Loft. Convidamos a comunidade a testar o XGBoost Survival Embeddings, usar em suas empresas e contribuir com o desenvolvimento da lib.

XGBoost Survival Embeddings é instalável via pip e compatível com a API do scikit-learn. A documentação do projeto se encontra disponível no https://loft-br.github.io/xgboost-survival-embeddings/index.html e o repositório no Github pode ser acessado por aqui. Recomendamos a seção de exemplos como uma boa forma de se familiarizar com as ferramentas e uso da lib.

Se tiverem qualquer comentário e/ou dúvidas sobre o XGBoost Survival Embeddings — xgbse, podem nos contatar através do email bandits@loft.com.br . Todo feedback é bem-vindo! Esperamos escutar de vocês o que acharam!

Por Davi Vieira, Guilherme Marmerola, Gabriel Gimenez e Vitor Estima

Quer fazer parte da Loft e nos ajudar a simplificar o mercado imobiliário?

Confira nossas vagas! — https://jobs.lever.co/loft (ps.: estamos com vagas abertas no nosso time de ciência de dados!)

--

--