Towards Open-World Feature Extrapolation:
An Inductive Graph Learning Approach

Learning a mapping from observation $\mathbf{x}$ (a vector of attribute features) to label $y$ is a fundamental and pervasive problem in ML community, with extensive applications spanning from classification/regression tasks for tabular data to advertisement click prediction feature-ctr ; lr ; wd ; fm ; fwl , item recommendation feature-rec ; feature-rec2 ; youtube ; pinsage , question answering qa2 ; qa1 , or AI more broadly. Existing approaches focus on a fixed input feature space shared by training and test data. Nevertheless, practical ML systems interact with a dynamic open-world where features are incrementally collected. For example, in recommender/advertisement systems, there often occur new user profile features unseen before for current prediction tasks. Also, with the advances of multi-modal recognition mm-4 and federated learning mm-2 , it is a requirement for a model trained with partial features to incorporate new features from other fields for decision-making on a target task.

A challenge stems from the fact that off-the-shelf neural network models cannot deal with new features without re-training on new data. As shown in Fig. 1(a), a neural network builds a mapping from input features to a hidden representation through a weight matrix in the first layer. Given new features as input, the network needs to be augmented with new weight parameters and cannot map the new features to a desirable position in latent space if without re-training.

However, model re-training would be tricky and bring up several issues. First, re-training a model with both previous and new features would be highly time-consuming and cannot meet the requirement for online systems. Alternatively, one can re-train a trained model only on new features, which will induce risks for over-fitting new data or forgetting previous data.

Different from machines, fortunately, humans are often equipped with the ability for extrapolating to unseen features and distill the knowledge in new information for solving a target task without any re-training on new data. The inherent gap between existing ML approaches and human intelligence raises a research question: Can we design a ML model that is trained on one set of features and able to generalize to combine new unseen features for the same task without further training?

To find desirable solutions to this non-trivial question is challenging. Let us resort to how humans think and act in physical world scenarios. Imagine that we are well trained on predicting a person’s income with one’s age and occupation from historical records. Now we are asked to estimate the income of a new person with his/her age, occ. and education (a new feature). We can modularize the process in brain systems that distill and leverage the knowledge in new observations into four steps, as shown in Fig. 1(b). 1) Perception: the observations are recognized by our senses; 2) Abstraction: the perceived information is aligned with concepts in our cognition; 3) Reasoning: we search similar concepts and observations in our memory to understand and absorb the new knowledge; 4) Decision: with new understanding and abstraction, we make final decisions for the prediction.

The above human’s thinking process inspires the methodology in our paper where we propose a new learning paradigm for open-world feature extrapolation. Our proposed framework contains two modules: a backbone network, which can be a feedforward neural network, and a graph neural network. 1) The backbone network first maps input features to embeddings, which can be seen as perception from observation. 2) We then treat observed data matrix (each row represents a feature vector for one instance) as a feature-data bipartite graph, which explicitly define proximity and locality structures among features and instances. Then a graph neural network is harnessed for neural message passing over adjacent features and instances in a latent space (abstraction). 3) The GNN will inductively compute embeddings for new features based on those of existing ones, mimicking the reasoning process from familiar concepts to new ones in our brain. 4) The newly obtained embeddings that capture both semantics and feature-level relations will be used to obtain hidden representations of new data with unseen features and make final decisions.

To endow the model with ability for extrapolation to new features, we propose two training algorithms, one by self-supervised learning and one by inductive learning. The proposed model and its learning algorithm can easily scale to large-scale dataset (with millions of features and instances) via mini-batch training. We also provide a theoretical analysis on the generalization error on test data with new features, which allows us to dissect the impact of training features and algorithms on generalization performance. To verify our approach, we conduct extensive experiments on several real-world datasets, including six classification datasets from biology, engineering and social domains with diverse features, as well as two large-scale datasets for advertisement click prediction. Our model is trained on training data using partial features and tested on test data with a mixture of seen and unseen features. The results demonstrate that 1) our approach consistently outperform models not using new features for inference; 2) our approach achieves averagely $29.8\%$ higher Accuracy than baseline methods using KNN, pooling or local aggregation for feature extrapolation; 3) our approach even exceeds models using incremental training on new features, yielding average $4.1\%$ higher Accuracy.

Our contributions are: 1) We formulate open-world feature extrapolation problem and show that it is feasible to extend neural models for extrapolating to new features without re-training; 2) We propose a new graph-learning model and two training algorithms for feature extrapolation problem; 3) Our theoretical analysis shows that the generalization error for data with new features relies on the number of training features and the randomness in training algorithms; 4) We conduct comprehensive experiments and show the effectiveness, applicality and scalability of proposed method.

We focus on attribute features as input in this paper. An input instance is a vector $\mathbf{r}_{i}=[r_{im}]_{m=1}^{d}\in\mathbb{R}^{d}$ where each entry $r_{im}$ denotes a raw feature (like age, occ., edu., etc.). If $r_{im}$ is a discrete/categorical raw feature with $R_{m}$ possible values, its space is an integer set $\{0,1,\cdots,R_{m}-1\}$. If $r_{im^{\prime}}$ is a continuous one, a common practice is to convert it into a discrete feature within space $\{0,1,\cdots,R_{m^{\prime}}-1\}$ by evenly division ffm or log transformation feature-ctr . We call $R_{m}$ as cardinality for $m$-th raw feature.

An effective way to handle attribute features is via one-hot encoding wd ; ffm ; deepfm . For $r_{im}$ with cardinality $R_{m}$ we convert it into a $R_{m}$-dimensional one-hot vector $\mathbf{x}_{i}^{m}$ where the unique 1 indexes the value. In this way, one can convert an input $\mathbf{r}_{i}$ into a concatenation of one-hot vectors:

Use $x_{ij}$ to denote the $j$-th entry of $\mathbf{x}_{i}$ and we call each $x_{ij}$ as feature in this paper. Assume $D=\sum_{m=1}^{d}R_{m}$ denotes the number of features and, as a reminder, $d$ is the number of raw features.

We next give a formal definition for open-world feature extrapolation problem in this paper: given training data $\{(\mathbf{x}_{i},y_{i})\}_{i\in I_{tr}}$ where $\mathbf{x}_{i}\in\mathcal{X}_{tr}=\{0,1\}^{D}$, $y_{i}\in\mathcal{Y}$ and $I_{tr}$ is a set of indices, we aim to learn a model that can generalize to test data $\{\overline{\mathbf{x}}_{i^{\prime}},y_{i^{\prime}}\}_{i^{\prime}\in I_{te}}$ where $\overline{\mathbf{x}}_{i^{\prime}}\in\mathcal{X}_{te}=\{0,1\}^{\overline{D}}$, $y_{i^{\prime}}\in\mathcal{Y}$ and $I_{te}$ is another set of indices. We term $\mathcal{X}_{tr}$ as training feature space and $\mathcal{X}_{te}$ as test feature space. We assume 1) the label space $\mathcal{Y}$ is shared by training and test data, and 2) $\mathcal{X}_{tr}\subset\mathcal{X}_{te}$, i.e., test feature space is an extension of training feature space. The feature space expansion stems from two possible causes: 1) there appear new raw features incrementally collected from other fields (i.e., $d$ increases) or 2) there appear new values out of the known space of existing raw features (i.e., $R_{m}$ increases).

In this section, we analyze the generalization error on test data with new features. We simplify the settings for analysis: 1) the backbone network is a two-layer FNN (an embedding layer $\mathbf{W}\in\mathbb{R}^{D\times H}$ plus a fully-connected layer $\Phi\in\mathbb{R}^{H\times 1}$) with sigmoid output; 2) the GNN network is a $L$-layer GCN which takes mean pooling aggregation over neighbored nodes without linear transformation and non-linearity in each layer; 3) the training algorithm is SGD. With above settings, the model can be written as $\hat{y}_{i}=\sigma(\sum_{i^{\prime}\in\mathcal{N}_{\tilde{L}}(i)\cup\{i\}}c_{ii^{\prime}}^{L}\mathbf{x}_{i^{\prime}}\mathbf{W}\Phi)$ where $\mathcal{N}_{\tilde{L}}(i)$ contains all the $\mathbf{x}_{i^{\prime}}$’s that appear in the ${\tilde{L}}$-hop neighbors of $\mathbf{x}_{i}$ in the feature-data graph, $\tilde{L}=2\cdot\lfloor\frac{L}{2}\rfloor$, and $c_{ii^{\prime}}^{L}\in\mathbb{R}^{+}$ is a weight that quantifies influence of $\mathbf{x}_{i}$ on $\mathbf{x}_{i^{\prime}}$ through $L$-layer mean-pooling graph convolution. More details for the derivation are in Appendix C. Also, we focus our analysis on the case of inductive learning and the results can be extended to self-supervised approach, which we leave for future work.

The data generation process can be described as follows. First, features $f_{j}$’s are sampled from an unknown distribution $\mathcal{F}$ and form a feature set $F_{tr}=\{f_{j}\}_{j=1}^{D}$. Then data $\{(\mathbf{x}_{i},y_{i})\}_{i\in I_{tr}}$ are sampled from a distribution $\mathcal{D}_{F_{tr}}$ whose support is over $\mathcal{X}_{F_{tr}}\times\mathcal{Y}$, and define $\mathbf{X}_{tr}=\{\mathbf{x}_{i}\}_{i\in I_{tr}}$ and $Y_{tr}=\{y_{i}\}_{i\in I_{tr}}$. Using $\psi=[\theta,\omega]$, the model can be denoted by $\hat{Y}=h(\mathbf{X};\psi)$ (with a simplification from $\hat{Y}=\{\hat{y}_{i}\}_{i\in I}=\{h(\mathbf{x}_{i};\psi)\}_{i\in I}$) with loss function $\mathcal{L}(Y,\hat{Y})=\frac{1}{|I|}\sum_{i\in I}l(y_{i},\hat{y}_{i})$.

Recall that in training stage, we randomly partition the instances in $\mathbf{X}_{tr}$ into mini-batches with size $B$ and each mini-batch further samples $k$ raw features to form a feature subset $F_{s}$. In each update, the model is exposed to a $B\times|F_{s}|$ sub-matrix from $\mathbf{X}_{tr}$ as proxy training data and uses it for once feedforward and backward computation. With given training data $(\mathbf{X}_{tr},Y_{tr})$, we define $\mathcal{S}$ as a set of all the proxy data sub-matrices that could be exposed to the model during the training

The training process can be seen as a sequence of operations each of which samples a sub-matrix from $\mathcal{S}$ as proxy data in an i.i.d. manner and computes gradients for one SGD update (more discussions are in Appendix C). Define $A_{\mathcal{S}}$ as a learning algorithm trained on $\mathcal{S}$, which gives a trained model $h(\mathbf{X};\psi_{\mathcal{S}})$ simplified as $h_{\mathcal{S}}$. The generalization error $R(h_{\mathcal{S}})$ can be defined as

where the expectation contains two stages of sampling: 1) a feature set $F=\{f_{j}\}$ is sampled according to $f_{j}\sim\mathcal{F}$, and 2) data $(\mathbf{X},Y)$ is sampled according to $(\mathbf{x}_{i},y_{i})\sim\mathcal{D}_{F}$. The empirical risk that our approach optimizes with the training data would be

We study the expected generalization gap

where the expectation is taken over the randomness of $A_{\mathcal{S}}$ stemming from sampling for SGD updates.

We assume that the loss function $l(y_{i},\hat{y}_{i})$ is Lipschitz-continuous and smooth w.r.t. the model output $\hat{y}_{i}$. Concretely, we have 1) $|l(y,f(\cdot))-l(y,f^{\prime}(\cdot))|\leq\beta|f(\cdot)-f^{\prime}(\cdot)|$ and 2) $|\nabla l(y,f(\cdot))-\nabla l(y,f^{\prime}(\cdot))|\leq\beta^{\prime}|\nabla f(\cdot)-\nabla f^{\prime}(\cdot)|$. Such condition can be satisfied by widely used loss functions such as cross-entropy and MSE. Then we have the following result (see Appendix C for proof).

The generalization gap depends on the number of raw features in training data and the size of $\mathcal{S}$. The latter is determined by configuration of proxy training data, particular, sampling over training features. If the sampling introduces more randomness (e.g. $k\approx d/2$), $\mathcal{S}$ would become larger, contributing to tighter gap. However, a large $\mathcal{S}$ would also lead to large variance in training and amplify optimization error. Therefore, there exists a trade-off w.r.t. how to sample/split observed features in training stage. Furthermore, the generalization gap also depends on $d$, and a larger $d$ would result in looser bound (since one often has $T>d$). This is because a larger $d$ would require to deal with more features. As the network becomes wider and complex, it would be more prone for over-fitting.

We apply our model FATE (for FeATure Extrapolation Networks) on real-world datasets. First, we consider six classification datasets from UCI Machine Learning Repository uci-dataset : Gene, Protein, Robot, Drive, Calls and Github, as collected from domains like biology, engineering and social networks. The feature numbers are ranged from 219 to 4006 and the instance numbers vary from 1080 to 58509. We consider two large-scale datasets Avazu and Criteo from real-world online advertisement system whose goal is to predict the Click-Through Rate (CTR) of exposed advertisement to users. The two datasets have $\sim 40$ million clicking/non-clicking records as instances and $\sim 2$ million features. More dataset information and implementation details are in Appendix D and E, respectively.

Our introduced problem setting, open-world feature extrapolation (OFE), can be treated as an instantiation of out-of-distribution generalization ood-classic-2 ; ood-classic-3 or domain shift problem, focusing on distribution shift led by feature space expansion. We next discuss the relationships of our problem setting and our model FATE with domain adaption (DA), continual learning (CL), open-set learning (OSL) and zero-shot learning (ZSL). In general, OFE is orthogonal to these problems and opens a new direction that can potentially have promising intersections with the well-established ones.

Domain Adaption adapts a model trained on source domain to target with different distribution da-old1 ; da-old2 ; da-old3 ; da-nn1 ; da-nn2 ; da-nn3 ; da-nn4 . Our problem OFE is different from DA in two aspects: 1) the label space/distribution of training and test data is the same for OFE, while DA often mostly considers different label distributions for source and target domains; 2) OFE focus on combining new features that are related to the current task, while DA considers different tasks from different domains.

Continual Learning, or lifelong learning, aims at enabling a single model to learn from a stream of data from different tasks that cannot be seen at one time cl-old1 ; cl-1 ; cl-2 ; cl-survey . By contrast, there are two-fold differences of OFE. First, OFE does not allow finetuning or further retraining on new data, which can be more challenging than CL. Second, CL mostly assume each piece of data in the stream is from different tasks with different labels to handle. The key challenge of CL is the catastrophic forgetting cl-old1 ; cl-survey that requires the model to balance a trade-off between previous and new tasks, while FATE is free from such issue in nature since we do not require incremental learning.

Open-set Learning is another line of researches that relate with us. Differently, open-set recognition mostly focus on expansion of label sets openset-cv1 ; openset-cv2 ; openset-cv3 ; openset-cv4 ; openset-cv5 ; openset-nlp . To our knowledge, we are the first to study feature set expansion, formulate it as OFE and further solve it via graph learning.

Zero-shot Learning. Our problem setting is also linked with few-shot/zero-shot learning. In NLP domains, some studies focus on dealing with rare entities exposed in limited times or new entities unseen by training nlp-1 ; nlp-3 . Similar problems are also encoutered and explored in cold-start recommender systems where there are also new users/items unseen before recsys-1 ; recsys-2 ; recsys-3 . One common nature of these works aim at inferring the embeddings for new entities based on some ‘held-out’ ones. With a similar end and distinct methodological aspect, a recent study recsys-3 explores a new possibility via learning a latent graph between existing entities (users) and newly arrived ones through attention mechanism for inductive representation learning. The core technical contributions of our work lie in the unique problem setting which stays focused on feature space expansion (domain shift) and the proposed feature-level sampling/partition training strategy, backed up with our theoretical insights.

Extension and Outlooks. Our work can be extended to solve more problems and push the development in broader areas. First, the input feature vectors can be replaced by feature maps given by CNN or word/sentence embeddings by Transformer transformer , for handling multi-modal data in the context of federated learning mm-2 or multi-view learning mm-4 ; mm-6 ; mm-7 . Admittedly, our formulation assumes multi-hot feature vectors as input, which is a common practice for handling attribute features but not often the case for other data format (like vision or texts). For practitioners who would like to apply FATE to broader areas, one can extend our graph representation and GNN model to directly handle continuous features by treating feature values as edge weights as is done by graphlearning-missing . Second, as shown in our experiments, the classification layers can be replaced by more complex models with inter-feature interactions ffm ; xdeepfm ; nfm ; fwl or advanced architectures danser ; featevolve for more sufficient feature-wise interaction and improving the expressiveness.

We present a new framework to address new features unseen in training, by formulating it as the open-world feature extrapolation problem. We target the problem via graph representation learning by treating observed data as a feature-data graph and further harness GNN to inductively compute embeddings for new features with those of existing ones, mimicking abstraction and reasoning process in human’s brain. We also propose two training strategies for effective feature extrapolation learning. Our theoretical results show that generalization error depends on training features and learning algorithms. Experiments verify its effectiveness and scalability to large-scale systems.

Potential Societal Impacts. When learning mapping from features to labels, the model is at risk of focusing on dominant features from majority groups and ignoring scarce features from minority ones. Potential extended works of much significance are to develop debiased methods for feature extrapolation. We believe AI models can be guided to promote social justice and well-being.