PSANet: Point-wise Spatial Attention Network for Scene Parsing

Idea

Apply point-wise spatial attention mechanism to scene parsing.

Method

Let ${\mathbf{z}}_i$ be the newly aggregated feature at position $i$, and ${\mathbf{x}}_i$ be the feature representation at position $i$ in the input feature map $\text{X}$, then we have the following bi-directional propagation formula:

$${\mathbf{z}}_i=\frac{1}{N}\sum _{\forall j\in \Omega (i)}\left[F_{{\Delta }_{ij}}({\mathbf{x}}_i){\mathbf{x}}_j+F_{{\Delta }_{ji}}({\mathbf{x}}_j){\mathbf{x}}_i\right]$$

where $\forall j\in \Omega (i)$ enumerates all positions in the region of interest associated with $i$, and ${\Delta }_{ij}$ represents the relative location of position $i$ and $j$

Bi-Direction Information Propagation

For the first term, $F_{{\Delta }_{ij}}({\mathbf{x}}_i)$ encodes to what extent the features at other positions can help prediction. Each position collects information from other positions. For the second term, $F_{{\Delta }_{ij}}({\mathbf{x}}_j)$ denotes the importance of the feature at one position to to features at other positions. Each position distributes information to others.

Specifically, in this model both $F_{{\Delta }_{ij}}({\mathbf{x}}_i)$ and $F_{{\Delta }_{ij}}({\mathbf{x}}_j)$ can be regarded as predicted attention values to aggregate feature ${\mathbf{x}}_j$, rewriting the formula above as

$${\mathbf{z}}_i=\frac{1}{N}\sum _{\forall j}{\mathbf{a}}_{i,j}^{\mathrm{c}}{\mathbf{x}}_i+\frac{1}{N}\sum _{\forall j}{\mathbf{a}}_{i,j}^{\mathrm{d}}{\mathbf{x}}_j$$

where ${\mathbf{a}}_{i,j}^{\mathrm{c}}$ and ${\mathbf{a}}_{i,j}^{\mathrm{d}}$ denote the predicted attention values in the point-wise attention maps ${\text{A}}^{\mathrm{c}}$ and ${\text{A}}^{\mathrm{d}}$ from collect and distribute branches, respectively.

Architecture

PSA Module

Attention Map Generation

In the collect branch, at each position $i$, with $k$-th row and $l$-th column, we predict how current position is related to other positions based on feature at position $i$.

Specifically, element at $s$-th row and $k$-th column in the attention mask ${\mathbf{a}}_i^{\mathrm{c}}$ (i.e. ${\mathbf{a}}_{[k,l]}^{\mathrm{c}}$) is

$${\mathbf{a}}_{[k,l],[s,t]}^c={\mathbf{h}}_{[k,l],[H-k+s,W-l+t]}^c,\forall s\in [0,H),t\in [0,W)$$

where $[\cdot ,\cdot ]$ indexed position in rows and columns.

Model