FR-TTS: Test-Time Scaling for NTP-based Image Generation with Effective Filling-based Reward Signal

TL;DR

FR-TTS introduces a novel filling-based reward to improve test-time scaling in NTP-based image generation, enabling more accurate intermediate sample evaluation and outperforming existing methods.

Contribution

The paper proposes FR, a filling-based reward metric, and FR-TTS, a scaling strategy that enhances image generation quality by better evaluating intermediate samples.

Findings

FR provides a reliable metric for intermediate sample quality.

FR-TTS outperforms multiple benchmarks and reward models.

The approach effectively balances diversity and quality in generated images.

Abstract

Test-time scaling (TTS) has become a prevalent technique in image generation, significantly boosting output quality by expanding the number of parallel samples and filtering them using pre-trained reward models. However, applying this powerful methodology to the next-token prediction (NTP) paradigm remains challenging. The primary obstacle is the low correlation between the reward of an image decoded from an intermediate token sequence and the reward of the fully generated image. Consequently, these incomplete intermediate representations prove to be poor indicators for guiding the pruning direction, a limitation that stems from their inherent incompleteness in scale or semantic content. To effectively address this critical issue, we introduce the Filling-Based Reward (FR). This novel design estimates the approximate future trajectory of an intermediate sample by finding and applying a reasonable filling scheme to complete the sequence. Both the correlation coefficient between rewards of intermediate samples and final samples, as well as multiple intrinsic signals like token confidence, indicate that the FR provides an excellent and reliable metric for accurately evaluating the quality of intermediate samples. Building upon this foundation, we propose FR-TTS, a sophisticated scaling strategy. FR-TTS efficiently searches for good filling schemes and incorporates a diversity reward with a dynamic weighting schedule to achieve a balanced and comprehensive evaluation of intermediate samples. We experimentally validate the superiority of FR-TTS over multiple established benchmarks and various reward models. Code is available at \href{https://github.com/xuhang07/FR-TTS}{https://github.com/xuhang07/FR-TTS}.