Deep Multi-Task Augmented Feature Learning via Hierarchical Graph Neural Network

The HGNN consists of two-level GNNs. The first-level GNN is an intra-task GNN to aggregate all the information contained in the data of a task to generate a task representation, which is called the task embedding. In the second level, based on the generated task embeddings in the first level, an inter-task GNN is used to update all the task embeddings by sharing information among all the tasks. Finally the task embeddings are used to augment the feature representation of the data to improve the learning performance. For classification tasks, we can learn augmented features in a fine granularity - the class level. That is, another intra-task GNN is used to aggregate all the information in a class of a task to generate a class embedding. Then based on class embeddings in all the tasks, another inter-class GNN is used to update them. Finally, both task embeddings and class embeddings are used to augment the feature representation. The whole architecture of HGNN is shown in Figure 1.

Suppose that there are $m$ multi-class classification tasks where each task has $k$ classes. The training dataset of the $i$th task consists of $n_{i}$ pairs of data samples and corresponding labels, i.e., $\mathcal{D}_{i}=\{(\mathbf{x}^{i}_{j},y^{i}_{j})\}_{j=1}^{n_{i}}$ where $y^{i}_{j}\in\{1,\ldots,k\}$.

For $\mathbf{x}^{i}_{j}$, we first define its hidden representation as

$$\hat{\mathbf{h}}^{i}_{j}=\sigma_{s}(\hat{\mathbf{W}}_{s}\mathbf{x}^{i}_{j}+\hat{\mathbf{b}}_{s}),$$

(1)

where $\sigma_{s}(\cdot)$ can be any activation function such as the ReLU function, and $\hat{\mathbf{W}}_{s},\hat{\mathbf{b}}_{s}$ are shared parameter among all the tasks.

For the intra-task GNN, we first construct an adjacency matrix $\mathbf{G}_{i}$ for the $i$th task based on the hidden representation and label information. Specifically, the $(j,l)$th entry in $\mathbf{G}_{i}$, $g^{i}_{jl}$, can be defined as

$$g^{i}_{jl}=\left\{\begin{array}[]{ll}\exp\{-\|\hat{\mathbf{h}}^{i}_{j}-\hat{\mathbf{h}}^{i}_{l}\|_{2}^{2}\}&\textrm{if $y^{i}_{j}=y^{i}_{l}$}\\ -\exp\{-\|\hat{\mathbf{h}}^{i}_{j}-\hat{\mathbf{h}}^{i}_{l}\|_{2}^{2}\}&\textrm{otherwise}\end{array}\right.,$$

where $\|\cdot\|_{2}$ denotes the $\ell_{2}$ norm of a vector. Then the intra-task GNN can be defined as

$$\mathbf{H}_{i}=\sigma_{h}(\mathbf{W}_{h}^{i}\mathbf{X}_{i}+\hat{\mathbf{H}}_{i}\mathbf{G}_{i}+\mathbf{b}_{h}^{i}\mathbf{1}),$$

(2)

where $\sigma_{h}(\cdot)$ can be any activation function such as the ReLU function, $\mathbf{X}_{i}=(\mathbf{x}^{i}_{1},\ldots,\mathbf{x}^{i}_{n_{i}})$, $\hat{\mathbf{H}}_{i}=(\hat{\mathbf{h}}^{i}_{1},\ldots,\hat{\mathbf{h}}^{i}_{n_{i}})$, $\mathbf{1}$ denotes a vector of all ones with an appropriate size, $\mathbf{W}_{h}$ and $\mathbf{b}_{h}$ are the parameters in the GNN. The $j$th column in $\mathbf{H}_{i}$ denoted by $\mathbf{h}^{i}_{j}$ is a new hidden representation for $\mathbf{x}^{i}_{j}$. $\mathbf{G}_{i}$ in Eq. (2) can make similar data points in the same class have similar representations in $\mathbf{H}_{i}$ and dissimilar data points from different classes have dissimilar representations. The intra-task GNN can have two or more layers each of which is defined as in Eq. (2).

Based on the intra-task GNN, the task embedding of the $i$th task is defined as

$$\mathbf{e}_{t}^{i}=\max_{l}\{\mathbf{h}^{i}_{l}\},$$

where the max operation is conducted elementwisely. So the task embedding is obtained via the max pooling on all the data points in the $i$th task based on the hidden representation learned by the intra-task GNN. Similarly, the class embedding of the $j$th class in $i$th task is defined as

$$\mathbf{e}_{c}^{i,j}=\max_{l:y^{i}_{l}=j}\{\mathbf{h}^{i}_{l}\},$$

which means that the class embedding of the $j$th class in the $i$th task is obtained via the max pooling on all the data points in the $j$th class of the $i$th task based on the intra-task GNN. We have tried other pooling methods such as the mean pooling but the performance is inferior to that of the max pooling. One reason is that the max pooling can bring some nonlinearity but the mean pooling is a linear operation.

The $m$ task embeddings $\{\mathbf{e}_{t}^{i}\}_{i=1}^{m}$ for the $m$ tasks can form a graph. The inter-task GNN is responsible of learning for the graph constructed by task embeddings $\{\mathbf{e}_{t}^{i}\}_{i=1}^{m}$ to generate new task embeddings $\{\hat{\mathbf{e}}_{t}^{i}\}_{i=1}^{m}$ by exchanging information among tasks. Here we use GAT as an implementation of the inter-task GNN.

In order to learn powerful task embeddings, each task embedding $\mathbf{e}_{t}^{i}$ is first transformed by a weight matrix $\mathbf{W}$. Then we perform self-attention on the task embeddings. That is, an attentional mechanism computes attention coefficients as

$$d_{ij}=a(\mathbf{W}\mathbf{e}_{t}^{i},\mathbf{We}_{t}^{j})$$

where the attentional mechanism $a(\cdot,\cdot)$ we use is the cosine function, which is different from the original GAT. To make coefficients comparable across different task embeddings, we normalize them over $j$ by using the softmax function as

$$\alpha_{ij}=\text{softmax}_{j}(d_{ij})=\frac{\exp(d_{ij})}{\sum_{k}\exp(d_{ik})}.$$

Attention values can be viewed as a measure of task relations between each pair of tasks. Once obtained, the normalized attention coefficients are used as combination coefficients to compute a linear combination of the new task embeddings before potentially applying a nonlinear activation function $\sigma$ as

$$\hat{\mathbf{e}}_{t}^{i}=\sigma\Big{(}\sum_{j=1}^{m}\alpha_{ij}\mathbf{W}\mathbf{e}_{t}^{j}\Big{)}.$$

Based on this equation, we can see that $\hat{\mathbf{e}}_{t}^{i}$ contains useful information from embeddings of other tasks. In experiments, the inter-task GNN adopts two such layers to generate the new task embeddings.

Similarly, the $mk$ class embeddings $\{\mathbf{e}_{c}^{i,j}\}$ also can form a graph. We uses another inter-class GNN to generate new class embeddings $\{\hat{\mathbf{e}}_{c}^{i,j}\}$ in a way similar to task embeddings.

The learned task embeddings and class embeddings can be used to augment the data feature representation to form a more expressive one as $\tilde{\mathbf{h}}^{i}_{l}=(\mathbf{\hat{h}}^{i}_{l},\hat{\mathbf{e}}_{t}^{i},\hat{\mathbf{e}}_{c}^{i,j})$, where $(\cdot,\cdot,\cdot)$ denotes the concatenation operation. Then data in such augmented representation can be fed into a deep multi-task learning model to learn class labels.

At the testing process, we do not know the true label, hence we cannot directly concatenate the class embedding to the hidden representation. We use the following method to solve this problem. For each testing sample, we concatenate the class embedding of each class $c$ to the hidden representation $\hat{\mathbf{h}}^{i}_{j}$ as its new hidden representation and then compute the prediction probability that the testing sample belongs to class $c$ via the softmax function used in the multi-task neural network. Finally we choose class $c$ with the largest prediction probability as the predicted label. In mathematics, we predict the class label of a testing sample as

$$c^{*}=\arg\max_{l\in[k]}\mathbb{P}(y=l|f((\mathbf{\hat{h}}^{i}_{*},\hat{\mathbf{e}}_{t}^{i},\hat{\mathbf{e}}_{c}^{i,l}))),$$

(3)

where $[k]$ denotes a set of positive integers no larger than $k$, $\mathbf{\hat{h}}^{i}_{*}$ denotes the transformed testing sample before feeding into the HGNN as shown in Eq. (1), and $f(\cdot)$ denotes the multi-task neural network used. Note that in the prediction rule (3), the concatenated class embedding $\hat{\mathbf{e}}_{c}^{i,l}$ changes with $l$.

For the regression problems, there are only continuous labels and we cannot define class embeddings. So we only use task embeddings as the augmenter feature. Furthermore, the adjacency matrix $\mathbf{G}_{i}$ for the $i$th task is different from classification tasks. Specifically, the $(j,l)$th entry in $\mathbf{G}_{i}$, $g^{i}_{jl}$, for a regression task is defined as

$$g^{i}_{jl}=\exp\{-\|\hat{\mathbf{h}}^{i}_{j}-\hat{\mathbf{h}}^{i}_{l}\|_{2}^{2}\}.$$

Since there is no class embedding, we do not need the prediction rule as in Eq. (3). The rest is identical to classification tasks.

We conduct experiments on several benchmark datasets, including ImageCLEF, Office-Caltech-10, Office-Home and SARCOS.

The ImageCLEF dataset is the benchmark for Image- CLEF domain adaptation challenge which contains about 2,400 images from 12 common categories shared by four tasks including Caltech-256 (C), ImageNet ILSVRC (I), Pascal VOC 2012 (P), and Bing (B). There are 50 images in each category and 600 images in each task.

The Office-Caltech-10 dataset includes 10 common categories shared by the Office-31 and Caltech-256 datasets. It contains four domains: Caltech (C) that is sampled from Caltech-256 dataset, Amazon (A) that contains images collected from the amazon website, Webcam (W) and DSLR (D) that are taken by the web camera and DSLR camera under the office environment. In our experiment, we regard each domain as a task.

The Office-Home dataset has 15,500 images across 65 classes in the office and home settings from four domains with a large domain discrepancy: Artistic images (Ar), Clip art (Cl), Product images (Pr), and Real-world images (Rw). In our experiment, we regard each domain as a task.

The SARCOS dataset studies a multi-output problem of learning the inverse dynamics of 7 SARCOS anthropomorphic robot arms, each of which corresponds to a task, based on 21 features, including seven joint positions, seven joint velocities, and seven joint accelerations. By following (Zhang & Yeung, 2010), we randomly sample 2000 data points from each output to construct the dataset.

Since the proposed HGNN can be combined with many deep multi-task learning models as discussed before, we incorporate the HGNN into the Deep Multi-Task Learning (DMTL) which shares the first several layers as the common hidden feature representation for all the tasks as did in (Caruana, 1997; Zhang et al., 2014; Liu et al., 2015; Zhang et al., 2015; Mrksic et al., 2015; Li et al., 2015), Deep Multi-Task Representation Learning (DMTRL) (Yang & Hospedales, 2017a), and Trace Norm Regularised Deep Multi-Task Learning (TNRMTL) (Yang & Hospedales, 2017b), respectively, to show the benefit of the learned augmented features.

In experiments, we leverage the VGG-19 network pretrained on the ImageNet dataset as the backbone of the feature generator $G$. After that, all the multi-task learning models adopt a two-layer fully-connected architecture (#data_dim $\times$ 600 $\times$ #classes) and the ReLU activation function is used. The first layer is shared by all tasks to learn a common representation and the second layer is for task-specific outputs. The HGNN learns 8-dimensional task embeddings and 8-dimensional class embeddings.

We use Adam with the learning rate varying as $\eta=\frac{0.02}{1+p}$, where $p$ is the number of the iteration. By following GAT, We fix $F^{\prime}=8$ in experiments. We adopt mini-batch SGD with $\text{batch\_size}=32$. Each experiment repeats for 5 times and we report the average performance as well as the standard deviation.

To study the effectiveness of task embeddings and class embeddings in the HGNN model, we study two variants of HGNN, including HGNN(T) that only augments with the task embedding and HGNN(C) that only augments with the class embedding. The comparison among baseline models, HGNN, variants of HGNN on the Office-Caltech-10 dataset is shown in Figure 6. According to the results, we can see that the use of only the class embedding in HGNN(C) or the task embedding in HGNN(T) can improve the performance of baseline models, which shows that augmented features learned in two ways are effective. HGNN(C) seems better than HGNN(T) in this experiment. One reason is that class embeddings may contain more discriminative features for the classification task. Figure 6 also indicates that using both task embeddings and class embeddings achieves the best performance, which again verifies the usefulness of the HGNN.

To dive deeper into the learned features, we plot in Figures 7 and 8 the t-SNE embeddings of the feature representations learned for the four tasks on Office-Caltech-10 dataset by TNRMTL and TNRMTL_HGNN, respectively, at the training and testing processes. We observe that the data based on the representation derived by the HGNN model are more separable among classes in each task during either the training process or the testing process. This phenomenon verifies the effectiveness of the augmented features learned in the HGNN to help discriminate data points in different classes of all the tasks.

We conduct the sensitivity analysis of the performance with respect to the dimension of task embedding ($F^{\prime}_{t}$) and class embedding ($F^{\prime}_{c}$), respectively, on the ImageCLEF dataset. The results are shown in Tables 1 and 2. According to the results, we can see that $F^{\prime}_{t}=8$ and $F^{\prime}_{c}=8$ are a good choice in most cases even though in some case, a lower value 4 performs better. When the dimension is not so large (e.g., not large than 32), the performance changes a little, making the choice of the dimension insensitive. However, when using a larger dimension (e.g., 64), the classification accuracy drops significantly, implying that the HGNN prefers a small dimension.

Proof. By setting the derivation of problems (4) and (5) to zero, these two problems have close form solutions as

	$\displaystyle\hat{\mathbf{W}}_{1}$	$\displaystyle=(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-1}\mathbf{X}\mathbf{Y}^{\intercal}$
	$\displaystyle\hat{\mathbf{W}}_{2}$	$\displaystyle=(\hat{\mathbf{X}}\hat{\mathbf{X}}^{\intercal}+\lambda\mathbf{I})^{-1}\hat{\mathbf{X}}\mathbf{Y}^{\intercal}$

Then

	$\displaystyle\mathbf{Y}-\hat{\mathbf{W}}_{1}^{\intercal}\mathbf{X}$	$\displaystyle=\mathbf{Y}-\mathbf{Y}\mathbf{X}^{\intercal}(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\mathbf{X}$
		$\displaystyle=\mathbf{Y}(\mathbf{I}-\mathbf{X}^{\intercal}(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\mathbf{X})$
		$\displaystyle\triangleq\mathbf{YA}$

	$\displaystyle\mathbf{Y}-\hat{\mathbf{W}}_{2}^{\intercal}\hat{\mathbf{X}}$	$\displaystyle=\mathbf{Y}-\mathbf{Y}\hat{\mathbf{X}}^{\intercal}(\hat{\mathbf{X}}\hat{\mathbf{X}}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\hat{\mathbf{X}}$
		$\displaystyle=\mathbf{Y}(\mathbf{I}-\hat{\mathbf{X}}^{\intercal}(\hat{\mathbf{X}}\hat{\mathbf{X}}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\hat{\mathbf{X}})$
		$\displaystyle\triangleq\mathbf{YB}.$

It is easy to show

$$(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{\intercal}\mathbf{X}=\mathbf{X}(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{\intercal}.$$

By left-multiplying by $(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-1}$ and right-multiplying by $(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-1}$, we can get

$$\mathbf{X}(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-\intercal}=(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\mathbf{X}.$$

So $\mathbf{A}$ can be simplified as

	$\displaystyle\mathbf{A}$	$\displaystyle=\mathbf{I}-\mathbf{X}^{\intercal}(\mathbf{X}\mathbf{X}^{\intercal}+\lambda\mathbf{I})^{-\intercal}\mathbf{X}$
		$\displaystyle=\mathbf{I}-\mathbf{X}^{\intercal}\mathbf{X}(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-\intercal}$
		$\displaystyle=\mathbf{I}-(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I}-\lambda\mathbf{I})(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-1}$
		$\displaystyle=\lambda(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-1}$

In order to prove

		$\displaystyle\|\mathbf{Y}-\hat{\mathbf{W}}_{1}^{\intercal}\mathbf{X}\|_{2}^{2}-\|\mathbf{Y}-\hat{\mathbf{W}}_{2}^{\intercal}\hat{\mathbf{X}}\|_{2}^{2}$
	$\displaystyle=$	$\displaystyle\|\mathbf{YA}\|_{2}^{2}-\|\mathbf{YB}\|_{2}^{2}$
	$\displaystyle\geq$	$\displaystyle 0,$

We need to require

$$\mathbf{A}^{\intercal}\mathbf{A}\succeq\mathbf{B}^{\intercal}\mathbf{B},$$

(7)

where $\mathbf{A}=\lambda(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})^{-1},\mathbf{B}=\lambda(\hat{\mathbf{X}}^{\intercal}\hat{\mathbf{X}}+\lambda\mathbf{I})^{-1}$. Note that

	$\displaystyle\mathbf{B}$	$\displaystyle=\lambda((\mathbf{X}^{\intercal},\mathbf{E}^{\intercal})(\mathbf{X}^{\intercal},\mathbf{E}^{\intercal})^{\intercal}+\lambda\mathbf{I})^{-1}$
		$\displaystyle=\lambda(\mathbf{X}^{\intercal}\mathbf{X}+\mathbf{E}^{\intercal}\mathbf{E}+\lambda\mathbf{I})^{-1}$

Eq. (7) is equivalent to

		$\displaystyle\big{(}(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})\big{)}^{-1}$
		$\displaystyle\succeq\big{(}(\mathbf{X}^{\intercal}\mathbf{X}+\mathbf{E}^{\intercal}\mathbf{E}+\lambda\mathbf{I})(\mathbf{X}^{\intercal}\mathbf{X}+\mathbf{E}^{\intercal}\mathbf{E}+\lambda\mathbf{I})\big{)}^{-1},$

which is also equivalent to

		$\displaystyle(\mathbf{X}^{\intercal}\mathbf{X}+\mathbf{E}^{\intercal}\mathbf{E}+\lambda\mathbf{I})(\mathbf{X}^{\intercal}\mathbf{X}+\mathbf{E}^{\intercal}\mathbf{E}+\lambda\mathbf{I})$
		$\displaystyle\succeq(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I})(\mathbf{X}^{\intercal}\mathbf{X}+\lambda\mathbf{I}).$

So we require

$\displaystyle\mathbf{X}^{\intercal}\mathbf{X}\mathbf{E}^{\intercal}\mathbf{E}+\mathbf{E}^{\intercal}\mathbf{E}\mathbf{X}^{\intercal}\mathbf{X}+2\lambda\mathbf{E}^{\intercal}\mathbf{E}+\mathbf{E}^{\intercal}\mathbf{E}\mathbf{E}^{\intercal}\mathbf{E}\succeq 0$

which is the condition. $\Box$

Before proving Theorem 2, we introduce a lemma from (Kakade et al., 2009).

Then we can give a proof of Theorem 2.

Proof. Since the loss function has a Lipschitz constant 2, according to Lemma 1, we let $p=q=2$ and then we can reach the conclusion. $\Box$