SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sebastian Salazar Escobedo
sebastian.salazar@cs.columbia.edu
Decision trees are a well-established tool in machine learning for classification and regression tasks. In this paper, we introduce a novel non-parametric Bayesian model that uses variational inference to approximate a posterior distribution over the space of stochastic decision trees. We evaluate the model's performance on 18 datasets and demonstrate its competitiveness with other state-of-the-art methods in regression tasks. We also explore its application to causal inference problems. We provide a fully vectorized implementation of our algorithm in PyTorch.
1 Introduction and Related Work
The use of binary trees for classification and regression problems is widespread in the Machine Learning community (see Breiman [2001], Criminisi [2011]). Chipman et al. [1998] proposed the Bayesian CART method, which offers a fully Bayesian treatment of the problem, aiming to learn a distribution over trees as opposed to a single tree. This innovation led to the development of more sophisticated algorithms, such as Bayesian Additive Regression Trees (BART) (see Chipman et al. [2010]), which has stood the test of time and still shows competitive performance to this day. Both BART and Bayesian CART approximate the posterior distribution over trees through the use of Markov Chain Monte Carlo (MCMC) methods.
In the period following the inception of the Bayesian CART methodology, the field has witnessed significant advancements in approximate Bayesian Inference. This paper explores approximating a posterior distribution over the space of trees via variational inference, a compelling MCMC alternative based on an optimization problem as elaborated in Blei et al. [2017]. Our approach introduces a unique stochastic process over the space of trees, serving as a model prior. Notably, despite the infinite parameter specification of this stochastic process, we show that variational inference remains feasible.
This paper is also related to work on soft decision trees (Kontschieder et al. [2015]), where the goal is to learn a stochastic decision tree of a fixed structure by minimizing some sort of loss function (see Section 2.1). In this paper, we adopt a fully Bayesian perspective to this problem and instead learn a distribution over soft decision trees of varying structure. A fully Bayesian approach has the advantage in that it allows us to naturally quantify the uncertainty of the predictions made by our model.
Overall, our contributions in this paper include a novel stochastic process over the space of trees, a variational inference algorithm for approximating the posterior distribution over trees, and a demonstration of the performance of our model on both regression and causal inference tasks. We believe that this work opens up new avenues for future research in the area of Bayesian Nonparametrics and decision trees.
2.1 Stochastic Decision Trees
A stochastic tree, denoted as  $T$ , is a binary tree where each of its  $b - 1$  internal nodes is associated with a soft decision rule. This rule indicates the probability of a particular data point taking the left branch of the tree. In this paper, we focus on decision rules of the form  $P(\text{datapoint } x \text{ at node } i \text{ takes the left branch}) = \sigma(\beta_i^T X)$ . Moreover, we introduce a prior probability measure, denoted as  $\Pi(\Theta, T)$ , that encompasses the space of trees  $T$  (defined by splits and decision rules) and the associated leaf parameters  $\Theta$ . The generative process for the data is divided into two parts:
1. Prior Distribution:
$$
(T, \Theta) \sim \Pi (T, \Theta) \tag {1}
$$
Unlike BART or Bayesian CART, our prior distribution is designed differently, incorporating a sampling algorithm that generates a collection of parameters used to fully and uniquely reconstruct a tree. The second part of the model is the likelihood that gives rise to the observed data
2. Likelihood:
$$
\left(y _ {i} \mid X _ {i}, T, \Theta\right) \sim p \left(y _ {i} \mid x _ {i}, T, \Theta\right) \tag {2}
$$
The likelihood represents the conditional distribution of  $y_{i}$  given  $x_{i}$ ,  $T$ , and  $\Theta$ . In our case, it is usually modeled as a Normal distribution. Further details regarding the likelihood and the overall model will be discussed in subsequent sections.
2.2 Likelihood
We associate each of the  $b$  leafs of a stochastic tree  $T$  with a parameter  $\Theta \coloneqq \{\theta_1,\dots,\theta_b\}$ . Every parameter  $\theta_{i}\in \Theta$  specifies a density over  $\mathcal{V}$ , denoted by  $f(y|\theta_i)$ . Throughout, we assume that the features  $x\in \mathcal{X}$  are fixed, and that the target variables  $y\in \mathcal{V}$  are sampled from a distribution whose parameters stochastically determined by having  $x\in \mathcal{X}$  traverse down the tree. In other words, if  $x\in \mathcal{X}$  is assigned to leaf  $i$ , then  $y\in \mathcal{V}$  is sampled from  $f(y|\theta_i)$ . If a feature  $x\in \mathcal{X}$  is assigned to leaf node  $i$  with a certain probability, this implies that the likelihood  $p(y|X,T,\Theta)$  is a mixture distribution over the densities  $f(y|\theta_i)$  and is given by
$$
p (y | T, \Theta , x) = \sum_ {i \in \text {L e a f}} P (x \text {t a k e s a p a t h l e a d i n g t o n o d e} i) f (y | \theta_ {i}). \tag {3}
$$
2.3 A Bayesian non-parametric process prior over the space of stochastic trees  $\Pi (\Theta ,T)$
We define a prior distribution over a stochastic decision trees by specifying a sampling algorithm. The joint distribution over decision trees  $T$  and it's corresponding leaf parameters  $\Theta$  is specified through an infinite number of parameters that have a one-to-one correspondence with characteristics of a stochastic decision tree. The parameters required to draw a sample from the joint distribution over  $(T, \Theta)$  —which we denote by  $\Pi(T, \Theta)$  are as follows<sup>2</sup>
1.  $\{\gamma_i\}_{i=1}^{\infty} \in [0,1]^{\infty}$  specify the probabilities of splitting a node in the tree. In other words, node  $i$  in the tree is split with probability  $\gamma_i$ .
2.  $\{\mu_i\}_{i = 1}^{\infty}$  are the parameters that define the distribution over the weights specifying the splitting rules.
3.  $\{\zeta_i\}_{i = 1}^{\infty}$  are the parameters that define the conditional likelihoods over the observations falling in node  $i$ , provided that it is a leaf.
<sup>1</sup>The densities are relative to an appropriate dominating measure  $\nu$ . In practice,  $\nu$  will usually be the Lebesgue or the counting measure.
In this paper, we adopt the following convention for numbering the nodes of a binary tree: (a) We start by labeling the root node as 1 (b) recursively label the left children of node  $j$  with  $2j$  and the right children of node  $j$  with  $2j + 1$ .
Figure 1: Taking a sample from the generative process
Figure 2: Graphical Model associated to the data generating process.
Figure 3: Graphical Model associated to the data generating process. The parameters  $s, \beta$ , and  $\theta$  are used to uniquely specify a stochastic decision tree and its parameters.
Figure 4: Variational approximation to the data generating process.
Given these parameters, a stochastic tree is sampled as follows:
1. For each node node  $i\in \{1,2,\ldots \}$
(a) Sample  $s_1,\ldots ,s_k,\ldots \sim$  Bernoulli  $(\gamma_{i})$
(b) Sample  $\beta_{1},\dots,\beta_{k},\dots\sim\mathrm{ExpFam}(\mu_{i})$
(c) Sample  $\theta_{1},\dots,\theta_{k},\dots\sim\mathrm{ExpFam}(\zeta_{i})$
Note that the parameters  $(s_1, \dots, \beta_1, \dots, \theta_1, \dots)$  suffice to completely specify a binary stochastic tree. We give a concrete example of this generative process in Figure 1 and the graphical model associated to this data-generating process is shown in 3.
2.4 A variational family over the space of trees and tree parameters  $q(\Theta, T)$
A variational family is obtained by truncating the Bayesian Nonparametric process described in section 2.3 at a fixed tree depth  $d$ . This truncation is depicted as a graphical model in figure 4. A similar pattern is observed in variational inference for Dirichlet Process Mixtures proposed by Blei and Jordan [2006], where the stick-breaking process of the Dirichlet Process mixture is also
truncated after a fixed number of steps. It is important to note that only the variational distribution is truncated, not the underlying model itself. Furthermore, we would like to emphasize that the choice of exponential family in the variational distribution aligns with those chosen for the prior distribution. To distinguish them from the prior parameters, we label the variational parameters with a tilde (i.e.  $\tilde{\gamma},\tilde{\mu},\tilde{\zeta}$ ) on top.
2.5 Maximizing the ELBO
We obtain an approximation of the posterior distribution by maximizing a weighted version of the evidence lower bound (ELBO) with respect to the model parameters:3
$$
\begin{array}{l} \operatorname {E L B O} (q) = \mathbb {E} _ {q} (\log p (Y | X, \Theta , T)) - \lambda \mathrm {K L} (q (\Theta , T) | | p (\Theta , T)) (4) \\ \operatorname {E L B O} (q) = \mathbb {E} _ {q} (\log p (Y | X, \Theta , T)) (5) \\ + \lambda \mathbb {E} _ {q} (\log p (\Theta , T)) (6) \\ - \lambda \mathbb {E} _ {q} (\log q (\Theta , T)). (7) \\ \end{array}
$$
Where  $\lambda$  is a hyperparameter controlling the amount of regularization. The gradients of the variational objective are approximated stochastically with Monte Carlo samples from the variational family. To obtain low-variance estimates of the gradient estimates, we make use of the reparametrization trick (see Kingma and Welling [2022], Ruiz et al. [2016]), where a Softmax-Gumbel approximation is used for the Bernoulli random variables (Jang et al. [2017]). Other approaches for variance-reduction like Rao-Blackwellization are possible but these ideas were not explored in this paper (see Ranganath et al. [2013]). Given a sample from the variational family  $(s_1,\dots,s_{2^{d - 1}},\beta_1,\dots,\beta_{2^{d - 1}},\theta_1,\dots,\theta_{2^d})$  we evaluate each of the terms in the ELBO as follows:
$$
\log p (y \mid T, \Theta , x) = \log \left(\sum_ {i \in \text {L e a f}} P (x \text {t a k e s a p a t h l e a d i n g t o n o d e} i) f (y \mid \theta_ {i})\right) \tag {8}
$$
$$
\begin{array}{l} \log p (\Theta , T) = \sum_ {i \in \text {I n t e r n a l N o d e s}} \log \left(\gamma_ {i}\right) + \sum_ {i \in \text {I n t e r n a l N o d e s}} \log \left(\exp \operatorname {F a m} \left(\beta_ {i}; \mu_ {i}\right)\right) (9) \\ + \sum_ {i \in \text {L e a v e s}} \log \left(1 - \gamma_ {i}\right) + \sum_ {i \in \text {L e a v e s}} \log \left(\exp \operatorname {F a m} \left(\theta_ {i}; \zeta_ {i}\right)\right) (10) \\ \end{array}
$$
$$
\begin{array}{l} \log q (\Theta , T) = \sum_ {i \in \text {I n t e r n a l N o d e s}} \log \left(\tilde {\gamma} _ {i}\right) + \sum_ {i \in \text {I n t e r n a l N o d e s}} \log \left(\exp \operatorname {F a m} \left(\beta_ {i}; \tilde {\mu} _ {i}\right)\right) (11) \\ + \sum_ {i \in \text {L e a v e s n o t i n l a s t l e v e l}} \log \left(1 - \tilde {\gamma} _ {i}\right) + \sum_ {i \in \text {L e a v e s}} \log \left(\exp \operatorname {F a m} \left(\theta_ {i}; \tilde {\zeta} _ {i}\right)\right) (12) \\ \end{array}
$$
In order to make use of reparametrization gradients, each of the three terms in the ELBO must be written in terms of the model parameters (i.e.  $\{\gamma, \mu, \zeta, \tilde{\gamma}, \tilde{\mu}, \tilde{\zeta}\}$ ) and a reparametrization for each of the samples must be specified. These calculations are deferred to the next section.
3 Variational Regression Trees
In this paper, we explore the use of Variational Regression Trees. We propose using the normal distribution as a natural choice for the splitting and prediction vectors of the prior and variational families, as follows:4
$$
\beta_ {i} \sim \mathcal {N} \left(\mu_ {i}, S _ {i}\right) \tag {13}
$$
$$
\theta_ {i} \sim \mathcal {N} \left(\zeta_ {i}, \Psi_ {i}\right) \tag {14}
$$
Where  $S_{i}$  and  $\Psi_{i}$  are diagonal covariance matrices. We also assign a normal distribution to each of the leaves in the tree by setting  $f(y|\theta_i)$  to a normal distribution with parameters  $\mathcal{N}(\theta_i^T x, \sigma_i^2)$ . The normal distribution is a natural choice for regression problems and allows for an easy application of
3 Whenever  $\lambda = 1.0$ , this is equivalent to minimizing the KL divergence between the variational family and the posterior distribution over the space of trees.
This shows the parameters of the prior, the corresponding parameters for the variational family are denoted with the same letters but with a tilde on top.
the reparametrization trick. For classification, adding a probit link function to the likelihood should result in a tractable variational inference algorithm. We believe including a logit link on the outcome  $y$  is also possible, but would require a Poly-Gamma augmentation (see Polson et al. [2013]). We leave all of these modifications to future work (see section 5 for an overview on potential future work).
3.1 Learning the Variational Family
The parameters of the variational family are learned by minimizing Monte Carlo approximations of the ELBO with gradient descent. The naive Monte Carlo Estimators for this types of problem are given by  $\nabla_{\phi}\mathbb{E}_{q_{\phi}(x)}(f(x)) = \mathbb{E}_{q_{\phi}}(f(x)\nabla_{\phi}\log q_{\phi}(x))\approx \frac{1}{N}\sum_{i = 1}^{N}f(x^{(i)})\nabla_{\phi}\log q_{\phi}(x^{(i)})$ . Although asymptotically consistent, this estimator exhibits extremely high variance, rendering it impractical for the purposes of stochastic optimization. Instead, we make use of a well-known variance-reduction trick through a reparametrization of the samples obtained from the variational family. In particular, we rewrite samples from the variational as a function of an auxiliary noise variable as follows
$$
\epsilon \sim p (\epsilon) \tag {15}
$$
$$
x = g (\phi , \epsilon). \tag {16}
$$
The reparametrization of the normal distribution of equations 13 and 14 is straightforward and is given below
$$
\epsilon_ {r u l e s}, \epsilon_ {l e a v e s} \sim_ {i i d} \mathcal {N} (0, I _ {d}) \tag {17}
$$
$$
\beta_ {i} = S _ {i} ^ {1 / 2} \epsilon_ {\text {r u l e s}} + \mu_ {i} \tag {18}
$$
$$
\theta_ {i} = \Psi_ {i} ^ {1 / 2} \epsilon_ {\text {l e a v e s}} + \zeta_ {i} \tag {19}
$$
Reparametrizing the Bernoulli random variables requires use of the Softmax-Gumbel trick of Jang et al. [2017]. The key insight relies on the fact that samples from a Bernoulli distribution with success probability  $\gamma$  may be obtained as follows
$$
v = \text {o n e} _ {\text {h o t}} (\arg \max  (g _ {1} + \log (\gamma), g _ {2} + \log (1 - \gamma))) \tag {20}
$$
$$
s = (1, 0) ^ {T} v \tag {21}
$$
Where  $g_{1}, g_{2}$  are independent samples from a Gumbel(0, 1) distribution. In practice, we apply equation 20 using the following limiting representation of the onehot argmax function
$$
v = \lim  _ {\tau \rightarrow 0 ^ {+}} \operatorname {s o f t m a x} \left(\frac {g _ {1} + \log (\gamma)}{\tau}, \frac {g _ {2} + \log (1 - \gamma)}{\tau}\right). \tag {22}
$$
During training, we set  $\tau$  to a fixed value and then anneal it towards zero as the training progresses. These reparametrizations allow us to use a single sample from the variational family to approximate the ELBO. To tie everything together and make use of the reparametrization trick, we need to rewrite each of the terms in equations 8 to 12 solely in terms of the sample parameters  $(s,\beta ,\theta)$  and the distributional parameters  $(\gamma ,\tilde{\gamma},\mu ,S,\tilde{\mu},\tilde{S},\zeta ,\Psi ,\tilde{\zeta},\tilde{\Psi})$ . This is done by replacing the sums over the leaves and internal nodes with indicator functions, which may be written as products over the node splitting indicator variables  $s = (s_{1},\dots,s_{2^{d - 1}})$ . If we denote the internal and leaf node indicator functions for node  $i$  by  $\psi_{i}$  and  $\kappa_{i}$  respectively, we may rewrite each of the terms in the ELBO as follows
$$
\begin{array}{l} \log p (\Theta , T) = \sum_ {i = 1} ^ {b - 1} \psi_ {i} \left(\log \left(\gamma_ {i}\right) + \log \mathcal {N} \left(\beta_ {i}; \mu_ {i}, S _ {i}\right)\right) (23) \\ + \sum_ {i = 1} ^ {2 b - 1} \kappa_ {i} \left(\log \left(1 - \gamma_ {i}\right) + \log \mathcal {N} \left(\theta_ {i}; \zeta_ {i}, \Psi_ {i}\right)\right) (24) \\ \end{array}
$$
$$
\begin{array}{l} \log q (\Theta , T) = \sum_ {i = 1} ^ {b - 1} \psi_ {i} \left(\log \left(\tilde {\gamma} _ {i}\right) + \log \mathcal {N} \left(\beta_ {i}; \mu_ {i}, S _ {i}\right)\right) (25) \\ + \sum_ {i = 1} ^ {b - 1} \kappa_ {i} \log (1 - \tilde {\gamma} _ {i}) + \sum_ {i = 1} ^ {2 b - 1} \kappa_ {i} \log \mathcal {N} \left(\theta_ {i}; \zeta_ {i}, \Psi_ {i}\right) (26) \\ \end{array}
$$
$$
\log p (Y \mid T, \Theta , X) = \log \left(\sum_ {i = 1} ^ {2 b - 1} \kappa_ {i} \prod_ {k = 1} ^ {\lfloor \log i \rfloor} \sigma \left((- 1) ^ {\mathrm {p a} _ {k - 1} (i)} \beta_ {\mathrm {p a} _ {k} (i)} ^ {T} x\right) f (y | \theta_ {i})\right) \tag {27}
$$
Where  $\psi_{i}$  and  $\kappa_{i}$  are given by
$$
\psi_ {i} = s _ {i} \prod_ {k = 1} ^ {\lfloor \log i \rfloor} s _ {\mathrm {p a} _ {k} (i)} \tag {28}
$$
$$
\kappa_ {i} = \left\{ \begin{array}{l} \left(1 - s _ {i}\right) \prod_ {k = 1} ^ {\lfloor \log i \rfloor} s _ {\mathrm {p a} _ {k} (i)}, i \in \{1, \dots , b - 1 \} \\ \prod_ {k = 1} ^ {\lfloor \log i \rfloor} s _ {\mathrm {p a} _ {k} (i)}, i \in \{b, \dots , 2 b - 1 \} \end{array} \right. \tag {29}
$$
3.2 Posterior Predictive Distribution and Bayesian Confidence Intervals
The fully Bayesian treatment of the problem enables us to obtain natural confidence intervals on our predictions, which is a significant advantage of the Bayesian approach compared to other black-box machine learning algorithms. Given a dataset  $\mathcal{D} = \{(x_i,y_i)\}_{i = 1}^N$ , confidence intervals for the predictions are obtained by approximating the predictive posterior distribution as follows:
$$
\begin{array}{l} p \left(y ^ {\prime} \mid \mathcal {D}, x ^ {\prime}\right) = \int p \left(y ^ {\prime}, T, \Theta \mid \mathcal {D}, x ^ {\prime}\right) d \Pi (T, \Theta) (30) \\ = \int p (T, \Theta | \mathcal {D}) p \left(y ^ {\prime} \mid T, \Theta , x ^ {\prime}\right) d \Pi (T, \Theta) (31) \\ \approx \int q (T, \Theta) p \left(y ^ {\prime} \mid T, \Theta , x ^ {\prime}\right) d \Pi (T, \Theta). (32) \\ \end{array}
$$
Equation 31 shows that the predictive posterior is a mixture of the posterior and the likelihood. Thus, to obtain samples from the predictive posterior distribution, we sample a tree with its parameters  $(\Theta, T)$  from the posterior and and then use  $(\Theta, T)$  to obtain a sample from the likelihood. Moreover, as we do not have access to the true posterior distribution, we use the learned variational family  $q(T, \Theta)$  as an approximation. For prediction purposes, we estimate the posterior predictive mean using the mixture representation of equation 32.
This section presents a comparison of the performance of Variational Regression Trees with three other widely used boosting algorithms: CatBoost, XGBoost, and RandomForest. A comparison with Bayesian Additive Regression Trees (BART, Chipman et al. [2010]) and a Multilayer Perceptron (MLP, LeCun et al. [2015]) is deferred to the Appendix. We conducted experiments on 18 distinct Datasets from the UCI Machine Learning Repository to benchmark our algorithm. In addition, we explored the use of VaRT for estimating of conditional treatment effects in causal inference problems, motivated by the fact that BART is a popular algorithm for such problems.
4.1 Performance on Regression Datasets
The results of our regression experiments are shown in Table 2. We compare the performance of VaRT to three popular Boosting Algorithms: CatBoost, XGBoost, and RandomForest. We compared an ensemble of fifty trees for each algorithm to single VaRT trees of depths 3,5,7, and 10. A comparison with BART and Multilayer Perceptrons is given in the Appendix.
In these experiments, VaRT was trained using gradient descent paired with a Clipped Adam optimizer in PyTorch (Paszke et al. [2019], Bingham et al. [2018]). A single VaRT tree for depths 3, 5, 7, and 10 were trained for each dataset. We report the average RMSE values for VaRT across all depths to provide a thorough representation of VaRT's performance across various configurations (Table A). The regularization parameters for all runs was set to  $10^{-3}$  and no hyperparameter tuning was performed to make a fair comparison on the off-the-shelf performance of each of the algorithms. In practice, both the regularization and truncation parameters could be chosen through cross-validation. Early stopping by evaluating the RMSE on a training or validation set every time a fixed number of epochs are completed should also increase performance.
Compared to other tree-based Bayesian algorithms, such as BART, VaRT offers distinct advantages that enhance its practicality. Notably, VaRT's framework is inherently parallelizable, a feature that sets it apart from MCMC-based methods like BART which are inherently sequential and lack this
Dataset	CatBoost	RandomForest	XGBoost	VaRT3	VaRT5	VaRT7	VaRT10
autompg	0.121	0.115	0.113	0.116	0.116	0.116	0.109
concreteslump	0.209	0.237	0.252	0.009	0.009	0.028	0.047
forest	0.396	0.403	0.403	0.392	0.394	0.387	0.387
solar	0.086	0.086	0.086	0.114	0.092	0.094	0.093
stock	0.074	0.067	0.072	0.065	0.068	0.065	0.059
yacht	0.036	0.026	0.018	0.041	0.035	0.035	0.046
airfoil	0.073	0.062	0.061	0.151	0.133	0.131	0.131
autos	0.092	0.072	0.076	0.121	0.146	0.134	0.140
breastcancer	0.403	0.362	0.398	0.390	0.476	0.380	0.360
concrete	0.095	0.093	0.081	0.149	0.141	0.158	0.148
fertility	0.402	0.375	0.358	0.481	0.454	0.456	0.387
housing	0.103	0.118	0.120	0.101	0.097	0.098	0.101
machine	0.089	0.112	0.097	0.110	0.123	0.124	0.138
pendulum	0.061	0.068	0.074	0.077	0.064	0.056	0.058
servo	0.066	0.062	0.050	0.096	0.077	0.078	0.109
skillcraft	0.118	0.116	0.120	0.114	0.114	0.114	0.114
sml	0.028	0.018	0.020	0.063	0.056	0.049	0.051
wine	0.108	0.103	0.096	0.122	0.119	0.117	0.121
Table A: Comparison of RMSE values for different boosting algorithms and VaRT. The RMSE values correspond to the average RMSE on the test sets of a random  $90 / 10\%$  train-test split of the data over ten runs. Bold values indicate the lowest RMSE. Each of the ensemble methods were trained using fifty trees. VaRT's performance is shown for trees of depth 3, 5, 7, and 10.
VaRT	CatBoost	RandomForest	XGBoost
tamielectric (45781 datapoints)	0.508	0.503	0.545	0.504
kin40k (40000 datapoints)	0.099	0.100	0.097	0.110
Table B: RMSE values for different boosting algorithms and VaRT on datasets containing more than 40,000 datapoints. The RMSE values correspond to the average RMSE on the test sets of a random  $90 / 10\%$  train-test split of the data over six runs. Bold values indicate the lowest RMSE. Each of the ensemble methods were trained using fifty trees, and a variational family of depth 5 was used for VaRT.
parallelization potential. This parallelizability not only accelerates the model training process but also empowers efficient exploration of large datasets, making VaRT particularly well-suited for modern computation demands. Furthermore, VaRT leverages gradient-based optimization techniques, making it compatible with widely used automatic differentiation engines such as PyTorch. This was validated by our experiments on datasets with over 40,000 datapoints, where training took less than 3 minutes.
We conducted all experiments on an ASUS Zephyrus G14 laptop with an RTX 2060 Max-Q GPU (6 GB VRAM), a 4900HS AMD CPU, 40 GB of RAM (although RAM usage was kept below 6 GB). Each experiment took roughly  $\sim$  1-3 minutes per dataset. Code with the random seeds needed to replicate the results of this section are provided with the supplementary material.
4.2 Causal Inference case-study