🤖 machine learning

ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control

ParetoSlider is a multi-objective reinforcement learning framework that trains a single diffusion model to approximate the entire Pareto front, enabling users to navigate optimal trade-offs between conflicting generative goals at inference time without retraining or maintaining multiple checkpoints.

Original authors: Shelly Golan, Michael Finkelson, Ariel Bereslavsky, Yotam Nitzan, Or Patashnik

Published 2026-04-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Shelly Golan, Michael Finkelson, Ariel Bereslavsky, Yotam Nitzan, Or Patashnik

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 a chef trying to cook the perfect dish. You have two competing goals: you want the food to be healthy (low calories) but also delicious (highly seasoned).

In the past, if you asked a cooking AI to help you, you'd have to pick a single recipe beforehand. You'd say, "Make me a dish that is 70% healthy and 30% tasty." The AI would cook that specific version, and if you later changed your mind and wanted 50/50, you'd have to throw away that pot and start cooking a completely new one from scratch. This is how current AI image generators work: they are locked into one specific "flavor" of compromise.

ParetoSlider is like giving that chef a magic slider that lets you adjust the taste while the food is being served, without ever needing to restart the cooking process.

Here is a simple breakdown of how it works, using everyday analogies:

1. The Problem: The "Fixed Recipe" Trap

Current AI models are trained using a method called "Reinforcement Learning." Think of this as a teacher giving a student a grade.

The Old Way: The teacher gives the student a single grade based on a fixed formula (e.g., 50% Health + 50% Taste). The student studies hard to get that specific grade. If you want a different balance later, the student has to go back to school and learn a whole new curriculum.
The Result: You end up with many different AI models, each good at one specific compromise, but none of them can adapt on the fly.

2. The Solution: The "Magic Slider" (ParetoSlider)

The researchers created a new system called ParetoSlider. Instead of teaching the AI one fixed recipe, they taught it every possible recipe at once.

The Analogy: Imagine a dimmer switch for a light, but instead of just "bright" or "dim," it controls a spectrum of colors.
How it works: The AI is trained with a "preference vector" (a set of dials). You can tell the AI, "I want 80% photorealism and 20% sketch style," or "I want 50% preservation of the original photo and 50% new editing."
The Magic: Because the AI learned the entire spectrum of possibilities during training, you can slide that dial back and forth at the very last second (inference time). The AI instantly adjusts its output to match your new preference without needing to retrain or switch models.

3. The Secret Sauce: "Late Scalarization"

One of the biggest challenges in teaching an AI to balance multiple goals is that some goals are "louder" than others.

The Analogy: Imagine a band where the drummer is playing at 100 decibels and the violinist is playing at 10 decibels. If you just mix the audio, the drums will drown out the violin, and the AI will only learn to be a drummer.
The Fix: The paper introduces a technique called Late Scalarization. Before mixing the signals, the AI "normalizes" them. It turns down the volume of the loud drums and turns up the quiet violin so they are on equal footing. Only after they are balanced does it mix them according to your slider settings. This ensures that no single goal hijacks the learning process.

4. Real-World Examples

The paper tested this on three different types of creative tasks:

Text-to-Image (The Art Gallery): You type "A cute dog." You can slide a control from Photorealistic (looks like a real photo) to Sketch (looks like a pencil drawing). As you slide, the image smoothly morphs from a photo to a drawing, hitting every style in between perfectly.
Image Editing (The Photo Editor): You upload a photo of yourself and say, "Turn me into a warrior."
- Slide to Preservation: You look exactly like yourself, just with a tiny sword.
- Slide to Editing: You become a full fantasy warrior, but you might lose some of your original facial features.
- The Slider: You can find the perfect middle ground where you look like a warrior but still look like you.
Text-to-Video (The Movie Maker): You can slide between Realistic (looks like a camera recording) and Animation (looks like a Disney movie). The video smoothly transitions between these two worlds.

Why This Matters

Before this, if you wanted to explore different creative balances, you had to train a new AI model for every single preference. That is slow, expensive, and requires storing hundreds of different models.

ParetoSlider is like having one universal remote control for creativity. It allows a single, smart AI to understand that "perfect" isn't one fixed point, but a whole range of possibilities. It gives non-expert users the power to intuitively navigate complex trade-offs, making AI art tools much more flexible and user-friendly.

In short: It turns a rigid, one-size-fits-all AI into a flexible, customizable tool that listens to your preferences in real-time.

1. Problem Statement

Current Reinforcement Learning (RL) post-training methods for generative models (diffusion and flow-matching) typically rely on single-objective optimization or early scalarization.

The Limitation: When multiple, often conflicting, objectives exist (e.g., adhering to a text prompt vs. preserving the source image in editing, or photorealism vs. sketch style), standard methods collapse these into a single fixed weighted sum of rewards during training.
The Consequence: This forces the model to converge to a single, static operating point (a specific trade-off). Users cannot adjust the balance between objectives at inference time without retraining the model or maintaining multiple checkpoints (one for each weight combination).
The Gap: Existing Multi-Objective RL (MORL) approaches for visual generation either lack inference-time control, require storing multiple models (scaling linearly with objectives), or suffer from high computational overhead during sampling.

2. Methodology: ParetoSlider

The authors propose ParetoSlider, a Multi-Objective Reinforcement Learning (MORL) framework that trains a single diffusion model to approximate the entire Pareto front (the set of optimal trade-offs where no objective can be improved without degrading another).

Core Components

Preference-Conditioned Policy:
- Instead of a fixed policy, the model is conditioned on a continuous preference vector $\omega = [\omega_1, \dots, \omega_M]$ , where $\omega$ lies on the probability simplex ( $\sum \omega_m = 1, \omega_m \ge 0$ ).
- During training, the model samples different $\omega$ vectors, learning to map any specific preference to its corresponding optimal point on the Pareto frontier.
- At inference, users can slide $\omega$ continuously to control the trade-off without retraining.
Late-Scalarization Strategy:
- Problem: Directly aggregating raw rewards (early scalarization) before normalization causes "reward hijacking," where objectives with larger magnitudes or variances dominate the gradient, ignoring the user's preference vector $\omega$ .
- Solution: ParetoSlider employs late scalarization.
  - Step 1: Compute advantages ( $A_m$ ) for each reward objective $m$ independently within a group of samples.
  - Step 2: Normalize these advantages independently to remove scale imbalances.
  - Step 3: Compute a separate DiffusionNFT loss ( $L_m$ ) for each objective.
  - Step 4: Aggregate the losses only at the final step using the preference weights: $L_{total} = \sum \omega_m L_m$ .
- This ensures the optimization landscape is strictly defined by the target preference $\omega$ , not the arbitrary scales of the reward functions.
Architecture & Conditioning:
- Built upon DiffusionNFT, an efficient online RL framework for flow-matching models that optimizes on the forward process rather than the reverse trajectory.
- Conditioning Mechanisms: The preference vector $\omega$ $ω$ is injected into the transformer backbone via lightweight modules (LoRA + specific conditioning):
  - Text-to-Image (SD3.5): Injects $\omega$ via timestep embeddings and shared residual modulation on the image stream.
  - Image-to-Image (FluxKontext): Modulates the AdaLN scale and shift of the context (text) stream.
  - Text-to-Video (LTX-2): Uses a shared block-residual mechanism targeting the video stream.

3. Key Contributions

Single-Model Pareto Approximation: Introduces a framework where a single set of parameters captures the entire Pareto front, eliminating the need for multiple checkpoints or model interpolation.
Inference-Time Control: Enables continuous, fine-grained control over competing rewards (e.g., sliding from "Photorealistic" to "Sketch") via a simple preference vector slider.
Novel Training Paradigm: Proposes Late-Scalarization with per-reward advantage normalization, which prevents reward hijacking and ensures faithful adherence to user-defined trade-offs.
Scalability: Demonstrates applicability across diverse tasks (Text-to-Image, Image Editing, Text-to-Video) and state-of-the-art backbones (SD3.5, FluxKontext, LTX-2).

4. Experimental Results

The authors evaluated ParetoSlider against baselines including Fixed-Weights (separate models for fixed weights), Flow-Multi (Pareto-based sample selection), Prompt Rewriting, and standard guidance scales (CFG).

Qualitative Performance:
- Text-to-Image: Smooth transitions between photorealism and styles (sketch, anime, watercolor). The model maintains coherence across the entire spectrum.
- Image Editing: Successfully balances instruction adherence (e.g., "turn into a warrior") with source preservation (keeping facial features). The "balanced" midpoint faithfully interpolates between extremes.
- Text-to-Video: Controls the trade-off between photorealism and 3D animation styles in generated videos.
Quantitative Performance:
- Pareto Frontier Coverage: ParetoSlider consistently traces a well-ordered, smooth Pareto frontier that dominates all baselines.
- Hypervolume (HV): Achieves significantly higher Hypervolume indicators compared to Fixed-Weights and other baselines, indicating better convergence to the optimal front and better spread across objectives.
- Efficiency: A single trained model replaces the need for $N$ separate models (where $N$ is the number of desired trade-off points), saving storage and training compute.

5. Significance

Paradigm Shift: Moves the field from "static alignment" (training a model for one specific goal) to "dynamic alignment" (training a model to understand and execute a spectrum of goals).
User Experience: Provides a practical interface for non-expert users to navigate complex, conflicting objectives (e.g., "I want the edit to be strong, but keep the face recognizable") without technical knowledge of reward weights or model architecture.
Scalability: Offers a scalable solution for multi-objective alignment in visual generation, addressing the "trilemma" of flexibility, efficiency, and scalability that plagued previous MORL attempts.
Robustness: The late-scalarization approach ensures that the model learns the relative importance of objectives as defined by the user, rather than being biased by the inherent numerical scales of the reward models.

In conclusion, ParetoSlider establishes a new standard for aligning generative models with human preferences, enabling flexible, continuous control over the trade-offs that define high-quality visual generation.