What's in a Latent? Leveraging Diffusion Latent Space for Domain Generalization

2025-03-09

AI Summary

This work addresses domain generalization (DG), aiming to enhance model generalization to unseen test domains. Existing DG methods either rely on domain labels or neglect the intrinsic domain structure in feature space. We make the novel observation that pretrained diffusion models’ latent spaces inherently encode fine-grained, domain-specific structural information. Leveraging this insight, we propose an unsupervised pseudo-domain discovery framework that automatically identifies pseudo-domains via latent-space clustering. We further design a plug-and-play pseudo-domain-aware classifier that jointly achieves feature disentanglement and domain-adaptive fusion. Evaluated on five standard DG benchmarks, our method improves upon the ERM baseline by up to 4.1% in accuracy and outperforms most state-of-the-art methods requiring domain labels—despite operating entirely without them.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Diffusion Models for Vision
• Application Domains: Misinformation & Fake News

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Domain Generalization aims to develop models that can generalize to novel and unseen data distributions. In this work, we study how model architectures and pre-training objectives impact feature richness and propose a method to effectively leverage them for domain generalization. Specifically, given a pre-trained feature space, we first discover latent domain structures, referred to as pseudo-domains, that capture domain-specific variations in an unsupervised manner. Next, we augment existing classifiers with these complementary pseudo-domain representations making them more amenable to diverse unseen test domains. We analyze how different pre-training feature spaces differ in the domain-specific variances they capture. Our empirical studies reveal that features from diffusion models excel at separating domains in the absence of explicit domain labels and capture nuanced domain-specific information. On 5 datasets, we show that our very simple framework improves generalization to unseen domains by a maximum test accuracy improvement of over 4% compared to the standard baseline Empirical Risk Minimization (ERM). Crucially, our method outperforms most algorithms that access domain labels during training.

Problem

• Develop models for generalization to unseen data distributions.
• Leverage pre-trained feature spaces for domain generalization.
• Improve test accuracy in unseen domains using diffusion models.

Innovation

• Leverages diffusion models for domain generalization
• Uses unsupervised pseudo-domains to capture variations
• Augments classifiers with pseudo-domain representations