Test-Time Debiasing with Probabilistic Prompts via Wasserstein Distance in Vision-Language Models

Vision-Language Models (VLMs) inherit social biases from large-scale pretraining data, and these biases can amplify in downstream tasks, leading to systematic performance disparities across sensitive groups. Due to the high training cost and the risk of catastrophic forgetting, recent research has focused more on lightweight \emph{test-time} debiasing, aiming to obtain an ideal fair embedding for each query. However, such point-based corrections are often unstable and become notably weaker in multi-class settings, where group structure cannot be adequately captured by a single point. Therefore, we propose W4D, a distributional debiasing framework that reframes fairness as aligning query embedding distributions to group reference distributions under the Wasserstein distance, which provides a geometry-aware notion of discrepancy beyond mean shifts. To make this alignment practical at test-time, W4D introduces probabilistic prompts that induce controlled distributional perturbations and optimizes a Wasserstein-based objective to reduce cross-group disparity while preserving task-relevant semantics. This distributional perspective improves robustness in multi-class debiasing and yields a stronger fairness--utility trade-off across diverse VLM downstream evaluations. Our code is available at https://github.com/QDRhhhh/W4D.

Vision-language models can connect images and text, but they can also inherit social biases from the data used to train them. For example, when asked to retrieve images for a neutral concept, a model may return results that overrepresent certain gender, race, or skin-tone groups. Existing debiasing methods often correct a single embedding point, which can miss the richer way different groups are distributed in the model’s representation space, especially when there are more than two groups. We propose W4D, a test-time debiasing method that treats fairness as a problem of aligning distributions rather than moving a single point. W4D compares query embeddings with group reference distributions using Wasserstein distance, a way to measure how different two distributions are while considering their geometry. It then uses probabilistic prompts to gently shift query embeddings so that they are more equally related to different sensitive groups, without changing the underlying vision-language model. Experiments on several face datasets show that W4D reduces bias while largely preserving task performance. This makes debiasing more practical for deployed models where retraining is expensive or risky.