Diffusion-Enhanced Test-time Adaptation with Text and Image Augmentation

TL;DR

RHFL+ tackles robust federated learning with noisy, model-heterogeneous clients by enabling mutual learning on a shared public dataset and by dynamically refining labels and weighting client contributions. It combines knowledge-distribution alignment on D0 via KL divergences, dynamic label refinement with a SL loss, and enhanced client confidence re-weighting to suppress noisy external feedback. Across CIFAR-10/100 and CelebA, RHFL+ consistently surpasses state-of-the-art heterogeneous FL baselines, especially under higher label-noise rates and non-IID settings. The approach offers a practical path to privacy-preserving, scalable collaboration among diverse client models in real-world noisy environments.

Abstract

Existing test-time prompt tuning (TPT) methods focus on single-modality data, primarily enhancing images and using confidence ratings to filter out inaccurate images. However, while image generation models can produce visually diverse images, single-modality data enhancement techniques still fail to capture the comprehensive knowledge provided by different modalities. Additionally, we note that the performance of TPT-based methods drops significantly when the number of augmented images is limited, which is not unusual given the computational expense of generative augmentation. To address these issues, we introduce IT3A, a novel test-time adaptation method that utilizes a pre-trained generative model for multi-modal augmentation of each test sample from unknown new domains. By combining augmented data from pre-trained vision and language models, we enhance the ability of the model to adapt to unknown new test data. Additionally, to ensure that key semantics are accurately retained when generating various visual and text enhancements, we employ cosine similarity filtering between the logits of the enhanced images and text with the original test data. This process allows us to filter out some spurious augmentation and inadequate combinations. To leverage the diverse enhancements provided by the generation model across different modals, we have replaced prompt tuning with an adapter for greater flexibility in utilizing text templates. Our experiments on the test datasets with distribution shifts and domain gaps show that in a zero-shot setting, IT3A outperforms state-of-the-art test-time prompt tuning methods with a 5.50% increase in accuracy.

Figures (8)

• Figure 1: Illustration of federated learning with noisy and heterogeneous clients, where clients possess heterogeneous local models and noisy datasets with different noise rates.
• Figure 2: Illustration of RHFL+, which performs the heterogeneous FL by aligning the knowledge distributions of individual models on the public dataset §\ref{['sec:hfl']}. Under the condition of noisy clients, DLR is proposed to refine the annotated labels while incorporating SL loss to mitigate overfitting to local noisy data §\ref{['sec:dlrsl']}. As for the noise generated in communication, the weight of noisy and inefficient clients is reduced by dynamically measuring the client confidence §\ref{['sec:eccr']}.
• Figure 3: of four scenarios under 20% symflip noise. We set up four heterogeneous clients, and their setup details are described in §\ref{['sec:expsetup']}. The label distribution of all clients is the same, that is, the client data is independently and identically distributed. Each scenario here means that one client is noisy and the others are clean. The loss curves of all clients are shown when using HFL for federated learning. The x-axis represents the training epoch, and the y-axis represents the corresponding CE loss. Dotted lines represent the learning curves for the noisy clients, while the solid lines denote the learning curve for the clean clients. For example, the leftmost figure represents adding 20% symflip label noise to the client with ResNet10 as the local model, and the other three clients all contain clean private datasets. In the process of using HFL for federated learning, the client with ResNet10 has the largest loss value and the smallest loss drop rate compared to other clean clients.
• Figure 4: The flow chart of RHFL.
• Figure 5: Comparison of RHFL+ under different public datasets proves that HFL does not rely on additional relevant data.