【論文5分まとめ】Ota: Optimal transport assignment for object detection

この記事は、論文の内容を5分くらいで読めるようにまとめた記事です。そのため、前提となる知識や関連研究に関する説明は大幅に省略しています。
基本的には筆者の備忘録ですが、面白そうと思ったら是非ご自身でも読んでみてください。

概要

物体検出モデルにおいて、予測された矩形と真の矩形（gt）の割り当ては、重要な課題の一つである。例えば、RetinaNetであれば、gtとアンカーのIoUが閾値以上であれば、そのアンカーから作られる予測矩形のターゲットとしてgtが割り当てられる。また、FCOSであれば、gtの中心に近い位置やgtの領域に対応する位置から作られる予測矩形にgtを割り当てる。
しかし、このような静的な割り当ては、以下の図のような複数のgtに所属してもよさそうな曖昧な領域の存在の扱いを難しくし、不適切なターゲットによる有害な勾配を生じさせる。

このような問題を回避するために、近年はさまざまな動的な割り当てが研究されており、本研究が提案する手法OTA(Optimal Transport Assignment)もその一つである。

書誌情報

• Ge, Zheng, et al. "OTA: Optimal Transport Assignment for Object Detection." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
• https://arxiv.org/abs/2103.14259
• 公式実装

ポイント

OTAの問題設定

OTAは、gtと予測矩形の割り当てを、最適輸送問題として捉える。最適輸送問題では、$m$個の輸送元と$n$個の輸送先があり、各輸送元から各輸送先への輸送コスト$c_{ij}$があるとする。各輸送元は$s_i$個の荷物を供給し、各輸送先は$d_j$個の荷物を需要しているとする。最適輸送問題は、このような条件のときに、どこからどこへ何個荷物を送るか、という最適な輸送計画$\pi^{*}=\left\{\pi_{i, j} \mid i=1,2, \ldots m, j=1,2, \ldots n\right\}$を求めよ、という問題である。

以上の問題設定は、以下のような最小化問題に整理できる。

$$\begin{array}{ll} \min _{\pi} & \sum_{i=1}^{m} \sum_{j=1}^{n} c_{i j} \pi_{i j} . \\ \text { s.t. } & \sum_{i=1}^{m} \pi_{i j}=d_{j}, \quad \sum_{j=1}^{n} \pi_{i j}=s_{i}, \\ & \sum_{i=1}^{m} s_{i}=\sum_{j=1}^{n} d_{j}, \\ & \pi_{i j} \geq 0, \quad i=1,2, \ldots m, j=1,2, \ldots n . \end{array}$$

OTAでは、輸送元をgt、輸送先を予測矩形と考える。それにより、各予測矩形に対して適切なgtを割り当てるという問題にうまく対応させることができる。OTAでは、各予測矩形に対してただ1つのgtまたは背景が割り当てられると考え、$\{d_j=1 \mid j=1,2, \ldots n\}$とする。また、gtと背景が供給する荷物の数を以下のように決める。$m+1$行目は背景を表す。

$$s_{i}= \begin{cases}k, & \text { if } \quad i \leq m \\ n-m \times k, & \text { if } \quad i=m+1 \end{cases}$$

コスト行列$c_{ij}$については、仮にgt$G_i$が予測矩形$P_j$のターゲットであるならばどのような分類損失やIoU損失になるのかをもとに決める。

$$\begin{aligned} c_{i j}^{f g}=& L_{c l s}\left(P_{j}^{c l s}(\theta), G_{i}^{c l s}\right)+ \alpha L_{r e g}\left(P_{j}^{b o x}(\theta), G_{i}^{b o x}\right) \end{aligned}$$

背景が輸送元の場合は、単に分類損失のみがコストとなる。

$$c_{j}^{b g}=L_{c l s}\left(P_{j}^{c l s}(\theta), \varnothing\right)$$

以上によって、gtを予測矩形に割り当てるという問題を定式化できた。この問題は、 Sinkhorn-Knoppアルゴリズムによって解くことができる。詳細は論文の付録に記載されている。以下に公式実装からの抜粋を示す（結構短い）。

Sinkhorn-Knopp

class SinkhornDistance(torch.nn.Module):
    r"""
        Given two empirical measures each with :math:`P_1` locations
        :math:`x\in\mathbb{R}^{D_1}` and :math:`P_2` locations :math:`y\in\mathbb{R}^{D_2}`,
        outputs an approximation of the regularized OT cost for point clouds.
        Args:
        eps (float): regularization coefficient
        max_iter (int): maximum number of Sinkhorn iterations
        reduction (string, optional): Specifies the reduction to apply to the output:
        'none' | 'mean' | 'sum'. 'none': no reduction will be applied,
        'mean': the sum of the output will be divided by the number of
        elements in the output, 'sum': the output will be summed. Default: 'none'
        Shape:
            - Input: :math:`(N, P_1, D_1)`, :math:`(N, P_2, D_2)`
            - Output: :math:`(N)` or :math:`()`, depending on `reduction`
    """

    def __init__(self, eps=1e-3, max_iter=100, reduction='none'):
        super(SinkhornDistance, self).__init__()
        self.eps = eps
        self.max_iter = max_iter
        self.reduction = reduction

    def forward(self, mu, nu, C):
        u = torch.ones_like(mu)
        v = torch.ones_like(nu)

        # Sinkhorn iterations
        for i in range(self.max_iter):
            v = self.eps * \
                (torch.log(
                    nu + 1e-8) - torch.logsumexp(self.M(C, u, v).transpose(-2, -1), dim=-1)) + v
            u = self.eps * \
                (torch.log(
                    mu + 1e-8) - torch.logsumexp(self.M(C, u, v), dim=-1)) + u

        U, V = u, v
        # Transport plan pi = diag(a)*K*diag(b)
        pi = torch.exp(
            self.M(C, U, V)).detach()
        # Sinkhorn distance
        cost = torch.sum(
            pi * C, dim=(-2, -1))
        return cost, pi

    def M(self, C, u, v):
        '''
        "Modified cost for logarithmic updates"
        "$M_{ij} = (-c_{ij} + u_i + v_j) / epsilon$"
        '''
        return (-C + u.unsqueeze(-1) + v.unsqueeze(-2)) / self.eps

Center Prior

直感的に、gtの位置は予測矩形の位置に近い方が良いことがわかる。この事前知識をコスト行列に反映させるために、gtの中心に近い$r^2$個の位置から得られる予測矩形については追加のコストを0とし、それより遠い位置から得られる予測矩形に対しては一定のコスト$c_{cp}$を加えることとしている。

最適な$r$については、実験では$r=5$で良い結果が得られている。

Dynamic $k$ Estimation

各gtが保持する割り当て数$s_i$は、固定値$k$にすることもできるが、各gtの大きさや、gt同士のオクルージョンの大きさなどの条件を考えると、各gtに応じて変化させるべきだろうと考えるのが自然な発想である。

具体的には、予測矩形のうち、$G_i$とのIoUが高いトップ$q$個のIoUを足し合わせた値を、$s_i(>=1)$とする。$q$は公式実装をざっと見る限り20個くらいみたい。

$k$を固定値にするべきか動的に決めるべきかについても実験をおこなっており、さまざまな固定の$k$よりもうまくいくことを示している。

アルゴリズムのまとめ

以上の流れをまとめると、以下のようなアルゴリズムとなる。

実験

実験については省略するが、他の動的な割り当て手法と比較して、OTAの割り当て結果が納得感のあるものであることがわかる図を確認する。

下図はATSS, PAAといういずれも動的な割り当ての手法であるが、赤い点線で囲まれた楕円を比較すると、OTAにおいては、複数のgtが重なっている割り当てが曖昧になりそうな箇所や、背景の領域に対する割り当てが少なく、gtと予測矩形の割り当てが適切になされていることがわかる。