RLMiniStyler: Light-weight RL Style Agent for Arbitrary Sequential Neural Style Generation

1 Introduction

The goal of style transfer is to alter the style of an image while preserving its content. Arbitrary style transfer (AST), a key task in this domain, involves the challenge of using a single model to apply any desired artistic style to any given content. Since the pioneering work of Gatys et al.[Gatys et al., 2015] in neural style transfer, subsequent research[An et al., 2021; Wu et al., 2022; Deng et al., 2022; Kwon et al., 2023; Lin et al., 2021; Liu et al., 2021; Park and Lee, 2019; Wang et al., 2023] has made significant strides in enhancing model generalization capabilities, optimizing result quality, and accelerating inference speeds. Due to the varying preferences for the degree of stylization among individuals, precisely controlling the level of stylization to meet diverse needs is a challenging task. Mainstream approaches typically rely on manual tuning of hyperparameters to balance content and style, achieving results with varying degrees of stylization [Gatys et al., 2015; Huang and Belongie, 2017; Liu et al., 2021; Park and Lee, 2019], including the mixing ratio of content features to style features, as well as the individual weightings of content loss and style loss. However, repetitive process of trial and adjustment for achieving suitable weighting parameters, along with the complexity of networks with over 7 million parameters, limits their applicability. To simplify network models, MicroAST [Wang et al., 2023] employs a streamlined model without pre-trained networks for faster inference, and AesFA [Kwon et al., 2023] decomposes images into frequency components for efficient stylization. Even though these lightweight AST methods ensure computational efficiency and can perform style transfer on any style, they also require manual adjustment of hyperparameters and retraining to achieve varying degrees of stylization for specific styles. Importantly, achieving a good balance between content and style through manual hyperparameter tuning is challenging and often results in under-stylization or over-stylization. Hence, it is necessary to develop a new arbitrary style transfer technique that not only facilitates the transfer of any style but also offers a rich array of style degree options for each particular style, relies less on manual hyperparameter tuning, and remains computationally efficient. Recently, RL-NST [Feng et al., 2023] pioneered the application of reinforcement learning to the single style transfer task, achieving precise control over the degree of stylization for one specific style. But it struggles to distinguish between diverse styles, requiring retraining when faced with new styles.

∗Corresponding Author

This paper proposed a novel framework named RLMiniStyler that leverages reinforcement learning for controlling the process of the arbitrary style transfer using a unified policy and uncertainty-aware automatic multi-task learning. Leveraging the autonomous exploration inherent in reinforcement learning, our proposed method refines style expression, resulting in a diverse range of stylized results. By integrating a unified policy capable of effectively encoding both content and style images within a single neural network without feature confusion, RLMiniStyler can use one encoder to extract content and style features, thereby reducing model complexity and ensuring a consistent approach to learning and adaptation. Compared to using two encoders, this design is more conducive to stable training in the reinforcement learning process. Additionally, the uncertainty-aware automatic multi-task learning allows for dynamic adjustment of learning priorities based on the current performance state. Capable of rapidly generating a diverse array of results with varying degrees of stylization under limited resources, our method offers a richer visual experience beyond a singular result, as shown in Fig. 1.

RLMiniStyler empowers the agent to autonomously learn and explore various style transformation strategies without being constrained by pre-defined rules, resulting in more diverse and innovative stylized images. In summary, we summarize the main contributions of this work as follows:

• We present the first method of arbitrary style transfer based on reinforcement learning. RLMiniStyler provides a stable and flexible control of stylization. It allows flexible control over the degree of stylization by progressively incorporating style patterns into the results over time.
• We propose a unified policy within RLMiniStyler to ensure it remains sufficiently lightweight to operate efficiently in resource-constrained environments, while still maintaining high performance.
• We propose an uncertainty-aware, multi-task learning optimization strategy within our RLMiniStyler to automatically balance style learning and content preservation.
• Through comprehensive experiments on diverse image resolutions, we show the effectiveness of RLMiniStyler in creating high-quality and varied artistic sequences, showcasing its lightweight model advantage and superior or comparable performance across various evaluation metrics relative to both existing lightweight and state-of-the-art style transfer methods.

2 RelatedWork

Arbitrary Style Transfer (AST). AST aims to enable style transfer using a single trained model, achieving a balance between content and style across various style images without requiring additional training. While recent advancements [Deng et al., 2022; Gu et al., 2018; Hu et al., 2023; Huang and Belongie, 2017; Liu et al., 2021; Park and Lee, 2019; Wang et al., 2020] have been made in this area, many methods have complex models and offer limited diversity in stylization results. Although recently proposed lightweight methods [Wang et al., 2023; Kwon et al., 2023] have employed lightweight models, they necessitate retraining to realize results with varying degrees of stylization for a particular style. Using pruning techniques [Wu et al., 2024] can also achieve lightweight style transfer models, but this approach inevitably leads to a decline in style transfer performance, such as insufficient stylization.

Deep Reinforcement Learning for Neural Style Transfer. The agent in reinforcement learning (RL) focuses on developing optimal strategies through continual exploration and exploitation to maximize cumulative rewards. Handling high-dimensional continuous state and action spaces is particularly challenging for RL agents. Maximum Entropy Reinforcement Learning (MERL) methods [Haarnoja et al., 2017; Haarnoja et al., 2018; Hu et al., 2023; Zhao et al., 2019] demonstrate robust performance in high-dimensional continuous RL tasks by encouraging exploration. However, they may face limitations when applied to generative tasks such as Image-to-Image Translation (I2IT), as they are not inherently designed for generative models. SAEC [Luo et al., 2021], a framework that extends the traditional MERL approach, introduces a generative component to effectively handle I2IT tasks, but the 1D action space limits its effectiveness when attempting to process images with resolution higher than 128 × 128. Recently, RL-NST [Feng et al., 2023] has successfully extended SAEC to the style transfer task by expanding the action space to 2D and 3D. However, as a singlestyle transfer method, it requires retraining for each new style, making it unsuitable for the AST task.

3 Method

Existing AST methods usually use complex neural networks for one-step inference, limiting stylization diversity and restricting user preferences. Based on MERL framework [Haarnoja et al., 2018], we propose a novel, lightweight RL method for AST to enhance the richness of artistic stylization. In our method, style transfer is regarded as a sequential decision-making problem. In the guidance of a welldefined reward function, our RL agent selects optimal actions at each time step, and generates intermediate stylized results with varying style degrees accordingly. The overview of our method is shown in Fig. 2. Our approach includes three key components: the actor Pκ characterized by parameters κ, the builder Bτ characterized by parameters τ, and the critic Qδ characterized by parameters δ. The actor serves as the unified policy network responsible for making decisions and style feature extraction based on the current state composed of both the moving image and the style image, the builder acts as the generation network responsible for executing the actor's stylized decisions, and the critic acts as the scoring network responsible for evaluating the actor's decisions. The actor and the critic constitute the RL learning path for style control, while the actor and the builder constitute the generative learning part for generating stylized image. We next describe our method in details.

3.1 Deep Reinforcement NST Framework

In our RL environment Υ, CD and SD represent the content dataset and the style dataset, respectively. The state yt ∈ Υ is composed of two parts: the moving image I tm and the style image Is ∈ SD. The moving image I tm is initialized using the content image Ic ∈ CD. The action xt is determined by the agent based on its observation of the current state. To extract high-level abstract actions from the actor, we establish a model of stochastic latent actions conditioned on the current state (that is, the action xt follows the conditional probability xt = Pκ(xt |yt )). In practice, we employ the reparameterization technique [Kingma and Welling, 2013] to obtain these actions. The moving image I t+1 m in state yt+1 is created by the builder based on xt and current state yt . The reward rt is derived from the measurement of style discrepancy between the moving image I tm and the style image Is. The reward is inversely proportional to this discrepancy, such that a smaller style difference results in a larger reward.

3.2 Unified Policy for Efficient Style and Content Representation

Existing AST models [Deng et al., 2022; Huang and Belongie, 2017; Johnson et al., 2016; Kwon et al., 2023; Liu et al., 2021; Park and Lee, 2019; Wang et al., 2023] widely employ an encoder-decoder network as backbone architecture. Most of them [Gu et al., 2023; Huang and Belongie, 2017;

Li et al., 2017; Liu et al., 2021; Park and Lee, 2019] use pre-trained VGG [Simonyan and Zisserman, 2014] as the encoder due to its strong capability to capture a wide range of features useful for representing both content and style in images, as shown in Fig. 3(b). But the the complexity of the pretrained VGG can lead to substantial computational expenses and may introduce unwanted style patterns, like "eyes". An alternative approach [Deng et al., 2022; Wang et al., 2023; Kwon et al., 2023] involves training two encoders to independently process content and style images, treating them as separate distributions, as shown in Fig. 3(c). In this way, more appropriate encoding of content and style images is achieved to avoid incorrect style patterns. However, as shown in Fig. 3(a), there are no clear boundaries between content images and style images in practice. Hence, due to the inherent overlap between these two types of images being overlooked, using two separate encoders complicates the AST model and slows down the training process.

In light of this, RLMiniStyler leverages a unified policy for modeling content and encoding style using a single encoder, as shown in Fig. 3(d). To enable the encoder to have the ability to simultaneously process both content and style images, we draw inspiration from the StyleBank [Chen et al., 2017], which decouples content and style images through explicit style representation. Specifically, we integrated two additional style space dedicated to style encoding at different positions within the encoder's architecture, while the other parts of the encoder were used for general feature extraction, as shown in Fig. 2. This design not only maintains the efficiency of a single encoder but also enhances our control over the subtleties between content and style images by processing style features at different levels. Compared to the design of two separate encoders, the actor, capable of perceiving content and style simultaneously, can make more precise decisions based on the current state. In other words, we can more accurately manipulate the outcome of style transfer to achieve a richer variety of stylistic fusion effects.

3.3 Joint Learning

Our framework employs a joint learning strategy, integrating two mutually coordinated optimization processes: control learning and generative learning. In the control learning, our model learns control policies, while in the generative learning, it learns stylized image generation. Training alternates between these two parts. The generative learning consists of the Actor Pκ and the Builder Bτ , and the design of the loss function helps in effectively propagating gradient information between the Actor Pκ and the Builder Bτ . The control learning consists of the Actor Pκ and the Critic Qδ, and the training of the Actor can be conducted jointly through the control learning and the generative learning to ensure rapid and stable convergence. Algorithm in appendix describes the RLMiniStyler algorithm. All parameters are optimized based on the samples from replay pool D.

Control Learning

In the control learning, adhering to the MERL framework [Haarnoja et al., 2018], we iteratively refine a stochastic policy Pκ utilizing reward signals rt and soft Q-values Qδ(yt , xt ). Here, the action xt is generated by the actor in response to the current state yt , following the policy Pκ(xt |yt ). The soft Q-function Qδ(yt , xt ), computed by the critic network, provides an estimation of the expected cumulative reward for the state-action pair (yt , xt ) under the current policy. During the evaluation phase, we guide the improvement of the stochastic policy through the minimization of the soft Bellman residual, defined as:

where D is the replay pool and V¯ δ(yt ) = Ext∼Pκ [Q¯ δ(yt , xt ) - α logPκ(xt |yt )]. We set the reward signal rt as the negative of the style loss to ensure that the agent learns to control the stylization process. Given our objective is style-related, such a choice of reward function is reasonable. Specifically, the style loss serves as a simple and effective means to assess the similarity between the stylized output and the target style image, hence we set the reward function as the negative of the style loss rt = -LST , where the detailed definition of LST is shown in Eq. (5).

The target critic network, denoted as Q¯ δ, plays a crucial role in stabilizing the training process. The parameters of this network, represented as ¯ δ, are determined by calculating the exponential moving average of the parameters from the critic network [Lillicrap et al., 2015]: ¯ δ → ωδ + (1 - ω)¯ δ, with hyperparameter ω ∈ [0, 1]. To optimize JQ(δ), we use the gradient descent with respect to parameters δ as:

where ρQ is the learning rate. In the RL framework, the critic evaluates the actions taken by the actor, which in turn influences the policy decisions of the actor. Consequently, the following objective can be applied to minimize the KullbackLeibler (KL) divergence between the policy induced by the actor and a Boltzmann distribution, as determined by the Qfunction:

The last equation holds because xt can be evaluated by fP (nt , yt ), where nt is a noise vector sampled from a 3D Gaussian distribution with mean µ = 0 and standard deviation Σ = 1. Note that hyperparameter α can be automatically adjusted by using the method proposed in [Haarnoja et al., 2018]. Similarly, we apply the gradient descent method with the learning rate ρκ to optimize parameters as:

Generative Learning

Generative learning, through specific training strategies, enhances our model's generation capability in style transfer, ensuring the production of high-quality stylized image results. We assess the similarity between the stylized image I t+1 m and the input images by comparing their high-level features. Specifically, we employ a pre-trained VGG [Simonyan and Zisserman, 2014] as a feature extraction backbone ϕ to extract features independently from the moving image I tm, the style image Is, and the stylized image I t+1 m . We calculate the content loss LCO by comparing the semantic similarity between I t+1 m and I tm and the style loss LST by comparing the style similarity between I t+1 m and Is.

Content Loss. We evaluate how closely the stylized image resembles the content image, by maximizing perceptual similarity using the widely adopted perceptual loss [Johnson et al., 2016]. Let ϕ(j) denote the activation of the j-th layer, producing a feature map with dimensions C j × Hj × W j , where C j , Hj , and W j represent the number of channels, height, and width of the feature map, respectively. The content loss LCO is calculated by:

Style Loss. The style loss LST estimates the style deviations between the stylized image I t+1 m and style image Is. Let J represent the layer number of the network ϕ. It calculates statistical measures of µ and standard deviation σ to penalize I t+1 m , inspired by [Huang and Belongie, 2017]:

Hierarchical Style Representation Contrastive Loss (HSRCL). Recent studies [Chen et al., 2021; Wang et al., 2023] have demonstrated that the lightweight network struggles to fully capture and express the style features of style images in a single inference process, and incorporating a contrastive learning loss can mitigate this issue. For instance, MicroAST [Wang et al., 2023] employs a style signal contrastive learning loss to deal with this, but it predominantly relys on deep-layer features for contrastive learning, overlooking the contributions of shallow-layer features to the overall style representation. To this end, we introduce a novel hierarchical style representation contrastive loss, which integrates contrastive learning between deep and shallow feature representations, so as to enhance the style representation. More specifically, when sampling a batch of data from the replay buffer D, we construct both positive and negative sets for each sample's deep and shallow features. And the feature contrastive loss respectively computed from the deep features and shallow features are combined to create a hierarchical style contrastive loss function LCT , which is defined as:

where K represents the number of feature layers in the unified policy network, N represents the batch size. The batch comprises N states Y = {y1, y2, ..., yN}. Each state yi ∈ Y consists of a moving image Im and a style image Is. For each Im, we consider the style image Is from yi as a positive sample and the style images Is from other yj as negative samples, where j ̸= i.

Uncertainty-aware Automatic Multi-task Learning. A common way to enhance the quality of style transfer involves quantifying the semantic similarity to the content image and the style similarity to the style image through content and style loss functions, along with auxiliary loss functions, such as adversarial loss. But the weights for these loss functions are usually heuristically selected before training and remain unchanged throughout the training process, which is not sufficient enough to handle images with different style and content.

To this end, as inspired by [Kendall et al., 2018], we propose to use a multi-task learning framework that treats content learning, style learning, and contrastive learning as distinct but interconnected tasks. Using homoscedastic uncertainty, we dynamically adjust the loss weights of each task derived from a principled probabilistic model, achieving a balanced optimization objective that adapts throughout training. Unlike traditional methods requiring manual tuning of loss weights, our approach learns the relative importance of each task's loss function directly from the data. This not only simplifies the training process but also enables the dynamic modulation of the content loss and the style loss ratios to find the optimal solution.

Let λc, λs, λct denote the loss weights for content loss, style loss, and contrastive loss, respectively. These weights can adapt based on homoscedastic uncertainty σ2 1 , σ2 2 , and σ2 3 , reflecting the noise level or task confidence. And they are inversely proportional to the noise parameters. The final loss is:

where log(σ1σ2σ3) acts as a regularizer to prevent excessive increase in noise. Lastly, we employ a gradient descent method with the learning rate η to update the Actor and Builder parameters (κ and τ ) as well as σi(i = 1, 2, 3):

4 Experiments

4.1 Experimental Setup

Datasets and evaluation metric: Like most AST methods [Deng et al., 2022; Huang and Belongie, 2017; Liu et al., 2021; Park and Lee, 2019; Wang et al., 2023], we utilize the MS-COCO dataset [Lin et al., 2014] for content and the WikiArt dataset [Phillips and Mackintosh, 2011] for style. During training, images are first scaled to 512×512 pixels, then randomly cropped to 256×256, while testing can handle any input size. Following MicroAST [Wang et al., 2023], we assess all algorithms across seven aspects: visual effect, inference time, parameter count, content loss, style loss, SSIM [Wang et al., 2004], and storage space.

Implementation details: We use the Adam optimizer [Kingma and Ba, 2014] with a learning rate 2e-4, the batch size in the environment set to 1, and the batch size sampled from the replay buffer set to 8. All experiments are conducted on a single NVIDIA Tesla P100 (16GB) GPU.

4.2 Comparisons with Prior Arts

Baselines: We compare our method with four light-weight AST methods: CAP-VSTNet [Wen et al., 2023], MicroAST [Wang et al., 2023], AesFA [Kwon et al., 2023], and ICCP [Wu et al., 2024], as well as four state-of-the-art AST methods: AdaAttN [Liu et al., 2021], EFDM [Zhang et al., 2022], AesPA-Net [Hong et al., 2023] and UniST [Gu et al., 2023]. All codes used in the experiment are sourced from their respective public repositories, and we use the default settings provided.

Qualitative comparison: We visually compare our method with all baseline methods in Fig. 4. AdaAttN shows a repetitive style pattern resembling the eyes (the third row), while EFDM and CAP-VSTNet lose a significant semantic and structural content (first and second rows). AesPA-Net produces inconsistent results, especially in the eye area (first row). UniST, MicroAST and ICCP show insufficient stylization (third row), and AesFA has severe boundary artifacts (third row). In contrast, our approach generates a sequence of results with increasing stylization levels while maintaining coherent content structure. Our method has also been compared with lightweight baselines at higher resolutions (512, 4K). Due to space constraints, the detailed comparison results are included in the supplementary materials.

Quantitative comparison: Table 1 provides a comprehensive comparison between our approach and baseline models. Our method consistently achieves competitive scores in content loss, SSIM, style loss, and inference time, demonstrating its efficiency and effectiveness in producing outputs that balance style expression with content preservation. As the sequence progresses, our method enhances style richness while maintaining content fidelity. In terms of model complexity, our model outperforms the minimally pruned model in performance and features a lower parameter count and reduced complexity compared to the smallest non-pruned model. Similarly, more comparative results with lightweight methods at high resolutions (1K, 2K, 4K) are included in the supplementary materials. Additionally, it is the first AST method capable of automatically controlling the degree of stylization on images ranging from 256 to 4K resolution.

4.3 Ablation Study

With and without RL: We discussed the effectiveness of RL in style control. In Fig. 5, without RL, the Actor-Builder (AB) in (a) initially preserves semantic information. But as shown in (e), at sequence 10, notable content information is lost. In contrast, our method in (d) produces smoother and clearer stylized images from the start, and stably maintains highquality results throughout the sequence in (h) at sequence 10. This consistent performance highlights the significant enhancement of RL provides to DL-based AST models.

Automatic multi-task learning (AML) vs. manual settings: We manually tuned the loss weights in our method, based on the settings of MicroAST and empirical adjustments. Specifically, we set the content loss weight λc = 1, the style loss weight λs = 3, and the HSRCL loss weight λct = 3, while keeping all other settings unchanged. As shown in Fig. 5 , compared to the fixed loss weight method in (b,f), our approach using AML demonstrates superior content preservation in both sequence 1 in (d) and sequence 10 in (h). Our study indicates that AML significantly enhances model performance and accelerates network convergence.

Hierarchical style representation contrastive loss (HSRCL) vs. style signal contrastive loss: We investigated the effectiveness of HSRCL by comparing with the deepfeature based contrastive loss proposed in MicroAST [Wang et al., 2023]. As shown in Fig. 5, for sequence 1, using only deep features for contrastive learning (c) exhibits less of style diversity as compared with the result in (d). Comparing with sequence 10 in (h), there is a noticeable decline in (g) in terms of content affinity due to incoherent style expression. This experiment demonstrates that HSRCL significantly enhances the model's capacity in style expression.

4.4 User Study

We conducted user study on nine different methods. We recruited 30 participants representing a diverse range of ages, genders, and professional backgrounds. Each participant was randomly presented with 20 ballots: 10 at a 256 × 256 resolution and 10 at a 512 × 512 resolution. Each ballot included the content image, the style image, and 11 randomly shuffled stylized results. Note that since our method produces sequential results, we present the outcomes at the first, fifth, and tenth sequences. We collected 300 valid ballots for each resolution, and the detailed results are shown in Table 1. It is evident that the majority of users prefer the stylized results generated by our method. In other words, although the assessment of stylized results is inherently subjective, our lightweight style transfer agent is designed to generate a diverse array of sequential outputs tailored to meet the varying preferences and requirements of different users.

5 Conclusion

In this paper, we introduce a lightweight Arbitrary Style Transfer method using reinforcement learning. Our approach employs a unified policy to simultaneously learn from content and style images through a coherent encoding and decoding process, thereby more effectively capturing the distinguishing information between content and style. Our novel hierarchical style representation contrastive loss differentiates between shallow and deep style representations, enriching the expressiveness of the style transfer. Furthermore, Automatic Multi-task Learning facilitates training across various stages, accelerating the convergence of the model. Extensive experiments have demonstrated that our method not only generates visually harmonious and aesthetically pleasing artistic images across different resolutions but also produces a diverse range of stylized outcomes. The simplicity and effectiveness of our approach are expected to accelerate the miniaturization of style transfer networks. Although this work has successfully achieved miniaturization and diversification in arbitrary style transfer for images, the challenge remains in applying it to video, which involves temporal processing. Our future goal is to extend our approach to video arbitrary style transfer.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China under Grant 42375148, Sichuan province Key Technology Research and Development project under Grants (Nos.2024ZHCG0190, No. 2024ZHCG0176), CUIT Science and Technology Innovation Capacity Enhancement Program project under Grant KYQN202305.

Contribution Statement

The contributions of Jing Hu and Chengming Feng to this paper were equal.