Distribution-Aware Margin Calibration for Semantic Segmentation in Images

Abstract

The Jaccard index, also known as Intersection-over-Union (IoU), is one of the most critical evaluation metrics in image semantic segmentation. However, direct optimization of IoU score is very difficult because the learning objective is neither differentiable nor decomposable. Although some algorithms have been proposed to optimize its surrogates, there is no guarantee provided for the generalization ability. In this paper, we propose a margin calibration method, which can be directly used as a learning objective, for an improved generalization of IoU over the data-distribution, underpinned by a rigid lower bound. This scheme theoretically ensures a better segmentation performance in terms of IoU score. We evaluated the effectiveness of the proposed margin calibration method on seven image datasets, showing substantial improvements in IoU score over other learning objectives using deep segmentation models.

A Appendix

1.1 A.1 Proof of Theorem 1

We first prove with probability $1-\frac{\eta }{K}$, the following inequality holds:

(19)

Assume the following inequality holds for non-negative $\epsilon _{k0}$ and $\epsilon _{0k}$:

(20)

Solving the above inequality, we can get:

$$\begin{aligned} \epsilon _k = (\frac{a_k}{b_k}\epsilon _{0k}+\epsilon _{k0})(b_k-\epsilon _{0k})^{-1}, \end{aligned}$$

(21)

where and .

Next, we should get the values of $\epsilon _{k0}$ and $\epsilon _{0k}$ to calculate $\epsilon _k$, which should satisfy the following inequality:

(22)

so we can simply substitute (22) into (20) to complete the proof.

The empirical label distribution $P_k$ is irrelevant to the model $\theta$ and we assume it is an accurate estimation of the label distribution $\mathcal {D}_{\mathcal {Y}}$, i.e.,

(23)

Based on the above approximation, a sufficient condition for (22) regarding $\epsilon _{0k}$ and $\epsilon _{k0}$ is:

(24)

Following the margin-based generalization bound in Mohri et al. (2018)[Theorem 9.2], for the $N_k$ foreground class pixels, with the probability at least $1-\frac{\eta }{2K}$, we have:

(25)

where $\mathfrak {R}_{N_k}(\Theta )$ is the Rademacher complexity for the hypothesis class $\Theta$ over the $N_k$ pixels of the k-th foreground class. For an input data batch with M pixels, we first apply the McDiarmid’s inequality for M-dependent data Liu et al. (2019) to the proof of Mohri et al. (2018)[Theorem 3.3]. Then we use it in the proof of Mohri et al. (2018)[Theorem 9.2] to get the formulation of (25).

The Rademacher complexity $\mathfrak {R}_{N_k}(\Theta )$ typically scales in $\sqrt{\frac{C(\Theta )}{N_k}}$ with $C(\Theta )$ being the a proper complexity measure of $\Theta$ Neyshabur et al. (2018), and such a scale has also been used in related work (see Cao et al. (2019) and the references therein). We can then rewrite (25) as:

(26)

where $\sigma (\frac{1}{\eta }) \triangleq \frac{\rho _{\max }}{4K} \sqrt{2M\log \frac{2K}{\eta }}$ is typically a low-order term in $\frac{1}{\eta }$.

Similarly, for the $N-N_k$ pixels of the background class, with the probability at least $1-\frac{\eta }{2}$,

(27)

for the $N-N_k$ background class pixels.

We then combine (26), (27), (24) and take a union bound over $\epsilon _{k0}$ and $\epsilon _{0k}$, to get following equations, with which (24) holds with the probability at least $1-\eta /K$:

$$\begin{aligned} \epsilon _{k0}&= \frac{\sqrt{N_k}}{N} \frac{4K}{\rho _{k0}}\mathcal {F},\nonumber \\ \epsilon _{0k}&= \frac{\sqrt{N-N_k}}{N}\frac{4K}{\rho _{0k}} \mathcal {F}. \end{aligned}$$

(28)

Then we substitute above equations into (21). Let $\mu _k = \frac{\rho _{k0}}{\rho _{0k}}$, we have:

$$\begin{aligned} \epsilon _k = \frac{\frac{a_k}{b_k}\sqrt{N-N_k} + \frac{\sqrt{N_k}}{\mu _k}}{\frac{b_k N}{4K\mathcal {F}}\rho _{0k}-\sqrt{N-N_k}}, \end{aligned}$$

(29)

so that with the probability at least $1-\eta /K$ the inequality (13) holds.

In practice, we do not know the values of $a_k$ and $b_k$ so that Eq.(29) has its own limitations. However, we know $\frac{a_k}{b_k}\le 1$ and so we can get a very useful bound:

$$\begin{aligned} \epsilon _k \le \frac{\sqrt{N-N_k} + \frac{\sqrt{N_k}}{\mu _k}}{\frac{N_k}{4K\mathcal {F}}\rho _{0k}-\sqrt{N-N_k}}. \end{aligned}$$

(30)

Averaging the union bound Eq.(13) over all classes, we can obtain the following inequality with probability at least $1-\eta$:

(31)

with

$$\begin{aligned} \epsilon = \frac{1}{K}\sum \limits _{k=1}^{K} \frac{ \sqrt{N-N_k}+ \frac{\sqrt{N_k}}{\mu _k}}{\frac{N_k}{4K\mathcal {F}}\rho _{0k}-\sqrt{N-N_k}}, \end{aligned}$$

(32)

where we complete the proof.

1.2 A.2 Proof of Colollary 1

We substitute $\mu _k$ in (32) with $\frac{\sqrt{N_k}}{r(N/N_k\!-\!1)-\sqrt{N\!-\!N_k}}$, where r is a hyper-parameter, we can get:

$$\begin{aligned} \epsilon&= \frac{1}{K}\sum _{k=1}^{K} (\frac{r(N-N_k)}{N_k})(\frac{N_k}{4K\mathcal {F}}\rho _{0k}-\sqrt{N-N_k})^{-1} \nonumber \\&= \frac{1}{K}\sum _{k=1}^{K} (\frac{r(N-N_k)}{N_k^2})(\frac{1}{4K\mathcal {F}}\rho _{0k}-\frac{\sqrt{N-N_k}}{N_k})^{-1}. \end{aligned}$$

(33)

Let $x_k = \frac{r(N-N_k)}{N_k^2}$ and $y_k = \frac{1}{4K\mathcal {F}}\rho _{0k}-\frac{\sqrt{N-N_k}}{N_k}$. According to Cauchy-Schwarz inequality we have:

$$\begin{aligned} \left( \sum _{k=1}^{K} \sqrt{\frac{x_k}{y_k}}\cdot \sqrt{y_k}\right) ^2 \le (\sum _{k=1}^{K}\frac{x_k}{y_k}) (\sum _{k=1}^{K} y_k), \end{aligned}$$

(34)

so that

$$\begin{aligned} \epsilon&\ge \frac{1}{K}\cdot \left( \sum \limits _{k=1}^{K} \sqrt{x_k}\right) ^2 (\sum \limits _{k=1}^{K} y_k)^{-1} \nonumber \\&= \frac{r}{K}\cdot \frac{\left( \sum \limits _{k=1}^{K} \frac{\sqrt{N-N_k}}{N_k}\right) ^2}{\frac{1}{4K\mathcal {F}}\sum \limits _{k=1}^{K} \rho _{0k}- \sum \limits _{k=1}^{K}\frac{\sqrt{N-N_k}}{N_k}}. \end{aligned}$$

(35)

The RHS of the equality is a constant because r is a given hyper parameter and we assume $\sum \limits _{k=1}^{K} \rho _{0k} = \text {some contant}$. The equality holds when $\frac{\sqrt{x_1}}{y_1}=\ldots =\frac{\sqrt{x_K}}{y_K}$, which yields Corollary 1.

Note that $\mu _k = \frac{\sqrt{N_k}}{r(N/N_k-1)-\sqrt{N-N_k}}$, while in Corollary 1$\mu _k = \frac{P_k\sqrt{N_k}}{\upsilon (N-N_k)-P_k\sqrt{N-N_k}}$. These two conditions are essentially equivalent when r and $\upsilon$ are hyper-parameters. To see this, simply let $r=N\upsilon$ and notice that $P_k=\frac{N_k}{N}$.