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

Technische Samenvatting van "ShrinkageTrees: Een R-pakket voor Bayesiaanse boomensembles voor overlevingsanalyse en causale inferentie"

Probleemstelling
Overlevingsanalyse in hoogdimensionale settings ($p \gg n$) staat voor drie primaire uitdagingen: (1) de complexe, onbekende relatie tussen covariaten en tijd-tot-gebeurtenis uitkomsten die flexibele niet-parametrische modellen vereist; (2) de prevalentie van onnauwkeurige observaties, specifiek rechts-gecensureerde en interval-gecensureerde data; en (3) de moeilijkheid van het schatten van causale behandelingseffecten onder censurering en confounding, met name wanneer behandelingseffecten variëren over subgroepen. Bestaande implementaties van Bayesian Additive Regression Trees (BART) slagen er niet in om alle drie tegelijkertijd aan te pakken. Specifiek biedt geen enkel R-pakket momenteel Bayesiaanse boomensembles voor interval-gecensureerde uitkomsten, de $\tau$-learner architectuur (die prognostische en behandelingseffecten scheidt) is nog niet geïmplementeerd voor tijd-tot-gebeurtenis eindpunten, en standaard BART-regularisatie is vaak ontoereikend voor hoogdimensionale genomische settings waar overfitting een risico vormt.

Methodologie
Het artikel introduceert ShrinkageTrees, een R-pakket dat Bayesiaanse boomensembles implementeert binnen een Accelerated Failure Time (AFT) raamwerk. Het kernmodel breidt de standaard BART-formulering uit:
$$\log T_i = \sum_{j=1}^m g(x_i; T_j, H_j) + \epsilon_i, \quad \epsilon_i \sim N(0, \sigma^2)$$
waarbij $T_i$ de tijd tot de gebeurtenis is, $x_i$ de covariaten zijn, en de som een ensemble is van $m$ ondiepe beslissingsbomen. Het pakket handelt rechts- en interval-gecensureerde data af via data-augmentatie in een Gibbs-sampler.

De methodologie maakt onderscheid tussen twee typen regularisatie, die beide geïmplementeerd en combineerbaar zijn:

1. Boomstructuur-regularisatie: Om splitsingen op een subset van relevante variabelen te stimuleren, implementeert het pakket de Dirichlet splitting prior (DART). Dit plaatst een Dirichlet-prior op splitsingskansen, wat structurele sparsiteit induceert.
2. Staphoogte-regularisatie: Om oninformatieve bladeren naar nul te laten krimpen terwijl sterke signalen behouden blijven, implementeert het pakket global–local shrinkage priors op de bladwaarden ($H_j$). Het artikel benadrukt het Horseshoe Forest, dat een horseshoe-prior op staphoogtes plaatst. Dit maakt gebruik van een scale mixture of normals-representatie met een globale shrinkage-parameter ($\tau$) en lokale scale-parameters ($\lambda_\ell$). De horseshoe-prior staat adaptieve shrinkage toe, wat effectief omgaat met het $p \gg n$ regime.

Voor causale inferentie implementeert het pakket de $\tau$-learner (Bayesian Causal Forests) aangepast voor overlevingsuitkomsten. Deze deelt de uitkomst op in een prognostische functie $\mu(x, \hat{e}(x))$ en een behandelingseffectfunctie $\tau(x)$, die worden gemodelleerd door afzonderlijke bossen. Dit maakt onderscheidende regularisatie en complexiteitscontrole mogelijk voor de prognostische en causale componenten.

Belangrijkste Bijdragen

• Eerste Implementatie van Horseshoe Forests voor Overleving: Het pakket biedt de eerste implementatie van de Horseshoe Forest, waarbij een horseshoe-prior op staphoogtes wordt geplaatst om hoogdimensionale overlevingsanalyse mogelijk te maken.
• Ondersteuning voor Interval-gecensureerde Data: Het is het eerste R-pakket dat Bayesiaanse boomensembles biedt voor interval-gecensureerde overlevingsuitkomsten.
• Causale Inferentie voor Overleving: Het implementeert de $\tau$-learner voor tijd-tot-gebeurtenis eindpunten, wat de schatting van heterogene behandelingseffecten (CATE) en gemiddelde behandelingseffecten (ATE) in overlevingscontexten mogelijk maakt.
• Flexibele Regularisatie: Het biedt een modulair raamwerk waar gebruikers structurele regularisatie (Dirichlet splitting) kunnen combineren met staphoogte-regularisatie (Horseshoe, Half-Cauchy, of Standaard Normaal).
• Software Implementatie: Het pakket beschikt over een efficiënte C++ backend via Rcpp, ondersteunt multi-chain MCMC, en biedt een volledige S3-workflow (fitting, predictie, convergentiediagnostiek via coda, en visualisatie).

Resultaten
Het artikel valideert de methodologie via een semi-synthetische ovariumkankerdataset en een simulatiestudie:

• Ovarium Dataset: In een hoogdimensionale setting ($n=357, p=1004$) bereikte standaard BART een bijna perfecte trainings C-index (0.989), wat duidt op overfitting. In contrast hiermee leverde de Horseshoe Forest een lagere trainings C-index (0.727), maar bood het beter gekalibreerde onzekerheid. De causale analyse schatte een klein, niet-significant gemiddeld behandelingseffect voor carboplatine versus cisplatine, met nauwe geloofwaardigheidsintervallen.
• Simulatiestudie: Een simulatie met bekende grondwaarheid en interval-gecensureerde uitkomsten vergeleek Standaard BART, DART en Horseshoe Forests over dimensies $p \in \{50, 500, 5000\}$. De Horseshoe Forest behaalde consequent de laagste Root Mean Squared Error (RMSE), de nauwste geloofwaardigheidsintervallen en bijna nominale 95% dekking rates over alle dimensies. Standaard BART produceerde overly brede intervallen (over-coverage), terwijl DART onder-dekte naarmate de dimensionaliteit toenam.
• Computationele Prestaties: Benchmarks tonen aan dat hoewel Horseshoe Forests computationeel duurder zijn dan conjugate-prior modellen (door reversible jump MCMC en het samplen van lokale scales), ze lineair schalen met de steekproefomvang en qua snelheid vergelijkbaar zijn met bestaande BART-implementaties voor conjugate modellen.

Betekenis en Claims
Het artikel claimt dat ShrinkageTrees kritieke hiaten in de bestaande Bayesiaanse boomensemble-literatuur aanpakt door overlevingsanalyse, causale inferentie en hoogdimensionale regularisatie te verenigen. Het positioneert het pakket als een instrument dat Bayesiaanse boomensembles uitbreidt naar hoogdimensionale settings waar standaard BART faalt om aan te passen. De auteurs benadrukken dat het pakket de eerste is die de $\tau$-learner voor overlevings-eindpunten biedt en de eerste die Horseshoe Forests voor dit doel implementeert. Het werk wordt gepresenteerd als een praktische implementatie van de methodologie ontwikkeld door Jacobs et al. (2025), gericht op zowel praktici die nieuw zijn in Bayesiaanse boomensembles als ervaren gebruikers die functionaliteit nodig hebben die niet beschikbaar is in bestaande pakketten zoals BART, dbarts, of bcf. De auteurs merken bescheiden beperkingen op, zoals de vereiste van numerieke covariaten (wat dummy-encoding voor ongeordende categorieën noodzakelijk maakt) en de aanname van normaal verdeelde fouten in het AFT-raamwerk.