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

AI summary

Overview

Research area: Multi-task learning (MTL) for computer vision, specifically Multi-Task Partially Supervised Learning (MTPSL), where only a subset of tasks is annotated.

Technical level: Advanced. The paper assumes familiarity with transformer attention mechanisms, cross-task regularization, vector quantization (codebooks), and dense prediction benchmarks.

Scope: The paper proposes a framework that replaces reliance on predictions from unlabeled tasks with a learned task prototype plus a knowledge retrieval transformer, evaluated on PASCAL-Context and NYUD-v2.

What This Paper Is About

In multi-task learning, training one model to do several visual tasks at once normally requires every image to be labeled for every task, which is expensive and often impractical. Partial-label settings exist, but the dominant strategy is to use the model's own predictions on unlabeled tasks as a stand-in for labels, and those predictions are noisy and incomplete, which can cause negative transfer. This paper's goal is to instead model task relationships from task-inherent characteristics stored in a prototype, so knowledge can be transferred to the target task without touching unlabeled-task predictions at all.

Key Contributions

1. A task prototype that captures task-specific characteristics and measures the degree of enhancement each task needs through quantified task associations.
2. A knowledge retrieval transformer that uses a task-affinity score to adaptively refine feature representations and align them with the target task's requirements.
3. A prototype-based knowledge retrieval learning method that leverages task-specific characteristics instead of predictions from unlabeled tasks, improving performance across tasks when annotations are missing.
4. Supporting loss design: an association knowledge generating (AKG) loss, composed of the task knowledge embedding (TKE) loss and the task consistency (TC) loss, plus a task-agnostic enhancement (TAE) loss over a vector-quantized codebook.

Main Findings

• PASCAL-Context, one-label setting (Table 1): the proposed method reaches 59.78 mIoU on semantic segmentation, 59.08 mIoU on human parsing, 78.62 maxF on saliency, 15.63 mErr on surface normal, and 65.10 odsF on boundary, compared with DiffusionMTL (Prediction, CVPR'24) at 59.43 / 56.79 / 77.57 / 16.20 / 64.00 and DiffusionMTL (Feature) at 57.78 / 58.98 / 77.82 / 16.11 / 64.50.
• PASCAL-Context, random-label setting (Table 1): the proposed method reports 64.30 mIoU semantic segmentation, 56.87 mIoU parsing, 80.51 maxF saliency, 14.48 mErr normal, and 67.30 odsF boundary. The paper notes DiffusionMTL's performance varied with whether the diffusion input was a prediction map or a feature map, whereas the proposed method stays consistent.
• NYUD-v2, one-label setting (Table 2): 45.95 mIoU semantic segmentation, 0.4865 absErr depth, 25.64 mErr normal, versus DiffusionMTL (Feature) at 44.47 / 0.5059 / 25.84.
• NYUD-v2, random-label setting (Table 2): 47.53 mIoU semantic segmentation, 0.4621 absErr depth, 24.67 mErr normal.
• Ablation on losses (Table 3, PASCAL-Context one-label): the baseline without any of the proposed losses reports 44.73 mIoU semantic segmentation, 57.03 mIoU parsing, 75.69 maxF saliency, 16.47 mErr normal, 64.38 odsF boundary. Adding the TAE loss gives 44.83 / 57.13 / 76.13 / 16.22 / 64.50. Adding TKE yields a large jump to 58.21 / 58.87 / 78.50 / 15.67 / 65.00. Adding the TC loss gives the final 59.78 / 59.08 / 78.62 / 15.63 / 65.10.
• Prototype slot dimension (Table 4, NYUD-v2 one-label): performance rises from 42.77 mIoU at the "-" configuration (146.5M parameters) to 44.91 at 256 dimensions (156.6M), 45.65 at 512 (157.5M), peaks at 45.95 at 1024 (159.4M), then drops to 45.33 at 2048 (163.0M). Optimal performance was achieved at 1024.
• Comparison with prompt-based methods (Table 5, NYUD-v2 one-label, ViT-L backbone): the proposed explicit-learning method reports 50.08 mIoU semantic segmentation, 0.3857 absErr depth, and 20.57 mErr normal, compared with the 45.41 / 0.4277 / 22.34 baseline, 48.68 / 0.4141 / 20.65 for the baseline plus TaskPrompter (ICLR'22), and 46.26 / 0.4239 / 21.00 for the baseline plus TSP-Transformer (WACV'24).
• Backbone generalization (Table 6, PASCAL-Context one-label): on the stated different backbone, the proposed method reports 62.23 mIoU semantic segmentation, 62.14 mIoU parsing, 78.10 maxF saliency, 16.19 mErr normal, and 66.70 odsF boundary, above both DiffusionMTL (Prediction) and DiffusionMTL (Feature).
• Prototype visualization: different target tasks activate different slots of the task prototype; individual slots hold per-task characteristics while some elements share properties across tasks, which the paper says is what enables task association capture.
• Stated limitation: the framework is designed for tasks seen during training; extending it to unseen tasks via zero-shot or meta-learning is described as an open challenge.

Methodology in Plain English

The framework has two parts and is trained end to end.

First, a backbone network encodes the input image into a feature, which is pushed through a vector quantization step: a codebook of 4,096 learnable slots, where each encoded feature element is snapped to its nearest codebook entry (Eq. 1). The quantized feature is added back to the encoded feature, and a convolutional decoder reconstructs the input image under a task-agnostic enhancement (TAE) loss (Eq. 2). The purpose is to widen the shared feature space so that diverse task cues survive even when many labels are missing.

Second, a task-specific decoder turns that feature into a task-specific feature. This feature is matched against a task prototype — a set of learnable slots, one per task (T = 5 for PASCAL-Context, T = 3 for NYUD-v2), each of dimension d. The match is computed as a cosine-similarity-based task-similarity (Eq. 3), and a softmax over it produces the task-affinity score, which is essentially "how much does this task need enhancement, and from which related task knowledge."

Teaching the prototype uses two losses combined into the association knowledge generating loss (Eq. 6). The task knowledge embedding (TKE) loss (Eq. 4) uses a one-hot target so that the affinity score is highest for the prototype slot corresponding to the actual target task, forcing each slot to memorize its task's characteristics. The task consistency (TC) loss (Eq. 5) works across a batch, pulling features of the same task together and pushing features of other tasks apart with a margin, so the characteristics stay consistent across scenarios and distinct between tasks.

For retrieval, the task-affinity feature is obtained by multiplying the affinity score with the prototype (Eq. 7). A knowledge retrieval transformer of multiple blocks then takes the flattened task-specific feature as the query (after self-attention, Eq. 8) and the task-affinity feature as key and value in cross-attention (Eq. 9), followed by a feed-forward network. This produces the task-refined feature that each task head uses for prediction. At inference, all prototype parameters are frozen so consistent task knowledge is recalled across scenarios.

The total objective adds the multi-task supervised loss, the summed TAE loss, and the AKG loss with balancing parameters λ1 and λ2 (Eq. 10). Training uses ResNet-18 as backbone, Adam at an initial learning rate of 2 × 10⁻⁵, 100 epochs with batch size 6 on PASCAL-Context and 200 epochs with batch size 4 on NYUD-v2, on a single RTX A6000 GPU. Each prototype slot has dimension 1024, with 8 heads for cross-attention and 1024 output channels; the task-specific decoder uses 3 × 3 convolution layers with ReLU and the task head is a 1 × 1 convolution.

Why This Matters

The work argues that task relationships can be modeled from what a task inherently is, rather than from possibly-wrong guesses about unlabeled tasks. That reframing matters for research on partially supervised learning because it removes a source of negative transfer at the root instead of trying to denoise the guesses. It also connects prompt-based and prototype-based approaches by showing prototypes can be learned explicitly rather than in a supervised manner over fully labeled data.

Real-world applications implied by the paper's framing (autonomous driving, robotics) and by the dense prediction tasks it evaluates:

• Autonomous driving perception stacks, where segmentation, depth, surface normals, and boundary cues all feed downstream planning but are rarely labeled together.
• Robotics, where a single unified model must handle multiple functional roles from limited annotated data.
• Indoor scene understanding (the NYUD-v2 setting), including 13-class semantic segmentation, depth estimation, and surface normal estimation for indoor scenarios.
• General dense prediction pipelines that use human parsing, saliency detection, and boundary detection as auxiliary or joint objectives (the PASCAL-Context task set).

Industry relevance: labeling budgets are the practical bottleneck for multi-task perception models. A method that improves every task in Table 1 and Table 2 under the one-label setting — where each training image carries an annotation for only one task — and that also generalizes to a different backbone (Table 6, and the ViT-L results in Table 5) is directly relevant to teams that want more tasks per labeled image.

Future Directions

• Extending the framework to unseen tasks through zero-shot or meta-learning, which the paper explicitly names as an open challenge.
• Determining how the task prototype scales beyond the evaluated settings: the prototype was configured with T = 5 for PASCAL-Context and T = 3 for NYUD-v2, and slot dimension peaked at 1024 in Table 4, so scaling behavior for larger task sets is not reported.
• Investigating whether the approach holds on backbones and architectures other than those tested (ResNet-18 in Tables 1–3, ResNet-50 in Table 6, ViT-L in Table 5) and on task sets beyond the five PASCAL-Context tasks and three NYUD-v2 tasks.
• Understanding how the learned shared elements across prototype slots could be exploited — the paper observes that certain prototype elements share properties across tasks, which raises the question of whether these shared properties can be transferred or reused directly.

Target Audience

Researchers and graduate students working on multi-task learning, partially supervised learning, and dense prediction; practitioners building perception systems who face partial annotation budgets; and readers interested in prototype- and prompt-based conditioning of transformer architectures for computer vision. The paper is not beginner-friendly: it assumes fluency with attention mechanisms, vector quantization, and multi-task evaluation metrics such as mIoU, maxF, mErr, odsF, and absErr.

Authors’ 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.