SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Camilo Ruiz $^{1,2,*}$ , Hongyu Ren $^{1,*}$ , Kexin Huang $^{1}$ , Jure Leskovec $^{1}$
$^{1}$ Department of Computer Science, Stanford University
$^{2}$ Department of Bioengineering, Stanford University
{caruiz, hyren, kexinh, jure}@cs.stanford.edu
Machine learning models exhibit strong performance on datasets with abundant labeled samples. However, for tabular datasets with extremely high  $d$ -dimensional features but limited  $n$  samples (i.e.  $d \gg n$ ), machine learning models struggle to achieve strong performance due to the risk of overfitting. Here, our key insight is that there is often abundant, auxiliary domain information describing input features which can be structured as a heterogeneous knowledge graph (KG). We propose PLATO, a method that achieves strong performance on tabular data with  $d \gg n$  by using an auxiliary KG describing input features to regularize a multilayer perceptron (MLP). In PLATO, each input feature corresponds to a node in the auxiliary KG. In the MLP's first layer, each input feature also corresponds to a weight vector. PLATO is based on the inductive bias that two input features corresponding to similar nodes in the auxiliary KG should have similar weight vectors in the MLP's first layer. PLATO captures this inductive bias by inferring the weight vector for each input feature from its corresponding node in the KG via a trainable message-passing function. Across  $6d \gg n$  datasets, PLATO outperforms 13 state-of-the-art baselines by up to  $10.19\%$ .
Machine learning models have reached state-of-the-art performance in domains with abundant labeled data like computer vision [76, 10] and natural language processing [70, 12, 52]. However, for tabular datasets in which the number  $d$  of features vastly exceeds the number  $n$  of samples, machine learning models struggle to achieve strong performance [24, 41]. Crucially, many tabular datasets from scientific domains [21, 30, 79, 17, 16, 33] have high-dimensional features but limited labeled samples due to the high time and labor costs of experiments. For these and other tabular datasets with  $d \gg n$ , the performance of machine learning models is currently limited.
The key challenge for machine learning models when  $d \gg n$  is the risk of overfitting. Indeed, deep models can have a large number of trainable weights, yet training is limited by the comparatively small number of labeled samples. As a result, tabular deep learning approaches so far have focused on data-rich regimes with far more samples than features ( $n \gg d$ ) [19, 18, 59]. In the low-data regime with far more features than samples ( $d \gg n$ ), the dominant approaches for single tabular datasets are still statistical methods [24]. These statistical methods reduce the dimensionality of the input space [1, 41, 66, 67], select features [64, 8, 14, 46], impose regularization penalties on parameter magnitudes [45], or use ensembles of weak tree-based models [15, 7, 34, 42, 51].
Here, we present a novel problem setting and framework that enables tabular deep learning when  $d \gg n$  (Figure 1). Our key insight is that there is often abundant, auxiliary domain information
(b) First layer of a MLP
Figure 1: PLATO is a method that uses auxiliary domain information describing input features to regularize a multilayer perceptron (MLP) and achieve strong performance on tabular data with  $d \gg n$ . (a) In PLATO, each input feature  $j$  corresponds to a node in an auxiliary KG of domain information. (b) In the first layer of a MLP with  $h$  hidden units, each input feature  $j$  corresponds to a vector of weights  $\Theta_{j}^{[1]} \in \mathbb{R}^{h}$  such that the weight vectors of all  $d$  features compose the weight matrix  $\Theta^{[1]} \in \mathbb{R}^{d \times h}$ . PLATO is based on the inductive bias that, if two input features  $j$  and  $k$  correspond to similar nodes in the auxiliary KG, they should have similar weight vectors  $\Theta_{j}^{[1]}$  and  $\Theta_{k}^{[1]}$  in the MLP. (c,d) PLATO captures this inductive bias by inferring the weight vector for each input feature  $j$  from its corresponding node in the KG. A trainable message-passing function  $\mathcal{Q}$  creates a low-dimensional embedding  $\mathbf{Q}_j \in \mathbb{R}^c$  for each input feature  $j$ . A neural network  $\mathcal{B}$  that is shared across all input features then infers the weight vector  $\Theta_{j}^{[1]}$  corresponding to input feature  $j$  from  $\mathbf{Q}_j$ . Input features with similar embeddings produce similar weight vectors, regularizing the MLP.
describing input features which can be structured as a heterogeneous knowledge graph (KG). We propose a novel problem setting in which each input feature of a tabular dataset corresponds to a node in an auxiliary KG (Figure 1a). To represent diverse domain information describing the input features, the KG contains feature and non-feature nodes as well as multiple node and edge types. For example, consider a tabular medical dataset in which each row is a cancer patient, each column is a gene, and each value is the amount of a gene in the patient's tumor. For this tabular dataset, there exists an auxiliary KG with each gene (i.e. input feature) as a node. Each gene node has edges to other gene nodes (i.e. other feature nodes) with diverse edge types like "activates" or "inhibits." Each gene node also has edges to other nodes (i.e. non-feature nodes) representing the gene's function in the body like "heart rate". Finally, the function nodes (i.e. non-feature nodes) have edges to each other representing their anatomical relationships like "heart rate"- "part of"- "cardiac system". Note that the KG does not capture the relationships between input data samples but instead captures the relationships between input features and other domain information.
Within our novel problem setting, we propose PLATO, a method that enables deep learning for tabular data with  $d \gg n$  by using an auxiliary KG describing input features (Figure 1). PLATO achieves strong performance by using the auxiliary KG to regularize a multilayer perceptron (MLP). In PLATO, each input feature corresponds to a node in the auxiliary KG (Figure 1a). In the first layer of the MLP, each input feature also corresponds to a weight vector such that the weight vectors of all features collectively compose the weight matrix (Figure 1b). PLATO is based on the inductive bias that two input features which correspond to similar nodes in the KG should have similar weight vectors in the first layer of the MLP. PLATO captures this inductive bias by inferring the weight vector for a feature from its corresponding node in the auxiliary KG with a trainable message-passing function (Figure 1c,d). Inferring the weights in the MLP's first layer also leads to a drastic reduction in the number of trainable weights, since most weights in a MLP are usually in the first layer when  $d \gg n$ .
We exhibit PLATO's performance on  $6d\gg n$  tabular datasets with 13 state-of-the-art baselines spanning dimensionality reduction, feature selection, statistical models, graph regularization, weight-inference, and tabular deep learning. Following a rigorous evaluation protocol from the tabular deep learning literature [19, 18], PLATO outperforms the prior state-of-the-art on all 6 datasets by up to  $10.19\%$ . Ablation studies demonstrate the importance of PLATO's trainable message-passing, the importance of non-feature nodes in the KG, and PLATO's robustness to missing edges in the KG. Ultimately, PLATO enables deep learning for tabular data with  $d\gg n$  by using an auxiliary KG describing the input features.
Tabular deep learning methods. In contrast to PLATO's setting, tabular deep learning methods have primarily been developed for settings with far more samples than features (i.e.  $n \gg d$ ). Indeed, recent tabular deep learning benchmarks ignore datasets with a large number of features and a small number of samples [19, 18, 59]. In the  $n \gg d$  setting, various categories of deep tabular models exist. First, decision tree models like NODE [50] make decision trees differentiable to enable gradient-based optimization [25, 36, 80]. Second, multilayer perceptrons (MLPs) apply sequential, non-linear transformations to input features [32, 31]. Third, tabular transformer architectures use an attention mechanism to select and learn interactions among features. Examples include TabNet [3], TabTransformer [29], FT-Transformer [18], TabPFN [26], SAINT [61], Non-Parametric Transformers [37], and AutoInt [62]. Finally, although PLATO focuses on single tabular datasets, transfer learning architectures can learn across multiple tabular datasets [39, 84, 73]. Ultimately, we compare PLATO to several benchmarked, state-of-the-art models for single, tabular datasets [19, 18, 59].
$d \gg n$  methods. For PLATO's setting in which  $d \gg n$ , various tabular machine learning approaches exist [24]. First, dimensionality reduction techniques like PCA [1] aim to reduce the dimensionality of the input data while preserving as much of the variance in the data as possible [41, 66, 67]. Second, feature selection approaches select a parsimonious set of features, leading to a smaller feature space. Feature selection approaches include LASSO [64] and its variants [8, 14, 46]. For feature selection with deep learning, Stochastic Gates [77] are among the best performing of many variants [4, 43]. Finally, tree-based models like XGBoost learn ensembles of weak decision trees models to make an overall prediction [15, 7, 34, 51].
Weight inference. Using one network to infer the weights of another has been studied extensively [11, 58, 6]. For example, [22] infers the weights in all layers of a sequential model (i.e. RNN, LSTM) by using information about the weights' structure. Diet Networks [54] infer weights by hand-crafting prior information about the input features or using random projections. By contrast, PLATO infers the weights in a MLP from prior information describing the input features in an auxiliary KG. PLATO's weight inference uniquely captures the inductive bias that two input features corresponding to similar nodes in a KG should have similar corresponding weight vectors in the first layer of a MLP (Figure 1).
Graph regularization. Graph regularization approaches regularize the weights of a linear model based on a simple graph between input features. The graph is typically constructed from the tabular data based on covariance relationships. Approaches then add a regularization penalty to the loss function which forces the weights of the linear model to vary smoothly over the corresponding feature nodes in the graph. State-of-the-art methods include GraphNet [20] and Network-Constrained LASSO [40] which are based on a Laplacian regularization [60, 2] as well as Network LASSO [23] which generalizes the Group LASSO [83] to a network setting. PLATO differs from graph regularization approaches in two key ways. First, PLATO's KG includes both feature and non-feature nodes and multiple edge types, thereby modeling diverse, prior domain information that is missing in graph regularization approaches. Second, PLATO infers the weights of a deep non-linear model (i.e. a MLP) rather than adding a regularization penalty to a loss, representing a distinct regularization mechanism.
Knowledge graph methods. Existing KG approaches are designed for tasks directly on the graph like link prediction or node classification [71, 65, 72, 78, 13]. By contrast, PLATO does not make any predictions on the KG. Instead, PLATO makes predictions on a separate, tabular dataset by using the KG as a prior. Graph classification methods also do not apply (Appendix B).
3 PLATO
PLATO is a machine learning method for tabular datasets with  $d \gg n$  and an auxiliary knowledge graph (KG) with input features as nodes (Section 3.1). PLATO's key insight is that there often exists abundant domain information describing input features which can be structured as an auxiliary KG  $G$  (Figure 1a). PLATO uses the auxiliary KG to regularize a multilayer perceptron (MLP) and achieve strong performance on tabular data when  $d \gg n$ .
Consider a tabular dataset  $\mathbf{X} \in \mathbb{R}^{n \times d}$  with labels  $\mathbf{y} \in \mathbb{R}^n$  and far more  $d$  features than  $n$  samples such that  $d \gg n$ . The goal is to train a model  $\mathcal{F}$  to predict labels  $\hat{\mathbf{y}}$  from the input  $\mathbf{X}$ . PLATO assumes the existence of an auxiliary knowledge graph  $G = (V, E)$  with  $|V|$  nodes and  $|E|$  edges such that each input feature  $j$  corresponds to a node in  $G$ . Formally,  $\forall j \in \{1, \dots, d\}, \exists v \in V$  s.t.  $j \mapsto v$ , as shown in Figure 1a.  $G$  also contains additional nodes which represent broader knowledge describing the domain. The edges in  $G$  are (head node, relation type, tail node) triplets.
3.2 PLATO's inductive bias
In PLATO, each input feature  $j$  corresponds to a node in the auxiliary KG (Figure 1a). In the first layer of a MLP with  $h$  hidden units, each input feature  $j$  also corresponds to a weight vector  $\Theta_{j}^{[1]} \in \mathbb{R}^{h}$  such that the weight vectors of all features collectively compose the weight matrix  $\Theta^{[1]} \in \mathbb{R}^{d \times h}$  (Figure 1b). PLATO is based on the inductive bias that two input features  $j$  and  $k$  which correspond to similar nodes in the KG should have similar weight vectors  $\Theta_{j}^{[1]}$  and  $\Theta_{k}^{[1]}$  in the first layer of the MLP. PLATO captures this inductive bias by inferring the weight vector for a feature from its corresponding node in the auxiliary KG with a trainable message-passing function (Figure 1c,d).
3.3 PLATO overview
PLATO has four key steps. First, PLATO uses a self-supervised objective on the auxiliary KG to pretrain an embedding for each input feature (Section 3.4). Second, PLATO updates each feature embedding with a trainable message-passing function that is trained on the supervised loss objective for the tabular data (Section 3.5, Figure 1c). Third, PLATO infers the weights in the first layer of the MLP from the feature embeddings with a small neural network that is shared across input features (Section 3.6, Figure 1d). Finally, the MLP predicts the label for the input sample.
3.4 Pretraining feature embeddings with self-supervision on the knowledge graph
First, PLATO learns general prior information about each input feature  $j$  from the auxiliary KG  $G$ . PLATO represents the general prior information about each input feature  $j$  as a low-dimensional embedding  $\mathbf{M}_j \in \mathbb{R}^c$ . Since each input feature  $j$  corresponds to a node in  $G$ , PLATO can learn  $\mathbf{M}_j$  by learning an embedding for the corresponding feature node in  $G$ . Any self-supervised node embedding method on  $G$  can be used within PLATO's framework.
Formal notation. Formally, PLATO uses self-supervision on  $G$  to pretrain an embedding for each input feature according to
$$
\mathbf {M} = \mathcal {H} (G). \tag {1}
$$
$\mathbf{M} \in \mathbb{R}^{d \times c}$  is the matrix of all feature embeddings.  $\mathcal{H}$  is a self-supervised node embedding method. We refer to Eq. (1) as pretraining since only the auxiliary KG  $G$  is used but the tabular data  $\mathbf{X}, \mathbf{y}$  is ignored. After pretraining, the feature embeddings  $\mathbf{M}$  are fixed.
For  $\mathcal{H}$ , we choose ComplEx as it is a prominent and highly scalable KG node embedding method [65]. ComplEx uses a self-supervised objective which learns an embedding for each node in  $G$  by classifying whether a proposed edge exists in  $G$ . ComplEx's proposed edges include both feature nodes and other nodes in  $G$ , thereby integrating prior information about the input features and the broader domain. We also test KG embedding methods DistMult [78] and TransE [72] in Appendix C.
3.5 Updating feature embeddings with a message-passing function trained on tabular data
PLATO next updates each feature embedding with a trainable message-passing function that is trained on the supervised loss for the tabular data (Figure 1c). During message-passing, PLATO updates the embedding of each input feature to be a weighted aggregation of its neighbors' embeddings.
Formal notation. Formally, PLATO uses a message-passing function  $\mathcal{Q}$  on the KG to update each pre-trained feature embedding  $\mathbf{M}_j \in \mathbb{R}^c$  to feature embedding  $\mathbf{Q}_j \in \mathbb{R}^c$  according to
$$
\mathbf {Q} = \mathcal {Q} (\mathbf {M}, G, \mathbf {X} _ {i}; \boldsymbol {\Pi}). \tag {2}
$$
As input, the message-passing function considers the pre-trained feature embeddings  $\mathbf{M}$ , the knowledge graph  $G$ , and the sample value  $\mathbf{X}_i$ .  $\mathcal{Q}$  uses an attention mechanism which considers the sample value  $\mathbf{X}_i$ . The only trainable weights in  $\mathcal{Q}$  are in the attention mechanism and are  $\Pi$ .
The message passing network  $\mathcal{Q}$ . Let  $\mathbf{Q}_j^{[r]}$  be the embedding of input feature  $j$  after round  $r\in \{1,\dots,R\}$  of message passing. For each input feature  $j$ ,  $\mathcal{Q}$  first initializes the updated feature embedding to the pretrained feature embedding.
$$
\mathbf {Q} _ {j} ^ {[ 0 ]} = \mathbf {M} _ {j}. \tag {2a}
$$
$\mathcal{Q}$  then conducts  $R$  rounds of message passing. In each round of message passing, the feature embedding  $\mathbf{Q}_j^{[r]}$  is updated from the feature embedding of each neighbor  $k$  in the prior round  $\mathbf{Q}_k^{[r-1]}$  and its own feature embedding in the prior round  $\mathbf{Q}_j^{[r-1]}$ . The "message" being passed is the embedding of each feature from the prior round.
$$
\mathbf {Q} _ {j} ^ {[ r ]} = \sigma \left[ \overbrace {\beta \left(\sum_ {k \in N _ {j}} \alpha_ {i j k} \mathbf {Q} _ {k} ^ {[ r - 1 ]}\right)} ^ {\text {W e i g h t e d m e a g e s f r o m n e i g h b o r s}} + \underbrace {(1 - \beta) \mathbf {Q} _ {j} ^ {[ r - 1 ]}} _ {\text {W e i g h t e d m e a g e f r o m s e l f}} \right]. \tag {2b}
$$
$\sigma$  is an optional nonlinearity.  $N_{j}$  are the neighbors of feature node  $j$  in  $G$
During message-passing,  $\mathcal{Q}$  uses two scalar values  $\beta \in \mathbb{R}$  and  $\alpha_{ijk} \in \mathbb{R}$  to control the weights of messages. First,  $\mathcal{Q}$  uses hyperparameter  $\beta \in \mathbb{R}$  to control the weight of the messages aggregated from the feature node's neighbors vs. from the feature node itself. Second,  $\mathcal{Q}$  calculates an attention coefficient  $\alpha_{ijk} \in \mathbb{R}$  to allow distinct nodes in the same neighborhood to have distinct weights. The coefficient  $\alpha_{ijk}$  specifies the weight of the message between feature  $j$  and neighbor  $k$  for sample  $i$ .
After  $R$  rounds of message-passing, the updated feature embeddings  $\mathbf{Q}_j$  are set.
$$
\mathbf {Q} _ {j} = \mathbf {Q} _ {j} ^ {[ R ]}. \tag {2c}
$$
The attention coefficient. PLATO's attention coefficient  $\alpha_{ijk}$  is inspired by [68] in which node attributes are used to calculate the weight of a message between neighboring nodes. For a sample  $i$  in PLATO, the node attributes for features  $j$  and  $k$  are their sample values  $\mathbf{X}_{ij} \in \mathbb{R}$  and  $\mathbf{X}_{ik} \in \mathbb{R}$ . PLATO thus uses the sample values  $\mathbf{X}_{ij}$  and  $\mathbf{X}_{ik}$  to calculate the attention coefficient. The attention coefficient  $e_{ijk}$  indicates the importance of node  $j$  to node  $k$  for sample  $i$ .
$$
e _ {i j k} = \mathcal {A} \left(\mathbf {X} _ {i j}, \mathbf {X} _ {i k}; \boldsymbol {\Pi}\right). \tag {2d}
$$
$\mathcal{A}$  is a shallow neural network parameterized by  $\Pi$  that is shared across samples and features. The number of trainable weights in  $\Pi$  is small since the input of  $\mathcal{A}$  is  $\mathbb{R}^2$  and the output of  $\mathcal{A}$  is a scalar  $\mathbb{R}$ .
To make the attention coefficients comparable across different nodes, PLATO normalizes the attention coefficients with a softmax function across the neighbors  $N_{j}$  of node  $j$ .
$$
\alpha_ {i j k} = \operatorname {s o f t m a x} _ {k} \left(e _ {i j k}\right) = \frac {\exp \left(e _ {i j k}\right)}{\sum_ {t \in N _ {j}} \exp \left(e _ {i j t}\right)}. \tag {2e}
$$
Algorithm 1: The PLATO Algorithm.
Input: A data sample  $\mathbf{X}_i\in \mathbb{R}^d$ , a knowledge graph  $G$  containing each input feature in  $\mathbf{X}$  as a node, a matrix of input feature embeddings  $\mathbf{M}\in \mathbb{R}^{d\times c}$  pre-trained over  $G$ .
Output: A predicted label  $\hat{\mathbf{y}}_i\in \mathbb{R}$
1 Use a trainable message-passing function  $\mathcal{Q}$  to update the pre-trained feature embeddings:  $\mathbf{Q} = \mathcal{Q}(\mathbf{M}, G, \mathbf{X}_i; \boldsymbol{\Pi})$ ,  $\mathbf{Q}_j \in \mathbb{R}^c$ ,  $\mathbf{Q} \in \mathbb{R}^{d \times c}$
2 Infer the weight vector in the first layer of a MLP that corresponds to an input feature  $j$  with a neural network  $\mathcal{B}$ :  $\hat{\Theta}_{j}^{[1]} = \mathcal{B}(\mathbf{Q}_{j}|\mathbf{X}_{i};\boldsymbol {\Phi}),\hat{\Theta}_{j}^{[1]}\in \mathbb{R}^{h}$
3 Repeat to infer the weight vectors corresponding to all input features by sharing the neural network  $\mathcal{B}$ :  $\hat{\Theta}^{[1]} \in \mathbb{R}^{d \times h}$
4 Concatenate the first layer inferred weights with the trainable weights in the rest of the MLP layers:
$$
\hat {\boldsymbol {\Theta}} = \{\hat {\boldsymbol {\Theta}} ^ {[ 1 ]} | \mathbf {X} _ {i} \} \cup \{\boldsymbol {\Theta} ^ {[ 2 ]}, \dots , \boldsymbol {\Theta} ^ {[ L ]} \}.
$$
5 Predict the label with a MLP  $\mathcal{F}$  that is parameterized by  $\hat{\Theta}$
$$
\hat {\mathbf {y}} _ {i} = \mathcal {F} \left(\mathbf {X} _ {i}; \hat {\boldsymbol {\Theta}} | \mathbf {X} _ {i}\right), \hat {\mathbf {y}} _ {i} \in \mathbb {R}
$$
Trainable weights:  $\Pi, \Phi, \Theta^{[2]}, \ldots, \Theta^{[L]}$ .
3.6 Inferring the first layer of weights in  $\mathcal{F}$  from the updated feature embeddings
Finally, PLATO infers the weights in the first layer of a MLP  $\mathcal{F}$  from the updated feature embeddings (Figure 1d). In the first layer of a MLP with  $h$  hidden units, each input feature  $j$  corresponds to a weight vector  $\Theta_{j}^{[1]} \in \mathbb{R}^{h}$  (Figure 1b). The weight matrix in the first layer of the MLP,  $\Theta^{[1]} \in \mathbb{R}^{d \times h}$ , is simply the concatenation of  $d$  weight vectors, one corresponding to each input feature. For each input feature  $j$ , PLATO infers the weight vector  $\hat{\Theta}_{j}^{[1]} \in \mathbb{R}^{h}$  from the feature embedding  $\mathbf{Q}_j \in \mathbb{R}^c$  by using a shallow neural network shared across input features. Input features with similar feature embeddings will produce similar weight vectors. Thus, PLATO captures the inductive bias that input features corresponding to similar nodes in the KG should have similar corresponding weight vectors in the MLP's first layer.
Formal notation. PLATO infers the weight vector associated with each input feature  $j$  in the first layer of  $\mathcal{F}$  with
$$
\hat {\boldsymbol {\Theta}} _ {j} ^ {[ 1 ]} = \mathcal {B} (\mathbf {Q} _ {j} | \mathbf {X} _ {i}; \boldsymbol {\Phi}). \tag {3}
$$
$\mathcal{B}$  is a shallow neural network with trainable weights  $\Phi$ .  $\mathbf{Q}_j$  is the updated feature embedding of  $j$  which is conditioned on the specific input sample  $\mathbf{X}_i$  since the input sample is used as an input in its calculation (Section 3.5, Equation 2).  $\Phi$  are the weights of  $\mathcal{B}$ .  $\mathcal{B}$  and its weights  $\Phi$  are shared across each feature  $j \in \{1, \dots, d\}$ .