ShrinkageTrees: An R Package for Bayesian Tree Ensembles for Survival Analysis and Causal Inference

Riepilogo Tecnico di "ShrinkageTrees: Un pacchetto R per l'insieme di alberi Bayesiani per l'analisi della sopravvivenza e l'inferenza causale"

Problematica
L'analisi della sopravvivenza in contesti ad alta dimensionalità ($p \gg n$) affronta tre sfide primarie: (1) la relazione complessa e sconosciuta tra covariate e outcome temporali che richiede modelli non parametrici flessibili; (2) la prevalenza di osservazioni imprecise, specificamente dati con censura a destra e censura a intervallo; e (3) la difficoltà di stimare gli effetti del trattamento causale sotto censura e confondimento, specialmente quando gli effetti del trattamento variano tra i sottogruppi. Le attuali implementazioni di Bayesian Additive Regression Trees (BART) non riescono ad affrontare tutti e tre questi aspetti simultaneamente. Nello specifico, nessun pacchetto R attualmente offre ensemble di alberi Bayesiani per outcome con censura a intervallo, l'architettura $\tau$-learner (che separa gli effetti prognostici da quelli del trattamento) non è stata implementata per endpoint tempo-evento, e la regolarizzazione standard di BART è spesso inadeguata per contesti genomici ad alta dimensionalità dove il rischio di overfitting è elevato.

Metodologia
Il documento presenta ShrinkageTrees, un pacchetto R che implementa ensemble di alberi Bayesiani all'interno di un framework di Accelerated Failure Time (AFT). Il modello centrale estende la formulazione standard di BART:
$$\log T_i = \sum_{j=1}^m g(x_i; T_j, H_j) + \epsilon_i, \quad \epsilon_i \sim N(0, \sigma^2)$$
dove $T_i$ è il tempo dell'evento, $x_i$ sono le covariate, e la sommatoria rappresenta un ensemble di $m$ alberi decisionali poco profondi. Il pacchetto gestisce dati con censura a destra e a intervallo tramite data augmentation in un campionatore Gibbs.

La metodologia distingue tra due tipi di regolarizzazione, entrambi implementati e combinabili:

1. Regolarizzazione della struttura dell'albero: Per incoraggiare la divisione (splitting) su un sottoinsieme di variabili rilevanti, il pacchetto implementa la prior di partizione di Dirichlet (DART). Questa applica una prior di Dirichlet alle probabilità di partizione, inducendo una sparsità strutturale.
2. Regolarizzazione dell'altezza dei gradini (Step Height): Per restringere (shrink) i nodi foglia non informativi verso lo zero pur preservando i segnali forti, il pacchetto implementa prior di shrinkage globale-locale sui valori delle foglie ($H_j$). Il documento evidenzia la Horseshoe Forest, che applica una prior di tipo horseshoe sulle altezze dei gradini. Questa utilizza una rappresentazione come miscela di normali con un parametro di shrinkage globale ($\tau$) e parametri di scala locali ($\lambda_\ell$). La prior di horseshoe permette uno shrinkage adattivo, gestendo efficacemente il regime $p \gg n$.

Per l'inferenza causale, il pacchetto implementa il $\tau$-learner (Bayesian Causal Forests) adattato per outcome di sopravvivenza. Questo decompone l'outcome in una funzione prognostica $\mu(x, \hat{e}(x))$ e una funzione dell'effetto del trattamento $\tau(x)$, modellate da foresta separate. Ciò consente un controllo distinto della regolarizzazione e della complessità per le componenti prognostica e causale.

Contributi Chiave

• Prima implementazione di Horseshoe Forests per la sopravvivenza: Il pacchetto fornisce la prima implementazione della Horseshoe Forest, applicando una prior di tipo horseshoe sulle altezze dei gradini per abilitare l'analisi della sopravvivenza ad alta dimensionalità.
• Supporto per la censura a intervallo: È il primo pacchetto R a fornire ensemble di alberi Bayesiani per outcome di sopravvivenza con censura a intervallo.
• Inferenza causale per la sopravvivenza: Implementa il $\tau$-learner per endpoint tempo-evento, consentendo la stima di effetti del trattamento eterogenei (CATE) ed effetti medi del trattamento (ATE) in contesti di sopravvivenza.
• Regolarizzazione flessibile: Offre un framework modulare in cui gli utenti possono combinare la regolarizzazione strutturale (partizione di Dirichlet) con la regolarizzazione dell'altezza dei gradini (Horseshoe, Half-Cauchy o Normale Standard).
• Implementazione software: Il pacchetto presenta un backend C++ efficiente tramite Rcpp, supporta MCMC multi-chain e fornisce un intero workflow S3 (fitting, predizione, diagnostica di convergenza tramite coda e visualizzazione).

Risultati
Il documento valida la metodologia attraverso un dataset semi-sintetico sull'ovario e uno studio di simulazione:

• Dataset sull'ovario: In un contesto ad alta dimensionalità ($n=357, p=1004$), la BART standard ha ottenuto un C-index di training quasi perfetto (0.989), indicando overfitting. Al contrario, la Horseshoe Forest ha prodotto un C-index di training inferiore (0.727) ma ha fornito un'incertezza meglio calibrata. L'analisi causale ha stimato un piccolo effetto medio del trattamento non significativo per carboplatino rispetto a cisplatino, con intervalli di credibilità stretti.
• Studio di simulazione: Una simulazione con verità fondamentale nota e outcome con censura a intervallo ha confrontato Standard BART, DART e Horseshoe Forests attraverso dimensioni $p \in \{50, 500, 5000\}$. La Horseshoe Forest ha costantemente ottenuto il RMSE più basso, gli intervalli di credibilità più stretti e tassi di copertura nominale del 95% quasi perfetti in tutte le dimensioni. La BART standard ha prodotto intervalli eccessivamente ampi (sovra-copertura), mentre la DART ha mostuto una sotto-copertura all'aumentare della dimensionalità.
• Performance computazionale: I benchmark mostrano che, sebbene le Horseshoe Forests siano computazionalmente più costose dei modelli con prior coniugata (a causa del MCMC a salto reversibile e del campionamento delle scale locali), scalano linearmente con la dimensione del campione e sono paragonabili in velocità alle implementazioni BART esistenti per i modelli con prior coniugata.

Significatività e Rivendicazioni
Il documento afferma che ShrinkageTrees colma lacune critiche nella letteratura esistente sugli ensemble di alberi Bayesiani, unificando l'analisi della sopravvivenza, l'inferenza causale e la regolarizzazione ad alta dimensionalità. Posiziona il pacchetto come uno strumento che estende gli ensemble di alberi Bayesiani a contesti ad alta dimensionalità dove la BART standard fallisce nell'adattarsi. Gli autori sottolineano che il pacchetto è il primo a offrire il $\tau$-learner per endpoint di sopravvivenza e il primo a implementare le Horseshoe Forests per questo scopo. Il lavoro è presentato come un'implementazione pratica della metodologia sviluppata da Jacobs et al. (2025), mirata sia a professionisti esperti di ensemble di alberi Bayesiani sia a utenti che richiedono funzionalità non disponibili nei pacchetti esistenti come BART, dbarts o bcf. Gli autori notano modestamente alcuni limiti, come la necessità di covariate numeriche (che richiede la codifica dummy per le categorie non ordinate) e l'assunzione di errori distribuiti normalmente nel framework AFT.