Graph-Conditioned MLP for High-Dimensional Tabular Biomedical Data

Genome-wide studies leveraging recent high-throughput sequencing technologies collect high-dimensional data. However, they usually include small cohorts of patients, and the resulting tabular datasets suffer from the"curse of dimensionality". Training neural networks on such datasets is typically unstable, and the models overfit. One problem is that modern weight initialisation strategies make simplistic assumptions unsuitable for small-size datasets. We propose Graph-Conditioned MLP, a novel method to introduce priors on the parameters of an MLP. Instead of randomly initialising the first layer, we condition it directly on the training data. More specifically, we create a graph for each feature in the dataset (e.g., a gene), where each node represents a sample from the same dataset (e.g., a patient). We then use Graph Neural Networks (GNNs) to learn embeddings from these graphs and use the embeddings to initialise the MLP's parameters. Our approach opens the prospect of introducing additional biological knowledge when constructing the graphs. We present early results on 7 classification tasks from gene expression data and show that GC-MLP outperforms an MLP.


Introduction
Genome-wide studies leveraging recent high-throughput sequencing technologies aim to better understand human diseases and facilitate adequate treatments through personalised medicine. Such studies collect high-dimensional biomedical data but usually include small cohorts of patients. The resulting tabular biomedical datasets typically suffer from the "curse of dimensionality", as the number of features is much larger than the number of samples (D >> N ). Even though well-regularised MLPs surpass specialised neural networks on tabular datasets [7], they are unstable to train and overfit when the datasets are small-size and high-dimensional.
We seek to address two problems of MLPs applied to such datasets. Firstly, MLPs applied to high-dimensional data contain a disproportionate number of parameters in the first layer. For example, on a typical genomics dataset with ∼ 20,000 features, a five-layer MLP with 100 neurons per layer contains 98% of the parameters in the first layer, and can lead to overfitting. Secondly, training MLPs on smallsize datasets is cumbersome because modern initialisation strategies [4,5] put 'weak priors' on the model's parameters. For instance, they assume independence between a layer's weights and independence between a layer's activations. Such simplistic assumptions are suitable for large datasets because they allow flexibility but are unsuitable when there are not enough samples to estimate the model's parameters robustly. We hypothesise that more robust priors on the parameters would alleviate both issues.
We propose Graph-Conditioned MLP (GC-MLP), a novel method to introduce priors on the parameters of an MLP. Instead of initialising the first layer W [1] at random, we seek to condition it directly on the training data X according to W [1] |X. More specifically, we create a graph for each feature in the dataset (e.g., a gene), where each node represents a sample from the same dataset (e.g., a patient). In turn, we use Graph Neural Networks (GNNs) to learn embeddings from these graphs, which we then use to initialise W [1] .  Figure 1. Overview of proposed method GC-MLP. Given a tabular dataset X, we create a graph Gj for each feature j. From each graph Gj we extract graph embeddings w (j) and concatenate them horizontally. The results matrix W [1] is used to initialise the first layer of a standard MLP applied on the tabular dataset X.
One benefit of our approach is conditioning the MLP's parameters on the graphs, thus opening the prospect of introducing additional biological knowledge when constructing the graphs. We present early results on 7 classification tasks from gene expression data (in real applications, this could mean identifying biomarkers for different diseases) and show that GC-MLP outperforms an MLP. Lastly, we discuss extending this method to provide model explanations. Figure 1 illustrates the proposed method. We define the data matrix as X ∈ R N ×D and create D graphs from it. Specifically, for each feature j we use its values X :,j to construct an undirected graph G j = (V, E j ). The nodes in each graph V represent the data samples (in our experiments -patients) with node-features as a one-hot vector.

Method
To this end, in our experiments we use a simple heuristic to create edges: connecting nodes with similar feature values. Namely, nodes k and l are connected in the graph G j if (|X k,j − X l,j | < 0.05 * (max(X :,j ) − min(X :,j ))). We connect each node with at most 25 closest neighbours.
GC-MLP acts as an initialisation scheme for the first layer, and training this method includes two steps. Firstly, we compute graph embeddings for each graph using a Graph Autoencoder (GAE) [9] with a Graph Convolutional Network GCN [10] backbone. Let A j be the adjacency ma- trix for graph G j . For each graph G j , the GAE computes node embeddings Z j = GCN(I N , A j ) and then reconstructs the adjacency matrixÂ j = σ Z j Z j . The GAE's objective is to minimise the negative log-likelihood of reconstructing the adjacency matrix. We compute D graph embeddings w (j) by applying global average pooling on the node embeddings of each graph. The graph embeddings are concatenated horizontally to obtain the weight matrix where K is the number of neurons on the first layer. Secondly, the MLP is trained as usual to minimise the cross-entropy loss.
What is the inductive bias of this method? Intuitively, GC-MLP pre-initialises the features learned by neurons in the first layer of the MLP. The GAE extracts latent features from the graphs, such that W [1] ij represents 'how present' is the learned latent feature i in the graph G j . Consequently, we use the values of the k-th latent feature across all graphs to initialise the weights connecting the k-th neuron with all input features (i.e., W [1] k,: ). In the beginning, the first layer's neurons will ignore the features where some latent features are not present (i.e., the weight W [1] ij ≈ 0). Why is this method particularly suitable for highdimensional and small-size datasets? On the one hand, the problem's high-dimensionality leads to many graphs, which serve as a large training set for the GAE. On the other hand, the computation is efficient since the graphs are very small as the number of samples N is small.

Experiments
We evaluate the method on 7 classification tasks using realworld gene expression datasets. These datasets contain between 3312−19993 features, between 111−200 samples, and between 2 − 4 classes (details in Appendix A). For each dataset, we perform 5-fold cross-validation repeated 5 times, resulting in 25 runs per model. We randomly select 10% of the training data for validation for each fold. As a benchmark, we consider an MLP with hidden layers of size 100, 100, 10 and softmax activation. The MLP and GC-MLP use the same architecture. The graphs are computed from the training data only. We train the MLP and GC-MLP with weighted cross-entropy loss. Appendix B contains complete training details.
We consider two modes of training GC-MLPs: A two-step and an end-to-end mode. In the former, we first train the autoencoder independently from the MLP, and subsequently initialise the MLP with learnt graph embeddings. In the latter, motivated by previous work on parameter prediction [16,18], we considered an end-to-end version of our method, called GC-MLP (joint training) in Table 1. The weight matrix W [1] is computed using the GAE at every iteration. Accordingly, the first layer of the MLP acts as a 'virtual layer' because the weights W [1] are computed by the GAE rather than learned directly (which reduces the number of learnable parameters by ∼ 90%). This approach jointly trains the GAE and the MLP to minimise the cross-entropy loss using back-propagation.
Results Table 1 compares the two variants of GC-MLP against an MLP baseline. We observe that GC-MLP outperforms an MLP on 6 out of 7 datasets, except 'toxicity', on which it performs similarly. Also, on 5 out of 7 datasets GC-MLP has a smaller standard deviation, and we hypothesise that it could potentially alleviate the MLP's training instabilities. We find these preliminary results encouraging. However, our early results of the jointly trained variant show that this approach is less successful than the vanilla MLP.

Discussion
One of the key advantages of our approach is that using graphs opens the possibility to further condition the model using biological knowledge, for example, using proteinprotein interaction (PPI) networks to guide the construction of predictive models [17]. Although the two-step GC-MLP approach for initialising MLPs showed good empirical results, we see potential in further investigating joint training of GC-MLP. The end-to-end nature of the model could provide valuable interpretability insights. From a medical perspective, this approach could provide explanations using a cohort of similar patients (e.g., by finding suitable patient subgraphs) or concept-based explanations for GNNs [14,15]. Such model explanations could be helpful in the pursuit of discovering new biological insights. From a machine learning perspective, GC-MLP may provide new insights into the data, allowing outliers detection, detecting mislabeled samples, or finding the influential training points.