MDFlow: Unsupervised Optical Flow Learning by Reliable Mutual Knowledge Distillation

Finding dense correspondences between a pair of time adjacent frames, namely optical flow estimation, has been studied for decades for its fundamental role in many downstream vision tasks. FlowNet [9] and its successor FlowNet2 [10] are the first attempt to apply deep learning methods for optical flow estimation, which directly regress flow field based on the encoder-decoder U-shape network or its cascaded version. Inspired by traditional coarse-to-fine pipeline [6, 31, 8], SPyNet [32], PWC-Net [11] and LiteFlowNet [33] employ pyramid, warping and cost volume into end-to-end learning and achieve impressive real-time performance. After that, IRR-PWC [12] iteratively and jointly estimates residual flow and occlusion with shared estimators across pyramid levels. The work [34] introduces a dual self-attention module to improve original pyramid flow network. MaskFlowNet [35] and OAS-Net [36] explore to resolve ambiguous matching caused by warping operation. As for more efficient architecture, FDFlowNet [37] introduces U-shape backbone and partial fully connected decoder for compact structure. FastFlowNet [38] constructs center dense dilated correlation layer and shuffle block decoder to reduce computation and first achieves real-time performance on embedded systems. Recently, DICL-Flow [14] transfers concatenated feature volume in stereo matching to optical flow and propose displacement-invariant cost learning. Different from above coarse-to-fine pipeline, RAFT [13] first calculates all pairs’ correspondences and then introduces a recurrent module to estimate residual flow and update multi-scale matching cost simultaneously, which achieves significant accuracy improvement. Later on, GMA [39] introduces a global motion aggregation module based on transformer to improve recurrent-based flow architecture and obtains state-of-the-art accuracy.

Due to prohibitive cost to acquire optical flow ground truth of real-world images, alternative approaches try to leverage countless unlabeled video sequences for unsupervised learning. These methods build objective functions based on brightness constancy assumption and local smoothness prior [18, 19]. Then, UnFlow [20] and OccAwareFlow [21] boost unsupervised performance by excluding residual calculation in reasoned occluded regions. However, the occluded areas are only guided by rigid smoothness constraint, that can damage overall accuracy. To handle this problem, DDFlow [40] and SelFlow [25] distill flow estimation from the teacher model to the student with random crop and occlusion hallucination in a data-driven manner, that further improves unsupervised accuracy. However, the teacher can not benefit from the improved student. STFlow [41] integrates variational refinement with unsuperivsed learning and propose a self-taught framework for continual improvement. SimFlow [26] explores learnable feature similarity for regulating previous census reconstruction loss. ARFlow [22] and UFlow [23] propose to integrate augmentation regularization into each iteration step for continuous improvement of a single model. CoT-AMFlow [42] develops an adaptive modulation network and adopts a co-teaching strategy for better accuracy. Later on, UPFlow [24] improves the upsampling unit of PWC-Net and proposes a better pyramid distillation loss. DistillFlow [43] trains multiple teacher models and introduces a confidence based two-stage distillation approach for improvement. OIFlow [44] puts up an occlusion-inpainting framework to make full use of occlusion regions. Recently, ASFlow [45] presents content-aware pooling and adaptive flow upsampling modules to improve pyramid-based unsupervised flow deep structure. MRDFlow [46] further introduces 4D correlation layer and recurrent decoder of RAFT to strengthen flow estimation network. At the same time, SMURF [27] replaces PWC-Net with a more powerful RAFT backbone and proposes a sequence-aware self-supervision loss to achieve state-of-the-art accruacy. However, it incurs huge computation cost and can not satisfy many real-time applications. Different from above methods, we separate and recombine multiple objective parts into a mutual distillation framework for reliable performance improvement without increasing model size and inference delay.

Another line of unsupervised optical flow methods resort to additional information beyond two input frames. Based on the setting of stereo video systems, UnOS [47] enforces geometry constraints among stereo depth, camera ego-motion and optical flow for mutual promotion. Flow2Stereo [48] introduces data distillation into the joint learning framework of optical flow and stereo matching. Most recently, the work [49] shows that feature-level collaboration of the networks for optical flow, stereo depth and camera motion can outperform previous methods that only consider loss-level joint optimization. To facilitate unsupervised optical flow in challenging scenes, such as fog, rain and night, GyroFlow [50] converts gyroscope readings into gyro field and fuse it with flow information for recovering more motion details. Different from these methods that resort to additional information, proposed MDFlow focuses on improving the basic setting by reducing matching noise and exploring better flow architecture during the decoupled bidirectional distillation procedure.

Knowledge Distillation (KD) is usually exploited to train a compact student network by mimicking the output distribution of a pre-trained complex teacher model as well as one-hot ground-truth labels [51]. Following variants [52, 53, 54, 55] mainly focus on utilizing intermediate information of teacher, such as feature maps or attention masks as extra hints for improvement. Different from above one-way transfer between a teacher and a student in knowledge distillation, Deep Mutual Learning (DML) [56] finds that mutual learning of a collection of simple student networks can outperforms distillation from a more powerful yet static teacher. The work EKD [57] introduces an evolutionary teacher that can enable more efficient knowledge transfer by minimizing the capability gap between teacher and student. Dense Cross-layer Mutual-distillation (DCM) [58] improves the two-way Knowledge Transfer (KT) scheme by adding auxiliary classifiers to hidden layers of both teacher and student networks and building dense bidirectional KD between these classifiers. The work [59] also adopts mutual distillation method to encourage information exchange between the local branch and the global branch of a single network for discovering new categories. There are also some research about applying mutual learning to computer vision tasks, such as object detection [60, 61] and video-based sign language recognition [62]. Existing mutual distillation approaches have achieved success in high-level classification tasks. However, whether mutual distillation can help the middle-level dense matching tasks has not been fully explored.

As for optical flow estimation, DDFlow [40] and SelFlow [25] distill flow estimation from teacher to student with several occlusion hallucination methods, which has been shown to be effective when inferring on occluded regions. DistillFlow [43] further improves above two-stage data distillation by introducing multiple teacher models and confidence mechanism. As for a related medical image registration task, the work CRD [63] distills knowledge from a feature-shifted teacher model with high resolution input to a student model with low resolution input for more efficiency. Different from their works, our MDFlow transfers knowledge between teacher and student networks mutually, so as to decouple matching outliers in augmentation regularization and exploit recent advanced architecture as student for better final prediction, while maintaining real-time inference.

As depicted in Figure 1, our proposed MDFlow algorithm trains a teacher model $\mathcal{T}$ and a student model $\mathcal{S}$ interactively for progressive improvement in a three-stage manner. In the first stage, our teacher network can be trained by any unsupervised approach for generating noisy pseudo labels, which will be used later. Specifically, we adopt PWC-Net [11] as teacher and use UFlow [23] loss functions for unsupervised initialization. In the second stage, we introduce a novel reliable matching selection mechanism, that can partly remove outliers in generated pseudo labels, which will then be used to train a student network with diverse data augmentation approaches, constituting our forward knowledge distillation process. Note that we can select a stronger student architecture, such as RAFT [13], for sufficient learning thanks to the decouple nature in mutual knowledge distillation. In the last stage, better pseudo labels predicted by the student will be adopted to supervise the learning of original teacher model enriched by diverse data augmentation. Moreover, we can improve the final teacher performance by adding a weak unsupervised learning objective, formulating it into a multi-target manner. More details are shown in Figure 2.

Given two consecutive frames $I_{1},I_{2}$ from unlabeled image sequences $\mathcal{I}$, unsupervised optical flow methods aim to train a flow network $\mathcal{N}(.)$ based on brightness constancy assumption and spatial smoothness prior. The first goal is usually realized as a photometric reconstruction loss:

where $\odot$ represents element-wise multiplication, $O_{f},O_{b}$ are occlusion masks inferred by forward-backward consistency check [64, 20], and $\rho(x)=(|x|+\epsilon)^{q}$ is a robust function with $\epsilon=0.01,q=0.4$ in all our experiments. Denoting $F_{f}=\mathcal{N}(I_{1},I_{2}),F_{b}=\mathcal{N}(I_{2},I_{1})$ as predicted forward and backward flow fields, and $I_{2}^{w}=\textit{w}(I_{2},F_{f}),I_{1}^{w}=\textit{w}(I_{1},F_{b})$ are flow warped reconstruction images. Following previous research [40, 23], we use the soft Hamming distance of census transformed [28] image patches to calculate reconstruction for its robust to illuminance variation.

Due to the well-known aperture problem, above solution can be ambiguous on textureless or repetitive pattern regions, so we further introduce a loss item to constrain the spatial smoothness of predicted flow fields, which is usually reweighted by local image gradients:

where $\lambda$ controls edge-aware weighting strength conditioned on reference image. $\nabla_{d}^{k},k=1,2$ stands for the $k$-th order gradient operator. In our experiments, we set $\lambda=150$. For smoothness, following the experience in UFlow [23], we use first order operator on Flying Chairs and Sintel, while adopting second order operator on KITTI for their better results. The reason is that there are more parallel motion about image plane in Chairs and Sintel, while more vertical motion in regard to image plane in driving scenes of KITTI. Flow fields projected from the parallel and vertical motion about image plane satisfy the first and second order smoothness respectively.

Inspired by recent work of ARFlow [22] and UFlow [23], we also add a self-supervised loss for continual improvement. Denoting $\mathcal{A}$ as the data augmentor, which can jointly augment input images $I_{1},I_{2}$, pseudo labels $F_{f},F_{b}$ and valid masks $M_{f},M_{b}$ by applying spatial and color transformations. Indicating $\overline{I_{1}},\overline{I_{2}},\overline{F_{f}},\overline{F_{b}}$ and $\overline{M_{f}},\overline{M_{b}}$ as corresponding augmented images, pseudo labels and valid masks, the surrogate supervision loss can be formulated as:

Note that there are two models $\mathcal{N}_{1}$ and $\mathcal{N}_{2}$ in Eq. 3, where $\mathcal{N}_{1}$ is optimized conditioned on the pseudo labels $F_{f},F_{b}$ predicted by $\mathcal{N}_{2}$ and augmentor $\mathcal{A}$. Following previous experience, augmented training samples are detached from gradient calculation diagram for stable learning. We adopt above loss notation for clarity of following expression. In regard to the first stage of MDFlow, we employ the off-the-shelf UFlow [23] method to initialize the teacher $\mathcal{T}$ with following objective:

where $\lambda_{1},\lambda_{2}$ are trade-off coefficients, while $\mathcal{L}_{sup}(\mathcal{T}|\mathcal{T},\mathcal{A})$ means that the second forward prediction is supervised by the augmented first prediction of the same model $\mathcal{T}$.

Due to shadow, exposure, textureless and repetitive patterns existed in natural images, optimal solution of photometric reconstruction will inevitably contain mismatching areas. On the other hand, state-of-the-art unsupervised optical flow approaches [22, 23] augment predicted flow fields for a second supervision to improve performance. Concretely, they select non-occluded regions as the valid mask when calculating above self-supervised loss. As shown in Figure 3, large error regions in predicted pseudo labels, for example, the background surrounded by moving persons, are not identified by occlusion map. Thus, the augmented surrogate supervision will hinder the network from learning true displacement. To alleviate this problem, we first propose a novel confidence matching selection mechanism to find relatively reliable prediction, and then separate $\mathcal{L}_{sup}$ in Eq. 4 with $\mathcal{L}_{ph},\mathcal{L}_{sm}$ into a supervised knowledge distillation process. Holding the opinion that better matched image patches should have lower residuals, we explore to set a threshold to keep only a portion of reliable pseudo labels during forward knowledge distillation. In this work, we employ the soft Hamming distance of census transformed image patches for confidence measurement, and also include a robust function $\rho$, which does not change confidence ranking. Denote $M_{f}$ as forward valid mask and $\tau$ as a threshold, the selection approach can be formulated as:

In order to set meaningful threshold for outlier removal, we analyse the variation curve of average endpoint error with removal rate on Sintel and KITTI training sets, which is controlled by $\tau$. Similar to uncertainty estimation [65], the sparsification curves are depicted in Figure 4, while several typical removal rates with corresponding thresholds and errors are listed in Table I. As can be seen, average endpoint error of remaining pseudo labels decreases as we gradually remove more prediction from original non-occluded mask, until removal rate reaches about $90\%$. On the other hand, larger percentage of removal means less pseudo labels for training, that can damage performance. Therefore, setting reasonable removal rate for better trade-off plays a key role in forward distillation process. We leave this question in experiment part for detailed discussion. Guided by the above confidence matching selection mask $M_{f},M_{b}$, reliable forward knowledge distillation can be formulated as:

Due to the decouple characteristic of proposed mutual distillation framework, we can employ a stronger flow architecture as student $\mathcal{S}$ than the original teacher $\mathcal{T}$, so as to learn from reliable pseudo labels sufficiently. This process is different from traditional knowledge distillation [51, 66] where a lightweight student usually distills knowledge from a large teacher model. Therefore, our approach further includes a backward distillation stage, aiming to transfer better student knowledge back to the efficient teacher model. As a result, proposed MDFlow does not add extra computation cost and inference delay in real deployment, compared with other unsupervised approaches.

To achieve this goal, one may follow the supervision loss of Eq. 6 used in stage 2, by exchanging $\mathcal{T}$ and $\mathcal{S}$. However, due to the limited learning ability of efficient optical flow network, and diverse strong augmentation imposed on pseudo labels, training data for the efficient teacher $\mathcal{T}$ is blended with augmentation noise, which deviates from the distribution of real-world dynamic scenes. On the other hand, we can not remove the augmentor $\mathcal{A}$ during backward distillation, since data augmentation has been proved to be one crucial step for excellent performance [67, 68]. To deal with this problem, we establish a multi-target learning manner for reliable backward knowledge distillation, integrating specific domain knowledge on optical flow task. Specifically, unsupervised photometric and smoothness losses based on original scene structure are introduced as regularization. In short, reliable backward distillation of stage 3 can be written as:

where $\lambda_{3},\lambda_{4}$ are weight coefficients controlling unsupervised regularization intensity, and non-occluded masks $1-O_{f},1-O_{b}$ of $\mathcal{T}$ are adopted as valid masks in $\mathcal{L}_{sup}$.

We employ the efficient PWC-Net [11] as teacher model, and the powerful RAFT [13] as student model, whose model size, running time and computation complexity are compared in Table II. It can be seen that even RAFT contains a little less parameters than PWC-Net, however, its running time and computation complexity is 4$\times$ larger. All our experiments are implemented in PyTorch and conducted on 2 NVIDIA Tesla V100 GPUs. We adopt the Adam [69] and AdamW [70] optimizers to train the PWC-Net and RAFT respectively. Batch size is set to 4 for all experiments. Following previous work [40, 26, 23], we first pre-train MDFlow on Flying Chairs, and then fine-tune above teacher and student networks on each target dataset. As for Sintel, we use the combination of Clean and Final training parts, and adopt the fine-tuned teacher model for online evaluation. In regard to KITTI, we train on the KITTI 2015 multi-view extension training set, while removing frame numbers 9-12 in each sequence. Results of the teacher on KITTI 2015 test set are uploaded to KITTI website for online comparison. It takes overall about 4 days to train the three-stage MDFlow first on Chairs and then on Sintel or KITTI datasets, which also relys on the model complexity, IO speed and code parallelism degree.

On each dataset, we perform proposed three-stage training pipeline for mutual knowledge distillation, where the first and last stage both take 300k iterations to train the teacher, while the second stage takes 100k iterations to train the student. When training on Flying Chairs, learning rate is initially set to $1e-4$ and decays half at 100k, 150k, 200k and 250k iterations in both stage 1 and stage 3. As for stage 2, it is initially set to $4e-4$ and decays half at 20k, 40k, 60k and 80k iterations. When fine-tuning on Sintel and KITTI, we use the same training schedule as Chairs, however, learning rate is divided by 2 in all corresponding phases. Across stages, both the teacher and the student load latest updated checkpoint for initialization, and are random initialized from scratch before training on Chairs. The data augmentor $\mathcal{A}$ includes operations of random flipping, resizing, rotating, cropping and color jittering. According to hyperparameter search, weight coefficients of $\lambda_{1},\lambda_{2},\lambda_{3},\lambda_{4}$ in Eq. 4 and Eq. 7 are respectively set to $1.0,0.05,0.1,0.1$ for Chairs and Sintel, while being $5.0,0.05,0.02,0.1$ for KITTI. Note that the smoothness coefficient $\lambda_{1}$ in KITTI is 5 $\times$ of that in Chairs and Sintel in the first stage, while the photometric coefficient $\lambda_{3}$ in KITTI is only $1/5$ of that in Chairs and Sintel in the third stage. It is because that KITTI uses the second order smoothness, a larger coefficient can reach a reasonable regularization strength in the first stage. And in the third stage, the non-occluded region in KITTI benefits more from the distillation term rather than the noisy reconstruction term.

To verify the effectiveness of proposed contributions, we conduct thorough ablation experiments of reliable mutual knowledge distillation algorithm considering all training stages and different training settings, and our findings are reported in Table IV. All ablation experiments are trained on Sintel training or KITTI 2015 multi-view extension datasets with corresponding loss functions in their training stage. We first report results of unsupervised initialization (stage 1) for teacher model on Sintel training and KITTI 2015 training sets using Eq. 4 as the baseline, which is denoted as E1 in Table IV.

Finally, we conduct cross domain experiments to compare with recent top-performing unsupervised methods. As illustrated in Table VI, our method consistently performs best when training on one dataset and testing on another one, which demonstrates the state-of-the-art generalization ability of MDFlow among pyramid-based approaches. We attribute its excellent generalization performance to the efficient utilization of data augmentation on reliable pseudo labels during the decoupled mutual knowledge distillation process. Figure 7 presents qualitative results of the efficient teacher network on the cross domain DAVIS dataset [73]. It can be seen that proposed unsupervised learning approaches can estimate relatively good optical flow with clear motion boundary in diverse dynamic scenes.