[论文解读] Unbalanced Optimal Transport: A Unified Framework for Object Detection

Henri De Plaen, Pierre-François De Plaen+4 more|arXiv (Cornell University)|Jul 5, 2023

Advanced Neural Network ApplicationsComputer Science被引用 3

一句话总结

本文提出不平衡最优传输（UOT）作为目标检测匹配的统一框架，涵盖传统策略如最近邻、匈牙利匹配和一对多匹配。通过引入可学习约束和熵正则化，UOT 实现了更快、更稳定的训练，在 COCO 和合成的 Color Boxes 数据集上均达到最先进（SOTA）的 AP 和 AR 指标，且在 GPU 上高度可并行化。

ABSTRACT

During training, supervised object detection tries to correctly match the predicted bounding boxes and associated classification scores to the ground truth. This is essential to determine which predictions are to be pushed towards which solutions, or to be discarded. Popular matching strategies include matching to the closest ground truth box (mostly used in combination with anchors), or matching via the Hungarian algorithm (mostly used in anchor-free methods). Each of these strategies comes with its own properties, underlying losses, and heuristics. We show how Unbalanced Optimal Transport unifies these different approaches and opens a whole continuum of methods in between. This allows for a finer selection of the desired properties. Experimentally, we show that training an object detection model with Unbalanced Optimal Transport is able to reach the state-of-the-art both in terms of Average Precision and Average Recall as well as to provide a faster initial convergence. The approach is well suited for GPU implementation, which proves to be an advantage for large-scale models.

研究动机与目标

• 将不同的目标检测匹配策略——如最近邻、匈牙利匹配和一对多匹配——统一到一个单一、连贯的框架中。
• 解决现有方法的局限性，包括冗余、非唯一性、收敛缓慢以及 GPU 并行化能力差的问题。
• 通过可调的预测和真实框质量约束，实现对匹配行为的细粒度控制。
• 通过 UOT 框架中的正则化，提升训练收敛速度和最终性能。
• 为大规模目标检测模型提供可扩展、GPU 友好的 UOT 实现。

提出的方法

• 将目标检测匹配建模为不平衡最优传输（UOT）问题，其中预测到真实框的分配通过最小化正则化传输成本获得。
• 引入两个关键约束：对预测质量的约束（由 τ₁ 控制）和对真实框质量的约束（由 τ₂ 控制），从而在不同匹配模式之间实现插值。
• 采用熵正则化（ε）以支持快速、可微分且可 GPU 并行化的 Sinkhorn 迭代，实现可扩展的推理。
• 使用 1 - GIoU 作为预测框与真实框之间的代价矩阵，背景代价（c∅）设为 0.8。
• 在训练流程中将 UOT 解作为可微分的匹配层使用，替代手工设计的匹配启发式规则。
• 实际应用中采用固定数量的 Sinkhorn 迭代（例如 10–20 次），因为收敛迅速，进一步迭代带来的收益微乎其微。

(a) Image №163 from the VOC training dataset. The ground truth boxes are colored, and the predictions are outlined in black.

实验结果

研究问题

• RQ1 不平衡最优传输能否将多种目标检测匹配策略统一到一个连续的单一框架中？
• RQ2 超参数 τ₁（预测质量约束）和 τ₂（真实框质量约束）如何影响匹配行为和模型性能？
• RQ3 与标准匹配基线相比，UOT 是否能提升训练收敛速度和最终检测精度？
• RQ4 UOT 是否能高效地在 GPU 上实现，以加速大规模目标检测器的训练？
• RQ5 熵正则化和 Sinkhorn 迭代次数对最终检测指标有何影响？

主要发现

• UOT 在 COCO 数据集上实现了最先进（SOTA）的平均精度（AP）和平均召回率（AR），优于标准基线。
• 在 Color Boxes 数据集上，当 τ₁ = 10 且 τ₂ = 0.01 时，UOT 在使用 NMS 时达到 47.8 AP 和 63.8 AR，不使用 NMS 时达到 45.0 AP 和 72.6 AR。
• 平衡的 UOT 情况（τ₁ → ∞）达到最高性能：使用 NMS 时为 48.1 AP 和 64.3 AR，表明平衡的质量约束可提升精度。
• 使用 UOT 训练相比匈牙利匹配实现了更快的初始收敛，尤其在 DETR 类模型中表现明显。
• Sinkhorn 迭代次数在超过 10 次后影响极小：性能迅速饱和，10–20 次之后仅有微小增益。
• 基于 GPU 的 UOT 将 SSD300 的训练时间减少了近 50%，匹配时间从 18.3ms 降低至 1.5ms。

(b) Costs between the predictions and the ground truth ( $1-\mathrm{GIoU}$ ). The background cost is $c_{\varnothing}=0.8$ .