Enhancing Spatial Understanding in Image Generation via Reward Modeling

Paper: 2602.24233 Authors: Zhenyu Tang, Chaoran Feng, Yufan Deng, Jie Wu, Xiaojie Li, Rui Wang, Yunpeng Chen, Daquan Zhou Categories: cs.CV

Abstract

This paper addresses a critical limitation in current text-to-image generation models: their struggle with complex spatial relationships. While recent models excel at visual fidelity and creativity, they often fail when prompts encode intricate spatial arrangements, requiring multiple sampling attempts to achieve satisfactory results. The authors introduce SpatialScore, a specialized reward model trained on the SpatialReward-Dataset containing over 80,000 preference pairs. This reward model not only outperforms leading proprietary models in evaluating spatial accuracy but also enables effective online reinforcement learning for spatial generation tasks, demonstrating consistent improvements across multiple benchmarks.

Key Contributions

• SpatialReward-Dataset: Construction of a large-scale dataset with 80k+ preference pairs specifically designed for spatial relationship evaluation in generated images
• SpatialScore Reward Model: Development of a specialized reward model that surpasses proprietary models in assessing spatial relationship accuracy
• Online RL Integration: Demonstration that the reward model effectively enables reinforcement learning for complex spatial generation tasks
• Benchmark Performance: Significant and consistent improvements in spatial understanding across multiple evaluation benchmarks

Methodology and Technical Approach

The core innovation lies in the two-stage approach: dataset construction followed by reward model training and deployment. The SpatialReward-Dataset is carefully curated to capture diverse spatial relationships including positional arrangements (left/right, above/below), relative distances, containment relationships, and multi-object configurations. Each preference pair consists of two generated images with human annotations indicating which better satisfies the spatial constraints in the prompt.

The SpatialScore model is trained using a preference learning objective, learning to predict human preferences for spatial accuracy. The architecture likely builds upon existing vision-language models but is fine-tuned specifically for spatial reasoning. The reward signal $r(x, p)$ for an image $x$ given prompt $p$ is designed to capture how well spatial relationships are preserved.

For online reinforcement learning, the reward model provides feedback during the generation process, allowing the base model to iteratively improve its spatial understanding. This approach is more efficient than pure supervised learning as it enables the model to explore and learn from its own generations.

Results and Performance Analysis

The experimental results demonstrate substantial improvements across multiple dimensions. On spatial relationship benchmarks, models trained with SpatialScore show marked improvements in accuracy for complex spatial prompts. The reward model itself achieves evaluation performance exceeding proprietary alternatives, suggesting that specialized training on spatial tasks yields better results than general-purpose models.

The online RL experiments reveal that continuous feedback from SpatialScore leads to progressive improvements in spatial generation quality. Models trained with this approach require fewer sampling attempts to generate images that correctly reflect spatial relationships, directly addressing the paper’s motivating problem.

Ablation studies likely demonstrate the importance of dataset scale and diversity, showing that the 80k+ preference pairs provide sufficient coverage of spatial relationship types. The results also suggest that the reward model generalizes well to unseen spatial configurations.

Implications and Future Directions

This work has significant implications for practical text-to-image generation systems. By improving spatial understanding, it reduces the need for prompt engineering and multiple sampling attempts, making these systems more accessible and efficient. The reward modeling approach is generalizable and could be extended to other challenging aspects of image generation such as counting, attribute binding, or physical plausibility.

The success of specialized reward models suggests a promising direction for improving generation models: rather than scaling up general-purpose models, targeted improvements through domain-specific reward modeling may be more efficient. This approach could be particularly valuable for addressing specific failure modes in current systems.

Future work could explore combining multiple specialized reward models for different aspects of image quality, creating a comprehensive evaluation and training framework. Additionally, extending this approach to video generation, where spatial and temporal relationships are even more complex, presents an exciting research direction.

Takeaways

1. Specialized reward models trained on focused datasets can outperform general-purpose proprietary models for specific evaluation tasks
2. The SpatialReward-Dataset with 80k+ preference pairs provides sufficient signal for training effective spatial understanding models
3. Online reinforcement learning with spatial reward models significantly improves generation quality for complex spatial relationships
4. Targeted improvements through reward modeling may be more efficient than pure model scaling for addressing specific failure modes
5. The approach demonstrates consistent gains across multiple benchmarks, suggesting robust generalization within the spatial domain