Multi-Axis Max@K Reinforcement Learning for Representative Diversity in Text-to-Image Generation
Ku Onoda, Paavo Parmas, Hiroki Furuta, Soichiro Nishimori, Yuta Oshima, Shohei Taniguchi, Yutaka Matsuo
原文摘要
Text-to-image (T2I) models produce realistic, prompt-aligned images, but samples for the same prompt often cover only a few visually distinct modes. For person-centric prompts, this lack of diversity can reflect or amplify demographic skew. We formalize representative diversity as target-mode coverage: how well a set of samples covers predefined modes, such as colors or perceived-race categories. We then propose MultiMaxPO, a reinforcement learning post-training method that improves this coverage in diffusion T2I models by optimizing a multi-axis max@K objective. Given one score per mode for each sample, multi-axis max@K takes the maximum over samples for each mode and then sums these maxima. A sample therefore receives credit for a mode only when it raises that mode's maximum, so different samples are rewarded for different modes. We verify this credit assignment on a synthetic Gaussian mixture and on Stable Diffusion 3.5 Medium with pixel-based color rewards, and then apply it to perceived-appearance fairness. On held-out prompts and under three automatic evaluators, MultiMaxPO improves the Fairness Score over the base model by 0.22-0.36 on Stable Diffusion 3.5 Medium and by 0.07-0.19 on FLUX.1-dev, outperforms matched RL baselines, and maintains image quality and text alignment. Code is available at https://github.com/KuOnoda/multi-axis-maxk.
cs.LG · cs.AI · cs.CV