Multi-Axis Max@K Reinforcement Learning for Representative Diversity in Text-to-Image Generation
Introduces multi-axis Max@K reinforcement learning to enhance target mode coverage in text-to-image models, significantly improving diversity and fairness.
Key Findings
Methodology
This paper proposes multi-axis Max@K as a group-based RL objective to improve mode coverage in diffusion models. It computes the maximum score per category across samples, then sums these maxima, allowing different samples to contribute to different categories. Validation includes synthetic mixture experiments and pixel-based rewards on SD3.5-M, demonstrating effective mode coverage. The approach is extended to perceived fairness, with training involving multiple samples per prompt, using advantage estimates derived from per-axis scores. The method optimizes the generator to produce diverse, fair outputs across multiple target modes.
Key Results
- In synthetic Gaussian experiments, multi-axis Max@K redistributes probability mass toward all target modes, with higher k values yielding more balanced distributions. On pixel-based color tasks, k=7 increased mode coverage by 57% over baseline. In fairness evaluations, the method improved fairness scores by 0.23-0.36 while maintaining image quality and text alignment. These results confirm the method’s ability to enhance target coverage and fairness across diverse scenarios.
- Compared to scalar rewards, multi-axis Max@K effectively encourages coverage of multiple mutually exclusive modes, reducing bias and increasing diversity. The experiments demonstrate that increasing k improves finite-batch coverage, especially for rare modes, and balances target distributions. The approach generalizes well across synthetic, pixel, and perceived fairness tasks, showing robustness and scalability.
- Ablation studies reveal that independent per-axis maximization is crucial for achieving balanced mode coverage. The method also outperforms traditional RL baselines and inference-time fairness interventions, indicating its effectiveness in multi-objective optimization. Overall, the results highlight the potential of multi-axis Max@K for fair, diverse image generation.
Significance
This work addresses fundamental limitations in current generative models, which often favor certain modes and exhibit demographic biases. By explicitly optimizing for target mode coverage, it advances the state-of-the-art in producing diverse, fair outputs. The approach bridges the gap between diversity and fairness, offering a scalable, theoretically grounded framework applicable to various multi-objective tasks. Its success on perception-based fairness metrics demonstrates practical relevance, especially for sensitive applications like human portrait generation. The methodology paves the way for more equitable AI systems that can serve broader societal needs while maintaining high-quality outputs.
Technical Contribution
The core technical innovation is the formulation of multi-axis Max@K as a set-based RL objective that preserves category-specific credit assignment. This contrasts with traditional scalar rewards, enabling the model to learn balanced representations across multiple mutually exclusive modes. The method integrates advantage estimation techniques like EI+L2O for efficient gradient computation, ensuring stable training. It extends prior set-based RL algorithms (e.g., ReMax) to high-dimensional, multi-objective image generation, providing theoretical guarantees for target mode coverage. The approach also introduces a flexible hyperparameter (k) to control diversity levels, and demonstrates its effectiveness across synthetic, pixel, and perception-based tasks.
Novelty
This is the first application of multi-axis Max@K in diffusion-based image generation, explicitly targeting mode coverage over predefined categories. Unlike existing methods that optimize a single scalar reward or rely on post-hoc fairness adjustments, this approach directly encourages balanced representation across multiple mutually exclusive modes. Its innovative use of maximum scores per category before aggregation ensures different samples contribute to different modes, addressing a key challenge in multi-objective generative modeling. This paradigm shift offers a new way to embed fairness and diversity into the training process itself, rather than relying solely on inference-time interventions.
Limitations
- The method depends on predefined categories, limiting flexibility in dynamic or unknown mode spaces. Extending to adaptive or learned categories remains a challenge.
- In high-dimensional mode spaces, the maximum operation may overly favor certain modes, potentially skewing fairness metrics.
- Training complexity and computational cost increase with the number of categories and sample group size, requiring further optimization for large-scale deployment.
Future Work
Future directions include developing adaptive category discovery mechanisms, reducing computational overhead, and extending the framework to continuous or hierarchical mode spaces. Incorporating user feedback or self-supervised signals could further enhance fairness and diversity. Additionally, applying the approach to other modalities such as video or 3D generation, and integrating it with larger foundation models, will broaden its impact. Exploring theoretical bounds and convergence guarantees for multi-axis Max@K in complex settings also presents promising research avenues.
AI Executive Summary
The rapid advancement of text-to-image (T2I) models has revolutionized visual content creation, enabling the synthesis of highly realistic images from textual prompts. Despite these breakthroughs, a persistent challenge remains: the generated samples tend to lack diversity, often covering only a narrow subset of possible visual modes. This limitation is especially problematic in applications involving human portraits or demographic-sensitive content, where biased or skewed representations can reinforce societal stereotypes. Traditional reinforcement learning (RL) approaches, which optimize scalar rewards like quality or preference, struggle to promote coverage across multiple mutually exclusive categories simultaneously.
To address this, the authors introduce multi-axis Max@K, a novel RL objective designed to enhance target-mode coverage by leveraging a group-based maximum scoring mechanism. Unlike conventional scalar rewards, this approach evaluates each target category independently, selecting the strongest sample per category before aggregating. This preserves category-specific credit, allowing different samples within a batch to represent different modes effectively. The method is validated through synthetic experiments, pixel-based color rewards, and perception-based fairness assessments, demonstrating significant improvements in coverage and fairness metrics.
Experimental results show that, on synthetic Gaussian mixtures, increasing the hyperparameter k leads to more balanced mode distributions. In pixel rewards, the model with k=7 achieves a 57% increase in mode coverage. In fairness tasks, the method boosts fairness scores by 0.23-0.36, without sacrificing image quality or text alignment. These findings highlight the potential of multi-axis Max@K to produce more equitable and diverse outputs, addressing societal biases and expanding the applicability of generative models.
Overall, this work offers a robust, theoretically grounded framework for multi-objective optimization in diffusion models. It opens new avenues for fair, diverse AI content generation, with broad implications for industry and research. Future work will focus on adaptive category definitions, reducing computational costs, and extending to more complex, real-world scenarios, ultimately aiming to create AI systems that are both inclusive and richly expressive.
Deep Dive
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.