NodeFlow: Towards End-to-end Flexible Probabilistic Regression on Tabular Data

1. Introduction

Tabular regression involves predicting a continuous target variable based on structured data arranged in a tabular format. It is a vital task in machine learning with applications in various domains, including finance, healthcare, and marketing. In these domains, making reliable and informed decisions is of utmost importance due to potential consequences or impacts and requires not only accurate predictions but also robust uncertainty quantification. These kinds of properties can be obtained by the usage of probabilistic methods that go beyond point estimation by modeling the entire conditional distribution. This approach offers several advantages, including the ability to quantify uncertainty, capture complex data distributions, and provide a more comprehensive understanding of the data.

Regarding deterministic tabular regression, there have been two distinct paths of research in the field of regression on tabular data without any clear conclusion of the best approach to the problem [1,2]. The first path focuses on gradient-boosted trees, exemplified by popular frameworks such as XGBoost [3], CatBoost [4], and LightGBM [5]. These methods have demonstrated remarkable performance in point estimation tasks, leveraging ensemble techniques to capture complex relationships in the data. The second research path explores deep learning techniques for regression on tabular data with models such as NODE [6], TabNet [7], or FT-Transfomer [8]. These methods, with their ability to capture intricate patterns and relationships, have shown promise in surpassing the performance of gradient-boosted trees. They offer flexibility in handling various data types, including categorical variables, and can capture complex interactions among features. However, challenges specific to tabular data, such as feature interactions and interpretability, continue to be active research areas.

In the context of probabilistic tabular regression, recent research predominantly centers on expanding tree-based methods. The development of the new methods resulted in models such as NGBoost [9], PGBM [10], and a probabilistic extension of CatBoost [11]. However, these methods are predominantly based on parametric distributions, with CatBoost limited to modeling only Gaussian distributions. As a result, a pressing need remains for more flexible approaches that can accurately capture a broader range of complex data distributions encountered in practical scenarios. The recent work on TreeFlow [12] showed that combining tree-based methods with normalizing flows can improve the modeling capabilities; however, a lack of end-to-end optimization might lead to suboptimal results.

To overcome the limitations associated with the absence of end-to-end optimization, we propose NodeFlow, a novel framework for flexible probabilistic regression on tabular data. NodeFlow combines the advantages of tree-based structures, deep learning approaches, and normalizing flows to provide an accurate probabilistic regression approach that can be learned end-to-end. By combining Neural Oblivious Decision Ensembles (NODE) and Conditional Continuous Normalizing Flows (CNF), NodeFlow offers a unique solution that enables the modeling of complex data distributions encountered in probabilistic tasks. Through extensive evaluations and comparative studies on benchmark datasets, we demonstrate the effectiveness of NodeFlow in capturing the underlying data distributions and providing state-of-the-art results for multivariate probabilistic regression problems and competitive performance in univariate regression tasks.

Concluding, our contributions are as follows:

• Introduced NodeFlow, to the best of our knowledge, the first framework to apply an end-to-end, tree-structured deep learning model for probabilistic regression on tabular data;
• Demonstrated superior performance in multivariate probabilistic regression and competitive results in univariate tasks on benchmark datasets, establishing NodeFlow’s effectiveness;
• Conducted a focused ablation study, hyperparameter sensitivity analysis, and computational efficiency assessment, validating NodeFlow’s design and scalability.

2. Literature Review

2.1. Tree-based Regression on Tabular Data

Standard tree-based regression approaches, including XGBoost [3], CatBoost [4], and LightGBM [5], have emerged as state-of-the-art methods for modeling tabular data in regression problems. These frameworks leverage ensemble techniques and advanced optimizations to achieve remarkable performance in various domains. XGBoost is an optimized gradient-boosting framework that combines decision trees to capture complex relationships in tabular data. CatBoost incorporates novel techniques to handle categorical features effectively, while LightGBM utilizes tree-based learning algorithms and efficient data processing strategies. Their widespread adoption and success in diverse applications highlight their effectiveness and prominence in the field of tabular regression modeling, enabling accurate point estimation and capturing intricate patterns within the data.

2.2. Tree-based Probabilistic Regression on Tabular Data

In recent years, several approaches have been developed for probabilistic regression on tabular data, including NGBoost [9], CatBoost with univariate Gaussian support [11], and the Probabilistic Gradient Boosting Machine (PGBM) [10], each offering unique methods to model probabilistic distributions and improve regression performance. NGBoost is a versatile algorithm that can model various probabilistic distributions using a defined probability density function. It estimates distribution parameters by optimizing scoring rules such as negative log-likelihood (NLL) or Continuous Ranked Probability Score (CRPS). RoNGBa [13] is an NGBoost extension that enhances performance through improved hyperparameter selection. CatBoost, a gradient-boosting framework, has also been adapted to probabilistic regression but supports only univariate Gaussian distributions. PGBM treats leaf weights as random variables and can model different posterior distributions, albeit limited to location and scale parameters.

2.3. Deep Learning Regression on Tabular Data

In recent years, deep neural networks have achieved remarkable success in handling unstructured data, but their effectiveness in dealing with tabular data remains inconclusive. Several research papers, including [6,7,8,14,15], have introduced new deep learning regression methods that demonstrate superiority over tree-based methods. However, recent surveys have produced conflicting results on this topic. Notably, Borisov et al. [1] conducted a study comparing deep models to traditional machine learning methods on selected datasets. They found that deep models consistently outperformed traditional methods, but no single deep model universally outperformed all others. These findings highlight the nuanced performance of deep learning models on tabular data. Additionally, recent benchmarks conducted by Grinsztajn et al. [2] compared tree-based models and deep learning methods, specifically on tabular data. The benchmarks revealed that tree-based models such as XGBoost and Random Forests remain state-of-the-art for medium-sized datasets (with fewer than 10,000 samples). Notably, even without considering their superior processing speed, tree-based models maintained a competitive edge over deep learning approaches.

Neural Oblivious Decision Ensembles (NODE) introduced by [6] is a deep learning architecture that extends ensembles of oblivious decision trees. It combines end-to-end gradient-based optimization with multi-layer hierarchical representation learning. DNF-Net proposed by [7] is a neural architecture incorporating a disjunctive normal form (DNF) structure, allowing efficient and interpretable feature selection. It promotes localized decisions over small feature subsets, enhancing interpretability and mitigating overfitting. TabNet [14] is a deep learning architecture specifically tailored for tabular data. It processes raw tabular data without preprocessing, facilitating seamless integration into end-to-end learning. Sequential attention mechanisms identify crucial features at each decision step, enhancing interpretability and learning efficiency. TabNet also provides interpretable feature attributions and insights into the model’s global behavior. [8] proposed FT-Transformer, a modified version of the Transformer architecture designed for tabular data. FT-Transformer incorporates both categorical and continuous features, employs self-attention mechanisms to capture feature relationships, and integrates residual connections akin to ResNet. In addition to these approaches, SAINT (Self-Attention and Intersample Attention Transformer) [15] is a hybrid deep learning approach designed to solve tabular data problems. SAINT integrates attention over both rows and columns, an enhanced embedding method, and a contrastive self-supervised pre-training technique.

2.4. Deep Learning Probabilistic Regression on Tabular Data

Recently, there has been limited research on Probabilistic Deep Learning for Tabular Data. One notable method in this area is Deep Ensemble [16], which involves training an ensemble of neural networks using negative log-likelihood optimization with a Gaussian distribution as the modeling choice. The authors also incorporate adversarial training to produce smoother predictive estimates. Another approach, MC-Dropout [17], extends the use of dropout to capture model uncertainty during inference. By sampling multiple dropout masks during inference and averaging the predictions over these masks, an ensemble of models is created to capture model uncertainty collectively. Probabilistic Backpropagation [18] treats the neural network weights as random variables and approximates their posterior distribution using a factorized Gaussian distribution. This approximation is updated iteratively utilizing a combination of variational inference and stochastic gradient descent. More recently, TreeFlow [12] introduced a tree-based approach that combines the advantages of tree ensembles with the flexibility of modeling probability distributions using normalizing flows. By using a tree-based model as a feature extractor and combining it with a conditional variant of normalizing flow, TreeFlow enables the modeling of complex distributions in regression outputs. While TreeFlow has shown superior performance in some cases, its lack of end-to-end training may result in suboptimal results.

In conclusion, the existing methods for probabilistic regression on tabular data often have limitations in terms of their modeling flexibility or end-to-end training. NodeFlow addresses these limitations by combining the tree-based NODE with the flexibility of CNF, offering end-to-end training and a unique solution for probabilistic regression on tabular data.

3. NodeFlow

The architecture of NodeFlow is provided in Figure 1. The real-valued input vector $\mathbf{x}$ of dimensionality D is initially processed using Neural Oblivious Decision Ensemble, consisting of NODE Layers (details of the layer are depicted in Figure 2) arranged in a multi-layer hierarchical structure. It allows extraction of rich hierarchical representation $\mathbf{w}$. We use that vector as a conditioning factor for conditional Continuous Normalizing Flow (CNF) in the next step. This component is responsible for flexible modeling of the conditional probabilistic distribution of vector $\mathbf{y}$. It is worth mentioning that there are no restrictions on response vector dimensionality. Thus, we could cover both uni- and multi-variate regression problems. The whole architecture is trained in an end-to-end fashion using gradient-based optimization.

3.1. Extracting Hierarchical Representation with NODE

In order to extract rich hierarchical representation for a given input $\mathbf{x}$ we utilize Neural Oblivious Decision Ensemble (NODE) ${\mathbf{h}}_{\varphi }\left(\mathbf{x}\right)$ parametrized by $\varphi$ which is a machine learning architecture that combines differentiable oblivious decision trees $\mathbf{f}\left(\mathbf{x}\right)$ (ODTs). In this section, we start with introducing the ODTs. Then, we discuss the composition of the ODTs into the NODE Layer, and finally, we present the NODE component responsible for hierarchical representation extraction in NodeFlow.

The single differentiable oblivious decision tree $\mathbf{f}\left(\mathbf{x}\right)$ of the depth d is defined as:

$$\mathbf{f}\left(\mathbf{x}\right)=\sum _{j=1}^{2^d}{r}_j·l_j\left(\mathbf{x}\right),$$

(1)

where $\mathbf{r}=[r_1,\cdots ,r_{2^d}]$ is a $2^d$-dimensional vector of real-valued trainable responses for each of the considered leaves in the tree, and $\mathbf{l}\left(\mathbf{x}\right)=[l_1\left(\mathbf{x}\right),\cdots ,l_{2^d}\left(\mathbf{x}\right)]$ is a $2^d$-dimensional vector of real-valued entries from range $[0,1]$. The vector is called a "choice vector" and corresponds to the probability of the sample ending up in the specific leaf.

To compute the choice vector, it is requisite to perform a multiplication of the probabilities associated with selecting either the left or right path across successive depth levels within the tree structure. It’s important to note that in an oblivious decision tree, only one decision is made at each level of depth, which is referred to as $c_i\left(\mathbf{x}\right)$ at depth i. The final choice vector $\mathbf{l}$ is derived using the formula:

$$\mathbf{l}\left(\mathbf{x}\right)=\left[\begin{array}{c}{c}_1\left(\mathbf{x}\right)\\ 1-c_1\left(\mathbf{x}\right)\end{array}\right]\otimes \left[\begin{array}{c}{c}_2\left(\mathbf{x}\right)\\ 1-c_2\left(\mathbf{x}\right)\end{array}\right]\otimes \cdots \otimes \left[\begin{array}{c}{c}_d\left(\mathbf{x}\right)\\ 1-c_d\left(\mathbf{x}\right)\end{array}\right],$$

(2)

where ⊗ denotes the Kronecker product.

To ensure differentiability during training in the tree split, we utilize the $\alpha$-entmax function [19], which generalizes Softmax ($\alpha =1$) and Sparsemax ($\alpha =2$) functions and allows for learning sparse choices through gradient-based learning methods. The feature choice function $c_i\left(\mathbf{x}\right)$ is then calculated as a two-class entmax function over the transformed output of the feature selection function $k_i\left(\mathbf{x}\right)$. This can be expressed formally as:

$$c_i\left(\mathbf{x}\right)={\mathrm{entmax}}_{\alpha }\left([\frac{k_i\left(\mathbf{x}\right)-b_i}{{\tau }_i},0]\right)$$

(3)

where $b_i$ and ${\tau }_i$ are learnable threshold and scale parameters, and $\alpha$ is the entmax function’s hyperparameter that controls the level of "sparsity" in the output. In addition, the function for differentiable selecting features can be written as follows:

$$k_i\left(\mathbf{x}\right)=\sum _{j=1}^D{x}_j·p_j^{\left(i\right)},$$

(4)

where ${\mathbf{p}}^{\left(i\right)}$ is the D-dimensional vector of feature selection weights given by the formula ${\mathbf{p}}^{\left(i\right)}={\mathrm{entmax}}_{\alpha }\left(F_{i,·}\right)$. Moreover, $\mathbf{F}\in {\mathbb{R}}^{d\times D}$ is called the Feature Selection Matrix, and it is a real-valued, learnable matrix.

In summary, the differentiable oblivious decision tree, denoted as $\mathbf{f}$, is parameterized by the response vector $\mathbf{r}$, threshold values $\tau$, scale factors $\mathbf{b}$, and the feature selection matrix $\mathbf{F}$, facilitating gradient-based learning.

To form the Neural Oblivious Decision Ensemble Layer $F_l$ (depicted in Figure 2), we need to concatenate all outputs of the T individual ${\mathbf{f}}_1,\dots ,{\mathbf{f}}_T$ ODTs forming the layer. The final output can be written as

$$F_l(·)=[{\mathbf{f}}_1(·),\dots ,{\mathbf{f}}_T(·)].$$

(5)

Finally, the NODE architecture ${\mathbf{h}}_{\varphi }\left(\mathbf{x}\right)$ is composed of L stacked NODE Layers in a similar fashion to the DenseNet model. It means that each layer takes the concatenated outputs of all previous layers as input, allowing the model to learn both low-level and high-level features. It can be written as:

$${\mathbf{w}}_0=\mathbf{x};{\forall }_{l\in [1,L]}{\mathbf{w}}_l=[{\mathbf{F}}_l\left({\mathbf{w}}_{l-1}\right),{\mathbf{w}}_{l-1}].$$

(6)

The outputs from each layer are concatenated to create the final representation extracted using NODE, $\mathbf{w}=[{\mathbf{w}}_{\mathbf{1}},\dots ,{\mathbf{w}}_{\mathbf{L}}]={\mathbf{h}}_{\varphi }\left(\mathbf{x}\right)$. The representation $\mathbf{w}$ is further delivered to CNF as a conditioning factor.

3.2. Probabilistic Modeling with CNF

We consider the conditional variant of CNF provided in [20] and [21], where the conditional factor $\mathbf{w}={\mathbf{h}}_{\varphi }\left(\mathbf{x}\right)$ is delivered to the function of dynamics of $\mathbf{z}\left(t\right)$, ${\mathbf{g}}_{\beta }(\mathbf{z}\left(t\right),t,\mathbf{w})$, parametrized by $\beta$. In the CNF setting, we aim at finding a solution $\mathbf{y}:=\mathbf{z}\left(t_1\right)$ for the differential equation, assuming the given initial state $\mathbf{z}:=\mathbf{z}\left(t_0\right)$ with a known prior, where $\mathbf{z}$ is a random variable, $\mathbf{z}\left(t_0\right)$ is a base distribution and $\mathbf{z}\left(t_1\right)$ constitutes our observable data. Moreover, $t_0$ and $t_1$ denote the start and end points, respectively, of the continuous transformation process. The transformation function between $\mathbf{y}$ and $\mathbf{z}$ is represented as:

$$\mathbf{y}={\mathbf{u}}_{\beta ,\varphi }(\mathbf{z},\mathbf{x})=\mathbf{z}+{\int }_{t_0}^{t_1}{\mathbf{g}}_{\beta }(\mathbf{z}\left(t\right),t,{\mathbf{h}}_{\varphi }\left(\mathbf{x}\right))dt.$$

(7)

The inverse form of the transformation ${\mathbf{u}}_{\beta ,\varphi }(·)$ is given by equation:

$$\mathbf{z}={\mathbf{u}}_{\beta ,\varphi }^{-1}(\mathbf{y},\mathbf{x})=\mathbf{y}-{\int }_{t_0}^{t_1}{\mathbf{g}}_{\beta }(\mathbf{z}\left(t\right),t,{\mathbf{h}}_{\varphi }\left(\mathbf{x}\right))dt.$$

(8)

Finally, we can calculate the log-probability of target variable $\mathbf{y}$ given the vector of features $\mathbf{x}$ by the following formula:

$$logp\left(\mathbf{y}\right|\mathbf{x})=logp\left(\mathbf{z}\right)-{\int }_{t_0}^{t_1}\mathrm{Tr}\left(\frac{\partial {\mathbf{g}}_{\beta }(\mathbf{z}\left(t\right),t,{\mathbf{h}}_{\varphi }\left(\mathbf{x}\right))}{\partial \mathbf{z}\left(t\right)}\right)dt,$$

(9)

which can be solved analogously to FFJORD [22] by employing the adjoint method to backpropagate through the solution of the neural ODE.

3.3. Training NodeFlow

Using the formula (9) that directly defines log-probability, we can train NodeFlow directly optimizing the negative log-likelihood function. Let’s assume we are given a dataset $\mathcal{D}={({\mathbf{x}}_n,y_n)}_{n=1..N}$, where ${\mathbf{x}}_n=(x_n^1,\dots ,x_n^D)$ represents a D-dimensional random feature vector, and ${\mathbf{y}}_n=(y_n^1,\dots ,y_n^P)$ is the P-dimensional vector of targets. The training of the probabilistic model involves minimizing the conditional negative log-likelihood function (NLL), defined as:

$$Q(\mathit{\beta },\mathit{\varphi })=-\sum _{n=1}^{N}logp\left({\mathbf{y}}_n\right|{\mathbf{x}}_n,\mathit{\beta },\mathit{\varphi }).$$

(10)

The goal during the training process is to find the optimal parameters ${\beta }^{*}$ and ${\varphi }^{*}$ such that:

$${\mathit{\beta }}^{*},{\mathit{\varphi }}^{*}=\underset{\mathit{\beta },\mathit{\varphi }}{arg\; min}Q(\mathit{\beta },\mathit{\varphi }).$$

(11)

All model parameters $\mathit{\beta },\mathit{\varphi }$ are trained end-to-end by optimizing the above-mentioned NLL using the standard gradient-based approach. Such an approach simplifies the modeling process by allowing the entire model to be trained using a single optimization algorithm. Moreover, the model can automatically learn relevant hierarchical representations of the data directly from the raw input data, capturing both low-level and high-level features. This eliminates the need for manual feature engineering, which can be time-consuming and require domain expertise.

4. Experiments

In this section, we present a comprehensive set of experiments to evaluate the performance and effectiveness of NodeFlow in the context of tabular regression problems. We aim to assess NodeFlow’s capabilities in capturing complex data distributions, generating accurate point estimates, and quantifying uncertainty. To achieve this, we conduct evaluations on univariate and multivariate benchmark datasets, comparing NodeFlow with other reference methods. We measure the performance using various evaluation metrics such as negative log-likelihood (NLL), Continuous Ranked Probability Score (CRPS), and Root Mean Squared Error (RMSE). Through these experiments, we aim to demonstrate the performance and flexibility of NodeFlow in probabilistic regression tasks, contributing to the advancement of the field and providing insights for practical applications.

4.1. Methodology

In our evaluation, we adhere to the established probabilistic regression benchmark, as delineated in previous studies [9,11,12], excluding the Boston dataset in consideration of ethical concerns [23]. For univariate regression, we employ nine datasets from the UCI Machine Learning Repository and six datasets for multivariate regression as suggested by [12], with comprehensive dataset details provided in the Appendix. In alignment with protocols from referenced literature, we generate 20 random folds for the univariate regression datasets (with the exception of Protein at five folds and Year MSD at a single fold), designating $10\%$ of the data for testing in each fold. The remainder is divided into an $80\%$$/20\%$ training/validation split for epoch selection. Our results present the mean and standard deviation across validation folds. We benchmark NodeFlow against a suite of models, including four tree-based probabilistic models (NGBoost, RoNGBa, CatBoost, PGBM), a deep learning approach (Deep Ensemble), and a hybrid model (TreeFlow) for univariate tasks. For multivariate regression challenges, we adopt training/testing splits as per the referenced protocols, comparing NodeFlow against NGBoost variants and TreeFlow. The architecture specifics and hyperparameter tuning methodology for NodeFlow are detailed in the Appendix.

4.2. Probabilistic Regression Framework

This segment evaluates NodeFlow’s performance within a probabilistic framework, analyzing its negative log-likelihood (NLL) scores against benchmark datasets for both univariate and multivariate regression tasks previously outlined.

In Table 1, we present the evaluation results for the univariate regression task, where NodeFlow exhibits competitive performance across a range of datasets, frequently achieving the best or second-best NLL scores. Notably, NodeFlow excels on the Year MSD dataset and secures commendable second-best results on the Wine, Protein, Power, and Kin8nm datasets. Our analysis extends to a detailed comparison of NodeFlow against various methodological approaches, including deep learning-based methods, tree-based ensemble methods, and the hybrid method, TreeFlow. Against the Deep Ensemble, NodeFlow consistently demonstrates superior or at least equivalent performance, with particularly noteworthy achievements on the Energy, Power, Protein, Wine, and Yacht datasets. This is especially significant for the Protein and Wine datasets, which are characterized by their underlying multimodal target distributions—a scenario where NodeFlow’s capabilities of flexible distribution modeling are especially advantageous (refer to [12] for details). When compared to tree-based methods such as CatBoost, NGBoost, RoNGBa, and PGBM, NodeFlow maintains a competitive edge, often outperforming or matching the best results, underscoring its robust ability to model complex data relationships within tabular datasets. In direct comparison with TreeFlow, NodeFlow and TreeFlow exhibit closely matched performance, with each method surpassing the other under different circumstances. This comparative analysis not only highlights NodeFlow’s versatile efficacy across a broad spectrum of univariate regression challenges but also its capacity to address the intricacies of tabular data modeling through its advanced, adaptive learning framework.

In Table 2, we detail NodeFlow’s performance across multivariate probabilistic regression tasks, where it consistently outperforms competing approaches in five of six datasets examined. Compared with TreeFlow, NodeFlow’s superiority is particularly evident in datasets with multiple target dimensions, such as scm20d (16 target dimensions) and Energy (17 target dimensions). For two-dimensional target datasets like Parkinsons and US Flight, NodeFlow continues to outperform, albeit with a narrower margin. The distinction becomes more nuanced with one-dimensional targets, as presented in prior analyses, where NodeFlow and TreeFlow show competitive yet comparable results. This differentiation underscores the strength of NodeFlow’s end-to-end learning model, which excels in complex, high-dimensional settings by providing finely tuned representations. Such comprehensive learning is absent in TreeFlow, limiting its effectiveness in comparison. This evidence reinforces the indispensable value of end-to-end learning in achieving optimal performance, particularly in addressing the intricate demands of multivariate regression problems.

4.3. Point Prediction Regression Setup

This section will assess the effectiveness of our method in a point-prediction context by comparing its Root Mean Squared Error (RMSE) scores on the univariate regression datasets. To calculate the RMSE results for TreeFlow and NodeFlow methods, we used the RMSE@K metric introduced in [12], where $K=2$. This metric is suitable for uni- and multi-variate regression problems with multiple-point predictions. We have presented the results in Table 3. We can observe that NodeFlow consistently delivers competitive performance, showcasing comparable RMSE scores to other methods.

4.4. Summary

In summary, our evaluation of NodeFlow across both probabilistic and point-prediction scenarios demonstrates its efficacy. While NodeFlow’s performance on tasks with one-dimensional targets aligns with existing benchmarks, it distinctly excels in handling problems with two or more target dimensions. The results unequivocally indicate that the greater the dimensionality of the target variable, the more pronounced NodeFlow’s superiority becomes. This superior performance is attributed to NodeFlow’s flexible probabilistic modeling and comprehensive end-to-end learning approach, ensuring highly tailored representations for complex problems. Consequently, NodeFlow stands out as a superior method for probabilistic regression tasks involving high-dimensional targets, affirming its suitability for addressing advanced modeling challenges.

5. Ablation Studies

In the pursuit of a comprehensive understanding of NodeFlow method, a series of ablation studies were undertaken to scrutinize the impacts of critical design choices therein. Specifically, this investigation focused on two integral constituents: the Feature Representation component, in NodeFlow attained by the usage of NODE, and the Probabilistic Modeling segment, which was realized through the utilization of CNF. The probabilistic results are presented in the main manuscript, while point prediction results and qualitative analysis of the learned representations are analyzed in the Appendix. Moreover, the results of the time complexity analysis are also included in the Appendix.

5.1. Feature Representation Component

In our ablation study, we assessed the critical role of the Neural Oblivious Decision Ensembles (NODE) component in enhancing feature extraction within our proposed framework, NodeFlow. To this end, we conducted both quantitative and qualitative analyses, employing two benchmarking variants for comparison: one with the NODE component removed, relying solely on min-max scaling (termed as CNF), and another replacing NODE with a shallow Multilayer Perceptron (MLP), labeled as CNF + MLP.

Quantitative results, detailed in Table 4, evaluate the performance across probabilistic and point prediction metrics: Negative Log-Likelihood (NLL), Continuous Ranked Probability Score (CRPS), and Root Mean Squared Error at 2 (RMSE@2), presented as mean values alongside their standard deviations. The experimental setup was kept consistent with the main experiments.

Our findings reveal that NodeFlow, with the NODE component integrated, consistently delivers the lowest NLL values across a majority of datasets, highlighting its exceptional data modeling and prediction accuracy capabilities. Additionally, NodeFlow surpasses comparative approaches in CRPS, indicating its enhanced precision in probabilistic forecasting. Furthermore, NodeFlow achieves the most favorable RMSE scores, underlining the NODE component’s pivotal role in achieving precise point predictions.

In our qualitative analysis, we visualize feature representations derived from the models, utilizing dimensionality reduction via the UMAP algorithm [23] and color-coding each point according to its target variable. Figure 3 illustrates these representations for the Energy dataset. The leftmost visualization corresponds to the CNF model, which, lacking additional processing layers, essentially reflects the rescaled raw dataset within the $(-1,1)$ range. The middle image depicts the representation from the CNF + MLP model, while the rightmost image shows the outcome of employing NODE within the NodeFlow method. Comparatively, the NodeFlow method’s representation, facilitated by NODE processing, showcases significantly enhanced separation and disentanglement of observations, with distinct clusters forming around similar target values. This level of disentanglement, absent in the CNF models’ representations, likely plays a crucial role in NodeFlow’s superior performance across quantitative metrics.

Collectively, these outcomes validate the NODE component’s indispensable contribution to NodeFlow’s architecture, ensuring competitive or superior performance in NLL, CRPS, and RMSE metrics and disentangled and more clearly separated representations compared to the alternatives examined.

5.2. Probabilistic Modeling Component

In this ablation study, we evaluate the effectiveness and fit of the probabilistic modeling component within our framework. Specifically, we substituted the CNF component with standard probabilistic distributions, labeling these variants as NodeGauss (using a Gaussian distribution) and NodeGMM (employing a Mixture of Gaussians). This experimental design mirrors the setup of our previous ablation studies.

The findings, detailed in Table 5, indicate that NodeFlow consistently surpasses both NodeGauss and NodeGMM in Negative Log-Likelihood (NLL) across the majority of the datasets, with NodeGMM outperforming only in a single dataset instance. In terms of the Continuous Ranked Probability Score (CRPS), NodeFlow attains the lowest scores universally, indicating a more accurate calibration of predictive uncertainty relative to the alternatives. Point prediction results further underscore NodeFlow’s superiority as the most effective approach. Notably, these outcomes underscore the benefit of integrating a versatile probabilistic modeling component, as evidenced by the enhanced performance across all evaluated metrics.

Overall, NodeFlow’s uniform advantage across diverse metrics and datasets robustly validates the integral role of the CNF component in its architecture, underscoring its indispensability for achieving optimal model performance.

6. Time Complexity Analysis

In this analysis, we evaluate the training durations of NodeFlow relative to benchmark models from ablation studies, including CNF, CNF + MLP from the feature representation study, and NodeGauss and NodeGMM from the probabilistic modeling investigation. Our objective is to elucidate the computational demands of training each model across various datasets, as detailed in Table 6. The table delineates mean training times and their standard deviations, offering insights into both average performance and variability.

In the feature representation study, the marginal difference in training times between NodeFlow, CNF, and CNF + MLP suggests that the NODE component’s integration is cost-effective, enhancing model output without a corresponding surge in training duration. Conversely, the probabilistic modeling study indicates a more pronounced disparity in training times, particularly between NodeFlow and the NodeGauss and NodeGMM variants, with NodeFlow achieving superior results with a proportional increase in computational time.

Overall, NodeFlow presents itself as a robust solution for probabilistic regression tasks on tabular data, adeptly balancing efficiency in training time with excellence in performance. This equilibrium makes NodeFlow a compelling option for both academic research and practical implementation, highlighting its potential as a preferred method in the domain.

7. Conclusions

In this study, we introduced NodeFlow, a novel framework for probabilistic regression on tabular data, leveraging Neural Oblivious Decision Ensembles (NODE) and Conditional Continuous Normalizing Flows (CNF). Our evaluations confirm NodeFlow’s exceptional capability in managing high-dimensional multivariate probabilistic regression tasks, effectively aligning with benchmarks for tasks with one-dimensional targets. Ablation studies elucidate the critical roles of NODE and CNF in NodeFlow’s architecture, enhancing feature processing and complex distribution modeling, respectively. Moreover, NodeFlow emerges as a robust solution for advanced modeling and uncertainty quantification in regression tasks, adeptly balancing performance with computational efficiency. It not only establishes a significant presence in the domain of probabilistic regression but also lays a foundation for future advancements in machine learning interpretability and robustness. The differentiability of NodeFlow’s architecture is particularly conducive to further research in interpretability techniques, including counterfactual explanations, feature attribution, and adversarial example generation, promising substantial contributions to the field’s evolution.