Task Prototype-Based Knowledge Retrieval for Multi-Task Learning from Partially Annotated Data

TL;DR

The paper tackles robust multi-task learning under partial annotations by introducing a prototype-based knowledge retrieval framework. It jointly learns a task prototype that encodes task-specific characteristics and a knowledge retrieval transformer that uses task-affinity guidance to adaptively refine feature representations, regulated by an association knowledge generating ($L_{akg}$) loss. Empirically, the method achieves state-of-the-art performance on PASCAL-Context and NYUD-v2 under one-label and random-label configurations, with ablations confirming the benefit of vector quantization, $L_{akg}$ components, and task-consistency. This approach enables reliable cross-task transfer without depending on predictions from unlabeled tasks, enhancing robustness and practical applicability for partially annotated real-world scenarios.

Abstract

Multi-task learning (MTL) is critical in real-world applications such as autonomous driving and robotics, enabling simultaneous handling of diverse tasks. However, obtaining fully annotated data for all tasks is impractical due to labeling costs. Existing methods for partially labeled MTL typically rely on predictions from unlabeled tasks, making it difficult to establish reliable task associations and potentially leading to negative transfer and suboptimal performance. To address these issues, we propose a prototype-based knowledge retrieval framework that achieves robust MTL instead of relying on predictions from unlabeled tasks. Our framework consists of two key components: (1) a task prototype embedding task-specific characteristics and quantifying task associations, and (2) a knowledge retrieval transformer that adaptively refines feature representations based on these associations. To achieve this, we introduce an association knowledge generating (AKG) loss to ensure the task prototype consistently captures task-specific characteristics. Extensive experiments demonstrate the effectiveness of our framework, highlighting its potential for robust multi-task learning, even when only a subset of tasks is annotated.

Figures (7)

• Figure 1: Conceptual diagram of the proposed method (semantic segmentation example). (1) Task prototype generates task-affinity score, and (2) prototype-based knowledge retrieval process utilizes this task-affinity to adaptively enhance task performance.
• Figure 2: Overview of the proposed multi-task learning with prototype-based knowledge retrieval. It consists of two main components: (1) multi-task learning module, and (2) prototype-based knowledge retrieval module using task prototype and knowledge retrieval transformer.
• Figure 3: (a) Illustration of the training process for embedding task-specific characteristics into the task prototype. (b) Illustration of the prototype-based knowledge retrieval process.
• Figure 4: Illustration of the Task Consistency (TC) loss.
• Figure 5: Qualitative comparison of our method and state-of-the-art method on PASCAL-Context under one-label setting.