A UTO T RANSFER : AutoML with Knowledge Transfer - An Application to Graph Neural Networks
Abstract
AutoML has demonstrated remarkable success in finding an effective neural architecture for a given machine learning task defined by a specific dataset and an evaluation metric.
However, most present AutoML techniques consider each task independently from scratch, which requires exploring many architectures, leading to high computational cost.
Here we propose A UTO T RANSFER , an AutoML solution that improves search efficiency by transferring the prior architectural design knowledge to the novel task of interest.
Our key insights are a task-model bank that captures the model performance over a diverse set of GNN architectures and tasks, and a computationally efficient task embedding that can accurately measure the similarity between different tasks.
Based on the task-model bank and the task embeddings, we estimate the design priors of desirable models of the novel task, by aggregating a similarity-weighted sum of the top-K design distributions on tasks that are similar to the task of interest.
The computed design priors can be used with any AutoML search algorithm.
We evaluate A UTO T RANSFER on six datasets in the graph machine learning domain.
Experiments demonstrate that (i) our proposed task embedding can be computed efficiently, and that tasks with similar embeddings have similar best-performing architectures; (ii) A UTO T RANSFER significantly improves search efficiency with the transferred design priors, reducing the number of explored architectures by an order of magnitude .
Finally, we release
GNN-B ANK -101, a large-scale dataset of detailed GNN training information of 120,000 task-model combinations to facilitate and inspire future research.
Introduction
Deep neural networks are highly modular, requiring many design decisions to be made regarding network architecture and hyperparameters.
These design decisions form a search space that is nonconvex and costly even for experts to optimize over, especially so when the optimization must be repeated from scratch for each new use case.
Automated machine learning (AutoML) is an active research area that aims to reduce the human effort required for architecture design that usually covers hyperparameter optimization and neural architecture search.
AutoML has demonstrated success [44, 26, 45, 6, 15, 11] in many application domains.
Finding a reasonably good model for a new learning tasks 1 in a computationally efficient manner is crucial for making deep learning accessible to domain experts with diverse backgrounds.
Efficient
AutoML is especially important in the graph learning domain for two reasons.
First, graph learning methods receive input data composed of a variety of data types and optimize over tasks that span
Task- records)
Task Embedding Space
Design Distributions
Task
Agg
Dim
Epoch ...
Val_loss
Cora sum  ...
Cora mean  ...
TU-DD sum  ...
TU-DD mean  ...
TU-DD max  ...
Arxiv mean  ...
Agg
Dim
Epoch
Agg
Novel Task of interest
Dim​
Epoch ... ​ ...
​ an equally diverse set of domains and modalities such as recommendation [37, 14], physical simulation [30, 25], and bioinformatics [43].
This differs from computer vision and natural language processing where the input data has a predefined, fixed structure that can be shared across different neural architectures.
Second, neural networks that operate on graphs come with a rich set of design choices and a large set of parameters to explore.
However, unlike other domains where a few pretrained architectures such as ResNet [13] and GPT-3 [5] dominate the benchmarks, it has been shown that the best graph neural network (GNN) design is highly task-dependent [38].
Although AutoML is promising, directly applying existing AutoML solutions to each graph learning task is infeasible due to its massive computational overhead.
Most present AutoML techniques consider each task independently and in isolation.
This leads to a high computational cost as many architectures have to be examined before a good one is found.
Here we propose A UTO T RANSFER , an AutoML solution that drastically improves AutoML architecture search by transferring previous architectural design knowledge to the task of interest.
Our key innovation is to introduce a task-model bank that stores the performance of a diverse set of
GNN architectures and tasks to guide the search algorithm.
To enable knowledge transfer, we define a task embedding space such that tasks close in the embedding space have similar corresponding top-performing architectures.
The challenge here is that the task embedding needs to capture the performance rankings of different architectures on different datasets, while being efficient to compute.
Our innovation here is to embed a task by using the condition number of its Fisher Information Matrix of various randomly initialized models and also a learning scheme with generalization guarantee.
This way we implicitly capture the properties of the dataset as well as the learning task, while being order of magnitudes faster (within seconds).
We then estimate the design prior of desirable models for the new task, by aggregating design distributions on tasks that are close to the task of interest.
Finally, we initiate a hyperparameter search algorithm with the task-informed design prior computed.
We evaluate A UTO T RANSFER on six datasets, including both node classification and graph classification tasks.
We show that our proposed task embeddings can be computed efficiently, and that the distance measured between tasks correlates highly with model performance rankings.
Furthermore, we present that A UTO T RANSFER significantly improves search efficiency when using the transferred design prior.
A UTO T RANSFER reduces the number of explored architectures needed to reach a target accuracy by an order of magnitude compared to SOTA baselines.
Finally, we release GNN-B ANK101—the first large-scale database containing detailed performance records for 120,000 task-model combinations which were trained using 16,128 GPU hours—as a tool to facilitate future research.
Related Work
In this section, we summarize the related work on AutoML regarding its applications on GNNs, the common search algorithms, and pioneering work regarding transfer learning and task embeddings.
AutoML for GNNs.
Neural architecture search (NAS), a unique and popular form of AutoML for deep learning, can be divided into two categories: multi-trial NAS and one-shot NAS.
During multi-trial NAS, each sampled architecture is trained separately.
GraphNAS [10] and Auto-GNN [42] are typical multi-trial NAS algorithms on GNNs which adopt an RNN controller that learns to suggest better sets of configurations through reinforcement learning.
One-shot NAS ( e.g. , [23, 27, 22]) involves encapsulating the entire model space in one super-model, training the super-model once, and then iteratively sampling sub-models from the super-model to find the best one.
In addition, there is work that explicitly studies fine-grained design choices such as data augmentation [39], message passing layer type [7, 9, 41], and graph pooling [33].
A UTO T RANSFER is the first AutoML solution for GNNs that efficiently transfer design knowledge across tasks.
HPO Algorithms.
Hyperparameter Optimization (HPO) algorithms search for the optimal model hyperparameters by iteratively suggesting a set of hyperparameters and evaluating their performance.
Random search samples hyperparameters from the search space with equal probability.
Despite not learning from previous trials, random search is commonly used for its simplicity and is much more efficient than grid search [2].
The TPE algorithm [3] builds a probabilistic model of task performance over the hyperparameter space and uses the results of past trials to choose the most promising next configuration to train, which the TPE algorithm defines as maximizing the Expected Improvement value [19].
Evolutionary algorithms [28, 18] train multiple models in parallel and replace poorly performing models with “mutated” copies of the current best models.
A UTO T RANSFER is a general
AutoML solution and can be applied in combination with any of these HPO algorithms.
Transfer Learning in AutoML.
Wong et al.
[34] proposed to transfer knowledge across tasks by reloading the controller of reinforcement learning search algorithms.
However, this method assumes that the search space on different tasks starts with the same learned prior.
Unlike A UTO T RANSFER , it cannot address the core challenge in GNN AutoML: the best GNN design is highly task-specific.
GraphGym([38]) attempts to transfer the best architecture design directly with a metric space that measures task similarity.
GraphGym([38]) computes task similarity by training a set of 12 ”anchor models” to convergence which is computationally expensive.
In contrast, A UTO T RANSFER designs light-weight task embeddings which require minimal computations overhead.
Task Embedding.
There is prior research trying to quantify task embeddings and similarities.
Similar to GraphGym, Taskonomy [40] estimates the task affinity matrix by summarizing final losses/evaluation metrics using an Analytic Hierarchy Process [29].
From a different perspective,
Task2Vec [1] generates task embeddings for a given task using the Fisher Information Matrix associated with a pretrained probe network.
This probe network is shared across tasks and allows
Task2Vec to estimate the Fisher Information Matrix of different image datasets.
It cannot be directly applied to GNNs as the inputs do not align across datasets.
A UTO T RANSFER avoids the bottleneck by using asymptotic statistics of the Fisher Information Matrix with randomly initiated weights.
Problem Formulation and Preliminaries
We first introduce formal definitions of data structures relevant to A UTO T RANSFER .
Definition 1 (Task)
We denote a task as T =
( D , L ( · )) , consisting of a dataset D and a loss function
L ( · ) related to the evaluation metric.
For each training attempt on a task T ( i ) , we can record its model architecture M j , hyperparameters
H j , and corresponding value of loss l j , i.e. , ( M j , H j , l j ) .
We propose to maintain a task-model bank to facilitate knowledge transfer to future novel tasks.
Definition 2 (Task-Model Bank)
A task-model bank B is defined as a collection of tasks, each with multiple training attempts, in the form of B = { ( T ( i ) , { ( M ( i ) j , H ( i ) j , l ( i ) j ) } ) } .
[Figure]
FIM
Scale-Invariant Measure
Task Feature
Task Embedding 𝑔
[Figure]
𝑓 𝑓
Task i
Training Objective
Task j x
Task i
Task k x
Margin Ranking Loss 𝑎" " 𝑎" # 𝑧 $ 𝑧 % 𝑧 %(’)
𝑧 %())
𝑧 %(’)
𝑧 %(*)
AutoML with Knowledge Transfer.
Suppose we have a task-model bank B .
Given a novel task
T ( n ) which has not been seen before, our goal is to quickly find a model that works reasonably well on the novel task by utilizing knowledge from the task-model bank.
In this paper, we focus on AutoML for graph learning tasks, though our developed technique is general and can be applied to other domains.
We define the input graph as G =
{ V, E } , where V is the node set and E ⊆
V × V is the edge set.
Furthermore, let y denote its output labels, which can be node-level, edge-level or graph-level.
A GNN parameterized by weights θ then outputs a posterior distribution P ( G, y, θ ) for label predictions.
Proposed Solution: A UTO T RANSFER
In this section, we introduce the proposed A UTO T RANSFER solution.
A UTO T RANSFER uses the task embedding space as a tool to understand the relevance of previous architectural designs to the target task.
The designed task embedding captures the performance rankings of different architectures on different tasks while also being efficient to compute.
We first introduce a theoretically motivated solution to extract a scale-invariant performance representation of each task-model pair.
We use these representations to construct task features (Section 4.2) and further learn task embeddings (Section 4.3).
These embeddings form the task embedding space that we finally use during the AutoML search (Section 4.4).
4.1 Basics of the Fisher Information Matrix (FIM)
Given a GNN defined above, its Fisher Information Matrix (FIM) F is defined as
[Equation]
which formally is the expected covariance of the scores with respect to the model parameters.
There are two popular geometric views for the FIM.
First, the FIM is an upper bound of the Hessian and coincides with the Hessian if the gradient is 0.
Thus, the FIM characterizes the local landscape of the loss function near the global minimum.
Second, similar to the Hessian, the FIM models the loss landscape with respect not to the input space, but to the parameter space.
In the information geometry view, if we add a small perturbation to the parameter space, we have
[Equation]
where KL ( · , · ) stands for Kullback–Leibler divergence.
It means that the parameter space of a model forms a Riemannian manifold and the FIM works as its Riemannian metric.
The FIM thus allows us to quantify the importance of a model’s weights in a ways that is applicable to a variety of different architectures.
4.2 FIM-based Task Features
Scale-invariant Representation of Task-Model Pairs.
We aim to find a scale-invariant representation for each task-model pair which will form the basis for constructing task features.
The major challenge in using the FIM to represent GNN performance is that graph datasets do not have a universal, fixed input structure, so it is infeasible to find a single pretrained model and extract its FIM.
However, training multiple networks poses a problem as the FIMs computed for different networks are not directly comparable.
We choose to use multiple networks but additionally propose to use asymptotic statistics of the FIM associated with randomly initialized weights.
The theoretical justification for the relationship between the asymptotic statistics of the FIM and the trainability of neural networks was studied in [20, 24] to which we refer the readers.
We hypothesize that such a measure of trainability encodes loss landscapes and generalization ability and thus correlates with final model performance on the task.
Another issue relates to input structures of graph datasets is that different models have a different number of parameters.
We note that each GNN architecture design can be represented as a sequence of pre-processing layers, message passing layers, and post-processing layers.
Pre-process layers and post-process layers are Multilayer Perceptron (MLP) layers, of which the dimensions vary across different tasks due to different input/output structures.
On the contrary, message passing layers are commonly regarded as the key design for GNNs and the number of weight parameters can remain the same across tasks.
In this light, we only consider the FIM with respect to the parameters in message passing layers so that the number of parameters considered stays the same for all datasets.
We further approximate the FIM by only considering the diagonal entries, which implicitly neglects the correlations between parameters.
We note that this is common practice when analyzing the FIMs of deep neural networks, as the full FIM is massive (quadratic in the number of parameters) and infeasible to compute even on modern hardware.
Similar to Pennington and Worah [24], we consider the first two moments of FIM m [Equation] [Equation] and m
[Equation]
= 1 n tr [Equation] and use α = m 2 /m 21 as the scale-invariant representation.
The computed α is lower bounded byand captures how concentrated the spectrum is.
A small α indicates the loss landscape is flat, and its corresponding model design enjoys fast first-order optimization and potentially better generalization.
To encode label space information into each task, we propose to train only the last linear layer of each model on a given task, which can be done efficiently.
The parameters in other layers are frozen after being randomly initialized.
We take the average over R initializations to estimate the average ¯ α .
Constructing Task Features.
We denote task features as measures extracted from each task that characterize its important traits.
The design of task features should reflect our final objective: to use these features to identify similar tasks and transfer the best design distributions.
Thus, we select U model designs as anchor models and concatenate the scale-invariant representations ¯ a u of each design as task features.
To retain only the relative ranking among anchor model designs, we normalize the concatenated feature vector to a scale of 1.
We let z f denote the normalized task feature.
4.3 From Task Features to Task Embeddings
The task feature z f introduced above can be regarded as a means of feature engineering.
We construct the feature vector with domain knowledge, but there is no guarantee it functions as anticipated.
We thus propose to learn a projection function g ( · ) : R U → R D that maps task feature z f to final task embedding z e = g ( z f ) .
We do not have any pointwise supervision that can be used as the training objective.
Instead, we consider the metric space defined by GraphGym.
The distance function in
GraphGym-computed using the Kendall rank correlation between performance rankings of anchor models trained on the two compared tasks-correlates nicely with our desired knowledge transfer objective.
It is not meaningful to enforce that task embeddings mimic GraphGym’s exact metric space, as GraphGym’s metric space can still contain noise, or does not fully align with the transfer objective.
We consider a surrogate loss that enforces only the rank order among tasks.
To illustrate, let us consider tasks T ( i ) , T ( j ) , T ( k ) and their corresponding task embeddings, z ( i ) e , z ( j ) e , z ( k ) e .
Note that z e is normalized to 1 so z ( i ) e ⊤ z ( j ) e measures the cosine similarity between tasks T ( i ) and T ( j ) .
Let d g ( · , · ) denote the distance estimated by GraphGym.
We want to enforce
[Equation]
To achieve this, we use the margin ranking loss as our surrogate supervised objective function:
[Equation]
Here if d g ( T ( i ) , T ( j ) ) <
d g ( T ( i ) , T ( k ) ) , then we have its corresponding label y
= 1 , and y = −otherwise.
Our final task embedding space is then a FIM-based metric space with cosine distance function, where the distance is defined as d e ( T ( i ) , T ( j ) )
= 1 − z
( i ) e ⊤ z ( j ) e .
Please refer to the detailed training pipeline at Algorithm 2 in the Appendix.
4.4 AutoML Search Algorithm with Task Embeddings
Most existing search algorithms do not leverage information about the optimal parameters found for previous tasks and usually require a few runs to warm up the search.
This is because it is nontrivial to estimate the performance of a set of configurations given its performance on a different task.
With our learned task embeddings, we are able to efficiently relate useful training trials in our task-model bank with the task of interest.
To transfer knowledge to a novel task, a naïve idea would be to directly carry over the best model configuration from the closest task in the bank.
However, even a high Kendall rank correlation between model performance rankings of two tasks T ( i ) , T ( j ) does not guarantee the best model configuration in task T ( i ) will also achieve the best performance on task T ( j ) .
In addition, since task similarities are subject to noise, this naïve solution may struggle when there exist multiple reference tasks that are all highly similar.
To make the knowledge transfer more robust to such failure cases, we introduce the notion of design distributions that depend on top performing model designs and propose to transfer design distributions rather than the best design configurations.
Formally, consider a task T ( i ) in the taskmodel bank B , associated with its trials { ( M ( i ) j , H ( i ) j , l ( i ) j ) } .
We can summarize its designs as a list of configurations C
= { c 1 , . .
. , c W } , such that all potential combinations of model architectures M and hyperparameters H fall under the Cartesian product of the configurations.
For example, c 1 might be the instantiation of aggregation layers, and c 2 might be the start learning rate.
We then define design distributions as random variables c 1 , c 2 , . .
c W each corresponding to a hyperparameter.
Each random variable c is defined as the frequency distribution of the design choices used in the top
K trials.
We multiply all distributions for the individual configurations { c 1 , . . .
, c W } to approximate the overall task’s design distribution P ( C | T ( i ) ) =
(cid:81) w P ( c w | T ( i ) ) .
During inference, given a novel task T ( n ) , we select a close task subset S by thresholding task embedding distances, i.e. , S =
{ T ( i )
|
d e ( T ( n ) , T ( i ) )
≤ d thres } .
We then derive the transferred design prior P t ( C | T ( n ) ) of the novel task by weighting design distributions from the close task subset S .
[Equation]
The inferred design prior for the novel task can then be used to guide various search algorithms.
The most natural choice for a few trial regime is random search.
Rather than sampling each design configuration following a uniform distribution, we propose to sample from the task-informed design prior P t ( C | T ( n ) ) .
Please refer to Appendix A to check how we augment other search algorithms.
For
A UTO T RANSFER , we can preprocess the task-model bank
B into
[Equation] T ( i ) ) } as our pipeline only requires using task embedding z ( i ) e and design distribution P ( C | T ( i ) ) rather than detailed training trials.
A detailed search pipeline is summarized in Algorithm 1.
Experiments 5.1 Experimental Setup
Task-Model Bank: GNN-B ANK -101.
To facilitate AutoML research with knowledge transfer, we collected GNN-B ANK -101 as the first large-scale graph database that records reproducible design
Require: A processed task-model bank B p = { ( D ( i ) , L ( i ) ( · )) , z ( i ) e , P ( C | T ( i ) ) } , a novel task
T ( n ) = (cid:0) D ( n ) , L ( n ) ( · ) (cid:1) , U anchor models M 1 , ..., M U , R specifies the number of repeats.
1: for u = 1 to U do 2: for r = 1 to R do 3:
Initialize weights θ for anchor model M u randomly 4:
Estimate FIM F ←− E D [ ∇ θ log P ( M u , y, θ ) ∇ θ log P ( M u , y, θ ) ⊤ ] 5:
Extract scale-invariant representation a ( v ) u
←− m 2 /m 21 following Eq.6: end for 7: ¯ a u ←− mean ( a (1) u , a (2) u , ..., a ( V ) u ) 8: end for 9: z ( n ) f ←− concat (¯ a 1 , ¯ a 2 , ..., ¯ a U ) 10: z ( n ) e ←− g ( z ( n ) f ) 11: Select close task subset S ←− { T ( i ) |
1 − z ( n ) e ⊤ z ( i ) e ≤ d thres } 12: Get design prior P t ( C | T ( n ) ) by aggregating subset S following Eq.13:
Start a HPO search algorithm with the task-informed design prior P t ( C | T ( n ) ) configurations and detailed training performance on a variety of tasks.
Specifically, GNN-B ANKcurrently includes six tasks for node classification (AmazonComputers [31], AmazonPhoto [31], CiteSeer [36], CoauthorCS [31], CoauthorPhysics [31], Cora [36]) and six tasks for graph classification (PROTEINS [17], BZR [17], COX2 [17], DD [17], ENZYMES [17], IMDB-M [17]).
Our design space follows [38], and we extend the design space to include various commonly adopted graph convolution and activation layers.
Please refer to Table 2 in the Appendix for details.
We extensively run 10,000 different models for each task, leading to 120,000 total task-model combinations, and record all training information including train/val/test loss.
Generating the benchmark takes [-16,
-]16,
GPU hours on an NVIDIA T4 GPU.
We plan to release GNN-B ANK -101 to inspire further research in this direction.
Benchmark Datasets.
We benchmark A UTO T RANSFER on six different datasets following prior work [27].
Our datasets include three standard node classification datasets (CoauthorPhysics [31],
CoraFull [4] and OGB-Arxiv [16]), as well as three standard benchmark graph classification datasets, (COX2 [17], IMDB-M [17] and PROTEINS [17]).
CoauthorPhysics and CoraFull are transductive node classification datasets, so we randomly assign nodes into train/valid/test sets following a 50%:25%:25% split [27].
We randomly split graphs following a 80%:10%:10% split for the three graph classification datasets [27].
We follow the default train/valid/test split for the OGB-Arxiv dataset [16].
To make sure there is no information leakage, we temporarily remove all records related to the task from our task-model bank if the dataset we benchmark was collected in the bank.
Baselines.
We compare our methods with the state-of-the-art approaches for GNN AutoML.
We use GCN and GAT with default architectures following their original implementation as baselines.
For multi-trial NAS methods, we consider GraphNAS [10].
For one-shot NAS methods, we include
DARTS [23] and GASSO [27].
GASSO is designed for transductive settings, so we omit it for graph classification benchmarks.
We further provide results of HPO algorithms based on our proposed search space as baselines: Random, Evolution, and TPE.
We by default allow searching 30 trials maximum for all the algorithms, i.e. , an algorithm can traindifferent models and collect the model with the best accuracy.
We use the default setting for one-shot
NAS algorithms (DARTS and GASSO), as they only train a super-model once and can efficiently evaluate different architectures.
We are mostly interested in studying the few-trial regime where most advanced search algorithms degrade to random search.
Thus we additionally include a random search (3 trials) baseline where we pick the best model out of only 3 trials.
Runtime analysis and additional experimental details can be found in Appendix A.
5.2 Experiments on Search Efficiency
[Table] We evaluate A UTO T RANSFER by reporting the average best test accuracy among all trials considered over ten runs of each algorithm in Table 1.
The test accuracy collected for each trial is selected at the epoch with the best validation accuracy.
By comparing results from random search (3 trials) and
A UTO T RANSFER (3 trials), we show that our transferred task-informed design prior significantly improves test accuracy in the few-trial regime, and can be very useful in environments that are where computationally constrained.
Even if we increase the number of search trials to 30, A UTO T RANSFER still demonstrates non-trivial improvement compared with TPE, indicating that our proposed pipeline has advantages even when computational resources are abundant.
Notably, with only 3 search trial,
A UTO T RANSFER surpasses most of the baselines, even those that use 30 trials.
To better understand the sample efficiency of A UTO T RANSFER , we plot the best test accuracy found at each trial in Figure 3 for OGB-Arxiv and TU-PROTEINS datasets.
We notice that the advanced search algorithms (Evolution and TPE) have no advantages over random search at the few-trial regime since the amount of prior search data is not yet sufficient to infer potentially better design configurations.
On the contrary, by sampling from the transferred design prior, A UTO T RANSFER achieves significantly better average test accuracy in the first few trials.
The best test accuracy at trial 3 of A UTO T RANSFER surpasses its counterpart at trial 10 for every other methods.
5.3 Analysis of Task Embeddings
Qualitative analysis of task features.
To examine the quality of the proposed task features, we visualize the proposed task similarity matrix (Figure 4 (b)) along with the task similarity matrix (Figure 4 (a)) proposed in GraphGym.
We show that our proposed task similarity matrix captures similar patterns as GraphGym’s task similarity matrix while being computed much more efficiently by omitting training.
We notice that the same type of tasks, i.e. , node classification and graph classification, share more similarities within each group.
As a sanity check, we examined that the [Figure] closest task in bank with respect to CoraFull is Cora.
The top 3 closest tasks for OGB-Arxiv are
AmazonComputers, AmazonPhoto, and CoauthorPhysics, all of which are node classification tasks.
Generalization of projection function g ( · ) .
We proposed the projection function g ( · ) with trainable weights with the goal of getting task embeddings that better capture our knowledge transfer objective through this surrogate supervision.
To investigate whether the trained g ( · ) can generalize to novel tasks, we conduct leave-one-out cross validation with all tasks in our task-model bank.
Concretely, for each task considered as a novel task T ( n ) , we use the rest of the tasks, along with their distance metric d g ( · , · ) estimated by GraphGym, to train the projection function g ( · ) .
For each method (Task
Feature, Task Embedding, and GraphGym), we sort all other tasks based on the defined similarity with respect to the central novel task.
We calculate Kendall rank correlation over similarities of tasks for Task Feature and Task Embedding against GraphGym.
We plot the average rank correlation and the standard deviation over ten runs on Figure 4 (c).
We conclude that with the trained projection function g ( · ) , our task embeddings correlate better with GraphGym’s estimation across the board compared to using task feature embeddings, indicating that our proposed learning pipeline for task embedding can generalize to novel tasks.
Conclusion
In this paper, we study how to improve AutoML search efficiency by transferring existing architectural design knowledge to novel tasks of interest.
We introduce a task-model bank that captures the performance over a diverse set of GNN architectures and tasks.
We also introduce a computationally efficient task embedding that can accurately measure the similarity between different tasks.
We release GNN-B ANK -101, a large-scale database that records detailed GNN training information of 120,000 task-model combinations.
We hope this work can facilitate and inspire future research in efficient AutoML to make deep learning more accessible to a general audience.
Acknowledgements
We thank Xiang Lisa Li, Hongyu Ren, Yingxin Wu for discussions and for providing feedback on our manuscript.
We also gratefully acknowledge the support of DARPA under Nos. HR (TAMI), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-10171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF- (Expeditions), NIH under No. 3U54HG010426-04S1 (HuBMAP), Stanford Data Science Initiative,
Wu Tsai Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper
Networks, KDDI, NEC, and Toshiba.
The content is solely the responsibility of the authors and does not necessarily represent the official views of the funding entities.
A Additional Implementation Details
Runtime analysis.
We have empirically shown that A UTO T RANSFER can significantly improve search efficiency by reducing the number of trials needed to achieve reasonably good accuracy.
The only overhead we introduced is the procedure of estimating task embeddings.
Since we use a randomly initialized architecture, extracting each task feature only requires at most one forward and a few backward passes from a single minibatch of data.
The wall-clock time depends on the size of the network and data structure.
In our experiments, it typically takes a few seconds on an NVIDIA [-T
-]T
GPU.
We repeat the task feature extraction process 5 times for each anchor model, and for a total of 12 anchor models.
Thus, the wall-clock time of the overhead for computing task embedding of novel task is within a few minutes.
We note the length of this process is generally comparable to one trial of training on a small-sized dataset, and the time saved is much more significant for large-scale datasets.
Details for task-model-bank training.
Our GNN model specifications are summarized in Table 2.
Our code base was developed based on GraphGym [38].
For all the training trials, We use the Adam optimizer and cosine learning rate scheduler (annealed to 0, no restarting).
We use L2 regularization with a weight decay of 5e-4.
We record losses and accuracies for training, validation and test splits every 20 epochs.
Training details for A UTO T RANSFER .
We summarize the training procedure for the projection function g ( · ) in Algorithm 2.
We set U = 12 and R = 5 throughout the paper.
We use the same set of anchor model designs as those in GraphGym.
We use a two-layer MLP with a hidden dimension of 16 to parameterize the projection function g ( · ) .
We use the Adam optimizer with a learning rate of 5e-3.
We use margin = 0 . 1 and train the network for 1000 iterations with a batch size of 128.
We adopt K = 16 when selecting the top K trials that are used to summarize design distributions.
Details for adapting TPE, evolution algorithms.
We hereby illustrate how to compose the search algorithms with the transferred design priors.
TPE is a Bayesian hyperparameter optimization method, meaning that it is initialized with a prior distribution to model the search space and updates the prior as it evaluates hyperparemeter configurations and records their performance.
We replace this prior distribution with the task-informed design priors.
As evolutionary algorithms generally initialize a large population and iteratively prune and mutate existing networks, we replace the random network initialization with the task-informed design priors.
As we mainly focus on the few trial search regime, we set the fully random warm-up trials to 5 for both TPE and evolution algorithms.
Code, dataset release.
We will make source code and collected task-model bank public at the time of publication.
[Table]
B Additional Discussion
Limitations.
In principle, A UTO T RANSFER leverages the correlation between model performance rankings among tasks to efficiently construct model priors.
Thus, it is less effective if the novel task has large task distances with respect to all tasks in the task-model bank.
In practice, users can
Require: Task features { z ( i ) f
| T ( i ) } extracted for each task from the task-model bank.
Distance measure d g ( · , · ) estimated in GraphGym.
1: for each iteration do 2:
Sample T ( i ) , T ( j ) , T ( k ) 3: z ( i ) e , z ( j ) e , z ( k ) e ←− g ( z ( i ) f ) , g ( z ( j ) f ) , g ( z ( k ) f ) 4: y ←− 1 if d g ( T ( i ) , T ( j ) ) <
d g ( T ( i ) , T ( k ) ) else −5:
Optimize objective function L r ( z ( i ) e , z ( j ) e , z ( k ) e , y ) in Eq.6: end for continuously add additional search trials to the bank.
As the bank size grows, it will be less likely that a novel task has a low correlation with all tasks in the bank.
Social Impact.
Our long-term goal is to provide a seamless GNN infrastructure that simplifies the training and deployment of ML models on structured data.
Efficient and robust AutoML algorithms are crucial to making deep learning more accessible to people who are interested but lack deep learning expertise as well as those who lack the massive computational budget AutoML traditionally requires.
We believe this paper is an important step to provide AI tools to a broader population and thus allow for AI to help enhance human productivity.
We will release a dataset containing GNN training information to promote future research.
The datasets we used for experiments are among the most widely-used benchmarks, which should not contain any undesirable bias.
Besides, training/test losses and accuracies are highly summarized statistics that we believe should not incur potential privacy issues.
C Additional Results
Search efficiency.
We summarize the average number of trials needed to surpass the average best accuracy found by TPE with 30 trials in Table 3.
We show that A UTO T RANSFER reduces the number of explored architectures by an order of magnitude .
[Table]
Ablation study on number of achor models.
Here we empirically demonstrate how the number of anchor models affect the rank correlation in Table 4.
While 3 anchor models is not enough to capture the task distance, we found that 9 and 12 has a satisfactory tradeoff between capturing task distance and computational efficiency
[Table]
[Table]