Two Birds, One Stone: A Unified Framework for Joint Learning of
Image and Video Style Transfers

As discussed above, Deng et al. [9] split the input images into patches directly for transformer input with the necessary position encoding, resulting in low efficiency. In this work, we take advantage of CNNs to strengthen locality and improve the efficiency. Similar to [16], we use the pre-trained VGG-19 network [30] to extract feature maps of input images. Then, the $512$-dim vector at each pixel in the $relu4\_1$ layer is treated as a token for further transformer encoders. In this way, we combine CNNs and transformer to exploit both local and long-range dependencies. Meanwhile, there is no need to maintain the position encoding.

As in Figure 2, $DIT$ consists of three modules in turn, namely intra-domain extraction, video-image and style-content interaction. Specifically, the domain-specific extraction module is stacked with $N_{c}$, $N_{s}$ transformer encoder blocks for the $Seq_{c}$ and $I_{s}$ respectively. While content-style interaction is stacked with $N_{t}$ transformer decoder blocks.

In-Domain Extraction Based on the tokenization, the local context has already been extracted. So $DIT$ further exploits domain-specific information with a number of consecutively stacked transformer encoder blocks in Figure 2(b). Taking either $Seq_{c}$ or $I_{s}$ as input, the in-domain extraction module simultaneously capture the long-range information within the content and style domains. Normally, the transformer encoder block consists of a multi-head self-attention (MSA) layer and a feed-forward network (FFN). For the better efficiency, we develop the Axial Multi-head Self-Attention (AMSA) mechanism to replace MSA, combined with a $1\times 1$ convolutional layer for locality strengthening. And this mechanism is applied to all of the transformer blocks mentioned below. In addition, residual connections and layer normalization are deployed after each layer. The transformer encoder block is defined as:

	$\displaystyle\mathcal{S}^{\prime}$	$\displaystyle=\mathcal{AMSA}(\text{Conv}(Q),\text{Conv}(K),\text{Conv}(V))+Q,$		(1)
	$\displaystyle\mathcal{S}$	$\displaystyle=\mathcal{FFN}(\mathcal{S}^{\prime})+\mathcal{S}^{\prime},$		(2)

where $\mathcal{S}$ is the output sequence.

Video-Image Interaction. After domain-specific excavation, $DIT$ presents a symmetric module based on two transformer encoder blocks to interact contextual information between two types of content modalities. Notably, we split the input content sequence into two parts: one half is video $Seq_{v}$ and the other half is image $Seq_{I}$. As illustrated in Figure 2, one of the blocks takes the $Seq_{v}$ as $Q$, the $Seq_{I}$ as $K$, $V$, and the other one does the opposite. In this way, joint learning is performed for style transfer. And two types of content sequence are concatenated after interaction.

Content-Style Interaction. To finally capture the relevance between the content and the style domain, this module consists of a set of stacked transformer decoder blocks. As in Figure 2(c), each transformer decoder block takes the content as $Q$ while the style as the $K$ and $V$, alternately containing two ASMA layers and one FFN layer. Similarly, layer normalization and residual connections are applied after each layer. The transformer decoder block is defined as:

	$\displaystyle\mathcal{S}^{{}^{\prime\prime}}$	$\displaystyle=\mathcal{AMSA}(\text{Conv}(Q),\text{Conv}(K),\text{Conv}(V))+Q,$		(3)
	$\displaystyle\mathcal{S}^{\prime}$	$\displaystyle=\mathcal{AMSA}(\text{Conv}(\mathcal{S}^{{}^{\prime\prime}}),\text{Conv}(K),\text{Conv}(V))+\mathcal{S}^{{}^{\prime\prime}},$		(4)
	$\displaystyle\mathcal{S}$	$\displaystyle=\mathcal{FFN}(\mathcal{S}^{\prime})+\mathcal{S}^{\prime}.$		(5)

Inspired by [32], $DIT$ computes self-attention along a separate axis, rather than in feature maps like other transformer models [9, 40, 14, 11, 35]. The improved AMSA layer is specialized for style transfer without position encoding.

An AMSA layer consists of two MSA [31] layers in total, operating sequentially along the height and width axes. As in Figure 3, we first reshape $Q$, $K$, $V$ to separate the width dimension information, then use the first MSA layer to calculate self-attention along the height axis, and the output $f$ is normalized using the layer normalization ($LN$) layer. Following almost the same process as before, except that what we separate in the second MSA layer is the height dimension. Notably, we use the previous output $f$ as $Q$ and $V$, while keeping $K$ unchanged. Experiments in Table 4 verify this dedicated AMSA layer reduces the memory consumption and meanwhile improves inference efficiency.

The reasons behind our AMSA are two-fold. First, the information contained in the previous $K$, $V$ is crucial for the second MSA layer to learn contextual dependencies either within or across domains. Besides, the $V$ of the second MSA contains the last axial information, which is critical to ensure the stability of stylization. Experiments in Figure 10 show that our design effectively prevents artifacts caused by side effects of axial attention in style transfer.

In inference, UniST takes bimodal content input and outputs both style transfer results simultaneously, achieving two birds with one stone. Since the model has learned complementary content knowledge, UniST could also take unimodal content input into account by replacing the cross-attention with self-attention in the video-image interaction without compromising the quality of results. As in Table 6, unimodality obtains the consistent scores with bimodality.

The overall loss function is the weighted summation of the content loss $\mathcal{L}_{c}$, style loss $\mathcal{L}_{s}$, identity loss $\mathcal{L}_{\text{id}}$ and temporal loss $\mathcal{L}_{t}$:

$$\mathcal{L}=\lambda_{c}\mathcal{L}_{c}+\lambda_{s}\mathcal{L}_{s}+\mathcal{L}_{\text{id}}+\lambda_{t}\mathcal{L}_{t},$$

(6)

where $\lambda_{c}$, $\lambda_{s}$ and $\lambda_{t}$ are balancing factors. UniST uses the pre-trained VGG-19 to extract feature maps as $\phi=\{relu{1\_1},relu{2\_1},relu{3\_1},relu{4\_1}\}$. In our experiment, we use the above four layers with equal weights to calculate the loss below.

We use Euclidean distance [10] to compute content loss $\mathcal{L}_{c}$:

$$\footnotesize\mathcal{L}_{c}=\sum_{i=1}^{4}\left\|\phi_{i}(Seq_{\text{cs}})-\phi_{i}(Seq_{c})\right\|_{2}.$$

(7)

Following [16], the style loss $\mathcal{L}_{s}$ is defined as:

	$\displaystyle\mathcal{L}_{s}$	$\displaystyle=\sum_{i=1}^{4}\big{(}\left\|\mu(\phi_{i}(Seq_{\text{cs}}))-\mu(\phi_{i}(I_{s}))\right\|_{2}$		(8)
		$\displaystyle+\left\|\sigma(\phi_{i}(Seq_{\text{cs}}))-\sigma(\phi_{i}(I_{s}))\right\|_{2}\big{)},$

where $\mu(\cdot)$ and $\sigma(\cdot)$ denote the mean and variance of features separately. We further use the identity loss $\mathcal{L}_{id}$ [27] to promote more accurate content and style representation:

		$\displaystyle\mathcal{L}_{\text{id}}=\lambda_{\text{id1}}(\left\|Seq_{\text{cc}}-Seq_{c}\right\|_{2}+\left\|I_{\text{ss}}-I_{s}\right\|_{2})$		(9)
		$\displaystyle+\lambda_{\text{id2}}(\sum_{i=1}^{4}\left\|\phi_{i}(Seq_{\text{cc}})-\phi_{i}(Seq_{c})\right\|_{2}+\left\|\phi_{i}(I_{\text{ss}})-\phi_{i}(I_{s})\right\|_{2}),$

where $\lambda_{\text{id1}}$ and $\lambda_{\text{id2}}$ both are balancing factors. $Seq_{\text{cc}}$/$I_{\text{ss}}$ denote the results synthesized from two identical content sequences or style images. Figure 4 illustrates the identity loss for better understanding.

Following the recent work AdaAttN [25], we preserve the temporal consistency via a cross-image temporal loss. This loss promotes the cosine distance $D_{cs}$ between adjacent stylization frames to be closer to the cosine distance $D_{c}$ between adjacent origin frames.

$\displaystyle\mathcal{L}_{t}=\sum_{i=3}^{4}\phi_{i}(\frac{1}{N_{c_{1}}N_{c_{2}}}\sum_{m,n}\left|\frac{D_{c_{1},c_{2}}^{m,n}}{\sum_{m}D_{c_{1},c_{2}}^{m,n}}-\frac{D_{cs_{1},cs_{2}}^{m,n}}{\sum_{m}D_{cs_{1},cs_{2}}^{m,n}}\right|),$

(10)

where $D_{u,v}^{m,n}=1-\frac{F_{u}^{m}\cdot F_{v}^{n}}{\left\|F_{u}^{m}\right\|\times\left\|F_{v}^{n}\right\|}$. $N$ is the spatial dimension of the current feature map. $D_{u,v}^{m,n,x}$ measures cosine distance, and $F^{k}$ represents the feature vector of the $k$-th entry. Note, we only adopt layer $relu3\_1$ and $relu4\_1$ to calculate $\mathcal{L}_{t}$. Meanwhile, in each training iteration, we compute the temporal loss between the consecutive video frames. Noted that temporal information is implicitly encoded in this way.

UniST is trained with MS-COCO [24] (image), MPI [2] (video) as the content datasets and WikiArt [28] as the style datasets. For $Seq_{c}$, the ratio of the two types of content is $1:1$, and the total number is $6$. In the training phase, all images are loaded with the resolution of $256\times 256$. While in the inference phase, UniST can be applied to any resolutions. Therefore, we manage to extend the UniST to multi-granularity style transfer (see supplementary material). In experiments, we adopt the resolutions of $1024\times 1024$ and $1024\times 2048$ for image and video respectively. For the loss function, the balancing factors $\lambda_{c}$, $\lambda_{s}$, $\lambda_{t}$, $\lambda_{\text{id1}}$, $\lambda_{\text{id2}}$ are set as $0.1$, $1.5$, $90$, $0.1$, $0.5$, empirically. The number of transformer blocks $N_{c},N_{s},N_{t}$ are empirically set as $2,1,3$. Our network is trained for $50K$ iterations on a NVIDIA Tesla V100 GPU with a batch size of $3$. We use the Ranger optimizer [34] with initial learning rate of $0.00005$.

We compare our method with $13$ state-of-the-art methods, including AdaIN [16], SANet [27], Linear [21], MCCNet [7], MAST [8], Artflow [1], AdaAttN [25], StyleFormer [37], IEST [4], Stytr² [9], CCPL [38], epsAM [5], MicroAST [33], MCCNetV2 [20].

Quantitative comparison. We randomly collect $17,124$ content images from ImageNet [6] and $17,124$ style images from WikiArt [28] which are separated from the training set to generate $17,124$ stylization results. Similar to Stytr² [9], we use the mean euclidean distance and the mean instance statistics difference mentioned in section 3.5 as metrics for content preservation and stylization degree. Furthermore, we conduct the color distribution experiments by adopting color loss in DSLR [17] and adopt the Gram matrices [13] for texture difference. As in Table 1, compared with existing methods, UniST achieves the best performance in both content and style differences and obtains promising results on texture and color differences of style aspects.

Qualitative comparison. As in Figure 5, AdaIN [16] transfers the style patterns but losses important content details (1st, 2nd rows). SANet [27] fails to align the distribution of the style patterns, leading to the distorted object boundaries (1st, 3rd, 5th rows) and inconsistent content background structures (2nd, 4th rows). The stylization of Linear [21] is not satisfactory enough, resulting relatively light migration effects such as pink in the 5th row. As with the linear transformation, MCCNet [7] learns the correlation of each channel through the transformed self-attention, slightly improving the transfer effect, but there are serious overflow problems around the object boundaries (1st, 3rd, 4th rows). MAST [8] distorts the content background structure with excessive style transfer (3rd, 4th rows). Based on WCT [22] in our setting, Artflow [1] leads to conspicuous vertical artifacts at the edges of the generated results(3rd, 5th rows). AdaAttN [25] does the decent foreground style transfer, but requires further rendering of the background (3rd, 5th rows). StyleFormer [37] provides vivid style patterns, but the colors of the results are inconsistent with the style reference images (1st, 3rd, 4th rows). On the contrary, IEST [4] presents stylization results with consistent colors, lacking more desirable style patterns (3rd, 4th, 5th rows). Stytr² [9] fails to preserve the background structure of the content inference image (2nd, 5th rows). CCPL [38] transfers the style patterns with a lot of vertical artifacts (1st, 3rd, 4th rows). In contrast, based on the inductive bias of CNNs, our method captures the relevance between content and style sequences jointly and efficiently. As a result, we provide vivid stylized results with desirable style pattern details, while keeping the content structure well-maintained.

For video style transfer, we compare our method with four state-of-the-art methods including Linear [21], MCCNet [7], AdaAttN [25], and CCPL [38]. Note, optical flow is not used for stabilization when conducting comparison.

Quantitative comparison. Following Liu et al. [25], we adopt the official optical flow [2] to wrap the output stylized frame and compute the per-pixel difference between warped and stylized frames. Meanwhile, we use the LPIPS [41] to measure the diversity of adjacent stylized frames, with smaller values indicating better consistency.

In practise, the style input is fixed at $512\times 512$, and the content input is fixed at $256\times 256$. Table 2 presents optical flow error and LPIPS metrics for $20$ styles over $23$ videos of compared methods. UniST achieves the best scores in both metrics and thus has the best temporal consistency.

Qualitative comparison. Figure 7 shows the results of the video qualitative comparison. To better verify long-term temporal consistency, we take the $5$-th, $15$-th and $25$-th frames of the example video as the reference content images, where the character in the video has large movements. The results show that all four methods visually satisfy the long-term temporal consistency. To be specific, Linear [21] uses the shallower feature map in video style transfer task, sacrificing the stylization effects for temporal consistency in some way. MCCNet [7] produces distorted results with severe artifacts along object contours, where the same drawback appears in image style transfer. AdaAttN [25] provides a well-transferred foreground leaving the background requires further rendering. Using the same model as image style transfer, CCPL [38] also shows a huge number of vertical artifacts in video stylization results. In contrast, benefiting from the joint learning framework, our video results are enhanced by the rich texture from images beyond the reference video. Besides, the difference map in column $4$ shows that our joint learning framework is stable enough when dealing with the motion blur. Based on the above analysis, our method can generate video results with more pleasing stylistic patterns while maintaining temporal consistency well. Notably, when applied to image style transfer, some of these compared methods [21, 25] need to be retrained with extra changes and consumption, while the rest [7, 38] use the same model directly with the flawed results. In contrast, our joint learning framework can perform well on both image and video style transfer in one go. Additional vivid stylization results are provided in our supplementary materials.

Furthermore, we carry out the user study experiment for comparison. Specifically, we use $27$ content images and $21$ style images to synthesize $567$ images in total. Given $20$ randomly combinations of content and style, the generated results obtained by $12$ image-based style transfer methods. Then, we ask $100$ participants to select their favorite one from three aspects: content preservation, style transfer, and overall preference. We collect $2,000$ votes and show results in Figure 6(a). The results verifies the superiority of our method over other models for image style transfer. Similarly, we take $12$ videos of $50$ frames and $21$ style images to synthesize $252$ video stylized clips in total. Given $4$ random combinations of video and style, the stylization clips obtained by $5$ video-based style transfer solutions. Then, we ask another $50$ subjects to select their favorite one from three views: temporal consistency, stylization effect (considering both content preservation and stylization degree), and overall preference. We collect $200$ votes and our method is selected as the best as shown in Figure 6(b).

For participants, there are 83 males and 67 and females (55/45 males/females for image, the other 28/22 males/females for video), aged from 23 to 42.

In Table 3, we report the inference time of our joint learning method and other approaches. Note that all the methods are run using a single TITAN XP GPU card. Although our joint learning framework consists of many stacked self-attention layers, our model can still achieves $35$ FPS at $256$px, which is comparable with attention-based methods such as AdaAttN [25] and Stytr² [9]. The main limitation of our work is the inference speed is limited when applied for high resolution input, which may hinder its usage.

Notably, our AMSA mechanism is crucial for the joint learning framework to address the expensive memory consumption and huge computation complexity. In particular, the computation complexity can be reduced from $\bm{O}$($\bm{H^{2}}\times\bm{W^{2}}$) to $\bm{O}$($\bm{H^{2}}+\bm{W^{2}}$) theoretically. As shown in Table 4, AMSA is much lighter compared to MSA under the same experiment setting, demonstrating the great potential for high-resolution image and long video style transfer. Although the framework consists of many self-attention layers, the UniST is still fully capable of practical applications.

Video-image interaction module. Since images and videos have different domain features, we design a cross-attention domain-interaction module for joint learning of two domains for better content knowledge, which improves the overall stylization results. We conduct an experiment by removing the cross-attention interaction module. As in the 2nd, 3rd columns of Table 5, both content preservation and style transfer are highly degraded without this module.

To further demonstrate the superiority of UniST, we compare 3 training strategies without video-image interaction module. As shown in the 4th, 5th, 6th columns of Table 5, UniST achieves the best performance compared with training separately and sequentially. Meanwhile, as shown in Figure 8, video style transfer is improved by learning complex appearances and textures from the image domain with UniST framework. Similarly, Figure 9 shows UniST enhances image stylization effects with more vivid results.

Inconsistent datasets concerns. In our method, we adopt two datasets from different modalities for learning. To ensure fairness, we have retrained five image style transfer models, while keeping the training dataset consistent with us. We provide the quantitative comparison in Table 7. Meanwhile, we conduct extra experiments by training the image-only UniST with different dataset scales in Table 8. The results show simply using more images for image-only UniST brings no obvious improvements, while gains come from different modalities and beneficial interactions, and further illustrating the superiority of our model.

Axial Multi-head Self-Attention. To verify the effectiveness of AMSA layer, we conduct two different design comparisons in Figure 10. First, the information contained in the previous $K$, $V$ is crucial for the second MSA layer to learn contextual dependencies either within or across domains. Therefore, we uniformly take the output of the first MSA as the $Q$, $K$, $V$ for the second MSA. The results in 4th column show that it is difficult for the model to capture the relevance between content and style domains in this way. Second, the first axial information is necessary for the $V$ of the second MSA to maintain the stable stylization results. To demonstrate this, we use the output of the first MSA as the $Q$ for the second MSA, while keeping the $K$ and $V$ unchanged. As in the 5th column, the style patterns transferred are mixed with bar artifacts. From the results in the 3th column, the original design in Figure 3 is necessary for our model to prevent side effects caused by axial attention.