🤖 machine learning

Multi-Axis Max@K Reinforcement Learning for Representative Diversity in Text-to-Image Generation

This paper introduces multi-axis max@K, a group-based reinforcement learning objective that enhances the diversity and demographic fairness of text-to-image generation by assigning credit to samples only when they uniquely improve coverage of specific semantic modes within a batch, thereby achieving higher fairness scores without compromising image quality or prompt alignment.

Original authors: Ku Onoda, Paavo Parmas, Hiroki Furuta, Soichiro Nishimori, Yuta Oshima, Shohei Taniguchi, Yutaka Matsuo

Published 2026-07-17

📖 9 min read🧠 Deep dive

CC BY 4.0

Original authors: Ku Onoda, Paavo Parmas, Hiroki Furuta, Soichiro Nishimori, Yuta Oshima, Shohei Taniguchi, Yutaka Matsuo

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 walking into a giant, magical art studio where a robot artist can instantly paint any picture you can describe. You say, "Draw a detective," and poof! A detective appears. You say it again, and another detective appears. But here's the catch: every single time you ask, the robot paints the exact same kind of detective—maybe always a tall man in a trench coat. It's like the robot has a favorite outfit and refuses to try anything else. This is a common problem with modern "text-to-image" AI. While these robots are incredibly good at making pictures that look real and match your words, they often get stuck in a rut, producing the same few variations over and over. This is a big deal because if you ask for a "doctor" or a "leader," and the robot only ever shows you one specific type of person, it reinforces a narrow and unfair view of the world. Scientists call this a lack of "diversity."

To fix this, researchers usually try to tell the robot, "Make it look different!" But often, when they force the robot to be different, the pictures start looking weird, blurry, or just plain wrong. The robot gets confused and stops listening to your instructions. This paper introduces a clever new way to teach the robot to be diverse without losing its cool. The authors propose a method called "Multi-Axis Max@K." Think of it like a talent show where you don't just pick the single best singer; instead, you look at a whole group of contestants and say, "Okay, this one is the best singer, this one is the best dancer, and this one is the best magician." The goal isn't for one person to do everything perfectly; it's for the whole group to cover all the bases. The researchers tested this idea and found that it helps the AI generate a much wider variety of people and colors, making the results fairer and more interesting, all while keeping the pictures sharp and true to your description.

The Problem: The Robot's "Favorite" Outfit

Imagine you have a friend who loves to draw. You ask them to draw a "cat." They draw a fluffy orange tabby. You ask again, and they draw the exact same orange tabby. You ask a third time, and it's still the orange tabby. Your friend isn't bad at drawing; they are just stuck in a loop. They have a "favorite mode" of drawing cats, and they keep hitting it.

In the world of Artificial Intelligence (AI), specifically "Text-to-Image" models, this happens all the time. These models are trained on billions of images from the internet. If the internet has more pictures of white men as doctors than women or people of other races, the AI learns that "doctor" equals "white man." When you ask for a doctor, it pulls from that narrow memory bank.

The usual fix is to tell the AI, "Be diverse!" But if you just shout "be diverse," the AI might get confused. It might start drawing a doctor who looks like a potato or a cat wearing a stethoscope. The AI tries to be different, but it forgets how to be good at following your instructions. The researchers wanted a way to make the AI show a whole group of different options—some doctors who are women, some who are men, some who are Black, some who are Asian—without making the pictures look bad.

The Solution: The "Best of the Batch" Talent Show

The authors of this paper came up with a new rule for how the AI learns. They call it Multi-Axis Max@K. That sounds complicated, but let's break it down with a simple story.

Imagine you are a teacher grading a class of students. You have a list of "skills" you want to see: Math, Art, Music, and Sports.

The Old Way (Scalar Reward): You look at each student individually. "Student A is great at Math but bad at everything else. Student B is great at Art but bad at everything else." If you just add up their scores, you might end up picking the one student who is "okay" at everything but amazing at nothing. Or, you might pick the Math genius and ignore the fact that the class has no musicians.
The New Way (Multi-Axis Max@K): You look at the whole class as a group. You ask: "Who is the best Math student in this group?" You pick Student A. "Who is the best Music student?" You pick Student B. "Who is the best Sports student?" You pick Student C.

The magic happens here: The group gets a high score because it has a Math genius, a Music genius, and a Sports genius, even if no single student is good at all three. The AI learns that it doesn't need to make one perfect picture that does everything. Instead, it learns that a "batch" of pictures is successful if different pictures in that batch represent different things.

The "Max" part means the AI looks for the best example of each category in the group. The "K" part means it looks at a group of K pictures (like a batch of 7 or 10). The "Multi-Axis" part means it checks for many different categories at once (like different skin tones, genders, or colors).

How They Tested It: From Pixels to People

The researchers didn't just guess this would work; they tested it in three different ways, getting more complex each time.

1. The Color Game (Synthetic Test)
First, they played a game with a simple computer program. They told the program to generate images that were either Red, Green, Blue, or other colors. They set up a rule where the program got points if the group of images it made covered all the colors.

What happened: When they used the old "single score" method, the program got lazy and just made a bunch of Red images because that was the easiest way to get points. But when they used the new "Multi-Axis Max@K" method, the program realized it needed to make a Red image, a Green image, and a Blue image to win. It started spreading its bets and making a colorful mix.

2. The Pixel Puzzle (Rule-Based Test)
Next, they used a real image generator called SD3.5-M (a popular AI art model). This time, they didn't use a human to judge the pictures. Instead, they used a computer program that looked at the pixels and counted colors. They asked the AI to make pictures with specific color themes (like "warm colors" or "dark colors").

The Result: The AI learned to produce a batch where one picture was bright and warm, another was cool and blue, and another was dark. It successfully covered all the "axes" of color without needing a human to tell it what was good.

3. The Fairness Test (Perceived Appearance)
Finally, they tackled the big, real-world problem: Fairness. They asked the AI to generate pictures of people with different jobs (like "doctor," "artist," "scientist"). They wanted to see if the AI would show a diverse mix of races and genders.

The Setup: They used automatic tools to "guess" the race and gender of the people in the generated pictures. They defined "target modes" like "Black woman," "Asian man," "Latino woman," etc.
The Outcome: When they used the new method, the AI started generating a much more balanced group.
- Before the fix, a batch of 16 pictures of a "doctor" might have had 13 white men and 3 people of other backgrounds.
- After the fix, a batch of 16 pictures might have had 4 white men, 3 Black people, 4 Asian people, 2 Indian people, and 3 Latino people.
- The researchers measured this using a "Fairness Score." Their new method improved the score by 0.23 to 0.36 compared to the base model. That's a huge jump! And the best part? The pictures still looked good. The "text alignment" (how well the picture matched the word "doctor") stayed high, and the image quality didn't drop.

What This Means (and What It Doesn't)

The paper shows that by changing how we reward the AI, we can make it naturally want to be diverse. Instead of forcing it to be different (which makes it weird), we reward the group for having different representatives.

However, the authors are careful to say what this doesn't do.

It doesn't create new people out of thin air. If the AI has never seen a picture of a specific type of person, it can't magically invent them. It can only spread out the people it already knows how to draw.
It relies on the "score" the AI gets. If the tool used to judge the pictures (the "evaluator") is biased, the AI might learn the wrong lesson. The researchers used automatic tools to judge the pictures, and while they checked their work with different tools to make sure it was consistent, they admit that human eyes would be the ultimate test.
It works best when you have a clear list of what "diversity" looks like (like a list of specific races or colors). It's a tool to help the AI cover a specific set of goals, not a magic wand that solves every fairness problem in the world.

The Takeaway

This paper is like giving the robot artist a new rulebook. Instead of saying, "You must make one perfect picture," it says, "Make a whole group of pictures, and I'll give you a gold star if the group has a little bit of everything."

The result? The AI becomes less boring and more fair. It stops showing us the same "default" person over and over and starts showing us the rich, colorful variety of the real world. And it does it without making the pictures look like a glitchy mess. It's a small change in the math, but it could lead to a much bigger change in how we see ourselves in the art we create with machines.

Problem Statement

Text-to-image (T2I) models, particularly diffusion-based generators, often produce samples for a single prompt that cover only a narrow subset of visually distinct modes. While these models synthesize realistic images, they frequently lack representative diversity. This limitation is critical when the goal is to provide a diverse set of alternatives rather than a single plausible image. In person-centric prompts, this concentration can reflect or amplify demographic skews (e.g., regarding race or gender), as the model fails to generate strong representatives for underrepresented categories.

Existing diversity interventions often reduce image quality, and standard Reinforcement Learning (RL) approaches using scalar rewards are ill-suited for this problem. Scalar rewards aggregate scores within a single image before comparing samples, which fails when the objective is to cover several mutually exclusive target categories (e.g., different color modes or demographic labels) across a group of samples. A single image cannot simultaneously represent all mutually exclusive modes; therefore, a generated group should be preferred if different samples within it serve as strong representatives for different modes.

Methodology: Multi-Axis Max@K

The authors formalize the problem as target-mode coverage and propose multi-axis max@K, a group-based reinforcement learning objective designed for diffusion models.

1. Objective Formulation:
Instead of aggregating reward dimensions within an image first, the method aggregates across samples first. For a prompt $c$ , a set of $k$ samples $S_k(c)$ is generated. Given $D$ predefined target modes (axes) and a score $r_d(x, c) \in [0, 1]$ for each mode $d$ returned by a fixed evaluator, the objective score $G_k$ is defined as:
$G_k(S_k, c) = \sum_{d=1}^{D} w_d(c) \max_{x_i \in S_k(c)} r_d(x_i, c)$
Here, the maximum score is selected per axis across the sampled set, and these maxima are then summed. This structure ensures that a sample receives credit for a specific category only if it improves the group's maximum score for that category, allowing different samples to contribute to different categories.

2. Credit Assignment Mechanism:
To optimize this set-level objective, the authors derive a policy gradient estimator.

Leave-One-Out (L2O) Marginal Contribution: For a sample $x_i$ and axis $d$ , the credit is determined by whether removing $x_i$ decreases the group's maximum score for that axis.
Expected Improvement (EI) + L2O: In practice, where the training batch size $m$ may differ from the objective window size $k$ , the authors use an Expected Improvement estimator with a leave-two-out baseline. This estimates the per-sample advantage by averaging improvements over all possible subsets of size $k$ drawn from the $m$ candidates.
Axis-Preserving Aggregation: The advantages are calculated independently for each axis and then summed to form a single policy weight for the sample. This preserves the separation of reward axes until the final aggregation step.

3. Training Loop:
The method employs a standard GRPO-style (Group Relative Policy Optimization) update for diffusion models.

Generate an on-policy group of $m$ candidates for a prompt.
Score all candidates on $D$ reward axes.
Compute axis-wise representative credits (max@K) and convert them to per-sample advantages.
Update the policy using these advantages with KL regularization against the base model.

Key Contributions

Formulation of Representative Diversity: The paper defines representative diversity as target-mode coverage, framing it as a set-level objective rather than an individual image property.
Multi-Axis Max@K Algorithm: They introduce a specific RL objective that performs max-selection per reward axis before aggregation. This is presented as a specialized case of the ReMax objective, tailored for generative modeling where modes are mutually exclusive.
Axis-Preserving Credit Assignment: The method implements a group-based diffusion RL algorithm that assigns credit to samples based on their ability to be the "best representative" for a specific mode, preventing the collapse of probability mass onto a single mode.
Comprehensive Validation: The method is validated in three distinct settings:
- A controlled synthetic mixture of Gaussian modes.
- A deterministic, rule-based color reward setting on Stable Diffusion 3.5 Medium (SD3.5-M).
- A perceived-appearance fairness setting (race and gender) on SD3.5-M.

Experimental Results

1. Synthetic and Color Mode Experiments:

In a synthetic setting with nine mutually exclusive Gaussian modes, multi-axis max@K successfully redistributed probability mass to cover all target modes. In contrast, scalar single-mode rewards collapsed the distribution onto a single preferred mode.
On SD3.5-M using seven rule-based color axes (e.g., red, green, warm, cool), the method significantly improved batch-level coverage. With a training window $k=7$ , the model achieved a 57% increase in batch-max coverage ( $C_{pix}$ ) compared to the base model, whereas a $k=1$ baseline (standard single-image GRPO) showed no improvement.

2. Perceived-Appearance Fairness:
The method was applied to improve fairness in person-centric prompts using perceived-race (5 categories) and perceived-gender (2 categories) as target modes.

Fairness Improvement: Multi-axis max@K improved the Fairness Score (measuring closeness to a uniform distribution) by 0.23–0.36 relative to the base SD3.5-M model across three automatic evaluators (CLIP, VLM, and FairFace).
Comparison to Baselines: The method outperformed matched RL baselines using count-based credit (Fair-GRPO) and inference-time interventions (FairImagen, Weak Guidance). Notably, while other methods improved fairness at the cost of text alignment (CLIP-T scores dropped), multi-axis max@K maintained or slightly improved image quality and text alignment metrics.
Qualitative Shift: Visual inspection of generated batches showed a shift from dominance by a single category (e.g., "White") to a balanced distribution across all five perceived-race categories within a single prompt's batch.

3. Ablation Studies:

Window Size ( $k$ ): The hyperparameter $k$ (the number of samples considered for max-selection) is critical. The Fairness Score improved sharply as $k$ increased from 1 to 5 (matching the number of race categories), saturating thereafter. A $k=1$ setting (no set-level credit) failed to improve diversity.
Credit Mechanism: The method outperformed hard-assignment count baselines, which the authors attribute to the ability of soft scores to capture sub-threshold evidence for rare categories and the explicit credit given to the "best available" representative per axis.

Significance and Claims

The paper claims that multi-axis max@K provides a general mechanism for improving diversity in generative modeling by treating diversity as a coverage problem over a predefined set of modes.

Mechanism Validity: The core contribution is the credit assignment rule that allows different samples in a batch to specialize in different reward axes, effectively solving the "mutually exclusive modes" problem that scalar rewards fail to address.
Practical Impact: The method demonstrates that it is possible to significantly improve demographic fairness in T2I generation without sacrificing image quality or text alignment, a trade-off often observed in previous fairness interventions.
Limitations and Scope: The authors modestly note that the method requires measurable scores for target modes and can only credit modes the generator is already capable of sampling; it does not create absent modes from scratch. Furthermore, the results rely on automatic evaluators, and the choice of target axes (e.g., specific race/gender categories) is a task-level design decision that determines what diversity is achieved. The method is presented as a set-level alignment primitive applicable to various domains where candidate sets should cover specific axes.