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

2026-07-16

AI Summary

This work addresses the limited diversity of existing text-to-image generation models under identical prompts, particularly the demographic biases in human image synthesis that hinder coverage of prescribed semantic modes. To tackle this, the authors propose a multi-axis Max@K reinforcement learning objective that decouples diversity and fidelity optimization during diffusion model fine-tuning. By integrating population-based rewards with a multi-axis credit assignment strategy, the method enables individual samples to independently contribute maximal values across distinct semantic dimensions. The approach combines pixel-level color rewards with an automated appearance fairness evaluator, achieving significant improvements over baseline methods—gaining 0.23–0.36 on three fairness metrics—while preserving image quality and text alignment, thereby substantially enhancing the representational diversity of generated outputs.

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.

Problem

• text-to-image generation
• diversity
• target-mode coverage
• demographic skew
• fairness

Innovation

• multi-axis max@K
• reinforcement learning
• text-to-image generation
• representative diversity
• fairness in AI