SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yijian Qin $^{1,2}$ , Xin Wang $^{1,2*}$ , Ziwei Zhang $^{1}$ , Hong Chen $^{1}$ , Wenwu Zhu $^{1,2*}$ $^{1}$ Department of Computer Science and Technology, Tsinghua University,  $^{2}$ BNRist† qinyj19@mails.tsinghua.edu.cn, {xin_wang,zwzhang}@tsinghua.edu.cn h-chen20@mails.tsinghua.edu.cn, wwzhu@tsinghua.edu.cn
Graph neural architecture search (GraphNAS) has shown great potential for automatically designing graph neural architectures for graph related tasks. However, multi-task GraphNAS, capable of handling multiple tasks simultaneously and capturing the complex relationships and dependencies between them, has been largely unexplored in literature. To tackle this problem, we propose a novel multi-task graph neural architecture search with task-aware collaboration and curriculum (MTGC $^3$ ), which is able to simultaneously discover optimal architectures for different tasks and learn the collaborative relationships among different tasks in a joint manner. Specifically, we design the structurally diverse supernet to manage multiple architectures and graph structures in a unified framework, which combines with our proposed soft task-collaborative module to learn the transferability relationships between tasks. To further improve the architecture search procedure, we develop the task-wise curriculum training strategy that reweighs the influence of different tasks based on their relative difficulties. Extensive experiments show that our proposed MTGC $^3$  model achieves state-of-the-art performance against several baselines in multi-task scenarios, demonstrating its ability to discover effective architectures and capture collaborative relationships for multiple tasks.
Graph-structured data has attracted lots of attentions in recent years for its flexible representation ability in various domains. Graph neural networks (GNNs) such as GCN [17], GAT [39], and GIN [47] have been proposed and achieved great success in many graph applications. To save human efforts on designing GNN architectures for different tasks, graph neural architecture search (GraphNAS) [19, 7, 42] has been proposed to search for optimal GNN architectures. These automatically designed GNN architectures have achieved competitive or better performances compared with manually designed GNNs for a single task.
On the other hand, multi-task scenarios are ubiquitous for graph data. For example, in drug discovery, predicting multiple properties of molecules can be regarded as multiple tasks. Properly capturing the relationship between these tasks allows for better learning of them. Nevertheless, there has been no work reported on automatically searching GNN architectures for multi-task scenarios in the literature.
In this paper, we study the important yet unexplored problem of multi-task GraphNAS for the first time. Given that different GNNs employ different architectures for adaption to various graph tasks, we tackle this problem through jointly customizing an architecture for each individual task and considering the task collaborations to share useful information among different tasks. However, discovering
†BNRist is the abbreviation of Beijing National Research Center for Information Science and Technology.
the optimal GNN architectures for multiple tasks poses several challenges: (1) It is non-trivial to manage multiple architectures within a supernet, which is crucial to guarantee the effectiveness of the neural architecture search; (2) It is challenging to capture the complex relationships among different tasks for information sharing during the searching process; (3) It remains a challenge to balance the influence of different tasks on architecture optimization since different tasks may show diverse difficulties and optimization patterns.
To address these challenges, we propose the multi-task graph neural architecture search with task-aware collaboration and curriculum (MTGC $^3$ ) in this paper. Our proposed MTGC $^3$  model is able to discover multiple optimal architectures within a unified supernet and simultaneously optimizes them through learning the collaborative relationships of different tasks in a joint manner. Specifically, we design the structurally diverse supernet to manage multiple architectures and graph structures within a supernet for effective architecture search. Then, we introduce the soft task-collaborative module to capture the complex relationships among tasks, enabling sharing useful information during the searching process. We further develop the task-wise curriculum training strategy to better optimize the supernet via reweighing the influence of different tasks on architecture optimization based on task difficulties. Extensive experiments on both synthetic and real-world datasets validate the superiority of our proposed MTGC $^3$  model over existing baselines via customizing the optimal architecture for each task and sharing useful information among them. Detailed ablation studies further verify the effective designs of MTGC $^3$ .
- We are the first to investigate the problem of multi-task graph neural architecture search via proposing the multi-task graph neural architecture search with task-aware collaboration and curriculum (MTGC $^3$ ), to the best of our knowledge.
- We propose to jointly discover optimal architectures for multiple tasks and capture their complex relations within a unified framework by designing i) the structurally diverse supernet, ii) the soft task-collaborative module, and iii) the task-wise curriculum training strategy.
- Extensive experimental results demonstrate that our proposed  $\mathrm{MTGC}^3$  model outperforms state-of-the-art baselines on both synthetic and real-world datasets.
2 Problem Formulation and Preliminaries
2.1 Message-passing graph neural network
GNNs are state-of-the-art models for graph machine learning, which typically follow a message passing scheme where nodes aggregate information from their neighbors in each layer formulated as:
$$
\mathbf {m} _ {i} ^ {(l)} = \operatorname {A g g} \left(\mathbf {h} _ {j} ^ {(l)} \mid j \in \mathcal {N} _ {i}\right) \tag {1}
$$
$$
\mathbf {h} _ {i} ^ {(l + 1)} = \operatorname {U p d a t e} \left(\mathbf {m} _ {i} ^ {(l)}\right), \tag {2}
$$
where  $\mathbf{h}_i^{(l)}$  is the representation of node  $i$  at the  $l$ -th layer,  $\mathcal{N}_i$  denotes the neighbors of node  $i$  derived from the adjacent matrix,  $\mathrm{Agg}(\cdot)$  is the aggregation function,  $\mathrm{Update}(\cdot)$  is an updating function between two node representations. Different GNNs have different aggregation and updating functions, which is the main search objective for GraphNAS.
2.2 Multi-task graph neural architecture search
Given a set of graph tasks  $\{T_m\}_{m=1}^M$ , where each task has a set of graphs  $G_m = \{g_m^1, g_m^2, \dots, g_m^{n_m}\}$  in a shared graph space  $\mathcal{G}$  and the labels  $Y_m = \{y_m^1, y_m^2, \dots, y_m^{n_m}\}$  in its own label space  $\mathcal{V}_m$ . We denote  $\mathcal{V}$  as the composition of all label spaces, i.e.,  $\mathcal{V} = \bigcup \{\mathcal{V}_m\}_{m=1}^M$ . The goal of multi-task graph learning is to design a model  $F: \mathcal{G} \to \mathcal{V}$  to map the graphs to the label space while minimizing the loss function, i.e.,
$$
\arg \min  _ {F} \mathcal {L} \left(\left\{F \left(G _ {m}\right), Y _ {m} \right\} _ {m = 1} ^ {M}\right), \tag {3}
$$
where  $\mathcal{L}$  is the multi-task loss function. Here we use multi-task graph classification as an example, but the definition can be easily extended to multi-task node classification problems. In GraphNAS,
<sup>3</sup>https://github.com/THUMNLab/AutoGL-light
Figure 1: Kendall correlation coefficient of the performance ranks of different architectures on different tasks.
Figure 2: An overview of our proposed MTGC $^3$  method. Step 1: Performing forward propagation through the structurally diverse supernet and the soft task-collaborative module. Step 2: Updating architecture parameters and the soft task-collaborative module. Step 3: Updating model weights through the task-wise curriculum learning.
we mainly focus on the function  $f$  being GNNs. A typical GNN consists of two parts: an architecture  $\alpha \in \mathcal{A}$  and learnable weights  $w \in \mathcal{W}$ , where  $\mathcal{A}$  and  $\mathcal{W}$  denotes the architecture space and the weight space, respectively. Therefore, we denote GNNs in Eq.(1) as the following mapping function:  $F_{\alpha ,w}: \mathcal{G} \to \mathcal{V}$ . In this paper, we focus on the aggregation functions  $\mathrm{Agg}(\cdot)$ . Specifically, we consider a search space of standard layer-by-layer architectures without sophisticated connections such as residual or jumping connections, though our proposed method can be easily generalized. We choose five widely used message-passing GNN layers as our operation candidate set  $\mathcal{O}$ , including GCN [17], GAT [39], GIN [47], SAGE [11], k-GNN [27], and ARMA [12]. Besides, we also adopt MLP, which does not consider graph structures.
To develop a GraphNAS method for multi-task learning, we first consider what designs should be included that benefit multi-task learning. We propose an assumption on which our method is based.
Assumption 3.1. Different GNN architectures are suitable for learning different graph tasks. For example, if we search architectures separately for different tasks, i.e., the following bi-level optimization problem:
$$
\alpha_ {m} = \arg \min _ {\alpha} \mathcal {L} (F _ {\alpha , w (\alpha)} (G _ {m}), Y _ {m})
$$
$$
w (\alpha) = \arg \min _ {w} \mathcal {L} (F _ {\alpha , w} (G _ {m}), Y _ {m}),
$$
where  $\alpha_{m}$  is the most suitable architecture for task  $T_{m}$ . Then, the optimal architecture  $\alpha_{m}$  can be different for different tasks.
Assumption 3.1 implies that different tasks have different characteristics, while these characteristics require different architectures. To validate this assumption, we randomly choose several architectures and evaluate them on different tasks individually on two graph datasets: Tox21 and ToxCast (refer to Section 4.1 for details). We rank these architectures according to their performance for each task. Then we use Kendall rank correlation to compute the similarity score between rankings of the architectures for all pairs of tasks. The result is shown in Figure 1. We can find that many tasks have a negative correlation between different architectures, indicating that different architectures behave very differently on different tasks. The result illustrates the necessity to customize an architecture for different tasks instead of using a shared architecture.
On the other hand, it is also important to share useful information between different tasks, which is a common and basic assumption in multi-task learning [3]. Combining these two aspects, our goal is to simultaneously search multiple architectures for different tasks while learning the collaborative relationships of them. Figure 2 shows the overall framework of our method. In the following subsections, we present our proposed method in detail.
3.1 Structurally diverse supernet
Based on Assumption 3.1, we need to search multiple architectures for different tasks simultaneously. To effectively manage multiple architectures, we propose the structurally diverse supernet, following the widely adopted supernets in NAS [29]. The design principle is that, although different tasks require different architectures, we cannot completely separate these architectures since they need to share useful knowledge. We propose to balance this tradeoff by disentangling the supernet. In our structurally diverse supernet, the node representation at each layer is divided into  $N$  chunks with different GNN operations. In addition, we make the graph structure to be diverse in different chunks to further disentangle the hidden features and help different chunks to learn more specific knowledge for the corresponding tasks. A disentangled layer can be formulated as follows:
$$
f (\mathbf {x}, \mathbf {A}) = \| _ {i = 1} ^ {N} f _ {i} (\mathbf {x}, \mathbf {A} _ {i}), \tag {5}
$$
where  $f$  one of the layers in the mapping function  $F$ ,  $\mathbf{x}$  is the hidden features given by the last layer, and  $\mathbf{A}$  is the adjacent matrix. Moreover,  $\| \cdot$  represents concatenation,  $\mathbf{A}_{\mathrm{i}}$  is the adjacent matrix of the graph used in the  $i$ -th chunk,  $f_{i}$  denotes the GNN operation of the  $i$ -th chunk. By stacking the disentangled layers into a GNN, we can capture a variety of latent factors of the graphs by different types of GNN operations and graph structures, which is beneficial for multi-task learning.
Following the NAS literature [24], we adopt supernets to search for the optimal architecture in a differentiable way. Specifically, in the supernet, all possible operations are jointly considered by mixing different operations into a mixed operation, i.e.,  $f(\mathbf{x}) = \sum_{o_k \in \mathcal{O}} \alpha_k o_k(\mathbf{x})$ , where  $o_k$  is a candidate operation and  $\mathcal{O}$  is the operation search space,  $\alpha$  denotes the architecture parameters to be optimized. We combine the DARTS manner with our structurally diverse supernet to enable differentiable optimization, which can be formulated as follows:
$$
f _ {i} (\mathbf {x}, \mathbf {A} _ {i}) = \sum_ {o _ {k} \in \mathcal {O}} \alpha_ {i k} o _ {i k} (\mathbf {x}, \mathbf {A} _ {i}), \tag {6}
$$
where  $\alpha_{ik}$  is the architecture parameter of the  $k$ -th operation at the  $i$ -th chunk. As a result, multiple architectures in different chunks are contained in the supernet and they can also transfer knowledge to each other.
To adopt diverse structures for different chunks, we use a parameter-efficient and differentiable manner to generate the weights of edges for the graph structures, which can be formulated as follows:
$$
A _ {i, (u, v)} = \delta A _ {(u, v)} + (1 - \delta) A _ {(u, v)} \cdot \operatorname {s i g m o i d} \left(\left(\mathbf {S} _ {i} ^ {u} \mathbf {x} _ {u}\right) ^ {\top} \left(\mathbf {S} _ {i} ^ {v} \mathbf {x} _ {v}\right)\right), \tag {7}
$$
where  $A_{i, (u, v)}$  is the weight of edge between node  $u$  and  $v$  in  $\mathbf{A}_i$ ,  $A_{(u, v)}$  indicates if there is an edge between node  $u$  and  $v$  in the original graph,  $\mathbf{S}_i^u$  and  $\mathbf{S}_i^v$  are learnable parameters to generate edge weights of the  $i$ -th chunk,  $\delta$  is the hyper-parameter. Consequently, we generate an adjacent matrix  $\mathbf{A}_i$  in a continuous space. We can conveniently use gradient based methods to optimize it.
3.2 Soft task-collaborative module
Using the structurally diverse supernet, different chunks can exchange information to share knowledge. However, the independence of chunks is not encouraged, and therefore representations could be potentially mixed together. To further ensure that different chunks of the supernet correspond to relatively independent architectures, we propose a soft task-collaborative module. First, we rewrite Eq. (6) as follows:
$$
f _ {i} (\mathbf {x}, \mathbf {A} _ {i}) = \sum_ {o _ {k} \in \mathcal {O}} \alpha_ {i k} \sum_ {j} o _ {i j k} \left(\mathbf {x} _ {j}, \mathbf {A} _ {i}\right), \tag {8}
$$
where  $\mathbf{x}_j$  is the output of the  $j$ -th chunk of the previous layer, and  $o_{ijk}$  is the GNN operation only containing the corresponding parameters. For example, many types of GNN operations in Eq. (6) contain a parameter matrix with the shape  $d\times d_c$  in the node updating function in Eq. (2), where
$d_{c}$  is the dimensionality of each chunk and  $d = N \times d_{c}$  is the overall dimensionality. For these operations, we can split the parameter matrix into  $N$  matrices with the shape  $d_{c} \times d_{c}$ , where the  $j$ -th matrix corresponds to the part of  $\mathbf{x}_{j}$ . We denote  $o_{ijk}$  as the GNN operation only containing this part of parameters.
Using Eq.(8), we introduce our proposed soft task-collaborative module. Our goal is to keep the independence of each chunk, i.e., the  $i$ -th chunk of node representation at the previous layer mainly contributes to the  $i$ -th chunk of node representation at the next layer. We formulate it as follows:
$$
f _ {i} (\mathbf {x}, \mathbf {A} _ {i}) = \sum_ {o _ {k} \in \mathcal {O}} \alpha_ {i k} \sum_ {j} p _ {i j} o _ {i j k} (\mathbf {x} _ {j}, \mathbf {A} _ {i}), \quad p _ {i j} = \left\{ \begin{array}{l l} 1 & \text {i f} i = j, \\ \tanh  \theta_ {i j} & \text {o t h e r w i s e ,} \end{array} \right. \tag {9}
$$
where  $\theta$  is learnable parameters.  $p_{ij}$  represents the impact of the  $j$ -th chunk of node representation at the previous layer to the  $i$ -th chunk of node representation at the next layer. We remark that the absolute value of  $\tanh \theta$  indicates the amount of contribution while the sign of  $\tanh \theta$  indicates the positivity or negativity of the correlation. Since  $\tanh \theta \in (-1, 1)$ , we ensure that node representation from the same chunk has the largest impact on the next layer. Therefore, the architectures of different chunks are relatively independent, which helps to make the hidden representation disentangled and contain more useful information for different tasks. Besides, we initialize  $\theta$  as 0 so that the initial structurally diverse supernet simply consists of several separated GNN architectures at the beginning. By updating  $\theta$  in the training procedure, our proposed method can automatically learn the transferability between architectures.
After learning representations containing diverse information by the supernet, our model makes predictions for multiple tasks. We propose two approaches to building the task head.
Task-separated head: We manually pair each chunk of the supernet to a task, i.e., the node representation of the last layer of the supernet is split according to the chunks, and each chunk is connected to the head of the corresponding task. In this case, the architecture of each chunk can be regarded as the customized architecture for the corresponding task, and  $p_{ij}$  represents the collaborative pattern between different tasks, i.e., whether to share useful information between them. Since the number of chunks needs to be equal to the number of tasks, this approach is only applicable when the number of tasks is not too large.
Cross-mixed head: When the number of tasks is large, using the task-separated head will result in the dimensionality of the representation growing too large. Therefore, we further propose the cross-mixed head by randomly pre-assigning a few dimensions to each task, i.e., the head of each task is only connected to the representations of these dimensions. In practice, we overlay a mask tensor sampled from a Bernoulli distribution on the weights of the classification heads. This approach allows each task to acquire different hidden features while preventing some tasks from dominating all the architectures in the supernet.
3.3 Task-wise curriculum training
Our framework is end-to-end and can be optimized by the multi-task loss. But using multi-task loss may cause difficulties in training due to task imbalances [5]. At first we give the proposition below to analyze the gradient of different chunks in our model. We denote  $w_{j'}$  as the parameters that generate  $\mathbf{x}_j$  in the last layer.
Proposition 3.2. When using the overall loss function of multiple tasks  $\mathcal{L}$  for gradient backpropagation in our framework, the partial derivative of  $\mathcal{L}$  with respect to  $w_{j'}$  is:
$$
\frac {\partial \mathcal {L}}{\partial w _ {j ^ {\prime}}} = \sum_ {i} p _ {i j ^ {\prime}} \frac {\partial \mathcal {L}}{\partial f _ {i}} \sum_ {o _ {k} \in \mathcal {O}} \alpha_ {i k} \frac {\partial o _ {i j ^ {\prime} k} \left(\mathbf {x} _ {j ^ {\prime}} , \mathbf {A} _ {i}\right)}{\partial w _ {j ^ {\prime}}}. \tag {10}
$$
For the proof of Proposition 3.2, please refer to Appendix B. In Eq. (10), we can regard the latter part of the right side  $\sum_{o_k\in \mathcal{O}}\alpha_{ik}\frac{\partial o_{ij'k}(\mathbf{x}_{j'},\mathbf{A}_i)}{\partial w_{j'}}$  as the partial derivative of the mixed operation in a classical supernet in Eq. (6), and  $\frac{\partial\mathcal{L}}{\partial w_{j'}}$  is the sum of these partial derivatives weighted by a score  $p_{ij'}\frac{\partial\mathcal{L}}{\partial f_i}$ . However, different tasks may have different difficulties and show diverse patterns in the optimization process, e.g., tasks are unbalanced in terms of backpropagated gradient scales, some tasks with large score  $p_{ij'}\frac{\partial\mathcal{L}}{\partial f_i}$  may dominate the training process. Besides, in our model, the optimization
Algorithm 1: MTGC3
Input: Training dataset  $\{G_i\}_{i=1}^M$  and  $\{Y_i\}_{i=1}^M$ , learning rates  $\eta_{\alpha}$ ,  $\eta_{\mathbf{S}}$ ,  $\eta_{\theta}$  and  $\eta_w$   
1 Initialize learnable parameters  $w$ ,  $\alpha$ ,  $\mathbf{S}$  and  $\theta$ . Set  $\gamma = 0$ ;  
2 while not converge do  
3 Forward propagate by Eq.(5) and Eq.(9) to calculate loss  $\mathcal{L}$ ;  
4  $\alpha \leftarrow \alpha - \eta_{\alpha} \nabla_{\alpha} \mathcal{L}(w, \alpha, \mathbf{S}, \theta)$ ,  $\mathbf{S} \leftarrow \mathbf{S} - \eta_{\mathbf{S}} \nabla_{\mathbf{S}} \mathcal{L}(w, \alpha, \mathbf{S}, \theta)$ ,  $\theta \leftarrow \theta - \eta_{\theta} \nabla_{\theta} \mathcal{L}(w, \alpha, \mathbf{S}, \theta)$ ;  
5 Using  $\frac{\partial \mathcal{L}}{\partial f_j}$  to calculate  $p'$  by Eq.(11);  
6 Forward propagate by Eq.(5) and Eq.(11) to calculate loss  $\mathcal{L}'$ ;  
7  $w \leftarrow w - \eta_w \nabla_w \mathcal{L}'(w, \alpha, \mathbf{S}, \theta)$ ;  
8 Increase  $\gamma$ ;  
9 end
curves of different chunks can also differ greatly, which makes the optimization of our model further challenging.
To tackle these problems, we propose the task-wise curriculum training strategy. Specifically, we modify Eq. (9) using a curriculum strategy as:
$$
f _ {i} ^ {\prime} (\mathbf {x}, \mathbf {A} _ {i}) = \sum_ {o _ {k} \in \mathcal {O}} \alpha_ {i k} \sum_ {j} p _ {i j} ^ {\prime} o _ {i j k} \left(\mathbf {x} _ {j}, \mathbf {A} _ {i}\right), \quad p _ {i j} ^ {\prime} = \gamma p _ {i j} + (1 - \gamma) \frac {\| \frac {\partial \mathcal {L}}{\partial f _ {j}} \| _ {2}}{\| \frac {\partial \mathcal {L}}{\partial f _ {i}} \| _ {2}} \tanh  \frac {p _ {i j} \| \frac {\partial \mathcal {L}}{\partial f _ {i}} \| _ {2}}{\| \frac {\partial \mathcal {L}}{\partial f _ {j}} \| _ {2}}, \tag {11}
$$
where  $\gamma$  is a parameter that increases from 0 to 1 during the training process, and  $\| * \|_2$  is the  $l_2$ -normalization of the target partial derivative for measuring its scale. We use Eq. (11) instead of Eq. (9) while calculating the gradients of  $w$ . Using this curriculum strategy benefits the optimization process of our framework. When  $\gamma = 0$ , we have  $p_{ij}' < \frac{\|\frac{\partial\mathcal{L}}{\partial f_j}\|_2}{\|\frac{\partial\mathcal{L}}{\partial f_i}\|_2}$ , thus the contribution of tasks with large partial derivatives to the loss function is suppressed. For tasks with small  $\|\frac{\partial\mathcal{L}}{\partial f_j}\|_2$ , the tanh function tends to be a linear function, thus  $p_{ij}' \approx p_{ij}$  and their contributions are not affected. As  $\gamma$  increases, Eq. (11) gradually converges to Eq. (9), keeping Eq. (11) an unbiased estimate of the original loss function. Note that gradients are often used in curriculum learning to measure the difficulty of data [40, 4, 55, 41, 53], but our method introduces this idea to measure the difficulty of tasks by gradients to implement task-aware curriculum learning and adjust the training process accordingly.
The overall training procedure is shown in Algorithm 1. The entire model can be learned end-to-end by gradient based approaches. We update the learnable parameters  $w$ ,  $\alpha$ ,  $\theta$ , and  $\mathbf{S}$  in an iterative way. After training, we keep the architecture and parameters in the supernet for evaluation without the architecture discretization step, enhancing flexibility on architecture search and simplifying the optimization strategy.
Datasets. We adopt both synthetic and real-world multi-task graph datasets.