Joint Enhancement and Classification using Coupled Diffusion Models of Signals and Logits

Robust classification in noisy environments remains a fundamental challenge in machine learning. Standard approaches typically treat signal enhancement and classification as separate, sequential stages: first enhancing the signal and then applying a classifier. This approach fails to leverage the semantic information in the classifier's output during denoising. In this work, we propose a general, domain-agnostic framework that integrates two interacting diffusion models: one operating on the input signal and the other on the classifier's output logits, without requiring any retraining or fine-tuning of the classifier. This coupled formulation enables mutual guidance, where the enhancing signal refines the class estimation and, conversely, the evolving class logits guide the signal reconstruction towards discriminative regions of the manifold. We introduce three strategies to effectively model the joint distribution of the input and the logit. We evaluated our joint enhancement method for image classification and automatic speech recognition. The proposed framework surpasses traditional sequential enhancement baselines, delivering robust and flexible improvements in classification accuracy under diverse noise conditions.

Modern AI systems often struggle when their input is noisy, for example, when an image is blurry or a voice recording is distorted. A common solution is to first “clean” the data and then apply a classifier. However, this two-step process can fail because the cleaning step does not know what information is important for the final decision. In this work, we propose a different approach where these two steps work together instead of separately. Our method simultaneously improves the input data and refines the model’s prediction, allowing each process to guide the other. Intuitively, the system not only cleans the data, but does so in a way that makes it easier to recognize what it represents. We show that this interaction leads to more accurate and robust predictions across different tasks, including image recognition and speech understanding, even under challenging noise conditions. Importantly, our method can be applied to existing AI systems without retraining them, making it a flexible tool for improving reliability in real-world applications.