Fig 1: Box embeddings of 150 datasets of iNaturalist + CUB and corresponding learned hierarchy for Class Arachnida. Each taxonomic category is treated as a separate dataset for which Task2Box embeddings are learned. (1) Shows the learned box embeddings where datasets from the same group (taxonomic class) have the same color. Datasets naturally cluster to their ground truth groups. (2) Shows the hierarchy learned through Task2Box for a specific class. The hierarchy matches the ground truth relationships based on biological classification. Orders that belong to Class Arachnida are learned as boxes (A, B, C) contained by the larger box for Arachnida; families under each of the orders are learned as smaller boxes contained by the corresponding orders they belong to.

Overview

Modeling and visualizing relationships between tasks or datasets is an important step towards solving various meta-tasks such as dataset discovery, multi-tasking, and transfer learning. However, many relationships (e.g., containment or transferability) are naturally asymmetric and current approaches for representation and visualization (e.g., t-SNE) do not readily support this. We propose Task2Box, an approach to represent tasks using box embeddings—axis-aligned hyperrectangles in low dimensional spaces–that can capture asymmetric relationships between tasks through volumetric overlaps. We show that Task2Box accurately predicts unseen hierarchical relationships between nodes in ImageNet and iNaturalist datasets, as well as transferability between tasks in the Taskonomy benchmark. We also show that box embeddings estimated from task representations (e.g., CLIP, Task2Vec, or attribute based ) can be used to predict relationships between unseen tasks more accurately than classifiers trained on the same representations, as well as handcrafted asymmetric distances (e.g., KL divergence). This suggests that low-dimensional box embeddings can effectively capture these task relationships and have the added advantage of being interpretable. We use the approach to visualize task relationships, and model image classification datasets on Hugging Face.

Key Idea

The success of deep learning has led to the proliferation of datasets for solving a wide range of computer vision problems. Yet, there are few tools available to enable practitioners to find datasets related to the task at hand, and to solve various meta-tasks related to it. We present Task2Box, a method to represent tasks using axis-aligned hyperrectangles (or box embeddings). Task2Box is framed as a learnable mapping from dataset representation to boxes, and can be trained to predict various relationships between novel tasks such as transferability, hierarchy, and overlap.

We test our framework to model relationships between nodes in iNaturalist, CUB and ImageNet datasets, as well as to predict transferability on the Taskonomy benchmark. Table 1 and 2 show that low-dimensional box embeddings accurately predicts novel relationships between datasets seen during training, as well as relationships with novel datasets. Remarkably, Task2Box outperforms classifiers trained to directly predict the relationships on the same representations, suggesting that the box embedding provides a good inductive bias for learning hierarchical relationships. We also outperform simple asymmetric distances proposed in prior work such as Kullback-Leibler (KL) divergence. To model the heterogeneous tasks in the Taskonomy benchmark we map each task to a set of attributes from which a box embedding is learned. Once again, we obtain significantly higher correlation between the true and predicted transferability for both existing and novel datasets compared to standard classifiers (Table 3). Such attribute-based representations can be readily derived from datasheets and modelcards.

Finally, the low-dimensional box embeddings have the added advantage of being interpretable. Fig. 1 and Fig. 2 show relationships on the iNaturalist+CUB and ImageNet categories, respectively. The 2D box representation allows us to readily visualize the strength and direction of task relationships based on the overlapping volumes, which is not possible using Euclidean representation (e.g., t-SNE). At the same time, new datasets can be embedded in constant time without needing to retrain or reoptimize. We use the framework to visualize overlap relationships among 131 publicly available datasets on Hugging Face (a platform for hosting datasets) in Fig. 3.

Table 1: Average F1 Score for predicting hierarchical relationships on iNaturalist + CUB.

Table 2: Average F1 Score for predicting hierarchical relationships on ImageNet instruments.

Table 3: Spearman correlation and standard deviation between the predicted and ground truth task affinities from Taskonomy. Our method shows higher correlation with the task affinities compared to other baseline methods.

Fig 2: Visualization of learned embeddings on instrument-related datasets in ImageNet. Datasets that belong to the same superset are shaded in the same color. Task2Box learns the hierarchy of various groups, and clusters similar datasets closer to each other.

Fig 3: Visualization of relationships between public image classification datasets in Hugging Face. The sample data points annotated on the highlighted datasets show that common tasks overlap with each other (e.g., sentiment classification and document classification datasets). Although labels could slightly differ between datasets, Task2Box can infer the level of similarity and represent it as the amount of overlap. The embedding size (box area) also shows the number of available data.

Code

The code is available at https://github.com/cvl-umass/task2box

BibTeX

@misc{daroya2024task2box,
      title={Task2Box: Box Embeddings for Modeling Asymmetric Task Relationships}, 
      author={Rangel Daroya and Aaron Sun and Subhransu Maji},
      year={2024},
      eprint={2403.17173},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Acknowledgements

This work was supported by awards from the National Science Foundation (2329927 and 1749833) and the NASA AIST program. The experiments were performed on the University of Massachusetts GPU cluster funded by the Mass. Technology Collaborative.