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

2026-03-20

AI Summary

This work addresses the systematic counting bias prevalent in text-to-image diffusion models when generating a specified number of objects. The authors propose a test-time adaptive guidance framework that requires neither architectural modifications nor retraining. By leveraging intermediate features during early denoising steps to estimate object counts in real time, the method applies count-aware noise correction and dynamically adjusts textual prompts to steer the generation process toward improved counting accuracy. This approach achieves, for the first time, model-agnostic and training-free dynamic count control. It significantly outperforms existing methods across multiple benchmarks and a newly curated complex dataset, particularly excelling in high-count scenarios by jointly preserving image quality and counting fidelity, while maintaining a favorable balance between accuracy and efficiency across diverse diffusion backbones.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation
• Machine Learning: Large Multimodal Models (LMMs)

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

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.

Problem

• count fidelity
• diffusion models
• numerical control
• object counting
• text-to-image generation

Innovation

• test-time adaptation
• count fidelity
• diffusion models
• noise correction
• object counting