Optical Flow Distillation: Towards Efficient and Stable Video Style Transfer

Neural style transfer is one of the most popular research hotspots in recent years. Gatys et al. [10, 11] used CNN to iteratively reconstruct a stylized image by minimizing the difference between the target image, the content image and the style image in high-level features. These methods solve the optimization by backward propagation and are computationally expensive. To make the inference more efficient, Johnson et al. [22] proposed a feed-forward network to stylize images, which replaces the iterative process of optimizing pictures with the optimization of CNNs via training.

Video style transfer is attracting more and more research interests. Researchers tried to utilize the inter-frame temporal relation to improve the visual stability of stylized videos, specifically motion estimation based on optical flow. Ruder et al. [33] initialized the optimization of the current frame with stylized output of the previous frame and proposed temporal loss which uses optical flow to maintain inter-frame consistency. This image based optimization algorithm outputs a very stable video but costs about 3 minutes to process a frame even with precomputed optical flow. Therefore, fast video style transfer is mainly based on model optimization. Ruder et al. [34] proposed a framework to use optical flow both in the training stage and in the inference stage to improve temporal consistency of output stylized videos. This framework contains two networks. The first network obtains the first frame of the video as input and outputs the stylized result. The second network obtains three inputs, including the current frame, the previous stylized frame warped by the optical flow and the mask which indicates motion boundaries and outputs the stylized result of current frame. Similarly, the architecture in [2] utilized optical flow both in training stage and inference stage. All these methods [34, 2] got stable stylized video but can not be used for real-time video style transfer. To address this problem, another family of video style transfer methods [8, 19] utilize optical flow only in the training stage thus speed up the inference. These methods utilize similar temporal loss to train the feed-forward transform network to improve the temporal consistency of output videos. They get rid of computing optical flows on the fly but produce less stable stylized results than those networks that adopt optical flows during inference stage [34, 19]. Our method aims to mitigate the gap between optical flow based and optical flow free methods for video style transfer via knowledge distillation.

Knowledge distillation is a technique that leverages intrinsic information of teacher network to train a smaller one, which is first pioneered by Hinton et al. [18]. Since then many algorithms have been proposed to improve knowledge distillation [20, 43, 6, 32, 16]. Wang et al. [40] exploited generative adversarial network to encourage the student network to learn similar feature distribution with teacher network. Zagoruyko and Komodakis [44] proposed to utilize spatial attention for distilling intermediate latent features of the network. Heo et al. [17] proposed a novel activation transfer loss to distill knowledge of activation boundaries from the teacher network. Chen et al. [5] introduced the locality preserving loss to preserve relationships between samples in high dimensional features from teacher network and low dimensional features from student network. Knowledge distillation has also been adopted in many other applications, e.g., object detection [23, 41, 4], semantic segmentation [21, 26, 14] and pose regression [35, 38]. However, distilling knowledge from a stable video style transfer to a lightweight student network is not yet explored, which is the main purpose of this paper.

Here we briefly revisit the perceptual loss introduced by Johnson et al. [22], which has been widely used for style transfer algorithms. The perceptual loss includes the content loss to encourage the output stylized image to have similar content representations with the input image, and the style loss to capture the style information. In addition, total variation regularization ($\mathcal{L}_{tv}$) is generally introduced to encourage spatial smoothness in the stylized images. Therefore, the basic style transfer perceptual loss contains three terms:

where $I_{s}$ is the given style image and $x$ is an input image. $\lambda_{c}$, $\lambda_{s}$ and $\lambda_{tv}$ are hyper parameters to balance three different losses. To simplify the notations, we omit the $I_{s}$ and denote the perceptual loss as $\mathcal{L}_{percep}(x)$ in the following sections.

This basic style transfer loss can be exploited to train a student network from scratch, which is denoted as $\tilde{\mathcal{N}}_{S}$. Since $\tilde{\mathcal{N}}_{S}$ is trained frame by frame, it suffers from flickering in consecutive frames and produces unstable stylized videos. Our goal is transferring the knowledge of a teacher network $\mathcal{N}_{T}$ to train the desired student network $\mathcal{N}_{S}$, so that student network produces coherent stylized videos.

A straightforward way to train the desired student network $\mathcal{N}_{S}$ via knowledge distillation is directly using the output stylized image of teacher network to teach the student network, as shown in the following loss function:

where $\mathcal{N}_{T}(x,f)$ and $\mathcal{N}_{S}(x)$ are the output stylized images for input image $x$ of teacher network and student network respectively, $f$ is the corresponding optical flow.

However, this strategy produces blurred stylized images, as shown in Fig. 2. The reasons behind are two folds. Firstly, it is a challenging optimization problem to force the output of student network to align with that of the teacher network at pixel level. More importantly, due to the difference of network architectures between student network and teacher network, the style of the output images from student network may slightly differ from those produced by teacher network. Therefore, directly distilling knowledge from the output stylized images of teacher network is hard to train a good student network.

To address the above problem, we propose the residual distillation loss for better knowledge distillation. Our goal is to train an optical flow free student network from the knowledge of an optical flow based teacher network. Therefore, the knowledge of the difference between teacher networks with and without optical flow is the key information to train a stable student network. Let $\tilde{\mathcal{N}}_{T}$ denotes the model similar to teacher model $\mathcal{N}_{T}$ but does not adopt optical flow for video style transfer and $\tilde{\mathcal{N}}_{T}(x)$ is the output of $\tilde{\mathcal{N}}_{T}$ for input image $x$. Obviously the output stylized video flickers since basically $\tilde{\mathcal{N}}_{T}$ just predicts stylized images frame by frame. The difference between the output of $\mathcal{N}_{T}$ and $\tilde{\mathcal{N}}_{T}$ is:

$\Delta\mathcal{T}(x)$ is attributed to additional temporal consistency that is brought by optical flows to the output stylized videos, which is the key information for student network to improve the stability.

$\tilde{\mathcal{N}_{S}}$ is the student network that is trained only using basic style transfer loss as in Eq. (1) and $\mathcal{N}_{S}$ is the stable student network we want to obtain. Thus the improvement of $\mathcal{N}_{S}$ over a vanilla baseline is:

We encourage the student network to learn how to improve the stability of output videos over the baseline model by forcing $\Delta\mathcal{S}$ to imitate $\Delta\mathcal{T}$. The residual distillation loss is formulated as follow:

The benefits of above residual distillation loss are two folds. Firstly, $\Delta\mathcal{T}(x)$ is the key information for the teacher network to become stable and could benefit the temporal consistency of student network. What’s more, $\Delta\mathcal{T}(x)$ and $\Delta\mathcal{S}(x)$ have eliminated the difference of stylized images brought by structures, thus it could alleviate the blurring effect as shown in Fig. 2.

In the above section, Eq. (5) is proposed for distilling knowledge from teacher network for one frame. We further develop a distillation loss by exploiting the temporal property of the video.

For a stylized video that is temporally consistent and stable, the same regions that are not located at the occluded regions or motion boundaries are supposed to have similar stylized patterns and strokes. Therefore, a basic assumption for a stable stylized video is that the non-occluded regions should be low rank representations.

Suppose we have $K$ consecutive frames $\{x_{t}\}_{t=1}^{K}$ and the corresponding optical flows $\{f_{t}\}_{t=1}^{K}$, the output stylized frames from the student network are calculated as $\{\mathcal{N}_{S}(x_{t})\}_{t=1}^{K}$. We warp all stylized frames to a specific frame $\mathcal{N}_{S}(x_{\tau})$ using the optical flows, where $x_{\tau}$ is often chosen as the middle frame, i.e., $\tau=\lfloor K/2\rfloor$. We denote the warped frames as $\mathcal{W}_{t\rightarrow\tau}({\mathcal{N}}_{S}(x_{t}),f_{t})$, where $\mathcal{W}_{t\rightarrow\tau}(\cdot)$ means warping the $t^{th}$ frame to $\tau^{th}$ frame. $\mathcal{M}_{t}$ is the occlusion mask for $x_{t}$, where $0$ indicates the motion regions or boundaries and $1$ indicates the traceable regions. In this way, if we put all traceable regions $\mathcal{R}_{t}=\mathcal{M}_{t}\odot\mathcal{W}_{t\rightarrow\tau}(\mathcal{N}_{S}(x_{t}),f_{t})$ into a matrix $\mathcal{X}$, then $\mathcal{X}$ is low rank, i.e.,

where $L=H\times W$ is the number of pixels in the images and $vec(\cdot)$ means transforming a two dimensional image into a one-dimensional vector.

We can get the Singular Value Decomposition (SVD) of $\mathcal{X}$ by:

where $\mathcal{U}\in\mathbb{R}^{K\times K}$ and $\mathcal{V}\in\mathbb{R}^{L\times L}$ are orthogonal matrices and $\Sigma\in\mathbb{R}^{K\times L}$ is the singular value matrix. The diagonal elements in $\Sigma$ are singular values $\Gamma=\{\gamma_{0},\ldots,\gamma_{K}\}$ of the matrix $\mathcal{X}$. The rank of $\mathcal{X}$ is calculated as:

where $\mathbb{I}(\cdot)$ is the indicator function.

However, the rank of a matrix is a non-differentiable function and can not directly be optimized via CNNs. Therefore, we instead adopt nuclear norm of $\mathcal{X}$, which is a convex relaxation of the rank function. The nuclear norm $||\cdot||_{*}$ is given by:

where $\gamma_{i}$ is the $i^{th}$ singular value of $\mathcal{X}$.

We expect that the rank of output stylized videos from student network should be similar with the stable videos. Intuitively, a straightforward way is to encourage the student network to imitate the nuclear norm of teacher network by the following low-rank distillation loss:

where $||\mathcal{X}_{T}||_{*}$ and $||\mathcal{X}_{S}||_{*}$ are nuclear norms of output stylized videos from teacher network and student network, respectively. Distilling knowledge of low rank from teacher network may help to improve the results of student network. However, the stability of output videos from the teacher network is still worse than the input videos. The stability of input video is a better teacher for the student network. To this end, we propose to distill the rank of input videos to train the student network by the following distillation loss:

where $||\mathcal{X}_{input}||_{*}$ is the nuclear norm of the input video. We utilize $\mathcal{L}_{rank}^{input}$ to train the student network and will further discuss the different low rank losses in experiments.

Temporal consistency loss imposes constraints on consecutive output frames and is widely used in prior methods [19, 2, 8], which is formulated as follows:

Using temporal loss improves the temporal consistency of student network and thus serves as a stronger baseline than vanilla perceptual loss. We will demonstrate the effectiveness of our proposed method on both two baselines in experiments.

To train student network with perceptual loss or residual distillation loss, the network only needs the current input frame. However, calculating low rank loss and temporal loss needs $K$ consecutive frames. Suppose the input video segments are $x=\{x_{t}\}_{t=1}^{K}$ and the corresponding optical flows are $f=\{f_{t}\}_{t=1}^{K}$. The desired student network can be optimized using the total distillation loss as follow:

where $\lambda_{res}$, $\lambda_{temp}$ and $\lambda_{rank}$ are hyper parameters to balance different terms of losses.

The $\mathcal{L}_{percep}$ and $\mathcal{L}_{temp}$ in Eq. (3.4) are used as our baseline for training the student network from scratch. We will explore the baseline with and without the temporal loss. The other two terms are knowledge distillation losses proposed in this paper and we will demonstrate in experiments that these distillation losses improve the stability of output stylized videos.

We use the Hollywood2 video scene dataset [31] as the training data and evaluate our method on MPI Sintel dataset [1]. We follow the same data preprocessing methods as in [34] and randomly sample 2000 tuples consisting of 5 consecutive frames from Hollywood2 dataset. MPI Sintel dataset provides ground truth optical flow and occlusion masks, which is widely used for the task of optical flow estimation and is also adopted to evaluate the temporal consistency of video style transfer. Following prior work [34, 8, 19], we evaluate our method on five videos in the MPI Sintel dataset.

During training all frames are downscaled to $640\times 360$ and the input size for evaluation is $1024\times 436$. We train the student network with learning rate of $10^{-3}$ and a batch size of 1 using Adam optimizer for 10 epochs. The learning rate is decayed by 1.2 in every 500 iterations. The hyper-parameters in Eq. (3.4) are set to be $K=5$, $\lambda_{res}=4\times 10^{8}$, $\lambda_{temp}=1\times 10^{6}$ and $\lambda_{rank}=1\times 10^{2}$.

Following the quantitative evaluation metric in prior methods [8, 19, 2], we utilize $e_{stab}$ to evaluate the temporal consistency of the output stylized videos. $e_{stab}$ is the square root of temporal errors between consecutive frames for the traceable regions of the videos:

where $N$ is the number of frames in the testing video, $y_{t}$ and $y_{t-1}$ are output stylized images for frame $t$ and $t-1$ respectively, $D$ is the number of pixels of output stylized image.