CountSteer: Steering Attention for Object Counting in Diffusion Models

AI summary

Overview

Research area: Computer vision and generative modeling — specifically inference-time control of text-to-image (T2I) diffusion models (Stable Diffusion v1.5) for object-count fidelity.

Technical level: Intermediate. Readers should have basic familiarity with diffusion denoising, cross-attention layers in UNet blocks, and concepts such as cosine similarity and latent representations.

Scope: The paper proposes CountSteer, a training-free inference-time method that steers cross-attention hidden states in a diffusion model so generated images better match the object count specified in the text prompt.

What This Paper Is About

Text-to-image diffusion models such as Stable Diffusion produce realistic images but frequently ignore the number of objects requested in the prompt, generating too few or too many. The authors observe that the model's internal cross-attention hidden states actually do separate correct from incorrect counting outcomes, suggesting the model already encodes a latent notion of numerical correctness that is simply not manifested in the final image. CountSteer exploits this signal by adapting the model's internal representations during inference, aiming to improve object-count accuracy without any retraining or architectural change.

Key Contributions

1. Discovery of latent numerical awareness. Using Kernel Density Estimation (KDE) on query hidden states from cross-attention layers, the authors show that samples with correct counts (Class 1) and incorrect counts (Class 0) form separable distributions, with a directional pattern that implicitly encodes counting behavior.

2. A training-free steering formulation. They define a steering vector as the difference between the mean hidden states of Class 1 and Class 0 samples, computed per denoising step t and per UNet block b (s_{t,b} = μ¹_{t,b} − μ⁰_{t,b}), and inject it into the hidden states during inference.

3. An adaptive scaling mechanism. Because prompt-dependent variability makes a fixed steering vector either too weak or too strong, CountSteer multiplies the base vector by an adaptive factor combining a distance ratio and a cosine-similarity term, with a global scaling constant empirically set to c = 100.

4. An empirical evaluation on Stable Diffusion v1.5. Using LLaVA-OneVision to automatically count objects in generated images, the method reports accuracy, mean absolute error, and CLIP-Score improvements, plus a catalog of failure modes.

Main Findings

• Quantitative gains over the baseline. On Stable Diffusion v1.5, ACC rose from 50.0% to 54.0% with CountSteer (a 4.0 percentage-point increase, described in the abstract and conclusion as roughly a 4% improvement), and MAE dropped from 1.125 to 0.940, a reduction of 0.185.

• Semantic alignment is essentially preserved. CLIP-Score was 30.99 for the baseline and 30.39 with CountSteer. The authors describe this as "comparable" and take it as evidence that counting improvements do not come at the cost of semantic integrity or image quality.

• Bidirectional control. The qualitative results indicate the baseline often under- or over-generates, whereas CountSteer adaptively increases or suppresses object production to correct the deviation, while preserving fine-grained detail and overall visual style.

• Counting collapses at higher quantities. The appendix reports that accuracy drops sharply as the target count increases, with generations almost entirely failing beyond a count of four. The paper's dataset is therefore restricted to the one-to-four range.

• Count is largely determined early in denoising. Visualizations of the denoising process show that approximate spatial arrangement and object count begin to emerge within the first few denoising steps, motivating the application of steering only during the initial phase.

• Three recurring failure modes. (1) Over-generation, when class-wise latent directions excessively amplify object features and cause duplication, especially for objects without a compact latent representation; (2) unreliable rendering, where steering has no effect because the model itself cannot generate the target object reliably; (3) degradation of initially correct output, where a previously accurate generation becomes incorrect after steering, indicating a risk of over-steering.

Methodology in Plain English

The researchers first asked whether a diffusion model internally "knows" when it has produced the right number of objects. They generated images from prompts built from a simple template, "{count} {object}" (for example, "three cats" or "one carrot"), and had a human annotate each image as correct (Class 1) or incorrect (Class 0). From those images they extracted the query vectors of cross-attention layers during the first k denoising timesteps across all UNet blocks — these are the "hidden states" the method operates on — and checked with KDE whether the two classes looked different. They did.

Next, they built a steering vector: for each denoising step and UNet block, they averaged the hidden states of the correct class and subtracted the average of the incorrect class. Adding this vector to a hidden state during generation should push it toward the "correct" region of the representation space. Because a single fixed vector cannot suit every prompt, they added an adaptive scale. That scale has two parts: a distance ratio measuring how far the current hidden state is from the correct-class mean (so steering is stronger when far away and gentler when close, to avoid overshooting), and a cosine similarity that checks whether the base steering vector points in a sensible direction — a negative value flips its direction. A constant multiplier, empirically set to 100, amplifies the resulting factor, which would otherwise be too small to matter. The final hidden state is the original plus the scaled steering vector.

Implementation details: Stable Diffusion v1.5 with 50 denoising steps and a guidance scale of 7.5; steering applied only during the first 10 denoising steps; all experiments run with randomly initialized seeds. The prompt set consists of 600 prompts generated with GPT-4o, with equal numbers for each count from one to four, split into 400 prompts for constructing the steering vector and 200 for evaluation with no overlap; the steering-vector construction images were manually annotated into balanced classes of 200 correct and 200 incorrect images, regenerating with different seeds where needed to balance classes. Evaluation uses the LLaVA-OneVision model to count objects, with the fixed instruction "How many {objects} are in the image? Reply with only a number.", and the reported test set comprises 100 test images generated with the same prompt template used for training. Metrics are Accuracy (ACC), Mean Absolute Error (MAE), and CLIP-Score.

Why This Matters

Impact on research. The work supports the hypothesis that diffusion models already contain correct semantic and quantitative representations internally but fail to express them in output. It demonstrates that a small, geometry-based intervention at inference time can unlock latent capability, paralleling findings in language models (the paper cites Inference-Time Intervention), and it offers a training-free alternative to fine-tuning or architectural modification.

Real-world applications (drawn from the domains the paper names):

• Synthetic data generation for downstream AI training, where label correctness depends on the requested object count.
• Creative design, where a designer specifies an exact number of elements and needs the image to match.
• Simulation, where scene contents must match specified quantities.
• Education, where illustrations of a given count (for example, "four soda cans") must be numerically correct.
• The paper also notes the framework could extend to compositional attributes such as color, spatial arrangement, and multi-object interactions.

Industry relevance. Any pipeline that generates images from text specifications and depends on quantitative accuracy — dataset construction, advertising and product imagery, content creation tooling — could apply CountSteer as an add-on to an existing diffusion model, since it requires no retraining, no new parameters, and no change to model structure.

Future Directions

• Extending beyond counting. The authors state that the same steering framework can be applied to compositional attributes such as color, spatial arrangement, and multi-object interactions.

• Scaling to higher counts. Because accuracy drops sharply and generations almost entirely fail beyond four objects, and because the dataset was restricted to one-to-four to stay within the model's "basic yet interpretable" numerical understanding, extending the approach to larger quantities remains open.

• Adaptive or selective steering. The failure analysis calls for mechanisms that dynamically respond to the model's intermediate state, rather than applying steering uniformly, to avoid over-generation and the regression of already-correct outputs.

• Handling objects the model cannot render. Cases where steering has no effect because the model generates the target object ambiguously or inconsistently reveal a dependence on the backbone's baseline generative capacity, raising the question of how to detect and respond to such prompts.

Target Audience

Researchers and practitioners working on controllable text-to-image generation, diffusion model interpretability, and inference-time intervention methods. It is also relevant to engineers building synthetic-data pipelines that require numerically faithful image generation, and to readers interested in whether generative models possess latent reasoning capabilities that can be elicited without retraining. Readers without prior exposure to diffusion sampling and cross-attention mechanics will need background reading, since the method is defined directly in terms of per-step, per-block hidden states.

Authors’ 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.