CountSteer: Steering Attention for Object Counting in Diffusion Models

TL;DR

This paper identifies latent numerical cues within diffusion models that correlate with counting accuracy and introduces CountSteer, a training-free inference-time method that steers cross-attention hidden states to improve object-count fidelity. By deriving an adaptive steering vector from the mean hidden-state differences between correct and incorrect counts and applying it during early denoising steps, the approach achieves about a 4% gain in counting accuracy while preserving image quality and semantic alignment. The method requires no fine-tuning and demonstrates potential for extending numerical and compositional control in text-to-image generation. Overall, CountSteer highlights latent quantitative reasoning in diffusion models and offers a lightweight pathway to more reliable, controllable synthesis.

Abstract

Text-to-image diffusion models generate realistic and coherent images but often fail to follow numerical instructions in text, revealing a gap between language and visual representation. Interestingly, we found that these models are not entirely blind to numbers-they are implicitly aware of their own counting accuracy, as their internal signals shift in consistent ways depending on whether the output meets the specified count. This observation suggests that the model already encodes a latent notion of numerical correctness, which can be harnessed to guide generation more precisely. Building on this intuition, we introduce CountSteer, a training-free method that improves generation of specified object counts by steering the model's cross-attention hidden states during inference. In our experiments, CountSteer improved object-count accuracy by about 4% without compromising visual quality, demonstrating a simple yet effective step toward more controllable and semantically reliable text-to-image generation.

Figures (8)

• Figure 1: CountSteer improves numerical fidelity in text-to-image diffusion. While baseline Stable Diffusion v1.5 often fails to generate the correct number of objects, our cross-attention steering framework generates images with more accurate object counts that better align with user-specified prompts across diverse object categories.
• Figure 2: CountSteer: During each denoising step $t$, adaptive steering vectors are injected into UNet blocks to steer hidden states based on latent distributional differences between correct ($\mu^1_{t,b}$) and incorrect ($\mu^0_{t,b}$) samples.
• Figure 3: Visualization of hidden state distributions analyzed via KDE. The two classes (correct and incorrect counts) exhibit clearly separable regions, supporting the formulation of our steering direction vector.
• Figure 4: We generate images from "{count} {object}" prompts, extract hidden states $h_{t,b}$, and manually annotate them into balanced classes of 200 correct and incorrect counting to construct the base steering vector.
• Figure 5: Overall distribution of correct and incorrect counts across different targets (one to ten).