TAFS: Task-aware Activation Function Search for Graph Neural Networks

GIN experiments

---

Dataset: ogbg-molhiv

Q3. * The authors should conduct a time complexity analysis and compare it with GReLU.*

Answer: The setting of GReLU is actually very different from ours. GReLU designs a hyperfunction which includes additional graph convolutional layers, as a result, GReLU is not a univariate function as commonly required for activation functions (on the other hand, ReLU, Swish, Tanh and our TAFS as univariate functions). This can also be perceived by GReLU’s parameters. GReLU has N * C * K extra parameters, where N is number of nodes, C is the node feature dimensions and K is number of segments. So GReLU is dependent on the graph dataset size. TAFS is independent. Our TAFS brings extra parameter only dependent to number of graph layers, which is basically M * L, where M is the size of MLP (usually dozens) and L is the applied layers.

Thank you for your appreciation of our work and for your insightful questions. We answer each below.

Q1: Graph classification is crucial in graph mining due to its emphasis on global topological structures. Given its distinct objectives from node classification, it's imperative for the author to include experiments on graph classification.

Answer: We add below the performance of graph classification on the ogbg-molhiv dataset. We include baselines of classical GNN as well as Top-1 ranking solution PAS+FPS. In all cases, TAFS could achieve a performance gain.

Graph level task Dataset: ogbg-molhiv, #Graphs 41K, #Nodes per graph 25.5

Q2. As a universal method, TAFS should be tested on as many GNN models as possible. GCN and GraphSAGE are quite similar, while GAT and GIN are different, which should also be used as backbone to evaluate the performance of TAFS.

Answer: We agree. Here are our detailed experiments on GAT/GIN across small and large datasets. As can be seen, TAFS achieves the best in most cases, showing the effectiveness of our method.

GAT experiments

---

Dataset: ogbn-proteins

Q5. In a previous work [1], the best-performing activation function is referred to as Swish, which has a specific functional form. Similarly, can the best-performing activation function identified by TAFS be represented using a generic function? [1] Prajit Ramachandran, Barret Zoph, and Quoc V . Le. Searching for activation functions. In International Conference on Learning Representations, ICLR, 2018.

Answer: This is an interesting question. We answer in two ways. Firstly, we could achieve the closed form equation by symbolic regression techniques. Because we have obtained the univariate function and we could simulate many (x, y) value pairs to symbolic regression to distill. Swish could do so because it uses a template based space which naturally leads to this symbolic form as a result. However, this has a big negative impact on the possible function space, i.e., it is very possible the suitable activation function is not in this manually-designed template space. Secondly, in terms of usage, we do not need to distill the explicit form because we open source the learned weights of the univariate function and for practitioners, they just need to load this small function and replace their ReLU. The usage is thus very easy.

We show below a case study of the symbolized version of TAFS searched results. On the Cora dataset with GCN-JK network baseline, the searched result is the blue line. We distill an explicit formula by symbolic regression: y = 0.6 Tanh(-x) and plug in this activation function back to the model. The performance is shown below and we provide a visualization in the appendix Sec H. There is a performance gap in the table below. This is because explicit symbolic space is not accurate, which shows further that our implicit functional space is expressive.

Q6-Q7: The sentence in the main text, "… how can we design GNN activation functions to adapt effectively to various graph-based tasks, creating task-aware activation functions??" seems to have an issue with the use of symbols.

In the text, "... and demote by ¯w_δ all the parameters of GNN, …", the word "demote" appears to be misspelled.

Answer: Thanks for your careful review. We have corrected these typos.

Q8: * The experimental section lacks clear explanations of the metrics used in the tables. For example, the description of Table 2 does not specify whether the results represent accuracy, AUC, or something else.*

Answer: In Table 2, the standard accuracy is used. In Table 3, both ROCAUC and PRAUC are used. In our extended OGB experiments, we use ROCAUC for node classification and graph classification, Hit@20 for link prediction, following the OGB standard.Thus, we have experimented on various tasks of different objective requirements.

Q9: In Figure 4, subfigure (a) shows that larger values of K yield better results but at a slower pace, demonstrating a "trade-off between accuracy and computation time." Is the Y-axis in (a) representing accuracy? And it's not clear how the relationship between K value and computation time is depicted.

Answer: Yes, Y-axis is the downstream performance, here classification accuracy. For example, moving from K=32 towards K=64 increases 20% more computational times. As a result, it is truly a trade-off that should depend on the user's own preference. We will add this discussion in the paper.

Thank the reviewer for insightful questions. We answer each below.

Q1. The search strategy presented in this paper does not support a larger GNN search space. For instance, if there is a need to search for aggregation functions or message passing functions, the optimal form of the activation function is likely to change.

Answer: Actually, this is totally integratable. Our method TAFS could be applied after the other GNN search method finds the backbone model (for example, GraphNAS, 2019 Gao et al). We experiment this strategy on the OGB graph classification top 1 winner PAS+FPS method. It is a neural architecture search baseline that does exactly GNN space search. In the searched 14 layers GNN checkpoint, we replace its ReLU functions by TAFS, and fine tune further. As can be seen, within an affordable time, the performance could boost further, showing that TAFS is easily integrable with the GNN NAS methods.

Reference: https://ogb.stanford.edu/docs/leader_graphprop/#ogbg-molhiv

Q2. If TAFS employs an MLP to approximate the optimal activation function, the performance of activation function would also depend on the number of layers in the MLP and the non-linear transformation functions. This has not been discussed in this paper.

Answer: We highlight that MLP is a universal approximator which we choose on purpose as an expressive functional space. As a result, a normal MLP (not too large) would be sufficient in finding the suitable functions.

We include the discussion on MLP hyperparameters in Figure 4(b)(c). It is true that this influences the performance. In order not to add too many extra parameters, we choose MLP with one-hidden layer (100 neurons) of Tanh activation function. No activation function is applied on the output neuron. We experiment below the differences of ReLU/Tanh used in MLP. Tanh is thus our final choice.

Dataset: DBLP

Q3. In experiments, although the paper introduces a Jacobian regularization term, the impact of this regularization has not been empirically tested.

Answer: Thanks for pointing this out. We give below the influence of TAFS w./w.o. regularization. We will add this part to the paper later.

Dataset: DBLP

Q4. In TAFS, a stochastic relaxation is used, involving the reparameterization of the activation function parameters using a Gaussian distribution. However, this paper does not discuss the advantage of this strategy.

Answer: The biggest advantage of the stochastic relaxation is to remove the differentiability requirement for the task evaluation metric (e.g. ROCAUC, Hit@20), which is very common in Graph tasks. By sampling the weights of activation functions from a probability distribution, we optimize the distribution parameters by proposition 1 in the paper. The proposition 1 gives the gradient of loss w.r.t distribution parameters, and there is no terms such as \nable M, avoiding the calculation of gradient of M. This makes the method more applicable in general tasks. The choice of Gaussian distribution is without loss of generality because we don’t have prior knowledge about this choice of distribution. In cases where we know certain probability distribution is better, we could definitely change to others such Poisson distribution.

We give below the influence of TAFS w./w.o. stochastic regularization. The full table has been updated in Table 3 in the paper.

Thanks for the questions. We answer each below.

Q1. Some of the experimental settings and results explanations in the paper are vague, such as Figure 1 in the Introduction section and Figure 3 in the Experiments section. I cannot understand how the experimental results in the figures were obtained, and what the data in the figures mean.

Answer: In a word, Figure 1 illustrates the impact of activation function in GNN and motivates our work to search activation functions. Figure 3 demonstrates the searched candidates from baselines and our proposed TAFS.

Figure 1 shows the experiments on two GNN baselines (GCN and GraphSage) on two typical graph datasets (Core and DBLP), with different choices of activation functions (ReLU, Tanh, Leaky-Relu, Swish). The evaluation metric is node classification accuracy. We show by Figure 1 that activation functions in GNN, which was previously paid little attention, could have a huge impact on the performance. This motivates our study to design efficient activation function search algorithms in this paper.

Figure 3 demonstrates multiple things. Fig 3(a) gives an idea on the search space of our compared baselines Swish and APL to show what functions in their search space looks like. Fig 3(b)(c) shows our searched activation functions tailoring two datasets Cora and DBLP. Specifically, we show that (1) our method finds adaptive activation functions for different datasets (2) due to the efficiency of our method, we could adapt layer-wise activation functions without exploding the search space. So different layers find different activation functions (3) our found activation functions are distinct to already published manually-designed ones (4) our activation functions are smooth functions which potentially reflects our smoothness regularization.

Q2. From the experimental results, the performance of TAFS needs to be improved. Table 2 shows that in some experimental scenarios on the DBLP, Cornell, Texas and Chameleon datasets, the results obtained by TAFS are only marginally better than directly using a specific activation function, or even worse. The authors did not explain this phenomenon.

Answer: Thank the reviewer for pointing this out. In our previously table version, we bold the best results by considering only the mean over multiple runs. Actually, to make it more precise, we now take into account both the mean and standard variance and bold best results of the same statistical significance level.

In the revised pdf, we see now that TAFS is best or tied for best in all cases. In some datasets which are small such as Texas, Chameleon, the variance is very high. Thus, TAFS archives the same performance level.

This phenomenon is much alleviated in large scale datasets such as OGB ones. In our experiments with OGB, we show that TAFS is the best as below.

Node level task Dataset: ogbn-proteins, Node 132,534, Edge 39,561,252

---

Link level task dataset: ogbl-ddi, Node 4,267, Edge 1,334,889

Q3. I wonder why in the drug and protein interaction prediction experiments, only the results of TAFS and ReLU were provided, instead of comparing with multiple activation functions as in the node classification experiments.

Answer: We add the other choices’ performance below. We didn’t include this to emphasize the impact by adopting TAFS. The table will be merged to the main paper later.

Q4. Since the authors compared the search efficiency with Swish and APL, why didn’t they involve the comparison with these two search-based methods in the effectiveness validation experiments (Table 2 and Table 3)?

Answer: Due to the inefficiency of Swish and APL, we could not practically obtain reasonable results by Searchable Swish and APL. For example, in the case of Table 4 in the paper, we can see that searchable Swish takes 50 Hours to run on Chameleon and APL costs 6 times more memory consumption than TAFS. In our table 2, we have included 5 datasets and 10 graph baselines, leading to a total of 50 sets of experiments. We find it not necessary to include the searchable Swish and APL in every single set of task. As a result, we included already in Table 2 and Table 3 the searched results of Swish paper (Swish), which is worse than TAFS. We also give below the APL results and searchable swish on two tasks for a general idea of their ineffectiveness. Note that both baselines could not be applied in OGB datasets due to inefficiency in time and memory, while TAFS could.

Dataset: DBLP

Graph level task Dataset: ogbg-molhiv, #Graphs 41K, #Nodes per graph 25.5

---

Q3. I'd appreciate if the authors can elaborate more on the difference of the proposed method versus the activation function search methods for CNNs and RNNs (as referenced in Sec. 1).

Answer: The two most important activation function search methods are Swish and APL, which is explained and compared in Sec 2 and especially in Table 1. We give a detailed introduction in Appendix B. Swish is inefficient in running time because every inner optimization will train the network until convergence. APL is inefficient in parameters and memory because it parameterized each neuron’s activation by multiple hinges. In the case of ogbg-molhiv datasets (above) and 1st ranking solution PAS on graph classification, the total parameter is 27M and the running time is over 17 hours. APL will add over 150M parameters which leads to OOM and Swish will cost over 70 days even if we train only 100 epochs. Our method TAFS, takes only 23 hours and adds only 0.01M parameters, which is significantly more efficient than baselines, and as a result, paves the way for further study in activation function search.

Thanks for the important suggestion on the scalability. We answer each question below.

Q1: The experimental results are not very convincing. The experiments are all conducted on very small datasets, and I don't think the datasets for link prediction experiments are the commonly used ones in literature. I'd recommend the authors to use more standardized and commonly accepted benchmarks such as the OGB ones. Although the proposed method is already much more efficient than the baselines Swish and APL, it still takes about 10x times of runtime when compared with fixed activation functions. It's hard to tell whether the performance improvements worth such huge overhead on the time cost.

Q2. Although the proposed method is already much more efficient than the baselines Swish and APL, it still takes about 10x times of runtime when compared with fixed activation functions. It's hard to tell whether the performance improvements worth such huge overhead on the time cost.

Answer: We answer Q1 and Q2 together. The inappropriate comparison in the manuscript was due to the fact that the baseline was too small and as a result the computational head seems to be 10x more, though our method is already much more efficient than literature. Actually, our method’s parameters have no dependency on the size of the graph, making it easily scalable. We experimented on OGB and much larger baselines as in the table below. Three OGB datasets in node/link/graph level tasks are chosen and we experiment both classical GNN baselines and top-ranking baselines. As can be seen, on larger datasets (ogbn-molhiv) and larger baselines (PAS) which take over 17 hours, our method only adds about 35% more computational time, which is much more acceptable than baselines. Actually, both Swish search and APL search could not be run on such larger datasets due to their inefficiency in running time and overparameterization, which leads to months’ GPU hours or leads to Out-of-Memory error.

---

Node level task Dataset: ogbn-proteins, Node 132,534, Edge 39,561,252

---

Link level task dataset: ogbl-ddi, Node 4,267, Edge 1,334,889