TOOD: Task-aligned One-stage Object Detection

TLDR

One‑stage detectors typically use parallel classification and localization heads, which can cause spatial misalignment between the two predictions. This paper introduces Task‑aligned One‑stage Object Detection (TOOD) to explicitly align classification and localization during training. TOOD employs a Task‑aligned Head that balances task‑interactive and task‑specific features, and a Task Alignment Learning scheme that pulls optimal anchors together via a sample‑assignment strategy and a task‑aligned loss. On MS‑COCO, TOOD attains 51.1 AP, outperforming recent one‑stage detectors such as ATSS, GFL, and PAA while using fewer parameters and FLOPs, and qualitative results confirm improved task alignment. Code is available at https://github.com/fcjian/TOOD.

Abstract

One-stage object detection is commonly implemented by optimizing two sub-tasks: object classification and localization, using heads with two parallel branches, which might lead to a certain level of spatial misalignment in predictions between the two tasks. In this work, we propose a Task-aligned One-stage Object Detection (TOOD) that explicitly aligns the two tasks in a learning-based manner. First, we design a novel Task-aligned Head (T-Head) which offers a better balance between learning task-interactive and task-specific features, as well as a greater flexibility to learn the alignment via a task-aligned predictor. Second, we propose Task Alignment Learning (TAL) to explicitly pull closer (or even unify) the optimal anchors for the two tasks during training via a designed sample assignment scheme and a task-aligned loss. Extensive experiments are conducted on MS-COCO, where TOOD achieves a 51.1 AP at single-model single-scale testing. This surpasses the recent one-stage detectors by a large margin, such as ATSS [30] (47.7 AP), GFL [14] (48.2 AP), and PAA [9] (49.0 AP), with fewer parameters and FLOPs. Qualitative results also demonstrate the effectiveness of TOOD for better aligning the tasks of object classification and localization. Code is available at https://github.com/fcjian/TOOD.