HiLo: A Learning Framework for Generalized Category Discovery Robust to Domain Shifts

Abstract: Generalized Category Discovery (GCD) is a challenging task in which, given a partially labelled dataset, models must categorize all unlabelled instances, regardless of whether they come from labelled categories or from new ones. In this paper, we challenge a remaining assumption in this task: that all images share the same domain. Specifically, we introduce a new task and method to handle GCD when the unlabelled data also contains images from different domains to the labelled set. Our proposed `HiLo' networks extract High-level semantic and Low-level domain features, before minimizing the mutual information between the representations. Our intuition is that the clusterings based on domain information and semantic information should be independent. We further extend our method with a specialized domain augmentation tailored for the GCD task, as well as a curriculum learning approach. Finally, we construct a benchmark from corrupted fine-grained datasets as well as a large-scale evaluation on DomainNet with real-world domain shifts, reimplementing a number of GCD baselines in this setting. We demonstrate that HiLo outperforms SoTA category discovery models by a large margin on all evaluations.

• The paper introduces a novel HiLo framework that disentangles high-level semantic and low-level domain features using vision transformers to address category discovery under domain shifts.
• It employs innovative PatchMix-based contrastive learning and a curriculum sampling strategy to regularize representations and mitigate domain-induced feature collapse.
• Experimental results demonstrate significant accuracy improvements, with gains up to 51.4% over prior methods on datasets with corrupted domains.

HiLo: Robust Generalized Category Discovery under Domain Shifts

Problem Formulation and Motivation

Generalized Category Discovery (GCD) presents an open-world learning challenge, where models must accurately categorize unlabelled data comprising both known and unknown classes. Existing GCD protocols rely on the strong, unrealistic assumption of shared domains between labelled and unlabelled samples. In practice, low-level covariate (domain) shift often emerges, significantly degrading the discriminative capacity of learned representations. This work introduces a formulation for GCD with domain shifts: only labelled samples from a source domain are annotated, while unlabelled data may include both known and unknown categories from mixed, potentially novel, domains. This protocol dramatically increases both practical relevance and difficulty, requiring methods that simultaneously handle semantic novelty and domain disparity.

HiLo Framework: Methodological Advances

The HiLo framework, underpinned by a vision transformer backbone, introduces three interlocking mechanisms to address the dual challenges of semantic novelty and domain shift: high/low-level feature disentanglement, PatchMix-based domain-robust contrastive learning, and a curriculum sampling regime.

Domain-Semantic Feature Disentanglement

HiLo explicitly decomposes representations into high-level (semantic) and low-level (domain) features, exploiting the implicit hierarchical structure of Vision Transformers: early layer activations generally encode covariate (style/domain) cues, while late layers encode semantic object information. Dedicated projection heads extract both feature types, and a mutual information (MI) minimization term between them enforces statistical independence, thereby promoting disentanglement of semantic and domain factors of variation for robust clustering.

Figure 1: PatchMix augmentation and high-/low-level feature design facilitate domain-semantic disentanglement via MI minimization.

The optimal placement ("shallow for domain, deep for semantics") is ablated and confirmed to yield the best performance:

Figure 2: Features from the first and last ViT layers, respectively, maximally benefit domain-semantic disentanglement for GCD under domain shift.

PatchMix-based Contrastive Learning

To mitigate domain-induced feature collapse, HiLo incorporates a modification of PatchMix in the embedding space. By stochastically mixing patch representations of labelled and unlabelled samples (potentially from different domains and classes), HiLo enables robust transfer and regularization under high covariate shift. Semantic contributions from each constituent patch are adaptively reweighted for both contrastive and classification losses, accounting for both attention—via computed weights—and the mixing coefficients.

Figure 3: PatchMix produces mixed patchwise embeddings; semantic components ($\alpha$) adjust the learning signal during contrastive and classification training for domain-agnostic discovery.

Curriculum Sampling

A curriculum sampling procedure assigns increasing weight to unlabelled samples predicted to derive from unknown domains as training progresses. Initially, learning focuses on the seen domain to stabilize semantic clustering; later, it incrementally addresses harder domain gaps. The curriculum is dynamically adapted via semi-supervised k-means over domain features.

Experimental Results

Comprehensive evaluations are conducted on DomainNet—a large-scale dataset with six distinct domains (Real, Painting, Sketch, Quickdraw, Clipart, Infograph)—and the newly constructed SSB-C benchmark, which injects 45 algorithmically generated corruptions across varied types and severities (Gaussian, shot, impulse noise, blur, fog, etc.) into fine-grained datasets (CUB, Scars, FGVC-Aircraft):

Figure 4: SSB-C benchmark summary—each fine-grained dataset is subjected to 9 corruption types $\times$ 5 severity levels, imposing challenging domain shifts.

Strong numerical results are obtained. On DomainNet, HiLo yields absolute accuracy improvements, outperforming prior SoTA GCD methods (e.g., SimGCD, UNO+) by wide margins under all settings (seen and unseen domains/classes). Notably, on CUB-C in SSB-C, HiLo improves over SimGCD by 43.8% on the original domain and by 51.4% on unlabelled, corrupted domains. The PatchMix and MI components are independently ablated, and each is shown to contribute significantly to final accuracy, especially when their synergy is leveraged.

The reweighted semantic attention enforced by MI minimization is further visualized through Transformer attention maps, which show HiLo focusing robustly on the foreground object regions even under strong domain perturbations, unlike SimGCD, where attention remains sensitive to background artifacts:

Figure 5: HiLo attention maps show that semantic focus is stable under large domain shifts.

Ablations confirm that PatchMix regularization, MI minimization, and curriculum sampling are all crucial for handling severe domain variations and semantic novelty.

Additional Analysis

• Applying legacy UDA algorithms (e.g., Mixstyle, EFDM, MCC, NWD) does not close the generalization gap on GCD with domain shifts when used as plug-in augmentations to GCD baselines.
• Zero-shot CLIP with access to an oracle vocabulary fails to surpass HiLo on synthetic domain shifts (CUB-C), suggesting that improved visual encoder pretraining is not a sufficient solution without explicit domain-semantic disentanglement and robust PatchMix augmentation.
• The MI regularization term is particularly beneficial for datasets where low-level statistics are highly variable between train and test domains, or where closely related semantics exist within a single corruption regime:

Figure 6: MI regularization yields clear improvements whenever unlabelled domain styles diverge from the labelled source or when fine-grained semantic separation is required under corruption.

• The learned feature space successfully clusters images by domain and semantics, as confirmed by PCAs over the corresponding branches:

Figure 7: Semantic and domain features are linearly separable: the semantic branch organizes images by class, the domain branch by covariate properties.

Theoretical and Practical Implications

HiLo's architecture is theoretically grounded: mutual information minimization tightens generalization bounds for open-world cluster discovery by ensuring that domain and semantic factors are non-interfering, thereby facilitating accurate category estimation even under high domain uncertainty. PatchMix creates a richer signal by blending statistical properties of all available domains and semantics, regularizing both heads simultaneously for robust discovery.

Practically, this framework offers a reliable path toward autonomous, scalable category discovery in high-variance, web-scale, or field-deployed systems, where annotation for all possible domains is infeasible. It also establishes principled design patterns for future research in open-world learning, suggesting that explicit representation disentanglement, contrastive regularization, and dynamic curriculum adaptation are essential for robust generalization.

Conclusion

HiLo provides a robust learning framework for generalized category discovery under pronounced domain shifts, coupling mutual information-driven disentanglement with PatchMix-based contrastive learning and adaptive curriculum sampling. Extensive benchmarks demonstrate large gains across both coarse- and fine-grained settings, in presence of substantial covariate and semantic novelty. These results establish HiLo as a strong reference architecture for open-world clustering, motivating further theoretical and empirical investigation into disentangled, domain-robust representation learning for real-world AI systems.

---

Reference:

"HiLo: A Learning Framework for Generalized Category Discovery Robust to Domain Shifts" (2408.04591)