ATHENA: Adaptive Test-Time Steering for Improving Count Fidelity in Diffusion Models

Mohammad Shahab Sepehri; Asal Mehradfar; Berk Tinaz; Salman Avestimehr; Mahdi Soltanolkotabi

cs.CVMar 20, 2026

ATHENA: Adaptive Test-Time Steering for Improving Count Fidelity in Diffusion Models

Authors: Mohammad Shahab Sepehri, Asal Mehradfar, Berk Tinaz, Salman Avestimehr, Mahdi Soltanolkotabi

Organizations: Department of Electrical and Computer Engineering, University of Southern California, Los Angeles, CA, USA

Abstract

Text-to-image diffusion models achieve high visual fidelity but surprisingly exhibit systematic failures in numerical control when prompts specify explicit object counts. To address this limitation, we introduce ATHENA, a model-agnostic, test-time adaptive steering framework that improves object count fidelity without modifying model architectures or requiring retraining. ATHENA leverages intermediate representations during sampling to estimate object counts and applies count-aware noise corrections early in the denoising process, steering the generation trajectory before structural errors become difficult to revise. We present three progressively more advanced variants of ATHENA that trade additional computation for improved numerical accuracy, ranging from static prompt-based steering to dynamically adjusted count-aware control. Experiments on established benchmarks and a new visually and semantically complex dataset show that ATHENA consistently improves count fidelity, particularly at higher target counts, while maintaining favorable accuracy-runtime trade-offs across multiple diffusion backbones.

Figures & tables

Figure 1: Count fidelity failures in T2I generative models. Despite explicit numerical specifications in the prompts, generated images frequently contain incorrect object counts across multiple generative backbones, even for simple objects and scenes, highlighting a persistent limitation in numerical control.

Figure 2: Count estimation across diffusion steps. At early sampling steps, object structure is weakly formed and decoded images are too noisy for reliable count estimation. At later steps, object instances become well defined, and counting is reliable, but requires additional diffusion steps and a higher inference cost. This trade-off motivates estimating object count at an intermediate diffusion step.

Figure 3: ATHENA steering block. The denoiser is evaluated with the original and control prompts to obtain $\epsilon_{t}$ and $\hat{\epsilon}_{t}$ , which are combined to obtain the latent state $z_{t-1}$ for test-time count control.

Figure 4: ATHENA-Adaptive pipeline. The method estimates object count at an intermediate diffusion step, applies early-stage prompt-based steering, adaptively adjusts the steering strength once based on the observed error direction, and completes generation using standard diffusion.

Level	Description	Example Prompt
L1	Hard object categories	“nine microphones”
L2	+ Scene context	“eight sneakers at the edge of a pond”
L3	+ Distractor objects	“five jars beside a river with a person”
L4	+ Relational constraints	“seven dumbbells lined up along a sidewalk”

Table 1: Structure of the ATHENA dataset with four levels of increasing prompt complexity.

Model	Method	CoCoCount				CoCoCount-E				ATHENA Dataset
Model	Method	Acc ( $\uparrow$ )	MAE ( $\downarrow$ )	RMSE ( $\downarrow$ )	Time ( $\downarrow$ )	Acc ( $\uparrow$ )	MAE ( $\downarrow$ )	RMSE ( $\downarrow$ )	Time ( $\downarrow$ )	Acc ( $\uparrow$ )	MAE ( $\downarrow$ )	RMSE ( $\downarrow$ )	Time ( $\downarrow$ )
FLUX.1-dev	Unsteered	58.4	0.98	1.96	45.8	46.5	1.12	1.93	22.6	39.4	1.78	2.96	45.5
	ATHENA-Static	73.3	0.56	1.35	54.8	58.5	0.84	1.73	27.4	48.9	1.31	2.44	54.7
	ATHENA-Feedback	71.4	0.67	1.72	56.0	58.3	0.98	2.00	29.8	52.2	1.38	2.60	60.8
	ATHENA-Adaptive	70.2	0.65	1.63	64.8	62.2	0.85	1.90	35.1	53.6	1.40	2.72	72.5
SD 3.5 Large	Unsteered	58.4	0.78	1.50	43.9	44.8	0.98	1.56	24.1	38.9	1.55	2.54	43.8

Table 2: Quantitative counting performance across diffusion backbones and datasets. Accuracy (%) is reported as exact-match count accuracy, with the best result for each model–dataset pair shown in bold . MAE and RMSE measure counting error, and Time denotes mean generation time per sample (seconds).

Figure 7: Qualitative results on complex, distractor-rich prompts. Guidance denotes Counting Guidance. Green borders indicate correct counts; red borders indicate incorrect counts. All ATHENA variants preserve scene structure and color consistency while handling relational and multi-object instructions. Compared to prior baselines, ATHENA avoids visual artifacts and background distortion, with adaptive steering yielding the most accurate counts.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: ATHENA-Static pipeline. Starting from the same initial noise, prompt-based steering is applied using a fixed, count-agnostic control prompt from diffusion step $T$ down to a cutoff $t_{\mathrm{steer}}$ . Standard diffusion then proceeds unmodified to completion, yielding improved count fidelity without intermediate feedback or adaptive adjustment.

Figure 9: ATHENA-Feedback pipeline. An intermediate count is estimated at step $t_{\mathrm{est}}$ without steering. If the estimate differs from the target, generation is restarted from the same initial noise using a feedback control prompt, and a single early-stage steering phase is applied before completing diffusion.

Figure 10: ATHENA-Adaptive pipeline. The method estimates object count at an intermediate diffusion step, applies early-stage prompt-based steering, adaptively adjusts the steering strength once based on the observed error direction, and completes generation using standard diffusion.

Dataset	Unique Prompts	Count Range	Total Prompts
CoCoCount (deduplicated)	161	2–10	161
CoCoCount-E	72	2–10	648

Appendix

Table 3: Dataset statistics for CoCoCount and CoCoCount-E.

	Diffusion Backbone
Parameter	FLUX.1-dev			SD 3.5 Large			SDXL
	Static	Feedback	Adaptive	Static	Feedback	Adaptive	Static	Feedback	Adaptive
$t_{\mathrm{est}}$	–	20	20	–	20	20	–	30	30
$t_{\mathrm{steer}}$	10	5	5	10	5	5	10	10	10
$\gamma$	4	4	4	4	4	4	5	5	5

Appendix

Table 4: Hyperparameter settings used for ATHENA across diffusion backbones and steering variants.

Panel of Figure 11 — Figure 11: Accuracy–runtime trade-offs on CoCoCount and CoCoCount-E. Colors denote the diffusion backbone, while marker shapes indicate the method. ATHENA-Adaptive consistently achieves higher accuracy at comparable or lower runtime than prior baselines.

Panel of Figure 12 — Figure 12: Counting accuracy versus target object count across datasets and diffusion backbones. ATHENA-Adaptive consistently maintains higher accuracy as the target count increases, demonstrating improved robustness to increasing counting difficulty.

Figure 13: Additional qualitative results on relational and multi-object prompts. Results are shown for unsteered generation (FLUX) and the three ATHENA variants. Green borders indicate correct object counts, while red borders denote counting errors. ATHENA improves count fidelity across diverse settings (e.g., grouping, circular layouts, and natural scenes) while preserving scene structure and visual coherence, with the adaptive variant yielding the most consistent corrections.

Figure 14: Additional qualitative results on relational and multi-object prompts. Results are shown for unsteered generation (SD 3.5) and the three ATHENA variants. Green borders indicate correct object counts, while red borders denote counting errors. ATHENA improves count fidelity across diverse settings (e.g., grouping, circular layouts, and natural scenes) while preserving scene structure and visual coherence, with the adaptive variant yielding the most consistent corrections.

Figure 15: Additional qualitative results on relational and multi-object prompts with the SDXL backbone. Results include unsteered generation, CountGen, Counting Guidance ( Guidance ), and the three ATHENA variants. Green borders indicate correct object counts, while red borders denote counting errors. Prior baselines frequently fail to enforce correct counts or introduce visual artifacts, whereas ATHENA improves count fidelity while preserving scene structure and visual coherence, with the adaptive variant yielding the most consistent corrections.

T2I Generation

Text-to-image generation creates images from natural-language prompts using generative models, with research focused on fidelity, control, efficiency, and safety.

See all 283 papers

CountLoop: Training-Free High-Instance Image Generation via Iterative Agent Guidance
Aug 18, 2025Anindya Mondal, Sauradip Nag, Ayan Banerjee +3Diffusion Model Guidance Object Counting
University of Surrey · Simon Fraser University · Universitat Autònoma de Barcelona
When Numbers Speak: Aligning Textual Numerals and Visual Instances in Text-to-Video Diffusion Models
Apr 9, 2026Zhengyang Sun, Yu Chen, Xin Zhou +4T2V Generation Video Diffusion Models
Huazhong University of Science and Technology · Afari Intelligent Drive · Zhejiang University
Improving Controllable Generation: Faster Training and Better Performance via $x_0$ -Supervision
Apr 7, 2026Amadou S. Sangare, Adrien Maglo, Mohamed Chaouch +1Controllable Image Generation Conditional Diffusion Models
Université Paris-Saclay, CEA, List, F-91120, Palaiseau, France

Aug 18, 2025J/K move · Enter open · S save

CountLoop: Training-Free High-Instance Image Generation via Iterative Agent Guidance

Anindya Mondal, Sauradip Nag, Ayan Banerjee, Josep Llados+2

University of Surrey · Simon Fraser University · Universitat Autònoma de Barcelona

Diffusion Model Guidance Object Counting Agentic Image Generation Controllable Image Generation

Open paper AskPDF arXiv

Diffusion models excel at photorealistic synthesis but struggle with object count fidelity, especially in high-density settings. We introduce COUNTLOOP, a training-free framework that achieves structured instance and count control through iterative, structured feedback. Our method alternates between synthesis and evaluation: a VLM-based planner generates structured scene layouts, while a VLM-based critic provides explicit feedback on object counts, spatial arrangements, and visual quality to refine the layout iteratively. Instance-driven attention masking and cumulative attention composition further prevent semantic leakage, ensuring clear object separation even in densely occluded scenes. Evaluations on COCO-Count, T2I-CompBench, and two newly introduced high instance benchmarks show that COUNTLOOP reduces counting error by up to 57% and achieves the highest or comparable spatial quality scores across all benchmarks, while maintaining photorealism. Project page is at https://mondalanindya.github.io/CountLoop/.

Abstract

T2I Generation

Explore similar work

CountLoop: Training-Free High-Instance Image Generation via Iterative Agent Guidance

When Numbers Speak: Aligning Textual Numerals and Visual Instances in Text-to-Video Diffusion Models

Improving Controllable Generation: Faster Training and Better Performance via x0x_0x0​-Supervision

Improving Controllable Generation: Faster Training and Better Performance via $x_0$ -Supervision