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

arXiv: https://arxiv.org/abs/2603.19676 · HTML: https://mlsanity.com/paper/2603.19676/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Image generation, text-to-image and editing, one of 14 directions in Diffusion models.
- Builds on 35 of its 37 references that are in the map.
- First posted on arXiv 20 Mar 2026.

## Placed in

- Diffusion models / [Image generation, text-to-image and editing](https://mlsanity.com/diffusion/image-generation-text-to-image-and-editing/)

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection](https://mlsanity.com/paper/2303.05499/) (2023)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Prompt-to-Prompt Image Editing with Cross Attention Control](https://mlsanity.com/paper/2208.01626/) (2022)
- [Grounded Language-Image Pre-training](https://mlsanity.com/paper/2112.03857/) (2021)
- [FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space](https://mlsanity.com/paper/2506.15742/) (2025)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Attend-and-Excite: Attention-Based Semantic Guidance for Text-to-Image Diffusion Models](https://mlsanity.com/paper/2301.13826/) (2023)
- [YOLO-World: Real-Time Open-Vocabulary Object Detection](https://mlsanity.com/paper/2401.17270/) (2024)
- [Perception Prioritized Training of Diffusion Models](https://mlsanity.com/paper/2204.00227/) (2022)
- [Training Deep Networks with Synthetic Data: Bridging the Reality Gap by Domain Randomization](https://mlsanity.com/paper/1804.06516/) (2018)
- [Localizing Object-level Shape Variations with Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.11306/) (2023)
- [Synthetic Data for Deep Learning](https://mlsanity.com/paper/1909.11512/) (2019)
- [ConceptMix: A Compositional Image Generation Benchmark with Controllable Difficulty](https://mlsanity.com/paper/2408.14339/) (2024)
- [Make It Count: Text-to-Image Generation with an Accurate Number of Objects](https://mlsanity.com/paper/2406.10210/) (2024)
- [Prompting Hard or Hardly Prompting: Prompt Inversion for Text-to-Image Diffusion Models](https://mlsanity.com/paper/2312.12416/) (2023)
- [All Seeds Are Not Equal: Enhancing Compositional Text-to-Image Generation with Reliable Random Seeds](https://mlsanity.com/paper/2411.18810/) (2024)
- [Counting Guidance for High Fidelity Text-to-Image Synthesis](https://mlsanity.com/paper/2306.17567/) (2023)
- [Emergence and Evolution of Interpretable Concepts in Diffusion Models](https://mlsanity.com/paper/2504.15473/) (2025)
- [Cycle Consistency as Reward: Learning Image-Text Alignment without Human Preferences](https://mlsanity.com/paper/2506.02095/) (2025)
- [Detector Guidance for Multi-Object Text-to-Image Generation](https://mlsanity.com/paper/2306.02236/) (2023)
- [MuLan: Multimodal-LLM Agent for Progressive and Interactive Multi-Object Diffusion](https://mlsanity.com/paper/2402.12741/) (2024)
- [YOLO-Count: Differentiable Object Counting for Text-to-Image Generation](https://mlsanity.com/paper/2508.00728/) (2025)

## Abstract (arXiv)

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.