Graph Neural Network contextual embedding for Deep Learning on Tabular
  Data
All industries are trying to leverage Artificial Intelligence (AI) based on
their existing big data which is available in so called tabular form, where
each record is composed of a number of heterogeneous continuous and categorical
columns also known as features. Deep Learning (DL) has constituted a major
breakthrough for AI in fields related to human skills like natural language
processing, but its applicability to tabular data has been more challenging.
More classical Machine Learning (ML) models like tree-based ensemble ones
usually perform better. This paper presents a novel DL model using Graph Neural
Network (GNN) more specifically Interaction Network (IN), for contextual
embedding and modelling interactions among tabular features. Its results
outperform those of a recently published survey with DL benchmark based on five
public datasets, also achieving competitive results when compared to
boosted-tree solutions.
Mario Villaizán-Vallelado, Matteo Salvatori, Belén Carro Martinez, Antonio Javier Sanchez Esguevillas
# Graph Neural Network Contextual Embedding for Deep Learning on Tabular Data

###### Abstract

All industries are trying to leverage Artificial Intelligence (AI) based on their existing big data which is available in so called tabular form, where each record is composed of a number of heterogeneous continuous and categorical columns also known as features. Deep Learning (DL) has constituted a major breakthrough for AI in fields related to human skills like natural language processing, but its applicability to tabular data has been more challenging. More classical Machine Learning (ML) models like tree-based ensemble ones usually perform better. In this manuscript a novel DL model that uses Graph Neural Network (GNN), more specifically Interaction Network (IN), for contextual embedding is introduced. Its results outperform those of the recently published survey with DL benchmark based on five public datasets, achieving also competitive results when compared to boosted-tree solutions.

 Deep Learning, Graph Neural Network, Interaction Network, Contextual Embedding, Tabular Data, Artificial Intelligence.

## I Introduction

Many practical real-world applications store data in tabular form, i.e. samples (rows) with the same set of attributes (columns). Medicine, finance or recommender systems are some common examples.

DL success in tasks involving texts, images or audio has sparked interest in its possible application to tabular data. Nevertheless, this success is often achieved when the input data are homogeneous and the structure used to organize the information provides insights about the data understanding. All tokens in a sentence are instances of the same categorical variable and their arrangement has semantic significance. Pixels in an image are continuous and usually have spatial correlation.

Tabular data have two characteristics that hinder DL performance. On the one hand, tabular features are heterogeneous, having a mix of continuous and categorical distributions that may correlate or be independent. On the other hand, the meaningfulness of tabular data row is independent of the column order, i.e. position is arbitrary and does not provide information.

Tree-based ensemble models such as XGBoost [1], CatBoost [2], and LightGBM [3] achieve the state of the art (SOTA) performances on tabular data: they have competitive prediction accuracy and are fast to train. Further research and development of DL models for tabular data is motivated, however, by the fact that standard tree-based approaches have limitations, for example, in case of continual learning, reinforcement learning or when tabular data is only part of the model input, which also includes data such as images, texts or audio.

Inspired by the success of contextual embedding in large language models (for example BERT [4]), several recent research [5, 6, 7] have investigated how to enhance tabular feature representation (and hence global DL model performances) by taking into consideration their context, that is, feature interaction. The results obtained in these works, as well as the outcomes of recent comparisons on many public datasets [8], illustrate how the contextual embedding approach tends to outperform not only standard Multi-Layer Perceptron (MLP) models, but also more complex models developed to solve complicated tasks [9, 10, 11, 12, 13] or models combining DL architectures with standard ML approaches [14, 15].

Many of the most recent studies employ Transformers [16] as a method for contextual embedding. However, in this paper, we look at how to use a GNN to improve contextual embedding for tabular data. GNNs are a special subset of neural networks that are capable of managing information organized in a graph which is a structure with variable shape or size and with complex topological relations. One of the most important features of a graph is that its meaning does not depend on the order of its nodes, just as the meaning of a tabular row does not depend on the order of its columns.

**Contributions.** The contributions of our paper are summarized as follows:

* original features and CLS virtual node
- and enhances their representation. The resulting CLS virtual node is sent into the final classifier/regressor. For sake of reproducibility, we share an implementation of INCE 2.
* We compare INCE against a wide range of deep tabular models and generally used tree-based approaches, using the tabular datasets provided in [8] as a benchmark. INCE outperforms all other DL methods on average, and it achieves competitive results when compared to boosted-tree solutions.
* We thoroughly investigate the differences between contextual embeddings based on Transformers and INs and analyze the influence of IN hyperparameters on model performance: quality of results, model size, computational time. Regardless of the dataset or task challenge, we gain a collection of patterns that aid in the establishment of a strong baseline.
* We investigate the interpretability of the IN ensuing contextual embeddings. On the one hand, we focus on the feature-feature relationship discovered by the IN, while on the other hand, we concentrate on how contextual embeddings improve traditional context-free embeddings.

## II Related Work

**Standard Tabular Models.** As already commented, when dealing with tabular data, tree-based ensemble models such as XGBoost, CatBoost and LightGBM are often a popular choice. They usually provide high performance regardless of the amount of data available, can handle many data types, are resilient in the case of null values, are fast to train and can be interpreted at least globally.

**Deep Tabular Models.** Due to the success of DL in task involving texts, sound or images, many efforts are being made to find the best approach to apply these models also to tabular data [5, 6, 7, 15, 20, 21]. Most of these efforts belong to one of the 3 categories described below.

_Modeling of multiplicative interactions between features_ Modeling explicitly the interaction between features of a tabular dataset [9, 10, 11, 12, 13] has been shown to have a significant impact on the performance of deep learning models in applications such as recommender systems and click-through-rate prediction. Nevertheless, recent comparisons [8, 6] show that these approaches produce worse outcomes than the rest of categories described below.

_Hybrid models._ Hybrid models transform the tabular data and combine deep neural networks with classical ML approaches, often decision trees. Such hybrid models can be designed to be optimized in a fully-differentiable end-to-end or to benefit from non-differentiable approaches combined with deep neural networks. NODE [14] is partially inspired by CatBoost [2] and provides an example of fully differentiable model based on an ensemble of oblivious decision trees [22]. Entmax transformation and soft splits allow to obtain a fully differentiable end-to-end optimization. Other examples of fully-differentiable hybrid architecture are [23, 24, 25]. On the other hand, DeepGBM model [26] is an example of how to take advantage from the combination of non-differentiable approaches with deep neural networks. It combines deep neural network flexibility with gradient boosting decision tree preprocessing capabilities. TabNN [27] first distills the knowledge from gradient boosting decision trees to retrieve feature groups and then constructs a neural network based on feature combinations produced by clustering the results of the previous step.

_Transformer-based models._ Many of DL recent successes have been driven by the use of transformer-based methods [4, 28, 29] inspiring the proposal of multiple approaches using deep attention mechanisms [16] for heterogeneous tabular data. The TabNet [15] design is inspired by decision trees: a set of subnetworks is processed in a hierarchical order and the results of all decision steps are aggregated in order to obtain the final prediction. A feature transformer module chooses which features should be transferred to the next decision step and which should be employed to get the output at the present decision phase. TabTransformer [5] uses Transformer to improve the contextual embeddings of tabular features. First, each categorical variable goes through a specific embedding layer. A stack of Transformers is then used to enhance the categorical feature representation. The final contextual embedding is given by the concatenation of the so obtained categorical representation and the initial continuous features. In FT-Transformer [6], columnar transformations (embeddings) are applied to both categorical and continuous features. As in BERT [4], a CLS token is added to the set of columnar embeddings and then, a stack of transformer layers is applied. The final CLS representation is employed as final contextual embedding, i. e. for predictions. SAINT [7] combines the self-attention between features of the same tabular row with inter-sample attention over multiple-rows. When handling missing or noisy data, this mechanism allows the model to borrow the corresponding information from similar samples.

As in [5, 6, 7], we investigate how contextual embedding affects the final model performance on supervised tasks. The main difference from the existing research is that in our approach, the contextual embedding is provided via GNNs and, more specifically, INs.

**Graph Neural Network and Interaction Network.** In case of neural networks such as Convolutional Neural Network or Transformer, the inputs must be structured data (grid and sequence, respectively). GNN are a special subset of neural networks that can cope with less structured data, such as a graph. This means that the input can have arbitrary shapes and sizes and can have complex topological relations. Permutation invariance is a crucial feature distinguishing GNN from the rest of neural networks. The order of nodes in a graph has no relevance, that is how we order the nodes in a graph does not impact the results produced by GNNs. In a tabular dataset, the order of features (columns) does not have any meaning, so GNN is a good candidate to model the interaction between them.

The flow of a GNN can be modeled using the Message-Passing scheme. a) For each pair of nodes \((u,v)\) in the graph, a message \(M(u,v,e_{u,v})\) from \(v\) to \(u\) is created. Here \(u\), \(v\) are the embedding of nodes and \(e_{u,v}\) is the (optional) embedding of edge. b) Each node aggregates the messages coming from all its neighbors. The aggregation must be permutation-invariant. c) The node is updated using its initial representation and the information obtained in point b.
It is simple to find a map between the Message-Passing scheme and the contextual embedding of tabular features. a) Initial node representation is given by columnar feature embeddings. b) Message-passing through edges is the pairwise interaction between features. c) The neighbor aggregation represents the effect of the interaction of current feature with all its neighbors. d) The update step provides the contextual representation of each feature.

In this paper, we investigate the benefits of using INs for contextual embeddings of tabular data. They are a low-biased family of GNN that have obtained enormous success when applied to simulation of complex physics or weather forecasting [30].

## III Interaction Network Contextual Embedding

This section introduces the INCE model and describes its components in depth.

**Problem Definition.** We focus on supervised learning problems with tabular datasets \(D=\left\{x_{i}^{cat},x_{i}^{num},y_{i}\right\}_{i=1}^{N}\) where \(x_{i}^{num}=\left\{x_{i}^{j_{n}}\right\}\) with \(j_{n}\in[1,M_{num}]\) is the set of numerical features, \(x_{i}^{cat}=\left\{x_{i}^{j_{c}}\right\}\) with \(j_{c}\in[1,M_{cat}]\) is the set of categorical features, \(y_{i}\) is the label, \(i\in[1,N]\) counts the dataset rows, \(N\) is the total number of rows and \(M=M_{num}+M_{cat}\) is total number of features.

**Encoder-Decoder Perspective.** As in [31], we use the encoder-decoder perspective, Fig. 1. First an encoder model maps each tabular dataset feature into a latent vector or embedding and then a decoder model takes the embeddings and uses them to solve the supervised learning task.

The encoder model is composed by two components: the _columnar_ and the _contextual_ embedding. The decoder model is given by a MLP tuned to the learning task to solve.

**Encoder - Columnar Embedding.** All of the original tabular heterogeneous features are projected in the same homogeneous and dense \(d\)-dimensional latent space by the _columnar_ embedding depicted in Fig. 2. As in the [6, 7], the columnar embedding \(c_{i}^{j_{n}},c_{i}^{j_{c}}\in\mathbb{R}^{d}\) of continuous and categorical features \(x_{i}^{j_{n}},x_{i}^{j_{c}}\) are obtained as follows:

\[c_{i}^{j_{n}} =\text{ReLU}\left(b_{num}^{j}+x_{i}^{j_{n}}\cdot W_{num}^{j} \right) W_{num}^{j}\in\mathbb{R}^{d} \tag{1}\] \[c_{i}^{j_{c}} =b_{cat}^{j}+e_{ij}^{j}W_{cat}^{j} W_{cat}^{j}\in\mathbb{R}^{|j_{c}|\times d} \tag{2}\]

where ReLU is the non-linear activation function for the continuous embedding, \(b^{j}\) is the \(j-\)th feature bias, \(W_{num}^{j}\in\mathbb{R}^{d}\) is a learnable vector, \(W_{cat}^{j}\in\mathbb{R}^{|j_{c}|\times d}\) is a learnable lookup table and \(|j_{c}|\) and \(e_{ij}^{T}\) are the size and the one-hot representation of the categorical feature \(x_{i}^{j_{c}}\), respectively.

**Encoder - Contextual Embedding.** The _columnar_ embedding works feature by feature and has trouble identifying correlation or more general relationships between features in tabular datasets. To overcome this limitation, a _contextual_ embedding is introduced. In contrast to recent research [5, 7, 6, 15] that use Transformers, we propose a _contextual_ embedding based on GNN and, more specifically, IN [17, 18, 19].

In this approach, the initial supervised learning task on tabular data is turned into a graph state estimation issue in which a categorical (classification task) or a continuous (regression task) graph state must be predicted. Taking into account the initial node representation (i. e. _columnar_ embedding) and graph edges, a stack of GNN has to model the interactions among nodes in the latent space and learn a richer representation of the entire graph capable of improving state estimation.

As shown in Fig. 3, the first step consists of building a fully-connected graph. For each original tabular feature, a node is created \(n_{j}\equiv x_{j}\) and for each pair of nodes \((n_{j_{1}},n_{j_{2}})\), two directed and independent edges are defined: \(e_{j_{1}j_{2}}:n_{j_{1}}\to n_{j_{2}}\) and \(e_{j_{2}j_{1}}:n_{j_{2}}\to n_{j_{1}}\). The dense \(d-\)dimensional vector \(c_{j}\in\mathbb{R}^{d}\) obtained from the _columnar_ embedding is used as initial node representation, giving rise to an homogeneous graph. No positional embedding is used to improve the node representation: the original tabular features are heterogeneous and each one is projected in the common

Fig. 1: The encoder-decoder perspective [31]: an encoder model maps each tabular dataset feature into a latent vector, a decoder model uses the embeddings to solve the supervised learning task. In the encoding step, first a _columnar_ embedding individually projects any feature in a common latent space and then a _contextual_ embedding improves these representations taking into account the relationships among features. The decoder MLP transforms the _contextual_ embedding output in the final model prediction.

Fig. 3: _Contextual_ embedding. (a) Homogeneous and fully-connected graph: it contains a node for each initial tabular features and a bidirectional-edge for each pair of nodes. The initial node representation is obtained by the _columnar_ embedding. A virtual CLS node is introduced to characterize the global graph state. (b) A stack of IN [17] models node interactions to create a more accurate representation of nodes (i.e. tabular features). (c) The final representation of the CLS virtual node is used as _contextual_ embedding.

Fig. 2: The _columnar_ embedding is responsible for projecting all the heterogeneous features in the tabular dataset in a common latent space. For each feature, a continuous or categorical transformation is defined. The _columnar_ embedding ignores any potential relationship or similarity between the tabular dataset features.
latent space using a separate _columnar_ embedding. This is enough to distinguish the nodes among them without explicitly modeling their position in the graph3. As in the BERT [4], a virtual CLS node connected to each existing node is added to the graph. The \(d-\)dimensional initial representation of the CLS virtual node is a vector of learnable parameters. No features are initially considered for the edges \(e_{ij}\).

Footnote 3: We have explicitly tested this hypothesis and the experiments confirm that the use of positional embedding does not improve the model performance.

In the following step, a stack of IN is used to improve the representation of each node and edge in the graph. The final CLS vector embedding produced by the stack of IN is used as global representation of the graph, i.e. as a _contextual_ embedding of the tabular row.

**Interaction Network.** The workflow of a standard IN layer [17, 18] is described in the Fig. 4. In the first step, the representation of each edge (i. e. interaction between each pair of tabular features) is updated using the information of the adjacent nodes (i. e. pair of tabular features):

\[e^{\prime}_{j_{1}\to j_{2}}=\text{MLP}_{\text{E}}\left(\text{ Concat}\left(n_{j_{1}},n_{j_{2}},e_{j_{1}\to j_{2}}\right)\right)\,, \tag{3}\]

where \(n_{j}\), \(e_{j_{1}j_{2}}\in\mathbb{R}^{d}\) are respectively node and edge representation, \(\text{MLP}_{\text{E}}\) is the shared neural network used to update all the graph edges. For simplicity of notation we have suppressed the row index.

In the second step, all the messages coming from the incoming edges are aggregated and used to update the node representation:

\[n^{\prime}_{j}=\text{MLP}_{\text{N}}\left(\text{Concat}\left(n_{j},\sum_{k\in \mathcal{N}}e_{k\to j}\right)\right)\,,\]

where \(\mathcal{N}\) is the set of \(n_{j}\) neighborhoods and \(\text{MLP}_{\text{N}}\) is the shared neural network used to update all the graph nodes.

The residual connection between the initial and updated representations yields the final node and edge representations:

\[n_{j} =n^{\prime}_{j}+n_{j}\] \[e_{j_{1}\to j_{2}} =e^{\prime}_{j_{1}\to j_{2}}+e_{j_{1}\to j_{2}} \tag{4}\]

**Decoder.** The decoder \(\text{MLP}_{\text{DEC}}\) receives the contextual embedding computed by the encoder. It is a MLP where the final output layer size and activation function are adapted to the supervised learning problem to solve - classification or regression.

## IV Experiments

[8] provides a detailed review on the literature of DL on tabular data together with an extensive empirical comparison of traditional ML methods and DL models on multiple real-world heterogeneous tabular datasets.

We consider the standard and deep models analyzed in [8] as baseline and evaluate INCE using the tabular benchmark presented therein.

**Data.** The main properties of datasets are summarized in Table I.

_HELOC_[32]: Home Equity Line of Credit (HELOC) provided by FICO (a data analytics company), contains anonymized credit applications of HELOC credit lines. The dataset contains 21 numerical and two categorical features characterizing the applicant to the HELOC credit line. The task is a binary classification and the goal is to predict whether the applicant will make timely payments over a two-year period.

_California Housing_[33]: The information refers to the houses located in a certain California district, as well as some basic statistics about them based on 1990 census data. This is a regression task, which requires to forecast the price of a property.

Fig. 4: Interaction Network layer
_Adult Incoming_[34]: Personal details such as age, gender or education level are used to predict whether an individual would earn more or less than \(50K\$\) per year.

_Forest Cover Type_[34]: Cartographic variables are used to predict the forest cover type: it is a multi-class (seven) classification task. The first eight features are continuous whereas the last two are categorical with four and 40 levels, respectively.

_HIGGS_[35]: The dataset contains 11M of rows and 28 features where the first 21 are kinematic properties measured by the particle detectors and the last seven are processed features built by physicists. The data has been produced using Monte Carlo simulations and the binary classification task is to distinguish between signals with Higgs bosons and a background process.

**Data Preprocessing.** In order to compare INCE with the results of [8], we reproduce the same data preprocessing. Zero-mean and unit-variance normalization is applied to the numerical features whereas an ordinal encoding is used for the categorical ones. The missing values were imputed with zeros.

**Baselines.** INCE is compared to the following models. _Standard methods_: Linear Model, KNN, Decision Tree, Random Forest [36], XGBoost [1], LightGBM [3], CatBoost [2]. _Deep learning models_: MLP [37], DeepFM [10], DeepGBM [26], RLN [38], TabNet [15], VIME [39], TabTrasformer [5], NODE [14], Net-DNF [25], SAINT [7], FT-Transformer [6].

**Setup.** For each tabular dataset, we use the Optuna library [40] with 50 iterations to tune INCE hyperparameters. Each hyperparameter configuration is cross-validated with five folds. The search space is the following: latent space size \(\in\{32,64,128\}\), number of stacked \(\text{IN}\in\{1,2,3,4\}\) and depth of \(\text{MLP}_{\text{E}}\), \(\text{MLP}_{\text{N}}\in\{1,2,3,4\}\).

In all the experiments, we consider a decoder \(\text{MLP}_{\text{DEC}}\) with two hidden layers and ReLU is the non-linear activation function used for \(\text{MLP}_{\text{E}}\), \(\text{MLP}_{\text{N}}\) and \(\text{MLP}_{\text{DEC}}\). Cross-Entropy and Mean Squared Error (MSE) are the loss functions used in classification and regression tasks, respectively. We train all the models \(200\) epochs using Adam optimizer with a learning rate of \(0.001\) and with batches of size \(256\). All the DL code is implemented using PyTorch [41] and PyTorch-Geometric [42] and parallelized with Ray [43].

### _Results_

In Table II and Fig. 5, we report the results on the tabular benchmark described above. In four of five datasets INCE outperforms all the DL baselines. In the fifth, HIGGS case, INCE obtains the second best performance behind SAINT model [7], but largely above the rest of DL models. In two of the five datasets, INCE outperforms tree-based models, while in the other three it achieves results that are competitive with them.

In terms of baseline performance, we carefully reproduced the findings for XGBoost, MLP, TabTransformer, and SAINT to ensure that our preprocessing and optimization approach was equivalent to [8] for all datasets in the benchmark. After demonstrating the comparability of [8] and our flows, the other baseline results are quoted from this paper. It should be noted that we include in our study the FT-Transformer that is subsequent to [8].

## V Deep Dive in Interaction Network

For each tabular dataset, we have studied how the choice of latent space size \(l\), \(\text{MLP}_{\text{N, E}}\) depth \(d\) and number \(n\) of stacked IN influences the model behavior: number of trainable parameters, performances and computational time. The findings from the various datasets reveal similar patterns, leading to consistent conclusions.

**Trainable parameters.** The number of trainable parameters \(\mathcal{TP}\left(\text{IN}\right)\) of a stack of \(n\) IN is given by:

\[\mathcal{TP}\left(\text{IN}\right) =\sum_{i=1}^{n}\mathcal{TP}\left(\text{IN}^{i}\right) \tag{5}\] \[=\sum_{i=1}^{n}\left[\mathcal{TP}\left(\text{MLP}_{\text{E}}^{i }\right)+\mathcal{TP}\left(\text{MLP}_{\text{N}}^{i}\right)\right]\] \[\mathcal{TP}\left(\text{MLP}_{\text{N}}^{i}\right) =\left(2\cdot l^{2}+l\right)+\left(d-1\right)\cdot\left(l^{2}+l\right)\] \[\mathcal{TP}\left(\text{MLP}_{\text{E}}^{i}\right) =\left(K_{i}\cdot l^{2}+l\right)+\left(d-1\right)\cdot\left(l^{2 }+l\right)\,,\]

where \(K_{i}=2\) if \(i=0\) and \(K_{i}=3\) otherwise. We consider all the hidden layers of \(\text{MLP}_{\text{E, N}}\) of the same size. The difference in the number of parameters between \(\text{MLP}_{\text{E}}^{i=0}\) and \(\text{MLP}_{\text{E}}^{i>0}\) is due to the fact that all IN with \(i>0\) receive the edge features computed by preceding layers, whilst the first IN does not use any initial edge features.

The quantity of trainable parameters increases quadratically with the size of the latent space and linearly with the number of stacked IN or the \(\text{MLP}_{\text{E, N}}\) depth, Fig. 6. The slope of the straight line corresponding to the number of stacked IN is steeper than the one relative to the \(\text{MLP}_{\text{E, N}}\) depth.

**Performances.** Our experiments suggest that whereas the latent space size needs to be fine-tuned for each dataset, the impact of \(\text{MLP}_{\text{E, N}}\) depth \(d\) and number \(n\) of stacked IN does not depend on the supervised learning problem to solve. The configuration with \(d=3\) and \(n=2\) is a solid baseline regardless of the underlying task.

To clarify this point, in Fig. 7 we show how the normalized metric changes as a function of the \(\text{MLP}_{\text{E, N}}\) depth and the

Fig. 5: The boxplots in red and blue illustrate the distribution of tree-based and DL baseline, respectively. The horizontal dotted line represents the INCE performance. Accuracy and MSE are the metrics used for classification and regression tasks. The presence of an up/down arrow near the dataset name indicates whether the metric must be maximized o minimized.
number of stacked IN. The normalized metric is a global performance measure (higher is better) generated using the findings from all of the datasets as described in Appendix A.

The left side plot in Fig. 7 depicts the normalized metric curves \(\mathcal{C}_{d}\) (blue line) and \(\mathcal{C}_{n}\) (orange line) obtained varying \(d\) and \(n\) respectively while the other parameters are kept constant. The information on the right side plot is the same as on the left, but it is compared to the normalized number of trainable parameters.

The depth \(d\) of the shared neural networks \(\text{MLP}_{\text{N, E}}\) has the most impact on the model performances and, at the same time, it has reduced effect on the number of learnable parameters. These results are coherent with the observed behavior of the Optuna [40] bayesian optimizer. Regardless of the supervised learning problem, after few attempts, it quickly reduces search space for \(d\) to \([3,4]\) and then it fine-tunes the number of stacked IN in the range \([2,3]\). The configuration with \(d=3\) and \(n=2\) is always a solid candidate regardless of the tabular dataset.

Why adding more than two layers does not improve the _contextual_ encoder capability? We interpret this as follows. a) The number of nodes in the graph is small. In our formulation there is a node for each tabular feature and the number of them goes from eight (California Housing) to 28 (HIGGS). After two IN layers, the information of a node has been transmitted to every other node in the graph. b) We are working with a fully connected graph, i. e. a trivial topology. The IN has to model the strength of each edge but the initial topological information seems to be poor. c) The size of datasets is limited (excludind HIGGS).

**Computational time.** Fig. 8 shows how the number of features in the tabular dataset as well as the INCE configuration (latent space size, number of stacked IN and \(\text{MLP}_{\text{N, E}}\) depth) impact on the training time. In particular, Fig. 8 presents the average training time for a batch size of \(256\). All the INCE training times are normalized by using the corresponding train time of a MLP with the same columnar embedding and the same decoder but without contextual embeddings. For each dataset, the three curves are obtained varying one parameter (for example \(n\in\{1,2,3,4\}\) for the orange line) while holding the other two constant (\(l=16\) and \(d=1\)).

* As expected, the number of features in the tabular dataset has an effect on the computational time: it grows from California Housing (eight features) to Heloc (\(23\) features) for a fixed INCE configuration. In our proposal, we are

Fig. 6: Growth of the normalized \(\mathcal{TP}\left(\text{IN}\right)\) as a function of \(\text{MLP}_{\text{R,N}}\) depth, number of stacked IN and latent space size. The plot on the left compares the evolution of \(\mathcal{TP}\left(\text{IN}\right)\) when two hyperparameters are fixed and the third is increased. The plot on right is a zoom on the contribution of \(\text{MLP}_{\text{E,N}}\) depth and number of stacked IN. The baseline used to normalize \(\mathcal{TP}\left(\text{IN}\right)\) is given by the number of trainable parameters of the simplest case: \(l=16\), \(d=1\), \(n=1\). It is trivial to show using Eq. 5 that the behavior of normalized \(\mathcal{TP}\left(\text{IN}\right)\) curve does not depend on the particular choice of the baseline latent space size \(l\).

Fig. 7: Average Normalized Metric. Left side plot depicts how the normalized metric changes when the \(\text{MLP}_{\text{R,N}}\) depth or the number of stacked IN is increased and the other is kept constant. The right side plot shows the same information but referenced to the normalized number of trainable parameters.
working with a fully-connected graph and the volume of operations increases quadratically with respect to the number of nodes (features).
* For a fixed dataset, the number \(n\) of stacked IN has the greatest impact on the amount of operations and, hence, on computational time.
* When the number of features is around \(20\), the impact of latent space size is comparable or even greater than the impact of \(\text{MLP}_{\text{N, E}}\) depth.

### _Interaction Network vs. Transformer_

Recent works [5, 7, 6] propose the Transformers encoder [16] as _contextual_ embedding. Here, we analyze similarities and differences between the two approaches.

**Performances.** For comparison purposes, we replace in our flow the IN with a Transformer encoder while keeping intact the rest of the model components - i.e. same _columnar_ embedding and same decoder. See [16] for the details about Transformer models and its components: Multi Head Self-Attention and FeedForward block. Using Optuna [40], we look for the best set of Transformer encoder hyperparameters in the following search space: number of attention heads \(h\in\{1,2,4,8\}\), FeedForward layer space size \(f\in\{512,1024,2048\}\), number of stacked Transformer encoders \(n\in\{1,2,3,4\}\) and latent space size \(l\in\{16,32,64,128\}\). As in the IN case, each hyperparameter configuration is cross-validated five folds and all the models are trained \(200\) epochs using Adam optimizer with a learning rate of \(0.001\) and batches of size \(256\).

Table III shows how the two approaches provide comparable results even though, at least on the selected benchmark, the IN encoder performs slightly better.

**Trainable parameters.** The size \(f\) of the latent space used by the Transformer FeedForward block has a significant impact on the number of trainable parameters in a Transformer encoder. In the comparison4 that follows, we take into account the setup where \(f=512\) since it achieves the best average results in the Optuna optimization.

Footnote 4: For the purpose of simplicity, we exclude the Normalization Layers parameters from our study in both cases, Transformer and IN.

The number of trainable parameters of a Transformer encoder is given by:

\[\begin{split}\mathcal{TP}\left(\text{Transformer}\right)& =n\cdot\left[\mathcal{TP}\left(Q,K,V\right)+\right.\\ &\qquad\qquad\qquad\qquad\qquad\qquad\left.\mathcal{TP}\left( \text{MultiAttention}\right)+\right.\\ &\qquad\qquad\qquad\left.\mathcal{TP}\left(\text{FeedForward} \right)\ \right]\\ \mathcal{TP}\left(Q,K,V\right)&=3\cdot h\cdot l \cdot\left(l+1\right)\\ \mathcal{TP}\left(\text{MultiAttention}\right)& =l\cdot\left(h+1\right)\\ \mathcal{TP}\left(\text{FeedForward}\right)&=2\cdot f \cdot l+f+l\\ \end{split} \tag{6}\]

where \(l\), \(h\), \(f\) and \(n\) are respectively the latent space size, the number of attention heads, the FeedForward latent space size and the number of stacked Transformer encoders.

Fig. 9 compares the behavior of \(\mathcal{TP}\left(\text{Transformer}\right)\) and \(\mathcal{TP}\left(\text{IN}\right)\). As in Fig. 6, the normalized number of trainable parameters \(\mathcal{TP}\) is obtained dividing by \(\mathcal{TP}\left(\text{IN}_{l,d=1,n=1}\right)\). The Fig. 9 presents the results for \(l=128\). IN has less trainable parameters than Transformers and the relative difference is even bigger when \(l\) decreases. When the number of attention heads is \(h\leq 2\), the difference is due to the FeedForward block parameters. For \(h>2\), Transformer has more parameters included, without taking into account the FeedForward block.

**Limitations.** When the number of tabular features increases, both IN and Transformers use greater resources. The vanilla Multi Head Self-Attention and IN on fully-connected graph share quadratic complexity with respect to the number of features. This issue can be alleviated by using efficient approximations of Multi Head Self-Attention [44] or a more complex graph topology with less edges in the Interaction Network case. Additionally, it is still possible to distill the final model into simpler architectures for better inference performance.

## VI Interpretability of contextual embedding

### _Columnar vs. Contextual embedding_

In subsection IV-A, the effect of _contextual_ embedding on the model performance has been shown. INCE outperforms solutions that just use _columnar_ embedding and, more generally, produces results that are on par with or even better than those of SOTA DL models when applied to tabular data.

Fig. 8: Average normalized training time. For each dataset the INCE training time is normalized using the time of the corresponding MLP with the same columnar embedding and decoder but without contextual embeddings. All the results are relative to a batch size of \(256\). Starting from the configuration base \(l=16\), \(n=1\) and \(d=1\), the different curves are computed varying one parameter while the others are kept constant.
In this subsection, we visually examine how this mechanism improves the features representation, enhancing the performance of the final model. For sake of simplicity, in the following discussion, we use the Titanic [34] dataset. The supervised learning problem is a binary classification. The preprocessed dataset contains eight features. Age and fare are the zero-mean and one-standard-deviation continuous variables. The categorical features are \(\text{sex}\in\{\text{female},\text{ male}\}\), \(\text{title}\in\{\text{Mr., Mrs., Rare}\}\), \(\text{pclass}\in\{1,\ 2,\ 3\}\), \(\text{family\_size}\in\{0,\ 1,\ 2,\ 3,\ 4,\ 6,\ 7,\ 8\}\), \(\text{is\_alone}\in\{0,\ 1\}\), \(\text{embark}\in\{\text{C}=\text{C}\text{Chebourg},\ \text{Q}=\text{Queenstown},\ \text{S}=\text{ Southampton}\}\). For this exercise, we consider a simple INCE model with latent space size \(l=2\), \(\text{MLP\_N,\ \text{E}}\) depth \(d=3\) and \(n=2\). The choice of \(l=2\) allows analyzing the representation in latent space without alleged artifacts introduced by the dimensional reduction.

The left side plot of Fig. 10 shows the output of _columnar_ embedding. Semantically related features like pclass-fare, sex-title, or family_size-is_alone are distributed without any discernible pattern. The representation does not depend on the context: regardless of pclass, age or family_size values, \(\text{title}=\text{Mrs}\) is always projected to the same point in the latent space size.

The _contextual_ embedding is depicted in the right side plot of Fig. 10. This represents the message sent from each node (i.e. tabular feature) to update the CLS representation in the last IN. Patterns are easily discernible: sex vs. title, family_size vs. is_alone and the feature that is closest to pclass is fare. Additionally, it is feasible to see that the latent projections of categorical features are not yet limited to a fixed number of points when the context is taken into consideration, as shown, for example, by title embedding.

### _Feature importance from Feature-Feature interaction_

The attention map for the CLS virtual node may be used to assess the feature relevance when the _contextual_ embedding is a Transformer [6, 7]. Here, we look into if the feature-feature interaction that the IN learns can reveal details about the significance of tabular features. We first explain our methodology using the Titanic dataset for the purpose of simplicity, and then we illustrate the findings we achieved using the same technique on the other tabular datasets.

In contrast to the Transformer case, we now have two new problems to resolve: 1) The feature-feature interaction is a \(l\)-dimensional vector (that is, it is not a scalar); 2) To assess the feature global significance, we must aggregate the feature-feature importance. The description of our process is provided below.

_First Step:_ We split data in train/test datasets. We train the model and use the trained INCE on the test dataset to produce the feature-feature interaction, i. e. \(e_{j_{1}\to j_{2}}^{i}\in\mathbb{R}^{l}\) in Eq. 4 returned by the last IN. In this notation, we have explicitly recovered the tabular row index \(i\).

_Second Step:_ We estimate mean \(\mu\) and covariance \(S\) of the entire population \(\{e_{j_{1}\to j_{2}}^{i}\ \forall\,i,j_{1},j_{2}\).

_Third Step:_ For each pair \((j_{1},j_{2})\) of features and for each test row \(i\), we compute the squared Mahalanobis distance:

\[D_{i}^{2}\left(j_{1}\to j_{2}\right)=\left(e_{j_{1}\to j_{2}}^{i}-\mu \right)S^{-1}\left(e_{j_{1}\to j_{2}}^{i}-\mu\right)\]

_Fourth Step:_ The squared Mahalanobis distance follows a Chi-Square distribution, so we can normalize the distance using p-value. The number of degrees of freedom of Chi-Square is given by the latent space size \(l\):

\[p_{i}\left(j_{1}\to j_{2}\right)=\text{Pr}\left(D_{i}^{2}\geq\chi_{l}^{2}\right)\]

_Fifth Step:_ The global interaction p-value \(p\left(j_{1}\to j_{2}\right)\) is obtained averaging the previous results over the test dataset:

\[p\left(j_{1}\to j_{2}\right)=\frac{1}{N_{test}}\sum_{i=1}^{N_{test}}p_{i} \left(j_{1}\to j_{2}\right)\]

The findings of the proposed methodology on the Titanic dataset are displayed in the heatmap of Fig. 11. The results are broken down at the feature-value level (i.e. \(\text{sex}=\text{female}\), \(\text{sex}=\text{male}\), \(\text{title}=\text{Mrs}\), \(\text{title}=\text{Mr}\), etc.). This is how the heatmap may be understood: the relevance of the message from the row-\(r\)-feature to the column-\(c\)-feature is represented by the element (row=\(r\), column=\(c\)) of the heatmap. A lower p-value implies more significance. The last column, "Mean", is created by averaging all of the row values and shows the average relevance of the messages sent by row-\(r\)-feature. In a

Fig. 10: **Left:**_Columnar_ embedding before the stack of IN. **Right:**_Contextual_ embedding from the last IN.

Fig. 9: Comparison of IN and Transformer trainable parameters. The normalized \(\mathcal{TP}\) is obtained using Eq. 5 and Eq. 6 and then normalizing with respect to \(\mathcal{TP}\left(IN_{l,d=1,n=1}\right)\). The plot shows the results for \(l=128\). Transformers have more trainable parameters than IN, and the relative difference grows when \(l\) decreases.
similar way, the last row (also known as "Mean") is derived by averaging all the values of the columns and it represents the mean relevance of the messages received by column-\(c\)-feature.

In order to quantitatively assess the quality of the heatmap, we compute the Spearman Rank correlation \(\rho\) between

\[p(j)=\frac{1}{|\mathcal{N}|}\sum_{j\in\mathcal{N}}p(j,\hat{j})=\frac{1}{2}\left[ p\left(j\to\hat{j}\right)+p\left(\hat{j}\to j\right)\right]\]

and the feature importance calculated by KernelShap [45]. In the formula above, \(\mathcal{N}\) and \(|\mathcal{N}|\) are the set of neighbors of node \(j\) and its size, respectively. The outcome for the Titanic dataset is \(\rho=0.81(\text{p-value}=0.05)\).

The heatmap and the Spearman Rank correlation provide the following insights.

**a)** The feature-feature interaction is not symmetric. In the fully connected graph we have two independent edges \(j_{1}\to j_{2}\) and \(j_{2}\to j_{1}\) and the Eq. 3 is not invariant by \(j_{1}\longleftrightarrow j_{2}\) interchange. Our experiments demonstrate that inducing \(j_{1}\longleftrightarrow j_{2}\) invariance in Eq. 3 results in a learning bias that negatively affects INCE performance.

**b)** From heatmap, it is possible to discern logical patterns. For example, is_alone \(=1\) does not add information (high p-value) when family_size is \(0\) or \(1\) and on the contrary, the value family_size is very relevant (low p-value) for any value of title.

**c)** Considering that KernelShap evaluates global model behavior (including the decoder) and that IN models separately \((j_{1}\to j_{2})\) and \((j_{2}\to j_{1})\) and that we have to aggregate and average them to compare with KernelShap, the Spearman Rank correlation analysis result can be considered encouraging.

Finally, Table IV summarizes the Spearman Rank correlation achieved on various datasets and demonstrates how the results are consistent regardless of the dataset under consideration.

## VII Conclusions

Let us highlight the main contributions of this article:

* To the best of our knowledge, this is the the first time that the model architecture proposes the use of GNN for contextual embedding to solve supervised tasks involving tabular data.
* Literature discusses mainly about the usage of Transformers. This manuscript shows that GNN, particularly IN, are a valid alternative. It shows better performance with a lower number of training parameters.
* As a matter of fact, this innovative architecture outperforms the state of the art DL benchmark based on 5 different diverse datasets. Moreover, it closes the gap with classical ML models (tree-based), outperforming them in 2 of these datasets, and being very close in two more. The tradeoff versus tree-based models is additional computational load in the form of training time, and scalability issues with the number of features (nodes) of the dataset, which constitute future lines of research to keep improving its practical implementation.
* Finally, the interpretability of GNN is explored. This is a key topic for industry environments, and apparently this is the first study for GNN and tabular data.

```
0:\(l\): latent space, \(r\): dataset
0:\(C_{d}\), \(C_{n}\) two list of normalized metric \(base\gets Metric_{r}(d=1,n=1,l,r)\) \(C_{d}\leftarrow Metric_{r}(d,n=1,l,r)\) \(\forall d\in\{1,2,3,4\}\) \(C_{n}\leftarrow Metric_{r}(d=1,n,l,r)\)\(\forall n\in\{1,2,3,4\}\) \(best\gets Best_{r}(C_{d},C_{n})\) \(C_{d}\leftarrow\dfrac{C_{d}-base}{best-base}\) \(C_{n}\leftarrow\dfrac{C_{n}-base}{best-base}\) \(\forall n\in\{1,2,3,4\}\) return\(C_{d}\), \(C_{n}\)
```

**Algorithm 1** Normalized metric

## Appendix A

Fig. 11: Titanic feature-feature interaction at feature-value level.
The curves of Fig. 7 are obtained by computing the average and the standard-deviation from results of Alg. 1.
