💻 computer science

Color Conditional Generation with Sliced Wasserstein Guidance

The paper introduces SW-Guidance, a training-free method that enhances color-conditional image generation by modifying the diffusion sampling process with Sliced Wasserstein distance to ensure generated images simultaneously match a reference color palette and maintain semantic coherence with the text prompt.

Original authors: Alexander Lobashev, Maria Larchenko, Dmitry Guskov

Published 2026-05-04

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Alexander Lobashev, Maria Larchenko, Dmitry Guskov

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are an artist trying to paint a picture based on a description like "a cat sitting on a rug." You have a great imagination, but you struggle with colors. You want the cat to be a specific shade of "midnight blue" and the rug to be "burnt orange," but when you ask your AI assistant to draw it, the cat ends up being a random gray and the rug turns out to be a muddy brown.

You might think, "Okay, I'll just take the picture the AI made and paint over it to match the colors I want." But here's the problem: if you just force the colors to match, the AI might get confused. It might turn the cat's fur into a blue texture that looks like water, or make the rug look like a field of orange flowers. The colors are right, but the feeling of the objects is lost.

This is the problem the paper "Color Conditional Generation with Sliced Wasserstein Guidance" (or SW-Guidance) tries to solve.

The Problem: "Style" vs. "Color"

Current AI art tools are great at following instructions, but they are terrible at describing specific colors. If you say "turquoise," the AI might give you anything from light blue to green. If you try to fix this by taking a reference photo (like a photo of a sunset) and forcing the AI to copy its colors, older methods often copy everything from that photo. They might copy the sunset's clouds, the texture of the sand, or the way the light hits the water, even if you only wanted the colors.

The Solution: A "Color Filter" that Doesn't Smear the Image

The authors created a new method called SW-Guidance. Think of it as a very smart, invisible filter that sits inside the AI's brain while it is drawing.

Here is how it works, using a simple analogy:

1. The "Sliced" Approach (The Bread Analogy)
Imagine you have two loaves of bread. One is your AI's current drawing, and the other is your reference photo (the one with the perfect colors you want).

Old methods tried to compare the whole loaves at once. This is messy and often leads to the AI copying the wrong things (like the crust texture).
SW-Guidance takes a knife and slices both loaves into thin strips. It compares the color distribution of just one slice at a time. It does this from many different angles (like slicing the bread vertically, horizontally, and diagonally).
By looking at these "slices" (which the paper calls Sliced Wasserstein distance), the AI can see exactly how the colors are distributed without getting confused by the shapes or textures. It's like checking if the flavor of the bread matches, rather than checking if the crust looks the same.

2. The "Guidance" (The GPS Analogy)
As the AI draws the picture step-by-step (starting from static noise and slowly revealing an image), this method acts like a GPS.

Every few steps, the GPS checks: "Is the color palette of this drawing getting closer to the reference photo?"
If the AI is drifting off course (e.g., making the sky too green), the GPS gently nudges the drawing back toward the correct color distribution.
Crucially, it only nudges the colors. It doesn't tell the AI to change the cat into a dog or the rug into a forest. It keeps the "cat-ness" and "rug-ness" intact while fixing the paint.

Why This is Special

The paper claims this method is training-free. This means you don't have to teach the AI a new language or show it thousands of examples of "blue cats." You just give it the text prompt and the reference photo, and the math does the rest.

The Results:

Better Color Match: The colors in the final image are much closer to the reference photo than with other methods.
No Unwanted Copying: Unlike other methods that might accidentally copy the "vibe" or texture of the reference photo (like copying a brushstroke style), this method only copies the colors.
Stays True to the Prompt: The image still looks exactly like what you asked for in your text description (e.g., a cat on a rug), not a copy of the reference photo.

In Summary

The authors built a tool that lets you say, "Draw me a cat, but use the exact color palette from this photo of a sunset," without the AI accidentally turning the cat into a sunset or copying the clouds. It does this by mathematically "slicing" the color information and guiding the AI to match those slices perfectly, ensuring the colors are right without messing up the picture.

1. Problem Statement

The paper addresses the challenge of color-conditional image generation in text-to-image diffusion models. While natural language prompts can describe objects and scenes, they are often imprecise for specifying exact color palettes (e.g., "turquoise blue" yields varying tones). Existing solutions face two main limitations:

Post-processing Color Transfer: Applying color transfer algorithms (e.g., WCT, Histogram Matching) to images generated by diffusion models often results in semantic mismatches. The color style is transferred, but the underlying content may become distorted, or the colors may not align semantically with the generated objects.
Stylized Generation: Methods that modify attention mechanisms or weights (e.g., IP-Adapter, InstantStyle) to control style often entangle color with other stylistic features (texture, brush strokes, composition). This makes it difficult to control color independently without transferring unwanted high-level features from the reference image.

The goal is to generate an image that strictly adheres to a reference color palette while maintaining the semantic content dictated by the text prompt and avoiding the transfer of non-color stylistic elements.

2. Methodology: SW-Guidance

The authors propose SW-Guidance, a training-free approach that modifies the sampling process of a diffusion model. Instead of fine-tuning weights or using attention adapters, it guides the denoising trajectory using a differentiable loss function based on the Sliced 1-Wasserstein (SW1) distance.

Core Mechanism

The method operates within the reverse diffusion process (specifically using DDIM schedulers). At each timestep $t$ , the algorithm performs an optimization step to shift the latent representation toward the reference color distribution.

Latent Prediction: Given a noisy latent $x_t$ , the model predicts a clean latent $\hat{x}_0$ . This is decoded to an image $\hat{y}_0$ .
Color Distribution Extraction: The color distribution of the generated image ( $\pi_{\hat{y}_0}$ ) and the reference image ( $\pi_{ref}$ ) are extracted in RGB space.
Sliced Wasserstein Loss: The algorithm computes the Sliced 1-Wasserstein distance between these two distributions.
- Why Sliced? Standard Wasserstein distance is computationally expensive and non-differentiable in high dimensions. Sliced Wasserstein reduces the problem to 1D projections, where the distance is easy to compute and differentiable.
- Why 1-Wasserstein? The authors utilize a theoretical proposition (Proposition 1) stating that for 1D distributions, the 1-Wasserstein distance can be computed as the $L_1$ distance between Cumulative Distribution Functions (CDFs), avoiding the need to compute inverse CDFs (quantile functions).
- Differentiability: The CDF is approximated by sorting pixel values, a differentiable operation, allowing gradients to flow back to the latent space.
Gradient-Based Guidance: An auxiliary control vector $u$ is added to the latent $x_t$ . The algorithm optimizes $u$ over $M$ inner steps to minimize the SW distance:
$\min_u \text{SW}_1(\pi_{\hat{y}_0(x_t+u)}, \pi_{ref})$
Sampling Update: The optimized shift $u^*$ is applied to the latent, and the standard denoising step (with classifier-free guidance) proceeds to $t-1$ .

Theoretical Justification

The paper provides theoretical proofs (Lemma 2) showing that minimizing the Sliced Wasserstein distance ensures weak convergence of the generated color distribution to the reference distribution, including the convergence of all moments (mean, variance, etc.).

3. Key Contributions

Novel Conditioning Mechanism: First integration of the differentiable Sliced Wasserstein distance into the conditioning of diffusion models for color control.
Decoupling Color and Style: The method successfully transfers only the color palette, avoiding the transfer of textures, brush strokes, or composition found in the reference image (unlike attention-based stylization methods).
Training-Free: The approach requires no fine-tuning of the diffusion model weights, making it applicable to any pre-trained latent diffusion model (e.g., SD1.5, SDXL).
Theoretical Foundation: Provides proofs for the convergence properties of the optimization procedure and the efficiency of using 1-Wasserstein distance via CDF differences.

4. Experimental Results

The method was evaluated on SD1.5 and SDXL using the ContraStyles dataset (1000 prompts) and Unsplash Lite (1000 reference images).

Quantitative Metrics

Color Similarity: Measured by 2-Wasserstein distance between the generated and reference color distributions.
- Result: SW-Guidance achieved the lowest distance (0.0297) compared to all baselines, including advanced color transfer algorithms (e.g., PhotoWCT2: 0.1028) and stylization methods (e.g., InstantStyle: 0.1758).
Content Fidelity: Measured by CLIP-T (similarity between text prompt and image).
- Result: SW-Guidance achieved the highest CLIP-T score (0.270) among stylized generation methods, indicating it follows the text prompt better than methods that entangle style and content.
Image Quality: Measured by CLIP-IQA.
- Result: SW-Guidance maintained high quality (0.285), ranking second only to unconditional generation, significantly outperforming many style transfer baselines.

Qualitative Analysis

Visual Comparison: Visual examples show that while methods like InstantStyle transfer the "vibe" (e.g., brush strokes, photorealism) along with colors, SW-Guidance strictly adheres to the color palette while preserving the original composition and texture of the generated image.
Compatibility: The method is compatible with ControlNets (e.g., Canny, Depth), allowing for layout control alongside color conditioning.

5. Significance and Limitations

Significance:

Precision Control: It solves the "color naming" problem in AI art, allowing users to define a palette via a reference image rather than vague text descriptions.
Modularity: It offers a clean separation between content generation (text prompt), layout (ControlNet), and color (SW-Guidance), which is difficult with current monolithic style transfer models.
Efficiency: It adds minimal computational overhead (approx. 30 seconds for SD1.5 on an RTX 4090) compared to training-based methods or complex post-processing pipelines.

Limitations:

Text-Color Conflict: If the text prompt explicitly describes a color that contradicts the reference palette (e.g., prompt says "red apple" but reference is blue), visual artifacts may occur.
Differentiation Requirement: The current implementation requires differentiating through the U-Net, which is computationally intensive compared to non-differentiable post-processing.
Interaction with Stylization: Combining SW-Guidance with other strong stylization methods (like InstantStyle) can lead to conflicts, as both attempt to modify the generation process.

In conclusion, SW-Guidance represents a significant advancement in controllable image generation, offering a mathematically grounded, training-free solution for precise color transfer that preserves semantic integrity and image quality.