Multi-Axis Max@K Reinforcement Learning for Representative Diversity in Text-to-Image Generation
Machine LearningArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
The authors study how text-to-image models often create images that cover only a few types of appearances for the same description, which can limit variety and cause biased representations, especially for people. They define this problem as 'target-mode coverage' and introduce a new method called multi-axis max@K, which helps the model produce diverse images across different categories by giving credit when a sample improves coverage in a particular category. They tested this method on synthetic data and a popular diffusion model, showing it makes the output fairer without losing image quality or relevance to the text. Their approach helps balance diversity and fairness in generated images.
Authors
Ku Onoda, Paavo Parmas, Hiroki Furuta, Soichiro Nishimori, Yuta Oshima, Shohei Taniguchi, Yutaka Matsuo
Abstract
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 the diversity of images, and for person-centric prompts, can reflect or amplify demographic skew. We formalize this problem as coverage of a predefined set of semantically specified modes, which we call target-mode coverage. We then propose multi-axis max@K, a group-based reinforcement learning objective for improving such coverage in diffusion-based T2I models. Given a group of samples and one score per target category, multi-axis max@K first takes the maximum score across samples for each category and then sums these category-wise maxima. The resulting credit assignment gives a sample positive weight on a category only when it increases that category's group-wise maximum, allowing different samples to contribute to different categories. We first validate the credit-assignment mechanism on a synthetic mixture and on SD3.5-M using deterministic pixel-based color rewards. We then evaluate the same objective on perceived-appearance fairness. Across three automatic evaluators on held-out prompts, multi-axis max@K improves the Fairness Score by 0.23-0.36 relative to the base model, while maintaining image quality and text alignment.