Association Graph Learning for Multi-Task Classification with Category Shifts
Abstract
In this paper, we focus on multi-task classification, where related classificationtasks share the same label space and are learned simultaneously.
In particular,we tackle a new setting, which is more realistic than currently addressed in theliterature, where categories shift from training to test data.
Hence, individual tasksdo not contain complete training data for the categories in the test set.
To generalizeto such test data, it is crucial for individual tasks to leverage knowledge fromrelated tasks.
To this end, we propose learning an association graph to transferknowledge among tasks for missing classes.
We construct the association graphwith nodes representing tasks, classes and instances, and encode the relationshipsamong the nodes in the edges to guide the knowledge transfer between them.
Bymessage passing on the association graph, our model enhances the categoricalinformation of each instance, making it more discriminative.
To avoid spuriouscorrelations between task and class nodes in the graph, we introduce an assignmententropy maximization that encourages each class node to balance its edge weights.
This enables all tasks to fully utilize the categorical information from related tasks.
An extensive evaluation on three general benchmarks and a medical dataset forskin lesion classification reveals that our method consistently performs better thanrepresentative baselines.
Introduction
Multi-task learning aims to simultaneously solve several related tasks [7, 49] by sharing informationand has attracted much attention in recent years.
In this paper, we focus on multi-task classificationunder the multi-input multi-output setting [50, 30, 37, 49].
In this setting, the label space is sharedbut each task operates on a different type of visual modality or the same modality collected fromdifferent environments or equipment.
As a consequence, each task follows different data distributionsfor the shared label.
The intuition behind multi-task classification is that tasks having the same labelspace provide partial knowledge of the distributions that can be shared among all tasks to reach abetter view of the full distribution which in turn benefits the individual tasks.
A real-world challenge for multi-task classification is category shift, where the categories in the testingphase are shared but during training not all classes are present for each individual task.
This challenge is common in various realistic scenarios, such as skin lesion classification [25, 47], fault diagnosis
[43, 10], or remote sensing scene classification [31].
For example, in skin lesion classification, data provided by different hospitals or healthcare facilities should lead to the same set of diagnoses [25, 47].
Unfortunately, due to patient populations or proprietary use regulations, these tasks do not share
Artistic
Clipart
Product
Real_world ? ? ? ? ? ? ? ? ? ? ? ?
Multi-Task Classification
Multi-Task Classification with Category Shifts ? ? ? ?
0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3, 2,3,4,5, 4,5,6,7, 0,1,6,7, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6, 0,1,2,3,4,5,6,7, 0,1,2,3,4,5,6,7, 0,1,2,3,4,5,6,7, 0,1,2,3,4,5,6,7, the same diagnosis categories at training time.
For example, some institutions miss instances frommelanoma and basal cell carcinoma while others lack dermatofibroma and benign keratosis [47].
Inthis case, it is beneficial to expand the diagnostic scope of different hospitals or healthcare facilities bysimultaneously learning their training data and improving the overall prediction accuracy.
Motivatedby these realistic scenarios, we propose a new multi-task setting, namely multi-task classificationwith category shifts.
Figure 1 shows comparisons between multi-task classification with and withoutcategory shifts.
The goal of the proposed setting is to explore task-relatedness with an incompletetraining label space to improve the generalization ability of the categorical information at test time.
To deal with category shifts, we propose to learn an association graph to transfer knowledge amongtasks.
The association graph is constructed over three different types of nodes: task, class and instancenodes.
To model the complex relationships among these heterogeneous nodes, we apply differentlearnable metric functions to construct the edge weights among different types of nodes.
To propagateknowledge between the nodes, we apply message passing to update each node in the association graphaccording to their relationships.
Essentially, the association graph stores the task and class-specificknowledge during training time and enhances each instance feature by associating it with other taskand class nodes during inference.
In the constructed association graph, the relationships betweenclass and task nodes tend to be biased due to the category shifts.
This hinders tasks in utilizing thecategorical information for missing classes.
To avoid these spurious correlations between task andclass nodes, we introduce assignment entropy maximization.
This regularization of the associationgraph encourages each class node to balance its edge weights with all tasks, enabling them to fullyutilize the categorical information.
We evaluate our model on three multi-task classification benchmarks and a medical dataset for skinlesion to demonstrate that the proposed model performs better in solving category shifts.
We alsoprovide detailed analyses to show how the proposed association graph enhances the categoricalinformation of each instance with transferred knowledge.
Problem Statement
We first formally introduce the new problem setting of multi-task classification with category shifts.
We consider T related classification tasks {D t } Tt =1 .
Each task contains a training set D trt and test set
D tet .
We define D trt = {X trt , Y trt } , where X trt denotes the set of training data from the t -th task and
Y trt is the corresponding label space in the training phase and t ∈ { 1 , 2 , ..., T } .
Likewise, we define
D tet = {X tet , Y tet } .
In addition, we define the entire label space of all tasks as Y .
The conventionalmulti-task classification setting, e.g., [30, 37], is a specific instantiation of our setting where all tasksshare the entire label space during both training and test time
Y trt = Y tet = Y .
Definition 1 (Category Shifts in Multi-Task Classification) .
For each task, the training label space isa subset of the test label space Y trt ⊂ Y tet , where t ∈ { 1 , 2 , .., T } .
The union of the training labelspaces of all tasks Y =
(cid:83)
Tt
=1 Y trt determines the label space for the test phase.
We provide a visual illustration of the difference between multi-task classification without and withcategory shifts in Figure 1.
The goal of the proposed setting is to explore task-relatedness in thepresence of missing classes to improve the generalization ability of the categorical information.
Tostudy the impact of different degrees of category shifts, we introduce the missing rate γ that formallymeasures the degree of category shifts, with higher missing rates yielding more severe category shifts,and therefore more challenging for each task.
Definition 2 (Missing Rate) .
Given T related classification tasks with entire label space Y , themissing rate γ is the average rate of the number of missing classes with respect to the size of theentire label space, γ = 1 T (cid:80) Tt
=1 ( |Y|−|Y trt | |Y| ) .
Having defined the problem setting, we are now ready to present the first multi-task classificationmethod tailored to handle category shifts.
Methodology3.1 Learning the association graph for knowledge transfer
To deal with category shifts, we propose to learn an association graph to transfer knowledge amongtasks for each class.
We construct an undirected graph over three types of nodes: task, class andinstance nodes.
The association graph stores the task and class specific knowledge from the trainingdata in the task and class nodes.
The edges encode the relationships between the nodes, enabling therelevant knowledge to be transferred to each instance node.
Since the nodes are heterogeneous, weapply different learnable metric functions to compute edge weights among different types of nodes.
To better explain the construction, we provide an illustration of the association graph in Figure 2.
The rationale behind our model is that relevant knowledge is transferred among nodes in the association graph to enhance the categorical information of each instance, making the instance morediscriminative.
During training, the model learns the ability to update all nodes in the graph by transferring the relevant knowledge.
At inference time, the ability is generalized to each test instance ofeither observed or missing categories.
Thus, each test instance is refined with the relevant knowledgestored in the association graph, which reduces the category shifts from the training to the test set.
Forclarity, we introduce the main components of the association graph: task and class graphs.
Task graph Given T related tasks, we introduce the task graph to model the relationships betweentasks.
The features of the nodes in the task graph are task-specific representations, each of whichaggregates all features from the corresponding task.
We define the node of t -th task as follows:
[Equation]
where x i is a training instance belonging to the t -th task and N trt is the number of training instancesfrom the corresponding task.
E is a feature extractor shared by all tasks, which embeds each instanceinto a d -dimensional feature space. E is a feature extractor shared by all tasks, which embeds eachinstance into a d -dimensional feature space.
With the task nodes, we further define the edges between the task nodes v i and v j as A T ( v i , v j ) .
The weight of the edge is determined by the learnable similarity between the task nodes, which isformulated as follows:
[Equation]
where W T and b T are the learnable parameters for the task graph.
α T is a scalar and σ is the sigmoidfunction that is used to normalize the edge weight between 0 and 1.
The edge weight indicatesthe proximity between the i -th and j -th tasks.
We denote the task graph as G T =( V T , A T ) .
In thetask graph, V T = { v t | t ∈ [1 , T ] } ∈ R T ×
d is the set of all task nodes and A T =
{ A T ( v i , v j ) | i, j ∈[1 , T ] } ∈ R T × T is the corresponding adjacency matrix, which characterizes task relationships.
Class graph Likewise, we define the class graph as G C =( V C , A C ) , where each node representsthe corresponding categorical information.
Here V C
=
{ k c | c ∈ [1 , C ]
} ∈ R C
×
d , where C denotesthe size of the entire label space of all tasks.
We define the features of the node of the c -th class as:
[Equation]
where x m is a training instance from c -th class and and N trc is the number of training instances ofthe corresponding class.
The edge weight between class nodes A C ( k i , k j ) is formulated as:
[Equation]
where W C and b C are the learnable parameters for the class graph.
α C is a fixed scalar.
The adjacencymatrix of the class graph is A C
= { A C
( k i , k j ) | i, j ∈ [1 , C ] } ∈ R C × C .
Association graph Having defined the task and class graphs, we build the association graph byconnecting both graphs with each instance node as shown in Figure 2.
We use the association graphto enhance instance feature learning by leveraging task representations (nodes in the task graph) andcategorical information (nodes in the class graph).
To do so, we query each instance in the task andclass graphs.
Formally, we define an instance node as V X =
{E ( x ) } ∈ R 1 × d .
We build up the connection of the instance node to each node of the task graph and the classgraph, which gives the instance node access to the knowledge stored in the task and class graph.
The edge between the instance node E ( x ) and a task node v t is denoted by A X∼T ( x , v t ) .
Theweight of the edge is obtained by the normalized similarity between the instance and task node,
A X∼T =softmax( E ( x ) ⊤ v t √ d ) .
The corresponding adjacency matrix is A X∼T
= { A X∼T ( x , v t ) | t ∈[1 , T ] } ∈ R 1 × T .
Likewise, we define the adjacency matrix between instance nodes and the classgraph as A X∼C = { A X∼C ( x , k c ) | c ∈ [1 , C ] } ∈ R 1 × C .
In the association graph, we further build theedges cross task and class nodes to enable knowledge transfer between them.
Formally, we define theedge connecting a class node and a task node as:
[Equation]
where α P is a fixed scaling factor.
The adjacency matrix between the task and class graphs is denotedby A C∼T =
{ A C∼T ( k c , v t ) | c ∈ [1 , C ] , t ∈ [1 , T ] } ∈ R C × T .
Thus, the whole association graph G =( V , A ) with three types of nodes can be formulated as: [Equation]
[Equation]
Knowledge transfer by message passing With the constructed association graph, we performmessage passing in the association graph via a multi-layer Graph Neural Network (GNN).
In general,our model can work with various GNN architectures.
In this work, we apply GraphSAGE [18].
Tosimplify, we use h i to represent one node in the association graph G , which could be a task, class orinstance node.
The l -th layer of GNNs can be written as:
[Equation]
where h ( l ) i denotes the node embedding by the l -th GNN layer and U l and W l are learnable weightmatrices of the l -th GNN layer.
l ∈ { 1 , 2 , ..., L } , where L denotes the number of the GNN layers.h (0) i is initialized as the nodes defined above.
N k
( h i ) denotes the top- k neighbors of the node
h i .
Through message passing on the graph by the stacked GNN layers, each node is updated with theknowledge transferred from its neighborhoods.
By doing so, we obtain the enhanced instance featureas ˆ V X
= {
ˆ E ( x ) } .
We make the prediction for the enhanced feature as p ( y | x )=
p
( y | ˆ E ( x ) , f t ) , wheref t denotes the corresponding task-specific classifier.
3.2 Assignment entropy maximization
Category shifts in multi-task classification yield spurious correlations between tasks and their corresponding observed classes.
This means the edges cross tasks and their observed classes havesignificantly higher weights than other edges.
The knowledge transfer between tasks and classeswill be dominated by this spurious correlations, hindering their missing classes from exploiting thecategorical information.
To tackle this problem, we propose assignment entropy maximization to encourage each class nodeto balance the weights of its edges with all tasks.
Formally, the assignment entropy for the c -th classis formulated as:
[Equation]
By maximizing the assignment entropy, weights are balanced for each class, which enables all tasksto fully utilize the categorical information for their missing classes.
Intuitively, each class node istask-agnostic when the assignment entropy reaches its maximum values.
In this ideal case, the modeleliminates the spurious correlations between the class and its corresponding tasks in the graph.
By combining the assignment entropy maximization and cross-entropy minimization of the classifiers,we have the final objective as follows:
[Equation]
where L CE =
E D t [ − log p ( D t |G )] and L AE ( k c )
= − H ( k c ) .
β is introduced to balance the importanceof the cross-entropy and assignment entropy losses.
We provide the training and inference algorithmsin the supplemental materials.
Related Works
Multi-task learning [7] aims to learn several related tasks simultaneously and improve their overallperformance.
The task relatedness is learned by many different aspects of the model, e.g., the lossfunctions [27, 23], gradient space [36, 48], parameter space [30, 3], or representation space [2, 32, 17].
The basic idea of sharing information from multiple tasks has been successfully applied underdifferent settings, including the single-input multi-output setting [36, 48, 38], the multi-input multioutput setting [2, 30, 37], and using meta-learning [12, 42, 1].
In the single-input multi-outputsetting [19, 11], tasks are defined by different supervision information included in the same input.
In meta-learning, tasks are sampled from one task distribution and different tasks have different category spaces [42].
In multi-input multi-output, tasks follow different data distributions since theyare collected from different visual modalities or equipment [30, 37, 49, 50].
Nevertheless, it remainsunexplored to investigate category shifts in multi-task classification, which is common in variousrealistic scenarios.
Category shifts denotes that training data collected from different domains may not completelyshare their categories, which is first proposed by [44] for domain adaptation.
[29] and [34] challengedomain generalization with category shift, where the change of domains is always followed by thechange of categories.
Category shift is a common scenario in real-world applications, since it relaxesthe requirement on the shared category set among any source domains [44, 8].
As a result, categoryshift has drawn increasing attention and has been applied to a wide range of learning tasks, includingfault diagnosis [10], skin lesion classification [25] and remote sensing image classification [31].
Inthis paper, we develop a new multi-task learning scenario, in which individual tasks do not containcomplete training data for the categories in the test set.
Unlike domain adaptation and generalizationonly focusing on the unidirectional knowledge transfer from source domains to a target domain,our multi-task classification encourages simultaneous bidirectional knowledge transfer between anypaired domains to enable efficient predictions for both tasks.
To the best of our knowledge, we arethe first to address category shifts in multi-task classification.
Exploring graph structure for classification Several recent works prove that graphs are effective inmodeling label correlation [9, 26, 24] or task relationships [46, 6, 20]. For multi-label classification,[26] formulates the multi-label predictions as a conditional graphical lasso inference problem, while[9] utilizes a graph convolutional network to propagate information between multiple labels andconsequently learn inter-dependent classifiers for each of the image labels.
For multi-task learning,some work learns the relationship between multiple tasks by message passing over a graph neuralnetwork [28, 16].
[33, 5] explore the graph structure to produce higher quality node embedding on thegraph structured data.
For few-shot learning, [45] and [46] design the hand-crafted and automaticallyconstructed meta-knowledge graph to provide meta knowledge for each task.
Different from thesemethods, we address the new challenge, category shifts in multi-task classification, which yields anovel graph construction.
Experiments and Results
Datasets We conduct experiments on three common multi-task classification benchmarks and a skinlesion classification dataset to evaluate the effectiveness of our proposed method. Office-Home [41] contains images from four domains/tasks: Artistic, Clipart, Product and Realworld.
Each task contains images from 65 object categories collected under office and home settings.
There are about 15 , 500 images in total.
Office-Caltech [13] contains the ten categories shared between Office-31 [35] and Caltech-256[14]. One task uses data from Caltech-256, and the other three tasks use data from Office-31, whoseimages were collected from three distinct domains/tasks, namely Amazon, Webcam and DSLR. Thereare 8 ∼ 151 samples per category per task, and 2 , 533 images in total. ImageCLEF [30], the benchmark for the ImageCLEF domain adaptation challenge, contains 12common categories shared by four public datasets/tasks:
Caltech-256, ImageNet ILSVRC 2012,
Pascal VOC 2012, and Bing.
There are 2 , 400 images in total.
Skin-Lesion contains three skin lesion classification tasks: HAM10000 [39], Dermofit [4] and
Derm7pt [22].
Tasks are collected from different hospitals or healthcare facilities.
In this dataset, eachtask contains a subset of the following classes: melanocytic nevus, melanoma, basal cell carcinoma,dermatofibroma, benign keratosis and vascular lesion.
Experimental setup We explore the effect of varying degrees of category shifts in these datasetsby different missing rates in the training label spaces.
For the three common multi-classificationbenchmarks, we set the missing rates for the three benchmarks as 75%, 50%, 25%, 0%, denoting that each task cannot access the training data of 75% (or 50%, 25%, 0%) categories.
For simplicity,we use the same missing rate for all tasks in each setting, which is achieved by assigning the samenumber of missing classes for each task.
Since the union of the training label space of all tasksequals the test label space in the proposed setting, 75% is the largest missing rate for the commonmulti-task classification datasets that have data from four tasks.
With the 75% missing rate, datafrom each category is only accessible in one task.
By contrast, when the missing rate is set to 0%,the problem degrades to the conventional multi-task classification without category shifts, i.e., eachtask has complete training data from all classes.
Since Skin-Lesion has three tasks, we set themissing rates as 67%, 33%, 0%.
For a fair comparison, the assignment of missing classes for differentmissing rates and datasets are shared for all methods.
We use ResNet-18 [21] as the backbone for allexperiments and deploy the graph in the output space of the backbone.
We provide the code and themissing class assignment of each dataset in the supplemental materials. Metrics
The average multi-task classification accuracy ( % , top-1) along with 95% confidenceintervals from five runs are reported across all tasks.
In order to evaluate the model’s generalizationfrom observed classes to missing classes of each task, test instances come from observed classes andmissing classes.
We report the average accuracy of missing classes and observed classes of all tasksas A m and A o , respectively.
Moreover, we apply the harmonic mean to show the overall performanceon both missing and observed classes, which is denoted by H =
× A m × A o A
m + A
o .
and proposed association graph. Our association graph consistently performs better than the attention graph with different message passing layers.
Benefit of the association graph To show thebenefit of the proposed association graph, weconduct experiments with increasing numbers ofthe message passing layers, where L =0 denotesthe model without the graph and knowledgetransferring from graph nodes.
As shown in Figure 3, we report the harmonic mean of the modelwith different numbers of the message passinglayers.
As the number L increases, our modelperforms increasingly better with the peak at
L =4 , surpassing the model without graphs bya large margin.
The results demonstrate thatthe task and class-specific knowledge stored inthe graph is important to enhance the categorical information in the features, which enablesthem to be more discriminative.
In the followingexperiments, we set L =4 for our model.
To better understand the benefit of the association graph in feature learning, we visualize the distribution of samples from observed and missing classes of each class on Office-Home .
Figure 4shows that the association graph reduces the overlap of the distributions of missing classes (shapesin red) and updates the features of each class to be more clustered.
We therefore conclude that theassociation graph enhance the categorical information of the instance features from both observedclasses and missing classes, making them more distinguishable.
Association graph vs. attention graph We compare the association graph with the attention graph[40], which incorporates the same self-attention strategy for all nodes.
Different from the attentiongraph, the proposed association graph adopts different learnable metric functions for different typesof nodes.
In Figure 3, the association graph performs consistently better than the attention graph withdifferent numbers of message passing layers.
The association graph with different learnable metricfunctions is suitable for modeling complex relationships among different types of nodes, leading toimprovements over the attention graph.
Benefit of assignment entropy maximization We investigate the benefit of the proposed assignmententropy maximization in our model.
As shown in Table 1, the assignment entropy maximizationregularization considerably improves the overall performance.
This is reasonable since the assignment [Figure] [Table] entropy maximization balances the edge weights between each class and all task nodes.
Thus, thespurious correlation between task and class nodes is reduced, which enables each task to fully utilizethe categorical information provided by other tasks.
Importance of knowledge transfer by message passing To investigate the importance ofknowledge transfer, we conduct experiments on Office-Home with different neighbor sizes foreach node during message passing, which reflects the amount of the transferred knowledge.
[Table]
With the size of 0, each node does not utilizethe transferred knowledge.
With the maximumsize (which is 70 in this dataset), each nodeaggregates the knowledge from all other nodeson the graph.
As shown in Table 2, we find thatthe best performance happens with the largestneighbor size, which indicates the importance ofpassing messages throughout the whole graphfor knowledge transfer.
Effect of different degrees of category shifts We evaluate the proposed method on three multi-taskclassification benchmarks with different missing rates in Table 3.
Our method achieves the bestoverall performance on all three benchmarks under each missing rate in terms of the harmonic meanof the accuracy of observed and missing classes.
We note that 75% is the most severe categoryshifts in multi-task classification, which demonstrates there are no overlap categories between tasks.
On Office-Home with the 75% missing rate, our model surpasses the second best method, i.e., WeighLosses [23], by a large margin of 10 .
33% , in terms of the harmonic mean.
The consistentimprovements on all benchmarks with different missing rates demonstrate that the association graphis effective in addressing the category shifts for multi-task classification.
More detailed results with95% confidence intervals are provided in Appendix.
Moreover, we also provide the results on therealistic medical dataset, i.e., Skin-Lesion with missing rates of 67% , 33% and 0% in Table 4.
Ourmodel achieves the best overall performance, which again confirms the effectiveness of our method.
[Table] [Table] It is worth mentioning that with γ =0% the setting reduces to traditional multi-task classificationwithout category shifts, where all classes are observed by each task during training.
In this case,the results of missing classes A m and the harmonic mean H are not available.
In this setting, ourmethod still outperforms other baselines in all datasets.
We conclude that our association graph canbetter utilize the shared knowledge to improve overall performance for all tasks, which also holds forsettings without category shifts.
Conclusion
We address category shifts in multi-task classification, which is challenging yet more realistic sinceindividual tasks do not contain complete training data for the categories in the test set.
To tackle this,we propose to learn an association graph to transfer knowledge among tasks for missing classes, whichenhances the categorical information of each instance, making them more discriminative.
To avoid thespurious correlations between task and class nodes, we introduce assignment entropy maximization,enabling all tasks to fully utilize the categorical information for missing classes.
To the best of ourknowledge, we are the first to address the challenge in multi-task classification.
We conduct ablationstudies to demonstrate the effectiveness of the proposed association graph and assignment entropymaximization in our model.
The superior performance on three multi-task classification benchmarksand the medical dataset for skin lesion classification further substantiates the benefits of our proposal.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[No]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]