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

arXiv: https://arxiv.org/abs/2607.14962 · HTML: https://mlsanity.com/paper/2607.14962/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RL with diffusion and flow models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 61 of its 66 references that are in the map.
- First posted on arXiv 16 Jul 2026.

## Placed in

- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 30 of 61)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Qwen2.5-VL Technical Report](https://mlsanity.com/paper/2502.13923/) (2025)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [ImageReward: Learning and Evaluating Human Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2304.05977/) (2023)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generation](https://mlsanity.com/paper/2305.01569/) (2023)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models](https://mlsanity.com/paper/2305.16381/) (2023)
- [Reliable Fidelity and Diversity Metrics for Generative Models](https://mlsanity.com/paper/2002.09797/) (2020)
- [DanceGRPO: Unleashing GRPO on Visual Generation](https://mlsanity.com/paper/2505.07818/) (2025)
- [Assessing Generative Models via Precision and Recall](https://mlsanity.com/paper/1806.00035/) (2018)
- [Unified Concept Editing in Diffusion Models](https://mlsanity.com/paper/2308.14761/) (2023)
- [Improving Video Generation with Human Feedback](https://mlsanity.com/paper/2501.13918/) (2025)
- [Easily Accessible Text-to-Image Generation Amplifies Demographic Stereotypes at Large Scale](https://mlsanity.com/paper/2211.03759/) (2022)
- [Stable Bias: Analyzing Societal Representations in Diffusion Models](https://mlsanity.com/paper/2303.11408/) (2023)
- [HIVE: Harnessing Human Feedback for Instructional Visual Editing](https://mlsanity.com/paper/2303.09618/) (2023)
- [Holistic Evaluation of Text-To-Image Models](https://mlsanity.com/paper/2311.04287/) (2023)
- [Fair Diffusion: Instructing Text-to-Image Generation Models on Fairness](https://mlsanity.com/paper/2302.10893/) (2023)
- [Debiasing Vision-Language Models via Biased Prompts](https://mlsanity.com/paper/2302.00070/) (2023)
- [Pass@k Training for Adaptively Balancing Exploration and Exploitation of Large Reasoning Models](https://mlsanity.com/paper/2508.10751/) (2025)
- [CADS: Unleashing the Diversity of Diffusion Models through Condition-Annealed Sampling](https://mlsanity.com/paper/2310.17347/) (2023)
- [TempFlow-GRPO: When Timing Matters for GRPO in Flow Models](https://mlsanity.com/paper/2508.04324/) (2025)
- [Finetuning Text-to-Image Diffusion Models for Fairness](https://mlsanity.com/paper/2311.07604/) (2023)
- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems](https://mlsanity.com/paper/2505.15201/) (2025)
- [Self-Discovering Interpretable Diffusion Latent Directions for Responsible Text-to-Image Generation](https://mlsanity.com/paper/2311.17216/) (2023)

## Built on by (most cited first, 2 of 2)

- [AutoRef: Harness Optimization for Agentic Multi-Reference Image Generation](https://mlsanity.com/paper/2609.35530/) (2026)
- [VIF-Bench: Evaluating Visual Instruction Following in Multi-Reference Image Generation](https://mlsanity.com/paper/2609.37709/) (2026)

## Abstract (arXiv)

Text-to-image (T2I) models can synthesize realistic, prompt-aligned images, yet samples generated for the same prompt often cover only a small subset of visually distinct modes. This limits diversity and, for person-centric prompts, can reflect or amplify demographic skew. We formalize this problem as target-mode coverage, the coverage of a predefined set of semantically specified modes, and propose multi-axis max@K, a group-based reinforcement learning objective for improving it in diffusion-based T2I models. Given a group of samples and one score per target mode, multi-axis max@K first takes the maximum score across samples for each mode and then sums these per-mode maxima. The resulting credit assignment gives a sample positive weight on a mode only when it raises that mode's group maximum, so different samples can contribute to different modes. We validate the credit-assignment mechanism on a synthetic mixture and on SD3.5-M with deterministic pixel-based color rewards, and then apply the same objective to perceived-appearance fairness. On held-out prompts, multi-axis max@K improves the Fairness Score by 0.23-0.36 over the base model under three automatic evaluators, while maintaining image quality and text alignment. Code is available at https://github.com/KuOnoda/multi-axis-maxk.