STC-Flow: Spatio-temporal Context-aware Optical Flow Estimation

Analysis on Attention Mechanism. To capture long-range dependencies and model contextual details of single images or video clips, the basic non-local network [17] aggregates pixel-wise information via self-attention mechanism. We denote $X$ and $Z$ as the input and output signals, such as the single image and the video clip. The non-local block can be expressed as follows:

where $i$ and $j$ are the indices of target position coordinates and all possible enumerated positions. $f(X_{i},X_{j})$ denotes the relationship between position $i$ and $j$, which is normalized by a factor $\mathcal{N}(X)$. The matrix multiplication operation is utilized to strengthen details of each query position. Embedded Gaussian is a widely-used instantiation of $f(X_{i},X_{j})$, to compute similarity in an embedding space, and normalized with a softmax function, which is a soft selection across channels in one position. the non-local block with Embedded Gaussian is shown in Figure 3(b), and is expressed as follows:

where $W_{q}$, $W_{k}$ and $W_{v}$ are linear transformation matrices.

Why attention for optical flow estimation? Here, we discuss the relation between correlation in optical flow estimation and matrix multiplication in self-attention mechanism. We aim to explore the contextual information from the input pairs. Denote the feature pairs by $F_{1}$ and $F_{2}$. To distinguish the features from different input paths, we mark the size of features with height $H_{k}$, width $W_{k}$, channel $C_{k}$ ($k=1,2$), and position coordinate $\mathbf{x}_{k}\in[1,H_{k}]\times[1,W_{k}]$, channel index $c_{k}\in C_{k}$, and here $H_{1}=H_{2}$, $W_{1}=W_{2}$ and $C_{1}=C_{2}$. As the key function in optical flow estimation, the “correlation” operation between two patches of feature pairs, $\mathbf{f}_{1}$ and $\mathbf{f}_{2}$, is defined as follows for temporal modeling:

where $Corr(\mathbf{f}_{1},\mathbf{f}_{2})$ denotes the cost volume calculated via correlation. $\mathbf{o}\in[-n,n]\times[-n,n]$ denotes the offset of correlation operation with search region. In consideration of matrix multiplication in the attention mechanism of $f(X_{i},X_{j})$ shown in Figure 4, the different order of the two matrices in multiplication leads to great disparity of explanation with different displacements of correlation.

Discussions. In Figure 4(a), the expression is defined as $F_{2}(\mathbf{x}_{2},c_{2})(F_{1}(\mathbf{x}_{1},c_{1}))^{\top}\in\mathbb{R}^{H_{2}W_{2}\times H_{1}W_{1}}$. If $F_{1}=F_{2}$, this operation strengthens the detail features of each position via aggregating information across channels from other positions, which would indicate the spatial attention integration at full resolution, and it is utilized to the basic non-local block [17]. However, if $F_{1}\neq F_{2}$, as the definition of cost volume, only the diagonal elements present the correlation with no displacement. On the contrary, the expression is defined as $(F_{1}(\mathbf{x}_{1},c_{1}))^{\top}F_{2}(\mathbf{x}_{2},c_{2})\in\mathbb{R}^{C_{1}\times C_{2}}$ in Figure 4(b), which is a global correlation representation at full resolution among channels, and is essential to the naive correlation operation between feature pairs. For different matrix multiplication approaches, the attention maps catch dependencies with corresponding concepts in spatial features and temporal dynamics, which enhance representation for input feature extraction and correlation calculation, respectively.

Lite matrix multiplication. Considering the runtime of the flow prediction, the matrix multiplication in contextual attention block needs to be simplified with less computational complexity. In Figure 5, according to the neighbor similarity of images or frame pairs, we propose the polyphase decomposition and reconstruction scheme to simplify matrix multiplication opreation, which would obtain better approximation than the naive downsampling-upsampling scheme, and reduce the computation complexity compared to the direct multiplication. Denote the polyphase decomposition factor as $s$ ($s>1$). Given a reshaped feature $\tilde{F}\in\mathbb{R}^{M\times N}$, the FLOPs of the entire multiplication is reduced from $O\left(NM^{2}\right)$ to $O\left(\frac{NM^{2}}{s}\right)$. The comparison of different factors is presented in Section IV.

Contextual Attention Framework. In general, the input of CNNs is not limited to the single feature through the single path, and the attention block needs to be adapted to more than one features, e.g. two input features of the correlation operation. As shown in Figure 3(a), the components of the attention block can be abstracted as follows:

where $\mathcal{F}(\cdot)$ and $\mathcal{G}(\cdot)$ are the fusion operations for attention aggregation and target fusion. $T(\cdot)$ and $A(\cdot)$ denote target operation and attention integration for the input features, $\omega$ is the factor of linear transformation. The non-local block or the other attention modules are the specific form of context attention block with the single input feature, e.g. $A_{ij}\left({X}\right)={{f(X_{i},X_{j})}}/{{\mathcal{N}(X)}}$, and $T\left({X}\right)$ is the all-pass function in the non-local block.

where $C_{P}^{(k)}$ and $C_{C}^{(k)}$ are contextual attention at stage $k$ fused with that of stage $k-1$, which is to aggregate context from different granularity:

where “$\Downarrow$” denotes max-pooling, and “$\left[\cdot\right]$” denotes the concatenation operator. $A_{P,ij}$ and $A_{C,ij}$ are attention integrations in position and channel, defined as follows to learn the spatial and channel interdependencies:

where $A_{T,ij}$ is the temporal attention integration with the “cross-attention” mechanism, which is defined as follows:

where “*” denotes the position-wise convolution operator, and here $W_{v}$ is a bilinear interpolation operator for $R^{(k)}$. $A_{U,ij}$ denotes the adaptive attention kernels to model the detail context defined as follows:

where $ps$ denotes the “Pixel Shuffle [18]” operator for sub-pixel convolution, to reconstruct the sub-pixel information and preserve edges and textures. $\sigma$ is the upsampling factor, and here $\sigma=2$.

Pyramidal spatial context module. STC-Flow utilizes PSC Module in level 3, 4 and 5. Table I(c)(a) demonstrates that using PSC Module can improve the performance on both the Sintel and KITTI datasets, since this module enhances the ability of discriminating feature texture in feature extraction stage, and PSC at stage 3 is more beneficial, for the low-level discriminative details matter.

Temporal context correlation module. TCC Module describes the relationship of correlation with the spatial and temporal context. In Table I(c)(b), we compare the performance of our network using TCC Module with naive correlation operator, and also compare with 2D non-local block for concatenated feature and 3D non-local block for feature pairs. It demonstrates that fusion of correlation with spatial and temporal context is better than single correlation. Notice that 3D non-local blocks perform better in Sintel, however, with heavy computational complexity. TCC can achieve the comparable performance with fewer FLOPs.

Recurrent residual contextual upsampling. We utilize the RRCU Module to learn high-frequency context features and preserve edges. In Table I(c)(c), we compare the quantity of our method using RRCU with single transpose convolution, which demonstrates that reconstruction context learning could preserve details and improve performance.

Lite matrix multiplication. Lite matrix multiplication is an efficient scheme to reduce the computational complexity. We compare the performance of this scheme with different polyphase decomposition factor $s$ on Sintel training. As shown in Table II, lite matrix multiplication has a margin influence on AEE, but increases the frame rate conspicuously. Considering the tradeoff between accuracy and time consumption, we select $s=2$ for the full model.