DIFF-NST: Diffusion Interleaving For deFormable Neural Style Transfer

Prior work [34] has shown that early diffusion timesteps affect an image’s global structural and compositional information, whereas later timesteps affect local fine details. Inspired by this, we set out to determine which timesteps of the diffusion process control style and which control content.

Given a lack of research around exemplar-based Neural Style Transfer with diffusion models, we use a prompt-based model, prompt-to-prompt [8], to carry out this visualization. We use ChatGPT [2] to generate 20 content prompts, and we further define 10 style modifier prompts. With the prompt-to-prompt pipeline (operating over the Stable Diffusion LDM weights), we use the content prompts to generate reference content images, and we combine each content prompt with each style modifier prompt to re-generate the content images with the different explicitly defined styles still using prompt-to-prompt [8]. At the end of the process, we have 20 reference content example images and 200 "stylized" images. During the generation process, we extract attention values for analysis. We average the differences between the content example images’ attention values and each of their 10 stylized variants’, at each timestep. Fig 1 in the supplementary materials visualizes the average differences between these attention values at the diffusion timesteps. The red indicates a larger difference between the original content image and its stylized versions. Given that the structural and compositional information of the example content and their "stylized" counterparts is similar, we can infer that the stylistic differences relate to the higher attention discrepancies found at the later timesteps. This preliminary exploratory experiment clarifies the different effects of diffusion timesteps across the LDM generation process.

An additional preliminary experiment using these prompt-to-prompt images is an analysis of where the style information is captured in the LDM activations. We explicitly focus on the attention mechanism, where $\mathcal{Q}$, $\mathcal{K}$, and $\mathcal{V}$ values are used in the attention process [32]. We generate a base non-stylized image with the content prompt and then stylized variants with style modifier prompts. We extract attention values from the content-only prompt generation and replace the attention values of the stylized generation with those from the content-only generation. Doing so re-generates the original, non-stylized image. However, in our analysis, we observe that interpolating between the $\mathcal{V}$ self-attention values of the content/style-modified generations (while using only the original content values for the rest) can provide control over the stylization strength. From this experiment, we can infer that most, if not all, style information is captured from just the $\mathcal{V}$ self-attention values in the LDM. We visualize examples of this style interpolation in the supplementary materials.

Our work aims to perform style transfer of existing real user-provided images. As such, the re-styled synthesized image must stay faithful to the provided content image in terms of overall composition and structure. This means we must edit the image rather than re-generate a semantically similar approximation. We invert the content image through the LDM, similar to previous works such as prompt-to-prompt [8] and diffusion disentanglement [34]. This inversion process extracts the predicted noise at each timestep, as predicted by the UNet modules. To reconstruct the same image using an LDM, this content noise can be injected into the reverse diffusion process, replacing the LDM noise predictions at multiple timesteps. The more timesteps the noises are applied to, the better the reconstruction fidelity, with less freedom of input from the LDM. As shown in the diffusion disentanglement work [34], applying changes to the diffusion values from an earlier timestep allows more significant change in image structure.

Similar to these previous works, we use 50 time steps for the forward (inversion) and reverse (re-generation) diffusion processes. However, unlike these previous works, we interleave this noise starting from an earlier time, step 5, rather than 16, to improve reconstruction quality. We apply noise until step 45 instead of 50 to allow the model to self-correct some artifacts. Also, unlike prior work, we do not set the LDM predicted noises to zero for timesteps where pre-extracted content noises are not injected into the diffusion process. We aim to allow the model to generate new details to leverage its learned priors.

A notable trait of image-to-image and image-inversion with diffusion models is that color information is not disentangled from overall image structure across timesteps, as it is with feature activation across layers of a VGG model, for example. Thus, color information must be explicitly handled before inversion. Similar to previous works [38, 25], we pre-adjust the color of the content image through mean and covariance matching. We do this dynamically during training before inversion.

A final consideration is that we aim to perform prompt-less execution of LDMs, given our use of exemplar images for both content and style. As such, we only need to use the model’s unconditional capabilities. Latent Diffusion Models execute two iterations of their model: one with no prompt conditioning and one with prompt conditioning. The output of both branches is joined at every time step via the classifier free guidance (CFG). This exposes prompt control via this adjustable strength. Given that we aim not to use any text prompts anywhere in the process, we, therefore, altogether disable the prompt-conditioned branch of the model execution and use only the un-conditional branch for both inversion and reverse diffusion. The process would function the same if the text prompt were fixed to a generic prompt throughout or if CFG was zero, but this approach saves on compute.

We train a set of MLPs across each self-attention module in the LDM UNet blocks. We do not wish to re-train or fine-tune the LDM weights due to large compute/financial requirements. Instead, we train several smaller modules to hijack part of the LDM process, similar to how content noises are injected into the diffusion process. We directly target the attention process’s $\mathcal{V}$ values, generating brand new values for the remaining process to use. We chose the $\mathcal{V}$ values following our initial exploratory experiments with existing text-prompt-based diffusion image editing techniques such as prompt-to-prompt, where we observed that interpolation between $\mathcal{V}$ values only is enough to induce stylistic changes between content prompts and style-modified prompts.

Before our reverse diffusion process, similar to the real content image inversion to collect the noise predictions for reconstruction, we additionally invert and fully reconstruct the real style image through the LDM. This time, instead of collecting the predicted noises, we collect the predicted attention $\mathcal{V}$ values at every location and timestep and interleave them into the reverse diffusion process. Here, the MLPs generate the new $\mathcal{V}$ values based on an input consisting of the current $\mathcal{V}$ values, the corresponding $\mathcal{V}$ values at the same location and timestep of the style image, and the ALADIN style code of the style image, which we also pre-extract. We use both the style attention values and ALADIN, as this provides both global and local style information. Using only the attention values induces a similar style transfer. Anecdotally, however, using both sources of style information leads to a higher overall perceived quality of style transfer. We use the more recent ALADIN-NST [26] variant of ALADIN, as it is more disentangled, capturing less content information. This helps to avoid semantic content creeping into the stylized image from the style image, as shown in Fig 3.

A final consideration is that we only apply this attention manipulation process to the UNet decoder/upscaling layers, as per ControlNet [39]. Similar to their findings, we notice no perceivable differences in the output quality, but the VRAM consumption and compute costs are lower.

Diffusion models are typically trained one random timestep at a time, given the nature of focusing the training on noise predictions at individual timesteps. In our case, however, such timestep-localized deltas are not as easy to isolate. We can only guide our model during training based on the final de-noised output image. Moreover, well known existing style losses have been designed to operate in pixel space. They are, therefore, not directly applicable to latent space - though this may be an area of potential future study.

Therefore, we build our training process around unrolling the entire diffusion process, from starting to ending timesteps. We then decode the latent values into pixel space, where we can finally apply standard NST losses amongst the stylized and real style images from our style dataset. We opt to keep these style learning losses similar to previous works to reduce variables and uncertainty from our work. We follow a similar training objective to recent works such as NeAT [25], ContraAST [4], and CAST [41] - described in detail in Sec 3.5. We can report some negative results in using the LDM UNet as a noised feature extractor for computing a VGG-like style loss to avoid the unrolling process - the features extracted by the UNet did not accurately model the image style features.

We train our model using well explored training objectives from traditional NST methods to focus solely on the model technique - we most similarly follow training objectives resembling those of NeAT [25], ContraAST [4], and CAST [41]. Between style and stylized images, we use a VGG [30] style loss (Eq. 1), identity loss (Eq. 4), contrastive loss (Eq. 7), sobel-guided patch discriminator (Eq. 9), domain-level discriminator (Eq. 2), and ALADIN loss (Eq. 6). Between the stylized and content images, we use a perceptual loss (Eq. 3), contrastive loss (Eq. 8), and identity loss (Eq. 5). We use Sobel guidance for the patch discriminator, as per NeAT.

Equation 1 shows the VGG style loss, with $\mu$ and $\sigma$ representing the mean and standard deviation of extracted feature maps, $I_{s}$ represents style image from the style dataset $S$, $I_{c}$ represents a content image from the content dataset $C$ after the color adjustments, and $I_{sc}$ represents the stylized image.

$$\mathcal{L}_{s}:=\lambda_{\text{vgg}}\left(\sum_{i=1}^{L}\left\|\mu\left(\phi_{i}\left(I_{sc}\right)\right)-\mu\left(\phi_{i}\left(I_{s}\right)\right)\right\|_{2}+\left\|\sigma\left(\phi_{i}\left(I_{sc}\right)\right)-\sigma\left(\phi_{i}\left(I_{s}\right)\right)\right\|_{2}\right)$$

(1)

Eq 2 represents the domain-level adversarial loss, as per ContraAST [4], learning to discriminate between generated stylized images and real artworks. Here, a discriminator $\mathcal{D}$ operates over the stylized image, following our model $M$ modules. Eq 3 details standard perceptual loss, where $\phi_{i}$ represents the pre-trained VGG-19 layer index.

$$\mathcal{L}_{adv}:=\lambda_{\text{adv}}\left(\underset{I_{s}\sim S}{\mathbb{E}}\left[\log\left(\mathcal{D}\left(I_{s}\right)\right)\right]+\underset{I_{c}\sim C,I_{s}\sim S}{\mathbb{E}}\left[\log\left(1-\mathcal{D}\left(M\left(I_{s},I_{c}\right)\right)\right)\right]\right)$$

(2)

$$\mathcal{L}_{\text{percep}}:=\lambda_{\text{percep}}\left(\left\|\phi_{\text{conv4\_2 }}\left(I_{sc}\right)-\phi_{\text{conv4\_2}}\left(I_{c}\right)\right\|_{2}\right)$$

(3)

Eqs 4 and 5 show MSE identity losses between the reconstructed images and the style or content images, respectively. Eq 6 shows the ALADIN loss, with $\mathcal{A}$ representing the ALADIN model.

$$\mathcal{L}_{\text{id\_s }}:=\lambda_{\text{identity}}\left(\left\|I_{ss}-I_{s}\right\|_{2}\right)$$

(4)

$$\mathcal{L}_{\text{id\_c }}:=\lambda_{\text{identity}}\left(\left\|I_{cc}-I_{c}\right\|_{2}\right)$$

(5)

$$\mathcal{L}_{\text{aladin}}:=\lambda_{\text{aladin}}\left(\left\|\mathcal{A}(I_{sc})-\mathcal{A}(I_{s})\right\|_{2}\right)$$

(6)

Eqs 7 and 8 show contrastive losses as detailed in Sec 4.1, similar to [4] and [41], where $l_{s}$ and $l_{c}$ are extracted style/content embeddings respectively, using a projection head, and $\tau$ is the temperature hyper-parameter. The contrastive losses are applied over the averaged attention values per timestep.

$$\mathcal{L}_{s\text{\_contra }}:=\lambda_{\text{c}}\left(-\log\left(\frac{\exp\left(l_{s}\left(s_{i}c_{j}\right)^{T}l_{s}\left(s_{i}c_{x}\right)/\tau\right)}{\exp\left(l_{s}\left(s_{i}c_{j}\right)^{T}l_{s}\left(s_{i}c_{x}\right)/\tau\right)+\sum\exp\left(l_{s}\left(s_{i}c_{j}\right)^{T}l_{s}\left(s_{m}c_{n}\right)/\tau\right)}\right)\right)$$

(7)

$$\mathcal{L}_{c\text{\_contra }}:=\lambda_{\text{c}}\left(-\log\left(\frac{\exp\left(l_{c}\left(s_{i}c_{j}\right)^{T}l_{c}\left(s_{y}c_{j}\right)/\tau\right)}{\exp\left(l_{c}\left(s_{i}c_{j}\right)^{T}l_{c}\left(s_{y}c_{j}\right)/\tau\right)+\sum\exp\left(l_{c}\left(s_{i}c_{j}\right)^{T}l_{c}\left(s_{m}c_{n}\right)/\tau\right)}\right)\right)$$

(8)

The $\mathcal{L}_{\text{p }}$ term defined in Eq 9 is our patch discriminator $D_{\text{patch }}$ loss, guided by Sobel Maps ($SM$).

$$\mathcal{L}_{\text{p }}=\lambda_{\text{patch}}\left(\underset{I_{s}\sim S}{\mathbb{E}}[-\log(D_{\text{patch }}(\operatorname{crop}(I_{sc},SM_{sc}),\operatorname{crops}(I_{s},SM_{s})))]\right)$$

(9)

Our final combined loss objective is shown in 10 where each term is weighted by their respective $\lambda$ term. The loss weights are as follows: $\lambda_{\text{vgg}}=0.5$, $\lambda_{\text{adv}}=5$, $\lambda_{\text{percep}}=6$, $\lambda_{\text{identity}}=100$, $\lambda_{\text{aladin}}=10$, $\lambda_{\text{c}}=1$, $\lambda_{\text{patch}}=10$, $\lambda_{\text{1}}=0.25$, $\lambda_{\text{2}}=0.75$.

$$\mathcal{L}_{\text{final }}:=\mathcal{L}_{s}+\mathcal{L}_{\text{adv}}+\mathcal{L}_{\text{percep}}+\mathcal{L}_{\text{id\_s}}+\mathcal{L}_{\text{id\_c}}+\mathcal{L}_{\text{aladin}}+\mathcal{L}_{\text{s\_contra }}+\mathcal{L}_{\text{c\_contra }}+\lambda_{1}\mathcal{L}_{\text{p\_simple}}+\lambda_{2}\mathcal{L}_{\text{p\_complex}}$$

(10)

We undertake a pair of user studies to gauge real life human preference amongst our method and the baselines. First, we carry out an individual rating exercise, measuring the content fidelity between the content image and the stylized image, and separately measuring the style consistence compared to the style image. Second, we carry out a 5-way comparison, where we ask workers to select their best preference from randomly shuffled samples. We bin the ratings in the individual exercise to five levels, and we explicitly instruct what each rating level should represent. We include the definitions in the supplementary material. We randomly sub-sample 750 stylized samples from the test set and compare our method against each baseline on Amazon Mechanical Turk (AMT). We collect and average our responses over 5 different workers for each comparison, and show our results in Table 2.

The results indicate that workers are scoring our DIFF-NST method low on the content information, in both the ratings and preference studies. This is a positive result, as it highlights our technique’s more substantial content deformation. The only model which scored lower is PARASOL. However, as seen in our visual comparison figures, PARASOL tends to make significant conceptual changes to the depicted content. It is not so much a technique for style transfer as it is for style-inspired re-generation of similar semantic content. The results for our style-focused experiments indicate that workers prefer our method to baselines in both individual ratings and 5-way preference studies, which signifies a successful transfer of style while still deforming the content.

One key strength of our diffusion-based NST method is control over the structural deformity in the represented content concerning the style image. The reference content information is injected into the diffusion process by applying noises at each time step, pre-extracted from the content image inversion. With diffusion models, the early time steps strongly affect the significant structural components of the image, whereas the later timesteps affect lower level textural information. Therefore, by varying the starting timestep at which these pre-extracted content noises are applied, we can adjust, at inference time, how much the style should deform the content structure. This effect is difficult to evaluate quantitatively, but we show two examples in Fig 5.

An alternative vector of inference-time control is varying the diffusion timesteps in which our method’s attention replacement happens. By stopping at earlier timesteps, less style information is injected into the diffusion process, reducing the stylization strength. Unlike reducing content noise injection, this approach maintains the content structure better and more directly targets the style properties instead of structure. We show examples of this second approach in Fig 5, using the same example images as in Fig 5 for clarity.